autoinference-utils 0.2.5__tar.gz → 0.2.7__tar.gz

{autoinference_utils-0.2.5 → autoinference_utils-0.2.7}/.gitignore

@@ -12,6 +12,7 @@ node_modules/
 build/
 develop-eggs/
 dist/
+!.github/actions/paths-filter/dist/
 downloads/
 eggs/
 .eggs/
@@ -149,8 +150,10 @@ activemq-data/
 
 # Environments
 .env
+.env.dev
 .envrc
 .venv
+.autoinference/
 env/
 venv/
 ENV/
@@ -239,10 +242,22 @@ Network Trash Folder
 Temporary Items
 .apdisk
 
+# Ephemeral files from benchmark runs
 benchmark_results/
-docs/*
-!docs/data-model.md
-!docs/dogfood-feedback.md
-!docs/main-architecture-explainer.md
-!docs/recipe-seeding.md
-!docs/sweep-learnings.md
+src/autoinference/deployments/
+
+# Modal Skills
+.agents/skills/modal/
+.agents/skills/.modal-skill-*
+
+# cook/etc is an extension point: local tools live there untracked.
+packages/cook/src/cook/etc/*
+# Local tools that existed before /etc are explicitly added back
+!packages/cook/src/cook/etc/__init__.py
+!packages/cook/src/cook/etc/tests/
+!packages/cook/src/cook/etc/inspect_hf_model.py
+!packages/cook/src/cook/etc/download.py
+!packages/cook/src/cook/etc/quantize.py
+!packages/cook/src/cook/etc/graft_mtp.py
+!packages/cook/src/cook/etc/tune_moe.py
+!packages/cook/src/cook/etc/tool_calling.py

{autoinference_utils-0.2.5 → autoinference_utils-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: autoinference-utils
-Version: 0.2.5
+Version: 0.2.7
 Summary: Shared endpoint abstractions for autoinference deployments
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -9,6 +9,12 @@ Description-Content-Type: text/markdown
 
 Shared endpoint abstractions (`SGLangEndpoint`, `VLLMEndpoint`) for autoinference deployments.
 
+For SGLang models whose custom Transformers code cannot load from symlinked
+Hugging Face snapshots, set `AUTOINFERENCE_MATERIALIZE_MODEL_PATH=1`. The endpoint
+copies local base/draft model code and metadata into unique temporary directories
+and symlinks top-level safetensors weights. Repository IDs pass through unchanged.
+Leave the setting unset to preserve the original paths.
+
 ## Publishing
 
 Bump `version` in `pyproject.toml` and merge to `main` — CI publishes automatically via trusted publishing.

autoinference_utils-0.2.7/README.md

@@ -0,0 +1,15 @@
+# autoinference-utils
+
+Shared endpoint abstractions (`SGLangEndpoint`, `VLLMEndpoint`) for autoinference deployments.
+
+For SGLang models whose custom Transformers code cannot load from symlinked
+Hugging Face snapshots, set `AUTOINFERENCE_MATERIALIZE_MODEL_PATH=1`. The endpoint
+copies local base/draft model code and metadata into unique temporary directories
+and symlinks top-level safetensors weights. Repository IDs pass through unchanged.
+Leave the setting unset to preserve the original paths.
+
+## Publishing
+
+Bump `version` in `pyproject.toml` and merge to `main` — CI publishes automatically via trusted publishing.
+
+To publish manually: `cd autoinference_utils && uv build && uv publish`

{autoinference_utils-0.2.5 → autoinference_utils-0.2.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "autoinference-utils"
-version = "0.2.5"
+version = "0.2.7"
 description = "Shared endpoint abstractions for autoinference deployments"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -10,7 +10,7 @@ requires = ["hatchling"]
 build-backend = "hatchling.build"
 
 [tool.hatch.build.targets.wheel]
-packages = ["autoinference_utils"]
+packages = ["src/autoinference_utils"]
 
 [tool.hatch.build.targets.sdist]
-include = ["autoinference_utils", "README.md", "pyproject.toml"]
+include = ["src", "README.md", "pyproject.toml"]

{autoinference_utils-0.2.5 → autoinference_utils-0.2.7/src}/autoinference_utils/endpoint.py

@@ -23,19 +23,62 @@ import urllib.request
 from abc import ABC
 from typing import Any, Callable, Literal, Mapping, Optional, Sequence
 
+from autoinference_utils.model_path import materialize_model_path
+
 BENCH_MODE_ENV = "ENDPOINT_BENCH_MODE"
 DUMMY_WEIGHTS_ENV = "ENDPOINT_DUMMY"
 BENCH_MODE_PORT_OFFSET = 10000
 
-MODAL_FLASH_REQUEST_UUID_HEADER= "modal-flash-request-uuid"
+MODAL_FLASH_REQUEST_UUID_HEADER = "modal-flash-request-uuid"
 MODAL_SESSION_ID_HEADER = "modal-session-id"
 
-# Endpoint class names whose servers emit OAI streaming chunks that
-# sglang.bench_serving can't parse. deploy_and_bench auto-injects
-# --openai-stream-compat for deployments importing any of these.
 ENDPOINTS_REQUIRING_OAI_STREAM_COMPAT: tuple[str, ...] = ("TRTLLMEndpoint",)
 
 
+def server_arg_tokens(flag: str, value: str) -> list[str]:
+    """Render one server arg as argv tokens.
+
+    A value splits on whitespace, so a repeatable flag takes its values as one
+    entry. A lone value starting with `-` keeps the attached `--flag=value`
+    spelling, the only one an argparse-style parser does not read as a flag.
+    """
+    if value and "=" in flag:
+        raise ValueError(f"server arg {flag!r} already carries a value: {value!r}")
+    if value == "":
+        return [flag]
+    tokens = value.split()
+    if len(tokens) == 1 and tokens[0].startswith("-"):
+        return [f"{flag}={tokens[0]}"]
+    return [flag, *tokens]
+
+
+def _flag_name(key: str) -> str:
+    """Return the flag a key sets, whose value may be attached as `--flag=value`."""
+    return key.partition("=")[0]
+
+
+def get_server_arg(server_args: Mapping[str, str], flag: str) -> tuple[str, str] | None:
+    """Return the key and value setting a flag under either spelling."""
+    for key, value in server_args.items():
+        if _flag_name(key) == flag:
+            return key, value
+    return None
+
+
+def has_server_arg(server_args: Mapping[str, str], flag: str) -> bool:
+    """Report whether a flag is set under either spelling, plain or attached."""
+    return get_server_arg(server_args, flag) is not None
+
+
+def _merge_server_args(
+    defaults: Mapping[str, str], overrides: Mapping[str, str]
+) -> dict[str, str]:
+    """Overlay overrides on defaults by flag name, so an attached spelling wins."""
+    merged = {_flag_name(key): (key, value) for key, value in defaults.items()}
+    merged.update((_flag_name(key), (key, value)) for key, value in overrides.items())
+    return dict(merged.values())
+
+
 class Endpoint(ABC):
     """A thing with a URL that you can start and stop."""
 
@@ -110,11 +153,15 @@ class SGLangEndpoint(Endpoint):
         super().__init__(base_url=f"http://localhost:{sglang_port}")
 
         self.worker_port = sglang_port
-        self.model_path = model_path
+        self.model_path = materialize_model_path(model_path)
         self.tp = tp
         self.ep = ep
         self.dp = dp
-        self.speculative_model_path = speculative_model_path
+        self.speculative_model_path = (
+            materialize_model_path(speculative_model_path)
+            if speculative_model_path is not None
+            else None
+        )
         self.load_format = load_format
         self.nnodes = nnodes
         self.node_rank = node_rank
@@ -123,9 +170,7 @@ class SGLangEndpoint(Endpoint):
         self.disaggregation_mode = disaggregation_mode
         self.prefill_bootstrap_port = prefill_bootstrap_port
         self.launcher_module = launcher_module
-        self.extra_server_args = (
-            dict(extra_server_args) if extra_server_args else {}
-        )
+        self.extra_server_args = dict(extra_server_args) if extra_server_args else {}
         self.health_timeout = health_timeout
         self.health_poll_interval = health_poll_interval
         self.health_request_timeout = health_request_timeout
@@ -140,28 +185,35 @@ class SGLangEndpoint(Endpoint):
         self._bench_server: Optional[http.server.ThreadingHTTPServer] = None
 
         if self.disaggregation_mode not in (None, "prefill", "decode"):
-            raise ValueError(
-                "disaggregation_mode must be None, 'prefill', or 'decode'"
-            )
+            raise ValueError("disaggregation_mode must be None, 'prefill', or 'decode'")
 
         if os.environ.get(DUMMY_WEIGHTS_ENV) == "1":
-            if (
-                self.load_format is None
-                and "--load-format" not in self.extra_server_args
+            if self.load_format is None and not has_server_arg(
+                self.extra_server_args, "--load-format"
             ):
                 self.load_format = "dummy"
-            self.extra_server_args.setdefault("--model-loader-extra-config", "{}")
-        
+            if not has_server_arg(
+                self.extra_server_args, "--model-loader-extra-config"
+            ):
+                self.extra_server_args["--model-loader-extra-config"] = "{}"
+
         # Request logging enabled, log only metadata by default
         if self.log_requests_level >= 0:
-            # Enable request logging    
-            if "--log-requests" not in self.extra_server_args:
-                self.extra_server_args.setdefault("--log-requests", "")
+            # Enable request logging
+            if not has_server_arg(self.extra_server_args, "--log-requests"):
+                self.extra_server_args["--log-requests"] = ""
             # Set request logging level
-            if "--log-requests-level" not in self.extra_server_args:
-                self.extra_server_args.setdefault("--log-requests-level", str(self.log_requests_level))
+            level = get_server_arg(self.extra_server_args, "--log-requests-level")
+            if level is None:
+                self.extra_server_args["--log-requests-level"] = str(
+                    self.log_requests_level
+                )
             else:
-                print(f"[endpoint] --log-requests-level already set to {self.extra_server_args['--log-requests-level']} in server args")
+                key, value = level
+                print(
+                    "[endpoint] --log-requests-level already set in server args: "
+                    f"{value or key}"
+                )
 
             # Add Modal-forwarded request headers to SGLang logged headers
             sglang_log_request_headers = [
@@ -177,19 +229,21 @@ class SGLangEndpoint(Endpoint):
             self.log_request_headers = ",".join(sglang_log_request_headers)
             os.environ["SGLANG_LOG_REQUEST_HEADERS"] = self.log_request_headers
 
-
     def _build_cmd(self) -> list[str]:
         cmd = [
-            "python", "-m", self.launcher_module,
-            "--host", "0.0.0.0",
-            "--port", str(self.worker_port),
-            "--model-path", self.model_path,
+            "python",
+            "-m",
+            self.launcher_module,
+            "--host",
+            "0.0.0.0",
+            "--port",
+            str(self.worker_port),
+            "--model-path",
+            self.model_path,
         ]
 
         if self.speculative_model_path is not None:
-            cmd.extend(
-                ["--speculative-draft-model-path", self.speculative_model_path]
-            )
+            cmd.extend(["--speculative-draft-model-path", self.speculative_model_path])
         if self.load_format is not None:
             cmd.extend(["--load-format", self.load_format])
         if self.tp is not None:
@@ -214,19 +268,20 @@ class SGLangEndpoint(Endpoint):
                 raise ValueError("dist_init_host is required when nnodes > 1")
             cmd.extend(
                 [
-                    "--nnodes", str(self.nnodes),
-                    "--node-rank", str(self.node_rank),
+                    "--nnodes",
+                    str(self.nnodes),
+                    "--node-rank",
+                    str(self.node_rank),
                     "--dist-init-addr",
                     f"{self.dist_init_host}:{self.dist_init_port}",
                 ]
             )
 
-        merged = {**self.DEFAULT_OPERATIONAL_ARGS, **self.extra_server_args}
+        merged = _merge_server_args(
+            self.DEFAULT_OPERATIONAL_ARGS, self.extra_server_args
+        )
         for key, value in merged.items():
-            if value == "":
-                cmd.append(key)
-            else:
-                cmd.extend([key, *value.split()])
+            cmd.extend(server_arg_tokens(key, value))
 
         return cmd
 
@@ -287,9 +342,7 @@ class VLLMEndpoint(Endpoint):
         super().__init__(base_url=f"http://localhost:{vllm_port}")
         self.model = model
         self.worker_port = vllm_port
-        self.extra_server_args = (
-            dict(extra_server_args) if extra_server_args else {}
-        )
+        self.extra_server_args = dict(extra_server_args) if extra_server_args else {}
         self.health_timeout = health_timeout
         self.health_poll_interval = health_poll_interval
         self.health_request_timeout = health_request_timeout
@@ -297,20 +350,23 @@ class VLLMEndpoint(Endpoint):
         self._bench_server: Optional[http.server.ThreadingHTTPServer] = None
 
         if os.environ.get(DUMMY_WEIGHTS_ENV) == "1":
-            self.extra_server_args.setdefault("--load-format", "dummy")
+            if not has_server_arg(self.extra_server_args, "--load-format"):
+                self.extra_server_args["--load-format"] = "dummy"
 
     def _build_cmd(self) -> list[str]:
         cmd = [
-            "python", "-m", "vllm.entrypoints.openai.api_server",
-            "--host", "0.0.0.0",
-            "--port", str(self.worker_port),
-            "--model", self.model,
+            "python",
+            "-m",
+            "vllm.entrypoints.openai.api_server",
+            "--host",
+            "0.0.0.0",
+            "--port",
+            str(self.worker_port),
+            "--model",
+            self.model,
         ]
         for key, value in self.extra_server_args.items():
-            if value == "":
-                cmd.append(key)
-            else:
-                cmd.extend([key, *value.split()])
+            cmd.extend(server_arg_tokens(key, value))
         return cmd
 
     def start(self):
@@ -361,9 +417,7 @@ class TRTLLMEndpoint(Endpoint):
         super().__init__(base_url=f"http://localhost:{worker_port}")
         self.model = model
         self.worker_port = worker_port
-        self.extra_server_args = (
-            dict(extra_server_args) if extra_server_args else {}
-        )
+        self.extra_server_args = dict(extra_server_args) if extra_server_args else {}
         self.health_timeout = health_timeout
         self.health_poll_interval = health_poll_interval
         self.health_request_timeout = health_request_timeout
@@ -373,14 +427,13 @@ class TRTLLMEndpoint(Endpoint):
         cmd = [
             "trtllm-serve",
             self.model,
-            "--host", "0.0.0.0",
-            "--port", str(self.worker_port),
+            "--host",
+            "0.0.0.0",
+            "--port",
+            str(self.worker_port),
         ]
         for key, value in self.extra_server_args.items():
-            if value == "":
-                cmd.append(key)
-            else:
-                cmd.extend([key, *value.split()])
+            cmd.extend(server_arg_tokens(key, value))
         return cmd
 
     def health_check(self) -> str | None:
@@ -442,18 +495,30 @@ class RouterEndpoint(Endpoint):
 
     def _build_cmd(self) -> list[str]:
         cmd = [
-            "python", "-m", "sglang_router.launch_router",
-            "--host", "0.0.0.0",
-            "--port", str(self.router_port),
-            "--prefill-policy", "cache_aware",
-            "--decode-policy", "round_robin",
-            "--max-concurrent-requests", "128",
-            "--rate-limit-tokens-per-second", "0",
-            "--queue-size", "0",
-            "--health-check-timeout-secs", "600",
-            "--log-level", "info",
+            "python",
+            "-m",
+            "sglang_router.launch_router",
+            "--host",
+            "0.0.0.0",
+            "--port",
+            str(self.router_port),
+            "--prefill-policy",
+            "cache_aware",
+            "--decode-policy",
+            "round_robin",
+            "--max-concurrent-requests",
+            "128",
+            "--rate-limit-tokens-per-second",
+            "0",
+            "--queue-size",
+            "0",
+            "--health-check-timeout-secs",
+            "600",
+            "--log-level",
+            "info",
             "--disable-circuit-breaker",
-            "--request-timeout-secs", "3600",
+            "--request-timeout-secs",
+            "3600",
         ]
 
         if self.api_key is not None:
@@ -465,9 +530,7 @@ class RouterEndpoint(Endpoint):
         for role, node_ip in self.pd_config:
             node_url = f"http://{node_ip}:{self.worker_port}"
             if role == "prefill":
-                cmd.extend(
-                    ["--prefill", node_url, str(self.prefill_bootstrap_port)]
-                )
+                cmd.extend(["--prefill", node_url, str(self.prefill_bootstrap_port)])
             elif role == "decode":
                 cmd.extend(["--decode", node_url])
             elif role == "worker":
@@ -511,6 +574,7 @@ class RouterEndpoint(Endpoint):
 # and this proxy listens on worker_port. POST /bench shells out to a benchmark
 # task via run_bench(); every other path is forwarded to the upstream server.
 
+
 def start_bench_proxy(
     *,
     listen_port: int,
@@ -521,17 +585,12 @@ def start_bench_proxy(
         (_BenchProxyHandler,),
         {"upstream_port": upstream_port},
     )
-    server = http.server.ThreadingHTTPServer(
-        ("0.0.0.0", listen_port), handler_cls
-    )
+    server = http.server.ThreadingHTTPServer(("0.0.0.0", listen_port), handler_cls)
     thread = threading.Thread(
         target=server.serve_forever, daemon=True, name="bench-proxy"
     )
     thread.start()
-    print(
-        f"[bench-proxy] listening on :{listen_port}, "
-        f"forwarding to :{upstream_port}"
-    )
+    print(f"[bench-proxy] listening on :{listen_port}, forwarding to :{upstream_port}")
     return server
 
 
@@ -560,26 +619,17 @@ class _BenchProxyHandler(http.server.BaseHTTPRequestHandler):
         try:
             payload = self._read_json_body()
         except (ValueError, json.JSONDecodeError) as exc:
-            return self._send_json(
-                400, {"ok": False, "error": f"bad body: {exc}"}
-            )
+            return self._send_json(400, {"ok": False, "error": f"bad body: {exc}"})
 
         benchmark = payload.get("benchmark") or ""
         args = payload.get("args") or []
-        target = (
-            payload.get("target")
-            or f"http://localhost:{self.upstream_port}"
-        )
+        target = payload.get("target") or f"http://localhost:{self.upstream_port}"
         output_dir = payload.get("output_dir") or "/tmp/bench-output"
 
         if not benchmark:
-            return self._send_json(
-                400, {"ok": False, "error": "benchmark is required"}
-            )
+            return self._send_json(400, {"ok": False, "error": "benchmark is required"})
         if not isinstance(args, list):
-            return self._send_json(
-                400, {"ok": False, "error": "args must be a list"}
-            )
+            return self._send_json(400, {"ok": False, "error": "args must be a list"})
 
         rc, body = run_bench(benchmark, list(args), target, output_dir)
         status = 200 if rc == 0 else 500
@@ -630,9 +680,7 @@ class _BenchProxyHandler(http.server.BaseHTTPRequestHandler):
         return json.loads(raw)
 
     def _send_json(self, status: int, obj: Mapping[str, Any]):
-        self._send_raw(
-            status, "application/json", json.dumps(obj).encode()
-        )
+        self._send_raw(status, "application/json", json.dumps(obj).encode())
 
     def _send_raw(self, status: int, content_type: str, body: bytes):
         self.send_response(status)
@@ -657,18 +705,21 @@ def run_bench(
         return 500, json.dumps(
             {
                 "ok": False,
-                "error": (
-                    "autoinference.benchmarks not importable in container"
-                ),
+                "error": ("autoinference.benchmarks not importable in container"),
                 "volume_path": output_dir,
             }
         ).encode()
 
     cmd = [
-        sys.executable, "-m", "invoke",
-        "--search-root", search_root,
-        "-c", "tasks",
-        benchmark, *[str(a) for a in args],
+        sys.executable,
+        "-m",
+        "invoke",
+        "--search-root",
+        search_root,
+        "-c",
+        "tasks",
+        benchmark,
+        *[str(a) for a in args],
     ]
     env = {
         **os.environ,
@@ -718,9 +769,7 @@ def wait_ready(
 
     while time.time() < deadline:
         try:
-            error = _health_check(
-                url, request_timeout=request_timeout, process=process
-            )
+            error = _health_check(url, request_timeout=request_timeout, process=process)
         except subprocess.CalledProcessError as exc:
             print(
                 f"[endpoint] !!! server process exited with code "
@@ -734,13 +783,11 @@ def wait_ready(
         time.sleep(poll_interval)
 
     print(
-        f"[endpoint] !!! health check timed out after {timeout}s; "
-        f"last={last_error}",
+        f"[endpoint] !!! health check timed out after {timeout}s; last={last_error}",
         flush=True,
     )
     raise TimeoutError(
-        f"Health check timed out after {timeout}s for {url}. "
-        f"Last error: {last_error}"
+        f"Health check timed out after {timeout}s for {url}. Last error: {last_error}"
     )
 
 
@@ -768,9 +815,7 @@ def warmup_chat_completions(
             request_timeout=request_timeout,
             max_attempts=max_attempts_per_request,
             retry_delay=retry_delay,
-            description=(
-                f"warmup request {request_idx + 1}/{successful_requests}"
-            ),
+            description=(f"warmup request {request_idx + 1}/{successful_requests}"),
         )
 
 
@@ -792,9 +837,7 @@ def validate_embeddings_endpoint(
             1 if all(isinstance(item, int) for item in inputs) else len(inputs)
         )
     else:
-        raise ValueError(
-            "embedding validation payload must contain non-empty input"
-        )
+        raise ValueError("embedding validation payload must contain non-empty input")
 
     request_headers = {"Content-Type": "application/json"}
     if headers:
@@ -857,10 +900,7 @@ def start_heartbeat_thread(
             try:
                 error = health_check_fn()
             except subprocess.CalledProcessError as exc:
-                print(
-                    f"[heartbeat] server process exited with code "
-                    f"{exc.returncode}"
-                )
+                print(f"[heartbeat] server process exited with code {exc.returncode}")
                 on_failure()
                 return
             if error is None:
@@ -872,10 +912,7 @@ def start_heartbeat_thread(
                 f"({consecutive_failures}/{max_consecutive_failures})"
             )
             if consecutive_failures >= max_consecutive_failures:
-                print(
-                    "[heartbeat] sustained health-check failure, "
-                    "invoking on_failure"
-                )
+                print("[heartbeat] sustained health-check failure, invoking on_failure")
                 on_failure()
                 return
 
@@ -960,8 +997,7 @@ def _wait_ready_url(
         last_error = error
         time.sleep(poll_interval)
     raise TimeoutError(
-        f"Timed out after {timeout}s waiting for {url}. "
-        f"Last error: {last_error}"
+        f"Timed out after {timeout}s waiting for {url}. Last error: {last_error}"
     )
 
 
@@ -979,9 +1015,7 @@ def _post_json(
     timeout: float | None = None,
 ) -> bytes:
     body = json.dumps(payload).encode("utf-8")
-    req = urllib.request.Request(
-        url, data=body, headers=dict(headers), method="POST"
-    )
+    req = urllib.request.Request(url, data=body, headers=dict(headers), method="POST")
     with urllib.request.urlopen(req, timeout=timeout) as resp:
         return resp.read()
 

autoinference_utils-0.2.7/src/autoinference_utils/model_path.py

@@ -0,0 +1,38 @@
+"""Prepare local model snapshots for Transformers' remote-code loader."""
+
+import os
+import shutil
+import tempfile
+from pathlib import Path
+
+MATERIALIZE_MODEL_PATH_ENV = "AUTOINFERENCE_MATERIALIZE_MODEL_PATH"
+
+
+def materialize_model_path(path: str) -> str:
+    """Copy snapshot metadata/code into a real directory; retain weight links.
+
+    Transformers can resolve Python-module symlinks into the HF blob directory,
+    losing the relative imports next to the original snapshot entry. This opt-in
+    keeps small files together without copying large safetensors shards. The
+    unique directory lasts for the container's lifetime.
+    """
+    source = Path(path)
+    if os.environ.get(MATERIALIZE_MODEL_PATH_ENV) != "1" or not source.is_dir():
+        return path
+
+    source = source.resolve()
+    destination = Path(tempfile.mkdtemp(prefix="autoinference-model-"))
+    try:
+        for entry in source.iterdir():
+            target = destination / entry.name
+            if entry.is_file() and entry.suffix == ".safetensors":
+                target.symlink_to(entry.resolve(strict=True))
+            elif entry.is_dir():
+                shutil.copytree(entry, target, symlinks=False)
+            else:
+                shutil.copy2(entry, target, follow_symlinks=True)
+    except BaseException:
+        shutil.rmtree(destination)
+        raise
+    print(f"Materialized model snapshot {source} -> {destination}")
+    return str(destination)

autoinference_utils-0.2.7/src/autoinference_utils/tests/test_model_path.py

@@ -0,0 +1,84 @@
+import shutil
+from pathlib import Path
+
+import pytest
+
+from autoinference_utils.endpoint import SGLangEndpoint
+from autoinference_utils.model_path import (
+    MATERIALIZE_MODEL_PATH_ENV,
+    materialize_model_path,
+)
+
+
+def snapshot(tmp_path):
+    blobs = tmp_path / "blobs"
+    source = tmp_path / "snapshots" / "revision"
+    blobs.mkdir()
+    source.mkdir(parents=True)
+    for name, content in {
+        "config.json": '{"model_type": "kimi"}',
+        "configuration_kimi.py": "from .helper import VALUE\n",
+        "helper.py": "VALUE = 42\n",
+        "model.safetensors": "weight placeholder",
+    }.items():
+        blob = blobs / ("hash-" + name)
+        blob.write_text(content)
+        (source / name).symlink_to(blob)
+    return source
+
+
+def test_opt_in_and_repository_ids(tmp_path, monkeypatch):
+    source = snapshot(tmp_path)
+    monkeypatch.delenv(MATERIALIZE_MODEL_PATH_ENV, raising=False)
+    assert materialize_model_path(str(source)) == str(source)
+    monkeypatch.setenv(MATERIALIZE_MODEL_PATH_ENV, "1")
+    assert materialize_model_path("moonshotai/Kimi-K3") == "moonshotai/Kimi-K3"
+
+
+def test_local_code_is_materialized_and_weights_stay_linked(tmp_path, monkeypatch):
+    source = snapshot(tmp_path)
+    monkeypatch.setenv(MATERIALIZE_MODEL_PATH_ENV, "1")
+    monkeypatch.setattr("tempfile.tempdir", str(tmp_path))
+    result = Path(materialize_model_path(str(source)))
+    assert result != source
+    for name in ("config.json", "configuration_kimi.py", "helper.py"):
+        assert not (result / name).is_symlink()
+        assert (result / name).resolve().parent == result
+        assert (result / name).read_bytes() == (source / name).read_bytes()
+    assert (result / "model.safetensors").is_symlink()
+    assert (result / "model.safetensors").resolve() == (
+        source / "model.safetensors"
+    ).resolve()
+    assert (source / "configuration_kimi.py").is_symlink()
+
+
+def test_endpoint_uses_materialized_base_and_draft(tmp_path, monkeypatch):
+    source = snapshot(tmp_path)
+    monkeypatch.setenv(MATERIALIZE_MODEL_PATH_ENV, "1")
+    monkeypatch.setattr("tempfile.tempdir", str(tmp_path))
+    endpoint = SGLangEndpoint(
+        model_path=str(source), speculative_model_path=str(source)
+    )
+    assert endpoint.model_path != str(source)
+    assert endpoint.speculative_model_path != str(source)
+    assert endpoint.model_path != endpoint.speculative_model_path
+    cmd = endpoint._build_cmd()
+    assert cmd[cmd.index("--model-path") + 1] == endpoint.model_path
+    assert (
+        cmd[cmd.index("--speculative-draft-model-path") + 1]
+        == endpoint.speculative_model_path
+    )
+
+
+def test_failed_copy_removes_partial_directory(tmp_path, monkeypatch):
+    source = snapshot(tmp_path)
+    monkeypatch.setenv(MATERIALIZE_MODEL_PATH_ENV, "1")
+    monkeypatch.setattr("tempfile.tempdir", str(tmp_path))
+
+    def fail(*args, **kwargs):
+        raise OSError("copy failed")
+
+    monkeypatch.setattr(shutil, "copy2", fail)
+    with pytest.raises(OSError, match="copy failed"):
+        materialize_model_path(str(source))
+    assert not list(tmp_path.glob("autoinference-model-*"))

autoinference_utils-0.2.5/README.md

@@ -1,9 +0,0 @@
-# autoinference-utils
-
-Shared endpoint abstractions (`SGLangEndpoint`, `VLLMEndpoint`) for autoinference deployments.
-
-## Publishing
-
-Bump `version` in `pyproject.toml` and merge to `main` — CI publishes automatically via trusted publishing.
-
-To publish manually: `cd autoinference_utils && uv build && uv publish`