skip downloading model weights if using mocker (only tokenizer)

Author: PeaBrane

launch/dynamo-run/src/lib.rs

@@ -40,7 +40,8 @@ pub async fn run(
         .http_port(Some(flags.http_port))
         .router_config(Some(flags.router_config()))
         .request_template(flags.request_template.clone())
-        .migration_limit(flags.migration_limit);
+        .migration_limit(flags.migration_limit)
+        .is_mocker(matches!(out_opt, Some(Output::Mocker)));
 
     // TODO: old, address this later:
     // If `in=dyn` we want the trtllm/sglang/vllm subprocess to listen on that endpoint.

lib/bindings/python/rust/llm/entrypoint.rs

@@ -156,7 +156,8 @@ pub fn make_engine<'p>(
         .request_template(args.template_file.clone())
         .kv_cache_block_size(args.kv_cache_block_size)
         .router_config(args.router_config.clone().map(|rc| rc.into()))
-        .http_port(args.http_port);
+        .http_port(args.http_port)
+        .is_mocker(matches!(args.engine_type, EngineType::Mocker));
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         let local_model = builder.build().await.map_err(to_pyerr)?;
         let inner = select_engine(distributed_runtime, args, local_model)

lib/llm/src/hub.rs

@@ -27,9 +27,19 @@ const IGNORED: [&str; 5] = [
 
 const HF_TOKEN_ENV_VAR: &str = "HF_TOKEN";
 
+/// Checks if a file is a model weight file
+fn is_weight_file(filename: &str) -> bool {
+    filename.ends_with(".bin")
+        || filename.ends_with(".safetensors")
+        || filename.ends_with(".h5")
+        || filename.ends_with(".msgpack")
+        || filename.ends_with(".ckpt.index")
+}
+
 /// Attempt to download a model from Hugging Face
 /// Returns the directory it is in
-pub async fn from_hf(name: impl AsRef<Path>) -> anyhow::Result<PathBuf> {
+/// If ignore_weights is true, model weight files will be skipped
+pub async fn from_hf(name: impl AsRef<Path>, ignore_weights: bool) -> anyhow::Result<PathBuf> {
     let name = name.as_ref();
     let token = env::var(HF_TOKEN_ENV_VAR).ok();
     let api = ApiBuilder::new()
@@ -66,6 +76,11 @@ pub async fn from_hf(name: impl AsRef<Path>) -> anyhow::Result<PathBuf> {
             continue;
         }
 
+        // If ignore_weights is true, skip weight files
+        if ignore_weights && is_weight_file(&sib.rfilename) {
+            continue;
+        }
+
         match repo.get(&sib.rfilename).await {
             Ok(path) => {
                 p = path;
@@ -83,8 +98,14 @@ pub async fn from_hf(name: impl AsRef<Path>) -> anyhow::Result<PathBuf> {
     }
 
     if !files_downloaded {
+        let file_type = if ignore_weights {
+            "non-weight"
+        } else {
+            "valid"
+        };
         return Err(anyhow::anyhow!(
-            "No valid files found for model '{}'.",
+            "No {} files found for model '{}'.",
+            file_type,
             model_name
         ));
     }

lib/llm/src/local_model.rs

@@ -47,6 +47,7 @@ pub struct LocalModelBuilder {
     kv_cache_block_size: u32,
     http_port: u16,
     migration_limit: u32,
+    is_mocker: bool,
 }
 
 impl Default for LocalModelBuilder {
@@ -62,6 +63,7 @@ impl Default for LocalModelBuilder {
             template_file: Default::default(),
             router_config: Default::default(),
             migration_limit: Default::default(),
+            is_mocker: Default::default(),
         }
     }
 }
@@ -119,6 +121,11 @@ impl LocalModelBuilder {
         self
     }
 
+    pub fn is_mocker(&mut self, is_mocker: bool) -> &mut Self {
+        self.is_mocker = is_mocker;
+        self
+    }
+
     /// Make an LLM ready for use:
     /// - Download it from Hugging Face (and NGC in future) if necessary
     /// - Resolve the path
@@ -169,7 +176,7 @@ impl LocalModelBuilder {
         let relative_path = model_path.trim_start_matches(HF_SCHEME);
         let full_path = if is_hf_repo {
             // HF download if necessary
-            super::hub::from_hf(relative_path).await?
+            super::hub::from_hf(relative_path, self.is_mocker).await?
         } else {
             fs::canonicalize(relative_path)?
         };

tests/router/test_router_e2e_with_mockers.py

@@ -5,11 +5,11 @@
 import json
 import logging
 import os
+import time
 
 import aiohttp
 import pytest
 
-from tests.conftest import download_models
 from tests.utils.managed_process import ManagedProcess
 
 pytestmark = pytest.mark.pre_merge
@@ -96,9 +96,6 @@ def test_mocker_kv_router(request, runtime_services):
     This test doesn't require GPUs and runs quickly for pre-merge validation.
     """
 
-    # Download only the Qwen model for this test
-    download_models([MODEL_NAME])
-
     # runtime_services starts etcd and nats
     logger.info("Starting mocker KV router test")
 
@@ -132,6 +129,9 @@ def test_mocker_kv_router(request, runtime_services):
         for mocker in mocker_processes:
             mocker.__enter__()
 
+        # Give 2 seconds for setups (download tokenizer)
+        time.sleep(2)
+
         # Send test requests
         test_payload = {
             "model": MODEL_NAME,