kiapi 0.1.0__py3-none-any.whl

kiapi/core/memory/_services/memory_manager.py

@@ -0,0 +1,257 @@
+"""Global, budget-bounded resident model cache shared by every capability.
+
+This generalizes mlx-vlm-server's per-server manager to the whole of kiapi.
+Because the worker is single-flight (one job at a time, see worker.py), the model
+acquired for the running job is the only one that will generate next, so it is
+correct to budget only *its* peak headroom. Before a model is used we ensure:
+
+    Σ(resident weights, excluding the target)
+        + target.weight + target.peak_headroom   ≤  memory_limit_gb
+
+evicting residents until it holds. Eviction order is **(priority asc, last_used
+asc)**: lower-priority models go first, and among equal priority the
+least-recently-used. So a small model you want resident gets a high ``priority``
+and survives even when a big model churns through.
+
+Two concepts are kept separate:
+  - **resident weights** — long-lived, reusable, evictable (tracked here),
+  - **peak headroom** — transient memory the *running* job needs on top of
+    weights; only one job runs at a time, so only the target's headroom matters.
+
+Framework specifics (so eviction actually frees memory) live in
+:func:`resolve_framework_strategy`:
+  - MLX: dropping Python refs is not enough — the Metal allocator keeps a buffer
+    cache. After deleting we ``gc.collect()`` then ``mx.clear_cache()``. Memory is
+    measured via ``mx.get_active_memory`` deltas.
+  - Other frameworks plug in via a new FrameworkStrategy and an optional
+    per-module ``release(payload)`` hook.
+
+All mutating methods run on the single worker thread; only ``stats()`` is read
+from the event-loop thread, so a small lock guards the dict.
+"""
+
+import gc
+import logging
+import threading
+import time
+import traceback
+
+from kiapi.core.model import ModelSpec
+
+from .._exceptions.memory_budget_error import MemoryBudgetError
+from .._operations.resolve_framework_strategy import resolve_framework_strategy
+from .._schemas.memory_stats import MemoryStats
+from .._schemas.resident_model import ResidentModel
+from .._schemas.resident_model_stats import ResidentModelStats
+from .._settings import MemorySettings
+from .._types.resident_key import ResidentKey
+
+_GB = 1024**3
+logger = logging.getLogger(__name__)
+
+
+class MemoryManager:
+    def __init__(self, settings: MemorySettings) -> None:
+        self._settings: MemorySettings = settings
+        self._loaded: dict[ResidentKey, ResidentModel] = {}
+        self._lock = threading.Lock()
+
+    @property
+    def memory_limit_gb(self) -> float:
+        limit = self._settings.memory_limit_gb
+
+        if limit is None:
+            raise RuntimeError("memory_limit_gb must be resolved before use")
+
+        return limit
+
+    @property
+    def default_ttl_s(self) -> float:
+        return self._settings.default_ttl_s
+
+    # -- TTL -------------------------------------------------------------------
+
+    def _effective_ttl(self, spec: ModelSpec) -> float | None:
+        ttl = spec.ttl_seconds
+
+        if ttl is None:
+            ttl = self.default_ttl_s
+        if ttl <= 0:
+            return None
+
+        return ttl
+
+    def sweep_expired(
+        self, *, exclude_key: ResidentKey | None = None, now: float | None = None
+    ) -> list[str]:
+        """Runs on the worker thread — it frees framework memory, so it must not
+        be called from the event loop. Returns the freed models' names."""
+        now = now if now is not None else time.monotonic()
+        expired: list[ResidentKey] = []
+        for key, ld in list(self._loaded.items()):
+            if key == exclude_key:
+                continue
+            ttl = self._effective_ttl(ld.spec)
+            if ttl is not None and (now - ld.last_used) > ttl:
+                expired.append(key)
+        freed = []
+        for key in expired:
+            name = self._loaded[key].spec.name if key in self._loaded else key
+            self._evict(key, reason="ttl")
+            freed.append(name)
+        return freed
+
+    # -- public (worker thread) ------------------------------------------------
+
+    def acquire(self, spec: ModelSpec) -> object:
+        # Free any idle-expired models first (don't touch the one we're about to
+        # use) so their memory is reclaimed before we budget/load.
+        self.sweep_expired(exclude_key=spec.key)
+        existing = self._loaded.get(spec.key)
+        if existing is not None:
+            # Re-check budget: another model may have loaded since, and this
+            # model's generation headroom must still fit.
+            self._ensure_budget(spec, existing.weight_gb)
+            existing.last_used = time.monotonic()
+            return existing.payload
+
+        self._ensure_budget(spec, spec.weight_gb)  # estimate (not loaded yet)
+        loaded = self._load(spec)
+        with self._lock:
+            self._loaded[spec.key] = loaded
+        return loaded.payload
+
+    def reserve(
+        self, need_gb: float, *, exclude_key: ResidentKey | None = None
+    ) -> None:
+        """Make budget room for a transient job that loads/frees its own weights.
+
+        For providers that don't keep a resident model (e.g. mlx-video's LTX-2,
+        which loads and frees per call), there's nothing to ``acquire``. Instead
+        the handler reserves the peak it will transiently need: we evict residents
+        by (priority asc, last_used asc) until ``resident_total + need ≤ limit``.
+        Single-flight guarantees nothing else runs while the job uses that room,
+        and the provider frees it when done."""
+        limit = self.memory_limit_gb
+        if need_gb > limit:
+            raise MemoryBudgetError(
+                f"transient job needs ~{need_gb:.1f} GB but the budget is only {limit:.1f} GB"
+            )
+        while self._resident_weight_gb(exclude_key=exclude_key) + need_gb > limit:
+            victim = self._victim(exclude_key=exclude_key)
+            if victim is None:
+                raise MemoryBudgetError(
+                    f"transient job needs ~{need_gb:.1f} GB but the budget is "
+                    f"{limit:.1f} GB with nothing left to evict"
+                )
+            self._evict(victim)
+
+    def stats(self) -> MemoryStats:
+        """Snapshot for /health (safe from the event-loop thread)."""
+        now = time.monotonic()
+        with self._lock:
+            loaded: list[ResidentModelStats] = []
+            resident_total = 0.0
+            for ld in self._loaded.values():
+                ttl = self._effective_ttl(ld.spec)
+                idle = round(now - ld.last_used, 1)
+                resident_total += ld.weight_gb
+                loaded.append(
+                    ResidentModelStats(
+                        name=ld.spec.name,
+                        family=ld.spec.family,
+                        domain=ld.spec.domain,
+                        repo=ld.spec.repo,
+                        weight_gb=round(ld.weight_gb, 1),
+                        priority=ld.spec.priority,
+                        idle_s=idle,
+                        ttl_s=ttl,
+                        expires_in_s=round(ttl - idle, 1) if ttl is not None else None,
+                    )
+                )
+        return MemoryStats(
+            loaded=loaded,
+            resident_gb=round(resident_total, 1),
+            budget_gb=self.memory_limit_gb,
+        )
+
+    def shutdown(self) -> None:
+        for key in list(self._loaded):
+            self._evict(key)
+
+    # -- budget / eviction -----------------------------------------------------
+
+    def _resident_weight_gb(self, exclude_key: ResidentKey | None = None) -> float:
+        return sum(
+            ld.weight_gb for key, ld in self._loaded.items() if key != exclude_key
+        )
+
+    def _ensure_budget(self, spec: ModelSpec, target_weight_gb: float) -> None:
+        limit = self.memory_limit_gb
+        need = target_weight_gb + spec.peak_headroom_gb
+        while self._resident_weight_gb(exclude_key=spec.key) + need > limit:
+            victim = self._victim(exclude_key=spec.key)
+            if victim is None:
+                raise MemoryBudgetError(
+                    f"model {spec.name!r} needs ~{need:.1f} GB but the budget is "
+                    f"{limit:.1f} GB with nothing left to evict"
+                )
+            self._evict(victim)
+
+    def _victim(self, *, exclude_key: ResidentKey | None) -> ResidentKey | None:
+        cands = [
+            (ld.spec.priority, ld.last_used, key)
+            for key, ld in self._loaded.items()
+            if key != exclude_key
+        ]
+        if not cands:
+            return None
+        cands.sort()  # priority asc, then last_used asc
+        return cands[0][2]
+
+    def _evict(self, key: ResidentKey, *, reason: str = "budget") -> None:
+        with self._lock:
+            loaded = self._loaded.pop(key, None)
+        if loaded is None:
+            return
+        spec = loaded.spec
+        strategy = resolve_framework_strategy(spec.framework)
+        before = strategy.active_bytes()
+        # Optional per-module cleanup (framework-specific extras) before dropping.
+        release = getattr(spec.module, "release", None)
+        if release is not None:
+            try:
+                release(loaded.payload)
+            except Exception:
+                traceback.print_exc()
+        del loaded
+        gc.collect()
+        strategy.free_caches()
+        freed = max(before - strategy.active_bytes(), 0) / _GB
+        logger.info(
+            "evicted %s (%s) [%s] - freed ~%.1f GB",
+            spec.name,
+            spec.key,
+            reason,
+            freed,
+        )
+
+    # -- loading ---------------------------------------------------------------
+
+    def _load(self, spec: ModelSpec) -> ResidentModel:
+        strategy = resolve_framework_strategy(spec.framework)
+        logger.info("loading %s (%s) ...", spec.name, spec.key)
+        before = strategy.active_bytes()
+        payload = spec.module.load(spec)
+        measured = max(strategy.active_bytes() - before, 0) / _GB
+        # Reconcile estimate with reality; fall back to the estimate if the delta
+        # looks bogus (e.g. weights still lazy / not yet materialized).
+        weight_gb = measured if measured > 0.5 else spec.weight_gb
+        logger.info(
+            "loaded %s - weight ~%.1f GB (estimate %.1f, measured %.1f)",
+            spec.name,
+            weight_gb,
+            spec.weight_gb,
+            measured,
+        )
+        return ResidentModel(spec, payload, weight_gb, time.monotonic())

kiapi/core/memory/_settings.py

@@ -0,0 +1,36 @@
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+
+class MemorySettings(BaseSettings):
+    """Settings for the model residency memory budget and idle TTL shared by all capabilities."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="KIAPI_",
+        extra="ignore",
+    )
+
+    memory_limit_gb: float | None = Field(
+        default=None,
+        title="Memory budget GB",
+        description=(
+            "Memory budget in GB shared by resident models and runtime peaks across all capabilities.\n"
+            "If unset, kiapi uses 80% of the machine's total memory.\n"
+            "Before a job starts, lower-priority and older resident models are "
+            "released until the job fits within this value."
+        ),
+    )
+
+    default_ttl_s: float = Field(
+        default=1800.0,
+        title="Default idle TTL seconds",
+        description=(
+            "Default number of seconds before an unused resident model is auto-released.\n"
+            "Set to 0.0 or less to disable TTL release and evict only under "
+            "memory pressure."
+        ),
+    )
+
+
+settings_manager = SettingsManager(MemorySettings)

kiapi/core/memory/_types/framework_name.py

@@ -0,0 +1,4 @@
+type FrameworkName = str
+"""The backend that owns a model's memory and knows how to measure/free it —
+``"mlx"`` (Metal allocator) or ``"rss"`` (framework-agnostic process-RSS
+fallback). Matches :class:`ModelSpec.framework`; selects a FrameworkStrategy."""

kiapi/core/memory/_types/resident_key.py

@@ -0,0 +1,4 @@
+type ResidentKey = str
+"""Identity of a resident model in the cache — equals :attr:`ModelSpec.key`
+(``"family:repo"``), unique per loaded set of weights. Used as the ``_loaded``
+dict key and as the ``exclude_key`` that spares the in-use model from eviction."""

kiapi/core/model/__init__.py

@@ -0,0 +1,36 @@
+"""The model registry: the catalog of servable models across all families.
+
+Public surface (unchanged from the former ``core/model.py``):
+
+  - :class:`ModelSpec` — a registry entry (variant + family/domain + repo +
+    handler module + memory/budget metadata).
+  - :class:`ModelRegistry` / the ``model_registry`` singleton — register,
+    list, group, and resolve specs.
+  - :class:`UnknownModelError` — raised by ``resolve`` on an unknown ``model``.
+
+The semantic axes are spelled out as types in :mod:`._types`
+(``ModelName`` / ``ModelFamily`` / ``ModelDomain`` / ``ModelRepo``).
+"""
+
+from ._exceptions.unknown_model_error import UnknownModelError
+from ._schemas.model_spec import ModelSpec
+from ._services.model_registry import ModelRegistry, model_registry
+from ._types.model_alias import ModelAlias
+from ._types.model_domain import ModelDomain
+from ._types.model_family import ModelFamily
+from ._types.model_key import ModelKey
+from ._types.model_name import ModelName
+from ._types.model_repo import ModelRepo
+
+__all__ = [
+    "ModelAlias",
+    "ModelDomain",
+    "ModelFamily",
+    "ModelKey",
+    "ModelName",
+    "ModelRegistry",
+    "ModelRepo",
+    "ModelSpec",
+    "UnknownModelError",
+    "model_registry",
+]

kiapi/core/model/_exceptions/unknown_model_error.py

@@ -0,0 +1,2 @@
+class UnknownModelError(ValueError):
+    """Requested ``model`` doesn't match any registry entry → HTTP 400."""

kiapi/core/model/_schemas/model_spec.py

@@ -0,0 +1,145 @@
+"""A registry entry: everything needed to discover, resolve, load, and budget a
+servable model variant.
+
+Each :class:`ModelSpec` ties together:
+
+  - the public ``name`` (the variant) used in a request's ``model`` field,
+  - its ``family`` (resolution namespace + cache-key prefix + OpenAPI topic) and
+    ``domain`` (grouping),
+  - the Hugging Face ``repo`` to load,
+  - the per-model **handler module** that owns its own flow,
+  - memory estimates (``weight_gb`` / ``peak_headroom_gb``), ``priority``,
+    ``framework``, ``resident``, ``ttl_seconds`` (see core/memory.py).
+
+A handler module must expose:
+
+    FEATURES : set[str]                      # modalities/features, family-specific
+    load(spec) -> payload                    # load weights, return an opaque payload
+    run(payload, req, settings) -> result    # one job's worth of work
+    # optional:
+    warmup(payload) -> None                  # prime kernels (else: load only)
+    release(payload) -> None                 # extra cleanup beyond framework default
+
+Estimates are intentionally rough; core/memory.py reconciles ``weight_gb`` with
+the real measured footprint after the first load.
+
+This is a Pydantic model so the family model-list endpoints
+(``/v1/{domain}/{family}/models``) can serialize specs directly and have their
+fields documented in OpenAPI. The handler ``module`` is excluded from output.
+"""
+
+from types import ModuleType
+
+from pydantic import BaseModel, ConfigDict, Field, computed_field
+from pydantic.json_schema import SkipJsonSchema
+
+from kiapi.core.setup import SetupResource
+
+from .._types.model_alias import ModelAlias
+from .._types.model_domain import ModelDomain
+from .._types.model_family import ModelFamily
+from .._types.model_name import ModelName
+from .._types.model_repo import ModelRepo
+
+
+class ModelSpec(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)
+
+    name: ModelName = Field(
+        ...,
+        description="Public model variant name passed in a request's model field; "
+        "named by the variant alone, never repeating the family.",
+        examples=["turbo"],
+    )
+    family: ModelFamily = Field(
+        ...,
+        description="Family that owns the operation vocabulary and resolves this "
+        "variant; also used as the API identifier and OpenAPI topic.",
+        examples=["zimage"],
+    )
+    domain: ModelDomain = Field(
+        ...,
+        description="Modality bucket used for discovery and grouping.",
+        examples=["image"],
+    )
+    repo: ModelRepo = Field(
+        ...,
+        description="Hugging Face repo id (or local path) the weights load from.",
+        examples=["mflux/z-image-turbo"],
+    )
+    module: SkipJsonSchema[ModuleType] = Field(
+        ...,
+        exclude=True,
+        description="Per-model handler module that owns this model's load/run flow "
+        "(internal; not serialized).",
+    )
+    weight_gb: float = Field(
+        ...,
+        ge=0.0,
+        description="Rough estimate of resident weight footprint in GB. "
+        "core/memory reconciles this with the measured footprint after first load.",
+        examples=[6.0],
+    )
+    peak_headroom_gb: float = Field(
+        ...,
+        ge=0.0,
+        description="Rough estimate of peak memory above the weights during a job, "
+        "in GB — the extra budget headroom kept free while running.",
+        examples=[4.0],
+    )
+    framework: str = Field(
+        default="mlx",
+        description='Backend that owns this model\'s memory (e.g. "mlx", or "rss" '
+        "for a subprocess whose footprint kiapi can only see as process RSS).",
+        examples=["mlx"],
+    )
+    priority: int = Field(
+        default=0,
+        description="Eviction priority. Residents are evicted by (priority asc, "
+        "last_used asc); a higher value keeps a model resident under churn.",
+        examples=[0],
+    )
+    aliases: tuple[ModelAlias, ...] = Field(
+        default_factory=tuple,
+        description="Extra names that also resolve to this spec.",
+        examples=[["omni", "qwen3-omni-30b"]],
+    )
+    default: bool = Field(
+        default=False,
+        exclude=True,
+        description="Whether this is the family's default variant when a request "
+        "omits model. If no spec sets it, the first-registered spec is the default "
+        "(internal; not serialized).",
+    )
+    resident: bool = Field(
+        default=True,
+        description="Whether the model is held resident after loading. Transient "
+        "specs (resident=False) load+free per call and reserve budget instead of "
+        "acquiring it.",
+        examples=[True],
+    )
+    ttl_seconds: float | None = Field(
+        default=None,
+        description="Optional idle TTL (seconds). None inherits the global default; "
+        "a value <= 0 means never expire (pin resident).",
+        examples=[1800.0],
+    )
+    setup_resources: tuple[SetupResource, ...] = Field(
+        default_factory=tuple,
+        description="Resources that must be activated before this model can run: "
+        "Hugging Face snapshots, Docker images, local paths, URL files, or Python "
+        "virtual environments.",
+    )
+
+    @property
+    def key(self) -> str:
+        """Stable identity in the resident cache (unique per loaded weights)."""
+        return f"{self.family}:{self.repo}"
+
+    @computed_field(  # type: ignore[prop-decorator]
+        description="Handler-declared modalities/features (NOT the family).",
+        examples=[["text", "image"]],
+    )
+    @property
+    def features(self) -> set[str]:
+        return set(getattr(self.module, "FEATURES", set()))

kiapi/core/model/_services/model_registry.py

@@ -0,0 +1,85 @@
+"""The model registry: the catalog of servable models across all families.
+
+kiapi organizes models on two axes:
+
+  - ``domain`` — the modality bucket for discovery/grouping: "chat", "embedding",
+    "audio", "video", "image".
+  - ``family`` — the provider/model family that owns an *operation vocabulary*,
+    spelled as a single lowercase token from the upstream's package/model name
+    ("acestep", "audiogen", "ltx2", "chat", "embedding"). The family is the
+    resolution namespace and the capability OpenAPI document.
+
+Generation endpoints are organized as ``/v1/<domain>/<family>/<op>`` — each
+family exposes exactly its own operations and request shape (no forced common
+schema). ``model`` selects a *variant within the family* (e.g. acestep
+``xl-base`` / ``turbo``), named by the variant alone (it must not repeat the
+family). chat/embedding stay standardized modality APIs (family == domain).
+"""
+
+from .._exceptions.unknown_model_error import UnknownModelError
+from .._schemas.model_spec import ModelSpec
+from .._types.model_domain import ModelDomain
+from .._types.model_family import ModelFamily
+from .._types.model_key import ModelKey
+from .._types.model_name import ModelName
+
+
+class ModelRegistry:
+    def __init__(self) -> None:
+        self._specs: list[ModelSpec] = []
+        self._index: dict[tuple[ModelFamily, ModelKey], ModelSpec] = {}
+
+    def register(self, spec: ModelSpec) -> None:
+        for k in (spec.name, spec.repo, *spec.aliases):
+            self._index[(spec.family, k.lower())] = spec
+        self._specs.append(spec)
+
+    def list_specs(self, family: ModelFamily | None = None) -> list[ModelSpec]:
+        if family is None:
+            return list(self._specs)
+        return [s for s in self._specs if s.family == family]
+
+    def families(self, domain: ModelDomain | None = None) -> list[ModelFamily]:
+        seen: list[ModelFamily] = []
+        for s in self._specs:
+            if domain is not None and s.domain != domain:
+                continue
+            if s.family not in seen:
+                seen.append(s.family)
+        return seen
+
+    def domains(self) -> list[ModelDomain]:
+        seen: list[ModelDomain] = []
+        for s in self._specs:
+            if s.domain not in seen:
+                seen.append(s.domain)
+        return seen
+
+    def domain_of(self, family: ModelFamily) -> ModelDomain | None:
+        for s in self._specs:
+            if s.family == family:
+                return s.domain
+        return None
+
+    def default_name(self, family: ModelFamily) -> ModelName | None:
+        specs = self.list_specs(family)
+        if not specs:
+            return None
+        for s in specs:
+            if s.default:
+                return s.name
+        return specs[0].name
+
+    def resolve(self, family: ModelFamily, name: ModelName | None) -> ModelSpec:
+        default = self.default_name(family)
+        key = (name or default or "").strip().lower()
+        spec = self._index.get((family, key))
+        if spec is None:
+            known = sorted({s.name for s in self.list_specs(family)})
+            raise UnknownModelError(
+                f"unknown {family} model {name!r}; available: {known}"
+            )
+        return spec
+
+
+model_registry = ModelRegistry()

kiapi/core/model/_types/model_alias.py

@@ -0,0 +1,2 @@
+type ModelAlias = str
+"""An extra name that resolves to a spec, alongside its ``name`` and ``repo``."""

kiapi/core/model/_types/model_domain.py

@@ -0,0 +1,6 @@
+from typing import Literal
+
+ModelDomain = Literal["chat", "embedding", "audio", "video", "image", "web"]
+"""The modality bucket used for discovery/grouping. A fixed, closed set — every
+family belongs to exactly one domain. Generation endpoints are organized as
+``/v1/<domain>/<family>/<op>``."""

kiapi/core/model/_types/model_family.py

@@ -0,0 +1,5 @@
+type ModelFamily = str
+"""The provider/model family that owns an *operation vocabulary*, spelled as a
+single lowercase token from the upstream's package/model name ("acestep",
+"audiogen", "ltx2", "chat", "embedding"). The family is the resolution
+namespace, the cache-key prefix, and the capability OpenAPI topic."""

kiapi/core/model/_types/model_key.py

@@ -0,0 +1,8 @@
+from .model_alias import ModelAlias
+from .model_name import ModelName
+from .model_repo import ModelRepo
+
+type ModelKey = ModelName | ModelRepo | ModelAlias
+"""A lookup token a request may use to address a model within its family — any of
+the spec's ``name``, ``repo``, or ``aliases``. The registry indexes specs by
+``(family, key.lower())``, so resolution is case-insensitive."""

kiapi/core/model/_types/model_name.py

@@ -0,0 +1,4 @@
+type ModelName = str
+"""The public ``model`` selector — a *variant within a family* (e.g. acestep
+``xl-base`` / ``turbo``). Named by the variant alone; it must not repeat the
+family. Used in a request's ``model`` field and as a registry lookup key."""

kiapi/core/model/_types/model_repo.py

@@ -0,0 +1,3 @@
+type ModelRepo = str
+"""The Hugging Face repo id (or local path) the model weights load from. Part of
+a :class:`ModelSpec`'s cache identity (``family:repo``)."""

kiapi/core/net/__init__.py

@@ -0,0 +1,16 @@
+"""Networking guards shared across capabilities.
+
+Currently a single SSRF guard for user-supplied URLs; see
+:func:`verify_public_url`.
+"""
+
+from ._exceptions.unsafe_url_error import UnsafeURLError
+from ._helpers.verify_public_url import verify_public_url
+from ._settings import NetSettings, settings_manager
+
+__all__ = [
+    "NetSettings",
+    "UnsafeURLError",
+    "settings_manager",
+    "verify_public_url",
+]

kiapi/core/net/_exceptions/unsafe_url_error.py

@@ -0,0 +1,7 @@
+class UnsafeURLError(ValueError):
+    """A user-supplied URL was rejected before fetching (SSRF guard).
+
+    Raised for non-http(s) schemes, unresolvable hosts, or hosts that resolve to
+    a non-public address (loopback, private, link-local, etc.). Callers map this
+    to a 400-class error.
+    """