fix: Fix disaggregated load balance test (NVIDIA#3689)

Signed-off-by: Iman Tabrizian <[email protected]>

Author: Tabrizian

tests/integration/defs/disaggregated/test_disaggregated.py

@@ -93,6 +93,8 @@ def get_test_config(test_desc, example_dir, test_root):
         (2,
          f"{test_configs_root}/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp_attention_dp_overlap.yaml"
          ),
+        "load_balance":
+        (4, f"{test_configs_root}/disagg_config_load_balance.yaml"),
     }
 
     if test_desc not in config_map:
@@ -345,11 +347,10 @@ def test_disaggregated_load_balance(disaggregated_test_root, llm_venv,
             os.makedirs(os.path.dirname(dst), exist_ok=True)
             os.symlink(src, dst, target_is_directory=True)
 
-    cmd = f"bash {disaggregated_test_root}/sanity_check.sh {disaggregated_example_root} load_balance"
-    check_call(cmd,
-               shell=True,
-               env=llm_venv._new_env,
-               cwd=llm_venv.get_working_directory())
+    run_disaggregated_test(disaggregated_example_root,
+                           "load_balance",
+                           env=llm_venv._new_env,
+                           cwd=llm_venv.get_working_directory())
 
 
 @pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-fp8'],

tests/integration/test_lists/waives.txt

@@ -456,6 +456,5 @@ examples/test_medusa.py::test_llama_medusa_1gpu[llama-3.1-8b] SKIP (https://nvbu
 examples/test_phi.py::test_llm_phi_quantization_1gpu[Phi-4-mini-instruct-fp8-bfloat16] SKIP (https://nvbugspro.nvidia.com/bug/5226339)
 examples/test_multimodal.py::test_llm_multimodal_general[neva-22b-pp:1-tp:1-bfloat16-bs:1-cpp_e2e:True-nb:1] SKIP (https://nvbugs/5227342)
 disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_ucx[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/5201168)
-disaggregated/test_disaggregated.py::test_disaggregated_load_balance[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/5232631)
 perf/test_perf.py::test_perf[gpt_350m-cppmanager-plugin-float16-bs:32-input_output_len:60,20] SKIP (https://nvbugs/5228840)
 perf/test_perf.py::test_perf[gpt_350m-cppmanager-static_batching-plugin-float16-bs:32-input_output_len:60,20] SKIP (https://nvbugs/5228840)