161 lines
6.5 KiB
Python
161 lines
6.5 KiB
Python
"""Register a verified native Shard through the ordinary capability contract.
|
|
|
|
This is intentionally an adapter, not a second tracker protocol. It converts
|
|
the native worker's immutable identity and enforced resource limits into the
|
|
same capability report every backend may submit. The tracker remains the sole
|
|
owner of certification and decides whether the visible registration is dark.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from .capability import ExecutionCapacity, build_capability_report
|
|
from .native_worker_supervisor import NativeWorkerProbe, NativeWorkerSpec, NativeWorkerSupervisor
|
|
from .runtime_recipe import ShardIdentity
|
|
|
|
|
|
class NativeRegistrationError(ValueError):
|
|
"""Native facts do not describe one coherent, registerable Shard."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class NativeShardRegistration:
|
|
"""One backend-neutral registration payload for a verified native Shard."""
|
|
|
|
endpoint: str
|
|
model_id: str
|
|
identity: ShardIdentity
|
|
worker: NativeWorkerSpec
|
|
probe: NativeWorkerProbe
|
|
device: str
|
|
capacity: ExecutionCapacity
|
|
duration_ms: int = 0
|
|
|
|
def __post_init__(self) -> None:
|
|
if not self.endpoint:
|
|
raise NativeRegistrationError("native registration requires an endpoint")
|
|
if not self.model_id:
|
|
raise NativeRegistrationError("native registration requires a model id")
|
|
if not self.device:
|
|
raise NativeRegistrationError("native registration requires a device label")
|
|
if self.identity.artifact.artifact_id != self.model_id:
|
|
raise NativeRegistrationError("native registration model does not match its identity")
|
|
if self.identity.fingerprint.model_artifact_digest != self.worker.artifact_digest:
|
|
raise NativeRegistrationError("native worker artifact digest does not match its identity")
|
|
if self.identity.fingerprint.runtime_recipe_digest != self.worker.recipe_digest:
|
|
raise NativeRegistrationError("native worker recipe digest does not match its identity")
|
|
expected = (
|
|
self.worker.artifact_digest,
|
|
self.worker.recipe_digest,
|
|
self.worker.recipe_id,
|
|
self.worker.recipe_version,
|
|
self.worker.catalogue_version,
|
|
self.worker.shard_start,
|
|
self.worker.shard_end,
|
|
)
|
|
actual = (
|
|
self.probe.artifact_digest,
|
|
self.probe.recipe_digest,
|
|
self.probe.recipe_id,
|
|
self.probe.recipe_version,
|
|
self.probe.catalogue_version,
|
|
self.probe.shard_start,
|
|
self.probe.shard_end,
|
|
)
|
|
if actual != expected:
|
|
raise NativeRegistrationError("native worker probe differs from its startup identity/range")
|
|
if not self.probe.serving:
|
|
raise NativeRegistrationError("native worker is not serving; it cannot register a capability")
|
|
if (
|
|
self.identity.shard_start,
|
|
self.identity.shard_end,
|
|
self.identity.recipe.recipe_id,
|
|
self.identity.recipe.recipe_version,
|
|
self.identity.recipe.catalogue_version,
|
|
) != (
|
|
self.worker.shard_start,
|
|
self.worker.shard_end,
|
|
self.worker.recipe_id,
|
|
self.worker.recipe_version,
|
|
self.worker.catalogue_version,
|
|
):
|
|
raise NativeRegistrationError("native identity differs from worker range or recipe labels")
|
|
if self.identity.recipe.axes["backend_id"] == "":
|
|
raise NativeRegistrationError("native identity must name its backend")
|
|
|
|
def payload(self) -> dict[str, Any]:
|
|
"""Return the existing tracker registration shape with no native branch."""
|
|
report = build_capability_report(
|
|
model_id=self.model_id,
|
|
shard_start=self.identity.shard_start,
|
|
shard_end=self.identity.shard_end - 1,
|
|
recipe_id=self.identity.recipe.recipe_id,
|
|
recipe_version=self.identity.recipe.recipe_version,
|
|
catalogue_version=self.identity.recipe.catalogue_version,
|
|
backend_id=self.identity.recipe.axes["backend_id"],
|
|
device=self.device,
|
|
quantization=self.identity.recipe.axes["weight_quantization"],
|
|
model_config="sha256:" + self.identity.artifact.architecture_digest,
|
|
revision=self.identity.artifact.revision,
|
|
status="passed",
|
|
duration_ms=self.duration_ms,
|
|
identity=self.identity,
|
|
capacity=self.capacity,
|
|
)
|
|
return {
|
|
"endpoint": self.endpoint,
|
|
"model": self.model_id.rsplit("/", 1)[-1],
|
|
"hf_repo": self.model_id,
|
|
"shard_start": self.identity.shard_start,
|
|
"shard_end": self.identity.shard_end - 1,
|
|
"recipe_id": self.identity.recipe.recipe_id,
|
|
"recipe_version": self.identity.recipe.recipe_version,
|
|
"capability_report": report.to_dict(),
|
|
# Existing tracker capacity fields are retained for placement views.
|
|
"ram_bytes": self.capacity.memory_capacity_bytes or 0,
|
|
"max_loaded_shards": 1,
|
|
}
|
|
|
|
|
|
RegistrationSender = Callable[[dict[str, Any]], None]
|
|
WithdrawalSender = Callable[[str], None]
|
|
|
|
|
|
class NativeCapabilityRegistrar:
|
|
"""Publish/withdraw a native capability through caller-owned transport.
|
|
|
|
The callbacks keep tracker HTTP, relay, billing, and provider mechanics out
|
|
of the native worker. A process supervisor calls ``withdraw`` on health
|
|
loss; the caller supplies the existing tracker registration/withdrawal
|
|
transport appropriate to its deployment.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
registration: NativeShardRegistration,
|
|
*,
|
|
register: RegistrationSender,
|
|
withdraw: WithdrawalSender,
|
|
) -> None:
|
|
self.registration = registration
|
|
self._register = register
|
|
self._withdraw = withdraw
|
|
|
|
def publish(self) -> None:
|
|
self._register(self.registration.payload())
|
|
|
|
def unavailable(self, reason: str) -> None:
|
|
self._withdraw(reason)
|
|
|
|
def bind(self, supervisor: NativeWorkerSupervisor) -> None:
|
|
"""Publish only after DGR-040 verification; withdraw on health loss."""
|
|
if supervisor.spec != self.registration.worker:
|
|
raise NativeRegistrationError("registrar and supervisor must own the same native worker")
|
|
supervisor.add_availability_callbacks(
|
|
on_available=lambda _reason: self.publish(),
|
|
on_unavailable=self.unavailable,
|
|
)
|