Add multi instance test

Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com>

Author: shuyixiong

examples/ray_orchestrator/rl_perf_repro.py

@@ -19,7 +19,7 @@
 
 
 @ray.remote
-class trtllm_instance:
+class TRTLLMInstance:
     def __init__(self, async_llm_kwargs: dict, sampling_kwargs: dict):
         self.async_llm_kwargs = async_llm_kwargs
         self.sampling_kwargs = sampling_kwargs
@@ -62,7 +62,7 @@ async def init_llm(self):
         )
 
     async def generate(self, prompt: list[int]):
-        """Generate for a single prompt"""
+        """Generate for a single prompt."""
         outputs = await self.llm.generate_async(inputs=prompt, sampling_params=self.sampling_params)
         token_ids = outputs.outputs[0].token_ids
         log_probs = None
@@ -111,7 +111,7 @@ async def setup_rl_llm(args):
         llm_instances = []
         for i in range(num_instances):
             llm_instances.append(
-                trtllm_instance.options(
+                TRTLLMInstance.options(
                     num_cpus=0,
                     num_gpus=0,
                     runtime_env=runtime_env,
@@ -130,7 +130,7 @@ async def setup_rl_llm(args):
                             "free_gpu_memory_fraction": args.kv_cache_fraction,
                         },
                         "cuda_graph_config": {
-                            "enable_padding": args.enable_padding,
+                            "enable_padding": args.enable_cuda_graph_padding,
                             "batch_sizes": args.batch_sizes,
                             "max_batch_size": 0 if args.batch_sizes else args.max_batch_size,
                         },
@@ -171,7 +171,7 @@ async def setup_rl_llm(args):
 
         # Helper function to wrap Ray remote call as async coroutine
         async def generate_single_prompt(instance, prompt):
-            """Generate a single prompt asynchronously"""
+            """Generate a single prompt asynchronously."""
             object_ref = instance.generate.remote(prompt=prompt)
             result = await asyncio.to_thread(ray.get, object_ref)
             return result
@@ -182,7 +182,7 @@ async def generate_single_prompt(instance, prompt):
             for idx, prompt in enumerate(prompts)
         ]
 
-        results = await asyncio.gather(*tasks)
+        await asyncio.gather(*tasks)
         end_time = time.time()
 
         print(f"Time taken: {end_time - start_time:.2f} seconds")
@@ -216,7 +216,10 @@ def add_rl_llm_args(parser):
         help="Sampler type.",
     )
     parser.add_argument(
-        "--trust_remote_code", type=bool, default=True, help="Whether to trust remote code."
+        "--trust_remote_code",
+        action="store_true",
+        default=False,
+        help="Whether to trust remote code.",
     )
 
     # KV Cache Config parameters
@@ -228,16 +231,16 @@ def add_rl_llm_args(parser):
     )
     parser.add_argument(
         "--enable_block_reuse",
-        type=bool,
-        default=True,
+        action="store_true",
+        default=False,
         help="Whether to enable block reuse for KV cache.",
     )
 
     # Cuda Graph Config parameters
     parser.add_argument(
-        "--enable_padding",
-        type=bool,
-        default=True,
+        "--enable_cuda_graph_padding",
+        action="store_true",
+        default=False,
         help="Whether to enable padding for CUDA graphs.",
     )
     parser.add_argument(

tests/unittest/_torch/ray_orchestrator/multi_gpu/test_multi_instance.py

@@ -0,0 +1,100 @@
+import os
+
+import pytest
+import ray
+import torch
+from ray.util.placement_group import (
+    PlacementGroupSchedulingStrategy,
+    placement_group,
+    remove_placement_group,
+)
+from utils.llm_data import llm_models_root
+
+from tensorrt_llm import AsyncLLM
+from tensorrt_llm.llmapi import KvCacheConfig
+
+
+@ray.remote
+class TRTLLMInstance:
+    def __init__(self, async_llm_kwargs: dict):
+        self.llm = AsyncLLM(
+            model=async_llm_kwargs["model"],
+            backend="pytorch",
+            orchestrator_type=async_llm_kwargs["orchestrator_type"],
+            kv_cache_config=KvCacheConfig(**async_llm_kwargs["kv_cache_config"]),
+            tensor_parallel_size=async_llm_kwargs["tensor_parallel_size"],
+            placement_groups=async_llm_kwargs["placement_groups"],
+            placement_bundle_indices=async_llm_kwargs["placement_bundle_indices"],
+            per_worker_gpu_share=async_llm_kwargs["per_worker_gpu_share"],
+        )
+
+    async def init_llm(self):
+        await self.llm.setup_async()
+
+
+@pytest.mark.gpu8
+@pytest.mark.parametrize(
+    "tp_size, num_instances", [(2, 4), (1, 8)], ids=["tp2_instances4", "tp1_instances8"]
+)
+def test_multi_instance(monkeypatch, tp_size, num_instances):
+    monkeypatch.setenv("RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", "1")
+
+    num_gpus = tp_size * num_instances
+    available_gpus = torch.cuda.device_count()
+    if num_gpus > 8:
+        raise ValueError(
+            f"Number of GPUs ({num_gpus}) is greater than 8. This script only supports single node."
+        )
+    if available_gpus < num_gpus:
+        raise ValueError(
+            f"Number of GPUs ({available_gpus}) is less than number of GPUs required ({num_gpus})."
+        )
+
+    excution_times = 5
+    for i in range(excution_times):
+        pg = None
+        try:
+            ray.init(address="local")
+            pg = placement_group(
+                [{"GPU": 1, "CPU": 2} for _ in range(num_gpus)], strategy="STRICT_PACK"
+            )
+
+            ray.get(pg.ready())
+
+            placement_group_list = [[pg] for _ in range(num_instances)]
+            placement_bundle_indices_list = [
+                [list(range(i * tp_size, (i + 1) * tp_size))] for i in range(num_instances)
+            ]
+
+            llm_instances = []
+            for i in range(num_instances):
+                llm_instances.append(
+                    TRTLLMInstance.options(
+                        num_cpus=0,
+                        num_gpus=0,
+                        scheduling_strategy=PlacementGroupSchedulingStrategy(
+                            placement_group=pg,
+                            placement_group_capture_child_tasks=True,
+                        ),
+                    ).remote(
+                        async_llm_kwargs={
+                            "model": os.path.join(
+                                llm_models_root(), "llama-models-v2", "TinyLlama-1.1B-Chat-v1.0"
+                            ),
+                            "kv_cache_config": {
+                                "free_gpu_memory_fraction": 0.1,
+                            },
+                            "tensor_parallel_size": tp_size,
+                            "orchestrator_type": "ray",
+                            "placement_groups": placement_group_list[i],
+                            "placement_bundle_indices": placement_bundle_indices_list[i],
+                            "per_worker_gpu_share": 0.5,
+                        }
+                    )
+                )
+            ray.get([llm.__ray_ready__.remote() for llm in llm_instances])
+            ray.get([llm.init_llm.remote() for llm in llm_instances])
+        finally:
+            if pg is not None:
+                remove_placement_group(pg)
+            ray.shutdown()