From 9726652f13c84a7de21079f43e1dae375a6350b5 Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Sun, 6 Sep 2026 23:24:24 -0700 Subject: [PATCH] fix(trtllm): export KV event hosts for MPI endpoints --- src/srtctl/cli/mixins/worker_stage.py | 12 +++++++++ tests/test_slurm.py | 35 +++++++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/src/srtctl/cli/mixins/worker_stage.py b/src/srtctl/cli/mixins/worker_stage.py index b2d34332e..5652c900c 100644 --- a/src/srtctl/cli/mixins/worker_stage.py +++ b/src/srtctl/cli/mixins/worker_stage.py @@ -343,6 +343,18 @@ def start_endpoint_worker(self, endpoint_processes: list["Process"]) -> ManagedP # Add config environment variables env_to_set.update(self.runtime.environment) + # Native TRT-LLM KV-event subscribers need routable publisher hosts for + # multi-node endpoints. Dynamo can otherwise fall back to + # SLURM_STEP_NODELIST, but that step-scoped variable is not guaranteed to + # be available inside every container-launch path. Set the endpoint's + # nodes explicitly, while preserving a recipe-provided override. + if ( + self.backend.type == "trtllm" + and len(endpoint_nodes) > 1 + and env_to_set.get("DYN_TRTLLM_PUBLISH_KV_EVENTS", "").lower() == "true" + ): + env_to_set.setdefault("DYN_TRTLLM_KV_EVENT_HOSTS", ",".join(endpoint_nodes)) + # Add profiling environment variables if profiling.enabled: profile_dir = str(self.runtime.log_dir / "profiles") diff --git a/tests/test_slurm.py b/tests/test_slurm.py index 3c2bf5d38..605fc9dd1 100644 --- a/tests/test_slurm.py +++ b/tests/test_slurm.py @@ -356,6 +356,41 @@ def test_start_endpoint_worker_request_plane_injected(tmp_path: Path) -> None: assert env["DYN_REQUEST_PLANE"] == "nats" +def test_trtllm_native_kv_events_receive_endpoint_hosts(tmp_path: Path) -> None: + mixin, process = _remap_worker_mixin(tmp_path, frontend_type="dynamo", dynamo_install=False) + mixin.config.backend.type = "trtllm" + mixin.runtime.environment = {"DYN_TRTLLM_PUBLISH_KV_EVENTS": "true"} + second_process = SimpleNamespace(**{**process.__dict__, "node": "node-b"}) + + with ( + patch("srtctl.cli.mixins.worker_stage.generate_capture_script", return_value="fingerprint || true"), + patch("srtctl.cli.mixins.worker_stage.start_srun_process") as mock_srun, + ): + mock_srun.return_value = MagicMock() + mixin.start_endpoint_worker([process, second_process]) + + assert mock_srun.call_args.kwargs["env_to_set"]["DYN_TRTLLM_KV_EVENT_HOSTS"] == "node-a,node-b" + + +def test_trtllm_native_kv_event_host_override_is_preserved(tmp_path: Path) -> None: + mixin, process = _remap_worker_mixin(tmp_path, frontend_type="dynamo", dynamo_install=False) + mixin.config.backend.type = "trtllm" + mixin.runtime.environment = { + "DYN_TRTLLM_PUBLISH_KV_EVENTS": "true", + "DYN_TRTLLM_KV_EVENT_HOSTS": "override-a,override-b", + } + second_process = SimpleNamespace(**{**process.__dict__, "node": "node-b"}) + + with ( + patch("srtctl.cli.mixins.worker_stage.generate_capture_script", return_value="fingerprint || true"), + patch("srtctl.cli.mixins.worker_stage.start_srun_process") as mock_srun, + ): + mock_srun.return_value = MagicMock() + mixin.start_endpoint_worker([process, second_process]) + + assert mock_srun.call_args.kwargs["env_to_set"]["DYN_TRTLLM_KV_EVENT_HOSTS"] == "override-a,override-b" + + def test_trtllm_sidecar_endpoint_kills_step_on_rank_failure(tmp_path: Path) -> None: mixin, process = _remap_worker_mixin(tmp_path, frontend_type="dynamo", dynamo_install=False) mixin.config.backend.type = "trtllm"