From 5a489649e15540719060878a211712b13dff68f8 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Mon, 13 Jul 2026 12:15:11 +0000 Subject: [PATCH 1/2] [TRTLLM-14040][fix] VisualGen: declare in-client-process workers at the launch site MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Deciding the response handle mode inside the worker via _detect_external_launch() misreads every mp-spawn worker as externally launched: run_diffusion_worker exports RANK/WORLD_SIZE for torch.distributed env:// init, after which spawn workers are indistinguishable from launcher-started ones. Rank 0 then ships a same-process (REBUILD_LOCAL) tensor handle across the process boundary and the client fails with 'REBUILD_LOCAL handle crossed a process boundary' on any single-node multi-GPU run. The launch site knows the fact the decision actually needs — whether the worker runs inside the client process — so declare it there: run_diffusion_worker and DiffusionExecutor take in_client_process (default False = cross-process handles), and only the external-launch branch that runs rank 0's worker as a client thread passes True. The worker-side environment sniffing is gone entirely; _detect_external_launch() remains only at its client-side call sites, which read the launcher-provided environment before anything mutates it. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/executor.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 2b81b8bd7a50..95a170120795 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/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, ) From d676a01ecdd5fbe4da1abac60922a108982a08c9 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Mon, 13 Jul 2026 12:15:24 +0000 Subject: [PATCH 2/2] [TRTLLM-14040][test] cover spawn-worker response handle mode end to end MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TestSpawnWorkerResponseHandleMode drives the real run_diffusion_worker entry, serve loop and ZMQ response path with world_size=2 (gloo, CPU-only) and asserts the client-side rebuild of the returned media succeeds — guarding the in_client_process=False default for spawn workers. Only the pipeline is stubbed, patched inside the spawned child. This is the missing combination the REBUILD_LOCAL regression slipped through: the worker entry's RANK/WORLD_SIZE env exports, a multi-worker world and the production handle mode decision were each tested separately but never together. Verified red/green: any wrongly-local handle mode fails this test with the production error ('REBUILD_LOCAL handle crossed a process boundary'). Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../test_executor_shared_tensor_ipc.py | 159 +++++++++++++++++- 1 file changed, 157 insertions(+), 2 deletions(-) diff --git a/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py b/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py index c5653c9e9d79..29fce319cdc7 100644 --- a/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py +++ b/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()