jevstiller 0.2.0__py3-none-any.whl

jevstiller/__init__.py

@@ -0,0 +1,19 @@
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _version
+
+from .core import Jevstiller, Result, Status, TeacherError, TrainReport
+from .encoders import BatchingEncoder, Encoder, HashEncoder, load_encoder
+from .manager import Admission, TaskManager
+from .scheduler import TrainScheduler
+from .task import Config, Task
+from .teachers import CachedTeacher, ReplayTeacher, SyntheticTeacher, SyntheticWorld, Teacher, TeacherOutput
+
+__all__ = ["Task", "Config", "Jevstiller", "Result", "Status", "TrainReport", "TeacherError", "Teacher",
+           "TeacherOutput", "SyntheticTeacher", "SyntheticWorld", "ReplayTeacher", "CachedTeacher", "Encoder",
+           "HashEncoder", "BatchingEncoder", "load_encoder", "TaskManager", "Admission",
+           "TrainScheduler"]
+
+try:
+    __version__ = _version("jevstiller")
+except PackageNotFoundError:  # pragma: no cover - running from a source tree that was never installed
+    __version__ = "0.0.0+unknown"

jevstiller/admin.py

@@ -0,0 +1,232 @@
+"""Admin API, under `/jevstiller/v1/` (never forwarded to Jev). Off unless an admin token is configured; every
+call needs `Authorization: Bearer <admin token>`.
+
+    GET    /jevstiller/v1/tasks[?tenant=]            list tasks
+    GET    /jevstiller/v1/tasks/{key}                status (+ text report) of one task
+    GET    /jevstiller/v1/tasks/{key}/versions       its student versions
+    POST   /jevstiller/v1/tasks/{key}/mode           {"mode": "auto" | "teacher_only" | "cascade"}   (persisted)
+    POST   /jevstiller/v1/tasks/{key}/target         {"target_agreement": 0.99}                      (persisted)
+    POST   /jevstiller/v1/tasks/{key}/train          train a candidate now (waits for it)
+    POST   /jevstiller/v1/tasks/{key}/promote        {"version": "student:v3"}
+    POST   /jevstiller/v1/tasks/{key}/rollback
+    DELETE /jevstiller/v1/tasks/{key}                delete the task and all its data
+    DELETE /jevstiller/v1/tenants/{tenant}           delete every task of a tenant
+    GET    /jevstiller/v1/stats                      proxy, manager and scheduler counters
+"""
+from __future__ import annotations
+
+import dataclasses
+import hmac
+import math
+import re
+from collections.abc import Callable
+from typing import Any
+
+import anyio
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response
+from starlette.routing import Route
+
+from .manager import TaskManager
+
+PREFIX = "/jevstiller/v1"
+_KEY = re.compile(r"^[0-9a-f]{20}$")
+
+
+def _json_safe(x: Any) -> Any:
+    if isinstance(x, float) and (math.isnan(x) or math.isinf(x)):
+        return None
+    if isinstance(x, dict):
+        return {str(k): _json_safe(v) for k, v in x.items()}
+    if isinstance(x, (list, tuple)):
+        return [_json_safe(v) for v in x]
+    if dataclasses.is_dataclass(x) and not isinstance(x, type):
+        return _json_safe(dataclasses.asdict(x))
+    return x
+
+
+def _ok(data: Any, status: int = 200) -> JSONResponse:
+    return JSONResponse(_json_safe(data), status_code=status)
+
+
+def _err(status: int, detail: str) -> JSONResponse:
+    return JSONResponse({"detail": detail}, status_code=status)
+
+
+class Admin:
+    def __init__(self, manager: TaskManager, token: str | None, stats: Callable[[], dict] | None = None):
+        self.manager, self.token, self.stats_fn = manager, token, stats
+
+    def authorized(self, request: Request) -> Response | None:
+        """None when the call may proceed, else the response to send."""
+        if not self.token:
+            return _err(404, "not found")                    # the admin API is off
+        auth = request.headers.get("authorization", "")
+        given = auth[7:].strip() if auth.lower().startswith("bearer ") else ""
+        if not hmac.compare_digest(given.encode(), self.token.encode()):
+            return _err(401, "admin token required")
+        return None
+
+    def _key(self, request: Request) -> str | None:
+        key = request.path_params.get("key", "")
+        return key if _KEY.match(key) and key in {i.key for i in self.manager.tasks()} else None
+
+    async def _with_engine(self, key: str, fn: Callable[[Any], Any]) -> Any:
+        """`fn(engine)` in a worker thread with the task's engine held open throughout (an unload or delete
+        closing it under a running call crashed the process: security audit run 2). None if the task is gone."""
+        def run():
+            try:
+                with self.manager.hold(key) as engine:
+                    return fn(engine)
+            except KeyError:                                # deleted meanwhile
+                return None
+        return await anyio.to_thread.run_sync(run)
+
+    async def _body(self, request: Request) -> dict:
+        try:
+            body = await request.json()
+        except ValueError:
+            return {}
+        return body if isinstance(body, dict) else {}
+
+    # ---- handlers -----------------------------------------------------------------
+    async def tasks(self, request: Request) -> Response:
+        if (denied := self.authorized(request)) is not None:
+            return denied
+        loaded = set(self.manager.loaded())
+        tenant = request.query_params.get("tenant")
+        return _ok([{"key": i.key, "tenant": i.tenant, "model": i.model, "classes": len(i.classes),
+                     "target_agreement": i.target_agreement, "mode": i.mode, "created": i.created,
+                     "last_seen": i.last_seen, "loaded": i.key in loaded} for i in self.manager.tasks(tenant)])
+
+    async def task(self, request: Request) -> Response:
+        if (denied := self.authorized(request)) is not None:
+            return denied
+        if (key := self._key(request)) is None:
+            return _err(404, "unknown task")
+        st = await self._with_engine(key, lambda e: e.status())
+        info = next((i for i in self.manager.tasks() if i.key == key), None)
+        if st is None or info is None:
+            return _err(404, "unknown task")
+        return _ok({"info": info, "status": st, "report": st.report()})
+
+    async def versions(self, request: Request) -> Response:
+        if (denied := self.authorized(request)) is not None:
+            return denied
+        if (key := self._key(request)) is None:
+            return _err(404, "unknown task")
+        versions = await self._with_engine(key, lambda e: e.versions())
+        return _ok(versions) if versions is not None else _err(404, "unknown task")
+
+    async def mode(self, request: Request) -> Response:
+        if (denied := self.authorized(request)) is not None:
+            return denied
+        if (key := self._key(request)) is None:
+            return _err(404, "unknown task")
+        mode = (await self._body(request)).get("mode")
+        try:
+            await anyio.to_thread.run_sync(self.manager.set_mode, key, mode)
+        except ValueError as e:
+            return _err(422, str(e))
+        return _ok({"key": key, "mode": mode})
+
+    async def target(self, request: Request) -> Response:
+        if (denied := self.authorized(request)) is not None:
+            return denied
+        if (key := self._key(request)) is None:
+            return _err(404, "unknown task")
+        value = (await self._body(request)).get("target_agreement")
+        if not isinstance(value, (int, float)) or isinstance(value, bool):
+            return _err(422, "target_agreement must be a number in [0.5, 1)")
+        try:
+            await anyio.to_thread.run_sync(self.manager.set_target, key, float(value))
+        except ValueError as e:
+            return _err(422, str(e))
+        return _ok({"key": key, "target_agreement": float(value)})
+
+    async def train(self, request: Request) -> Response:
+        if (denied := self.authorized(request)) is not None:
+            return denied
+        if (key := self._key(request)) is None:
+            return _err(404, "unknown task")
+        try:
+            report = await self._with_engine(key, lambda e: e.train_now())
+        except Exception as e:
+            return _err(500, f"training failed: {type(e).__name__}")
+        return _ok(report) if report is not None else _err(404, "unknown task")
+
+    async def promote(self, request: Request) -> Response:
+        if (denied := self.authorized(request)) is not None:
+            return denied
+        if (key := self._key(request)) is None:
+            return _err(404, "unknown task")
+        version = (await self._body(request)).get("version")
+        try:
+            if await self._with_engine(key, lambda e: e.promote(str(version)) or True) is None:
+                return _err(404, "unknown task")
+        except ValueError as e:
+            return _err(422, str(e))
+        return _ok({"key": key, "production": version})
+
+    async def rollback(self, request: Request) -> Response:
+        if (denied := self.authorized(request)) is not None:
+            return denied
+        if (key := self._key(request)) is None:
+            return _err(404, "unknown task")
+        target = await self._with_engine(key, lambda e: (e.rollback(),))
+        if target is None:
+            return _err(404, "unknown task")
+        return _ok({"key": key, "production": target[0]})
+
+    async def delete_task(self, request: Request) -> Response:
+        if (denied := self.authorized(request)) is not None:
+            return denied
+        if (key := self._key(request)) is None:
+            return _err(404, "unknown task")
+        if not await anyio.to_thread.run_sync(self.manager.delete, key, "admin API"):
+            return _err(409, "a request is in flight for this task; retry")
+        return _ok({"deleted": key})
+
+    async def delete_tenant(self, request: Request) -> Response:
+        if (denied := self.authorized(request)) is not None:
+            return denied
+        tenant = request.path_params["tenant"]
+        deleted = await anyio.to_thread.run_sync(self.manager.delete_tenant, tenant)
+        left = [i.key for i in self.manager.tasks(tenant)]
+        return _ok({"deleted": deleted, "remaining": left}, 200 if not left else 409)
+
+    async def stats(self, request: Request) -> Response:
+        if (denied := self.authorized(request)) is not None:
+            return denied
+        m = self.manager
+        data = {"tasks": len(m.tasks()), "loaded": len(m.loaded()), "loads": m.loads, "unloads": m.unloads,
+                "student_memory_mb": m.memory_mb()}
+        sched = getattr(m.train_executor, "stats", None)
+        if callable(sched):
+            data["training"] = sched()
+        if self.stats_fn:
+            data.update(self.stats_fn())
+        return _ok(data)
+
+    def routes(self) -> list[Route]:
+        p = PREFIX
+        return [
+            Route(f"{p}/tasks", self.tasks, methods=["GET"]),
+            Route(f"{p}/tasks/{{key}}", self.task, methods=["GET"]),
+            Route(f"{p}/tasks/{{key}}", self.delete_task, methods=["DELETE"]),
+            Route(f"{p}/tasks/{{key}}/versions", self.versions, methods=["GET"]),
+            Route(f"{p}/tasks/{{key}}/mode", self.mode, methods=["POST"]),
+            Route(f"{p}/tasks/{{key}}/target", self.target, methods=["POST"]),
+            Route(f"{p}/tasks/{{key}}/train", self.train, methods=["POST"]),
+            Route(f"{p}/tasks/{{key}}/promote", self.promote, methods=["POST"]),
+            Route(f"{p}/tasks/{{key}}/rollback", self.rollback, methods=["POST"]),
+            Route(f"{p}/tenants/{{tenant}}", self.delete_tenant, methods=["DELETE"]),
+            Route(f"{p}/stats", self.stats, methods=["GET"]),
+            Route(f"{p}/{{rest:path}}", self._not_found,
+                  methods=["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"]),
+        ]
+
+    async def _not_found(self, request: Request) -> Response:
+        if (denied := self.authorized(request)) is not None:
+            return denied
+        return _err(404, "not found")

jevstiller/backup.py

@@ -0,0 +1,112 @@
+"""`jevstiller backup` / `jevstiller restore`: a consistent copy of a data directory, safe while serving.
+
+Per task: SQLite's online backup API for `samples.sqlite` (a consistent snapshot even with the writer running),
+`task.json`, and `versions/` (the registry index is copied first, so every version it lists is present in the
+copy). The deployment's `key-salt` is copied too: without it, key hashes in a tenants file would no longer match.
+
+Backups hold request text like the data dir, so they get the same modes (files 0600, directories 0700), and a
+restored data dir is set to them too, whatever the source's modes (security audit run 2).
+"""
+from __future__ import annotations
+
+import contextlib
+import json
+import os
+import shutil
+import sqlite3
+import time
+from collections.abc import Iterator
+from pathlib import Path
+
+
+@contextlib.contextmanager
+def _private() -> Iterator[None]:
+    """Everything created inside is owner-only (umask 077), as `jevstiller serve` does."""
+    old = os.umask(0o077)
+    try:
+        yield
+    finally:
+        os.umask(old)
+
+
+def _lock_down(root: Path) -> None:
+    """Owner-only modes for `root` and everything below it (copies keep their source's modes)."""
+    if not root.exists():
+        return
+    os.chmod(root, 0o700 if root.is_dir() else 0o600)
+    for dirpath, dirnames, filenames in os.walk(root):
+        for d in dirnames:
+            os.chmod(os.path.join(dirpath, d), 0o700)
+        for f in filenames:
+            os.chmod(os.path.join(dirpath, f), 0o600)
+
+
+def _copy_versions(src: Path, dst: Path) -> None:
+    dst.mkdir(parents=True, exist_ok=True)
+    if (src / "registry.json").exists():
+        shutil.copy2(src / "registry.json", dst / "registry.json")
+    for d in src.iterdir():
+        if d.is_dir() and not d.name.startswith("."):          # skip staging dirs of running jobs
+            shutil.copytree(d, dst / d.name, dirs_exist_ok=True)
+
+
+def backup(data_dir: str | Path, out: str | Path) -> dict:
+    with _private():
+        manifest = _backup(Path(data_dir), Path(out))
+        _lock_down(Path(out))
+    return manifest
+
+
+def _backup(src: Path, dst: Path) -> dict:
+    if dst.exists() and any(dst.iterdir()):
+        raise FileExistsError(f"{dst} is not empty")
+    dst.mkdir(parents=True, exist_ok=True)
+    tasks = 0
+    for tdir in sorted((src / "tasks").glob("*")) if (src / "tasks").exists() else []:
+        if not (tdir / "task.json").exists():
+            continue
+        out_t = dst / "tasks" / tdir.name
+        out_t.mkdir(parents=True)
+        shutil.copy2(tdir / "task.json", out_t / "task.json")
+        if (tdir / "samples.sqlite").exists():
+            s = sqlite3.connect(f"{(tdir / 'samples.sqlite').resolve().as_uri()}?mode=ro", uri=True)
+            d = sqlite3.connect(out_t / "samples.sqlite")
+            try:
+                s.backup(d)
+            finally:
+                d.close()
+                s.close()
+        if (tdir / "versions").exists():
+            _copy_versions(tdir / "versions", out_t / "versions")
+        tasks += 1
+    if (src / "key-salt").exists():
+        shutil.copy2(src / "key-salt", dst / "key-salt")
+    manifest = {"created": time.time(), "source": str(src.resolve()), "tasks": tasks}
+    (dst / "backup.json").write_text(json.dumps(manifest, indent=2))
+    return manifest
+
+
+def restore(backup_dir: str | Path, data_dir: str | Path, force: bool = False) -> dict:
+    """Copy a backup into an empty data directory (or over an existing one with `force`). Stop the server first."""
+    with _private():
+        manifest = _restore(Path(backup_dir), Path(data_dir), force)
+        dst = Path(data_dir)
+        os.chmod(dst, 0o700)
+        _lock_down(dst / "tasks")
+        _lock_down(dst / "key-salt")
+    return manifest
+
+
+def _restore(src: Path, dst: Path, force: bool) -> dict:
+    if not (src / "backup.json").exists():
+        raise FileNotFoundError(f"{src} is not a jevstiller backup (no backup.json)")
+    if dst.exists() and any(dst.iterdir()) and not force:
+        raise FileExistsError(f"{dst} is not empty (use force to overwrite)")
+    dst.mkdir(parents=True, exist_ok=True)
+    if (dst / "tasks").exists() and force:
+        shutil.rmtree(dst / "tasks")
+    if (src / "tasks").exists():
+        shutil.copytree(src / "tasks", dst / "tasks")
+    if (src / "key-salt").exists():
+        shutil.copy2(src / "key-salt", dst / "key-salt")
+    return json.loads((src / "backup.json").read_text())

jevstiller/calibrate.py

@@ -0,0 +1,133 @@
+"""Threshold selection under a disagreement budget, with a finite-sample bound."""
+from __future__ import annotations
+
+import json
+import math
+from collections.abc import Sequence
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+
+import numpy as np
+
+
+def _log_binom_cdf(k: int, n: int, p: float) -> float:
+    """log P(X <= k) for X ~ Binomial(n, p)."""
+    if p <= 0:
+        return 0.0
+    if p >= 1:
+        return 0.0 if k >= n else -math.inf
+    lp, lq = math.log(p), math.log1p(-p)
+    terms = [math.lgamma(n + 1) - math.lgamma(i + 1) - math.lgamma(n - i + 1) + i * lp + (n - i) * lq
+             for i in range(k + 1)]
+    m = max(terms)
+    return m + math.log(sum(math.exp(t - m) for t in terms))
+
+
+def clopper_pearson_upper(k: int, n: int, delta: float = 0.05) -> float:
+    """One-sided upper confidence bound on a binomial proportion: smallest p with P(X<=k|n,p) <= delta."""
+    if n <= 0:
+        return 1.0
+    if k >= n:
+        return 1.0
+    lo, hi = k / n, 1.0
+    log_delta = math.log(delta)
+    for _ in range(60):
+        mid = (lo + hi) / 2
+        if _log_binom_cdf(k, n, mid) > log_delta:
+            lo = mid
+        else:
+            hi = mid
+    return hi
+
+
+def clopper_pearson_lower(k: int, n: int, delta: float = 0.05) -> float:
+    """One-sided lower confidence bound on a binomial proportion."""
+    if n <= 0 or k <= 0:
+        return 0.0
+    return 1.0 - clopper_pearson_upper(n - k, n, delta)
+
+
+@dataclass
+class RoutingPolicy:
+    conf_threshold: float | None      # None -> student never answers
+    ood_threshold: float
+    expected_coverage: float          # share of calibration rows the student answers
+    disagreement_ub: float            # upper bound (at 1 - delta) on P(student answers and disagrees), per request
+    expected_system_disagreement: float   # that rate on the calibration rows (the bound is what meets the budget)
+    budget: float
+    delta: float
+    n_calib: int
+    deferred_labels: list[str] = field(default_factory=list)   # the student never answers these (rare classes)
+
+    @property
+    def usable(self) -> bool:
+        return self.conf_threshold is not None
+
+    def accepts(self, conf: np.ndarray, ood: np.ndarray, pred: np.ndarray | None = None) -> np.ndarray:
+        """Which rows the student answers. `pred` (the student's predicted labels) is required when the
+        policy defers labels."""
+        if self.conf_threshold is None:
+            return np.zeros(len(conf), dtype=bool)
+        ok = (conf >= self.conf_threshold) & (ood <= self.ood_threshold)
+        if self.deferred_labels:
+            if pred is None:
+                raise ValueError("this policy defers labels; pass the predicted labels")
+            ok &= ~np.isin(np.asarray(pred, dtype=object), self.deferred_labels)
+        return ok
+
+    def save(self, path: Path) -> None:
+        path.write_text(json.dumps(asdict(self), indent=2))
+
+    @classmethod
+    def load(cls, path: Path) -> RoutingPolicy:
+        return cls(**json.loads(path.read_text()))
+
+
+def threshold_grid(n_labels: int, size: int = 400) -> np.ndarray:
+    """Candidate confidence thresholds, strictest first. Fixed before any calibration row is seen, so testing
+    them in sequence costs no confidence. Dense near 1, where a student's usable answers are."""
+    floor = 1.0 / max(n_labels, 2)
+    return 1.0 - np.geomspace(1e-4, 1.0 - floor, size)
+
+
+def fit_policy(conf: np.ndarray, agree: np.ndarray, ood: np.ndarray, budget: float, delta: float = 0.05,
+               ood_threshold: float = math.inf, candidates: np.ndarray | None = None,
+               eligible: np.ndarray | None = None, deferred_labels: Sequence[str] = ()) -> RoutingPolicy:
+    """The loosest confidence threshold whose rate of *answered and disagreeing* requests is at most `budget`,
+    with probability at least 1 - `delta`.
+
+    Rows come from an IID calibration set: the student's max-prob, whether it agrees with the teacher, and its
+    OOD score. The loss per row is 1[the student answers and disagrees], so its rate over all rows is exactly
+    what the budget limits (disagreement over all requests).
+
+    Candidates are tested strictest first with an exact Clopper-Pearson bound at `delta`, and the scan stops at
+    the first failure (fixed-sequence testing, as in "Learn Then Test"). The loss can only grow as the
+    threshold loosens, so the chance that the chosen threshold breaks the budget is at most `delta`, with no
+    multiple-testing penalty. That needs everything else fixed without these rows: `candidates` (default: a
+    fixed grid) and `ood_threshold` (the caller picks it on training data). `eligible` marks rows the student
+    may answer at all (False where it predicts one of `deferred_labels`); the rate still counts every row.
+    """
+    conf = np.asarray(conf, dtype=float)
+    agree = np.asarray(agree, dtype=bool)
+    ood = np.asarray(ood, dtype=float)
+    N = len(conf)
+    in_dist = ood <= ood_threshold
+    if eligible is not None:
+        in_dist &= np.asarray(eligible, dtype=bool)
+    grid = threshold_grid(2) if candidates is None else np.asarray(candidates, dtype=float)
+    deferred = list(deferred_labels)
+    best = None
+    if N:
+        for t in sorted(grid, reverse=True):
+            sel = in_dist & (conf >= t)
+            k = int((sel & ~agree).sum())
+            ub = clopper_pearson_upper(k, N, delta)
+            if ub > budget:
+                break                                    # fixed sequence: every looser threshold is untested
+            if sel.any():                                # a threshold that answers nothing is no policy
+                best = (float(t), int(sel.sum()) / N, ub, k / N)
+    ood_thr = float(ood_threshold) if math.isfinite(ood_threshold) else float(np.max(ood, initial=0.0))
+    if best is None:
+        return RoutingPolicy(None, ood_thr, 0.0, 1.0, 0.0, budget, delta, N, deferred)
+    t, cov, ub, rate = best
+    return RoutingPolicy(t, ood_thr, cov, ub, rate, budget, delta, N, deferred)