autoinference-utils 0.2.0__tar.gz → 0.2.2__tar.gz

{autoinference_utils-0.2.0 → autoinference_utils-0.2.2}/.gitignore

@@ -8,13 +8,13 @@ __pycache__/
 
 # Distribution / packaging
 .Python
+node_modules/
 build/
 develop-eggs/
 dist/
 downloads/
 eggs/
 .eggs/
-lib/
 lib64/
 parts/
 sdist/
@@ -195,9 +195,9 @@ cython_debug/
 .abstra/
 
 # Visual Studio Code
-#   Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#   Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
 #   that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
-#   and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#   and can be added to the global gitignore or merged into this file. However, if you prefer,
 #   you could uncomment the following to ignore the entire vscode folder
 # .vscode/
 
@@ -219,8 +219,6 @@ __marimo__/
 __MACOSX/
 .AppleDouble
 .LSOverride
-Icon[
-]
 
 # Thumbnails
 ._*
@@ -242,4 +240,9 @@ Temporary Items
 .apdisk
 
 benchmark_results/
-docs/
+docs/*
+!docs/data-model.md
+!docs/dogfood-feedback.md
+!docs/main-architecture-explainer.md
+!docs/recipe-seeding.md
+!docs/sweep-learnings.md

{autoinference_utils-0.2.0 → autoinference_utils-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: autoinference-utils
-Version: 0.2.0
+Version: 0.2.2
 Summary: Shared endpoint abstractions for autoinference deployments
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

{autoinference_utils-0.2.0 → autoinference_utils-0.2.2}/autoinference_utils/endpoint.py

@@ -238,9 +238,17 @@ class VLLMEndpoint(Endpoint):
         health_poll_interval: float = 5.0,
         health_request_timeout: float = 5.0,
     ):
-        super().__init__(base_url=f"http://localhost:{worker_port}")
+        # In bench mode vLLM runs on worker_port + 10000 and a /bench proxy
+        # listens on the original worker_port; both fall through to the same
+        # port when bench mode is off.
+        self.bench_mode = os.environ.get(BENCH_MODE_ENV) == "1"
+        self.listen_port = worker_port
+        vllm_port = (
+            worker_port + BENCH_MODE_PORT_OFFSET if self.bench_mode else worker_port
+        )
+        super().__init__(base_url=f"http://localhost:{vllm_port}")
         self.model = model
-        self.worker_port = worker_port
+        self.worker_port = vllm_port
         self.extra_server_args = (
             dict(extra_server_args) if extra_server_args else {}
         )
@@ -248,6 +256,10 @@ class VLLMEndpoint(Endpoint):
         self.health_poll_interval = health_poll_interval
         self.health_request_timeout = health_request_timeout
         self._proc: Optional[subprocess.Popen] = None
+        self._bench_server: Optional[http.server.ThreadingHTTPServer] = None
+
+        if os.environ.get(DUMMY_WEIGHTS_ENV) == "1":
+            self.extra_server_args.setdefault("--load-format", "dummy")
 
     def _build_cmd(self) -> list[str]:
         cmd = [
@@ -264,13 +276,6 @@ class VLLMEndpoint(Endpoint):
         return cmd
 
     def start(self):
-        if (
-            os.environ.get(BENCH_MODE_ENV) == "1"
-            or os.environ.get(DUMMY_WEIGHTS_ENV) == "1"
-        ):
-            raise NotImplementedError(
-                "bench_mode and dummy weights are not supported for VLLMEndpoint"
-            )
         cmd = self._build_cmd()
         print(f"[vllm] starting: {shlex.join(cmd)}")
         self._proc = subprocess.Popen(cmd)
@@ -281,8 +286,16 @@ class VLLMEndpoint(Endpoint):
             poll_interval=self.health_poll_interval,
             request_timeout=self.health_request_timeout,
         )
+        if self.bench_mode:
+            self._bench_server = start_bench_proxy(
+                listen_port=self.listen_port,
+                upstream_port=self.worker_port,
+            )
 
     def stop(self):
+        if self._bench_server is not None:
+            self._bench_server.shutdown()
+            self._bench_server = None
         terminate_process(self._proc)
         self._proc = None
 
@@ -371,15 +384,23 @@ class RouterEndpoint(Endpoint):
         health_timeout: float = 10 * 60,
         health_poll_interval: float = 5.0,
     ):
-        super().__init__(base_url=f"http://localhost:{router_port}")
+        self.bench_mode = os.environ.get(BENCH_MODE_ENV) == "1"
+        self.listen_port = router_port
+        actual_router_port = (
+            router_port + BENCH_MODE_PORT_OFFSET if self.bench_mode else router_port
+        )
+        super().__init__(base_url=f"http://localhost:{actual_router_port}")
         self.pd_config = list(pd_config)
-        self.worker_port = worker_port
-        self.router_port = router_port
+        self.worker_port = (
+            worker_port + BENCH_MODE_PORT_OFFSET if self.bench_mode else worker_port
+        )
+        self.router_port = actual_router_port
         self.prefill_bootstrap_port = prefill_bootstrap_port
         self.api_key = api_key
         self.health_timeout = health_timeout
         self.health_poll_interval = health_poll_interval
         self._proc: Optional[subprocess.Popen] = None
+        self._bench_server: Optional[http.server.ThreadingHTTPServer] = None
 
     def _build_cmd(self) -> list[str]:
         cmd = [
@@ -434,8 +455,16 @@ class RouterEndpoint(Endpoint):
             timeout=self.health_timeout,
             poll_interval=self.health_poll_interval,
         )
+        if self.bench_mode:
+            self._bench_server = start_bench_proxy(
+                listen_port=self.listen_port,
+                upstream_port=self.router_port,
+            )
 
     def stop(self):
+        if self._bench_server is not None:
+            self._bench_server.shutdown()
+            self._bench_server = None
         terminate_process(self._proc)
         self._proc = None
 
@@ -579,7 +608,7 @@ def run_bench(
     benchmark: str, args: list[str], target: str, output_dir: str
 ) -> tuple[int, bytes]:
     """Run a benchmark task (an ``invoke`` @task in
-    autoinference.tools.benchmarks) as a subprocess. Used by the in-container
+    autoinference.benchmarks) as a subprocess. Used by the in-container
     /bench proxy and the Modal bench worker (--target mode). stdout/stderr
     stream to the container's own stdout so Modal logs capture them live;
     returns (rc, run_report.json body) at the end."""
@@ -591,7 +620,7 @@ def run_bench(
             {
                 "ok": False,
                 "error": (
-                    "autoinference.tools.benchmarks not importable in container"
+                    "autoinference.benchmarks not importable in container"
                 ),
                 "volume_path": output_dir,
             }
@@ -629,7 +658,7 @@ def run_bench(
 
 def _benchmarks_search_root() -> str | None:
     try:
-        import autoinference.tools.benchmarks.tasks as tasks_mod
+        import autoinference.benchmarks.tasks as tasks_mod
     except Exception:
         return None
     return os.path.dirname(os.path.abspath(tasks_mod.__file__))

{autoinference_utils-0.2.0 → autoinference_utils-0.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "autoinference-utils"
-version = "0.2.0"
+version = "0.2.2"
 description = "Shared endpoint abstractions for autoinference deployments"
 readme = "README.md"
 requires-python = ">=3.10"