[TRTLLM-14040][fix] Fix multi-GPU VisualGen response IPC: REBUILD_LOCAL handle crossed a process boundary

tensorrt_llm/_torch/visual_gen/executor.py

@@ -277,12 +277,14 @@ def __init__(
         visual_gen_args: "VisualGenArgs",
         req_hmac_key: Optional[bytes] = None,
         resp_hmac_key: Optional[bytes] = None,
+        in_client_process: bool = False,
     ):
         self.request_queue_addr = request_queue_addr
         self.response_queue_addr = response_queue_addr
         self.device_id = device_id
         self.visual_gen_args = visual_gen_args
         self.resp_hmac_key = resp_hmac_key
+        self.in_client_process = in_client_process
 
         self.pipeline = None  # initialized in _load_pipeline
         self.requests_ipc = None
@@ -435,10 +437,9 @@ def process_request(self, req: DiffusionRequest):
             output = self.pipeline.infer(req)
             generation = time.perf_counter() - generation_start  # seconds
             if self.rank == 0:
-                # External launch co-locates this worker with the coordinator in one
-                # process; to_handle(local=True) hands the tensor over in-process
-                # instead of cross-process CUDA IPC (invalid same-process).
-                output.to_handle(local=_detect_external_launch() is not None)
+                # CUDA IPC handles are invalid within the producing process, so
+                # a same-process client takes the media via in-process handoff.
+                output.to_handle(local=self.in_client_process)
                 self.response_queue.put(
                     DiffusionResponse(
                         request_id=req.request_id,
@@ -467,8 +468,14 @@ def run_diffusion_worker(
     req_hmac_key: Optional[bytes] = None,
     resp_hmac_key: Optional[bytes] = None,
     local_rank: Optional[int] = None,
+    in_client_process: bool = False,
 ):
-    """Entry point for worker process."""
+    """Entry point for worker process.
+
+    ``in_client_process``: True only when this worker runs inside the client
+    process. Declared by the launch site — never derive it from the
+    environment here, the env writes below make every worker look external.
+    """
     try:
         # Set log level before any other work so loading logs are visible
         logger.set_level(log_level)
@@ -530,6 +537,7 @@ def run_diffusion_worker(
             visual_gen_args=visual_gen_args,
             req_hmac_key=req_hmac_key,
             resp_hmac_key=resp_hmac_key,
+            in_client_process=in_client_process,
         )
         executor.serve_forever()
         if executor.pipeline is not None:
@@ -690,6 +698,7 @@ def __init__(
                     "resp_hmac_key": self.resp_hmac_key,
                     "log_level": logger.level,
                     "local_rank": local_rank,
+                    "in_client_process": True,
                 },
                 daemon=True,
             )

tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py

@@ -9,12 +9,23 @@
 """
 
 import multiprocessing as mp
+import os
 import unittest
+from unittest import mock
 
 import torch
-
-from tensorrt_llm._torch.visual_gen.executor import DiffusionResponse
+import zmq
+
+from tensorrt_llm._torch.visual_gen.executor import (
+    DiffusionExecutor,
+    DiffusionRequest,
+    DiffusionResponse,
+    find_free_port,
+    run_diffusion_worker,
+)
 from tensorrt_llm._torch.visual_gen.output import PipelineOutput
+from tensorrt_llm.executor.ipc import ZeroMqQueue
+from tensorrt_llm.visual_gen import VisualGenArgs, VisualGenParams
 
 
 def _producer(q, device):
@@ -94,6 +105,150 @@ def test_roundtrip_local_same_process(self):
         self.assertEqual(resp.output.frame_rate, 24.0)
 
 
+def _expected_video() -> torch.Tensor:
+    return torch.arange(2 * 3 * 4 * 5 * 3, dtype=torch.uint8).reshape(2, 3, 4, 5, 3)
+
+
+class _StubPipeline:
+    """Stand-in for BasePipeline covering exactly what the executor touches."""
+
+    default_generation_params: dict = {}
+    extra_param_specs: dict = {}
+    _warmed_up_shapes: set = set()
+
+    def warmup_cache_key(self, height, width, num_frames):
+        return (height, width, num_frames)
+
+    def infer(self, req):
+        return PipelineOutput(video=_expected_video(), frame_rate=24.0)
+
+    def cleanup(self):
+        pass
+
+
+def _stub_load_pipeline(self):
+    self.pipeline = _StubPipeline()
+
+
+def _diffusion_worker_entry(rank, world_size, master_port, req_addr, resp_addr, req_key, resp_key):
+    """Run the real ``run_diffusion_worker`` with only the pipeline stubbed out.
+
+    The env scrub mirrors a worker spawned by a plain single-process client;
+    CUDA is forced off so dist init picks gloo and the stub stays on CPU.
+    """
+    for var in (
+        "RANK",
+        "WORLD_SIZE",
+        "LOCAL_RANK",
+        "MASTER_ADDR",
+        "MASTER_PORT",
+        "SLURM_PROCID",
+        "SLURM_NTASKS",
+    ):
+        os.environ.pop(var, None)
+
+    with (
+        mock.patch("torch.cuda.is_available", return_value=False),
+        mock.patch.object(DiffusionExecutor, "_load_pipeline", _stub_load_pipeline),
+    ):
+        run_diffusion_worker(
+            rank=rank,
+            world_size=world_size,
+            master_addr="127.0.0.1",
+            master_port=master_port,
+            request_queue_addr=req_addr,
+            response_queue_addr=resp_addr,
+            visual_gen_args=VisualGenArgs(model="/tmp/model"),
+            req_hmac_key=req_key,
+            resp_hmac_key=resp_key,
+            local_rank=rank,
+        )
+
+
+class TestSpawnWorkerResponseHandleMode(unittest.TestCase):
+    """A spawned worker must return cross-process response handles.
+
+    Drives the real ``run_diffusion_worker`` + serve loop + ZMQ response path
+    with world_size=2 (gloo, CPU) and asserts the client-side rebuild
+    succeeds. Any wrongly-local handle mode (e.g. derived from the
+    environment, which the worker entry mutates for dist init) ships a
+    same-process REBUILD_LOCAL handle the client cannot rebuild.
+    """
+
+    def test_spawn_worker_response_rebuilds_in_client(self):
+        ctx = mp.get_context("spawn")
+        world_size = 2
+        master_port = find_free_port()
+        req_port, resp_port = find_free_port(), find_free_port()
+        req_key, resp_key = os.urandom(32), os.urandom(32)
+
+        # Client-side server sockets, mirroring DiffusionRemoteClient.
+        requests_ipc = ZeroMqQueue(
+            (f"tcp://0.0.0.0:{req_port}", req_key),
+            is_server=True,
+            socket_type=zmq.PUSH,
+            use_hmac_encryption=True,
+        )
+        responses_ipc = ZeroMqQueue(
+            (f"tcp://0.0.0.0:{resp_port}", resp_key),
+            is_server=True,
+            socket_type=zmq.PULL,
+            use_hmac_encryption=True,
+        )
+
+        workers = [
+            ctx.Process(
+                target=_diffusion_worker_entry,
+                args=(
+                    rank,
+                    world_size,
+                    master_port,
+                    f"tcp://127.0.0.1:{req_port}",
+                    f"tcp://127.0.0.1:{resp_port}",
+                    req_key,
+                    resp_key,
+                ),
+            )
+            for rank in range(world_size)
+        ]
+        try:
+            for p in workers:
+                p.start()
+
+            requests_ipc.put(
+                DiffusionRequest(
+                    request_id=7,
+                    prompt=["stub"],
+                    params=VisualGenParams(height=8, width=8, num_frames=2),
+                )
+            )
+
+            resp = None
+            while resp is None or resp.request_id != 7:
+                self.assertTrue(responses_ipc.poll(timeout=180), "no response from spawned workers")
+                resp = responses_ipc.get()
+
+            self.assertIsNone(resp.error_msg, msg=str(resp.error_msg))
+            # A wrongly-local handle raises here (crossed a process boundary).
+            resp.output.to_tensor()
+            self.assertTrue(torch.equal(resp.output.video, _expected_video()))
+
+            requests_ipc.put(None)  # shutdown broadcast
+            for p in workers:
+                p.join(timeout=60)
+                self.assertEqual(p.exitcode, 0)
+        finally:
+            for p in workers:
+                if p.is_alive():
+                    p.terminate()
+            for p in workers:
+                p.join(timeout=30)
+            for q in (requests_ipc, responses_ipc):
+                if q.socket:
+                    q.socket.setsockopt(zmq.LINGER, 0)
+                q.close()
+
+
 if __name__ == "__main__":
     mp.set_start_method("spawn", force=True)
     unittest.main()