Update to support serverless

mypy.ini

@@ -18,4 +18,7 @@ implicit_reexport = False
 ignore_missing_imports = True
 
 [mypy-google.*]
+ignore_missing_imports = True
+
+[mypy-vastai.*]
 ignore_missing_imports = True

pyproject.toml

@@ -11,10 +11,12 @@ dependencies = [
     "google-cloud-core>=2.4.3",
     "google-genai>=1.59.0",
     "json-repair~=0.40.0",
-    "openai>=2.44.0",
     "pydantic>=2.11.7",
     "rapidfuzz>=3.13.0",
     "requests>=2.32.4",
+    "tenacity>=9.1.2",
+    "openai>=2.44.0",
+    "vastai>=1.2.1",
 ]
 
 [build-system]

src/genai_utils/local_models/__init__.py

@@ -0,0 +1,2 @@
+class NoOutputError(Exception):
+    pass

src/genai_utils/openai_llm.py → src/genai_utils/local_models/llm.py

@@ -16,6 +16,9 @@
     wait_exponential,
 )
 
+from genai_utils.local_models import NoOutputError
+from genai_utils.local_models.throughput import log_tokens_per_second
+
 _logger = logging.getLogger(__name__)
 
 
@@ -29,47 +32,6 @@ def _is_model_loading(exception: BaseException) -> bool:
     )
 
 
-def log_tokens_per_second(response: object, elapsed: float) -> None:
-    """
-    Logs generation throughput (tokens/second) at info level.
-
-    Prefers the server's own decode timing when available (llama.cpp reports a
-    ``timings`` object with ``predicted_per_second``), otherwise falls back to
-    ``completion_tokens`` over the measured wall-clock time, which works for any
-    OpenAI-compatible server such as vLLM.
-    """
-    if not _logger.isEnabledFor(logging.INFO):
-        return
-
-    # llama.cpp: authoritative decode rate from the server, no network overhead.
-    # The openai client stashes unknown response fields on ``model_extra``.
-    extra = getattr(response, "model_extra", None) or {}
-    timings = extra.get("timings") if isinstance(extra, dict) else None
-    if isinstance(timings, dict) and timings.get("predicted_per_second") is not None:
-        _logger.info(
-            "Generation throughput: %.1f tok/s (%d tokens in %.2fs, server timings)",
-            timings["predicted_per_second"],
-            timings.get("predicted_n", 0),
-            timings.get("predicted_ms", 0) / 1000,
-        )
-        return
-
-    # vLLM / generic fallback: completion tokens over wall-clock time.
-    usage = getattr(response, "usage", None)
-    completion_tokens = getattr(usage, "completion_tokens", None)
-    if completion_tokens and elapsed > 0:
-        _logger.info(
-            "Generation throughput: %.1f tok/s (%d completion tokens in %.2fs)",
-            completion_tokens / elapsed,
-            completion_tokens,
-            elapsed,
-        )
-
-
-class NoOutputError(Exception):
-    pass
-
-
 class LLM:
     """
     An LLM using the Open AI API.
@@ -141,7 +103,7 @@ def run_message(self, message_content: str, use_thinking: bool) -> str:
             ],
             extra_body=extra_model_config,
         )
-        log_tokens_per_second(response, time.perf_counter() - start)
+        log_tokens_per_second(response.model_dump(), time.perf_counter() - start)
         response_text = response.choices[0].message.content
         if response_text is None:
             raise NoOutputError()

src/genai_utils/local_models/serverless_llm.py

@@ -0,0 +1,113 @@
+"""
+This file contains basic functionality for making requests to a model deployed
+as a vast.ai *serverless* endpoint.
+
+Unlike a plain instance (where you'd POST straight to the vLLM/llama.cpp server,
+see ``llm.LLM``), a serverless endpoint sits behind vast's autoscaler:
+you ask the autoscaler for a ready worker and it routes the request. The
+``vastai`` SDK's async ``Serverless`` client handles that for you, so this is an
+async client.
+
+Set ``VAST_API_KEY`` in the environment (the SDK reads it; auth is handled by
+vast's routing layer, not by the model itself).
+"""
+
+import time
+
+from vastai import Serverless
+
+from genai_utils.local_models import NoOutputError
+from genai_utils.local_models.throughput import log_tokens_per_second
+
+
+class ServerlessLLM:
+    """
+    A model served via a vast.ai serverless endpoint.
+
+    The vast SDK is async, so this is an async client: await
+    :meth:`run_message`.
+    """
+
+    def __init__(
+        self,
+        endpoint_name: str,
+        model_name: str = "local-model",
+        max_tokens: int = 512,
+        temperature: float = 0.7,
+    ):
+        """
+        Args:
+            endpoint_name:
+                The name of your vast.ai serverless endpoint.
+            model_name:
+                The model or LoRA adapter to run (e.g. an adapter name like
+                "claimants", or the base model name set when the endpoint was
+                deployed).
+            max_tokens:
+                Default generation length. Also used as the autoscaler's load
+                estimate (``cost``) for each request, since generation length
+                drives GPU time.
+            temperature:
+                Default sampling temperature.
+        """
+        self.endpoint_name = endpoint_name
+        self.model_name = model_name
+        self.max_tokens = max_tokens
+        self.temperature = temperature
+
+    async def run_message(self, message_content: str, use_thinking: bool) -> str:
+        """
+        Sends the message to the serverless endpoint and returns the result.
+        Will use thinking if you ask it to.
+
+        The autoscaler waits for a ready worker before routing, so — unlike the
+        direct client — there's no "model loading" (503) retry to ride out here.
+        """
+        extra_model_config = (
+            {
+                "chat_template_kwargs": {"enable_thinking": False},
+                "reasoning_budget": 0,
+            }
+            if not use_thinking
+            else {
+                "chat_template_kwargs": {"enable_thinking": True},
+            }
+        )
+        # Ask llama.cpp to report per-request timings. Ignored by servers that
+        # don't support it (e.g. vLLM), so it's safe to always send.
+        extra_model_config["timings_per_token"] = True
+
+        payload = {
+            "model": self.model_name,
+            "messages": [
+                {
+                    "role": "user",
+                    "content": str(message_content),
+                },
+            ],
+            "max_tokens": self.max_tokens,
+            "temperature": self.temperature,
+            **extra_model_config,
+        }
+
+        start = time.perf_counter()
+        async with Serverless() as client:
+            endpoint = await client.get_endpoint(name=self.endpoint_name)
+            # `cost` is the autoscaler's load estimate for this request;
+            # max_tokens is the right proxy since generation length drives GPU
+            # time.
+            resp = await endpoint.request(
+                "/v1/chat/completions",
+                payload,
+                cost=self.max_tokens,
+                stream=False,
+            )
+
+        response = resp["response"]
+        log_tokens_per_second(response, time.perf_counter() - start)
+        response_text = (
+            (response.get("choices") or [{}])[0].get("message", {}).get("content")
+        )
+        if response_text is None:
+            raise NoOutputError()
+        return response_text

src/genai_utils/local_models/throughput.py

@@ -0,0 +1,41 @@
+import logging
+from typing import Any
+
+_logger = logging.getLogger(__name__)
+
+
+def log_tokens_per_second(response: dict[str, Any], elapsed: float) -> None:
+    """
+    Logs generation throughput (tokens/second) at info level.
+
+    Takes a plain OpenAI-shaped ``dict``. The serverless endpoint returns one
+    directly; the direct :class:`LLM` client passes ``response.model_dump()``.
+    Prefers the server's own decode timing when available (llama.cpp reports a
+    ``timings`` object with ``predicted_per_second``), otherwise falls back to
+    ``completion_tokens`` over the measured wall-clock time, which works for any
+    OpenAI-compatible server such as vLLM.
+    """
+    if not _logger.isEnabledFor(logging.INFO):
+        return
+
+    # llama.cpp: authoritative decode rate from the server, no network overhead.
+    timings = response.get("timings")
+    if isinstance(timings, dict) and timings.get("predicted_per_second") is not None:
+        _logger.info(
+            "Generation throughput: %.1f tok/s (%d tokens in %.2fs, server timings)",
+            timings["predicted_per_second"],
+            timings.get("predicted_n", 0),
+            timings.get("predicted_ms", 0) / 1000,
+        )
+        return
+
+    # vLLM / generic fallback: completion tokens over wall-clock time.
+    usage = response.get("usage") or {}
+    completion_tokens = usage.get("completion_tokens")
+    if completion_tokens and elapsed > 0:
+        _logger.info(
+            "Generation throughput: %.1f tok/s (%d completion tokens in %.2fs)",
+            completion_tokens / elapsed,
+            completion_tokens,
+            elapsed,
+        )

tests/genai_utils/test_openai_llm.py

@@ -1,25 +1,25 @@
-import logging
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import openai
 from pytest import mark, param, raises
 
-from genai_utils.openai_llm import (
-    LLM,
-    NoOutputError,
-    _is_model_loading,
-    log_tokens_per_second,
-)
-
-LOGGER_NAME = "genai_utils.openai_llm"
+from genai_utils.local_models.llm import LLM, NoOutputError, _is_model_loading
 
 
 def _make_response(content):
-    """Build a minimal stand-in for an OpenAI chat completion response."""
+    """Build a minimal stand-in for an OpenAI chat completion response.
+
+    ``run_message`` reads the text off the object and hands
+    ``response.model_dump()`` to ``log_tokens_per_second``, so the stub needs
+    both.
+    """
     message = SimpleNamespace(content=content)
     choice = SimpleNamespace(message=message)
-    return SimpleNamespace(choices=[choice], model_extra={}, usage=None)
+    return SimpleNamespace(
+        choices=[choice],
+        model_dump=lambda: {"choices": [{"message": {"content": content}}]},
+    )
 
 
 class _ModelLoadingError(openai.InternalServerError):
@@ -92,42 +92,3 @@ def test_run_message_passes_thinking_config(use_thinking, expected_enabled) -> N
     assert extra_body["timings_per_token"] is True
     # the model name is forwarded to the API
     assert kwargs["model"] == "test-model"
-
-
-def test_log_tokens_per_second_silent_when_info_disabled(caplog) -> None:
-    """Should be a no-op (and never raise) when info logging is off."""
-    caplog.set_level(logging.WARNING, logger=LOGGER_NAME)
-    response = SimpleNamespace(model_extra={}, usage=None)
-    log_tokens_per_second(response, elapsed=1.0)
-    assert caplog.records == []
-
-
-def test_log_tokens_per_second_prefers_server_timings(caplog) -> None:
-    caplog.set_level(logging.INFO, logger=LOGGER_NAME)
-    response = SimpleNamespace(
-        model_extra={
-            "timings": {
-                "predicted_per_second": 42.0,
-                "predicted_n": 10,
-                "predicted_ms": 250,
-            }
-        },
-        usage=SimpleNamespace(completion_tokens=999),
-    )
-    log_tokens_per_second(response, elapsed=5.0)
-
-    assert "throughput" in caplog.text
-    # the server timing rate wins over the wall-clock fallback
-    assert "42.0 tok/s" in caplog.text
-
-
-def test_log_tokens_per_second_falls_back_to_usage(caplog) -> None:
-    caplog.set_level(logging.INFO, logger=LOGGER_NAME)
-    response = SimpleNamespace(
-        model_extra={},
-        usage=SimpleNamespace(completion_tokens=20),
-    )
-    log_tokens_per_second(response, elapsed=2.0)
-
-    # 20 tokens over 2 seconds => 10 tok/s
-    assert "10.0 tok/s" in caplog.text