feat: checkpoint distributed gguf runtime stories

This commit is contained in:
Dobromir Popov
2026-07-15 23:42:58 +03:00
parent eaf00f6add
commit 1fe31ef38d
60 changed files with 8478 additions and 105 deletions

View File

@@ -20,6 +20,16 @@ import time
from dataclasses import dataclass, field
from typing import Any, Mapping
from . import __version__ as _PACKAGE_VERSION
from .runtime_recipe import (
ArtifactIdentity,
RuntimeRecipeIdentity,
build_artifact_identity,
build_runtime_recipe_identity,
compatibility_fingerprint,
fingerprint_payload,
)
# Layout of the serialized report. Bump when the JSON shape changes.
CAPABILITY_SCHEMA_VERSION = 1
@@ -172,6 +182,14 @@ def _optional_text(value: Any, field_name: str) -> str | None:
return _require_text(value, field_name)
def _optional_bool(value: Any, field_name: str) -> bool:
if value is None:
return False
if isinstance(value, bool):
return value
raise CapabilityReportError(f"{field_name!r} must be a boolean")
def _require_int(value: Any, field_name: str, minimum: int) -> int:
if isinstance(value, bool) or not isinstance(value, int):
raise CapabilityReportError(f"{field_name!r} must be an integer")
@@ -218,6 +236,8 @@ class ShardRange:
start: int
end: int
owns_embedding: bool = False
owns_final_head: bool = False
def __post_init__(self) -> None:
_require_int(self.start, "shard.start", 0)
@@ -226,9 +246,18 @@ class ShardRange:
raise CapabilityReportError(
f"'shard.end' ({self.end}) must be >= 'shard.start' ({self.start})"
)
if not isinstance(self.owns_embedding, bool):
raise CapabilityReportError("'shard.owns_embedding' must be a boolean")
if not isinstance(self.owns_final_head, bool):
raise CapabilityReportError("'shard.owns_final_head' must be a boolean")
def to_dict(self) -> dict:
return {"start": self.start, "end": self.end}
return {
"start": self.start,
"end": self.end,
"owns_embedding": self.owns_embedding,
"owns_final_head": self.owns_final_head,
}
@classmethod
def from_dict(cls, data: Any) -> ShardRange:
@@ -236,6 +265,12 @@ class ShardRange:
return cls(
start=_require_int(doc.get("start"), "shard.start", 0),
end=_require_int(doc.get("end"), "shard.end", 0),
owns_embedding=_optional_bool(
doc.get("owns_embedding"), "shard.owns_embedding"
),
owns_final_head=_optional_bool(
doc.get("owns_final_head"), "shard.owns_final_head"
),
)
@@ -336,6 +371,8 @@ class CapabilityReport:
shard: ShardRange
recipe: RecipeIdentity
backend: BackendIdentity
artifact: ArtifactIdentity
runtime_recipe: RuntimeRecipeIdentity
status: str
validated_at: float
duration_ms: int
@@ -376,6 +413,20 @@ class CapabilityReport:
self.backend.device,
)
@property
def compatibility_fingerprint(self) -> str:
"""Stable compatibility digest over the full routable identity."""
return compatibility_fingerprint(
fingerprint_payload(
model=self.model.to_dict(),
shard=self.shard.to_dict(),
recipe=self.recipe.to_dict(),
backend=self.backend.to_dict(),
artifact=self.artifact.to_dict(),
runtime_recipe=self.runtime_recipe.to_dict(),
)
)
def age_seconds(self, now: float | None = None) -> float:
return max(0.0, (time.time() if now is None else now) - self.validated_at)
@@ -386,6 +437,9 @@ class CapabilityReport:
"shard": self.shard.to_dict(),
"recipe": self.recipe.to_dict(),
"backend": self.backend.to_dict(),
"artifact": self.artifact.to_dict(),
"runtime_recipe": self.runtime_recipe.to_dict(),
"compatibility_fingerprint": self.compatibility_fingerprint,
"status": self.status,
"validated_at": self.validated_at,
"duration_ms": self.duration_ms,
@@ -398,6 +452,9 @@ class CapabilityReport:
@classmethod
def from_dict(cls, data: Any) -> CapabilityReport:
doc = _as_mapping(data, "report")
declared_compatibility_fingerprint = _optional_text(
doc.get("compatibility_fingerprint"), "compatibility_fingerprint"
)
if "schema_version" not in doc:
raise CapabilityReportError(
@@ -417,7 +474,13 @@ class CapabilityReport:
):
raise CapabilityReportError("'validated_at' must be a Unix timestamp")
return cls(
try:
artifact = ArtifactIdentity.from_dict(doc.get("artifact"))
runtime_recipe = RuntimeRecipeIdentity.from_dict(doc.get("runtime_recipe"))
except ValueError as exc:
raise CapabilityReportError(str(exc)) from exc
report = cls(
schema_version=schema_version,
model=ModelIdentity.from_dict(doc.get("model")),
shard=ShardRange.from_dict(doc.get("shard")),
@@ -427,7 +490,18 @@ class CapabilityReport:
validated_at=float(validated_at),
duration_ms=_require_int(doc.get("duration_ms"), "duration_ms", 0),
diagnostics=sanitize_diagnostics(doc.get("diagnostics")),
artifact=artifact,
runtime_recipe=runtime_recipe,
)
if (
declared_compatibility_fingerprint is not None
and report.compatibility_fingerprint != declared_compatibility_fingerprint
):
raise CapabilityReportError(
"report declares a compatibility fingerprint that does not match "
"its artifact/runtime recipe"
)
return report
@classmethod
def from_json(cls, text: str) -> CapabilityReport:
@@ -458,6 +532,19 @@ def build_capability_report(
device_name: str | None = None,
quantization: str | None = None,
runtime: Mapping[str, str] | None = None,
artifact_hash: str | None = None,
runtime_recipe: RuntimeRecipeIdentity | None = None,
owns_embedding: bool = False,
owns_final_head: bool = False,
activation_dtype: Any = None,
compute_dtype: Any = None,
kv_dtype: Any = None,
kv_layout: str | None = None,
tokenizer_revision: str | None = None,
architecture_adapter: str | None = None,
boundary_schema_version: int = 1,
cache_layout: str | None = None,
recipe_params: Mapping[str, Any] | None = None,
diagnostics: Any = None,
validated_at: float | None = None,
environ: Mapping[str, str] | None = None,
@@ -468,25 +555,62 @@ def build_capability_report(
or an already-computed ``sha256:…`` string. `validated_at` defaults to now,
so callers that need determinism pass it explicitly.
"""
return CapabilityReport(
model=ModelIdentity(
model_identity = ModelIdentity(
model_id=model_id,
revision=revision,
config_fingerprint=config_fingerprint(model_config),
)
shard = ShardRange(
start=shard_start,
end=shard_end,
owns_embedding=owns_embedding,
owns_final_head=owns_final_head,
)
recipe_identity = RecipeIdentity(
recipe_id=recipe_id,
recipe_version=recipe_version,
catalogue_version=catalogue_version,
)
backend_identity = BackendIdentity(
backend_id=backend_id,
device=device,
device_name=device_name,
quantization=quantization,
runtime=dict(runtime or {}),
)
artifact = build_artifact_identity(
model_id=model_id,
revision=revision,
model_config=model_config,
artifact_hash=artifact_hash,
shard_start=shard_start,
shard_end=shard_end,
)
if runtime_recipe is None:
runtime_recipe = build_runtime_recipe_identity(
model_id=model_id,
revision=revision,
config_fingerprint=config_fingerprint(model_config),
),
shard=ShardRange(start=shard_start, end=shard_end),
recipe=RecipeIdentity(
recipe_id=recipe_id,
recipe_version=recipe_version,
catalogue_version=catalogue_version,
),
backend=BackendIdentity(
model_config=model_config,
recipe_params=recipe_params,
weight_quantization=quantization or "unknown",
backend_id=backend_id,
device=device,
device_name=device_name,
quantization=quantization,
runtime=dict(runtime or {}),
),
runtime_version=_PACKAGE_VERSION,
activation_dtype=activation_dtype,
compute_dtype=compute_dtype,
kv_dtype=kv_dtype,
kv_layout=kv_layout,
tokenizer_revision=tokenizer_revision,
architecture_adapter=architecture_adapter,
boundary_schema_version=boundary_schema_version,
cache_layout=cache_layout,
)
return CapabilityReport(
model=model_identity,
shard=shard,
recipe=recipe_identity,
backend=backend_identity,
artifact=artifact,
runtime_recipe=runtime_recipe,
status=status,
validated_at=time.time() if validated_at is None else validated_at,
duration_ms=duration_ms,