From 4f759d8de30424fc5ade232093cf3d78cc2670fe Mon Sep 17 00:00:00 2001
From: Dom Brown <3886319+DomBrown@users.noreply.github.com>
Date: Mon, 19 May 2025 16:10:15 +0000
Subject: [PATCH] test: Split test_simple into mpi_utils and cache transceiver
 tests for DGX

Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
---
 tests/integration/defs/cpp/test_multi_gpu.py  | 79 ++++++++-----------
 .../test_lists/test-db/l0_dgx_h100.yml        |  6 +-
 2 files changed, 39 insertions(+), 46 deletions(-)

diff --git a/tests/integration/defs/cpp/test_multi_gpu.py b/tests/integration/defs/cpp/test_multi_gpu.py
index be15906cae90..c6e5234a2c30 100644
--- a/tests/integration/defs/cpp/test_multi_gpu.py
+++ b/tests/integration/defs/cpp/test_multi_gpu.py
@@ -44,10 +44,11 @@ def get_multi_gpu_env(kv_cache_type=KVCacheType.NONE, llama_multi_gpu=False):
     return env
 
 
-def run_simple_multi_gpu_tests(build_dir: _pl.Path, timeout=1500):
+def run_mpi_utils_tests(build_dir, timeout=300):
+
     tests_dir = build_dir / "tests"
-    cpp_env = {**_os.environ}
-    # Utils tests
+    mgpu_env = get_multi_gpu_env()
+
     mpi_utils_test = [
         "mpirun",
         "-n",
@@ -55,56 +56,31 @@ def run_simple_multi_gpu_tests(build_dir: _pl.Path, timeout=1500):
         "--allow-run-as-root",
         "mpiUtilsTest",
     ]
-    _cpp.run_command(mpi_utils_test, cwd=tests_dir, env=cpp_env, timeout=300)
+    _cpp.run_command(mpi_utils_test,
+                     cwd=tests_dir,
+                     env=mgpu_env,
+                     timeout=timeout)
 
-    # Cache transceiver MPI tests
-    new_env = get_multi_gpu_env(kv_cache_type=KVCacheType.MPI)
 
-    cache_trans_test = [
-        "mpirun",
-        "-n",
-        "2",
-        "--allow-run-as-root",
-        "batch_manager/cacheTransceiverTest",
-    ]
-    _cpp.run_command(cache_trans_test, cwd=tests_dir, env=new_env, timeout=300)
-
-    cache_trans_test_8_proc = [
-        "mpirun",
-        "-n",
-        "8",
-        "--allow-run-as-root",
-        "batch_manager/cacheTransceiverTest",
-    ]
-    _cpp.run_command(cache_trans_test_8_proc,
-                     cwd=tests_dir,
-                     env=new_env,
-                     timeout=600)
+def run_cache_transceiver_tests(build_dir: _pl.Path,
+                                nprocs=2,
+                                kv_cache_type=KVCacheType.MPI,
+                                timeout=600):
 
-    # Cache transceiver tests with UCX
-    new_env = get_multi_gpu_env(kv_cache_type=KVCacheType.UCX)
+    tests_dir = build_dir / "tests"
+    mgpu_env = get_multi_gpu_env(kv_cache_type=kv_cache_type)
 
     cache_trans_test = [
         "mpirun",
         "-n",
-        "2",
-        "--allow-run-as-root",
-        "batch_manager/cacheTransceiverTest",
-    ]
-    _cpp.run_command(cache_trans_test, cwd=tests_dir, env=new_env, timeout=300)
-
-    # Cache transceiver tests
-    cache_trans_test_8_proc = [
-        "mpirun",
-        "-n",
-        "8",
+        f"{nprocs}",
         "--allow-run-as-root",
         "batch_manager/cacheTransceiverTest",
     ]
-    _cpp.run_command(cache_trans_test_8_proc,
+    _cpp.run_command(cache_trans_test,
                      cwd=tests_dir,
-                     env=new_env,
-                     timeout=600)
+                     env=mgpu_env,
+                     timeout=timeout)
 
 
 def run_llama_executor_leader_tests(build_dir: _pl.Path, timeout=1500):
@@ -447,12 +423,25 @@ def multi_gpu_model(request, prepare_model_multi_gpu):
 
 @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"],
                          indirect=True)
-def test_simple(build_google_tests, build_dir):
+def test_mpi_utils(build_google_tests, build_dir):
+
+    if platform.system() != "Windows":
+        run_mpi_utils_tests(build_dir, timeout=300)
+
+
+@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"],
+                         indirect=True)
+@pytest.mark.parametrize("kvcache_type", [KVCacheType.MPI, KVCacheType.UCX],
+                         ids=["mpi_kvcache", "ucx_kvcache"])
+@pytest.mark.parametrize("nprocs", [2, 8], ids=["2proc", "8proc"])
+def test_cache_transceiver(build_google_tests, nprocs, kvcache_type, build_dir):
 
     if platform.system() != "Windows":
 
-        run_simple_multi_gpu_tests(build_dir=build_dir,
-                                   timeout=_cpp.default_test_timeout)
+        run_cache_transceiver_tests(build_dir=build_dir,
+                                    nprocs=nprocs,
+                                    kv_cache_type=kvcache_type,
+                                    timeout=600)
 
 
 @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"],
diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml
index 58b7464fd000..c24f807adbd9 100644
--- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml
+++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml
@@ -125,7 +125,11 @@ l0_dgx_h100:
       auto_trigger: others
   tests:
   # ------------- CPP tests ---------------
-  - cpp/test_multi_gpu.py::test_simple[90]
+  - cpp/test_multi_gpu.py::test_mpi_utils[90]
+  - cpp/test_multi_gpu.py::test_cache_transceiver[2proc-mpi_kvcache-90]
+  - cpp/test_multi_gpu.py::test_cache_transceiver[2proc-ucx_kvcache-90]
+  - cpp/test_multi_gpu.py::test_cache_transceiver[8proc-mpi_kvcache-90]
+  - cpp/test_multi_gpu.py::test_cache_transceiver[8proc-ucx_kvcache-90]
   - cpp/test_multi_gpu.py::test_enc_dec[t5-90]
   - cpp/test_multi_gpu.py::test_llama_executor[llama-orchestrator-90]
   - cpp/test_multi_gpu.py::test_llama_executor[llama-leader-90]