diff --git a/components/src/dynamo/trtllm/backend_args.py b/components/src/dynamo/trtllm/backend_args.py index 002d9cd98ccd..5ebbd02bb549 100644 --- a/components/src/dynamo/trtllm/backend_args.py +++ b/components/src/dynamo/trtllm/backend_args.py @@ -91,6 +91,14 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None: default=False, help="Enable attention data parallelism. When enabled, attention_dp_size equals tensor_parallel_size.", ) + add_negatable_bool_argument( + g, + flag_name="--conversation-affinity", + env_var="DYN_ENGINE_CONV_AFFINITY", + default=False, + help="Force engine-owned conversation-affinity ADP routing: the engine picks the " + "attention-DP rank from the conversation id, even if the router selects a rank.", + ) add_argument( g, flag_name="--kv-block-size", @@ -471,6 +479,7 @@ class DynamoTrtllmConfig(ConfigBase): pipeline_parallel_size: int expert_parallel_size: Optional[int] enable_attention_dp: bool + conversation_affinity: bool kv_block_size: int gpus_per_node: Optional[int] = None max_batch_size: int diff --git a/components/src/dynamo/trtllm/llm_engine.py b/components/src/dynamo/trtllm/llm_engine.py index e749d55e1654..d2573274e2dd 100644 --- a/components/src/dynamo/trtllm/llm_engine.py +++ b/components/src/dynamo/trtllm/llm_engine.py @@ -142,6 +142,7 @@ def __init__( disaggregation_mode: DisaggregationMode = DisaggregationMode.AGGREGATED, component: str = "backend", publish_events_and_metrics: bool = False, + conversation_affinity: bool = False, ): self.engine_args = engine_args self.model_name = model_name @@ -195,6 +196,9 @@ def __init__( # Default False so generate() is safe if reached before initialize() # (e.g. unit tests that drive generate directly). self._conversation_affinity: bool = False + # Manual override (--conversation-affinity / DYN_ENGINE_CONV_AFFINITY) to force + # engine-side conversation-affinity assignment regardless of engine detection. + self._engine_conversation_affinity_override: bool = conversation_affinity # Worker invokes on_ready callbacks serially during setup (see # `setup_kv_publishers` in lib/backend-common/src/publisher.rs); the # dict is fully populated before `_kv_events_thread` starts and @@ -330,6 +334,7 @@ async def from_args( disaggregation_mode=config.disaggregation_mode, component=config.component, publish_events_and_metrics=config.publish_events_and_metrics, + conversation_affinity=config.conversation_affinity, ) worker_config = WorkerConfig.from_runtime_config( config, @@ -879,8 +884,24 @@ async def _generate_started( # the engine side (always record the binding, whether the rank came # from explicit placement or fresh selection). Track upstream and drop # this note when the router records unconditionally. + # Force engine-side conversation-affinity assignment when the operator set + # --conversation-affinity / DYN_ENGINE_CONV_AFFINITY, even if the engine did + # not advertise the capability at initialize() time. Mirrors the legacy + # HandlerBase path so both entry points behave identically. + conv_affinity = ( + self._conversation_affinity or self._engine_conversation_affinity_override + ) + if ( + self._engine_conversation_affinity_override + and not CONVERSATION_PARAMS_AVAILABLE + ): + raise RuntimeError( + "--conversation-affinity / DYN_ENGINE_CONV_AFFINITY is set but " + "the installed TensorRT-LLM build has no ConversationParams API (requires a " + "release newer than 1.3.0rc20)." + ) conversation_params = None - if self._conversation_affinity: + if conv_affinity: scheduling_params = None conversation_params = conversation_params_for( session_id_from_request(request) @@ -908,9 +929,7 @@ async def _generate_started( # Only pass conversation_params when affinity is active — older TRT-LLM builds' # generate_async does not accept the keyword. conv_kwargs = ( - {"conversation_params": conversation_params} - if self._conversation_affinity - else {} + {"conversation_params": conversation_params} if conv_affinity else {} ) generation_result = self._engine.llm.generate_async( inputs=token_ids, diff --git a/components/src/dynamo/trtllm/request_handlers/handler_base.py b/components/src/dynamo/trtllm/request_handlers/handler_base.py index 0c37e260dd7e..5996eb6c46d9 100644 --- a/components/src/dynamo/trtllm/request_handlers/handler_base.py +++ b/components/src/dynamo/trtllm/request_handlers/handler_base.py @@ -247,6 +247,8 @@ class RequestHandlerConfig: additional_metrics: Optional["AdditionalMetricsCollector"] = None max_seq_len: Optional[int] = None disagg_machine_id: int = 0 # 10-bit machine_id for snowflake disagg_request_id + # Force engine-owned conversation-affinity ADP routing regardless of engine detection. + conversation_affinity: bool = False class HandlerBase(BaseGenerativeHandler): @@ -271,6 +273,9 @@ def __init__(self, config: RequestHandlerConfig): # request (engine may not be initialized at handler construction). See # conversation_affinity.py. None = not yet resolved. self._conversation_affinity: Optional[bool] = None + # Manual override (--conversation-affinity / DYN_ENGINE_CONV_AFFINITY) to force + # engine-side assignment of conversation-affinity regardless of engine detection. + self._engine_conversation_affinity_override: bool = config.conversation_affinity self.encode_client = config.encode_client self.multimodal_processor = config.multimodal_processor self.first_generation = True @@ -1110,8 +1115,18 @@ async def _generate_locally_impl( # initialized by first request). When on, the engine's ConversationAwareADPRouter # picks the attention-DP rank from the conversation id, so we must NOT force a rank. if self._conversation_affinity is None: - self._conversation_affinity = engine_conversation_affinity_enabled( - self.engine.llm + if ( + self._engine_conversation_affinity_override + and not CONVERSATION_PARAMS_AVAILABLE + ): + raise RuntimeError( + "--conversation-affinity / DYN_ENGINE_CONV_AFFINITY is set but " + "the installed TensorRT-LLM build has no ConversationParams API (requires a " + "release newer than 1.3.0rc20)." + ) + self._conversation_affinity = ( + self._engine_conversation_affinity_override + or engine_conversation_affinity_enabled(self.engine.llm) ) if self._conversation_affinity and not CONVERSATION_PARAMS_AVAILABLE: raise RuntimeError( diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_engine_conversation_affinity.py b/components/src/dynamo/trtllm/tests/test_trtllm_engine_conversation_affinity.py index 0ba1291b785a..fb697909612a 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_engine_conversation_affinity.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_engine_conversation_affinity.py @@ -14,6 +14,11 @@ ``conversation_params.conversation_id == session_id``. * on + no session id → ``scheduling_params=None`` and ``conversation_params=None`` (the kwarg is still passed so the engine's no-id balancing fallback runs). + +Plus the manual ``--conversation-affinity`` / ``DYN_ENGINE_CONV_AFFINITY`` override +(``_engine_conversation_affinity_override``), which forces the affinity branch even when +engine detection is off and raises when the ConversationParams API is missing — matching +the legacy ``HandlerBase`` override tests. """ from __future__ import annotations @@ -81,6 +86,9 @@ def _make_engine( engine._no_inflight_requests.set() engine._reject_new_requests = False engine._conversation_affinity = conversation_affinity + # Manual override (--conversation-affinity / DYN_ENGINE_CONV_AFFINITY); off by + # default, individual tests flip it to exercise the override path. + engine._engine_conversation_affinity_override = False engine._attention_dp_size = attention_dp_size engine._override_sampling_params = lambda default, request: SimpleNamespace( max_tokens=None @@ -189,3 +197,72 @@ def fake_generate_async(**kwargs): assert captured["scheduling_params"] is None assert "conversation_params" in captured assert captured["conversation_params"] is None + + +async def test_override_suppresses_dp_rank_and_forwards_conversation_params( + monkeypatch, +): + """DYN_ENGINE_CONV_AFFINITY override=True + engine detection disabled → + dp_rank suppressed, conversation_params forwarded. Mirrors the legacy + HandlerBase override test so both entry points stay in lockstep.""" + monkeypatch.setattr( + "dynamo.trtllm.llm_engine.CONVERSATION_PARAMS_AVAILABLE", + True, + ) + monkeypatch.setattr( + "dynamo.trtllm.conversation_affinity.ConversationParams", + _FakeConversationParams, + ) + + captured: dict = {} + + def fake_generate_async(**kwargs): + captured.update(kwargs) + return _empty_async_iter() + + # Engine detection returns False (no engine-side affinity config)... + engine = _make_engine( + fake_generate_async, conversation_affinity=False, attention_dp_size=8 + ) + # ...but the operator forced it via DYN_ENGINE_CONV_AFFINITY. + engine._engine_conversation_affinity_override = True + # Router still stamps a rank; the override must ignore it just like + # auto-detection does. + await _drain( + engine, + { + "token_ids": [1, 2, 3], + "routing": {"dp_rank": 3}, + "agent_context": {"session_id": "run-99:agent-0"}, + }, + ) + + assert captured["scheduling_params"] is None + conv_params = captured["conversation_params"] + assert conv_params is not None + assert conv_params.conversation_id == "run-99:agent-0" + + +async def test_override_raises_when_conversation_params_api_missing(monkeypatch): + """DYN_ENGINE_CONV_AFFINITY=true on a build without ConversationParams → + RuntimeError in _generate_started, before generate_async is called.""" + monkeypatch.setattr( + "dynamo.trtllm.llm_engine.CONVERSATION_PARAMS_AVAILABLE", + False, + ) + + called = False + + def fake_generate_async(**kwargs): + nonlocal called + called = True + return _empty_async_iter() + + engine = _make_engine( + fake_generate_async, conversation_affinity=False, attention_dp_size=8 + ) + engine._engine_conversation_affinity_override = True + + with pytest.raises(RuntimeError, match="DYN_ENGINE_CONV_AFFINITY"): + await _drain(engine, {"token_ids": [1, 2, 3]}) + assert called is False diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py b/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py index 14cd94d673b6..030d5a0d0dcd 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py @@ -470,6 +470,7 @@ class TestDeferredAbortGuard: def _make_handler(self) -> HandlerBase: config = MagicMock() config.shutdown_event = None + config.conversation_affinity = False return _ConcreteHandler(config) @pytest.mark.asyncio @@ -591,6 +592,7 @@ def _make_handler(self, multimodal_processor=None) -> HandlerBase: config = MagicMock() config.multimodal_processor = multimodal_processor config.shutdown_event = None + config.conversation_affinity = False return _ConcreteHandler(config) async def _prepare(self, handler, request, epd_metadata=None): @@ -647,6 +649,7 @@ def _make_handler(self) -> HandlerBase: config = MagicMock() config.multimodal_processor = MagicMock() config.shutdown_event = None + config.conversation_affinity = False return _ConcreteHandler(config) def _pack_metadata(self, request, processed_input, prompt, prompt_token_ids): @@ -717,6 +720,7 @@ def _make_prefill_handler(self, machine_id: int = 42) -> HandlerBase: config = MagicMock() config.shutdown_event = None config.disagg_machine_id = machine_id + config.conversation_affinity = False handler = _ConcreteHandler(config) handler.disaggregation_mode = DisaggregationMode.PREFILL return handler @@ -800,6 +804,7 @@ def _make_handler(self) -> HandlerBase: config = MagicMock() config.shutdown_event = None config.disaggregation_mode = DisaggregationMode.AGGREGATED + config.conversation_affinity = False handler = _ConcreteHandler(config) handler.publisher = None handler.multimodal_processor = None @@ -1128,6 +1133,7 @@ def _make_handler(self, *, conversation_affinity: bool) -> HandlerBase: config = MagicMock() config.shutdown_event = None config.disaggregation_mode = DisaggregationMode.AGGREGATED + config.conversation_affinity = False handler = _ConcreteHandler(config) handler.publisher = None handler.multimodal_processor = None @@ -1284,3 +1290,69 @@ async def test_lazy_resolution_raises_when_api_missing(self, monkeypatch): async for _ in handler.generate_locally(request, self._make_context()): pass handler.engine.llm.generate_async.assert_not_called() + + @pytest.mark.asyncio + async def test_override_suppresses_dp_rank_and_forwards_conversation_params( + self, monkeypatch + ): + """DYN_ENGINE_CONV_AFFINITY override=True + engine detection disabled → + dp_rank suppressed, conversation_params forwarded on first request.""" + monkeypatch.setattr( + "dynamo.trtllm.request_handlers.handler_base.CONVERSATION_PARAMS_AVAILABLE", + True, + ) + monkeypatch.setattr( + "dynamo.trtllm.conversation_affinity.ConversationParams", + _FakeConversationParams, + ) + monkeypatch.setattr( + "dynamo.trtllm.request_handlers.handler_base.engine_conversation_affinity_enabled", + lambda _: False, + ) + handler = self._make_handler(conversation_affinity=False) + # Reset to None so lazy init runs on first request and folds in the override. + handler._conversation_affinity = None + handler._engine_conversation_affinity_override = True + kwargs = await self._drive( + handler, + { + "token_ids": [1, 2, 3], + "stop_conditions": {"max_tokens": 10}, + "sampling_options": {"temperature": 0.7}, + "routing": {"dp_rank": 3}, + "agent_context": {"session_id": "run-99:agent-0"}, + }, + ) + # Override must suppress the router rank just like auto-detection does. + assert kwargs["scheduling_params"] is None + conv_params = kwargs["conversation_params"] + assert conv_params is not None + assert conv_params.conversation_id == "run-99:agent-0" + + @pytest.mark.asyncio + async def test_override_raises_when_conversation_params_api_missing( + self, monkeypatch + ): + """DYN_ENGINE_CONV_AFFINITY=true on a build without ConversationParams → + RuntimeError on first request during lazy init.""" + monkeypatch.setattr( + "dynamo.trtllm.request_handlers.handler_base.CONVERSATION_PARAMS_AVAILABLE", + False, + ) + handler = self._make_handler(conversation_affinity=False) + # Reset to None so lazy init runs and hits the guard. + handler._conversation_affinity = None + handler._engine_conversation_affinity_override = True + handler.engine.llm.generate_async = MagicMock() + + with pytest.raises(RuntimeError, match="DYN_ENGINE_CONV_AFFINITY"): + async for _ in handler.generate_locally( + { + "token_ids": [1, 2, 3], + "stop_conditions": {"max_tokens": 10}, + "sampling_options": {"temperature": 0.7}, + }, + self._make_context(), + ): + pass + handler.engine.llm.generate_async.assert_not_called() diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_trace_propagation.py b/components/src/dynamo/trtllm/tests/test_trtllm_trace_propagation.py index e4334ff31b88..29ba8c0139c7 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_trace_propagation.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_trace_propagation.py @@ -68,8 +68,9 @@ def _make_engine(generate_async) -> TrtllmLLMEngine: engine._no_inflight_requests = asyncio.Event() engine._no_inflight_requests.set() engine._reject_new_requests = False - # Match __init__ default; generate_locally reads this before dispatch. + # Match __init__ defaults; _generate_started reads both before dispatch. engine._conversation_affinity = False + engine._engine_conversation_affinity_override = False # Single-rank default: validate_global_dp_rank(None, 0, 1, ...) -> None, # so scheduling_params stays None and never touches a real SchedulingParams. engine._attention_dp_size = 1 diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_unit.py b/components/src/dynamo/trtllm/tests/test_trtllm_unit.py index e041efdbb4c2..1a377d48f9be 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_unit.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_unit.py @@ -249,6 +249,27 @@ def test_disaggregation_mode_legacy_aggregated_value_warns(): assert config.disaggregation_mode == DisaggregationMode.AGGREGATED +def test_conversation_affinity_cli_flag(monkeypatch): + """--conversation-affinity sets conversation_affinity=True in Config.""" + monkeypatch.delenv("DYN_ENGINE_CONV_AFFINITY", raising=False) + config = parse_args(["--model", "fake-model", "--conversation-affinity"]) + assert config.conversation_affinity is True + + +def test_conversation_affinity_env_var(monkeypatch): + """DYN_ENGINE_CONV_AFFINITY=true is read and sets conversation_affinity=True.""" + monkeypatch.setenv("DYN_ENGINE_CONV_AFFINITY", "true") + config = parse_args(["--model", "fake-model"]) + assert config.conversation_affinity is True + + +def test_conversation_affinity_defaults_false(monkeypatch): + """conversation_affinity defaults to False when neither flag nor env var is set.""" + monkeypatch.delenv("DYN_ENGINE_CONV_AFFINITY", raising=False) + config = parse_args(["--model", "fake-model"]) + assert config.conversation_affinity is False + + def test_enable_multimodal_rejects_diffusion_modality(): with pytest.raises(ValueError, match="--enable-multimodal cannot be combined"): parse_args( diff --git a/components/src/dynamo/trtllm/tests/utils.py b/components/src/dynamo/trtllm/tests/utils.py index aed2cd663d00..a9157020370a 100644 --- a/components/src/dynamo/trtllm/tests/utils.py +++ b/components/src/dynamo/trtllm/tests/utils.py @@ -32,5 +32,6 @@ def create_mock_request_handler_config( config.runtime = None config.kv_block_size = 32 config.shutdown_event = None + config.conversation_affinity = False config.encoder_cache_capacity_gb = encoder_cache_capacity_gb return config diff --git a/components/src/dynamo/trtllm/workers/llm_worker.py b/components/src/dynamo/trtllm/workers/llm_worker.py index 5f71ca4ef585..f3b8001a7b35 100644 --- a/components/src/dynamo/trtllm/workers/llm_worker.py +++ b/components/src/dynamo/trtllm/workers/llm_worker.py @@ -678,6 +678,7 @@ async def init_llm_worker( additional_metrics=additional_metrics, max_seq_len=config.max_seq_len, disagg_machine_id=int(endpoint.connection_id()) % 1021, + conversation_affinity=config.conversation_affinity, ) media_decoder = None