Initial public ModelForge release
This commit is contained in:
Jens
@@ -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()
|
||||
Reference in New Issue
Block a user