use get_free_port_in_ci

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>

Author: reasonsolo

tests/integration/defs/disaggregated/test_auto_scaling.py

@@ -12,9 +12,9 @@
 import pytest
 import requests
 import yaml
+from defs.common import get_free_port_in_ci as get_free_port
 from defs.conftest import llm_models_root
 
-from tensorrt_llm._utils import get_free_port
 from tensorrt_llm.logger import logger
 
 HEARTBEAT_INTERVAL = 1
@@ -454,7 +454,7 @@ async def test_worker_restart(model_name, disagg_server_config, worker_config,
                                       port=disagg_port)
         print(response)
         # kill gen1, the request should fail
-        terminate(gen_worker1, release_port=False)
+        terminate(gen_worker1, release_port=True)
         await asyncio.sleep(CHECK_STATUS_INTERVAL)
         verify_cluster_info(False, 1, 0, port=disagg_port)
         with pytest.raises(Exception):
@@ -480,7 +480,7 @@ async def test_worker_restart(model_name, disagg_server_config, worker_config,
         assert len(response.choices[0].text) >= 1
 
         # kill ctx1, the request should fail
-        terminate(ctx_worker1, release_port=False)
+        terminate(ctx_worker1, release_port=True)
         await asyncio.sleep(CHECK_STATUS_INTERVAL)
         verify_cluster_info(False, 0, 1, port=disagg_port)
         with pytest.raises(Exception):
@@ -500,16 +500,16 @@ async def test_worker_restart(model_name, disagg_server_config, worker_config,
         assert len(response.choices[0].text) >= 1
 
         # start ctx1 and gen1 again, we have 2 ctxs and 2 gens now
-        await wait_for_port_released(ctx_worker1.port)
-        await wait_for_port_released(gen_worker1.port)
         ctx_worker1 = run_ctx_worker(model_name,
                                      worker_config,
                                      work_dir,
-                                     port=ctx_worker1.port)
+                                     port=0,
+                                     device=0)
         gen_worker1 = run_gen_worker(model_name,
                                      worker_config,
                                      work_dir,
-                                     port=gen_worker1.port)
+                                     port=0,
+                                     device=1)
         await wait_for_worker_ready(ctx_worker1.port)
         await wait_for_worker_ready(gen_worker1.port)
         await asyncio.sleep(CHECK_STATUS_INTERVAL)
@@ -556,6 +556,7 @@ async def test_disagg_server_restart(model_name, disagg_server_config,
         terminate(disagg_server)
         # wait for the port to be released, so we can rebind the new process to the same port
         await wait_for_port_released(disagg_port)
+        await asyncio.sleep(CHECK_STATUS_INTERVAL)
 
         with pytest.raises(requests.exceptions.RequestException):
             verify_cluster_info(False,

tests/integration/test_lists/qa/llm_function_core_sanity.txt

@@ -251,11 +251,6 @@ disaggregated/test_workers.py::test_workers_conditional_disaggregation[TinyLlama
 disaggregated/test_workers.py::test_workers_kv_cache_aware_router_eviction[TinyLlama-1.1B-Chat-v1.0]
 disaggregated/test_workers.py::test_workers_kv_cache_aware_router[TinyLlama-1.1B-Chat-v1.0]
 disaggregated/test_workers.py::test_workers_kv_cache_events[TinyLlama-1.1B-Chat-v1.0]
-disaggregated/test_auto_scaling.py::test_service_discovery[etcd-load_balancing]
-disaggregated/test_auto_scaling.py::test_service_discovery[http-load_balancing]
-disaggregated/test_auto_scaling.py::test_worker_restart[etcd-round_robin]
-disaggregated/test_auto_scaling.py::test_worker_restart[http-round_robin]
-disaggregated/test_auto_scaling.py::test_worker_restart[http-kv_cache_aware]
 test_e2e.py::test_openai_chat_harmony
 test_e2e.py::test_openai_consistent_chat
 test_e2e.py::test_openai_multi_chat_example

tests/integration/test_lists/waives.txt

@@ -454,7 +454,6 @@ accuracy/test_cli_flow.py::TestPhi3Small128kInstruct::test_auto_dtype SKIP (http
 accuracy/test_cli_flow.py::TestPhi3_5MiniInstruct::test_auto_dtype SKIP (https://nvbugs/5744293)
 unittest/_torch/auto_deploy/unit/singlegpu/models/test_llama4_vlm_patch.py::test_build_run_llama4_vlm SKIP (https://nvbugs/5747878)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_online_eplb[fp8kv=True-moe_backend=TRTLLM] SKIP (https://nvbugs/5740377)
-disaggregated/test_auto_scaling.py::test_minimal_instances[etcd-round_robin] SKIP (https://nvbugs/5748564)
 unittest/_torch/attention/test_trtllm_flashinfer_symbol_collision.py::test_flashinfer_fused_moe_matches_torch_moe SKIP (https://nvbugs/5752521)
 cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[gpt-2proc-mpi_kvcache-90] SKIP (https://nvbugs/5755941)
 accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_guided_decoding[xgrammar-mtp_nextn=0] SKIP (https://nvbugs/5748600)