[None][test] Organize multi-GPU tests into gpu2/gpu4 classes with class-scoped sessions

Group test_llm_multi_gpu_pytorch.py into TestLlmMultiGpu2gpu (@gpu2) and
TestLlmMultiGpu4gpu (@gpu4), sharing a small _MultiGpuLlmTests base whose
class-scoped mpi_session/mpi_kwargs/shared_llm fixtures build one MpiPoolSession
of self.n_gpus workers. The session is created lazily and torn down at class
end, so the 2-GPU and 4-GPU pools no longer coexist (as they did with the
previous module-scoped fixtures). Update the two explicit nodeid references for
test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 (l0_dgx_h100.yml,
waives.txt) to include the class; -m gpu2/gpu4 and -k filters are unaffected.

Signed-off-by: qgai <qgai@nvidia.com>

Author: sunnyqgg

tests/integration/test_lists/test-db/l0_dgx_h100.yml

@@ -15,7 +15,7 @@ l0_dgx_h100:
       auto_trigger: others
       orchestrator: mpi
   tests:
-  - unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1
+  - unittest/llmapi/test_llm_multi_gpu_pytorch.py::TestLlmMultiGpu2gpu::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1
   - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu2" -k "not test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1"
   - unittest/llmapi/test_additional_model_outputs.py -m "gpu2"
   - unittest/_torch/multi_gpu -m "not post_merge" TIMEOUT (90)

tests/integration/test_lists/waives.txt

@@ -488,7 +488,7 @@ unittest/executor/test_rpc.py::TestRpcCorrectness::test_incremental_task_async S
 unittest/executor/test_rpc_proxy.py SKIP (https://nvbugs/5605741)
 unittest/executor/test_rpc_worker.py SKIP (https://nvbugs/5605741)
 unittest/llmapi/test_llm_multi_gpu.py -m "gpu4 and part0" SKIP (https://nvbugs/5348958)
-unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 SKIP (https://nvbugs/6109745)
+unittest/llmapi/test_llm_multi_gpu_pytorch.py::TestLlmMultiGpu2gpu::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 SKIP (https://nvbugs/6109745)
 unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[None] SKIP (https://nvbugs/6162504)
 unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[cuda_graph_config0] SKIP (https://nvbugs/6162504)
 unittest/llmapi/test_memory_profiling.py::test_profile_kvcache SKIP (https://nvbugs/5580781)

tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py

@@ -27,110 +27,6 @@
 global_kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.4)
 
 
-@pytest.fixture(scope="module")
-def shared_mpi_session_2gpu():
-    yield from _shared_mpi_session(2)
-
-
-@pytest.fixture(scope="module")
-def shared_mpi_session_4gpu():
-    yield from _shared_mpi_session(4)
-
-
-@pytest.fixture(scope="module")
-def shared_llm_2gpu(shared_mpi_session_2gpu):
-    # Factory for tests that construct the LLM directly (see RPC tests below).
-    return _make_shared_llm(shared_mpi_session_2gpu)
-
-
-@pytest.fixture(scope="module")
-def mpi_kwargs_2gpu(shared_mpi_session_2gpu):
-    # Ready-to-spread LLM kwargs that inject the shared 2-GPU session (or {} when
-    # MPI is disabled). Harness-based tests just spread this; they never touch the
-    # raw session object.
-    return _mpi_session_kwargs(shared_mpi_session_2gpu)
-
-
-@pytest.fixture(scope="module")
-def mpi_kwargs_4gpu(shared_mpi_session_4gpu):
-    return _mpi_session_kwargs(shared_mpi_session_4gpu)
-
-
-@pytest.mark.gpu2
-def test_llm_capture_request_error(mpi_kwargs_2gpu):
-    _test_llm_capture_request_error(pytorch_backend=True,
-                                    tp_size=2,
-                                    **mpi_kwargs_2gpu)
-
-
-@pytest.mark.gpu4
-def test_tinyllama_logits_processor_tp2pp2(mpi_kwargs_4gpu):
-    tinyllama_logits_processor_test_harness(backend="pytorch",
-                                            tensor_parallel_size=2,
-                                            pipeline_parallel_size=2,
-                                            **mpi_kwargs_4gpu)
-
-
-@pytest.mark.gpu2
-@pytest.mark.part0
-@pytest.mark.parametrize("tp_size, pp_size", [(1, 2), (2, 1)])
-def test_tinyllama_logits_processor_2gpu(tp_size: int, pp_size: int,
-                                         mpi_kwargs_2gpu):
-    tinyllama_logits_processor_test_harness(backend="pytorch",
-                                            tensor_parallel_size=tp_size,
-                                            pipeline_parallel_size=pp_size,
-                                            **mpi_kwargs_2gpu)
-
-
-@pytest.mark.gpu2
-def test_llama_7b_lora_tp2(mpi_kwargs_2gpu):
-    llama_7b_lora_from_dir_test_harness(tensor_parallel_size=2,
-                                        kv_cache_config=global_kv_cache_config,
-                                        **mpi_kwargs_2gpu)
-
-
-@pytest.mark.gpu4
-@skip_ray  # https://nvbugs/5682551
-@test_lora_with_and_without_cuda_graph
-def test_llama_7b_multi_lora_tp4(cuda_graph_config, mpi_kwargs_4gpu):
-    # For LoRA checkpoints without finetuned embedding and lm_head, we can either:
-    # (1) specify lora_target_modules, or
-    # (2) provide a lora_dir to infer the lora_target_modules.
-    lora_config = LoraConfig(lora_target_modules=['attn_q', 'attn_k', 'attn_v'],
-                             max_lora_rank=8,
-                             max_loras=1,
-                             max_cpu_loras=8)
-    check_llama_7b_multi_lora_from_request_test_harness(
-        LLM,
-        lora_config=lora_config,
-        tensor_parallel_size=4,
-        kv_cache_config=global_kv_cache_config,
-        cuda_graph_config=cuda_graph_config,
-        **mpi_kwargs_4gpu)
-
-
-@skip_ray  # https://nvbugs/5727075
-@pytest.mark.gpu2
-@test_lora_with_and_without_cuda_graph
-def test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1(
-        cuda_graph_config) -> None:
-    check_phi3_lora_fused_modules_output_tp2_identical_to_tp1(
-        LLM, cuda_graph_config=cuda_graph_config)
-
-
-@skip_ray
-@pytest.mark.gpu2
-def test_llm_rpc_tp2(shared_llm_2gpu):
-    _run_llm_rpc_tp2(shared_llm_2gpu)
-
-
-@skip_ray
-@pytest.mark.gpu2
-@pytest.mark.asyncio
-async def test_llm_rpc_streaming_tp2(shared_llm_2gpu):
-    await _run_llm_rpc_streaming_tp2(shared_llm_2gpu)
-
-
 def _run_llm_rpc_tp2(make_llm=LLM):
     with make_llm(model=llama_model_path,
                   kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4),
@@ -160,91 +56,178 @@ async def _run_llm_rpc_streaming_tp2(make_llm=LLM):
             print(f"get result: {output}")
 
 
-@skip_ray
-@pytest.mark.gpu2
-@pytest.mark.parametrize(
-    "prompt_logprobs, logprobs, return_context_logits, return_generation_logits",
-    [
-        (None, 1, False,
-         False),  # generation logprobs only (top-1, PyTorch limit)
-    ])
-def test_llm_return_logprobs_streaming_tp2(prompt_logprobs, logprobs,
-                                           return_context_logits,
-                                           return_generation_logits,
-                                           mpi_kwargs_2gpu):
-    llm_return_logprobs_test_harness(prompt_logprobs,
-                                     logprobs,
-                                     return_context_logits,
-                                     return_generation_logits,
-                                     streaming=True,
-                                     backend="pytorch",
-                                     tp_size=2,
-                                     **mpi_kwargs_2gpu)
-
-
-@skip_ray
-@pytest.mark.gpu2
-@pytest.mark.parametrize(
-    "return_context_logits, enable_chunked_prefill, enable_iter_req_stats",
-    [
-        (False, False, True),
-        (False, True, True),
-    ],
-)
-def test_llm_get_stats_pp2(return_context_logits, enable_chunked_prefill,
-                           enable_iter_req_stats, mpi_kwargs_2gpu):
-    llm_get_stats_test_harness(
-        tp_size=1,
-        pp_size=2,
-        return_context_logits=return_context_logits,
-        pytorch_backend=True,
-        enable_chunked_prefill=enable_chunked_prefill,
-        enable_iter_req_stats=enable_iter_req_stats,
-        **mpi_kwargs_2gpu,
-    )
+class _MultiGpuLlmTests:
+    """Base for the multi-GPU LLM API test classes.
 
+    Subclasses set ``n_gpus``; the class-scoped fixtures below build one shared
+    MpiPoolSession (of that many workers) reused across the subclass's tests and
+    torn down when the subclass finishes -- so the 2-GPU and 4-GPU pools never
+    coexist. Fixtures are lazy, so tests that opt out (e.g. constructing a bare
+    ``LLM``) simply don't request them and pay no session cost.
+    """
+    n_gpus: int
 
-@skip_ray
-@pytest.mark.gpu4
-@pytest.mark.parametrize(
-    "return_context_logits, enable_chunked_prefill, enable_iter_req_stats",
-    [
-        (False, False, True),
-        (False, True, True),
-    ],
-)
-def test_llm_get_stats_pp4(return_context_logits, enable_chunked_prefill,
-                           enable_iter_req_stats, mpi_kwargs_4gpu):
-    llm_get_stats_test_harness(
-        tp_size=1,
-        pp_size=4,
-        return_context_logits=return_context_logits,
-        pytorch_backend=True,
-        enable_chunked_prefill=enable_chunked_prefill,
-        enable_iter_req_stats=enable_iter_req_stats,
-        **mpi_kwargs_4gpu,
-    )
+    @pytest.fixture(scope="class")
+    def mpi_session(self):
+        yield from _shared_mpi_session(self.n_gpus)
 
+    @pytest.fixture(scope="class")
+    def mpi_kwargs(self, mpi_session):
+        return _mpi_session_kwargs(mpi_session)
 
-@skip_ray
-@pytest.mark.gpu2
-def test_llm_get_stats_tp2(mpi_kwargs_2gpu):
-    llm_get_stats_test_harness(tp_size=2,
-                               pytorch_backend=True,
-                               **mpi_kwargs_2gpu)
+    @pytest.fixture(scope="class")
+    def shared_llm(self, mpi_session):
+        return _make_shared_llm(mpi_session)
 
 
-@skip_ray
 @pytest.mark.gpu2
-def test_llm_get_stats_async_tp2(mpi_kwargs_2gpu):
-    llm_get_stats_async_test_harness(tp_size=2,
-                                     pytorch_backend=True,
-                                     **mpi_kwargs_2gpu)
+class TestLlmMultiGpu2gpu(_MultiGpuLlmTests):
+    """2-GPU multi-GPU LLM API tests (shared session from _MultiGpuLlmTests)."""
+    n_gpus = 2
+
+    def test_llm_capture_request_error(self, mpi_kwargs):
+        _test_llm_capture_request_error(pytorch_backend=True,
+                                        tp_size=2,
+                                        **mpi_kwargs)
+
+    @pytest.mark.part0
+    @pytest.mark.parametrize("tp_size, pp_size", [(1, 2), (2, 1)])
+    def test_tinyllama_logits_processor_2gpu(self, tp_size, pp_size,
+                                             mpi_kwargs):
+        tinyllama_logits_processor_test_harness(backend="pytorch",
+                                                tensor_parallel_size=tp_size,
+                                                pipeline_parallel_size=pp_size,
+                                                **mpi_kwargs)
+
+    def test_llama_7b_lora_tp2(self, mpi_kwargs):
+        llama_7b_lora_from_dir_test_harness(
+            tensor_parallel_size=2,
+            kv_cache_config=global_kv_cache_config,
+            **mpi_kwargs)
+
+    @skip_ray  # https://nvbugs/5727075
+    @test_lora_with_and_without_cuda_graph
+    def test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1(
+            self, cuda_graph_config) -> None:
+        check_phi3_lora_fused_modules_output_tp2_identical_to_tp1(
+            LLM, cuda_graph_config=cuda_graph_config)
+
+    @skip_ray
+    def test_llm_rpc_tp2(self, shared_llm):
+        _run_llm_rpc_tp2(shared_llm)
+
+    @skip_ray
+    @pytest.mark.asyncio
+    async def test_llm_rpc_streaming_tp2(self, shared_llm):
+        await _run_llm_rpc_streaming_tp2(shared_llm)
+
+    @skip_ray
+    @pytest.mark.parametrize(
+        "prompt_logprobs, logprobs, return_context_logits, return_generation_logits",
+        [
+            (None, 1, False,
+             False),  # generation logprobs only (top-1, PyTorch limit)
+        ])
+    def test_llm_return_logprobs_streaming_tp2(self, prompt_logprobs, logprobs,
+                                               return_context_logits,
+                                               return_generation_logits,
+                                               mpi_kwargs):
+        llm_return_logprobs_test_harness(prompt_logprobs,
+                                         logprobs,
+                                         return_context_logits,
+                                         return_generation_logits,
+                                         streaming=True,
+                                         backend="pytorch",
+                                         tp_size=2,
+                                         **mpi_kwargs)
+
+    @skip_ray
+    @pytest.mark.parametrize(
+        "return_context_logits, enable_chunked_prefill, enable_iter_req_stats",
+        [
+            (False, False, True),
+            (False, True, True),
+        ],
+    )
+    def test_llm_get_stats_pp2(self, return_context_logits,
+                               enable_chunked_prefill, enable_iter_req_stats,
+                               mpi_kwargs):
+        llm_get_stats_test_harness(
+            tp_size=1,
+            pp_size=2,
+            return_context_logits=return_context_logits,
+            pytorch_backend=True,
+            enable_chunked_prefill=enable_chunked_prefill,
+            enable_iter_req_stats=enable_iter_req_stats,
+            **mpi_kwargs,
+        )
+
+    @skip_ray
+    def test_llm_get_stats_tp2(self, mpi_kwargs):
+        llm_get_stats_test_harness(tp_size=2,
+                                   pytorch_backend=True,
+                                   **mpi_kwargs)
+
+    @skip_ray
+    def test_llm_get_stats_async_tp2(self, mpi_kwargs):
+        llm_get_stats_async_test_harness(tp_size=2,
+                                         pytorch_backend=True,
+                                         **mpi_kwargs)
+
+    @skip_ray
+    def test_llm_get_stats_async_pp2(self, mpi_kwargs):
+        llm_get_stats_async_test_harness(pp_size=2,
+                                         pytorch_backend=True,
+                                         **mpi_kwargs)
 
 
-@skip_ray
-@pytest.mark.gpu2
-def test_llm_get_stats_async_pp2(mpi_kwargs_2gpu):
-    llm_get_stats_async_test_harness(pp_size=2,
-                                     pytorch_backend=True,
-                                     **mpi_kwargs_2gpu)
+@pytest.mark.gpu4
+class TestLlmMultiGpu4gpu(_MultiGpuLlmTests):
+    """4-GPU multi-GPU LLM API tests (shared session from _MultiGpuLlmTests)."""
+    n_gpus = 4
+
+    def test_tinyllama_logits_processor_tp2pp2(self, mpi_kwargs):
+        tinyllama_logits_processor_test_harness(backend="pytorch",
+                                                tensor_parallel_size=2,
+                                                pipeline_parallel_size=2,
+                                                **mpi_kwargs)
+
+    @skip_ray  # https://nvbugs/5682551
+    @test_lora_with_and_without_cuda_graph
+    def test_llama_7b_multi_lora_tp4(self, cuda_graph_config, mpi_kwargs):
+        # For LoRA checkpoints without finetuned embedding and lm_head, we can
+        # either: (1) specify lora_target_modules, or (2) provide a lora_dir to
+        # infer the lora_target_modules.
+        lora_config = LoraConfig(
+            lora_target_modules=['attn_q', 'attn_k', 'attn_v'],
+            max_lora_rank=8,
+            max_loras=1,
+            max_cpu_loras=8)
+        check_llama_7b_multi_lora_from_request_test_harness(
+            LLM,
+            lora_config=lora_config,
+            tensor_parallel_size=4,
+            kv_cache_config=global_kv_cache_config,
+            cuda_graph_config=cuda_graph_config,
+            **mpi_kwargs)
+
+    @skip_ray
+    @pytest.mark.parametrize(
+        "return_context_logits, enable_chunked_prefill, enable_iter_req_stats",
+        [
+            (False, False, True),
+            (False, True, True),
+        ],
+    )
+    def test_llm_get_stats_pp4(self, return_context_logits,
+                               enable_chunked_prefill, enable_iter_req_stats,
+                               mpi_kwargs):
+        llm_get_stats_test_harness(
+            tp_size=1,
+            pp_size=4,
+            return_context_logits=return_context_logits,
+            pytorch_backend=True,
+            enable_chunked_prefill=enable_chunked_prefill,
+            enable_iter_req_stats=enable_iter_req_stats,
+            **mpi_kwargs,
+        )