kiapi 0.1.0__py3-none-any.whl

kiapi/core/job/_services/job_store.py

@@ -0,0 +1,33 @@
+import threading
+from typing import Any
+
+from .._models.job import Job
+from .._types.job_id import JobID
+from .._types.job_type import JobType
+
+
+class JobStore:
+    def __init__(self) -> None:
+        self._jobs: dict[JobID, Job] = {}
+        self._lock = threading.Lock()
+
+    def create(self, type: JobType, params: dict[str, Any] | None = None) -> Job:
+        job = Job(type=type, params=params or {})
+        with self._lock:
+            self._jobs[job.id] = job
+        return job
+
+    def get(self, job_id: JobID) -> Job | None:
+        with self._lock:
+            return self._jobs.get(job_id)
+
+    def list_all(self) -> list[Job]:
+        with self._lock:
+            return sorted(self._jobs.values(), key=lambda j: j.created_at, reverse=True)
+
+    def delete(self, job_id: JobID) -> Job | None:
+        """A running job cannot be interrupted (MLX/generation is not
+        preemptible) — deleting it just forgets it once it finishes. Queued jobs
+        are effectively canceled."""
+        with self._lock:
+            return self._jobs.pop(job_id, None)

kiapi/core/job/_services/progress_reporter.py

@@ -0,0 +1,79 @@
+"""Ambient progress reporting for the running job.
+
+kiapi runs exactly one job at a time on one worker thread (see core/worker), so
+"the progress of the currently running job" is a well-defined ambient rather than
+something every call site must thread through by hand. The worker binds a
+:class:`ProgressReporter` for the job it is about to run; capability code reaches
+it with :meth:`ProgressReporter.current` and calls :meth:`update` / :meth:`step`
+to push coarse progress onto the :class:`Job` (surfaced by /v1/jobs).
+
+A reporter bound to no job (the default outside a running job) is a silent no-op,
+so capability code can call ``ProgressReporter.current().update(...)``
+unconditionally. Updates are throttled (time + delta) so a hot denoising loop
+calling ``step`` per iteration stays cheap; the terminal ``1.0`` from
+``Job.mark_succeeded`` is what guarantees a clean final value.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from contextvars import ContextVar
+from typing import Self, TypeVar
+
+from .._models.job import Job
+
+T = TypeVar("T")
+
+
+class ProgressReporter:
+    _current: ContextVar[ProgressReporter | None] = ContextVar(
+        "kiapi_progress_reporter", default=None
+    )
+
+    def __init__(
+        self,
+        job: Job | None = None,
+        *,
+        min_interval_s: float = 0.5,
+        min_delta: float = 0.01,
+    ) -> None:
+        self._job = job
+        self._min_interval_s = min_interval_s
+        self._min_delta = min_delta
+        self._last_at = 0.0
+        self._last_fraction = -1.0
+
+    def update(self, fraction: float, label: str | None = None) -> None:
+        if self._job is None:
+            return
+
+        now = time.monotonic()
+        forced = label is not None or fraction >= 1.0
+
+        if not forced:
+            if now - self._last_at < self._min_interval_s:
+                return
+            if abs(fraction - self._last_fraction) < self._min_delta:
+                return
+
+        self._last_at = now
+        self._last_fraction = fraction
+        self._job.update_progress(fraction, label)
+
+    def step(self, i: int, total: int, label: str | None = None) -> None:
+        if total > 0:
+            self.update(i / total, label)
+
+    @classmethod
+    def current(cls) -> Self:
+        reporter = cls._current.get()
+        return reporter if isinstance(reporter, cls) else cls()
+
+    def run_bound(self, fn: Callable[[], T]) -> T:
+        token = self._current.set(self)
+
+        try:
+            return fn()
+        finally:
+            self._current.reset(token)

kiapi/core/job/_types/job_id.py

@@ -0,0 +1,3 @@
+type JobID = str
+"""The address of a job (``"job_<hex>"``). Async requests return it immediately;
+clients poll the job by this id. Jobs are in-memory, so it is not persistent."""

kiapi/core/job/_types/job_result.py

@@ -0,0 +1,5 @@
+from typing import Any
+
+type JobResult = dict[str, Any]
+"""A job's result payload. Free-form — its shape depends on the job's
+:class:`JobType` (the caller dispatches on the type to read it)."""

kiapi/core/job/_types/job_type.py

@@ -0,0 +1,4 @@
+type JobType = str
+"""Names the operation a job runs (e.g. ``"chat"``, ``"acestep.generate"``,
+``"ltx2"``). The caller dispatches on it, and it determines the shape of the
+job's ``result``."""

kiapi/core/logging/__init__.py

@@ -0,0 +1,14 @@
+from ._helpers.get_log_level import get_log_level
+from ._helpers.setup_logger import setup_logger
+from ._settings import settings_manager
+from ._types.log_level import LogLevel
+
+__all__ = [
+    # ._types
+    "LogLevel",
+    # ._helpers
+    "get_log_level",
+    # ._settings
+    "settings_manager",
+    "setup_logger",
+]

kiapi/core/logging/_helpers/get_log_level.py

@@ -0,0 +1,8 @@
+import logging
+
+from .._settings import settings_manager
+
+
+def get_log_level() -> int:
+    settings = settings_manager.get_settings()
+    return logging.getLevelNamesMapping()[settings.log_level]

kiapi/core/logging/_helpers/setup_logger.py

@@ -0,0 +1,14 @@
+import logging
+
+from .._settings import settings_manager
+
+_LOG_FORMAT = "%(asctime)s %(levelname)s [%(name)s] %(message)s"
+
+
+def setup_logger() -> None:
+    settings = settings_manager.get_settings()
+    log_level = logging.getLevelNamesMapping()[settings.log_level]
+
+    logging.basicConfig(level=log_level, format=_LOG_FORMAT)
+    logging.getLogger().setLevel(log_level)
+    logging.getLogger("kiapi").setLevel(log_level)

kiapi/core/logging/_settings.py

@@ -0,0 +1,34 @@
+from pydantic import Field, field_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+from ._types.log_level import LogLevel
+
+
+class LoggingSettings(BaseSettings):
+    """Settings for the global kiapi logging level."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="KIAPI_",
+        extra="ignore",
+    )
+
+    log_level: LogLevel = Field(
+        default="INFO",
+        title="Log level",
+        description=(
+            "Verbosity of emitted logs.\n"
+            "Standard Python logging levels such as DEBUG, INFO, WARNING, ERROR, "
+            "and CRITICAL are accepted."
+        ),
+    )
+
+    @field_validator("log_level", mode="before")
+    @classmethod
+    def _to_upper(cls, v: str) -> str:
+        if isinstance(v, str):
+            return v.upper()
+        return v  # type: ignore
+
+
+settings_manager = SettingsManager(LoggingSettings)

kiapi/core/logging/_types/log_level.py

@@ -0,0 +1,3 @@
+from typing import Literal
+
+type LogLevel = Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]

kiapi/core/memory/__init__.py

@@ -0,0 +1,21 @@
+from ._exceptions.memory_budget_error import MemoryBudgetError
+from ._helpers.create_memory_manager import create_memory_manager
+from ._schemas.memory_stats import MemoryStats
+from ._schemas.resident_model_stats import ResidentModelStats
+from ._services.memory_manager import MemoryManager
+from ._settings import MemorySettings, settings_manager
+
+__all__ = [
+    # ._exceptions
+    "MemoryBudgetError",
+    # ._services
+    "MemoryManager",
+    # ._settings
+    "MemorySettings",
+    # ._schemas
+    "MemoryStats",
+    "ResidentModelStats",
+    # ._helpers
+    "create_memory_manager",
+    "settings_manager",
+]

kiapi/core/memory/_exceptions/memory_budget_error.py

@@ -0,0 +1,2 @@
+class MemoryBudgetError(RuntimeError):
+    """Target model can't fit within the budget even after evicting all others."""

kiapi/core/memory/_helpers/create_memory_manager.py

@@ -0,0 +1,17 @@
+from .._operations.log_memory_limit import log_memory_limit
+from .._operations.resolve_effective_memory_limit_gb import (
+    resolve_effective_memory_limit_gb,
+)
+from .._services.memory_manager import MemoryManager
+from .._settings import settings_manager
+
+
+def create_memory_manager() -> MemoryManager:
+    settings = settings_manager.get_settings()
+
+    effective_limit_gb = resolve_effective_memory_limit_gb(settings)
+    log_memory_limit(settings, effective_limit_gb)
+
+    return MemoryManager(
+        settings.model_copy(update={"memory_limit_gb": effective_limit_gb})
+    )

kiapi/core/memory/_operations/log_memory_limit.py

@@ -0,0 +1,29 @@
+import logging
+
+import psutil  # type: ignore[import-untyped]
+
+from .._settings import MemorySettings
+
+_GB = 1024**3
+logger = logging.getLogger(__name__)
+
+
+def log_memory_limit(settings: MemorySettings, effective_limit_gb: float) -> None:
+    memory = psutil.virtual_memory()
+    total_gb = memory.total / _GB
+    available_gb = memory.available / _GB
+    mode = "configured" if settings.memory_limit_gb is not None else "auto"
+
+    logger.info(
+        "memory budget %s: effective %.1f GB (total %.1f GB, available %.1f GB)",
+        mode,
+        effective_limit_gb,
+        total_gb,
+        available_gb,
+    )
+    if available_gb < effective_limit_gb:
+        logger.warning(
+            "memory budget %.1f GB exceeds currently available memory %.1f GB",
+            effective_limit_gb,
+            available_gb,
+        )

kiapi/core/memory/_operations/resolve_effective_memory_limit_gb.py

@@ -0,0 +1,13 @@
+import psutil  # type: ignore[import-untyped]
+
+from .._settings import MemorySettings
+
+_GB = 1024**3
+_AUTO_MEMORY_LIMIT_RATIO = 0.8
+
+
+def resolve_effective_memory_limit_gb(settings: MemorySettings) -> float:
+    if settings.memory_limit_gb is not None:
+        return settings.memory_limit_gb
+
+    return float(psutil.virtual_memory().total) / _GB * _AUTO_MEMORY_LIMIT_RATIO

kiapi/core/memory/_operations/resolve_framework_strategy.py

@@ -0,0 +1,42 @@
+"""Resolve a :class:`FrameworkName` to its :class:`FrameworkStrategy`.
+
+A new framework plugs in by adding a ``(active_bytes, free_caches)`` pair to
+``_FRAMEWORKS`` and, if it needs extra teardown, a per-module ``release(payload)``
+hook (called by the manager before dropping refs).
+"""
+
+from .._schemas.framework_strategy import FrameworkStrategy
+from .._types.framework_name import FrameworkName
+
+
+def _mlx_active() -> int:
+    import mlx.core as mx
+
+    return int(mx.get_active_memory())
+
+
+def _mlx_free() -> None:
+    import mlx.core as mx
+
+    mx.clear_cache()
+
+
+def _rss_active() -> int:
+    import psutil  # type: ignore[import-untyped]
+
+    return int(psutil.Process().memory_info().rss)
+
+
+def _noop_free() -> None:
+    pass
+
+
+_FRAMEWORKS: dict[FrameworkName, FrameworkStrategy] = {
+    "mlx": FrameworkStrategy(_mlx_active, _mlx_free),
+    # framework-agnostic fallback (e.g. torch/MPS until a dedicated entry lands)
+    "rss": FrameworkStrategy(_rss_active, _noop_free),
+}
+
+
+def resolve_framework_strategy(framework: FrameworkName) -> FrameworkStrategy:
+    return _FRAMEWORKS.get(framework, _FRAMEWORKS["rss"])

kiapi/core/memory/_schemas/framework_strategy.py

@@ -0,0 +1,14 @@
+from collections.abc import Callable
+from typing import NamedTuple
+
+
+class FrameworkStrategy(NamedTuple):
+    """How to measure and reclaim a framework's memory, so eviction actually
+    frees it. Resolved per :class:`FrameworkName`."""
+
+    active_bytes: Callable[[], int]
+    """Current active allocation, in bytes. Deltas around load/evict give the
+    measured footprint."""
+    free_caches: Callable[[], None]
+    """Release the framework's cached buffers after dropping refs (e.g. MLX's
+    Metal buffer cache); a no-op where the runtime needs none."""

kiapi/core/memory/_schemas/memory_stats.py

@@ -0,0 +1,24 @@
+from pydantic import BaseModel, Field
+
+from .resident_model_stats import ResidentModelStats
+
+
+class MemoryStats(BaseModel):
+    """Resident model memory budget snapshot, served by ``/health``."""
+
+    loaded: list[ResidentModelStats] = Field(
+        default_factory=list,
+        description="Resident models currently held in memory.",
+    )
+    resident_gb: float = Field(
+        ...,
+        ge=0.0,
+        description="Total resident model weight currently counted against the budget.",
+        examples=[12.0],
+    )
+    budget_gb: float = Field(
+        ...,
+        gt=0.0,
+        description="Effective global memory budget in GB shared by all resident models and the running job's peak headroom.",
+        examples=[100.0],
+    )

kiapi/core/memory/_schemas/resident_model.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+
+from kiapi.core.model import ModelSpec
+
+
+@dataclass
+class ResidentModel:
+    spec: ModelSpec
+    payload: object
+    """Whatever the handler's ``load()`` returned (opaque to the manager)."""
+    weight_gb: float
+    """Resident footprint in GB — measured at load, else the spec estimate."""
+    last_used: float
+    """``time.monotonic()`` of the most recent acquire; drives LRU + idle TTL."""

kiapi/core/memory/_schemas/resident_model_stats.py

@@ -0,0 +1,49 @@
+from pydantic import BaseModel, Field
+
+
+class ResidentModelStats(BaseModel):
+    """Snapshot of one resident model held in the memory budget."""
+
+    name: str = Field(..., description="Model variant name.", examples=["turbo"])
+    family: str = Field(
+        ...,
+        description="Model family and API identifier.",
+        examples=["zimage"],
+    )
+    domain: str = Field(
+        ...,
+        description="Capability domain such as chat, embedding, image, audio, video, or web.",
+        examples=["image"],
+    )
+    repo: str = Field(
+        ...,
+        description="Upstream model repository, Docker image, or local resource label.",
+        examples=["mflux/z-image-turbo"],
+    )
+    weight_gb: float = Field(
+        ...,
+        ge=0.0,
+        description="Current resident weight estimate in GB.",
+        examples=[6.0],
+    )
+    priority: int = Field(
+        ...,
+        description="Eviction priority. Lower values are evicted before higher values under budget pressure.",
+        examples=[0],
+    )
+    idle_s: float = Field(
+        ...,
+        ge=0.0,
+        description="Seconds since this resident model was last used.",
+        examples=[12.5],
+    )
+    ttl_s: float | None = Field(
+        default=None,
+        description="Idle TTL in seconds. Null means the resident is pinned until budget pressure or shutdown.",
+        examples=[1800.0],
+    )
+    expires_in_s: float | None = Field(
+        default=None,
+        description="Seconds until TTL eviction. Null when ttl_s is null.",
+        examples=[1787.5],
+    )