fix(backends): backport ModelStreamer object-storage fixes (#10674)

Signed-off-by: Zheng Luo <zheluo@nvidia.com>
Signed-off-by: nnshah1 <neelays@nvidia.com>
Co-authored-by: Neelay Shah <neelays@nvidia.com>

Author: zhengluo-nv

components/src/dynamo/sglang/register.py

@@ -28,6 +28,38 @@
 SGLANG_HICACHE_MOONCAKE_RUNTIME_KEY = "sglang_hicache_mooncake"
 
 
+def _register_model_source_path(engine: sgl.Engine, server_args: ServerArgs) -> str:
+    """Pick the path passed to `register_model` for MDC construction.
+
+    When `--model-path` is a remote URI (`s3://...`, `gs://...`), SGLang's
+    `ModelConfig.maybe_pull_model_tokenizer_from_remote` (sglang/srt/configs/
+    model_config.py) pulls metadata files (`*config.json`) to a local temp dir
+    and rewrites the engine's ModelConfig:
+
+      - `.model_weights = <original URI>`  (preserved for the weight loader)
+      - `.model_path    = <local temp dir>` (now contains the metadata)
+
+    `server_args.model_path` itself is NOT mutated. Dynamo's `register_model`
+    would otherwise pass the raw URI through `hub.rs` -> ModelExpress, which
+    has no S3 provider and 404s. Returning the post-pull local dir lets
+    `register_model` take its `fs::exists` shortcut and skip the broken MX
+    path.
+
+    Temporary LLM-only workaround mirroring the vLLM fix in
+    `components/src/dynamo/vllm/main.py`. Diffusion paths (Images/Videos) skip
+    the Rust-side HF download entirely (lib/bindings/python/rust/lib.rs:314)
+    and don't need this rewrite.
+    """
+    try:
+        mc = engine.tokenizer_manager.model_config
+    except AttributeError:
+        return server_args.model_path
+    weights = getattr(mc, "model_weights", None)
+    if weights:
+        return mc.model_path
+    return server_args.model_path
+
+
 def _build_media_decoder_and_fetcher():
     """Construct MediaDecoder/MediaFetcher for frontend-decoded multimodal.
 
@@ -89,7 +121,7 @@ async def _register_model_with_runtime_config(
             input_type,
             output_type,
             endpoint,
-            server_args.model_path,
+            _register_model_source_path(engine, server_args),
             server_args.served_model_name,
             context_length=server_args.context_length,
             kv_cache_block_size=server_args.page_size,

components/src/dynamo/vllm/llm_engine.py

@@ -89,14 +89,13 @@ async def start(self, worker_id: int) -> EngineConfig:
             )
             os.environ["PROMETHEUS_MULTIPROC_DIR"] = self._prometheus_temp_dir.name
 
-        self._default_sampling_params = (
-            self.engine_args.create_model_config().get_diff_sampling_param()
-        )
-
         vllm_config = self.engine_args.create_engine_config(
             usage_context=UsageContext.OPENAI_API_SERVER
         )
         self._vllm_config = vllm_config
+        self._default_sampling_params = (
+            vllm_config.model_config.get_diff_sampling_param()
+        )
 
         self.engine_client = AsyncLLM.from_vllm_config(
             vllm_config=vllm_config,

components/src/dynamo/vllm/main.py

@@ -70,6 +70,29 @@ def should_register_model_ignore_weights(config: Config) -> bool:
     return uses_modelexpress_load_format(config)
 
 
+def _register_model_source_path(config: Config, vllm_config: VllmConfig) -> str:
+    """Pick the path passed to `register_model` for MDC construction.
+
+    When `--model` is an object-storage URI (`s3://...`, `gs://...`, `az://...`),
+    vLLM's `maybe_pull_model_tokenizer_for_runai` (vllm/config/model.py) pulls
+    metadata files to a local temp dir and rewrites `vllm_config.model_config`:
+
+      - `.model_weights = <original URI>`  (used by runai-streamer / mx plugin)
+      - `.model = <local temp dir>`        (contains config.json, tokenizer, …)
+
+    Dynamo's `register_model` would otherwise try to resolve the raw URI via
+    `hub.rs` → ModelExpress, which has no S3 provider and 404s. Returning the
+    local dir lets `register_model` take its `fs::exists` shortcut.
+
+    Temporary vLLM-only workaround until `hub.rs` learns object-storage routing.
+    Falls back to `config.model` whenever vLLM did not pull (HF id, local path,
+    or older vLLM without `model_weights`).
+    """
+    if getattr(vllm_config.model_config, "model_weights", ""):
+        return vllm_config.model_config.model
+    return config.model
+
+
 def build_headless_namespace(config: Config) -> argparse.Namespace:
     """Build an argparse Namespace from engine_args for vLLM's run_headless().
 
@@ -490,14 +513,6 @@ def setup_vllm_engine(
             # Prevents deadlock during TP>1 failover.
             configure_gms_lock_mode(engine_args)
 
-    # ModelExpress uses vLLM's plugin path with --load-format=modelexpress.
-    # Dynamo does not register loaders or set a custom worker class here.
-
-    # Load default sampling params from `generation_config.json`
-    default_sampling_params = (
-        engine_args.create_model_config().get_diff_sampling_param()
-    )
-
     # Configure ec_both mode with DynamoMultimodalEmbeddingCacheConnector.
     # Must happen BEFORE engine setup so vLLM sees ec_transfer_config.
     if (
@@ -527,6 +542,7 @@ def setup_vllm_engine(
     # Taken from build_async_engine_client_from_engine_args()
     usage_context = UsageContext.OPENAI_API_SERVER
     vllm_config = engine_args.create_engine_config(usage_context=usage_context)
+    default_sampling_params = vllm_config.model_config.get_diff_sampling_param()
 
     # Set up consolidator endpoints if KVBM (DynamoConnector) is enabled
     consolidator_endpoints = None
@@ -680,7 +696,7 @@ async def register_vllm_model(
         model_input,
         model_type,
         generate_endpoint,
-        config.model,
+        _register_model_source_path(config, vllm_config),
         config.served_model_name,
         context_length=vllm_config.model_config.max_model_len,
         kv_cache_block_size=runtime_values["kv_event_block_size"],

components/src/dynamo/vllm/tests/test_vllm_unit.py

@@ -336,6 +336,69 @@ def test_should_register_model_fetch_weights_for_default_load_format():
     assert should_register_model_ignore_weights(config) is False
 
 
+def test_setup_vllm_engine_reuses_engine_config_model_config(monkeypatch):
+    from dynamo.vllm import main as vllm_main
+
+    class FakeModelConfig:
+        def get_diff_sampling_param(self):
+            return {"temperature": 0.7}
+
+    vllm_config = SimpleNamespace(
+        additional_config={},
+        cache_config=SimpleNamespace(block_size=None),
+        model_config=FakeModelConfig(),
+    )
+
+    class FakeEngineArgs:
+        enable_log_requests = False
+        enable_lora = False
+        disable_log_stats = True
+        load_format = "modelexpress"
+
+        def create_model_config(self):
+            raise AssertionError("setup_vllm_engine must not create ModelConfig twice")
+
+        def create_engine_config(self, usage_context):
+            return vllm_config
+
+    engine_client = SimpleNamespace(vllm_config=vllm_config)
+
+    class FakeAsyncLLM:
+        @staticmethod
+        def from_vllm_config(**_kwargs):
+            return engine_client
+
+    class FakeMetrics:
+        def __init__(self, **_kwargs):
+            pass
+
+        def set_model_load_time(self, _load_time):
+            pass
+
+    monkeypatch.setattr(vllm_main, "setup_multiprocess_prometheus", lambda: None)
+    monkeypatch.setattr(vllm_main, "LLMBackendMetrics", FakeMetrics)
+    monkeypatch.setattr(vllm_main, "_uses_dynamo_connector", lambda _args: False)
+    monkeypatch.setattr(vllm_main, "AsyncLLM", FakeAsyncLLM)
+    monkeypatch.setattr(
+        vllm_main,
+        "get_engine_cache_info",
+        lambda _engine: {"block_size": 16},
+    )
+
+    config = SimpleNamespace(
+        component="backend",
+        engine_args=FakeEngineArgs(),
+        gms_shadow_mode=False,
+        multimodal_embedding_cache_capacity_gb=0,
+        route_to_encoder=False,
+        served_model_name="Qwen/Qwen3-0.6B",
+    )
+
+    _, _, default_sampling_params, _, _ = vllm_main.setup_vllm_engine(config)
+
+    assert default_sampling_params == {"temperature": 0.7}
+
+
 # --disaggregation-mode tests