172 lines
7.1 KiB
Python
172 lines
7.1 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, RoutingMeasurements, 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
|
||
routing: RoutingMeasurements | None = None
|
||
|
||
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,
|
||
routing=self.routing,
|
||
)
|
||
payload = {
|
||
"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,
|
||
}
|
||
# These are the tracker’s established dynamic scoring inputs. The
|
||
# exact same optional report can be sent by any backend; no native
|
||
# route or balancing branch is introduced here.
|
||
if self.routing is not None:
|
||
if self.routing.tokens_per_second is not None:
|
||
payload["benchmark_tokens_per_sec"] = self.routing.tokens_per_second
|
||
if self.routing.queue_depth is not None:
|
||
payload["queue_depth"] = self.routing.queue_depth
|
||
return payload
|
||
|
||
|
||
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,
|
||
)
|