feat(router): add dynamic worker taint updates

Cargo.toml

@@ -107,6 +107,7 @@ either = { version = "1.13", features = ["serde"] }
 etcd-client = { version = "0.17.0", features = ["tls"] }
 futures = { version = "0.3" }
 futures-util = { version = "0.3.32" }
+fs4 = { version = "0.13.1" }
 hf-hub = { version = "0.4.2", default-features = false }
 
 # ModelExpress for model downloading

components/src/dynamo/common/model_taints.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Worker-local HTTP route for updating model routing taints."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from dynamo.llm import update_model_taints
+from dynamo.runtime import DistributedRuntime, Endpoint
+
+MODEL_TAINT_ROUTE = "update/model_taints"
+TOPOLOGY_TAINT_PREFIX = "dynamo.topology/"
+
+
+def register_model_taint_route(runtime: DistributedRuntime, endpoint: Endpoint) -> None:
+    """Register POST /engine/update/model_taints on the system status server."""
+
+    async def _update_model_taints(body: dict[str, Any]) -> dict[str, Any]:
+        if not isinstance(body, dict):
+            raise ValueError("request body must be a JSON object")
+
+        taints = body.get("taints")
+        if not isinstance(taints, list) or not all(
+            isinstance(taint, str) for taint in taints
+        ):
+            raise ValueError("'taints' must be a JSON array of strings")
+        if reserved := next(
+            (taint for taint in taints if taint.startswith(TOPOLOGY_TAINT_PREFIX)),
+            None,
+        ):
+            raise ValueError(
+                f"taint '{reserved}' uses reserved prefix '{TOPOLOGY_TAINT_PREFIX}'"
+            )
+
+        unique_taints = set(taints)
+        await update_model_taints(endpoint, unique_taints)
+        return {
+            "status": "ok",
+            "taints": sorted(unique_taints),
+        }
+
+    runtime.register_engine_route(MODEL_TAINT_ROUTE, _update_model_taints)

components/src/dynamo/common/tests/test_model_taints.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import AsyncMock
+
+import pytest
+
+from dynamo.common import model_taints
+
+pytestmark = [pytest.mark.pre_merge, pytest.mark.unit, pytest.mark.gpu_0]
+
+
+class _Runtime:
+    def __init__(self) -> None:
+        self.route_name: str | None = None
+        self.handler = None
+
+    def register_engine_route(self, name, handler) -> None:
+        self.route_name = name
+        self.handler = handler
+
+
+def test_model_taint_route_updates_worker(monkeypatch: pytest.MonkeyPatch) -> None:
+    update = AsyncMock()
+    monkeypatch.setattr(model_taints, "update_model_taints", update)
+    runtime = _Runtime()
+    endpoint = object()
+
+    model_taints.register_model_taint_route(runtime, endpoint)
+
+    assert runtime.route_name == "update/model_taints"
+    response = asyncio.run(
+        runtime.handler({"taints": ["capacity/fast", "capacity/fast"]})
+    )
+    assert response == {
+        "status": "ok",
+        "taints": ["capacity/fast"],
+    }
+    update.assert_awaited_once_with(endpoint, {"capacity/fast"})
+
+
+@pytest.mark.parametrize(
+    ("body", "message"),
+    [
+        ({}, "'taints' must be a JSON array of strings"),
+        ({"taints": "fast"}, "'taints' must be a JSON array of strings"),
+        ({"taints": [1]}, "'taints' must be a JSON array of strings"),
+        (
+            {"taints": ["dynamo.topology/zone=west"]},
+            "uses reserved prefix 'dynamo.topology/'",
+        ),
+    ],
+)
+def test_model_taint_route_rejects_invalid_requests(body, message) -> None:
+    runtime = _Runtime()
+    model_taints.register_model_taint_route(runtime, object())
+
+    with pytest.raises(ValueError, match=message):
+        asyncio.run(runtime.handler(body))

components/src/dynamo/sglang/init_diffusion.py

@@ -8,6 +8,7 @@
 
 import sglang as sgl
 
+from dynamo.common.model_taints import register_model_taint_route
 from dynamo.common.storage import get_fs
 from dynamo.common.utils.endpoint_types import parse_endpoint_types
 from dynamo.llm import WorkerType
@@ -186,6 +187,7 @@ async def init_image_diffusion(
             "Overriding output_modalities to ['image'] for image diffusion worker"
         )
 
+    register_model_taint_route(runtime, generate_endpoint)
     try:
         await asyncio.gather(
             generate_endpoint.serve_endpoint(
@@ -261,6 +263,7 @@ async def init_video_diffusion(
 
     ready_event = asyncio.Event()
 
+    register_model_taint_route(runtime, generate_endpoint)
     try:
         await asyncio.gather(
             generate_endpoint.serve_endpoint(

components/src/dynamo/sglang/init_embedding.py

@@ -7,6 +7,7 @@
 
 import sglang as sgl
 
+from dynamo.common.model_taints import register_model_taint_route
 from dynamo.common.utils.prometheus import register_engine_metrics_callback
 from dynamo.llm import ModelInput, ModelType, WorkerType
 from dynamo.runtime import DistributedRuntime
@@ -71,6 +72,7 @@ async def init_embedding(
         engine, use_text_input=dynamo_args.use_sglang_tokenizer
     ).to_dict()
 
+    register_model_taint_route(runtime, generate_endpoint)
     try:
         await asyncio.gather(
             generate_endpoint.serve_endpoint(

components/src/dynamo/sglang/init_multimodal.py

@@ -9,6 +9,7 @@
 
 from dynamo import prometheus_names
 from dynamo.common.constants import DisaggregationMode
+from dynamo.common.model_taints import register_model_taint_route
 from dynamo.common.utils.prometheus import register_embedding_cache_metrics
 from dynamo.llm import (
     ModelInput,
@@ -78,6 +79,7 @@ async def init_multimodal_encode_worker(
 
     ready_event = asyncio.Event()
 
+    register_model_taint_route(runtime, generate_endpoint)
     try:
         _ = await asyncio.gather(
             generate_endpoint.serve_endpoint(
@@ -172,6 +174,7 @@ async def init_multimodal_worker(
         readiness_worker_type = WorkerType.Aggregated
         readiness_needs = [[WorkerType.Encode]]
 
+    register_model_taint_route(runtime, generate_endpoint)
     try:
         await asyncio.gather(
             generate_endpoint.serve_endpoint(
@@ -223,6 +226,7 @@ async def init_multimodal_prefill_worker(
 
     health_check_payload = SglangPrefillHealthCheckPayload(engine).to_dict()
 
+    register_model_taint_route(runtime, generate_endpoint)
     # No OpenAI surface (ModelType.Empty): internal prefill worker, reached via
     # the decode worker / prefill router, never by the frontend. Registers a
     # topology card so the serving-readiness gate counts it.

components/src/dynamo/sglang/request_handlers/handler_base.py

@@ -27,6 +27,7 @@
 from dynamo._core import Context
 from dynamo.common.constants import DisaggregationMode
 from dynamo.common.lora.manager import get_lora_manager
+from dynamo.common.model_taints import MODEL_TAINT_ROUTE, register_model_taint_route
 from dynamo.common.utils.endpoint_types import parse_endpoint_types
 from dynamo.common.utils.guided_json import reject_nonprogressing_guided_json_ref_cycles
 from dynamo.common.utils.input_params import InputParamManager
@@ -893,13 +894,15 @@ def register_engine_routes(self, runtime: DistributedRuntime) -> None:
             "control/update_weights_from_ipc": self.update_weights_from_ipc,
             "control/update_weight_version": self.update_weight_version,
         }
+        reserved_routes = {*built_in_routes, MODEL_TAINT_ROUTE}
         for path, _ in configured_routes:
-            if path in built_in_routes:
+            if path in reserved_routes:
                 raise ValueError(
                     f"Configured SGLang engine route /engine/{path} collides "
                     "with a built-in route"
                 )
 
+        register_model_taint_route(runtime, self.generate_endpoint)
         for path, handler in built_in_routes.items():
             runtime.register_engine_route(path, handler)
         for path, configured_handler in configured_routes:

components/src/dynamo/sglang/tests/test_sglang_unit.py

@@ -97,13 +97,14 @@ def _cpu_engine_when_no_accelerator(monkeypatch):
     _sgl_common.get_device.cache_clear()
 
 
-def test_configured_engine_route_cannot_replace_built_in_route():
+@pytest.mark.parametrize(
+    "reserved_path", ["control/start_profile", "update/model_taints"]
+)
+def test_configured_engine_route_cannot_replace_built_in_route(reserved_path):
     handler = object.__new__(DecodeWorkerHandler)
     handler.engine = SimpleNamespace(custom_method=lambda: None)
     handler.config = SimpleNamespace(
-        dynamo_args=SimpleNamespace(
-            engine_routes=["control/start_profile=custom_method"]
-        )
+        dynamo_args=SimpleNamespace(engine_routes=[f"{reserved_path}=custom_method"])
     )
 
     registered_routes = []
@@ -115,15 +116,45 @@ def register_engine_route(self, path, route_handler):
     with pytest.raises(
         ValueError,
         match=(
-            "Configured SGLang engine route /engine/control/start_profile "
-            "collides with a built-in route"
+            rf"Configured SGLang engine route /engine/{reserved_path} "
+            r"collides with a built-in route"
         ),
     ):
         handler.register_engine_routes(Runtime())
 
     assert registered_routes == []
 
 
+def test_builtin_engine_routes_include_model_taint_update(monkeypatch):
+    handler = object.__new__(DecodeWorkerHandler)
+    handler.engine = SimpleNamespace()
+    handler.generate_endpoint = object()
+    handler.config = SimpleNamespace(dynamo_args=SimpleNamespace(engine_routes=[]))
+
+    registered_routes = []
+    taint_route_endpoints = []
+
+    class Runtime:
+        def register_engine_route(self, path, route_handler):
+            registered_routes.append((path, route_handler))
+
+    runtime = Runtime()
+    monkeypatch.setattr(
+        "dynamo.sglang.request_handlers.handler_base.register_model_taint_route",
+        lambda candidate_runtime, endpoint: taint_route_endpoints.append(
+            (candidate_runtime, endpoint)
+        ),
+    )
+
+    handler.register_engine_routes(runtime)
+
+    assert taint_route_endpoints == [(runtime, handler.generate_endpoint)]
+    assert {path for path, _ in registered_routes} >= {
+        "control/start_profile",
+        "control/stop_profile",
+    }
+
+
 def _make_sglang_config(**overrides):
     config = DynamoSGLangConfig()
     config.use_sglang_tokenizer = False

components/src/dynamo/trtllm/workers/image_diffusion_worker.py

@@ -11,6 +11,7 @@
 import logging
 from typing import Optional
 
+from dynamo.common.model_taints import register_model_taint_route
 from dynamo.llm import ModelInput, ModelType, WorkerType, register_model
 from dynamo.runtime import DistributedRuntime
 from dynamo.trtllm.args import Config
@@ -101,6 +102,7 @@ async def init_image_diffusion_worker(
         worker_type=WorkerType.Aggregated,
         needs=[],
     )
+    register_model_taint_route(runtime, endpoint)
 
     logging.info(f"Model registered, serving endpoint: {config.endpoint}")
 

components/src/dynamo/trtllm/workers/llm_worker.py

@@ -38,6 +38,7 @@
 import dynamo.nixl_connect as nixl_connect
 from dynamo import prometheus_names
 from dynamo.common.config_dump import dump_config
+from dynamo.common.model_taints import register_model_taint_route
 from dynamo.common.utils.endpoint_types import parse_endpoint_types
 from dynamo.common.utils.prometheus import (
     LLMBackendMetrics,
@@ -875,6 +876,7 @@ async def init_llm_worker(
             worker_type=worker_type,
             needs=needs,
         )
+        register_model_taint_route(runtime, endpoint)
 
         health_check_payload = TrtllmHealthCheckPayload(
             tokenizer=tokenizer,

components/src/dynamo/trtllm/workers/video_diffusion_worker.py

@@ -11,6 +11,7 @@
 import logging
 from typing import Optional
 
+from dynamo.common.model_taints import register_model_taint_route
 from dynamo.llm import ModelInput, ModelType, WorkerType, register_model
 from dynamo.runtime import DistributedRuntime
 from dynamo.trtllm.args import Config
@@ -100,6 +101,7 @@ async def init_video_diffusion_worker(
         worker_type=WorkerType.Aggregated,
         needs=[],
     )
+    register_model_taint_route(runtime, endpoint)
 
     logging.info(f"Model registered, serving endpoint: {config.endpoint}")
 

components/src/dynamo/vllm/omni/main.py

@@ -11,6 +11,7 @@
 
 from dynamo import prometheus_names
 from dynamo.common.config_dump import dump_config
+from dynamo.common.model_taints import register_model_taint_route
 from dynamo.common.rl import first_endpoint_response
 from dynamo.common.storage import get_fs
 from dynamo.common.utils.graceful_shutdown import install_signal_handlers
@@ -132,6 +133,7 @@ async def init_omni(
         await shutdown_event.wait()
         return
 
+    register_model_taint_route(runtime, generate_endpoint)
     model_type = get_output_modalities(config.output_modalities, config.model)
     if model_type is None:
         model_type = ModelType.Images

components/src/dynamo/vllm/omni/realtime_utils.py

@@ -13,6 +13,7 @@
 import logging
 
 from dynamo import prometheus_names
+from dynamo.common.model_taints import register_model_taint_route
 from dynamo.llm import ModelInput, ModelType, WorkerType, register_model
 from dynamo.runtime import DistributedRuntime
 from dynamo.vllm.main import setup_metrics_collection
@@ -69,6 +70,7 @@ async def init_omni_realtime(
         await shutdown_event.wait()
         return
 
+    register_model_taint_route(runtime, generate_endpoint)
     model_label = config.served_model_name or config.model
     try:
         await register_model(

components/src/dynamo/vllm/omni/stage_router.py

@@ -12,6 +12,7 @@
 from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs
 
 from dynamo import prometheus_names
+from dynamo.common.model_taints import register_model_taint_route
 from dynamo.common.storage import get_fs
 from dynamo.common.utils.output_modalities import (
     RequestType,
@@ -322,6 +323,7 @@ async def init_omni_stage_router(
         worker_type=WorkerType.Aggregated,
         needs=[],
     )
+    register_model_taint_route(runtime, generate_endpoint)
     logger.info("OmniStageRouter registered at '%s'", generate_endpoint)
 
     try: