kiapi 0.1.0__py3-none-any.whl

kiapi/core/setup/_services/setup_manager.py

@@ -0,0 +1,331 @@
+import shutil
+import subprocess
+import sys
+import tempfile
+import urllib.request
+from pathlib import Path
+
+from huggingface_hub import scan_cache_dir, snapshot_download
+
+from .._exceptions.setup_required_error import SetupRequiredError
+from .._schemas.docker_image_resource import DockerImageResource
+from .._schemas.hf_snapshot_resource import HfSnapshotResource
+from .._schemas.local_path_resource import LocalPathResource
+from .._schemas.python_package_resource import PythonPackageResource
+from .._schemas.python_venv_resource import PythonVenvResource
+from .._schemas.setup_status import SetupStatus
+from .._schemas.url_file_resource import UrlFileResource
+from .._types.setup_resource import SetupResource
+from .._types.setup_target import SetupTarget
+
+
+class SetupManager:
+    def status(self, resource: SetupResource) -> SetupStatus:
+        if isinstance(resource, HfSnapshotResource):
+            return self._hf_status(resource)
+        if isinstance(resource, DockerImageResource):
+            return self._docker_status(resource)
+        if isinstance(resource, LocalPathResource):
+            path = Path(resource.path).expanduser()
+            if path.exists():
+                return SetupStatus(True, str(path))
+            return SetupStatus(False, f"missing: {path}")
+        if isinstance(resource, UrlFileResource):
+            path = Path(resource.path).expanduser()
+            if path.exists() and path.stat().st_size > 0:
+                return SetupStatus(True, str(path))
+            return SetupStatus(False, f"missing: {path}")
+        if isinstance(resource, PythonVenvResource):
+            return self._python_venv_status(resource)
+        if isinstance(resource, PythonPackageResource):
+            return self._python_package_status(resource)
+
+    def activate(self, resource: SetupResource) -> SetupStatus:
+        if isinstance(resource, HfSnapshotResource):
+            snapshot_path = snapshot_download(
+                repo_id=resource.repo,
+                revision=resource.revision,
+                local_dir=resource.local_dir,
+            )
+            return SetupStatus(True, str(snapshot_path))
+        if isinstance(resource, DockerImageResource):
+            self._docker("pull", resource.image)
+            return self.status(resource)
+        if isinstance(resource, LocalPathResource):
+            return self.status(resource)
+        if isinstance(resource, UrlFileResource):
+            file_path = Path(resource.path).expanduser()
+            file_path.parent.mkdir(parents=True, exist_ok=True)
+            if not file_path.exists():
+                with tempfile.NamedTemporaryFile(
+                    dir=file_path.parent, delete=False
+                ) as tmp_file:
+                    tmp_path = Path(tmp_file.name)
+                try:
+                    urllib.request.urlretrieve(resource.url, tmp_path)
+                    tmp_path.replace(file_path)
+                finally:
+                    tmp_path.unlink(missing_ok=True)
+            return self.status(resource)
+        if isinstance(resource, PythonVenvResource):
+            return self._python_venv_activate(resource)
+        if isinstance(resource, PythonPackageResource):
+            return self._python_package_activate(resource)
+
+    def deactivate(self, resource: SetupResource) -> SetupStatus:
+        if isinstance(resource, HfSnapshotResource):
+            return self._hf_deactivate(resource)
+        if isinstance(resource, DockerImageResource):
+            self._docker("image", "rm", resource.image)
+            return self.status(resource)
+        if isinstance(resource, LocalPathResource):
+            path = Path(resource.path).expanduser()
+            if path.exists():
+                if path.is_dir():
+                    shutil.rmtree(path)
+                else:
+                    path.unlink()
+            return self.status(resource)
+        if isinstance(resource, UrlFileResource):
+            Path(resource.path).expanduser().unlink(missing_ok=True)
+            return self.status(resource)
+        if isinstance(resource, PythonVenvResource):
+            path = Path(resource.path).expanduser()
+            if path.exists():
+                shutil.rmtree(path)
+            return self.status(resource)
+        if isinstance(resource, PythonPackageResource):
+            return self._python_package_deactivate(resource)
+
+    def ensure_ready(self, spec: SetupTarget) -> None:
+        missing = [
+            resource
+            for resource in spec.setup_resources
+            if not self.status(resource).ready
+        ]
+        if not missing:
+            return
+
+        labels = ", ".join(resource.label for resource in missing)
+        repos = " ".join(f"--repo {resource.label}" for resource in missing)
+        raise SetupRequiredError(
+            f"model {spec.name!r} is not activated; missing: {labels}. "
+            f"Run: kiapi activate {repos}"
+        )
+
+    def _hf_status(self, resource: HfSnapshotResource) -> SetupStatus:
+        if resource.local_dir is not None:
+            local_path = Path(resource.local_dir).expanduser()
+            if local_path.exists() and any(local_path.iterdir()):
+                return SetupStatus(True, str(local_path))
+            return SetupStatus(False, f"missing: {local_path}")
+
+        try:
+            snapshot_path = snapshot_download(
+                repo_id=resource.repo,
+                revision=resource.revision,
+                local_dir=resource.local_dir,
+                local_files_only=True,
+            )
+        except Exception as exc:
+            return SetupStatus(False, str(exc))
+        return SetupStatus(True, str(snapshot_path))
+
+    def _hf_deactivate(self, resource: HfSnapshotResource) -> SetupStatus:
+        if resource.local_dir is not None:
+            path = Path(resource.local_dir).expanduser()
+            if path.exists():
+                shutil.rmtree(path)
+            return self.status(resource)
+
+        info = scan_cache_dir()
+        revisions: list[str] = []
+        for repo in info.repos:
+            if repo.repo_id != resource.repo:
+                continue
+            for revision in repo.revisions:
+                revisions.append(revision.commit_hash)
+
+        if revisions:
+            strategy = info.delete_revisions(*revisions)
+            strategy.execute()
+
+        return self.status(resource)
+
+    def _docker_status(self, resource: DockerImageResource) -> SetupStatus:
+        try:
+            self._docker("image", "inspect", resource.image)
+        except subprocess.CalledProcessError as exc:
+            detail = (exc.stderr or exc.stdout or str(exc)).strip()
+            return SetupStatus(False, detail)
+        return SetupStatus(True, "present")
+
+    def _docker(self, *args: str) -> subprocess.CompletedProcess[str]:
+        return subprocess.run(
+            ["docker", *args],
+            check=True,
+            text=True,
+            capture_output=True,
+        )
+
+    def _python_venv_status(self, resource: PythonVenvResource) -> SetupStatus:
+        python_path = self._python_venv_python_path(resource)
+        if not python_path.exists():
+            return SetupStatus(False, f"missing: {python_path}")
+        try:
+            self._run_python_import(python_path, resource.import_name)
+        except subprocess.CalledProcessError as exc:
+            detail = (exc.stderr or exc.stdout or str(exc)).strip()
+            return SetupStatus(False, detail)
+        return SetupStatus(True, str(python_path))
+
+    def _python_venv_activate(self, resource: PythonVenvResource) -> SetupStatus:
+        state = self.status(resource)
+        if state.ready:
+            return state
+
+        path = Path(resource.path).expanduser()
+        python_path = self._python_venv_python_path(resource)
+        if not python_path.exists():
+            self._create_python_venv(path, resource.python)
+        self._install_python_packages(python_path, resource.packages)
+        return self.status(resource)
+
+    def _python_venv_python_path(self, resource: PythonVenvResource) -> Path:
+        return Path(resource.path).expanduser() / "bin" / "python"
+
+    def _run_python_import(self, python_path: Path, import_name: str) -> None:
+        subprocess.run(
+            [str(python_path), "-c", f"import {import_name}"],
+            check=True,
+            text=True,
+            capture_output=True,
+        )
+
+    def _python_package_status(self, resource: PythonPackageResource) -> SetupStatus:
+        try:
+            self._run_python_package_verify(resource)
+        except subprocess.CalledProcessError as exc:
+            detail = (exc.stderr or exc.stdout or str(exc)).strip()
+            return SetupStatus(False, detail)
+        return SetupStatus(True, resource.import_name)
+
+    def _python_package_activate(self, resource: PythonPackageResource) -> SetupStatus:
+        state = self.status(resource)
+        if state.ready:
+            return state
+
+        self._install_python_packages(Path(sys.executable), (resource.spec,))
+        return self.status(resource)
+
+    def _python_package_deactivate(
+        self, resource: PythonPackageResource
+    ) -> SetupStatus:
+        self._uninstall_python_package(Path(sys.executable), resource.package)
+        return self.status(resource)
+
+    def _run_python_package_verify(self, resource: PythonPackageResource) -> None:
+        code = "\n".join(
+            [
+                "import importlib",
+                f"module = importlib.import_module({resource.import_name!r})",
+                f"attrs = {list(resource.verify_attrs)!r}",
+                "missing = [attr for attr in attrs if not hasattr(module, attr)]",
+                "if missing:",
+                "    raise RuntimeError(f'missing attrs: {missing}')",
+            ]
+        )
+        subprocess.run(
+            [sys.executable, "-c", code],
+            check=True,
+            text=True,
+            capture_output=True,
+        )
+
+    def _create_python_venv(self, path: Path, python: str) -> None:
+        if shutil.which("uv"):
+            subprocess.run(
+                ["uv", "venv", "--python", python, str(path)],
+                check=True,
+                text=True,
+                capture_output=True,
+            )
+            return
+
+        python_bin = self._resolve_python_executable(python)
+        subprocess.run(
+            [str(python_bin), "-m", "venv", str(path)],
+            check=True,
+            text=True,
+            capture_output=True,
+        )
+
+    def _install_python_packages(
+        self, python_path: Path, packages: tuple[str, ...]
+    ) -> None:
+        if shutil.which("uv"):
+            subprocess.run(
+                ["uv", "pip", "install", "--python", str(python_path), *packages],
+                check=True,
+                text=True,
+                capture_output=True,
+            )
+            return
+
+        subprocess.run(
+            [str(python_path), "-m", "pip", "install", *packages],
+            check=True,
+            text=True,
+            capture_output=True,
+        )
+
+    def _uninstall_python_package(self, python_path: Path, package: str) -> None:
+        if shutil.which("uv"):
+            subprocess.run(
+                ["uv", "pip", "uninstall", "--python", str(python_path), package],
+                check=True,
+                text=True,
+                capture_output=True,
+            )
+            return
+
+        subprocess.run(
+            [str(python_path), "-m", "pip", "uninstall", "-y", package],
+            check=True,
+            text=True,
+            capture_output=True,
+        )
+
+    def _resolve_python_executable(self, python: str) -> Path:
+        candidates = [python]
+        if python.startswith("3."):
+            candidates.append(f"python{python}")
+        candidates.extend(["python3", sys.executable])
+
+        for candidate in candidates:
+            executable = (
+                shutil.which(candidate) if candidate != sys.executable else candidate
+            )
+            if executable and self._python_matches_requirement(
+                Path(executable), python
+            ):
+                return Path(executable)
+
+        raise RuntimeError(f"Python {python} executable was not found")
+
+    def _python_matches_requirement(self, python_path: Path, python: str) -> bool:
+        if not python.startswith("3."):
+            return True
+
+        code = "\n".join(
+            [
+                "import sys",
+                f"raise SystemExit(0 if sys.version_info[:2] == {tuple(map(int, python.split('.')))!r} else 1)",
+            ]
+        )
+        result = subprocess.run(
+            [str(python_path), "-c", code],
+            text=True,
+            capture_output=True,
+        )
+        return result.returncode == 0

kiapi/core/setup/_types/setup_resource.py

@@ -0,0 +1,21 @@
+from typing import Annotated
+
+from pydantic import Field
+
+from .._schemas.docker_image_resource import DockerImageResource
+from .._schemas.hf_snapshot_resource import HfSnapshotResource
+from .._schemas.local_path_resource import LocalPathResource
+from .._schemas.python_package_resource import PythonPackageResource
+from .._schemas.python_venv_resource import PythonVenvResource
+from .._schemas.url_file_resource import UrlFileResource
+
+SetupResource = Annotated[
+    HfSnapshotResource
+    | DockerImageResource
+    | LocalPathResource
+    | UrlFileResource
+    | PythonPackageResource
+    | PythonVenvResource,
+    Field(discriminator="kind"),
+]
+"""A resource that must be activated before a model can run, tagged by ``kind``."""

kiapi/core/setup/_types/setup_target.py

@@ -0,0 +1,11 @@
+from typing import Protocol
+
+from .setup_resource import SetupResource
+
+
+class SetupTarget(Protocol):
+    @property
+    def name(self) -> str: ...
+
+    @property
+    def setup_resources(self) -> tuple[SetupResource, ...]: ...

kiapi/core/workdir/__init__.py

@@ -0,0 +1,8 @@
+from ._helpers.create_work_dir import create_work_dir
+from ._settings import WorkDirSettings, settings_manager
+
+__all__ = [
+    "WorkDirSettings",
+    "create_work_dir",
+    "settings_manager",
+]

kiapi/core/workdir/_helpers/create_work_dir.py

@@ -0,0 +1,11 @@
+import tempfile
+from pathlib import Path
+
+from .._settings import settings_manager
+
+
+def create_work_dir(purpose: str) -> Path:
+    root = Path(settings_manager.get_settings().tmp_root).expanduser()
+    parent = root / purpose
+    parent.mkdir(parents=True, exist_ok=True)
+    return Path(tempfile.mkdtemp(dir=parent))

kiapi/core/workdir/_settings.py

@@ -0,0 +1,21 @@
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+
+class WorkDirSettings(BaseSettings):
+    """Settings for per-request temporary work directories."""
+
+    model_config = SettingsConfigDict(env_prefix="KIAPI_")
+
+    tmp_root: str = Field(
+        default="/tmp/kiapi/work",
+        title="Temporary work directory root",
+        description=(
+            "Root directory for temporary work files created while handling requests.\n"
+            "Each operation creates a purpose-specific child directory under this root."
+        ),
+    )
+
+
+settings_manager = SettingsManager(WorkDirSettings)

kiapi/core/worker/__init__.py

@@ -0,0 +1,23 @@
+"""The single-flight worker: one dedicated thread runs every job, one at a time
+(MLX thread-affinity + sound memory budgeting). See :mod:`._services.worker`.
+
+Public surface (unchanged from the former ``core/worker.py``): :class:`Worker`,
+:data:`JobThunk`.
+"""
+
+from ._helpers.create_worker import create_worker
+from ._services.worker import Worker
+from ._settings import WorkerSettings, settings_manager
+from ._types.job_thunk import JobThunk
+
+__all__ = [
+    # ._types
+    "JobThunk",
+    # ._services
+    "Worker",
+    # ._settings
+    "WorkerSettings",
+    # ._helpers
+    "create_worker",
+    "settings_manager",
+]

kiapi/core/worker/_helpers/create_worker.py

@@ -0,0 +1,9 @@
+from kiapi.core.app import AppContext
+
+from .._services.worker import Worker
+from .._settings import settings_manager
+
+
+def create_worker(ctx: AppContext) -> Worker:
+    settings = settings_manager.get_settings()
+    return Worker(settings, ctx)

kiapi/core/worker/_services/worker.py

@@ -0,0 +1,165 @@
+"""The single-flight worker — kiapi's one hard rule: only one job runs at a time.
+
+All heavy work (model load, eviction, warmup, every generation, across every
+capability) runs on ONE dedicated thread (``ThreadPoolExecutor(max_workers=1)``).
+This is required for MLX correctness (arrays/streams are thread-affine) AND it is
+what makes the global memory budget sound: the model being acquired for the
+running job is the only one generating, so budgeting just its peak headroom is
+correct (see core/memory). Concurrent HTTP requests are accepted and simply queue.
+
+If you need real parallelism, run multiple kiapi processes, each scoped to a
+subset of capabilities with its own slice of the memory budget.
+
+A job's work is supplied as a JobThunk. The worker drives the job's lifecycle
+(running → succeeded/failed) and resolves the caller's future. Sync callers await
+the future (with a timeout); async callers drop it and poll the job store.
+"""
+
+import asyncio
+import logging
+import traceback
+from concurrent.futures import ThreadPoolExecutor
+
+from kiapi.core.app import AppContext
+from kiapi.core.job import Job, ProgressReporter
+from kiapi.core.model import ModelKey
+from kiapi.core.setup import SetupRequiredError
+
+from .._settings import WorkerSettings
+from .._types.job_thunk import JobThunk
+
+logger = logging.getLogger(__name__)
+
+
+class Worker:
+    def __init__(
+        self,
+        settings: WorkerSettings,
+        ctx: AppContext,
+    ) -> None:
+        self.settings: WorkerSettings = settings
+        self.ctx: AppContext = ctx
+        self.queue: asyncio.Queue[tuple[Job, JobThunk, asyncio.Future]] = (
+            asyncio.Queue()
+        )
+        self.executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="kiapi")
+        self.warm = False
+        self._task: asyncio.Task | None = None
+        self._ttl_task: asyncio.Task | None = None
+
+    def start(self) -> None:
+        self._task = asyncio.create_task(self._run(), name="kiapi-worker")
+        if self.settings.ttl_sweep_interval_s > 0:
+            self._ttl_task = asyncio.create_task(
+                self._ttl_loop(), name="kiapi-ttl-sweep"
+            )
+
+    async def stop(self) -> None:
+        if self._task is not None:
+            self._task.cancel()
+        if self._ttl_task is not None:
+            self._ttl_task.cancel()
+        # Free resident models on the worker thread, then tear it down.
+        try:
+            loop = asyncio.get_running_loop()
+            await loop.run_in_executor(self.executor, self.ctx.memory_manager.shutdown)
+        except Exception:
+            traceback.print_exc()
+        self.executor.shutdown(wait=False, cancel_futures=True)
+
+    async def submit(self, job: Job, fn: JobThunk) -> asyncio.Future:
+        loop = asyncio.get_running_loop()
+        fut: asyncio.Future = loop.create_future()
+        await self.queue.put((job, fn, fut))
+        return fut
+
+    async def _run(self) -> None:
+        loop = asyncio.get_running_loop()
+        while True:
+            job, fn, fut = await self.queue.get()
+            job.mark_running()
+            # Bind a reporter for this job so capability code can push coarse
+            # progress via ProgressReporter.current(); run_bound installs it on
+            # the worker thread, where the thunk actually executes.
+            reporter = ProgressReporter(job)
+            try:
+                result, artifacts = await loop.run_in_executor(
+                    self.executor, reporter.run_bound, fn
+                )
+                job.mark_succeeded(result, artifacts)
+                if not fut.done():
+                    fut.set_result(result)
+            except Exception as exc:
+                traceback.print_exc()
+                job.mark_failed(str(exc))
+                if not fut.done():
+                    fut.set_exception(exc)
+
+    async def _ttl_loop(self) -> None:
+        """The sweep runs on the worker thread (shared executor), so it serializes
+        with generations and respects MLX thread-affinity — it frees models past
+        their TTL even when no requests arrive."""
+        loop = asyncio.get_running_loop()
+        interval = self.settings.ttl_sweep_interval_s
+        while True:
+            await asyncio.sleep(interval)
+            try:
+                freed = await loop.run_in_executor(
+                    self.executor, self.ctx.memory_manager.sweep_expired
+                )
+                if freed:
+                    logger.info("idle sweep freed: %s", freed)
+            except Exception:
+                traceback.print_exc()
+
+    async def warmup(self) -> None:
+        """Each model is loaded + primed on the worker thread; acquire evicts as
+        needed to respect the memory budget."""
+        model_names = self.settings.warmup_models
+        if not model_names:
+            self.warm = True
+            return
+        loop = asyncio.get_running_loop()
+        for name in model_names:
+            try:
+                await loop.run_in_executor(self.executor, self._warm_one, name)
+            except Exception:
+                traceback.print_exc()
+        self.warm = True
+
+    def _warm_one(self, name: ModelKey) -> None:
+        """Runs on the worker thread."""
+        from kiapi.core.model import model_registry
+
+        # A warmup name is a model name; find which family owns it.
+        spec = None
+        for family in model_registry.families():
+            try:
+                spec = model_registry.resolve(family, name)
+                break
+            except Exception:
+                continue
+        if spec is None:
+            logger.warning("unknown warmup model %r; skipping", name)
+            return
+        if not spec.resident:
+            # Transient providers (e.g. video) load/free per call — nothing to
+            # keep resident. Prime via the module's own warmup if it has one.
+            warm = getattr(spec.module, "warmup", None)
+            if warm is not None:
+                try:
+                    self.ctx.ensure_model_ready(spec)
+                except SetupRequiredError as exc:
+                    logger.warning("skipping warmup for %s: %s", spec.name, exc)
+                    return
+                warm(None)
+            return
+        try:
+            self.ctx.ensure_model_ready(spec)
+        except SetupRequiredError as exc:
+            logger.warning("skipping warmup for %s: %s", spec.name, exc)
+            return
+        payload = self.ctx.memory_manager.acquire(spec)
+        warm = getattr(spec.module, "warmup", None)
+        if warm is not None:
+            warm(payload)

kiapi/core/worker/_settings.py

@@ -0,0 +1,36 @@
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+from kiapi.core.model import ModelKey
+
+
+class WorkerSettings(BaseSettings):
+    """Settings for the single worker queue, TTL sweeping, and startup warmup."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="KIAPI_",
+        extra="ignore",
+    )
+
+    ttl_sweep_interval_s: float = Field(
+        default=60.0,
+        title="TTL sweep interval seconds",
+        description=(
+            "Interval for background checks of resident models that exceeded their idle TTL.\n"
+            "Set to 0.0 to completely disable the background sweep."
+        ),
+    )
+
+    warmup_models: list[ModelKey] = Field(
+        default_factory=list,
+        title="Startup warmup models",
+        description=(
+            "List of model keys to preload when the server starts.\n"
+            "Each item uses the model registry key format. An empty list keeps "
+            "models lazily loaded."
+        ),
+    )
+
+
+settings_manager = SettingsManager(WorkerSettings)

kiapi/core/worker/_types/job_thunk.py

@@ -0,0 +1,6 @@
+from collections.abc import Callable
+
+from kiapi.core.file import FileID
+from kiapi.core.job import JobResult
+
+type JobThunk = Callable[[], tuple[JobResult, list[FileID]]]