376 lines
14 KiB
Python
376 lines
14 KiB
Python
"""Exact artifact and runtime-recipe identity helpers.
|
|
|
|
The runtime recipe is the compatibility contract for one routable shard. It is
|
|
kept separate from the user-facing recipe catalogue so the tracker can compare
|
|
the exact execution footprint that was validated, not just a named recipe.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
from dataclasses import dataclass
|
|
from typing import Any, Mapping
|
|
|
|
|
|
def _require_text(value: Any, field_name: str) -> str:
|
|
if not isinstance(value, str) or not value.strip():
|
|
raise ValueError(f"{field_name!r} must be a non-empty string")
|
|
return value
|
|
|
|
|
|
def _optional_text(value: Any, field_name: str) -> str | None:
|
|
if value is None:
|
|
return None
|
|
return _require_text(value, field_name)
|
|
|
|
|
|
def _sha256_text(text: str) -> str:
|
|
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def _stable_json(data: Any) -> str:
|
|
return json.dumps(
|
|
data,
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
ensure_ascii=False,
|
|
default=str,
|
|
)
|
|
|
|
|
|
def _normalise_dtype(value: Any, default: str) -> str:
|
|
if value is None:
|
|
return default
|
|
if isinstance(value, str):
|
|
text = value.strip()
|
|
if not text:
|
|
return default
|
|
return text.removeprefix("torch.")
|
|
return str(value).removeprefix("torch.")
|
|
|
|
|
|
def _architecture_adapter_from_config(model_config: Any, default: str) -> str:
|
|
if not isinstance(model_config, Mapping):
|
|
return default
|
|
for key in ("architecture_adapter", "model_type"):
|
|
value = model_config.get(key)
|
|
if isinstance(value, str) and value.strip():
|
|
return value
|
|
architectures = model_config.get("architectures")
|
|
if isinstance(architectures, list) and architectures:
|
|
first = architectures[0]
|
|
if isinstance(first, str) and first.strip():
|
|
return first
|
|
text_config = model_config.get("text_config")
|
|
if isinstance(text_config, Mapping):
|
|
return _architecture_adapter_from_config(text_config, default)
|
|
return default
|
|
|
|
|
|
def _tokenizer_revision_from_config(
|
|
model_id: str,
|
|
revision: str | None,
|
|
model_config: Any,
|
|
) -> str:
|
|
if isinstance(model_config, Mapping):
|
|
for key in ("tokenizer_revision", "tokenizer_version", "_commit_hash"):
|
|
value = model_config.get(key)
|
|
if isinstance(value, str) and value.strip():
|
|
return value
|
|
if revision:
|
|
return revision
|
|
return model_id
|
|
|
|
|
|
def _cache_layout_from_recipe_params(recipe_params: Mapping[str, Any] | None) -> str:
|
|
if not recipe_params:
|
|
return "local-hot-kv"
|
|
use_cache = recipe_params.get("use_cache")
|
|
if use_cache is False:
|
|
return "stateless"
|
|
if "cache_layout" in recipe_params:
|
|
value = recipe_params.get("cache_layout")
|
|
if isinstance(value, str) and value.strip():
|
|
return value
|
|
return "local-hot-kv"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ArtifactIdentity:
|
|
"""Exact source artifact binding for a routable shard."""
|
|
|
|
model_id: str
|
|
revision: str | None = None
|
|
artifact_hash: str | None = None
|
|
shard_start: int | None = None
|
|
shard_end: int | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
_require_text(self.model_id, "artifact.model_id")
|
|
_optional_text(self.revision, "artifact.revision")
|
|
_optional_text(self.artifact_hash, "artifact.artifact_hash")
|
|
if self.shard_start is not None and self.shard_start < 0:
|
|
raise ValueError("'artifact.shard_start' must be >= 0")
|
|
if self.shard_end is not None and self.shard_end < 0:
|
|
raise ValueError("'artifact.shard_end' must be >= 0")
|
|
if (
|
|
self.shard_start is not None
|
|
and self.shard_end is not None
|
|
and self.shard_end < self.shard_start
|
|
):
|
|
raise ValueError("'artifact.shard_end' must be >= 'artifact.shard_start'")
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"model_id": self.model_id,
|
|
"revision": self.revision,
|
|
"artifact_hash": self.artifact_hash,
|
|
"shard_start": self.shard_start,
|
|
"shard_end": self.shard_end,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: Any) -> "ArtifactIdentity":
|
|
if not isinstance(data, Mapping):
|
|
raise ValueError(f"'artifact' must be a JSON object, got {type(data).__name__}")
|
|
return cls(
|
|
model_id=_require_text(data.get("model_id"), "artifact.model_id"),
|
|
revision=_optional_text(data.get("revision"), "artifact.revision"),
|
|
artifact_hash=_optional_text(
|
|
data.get("artifact_hash"), "artifact.artifact_hash"
|
|
),
|
|
shard_start=_optional_int(data.get("shard_start"), "artifact.shard_start"),
|
|
shard_end=_optional_int(data.get("shard_end"), "artifact.shard_end"),
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RuntimeRecipeIdentity:
|
|
"""Exact runtime recipe used for admission and handshake compatibility."""
|
|
|
|
weight_quantization: str
|
|
activation_dtype: str
|
|
compute_dtype: str
|
|
kv_dtype: str
|
|
kv_layout: str
|
|
tokenizer_revision: str
|
|
architecture_adapter: str
|
|
backend_id: str
|
|
runtime_version: str
|
|
boundary_schema_version: int = 1
|
|
cache_layout: str = "local-hot-kv"
|
|
fingerprint: str | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
_require_text(self.weight_quantization, "runtime_recipe.weight_quantization")
|
|
_require_text(self.activation_dtype, "runtime_recipe.activation_dtype")
|
|
_require_text(self.compute_dtype, "runtime_recipe.compute_dtype")
|
|
_require_text(self.kv_dtype, "runtime_recipe.kv_dtype")
|
|
_require_text(self.kv_layout, "runtime_recipe.kv_layout")
|
|
_require_text(self.tokenizer_revision, "runtime_recipe.tokenizer_revision")
|
|
_require_text(self.architecture_adapter, "runtime_recipe.architecture_adapter")
|
|
_require_text(self.backend_id, "runtime_recipe.backend_id")
|
|
_require_text(self.runtime_version, "runtime_recipe.runtime_version")
|
|
_require_text(self.cache_layout, "runtime_recipe.cache_layout")
|
|
if self.boundary_schema_version < 1:
|
|
raise ValueError("'runtime_recipe.boundary_schema_version' must be >= 1")
|
|
expected = compatibility_fingerprint(self._fingerprint_payload())
|
|
if not self.fingerprint:
|
|
object.__setattr__(self, "fingerprint", expected)
|
|
elif self.fingerprint != expected:
|
|
raise ValueError(
|
|
"'runtime_recipe.fingerprint' does not match the encoded fields"
|
|
)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"weight_quantization": self.weight_quantization,
|
|
"activation_dtype": self.activation_dtype,
|
|
"compute_dtype": self.compute_dtype,
|
|
"kv_dtype": self.kv_dtype,
|
|
"kv_layout": self.kv_layout,
|
|
"tokenizer_revision": self.tokenizer_revision,
|
|
"architecture_adapter": self.architecture_adapter,
|
|
"backend_id": self.backend_id,
|
|
"runtime_version": self.runtime_version,
|
|
"boundary_schema_version": self.boundary_schema_version,
|
|
"cache_layout": self.cache_layout,
|
|
"fingerprint": self.fingerprint,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: Any) -> "RuntimeRecipeIdentity":
|
|
if not isinstance(data, Mapping):
|
|
raise ValueError(
|
|
f"'runtime_recipe' must be a JSON object, got {type(data).__name__}"
|
|
)
|
|
boundary_schema_version = data.get("boundary_schema_version", 1)
|
|
if isinstance(boundary_schema_version, bool) or not isinstance(
|
|
boundary_schema_version, int
|
|
):
|
|
raise ValueError(
|
|
"'runtime_recipe.boundary_schema_version' must be an integer"
|
|
)
|
|
return cls(
|
|
weight_quantization=_require_text(
|
|
data.get("weight_quantization"), "runtime_recipe.weight_quantization"
|
|
),
|
|
activation_dtype=_require_text(
|
|
data.get("activation_dtype"), "runtime_recipe.activation_dtype"
|
|
),
|
|
compute_dtype=_require_text(
|
|
data.get("compute_dtype"), "runtime_recipe.compute_dtype"
|
|
),
|
|
kv_dtype=_require_text(data.get("kv_dtype"), "runtime_recipe.kv_dtype"),
|
|
kv_layout=_require_text(data.get("kv_layout"), "runtime_recipe.kv_layout"),
|
|
tokenizer_revision=_require_text(
|
|
data.get("tokenizer_revision"), "runtime_recipe.tokenizer_revision"
|
|
),
|
|
architecture_adapter=_require_text(
|
|
data.get("architecture_adapter"),
|
|
"runtime_recipe.architecture_adapter",
|
|
),
|
|
backend_id=_require_text(data.get("backend_id"), "runtime_recipe.backend_id"),
|
|
runtime_version=_require_text(
|
|
data.get("runtime_version"), "runtime_recipe.runtime_version"
|
|
),
|
|
boundary_schema_version=boundary_schema_version,
|
|
cache_layout=_require_text(data.get("cache_layout"), "runtime_recipe.cache_layout"),
|
|
fingerprint=_optional_text(data.get("fingerprint"), "runtime_recipe.fingerprint"),
|
|
)
|
|
|
|
def _fingerprint_payload(self) -> dict[str, Any]:
|
|
return {
|
|
"weight_quantization": self.weight_quantization,
|
|
"activation_dtype": self.activation_dtype,
|
|
"compute_dtype": self.compute_dtype,
|
|
"kv_dtype": self.kv_dtype,
|
|
"kv_layout": self.kv_layout,
|
|
"tokenizer_revision": self.tokenizer_revision,
|
|
"architecture_adapter": self.architecture_adapter,
|
|
"backend_id": self.backend_id,
|
|
"runtime_version": self.runtime_version,
|
|
"boundary_schema_version": self.boundary_schema_version,
|
|
"cache_layout": self.cache_layout,
|
|
}
|
|
|
|
|
|
def _optional_int(value: Any, field_name: str) -> int | None:
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, bool) or not isinstance(value, int):
|
|
raise ValueError(f"{field_name!r} must be an integer")
|
|
if value < 0:
|
|
raise ValueError(f"{field_name!r} must be >= 0")
|
|
return value
|
|
|
|
|
|
def build_artifact_identity(
|
|
*,
|
|
model_id: str,
|
|
revision: str | None = None,
|
|
model_config: Any = None,
|
|
artifact_hash: str | None = None,
|
|
shard_start: int | None = None,
|
|
shard_end: int | None = None,
|
|
) -> ArtifactIdentity:
|
|
"""Build a stable artifact binding from the locally loaded artifact."""
|
|
resolved_hash = artifact_hash
|
|
if resolved_hash is None:
|
|
if isinstance(model_config, Mapping):
|
|
resolved_hash = _hash_mapping(model_config)
|
|
elif model_config is not None:
|
|
resolved_hash = _sha256_text(_stable_json(model_config))
|
|
if resolved_hash is None:
|
|
resolved_hash = _sha256_text(
|
|
_stable_json(
|
|
{
|
|
"model_id": model_id,
|
|
"revision": revision,
|
|
"shard_start": shard_start,
|
|
"shard_end": shard_end,
|
|
}
|
|
)
|
|
)
|
|
return ArtifactIdentity(
|
|
model_id=model_id,
|
|
revision=revision,
|
|
artifact_hash=resolved_hash,
|
|
shard_start=shard_start,
|
|
shard_end=shard_end,
|
|
)
|
|
|
|
|
|
def build_runtime_recipe_identity(
|
|
*,
|
|
model_id: str,
|
|
weight_quantization: str,
|
|
backend_id: str,
|
|
runtime_version: str,
|
|
revision: str | None = None,
|
|
model_config: Any = None,
|
|
recipe_params: Mapping[str, Any] | None = None,
|
|
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,
|
|
) -> RuntimeRecipeIdentity:
|
|
"""Build the exact runtime recipe used for compatibility admission."""
|
|
activation = _normalise_dtype(activation_dtype, "bfloat16")
|
|
compute = _normalise_dtype(compute_dtype, activation)
|
|
kv_dtype_text = _normalise_dtype(kv_dtype, compute)
|
|
kv_layout_text = kv_layout or "session-cache"
|
|
tokenizer = tokenizer_revision or _tokenizer_revision_from_config(
|
|
model_id, revision, model_config
|
|
)
|
|
architecture = architecture_adapter or _architecture_adapter_from_config(
|
|
model_config, backend_id
|
|
)
|
|
cache_layout_text = cache_layout or _cache_layout_from_recipe_params(recipe_params)
|
|
return RuntimeRecipeIdentity(
|
|
weight_quantization=weight_quantization,
|
|
activation_dtype=activation,
|
|
compute_dtype=compute,
|
|
kv_dtype=kv_dtype_text,
|
|
kv_layout=kv_layout_text,
|
|
tokenizer_revision=tokenizer,
|
|
architecture_adapter=architecture,
|
|
backend_id=backend_id,
|
|
runtime_version=runtime_version,
|
|
boundary_schema_version=boundary_schema_version,
|
|
cache_layout=cache_layout_text,
|
|
)
|
|
|
|
|
|
def compatibility_fingerprint(data: Mapping[str, Any]) -> str:
|
|
"""Return a stable SHA256 compatibility fingerprint for an exact route."""
|
|
return "sha256:" + _sha256_text(_stable_json(data))
|
|
|
|
|
|
def fingerprint_payload(
|
|
*,
|
|
model: Mapping[str, Any],
|
|
shard: Mapping[str, Any],
|
|
recipe: Mapping[str, Any],
|
|
backend: Mapping[str, Any],
|
|
artifact: Mapping[str, Any],
|
|
runtime_recipe: Mapping[str, Any],
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"model": dict(model),
|
|
"shard": dict(shard),
|
|
"recipe": dict(recipe),
|
|
"backend": dict(backend),
|
|
"artifact": dict(artifact),
|
|
"runtime_recipe": dict(runtime_recipe),
|
|
}
|
|
|
|
|
|
def _hash_mapping(data: Mapping[str, Any]) -> str:
|
|
return "sha256:" + _sha256_text(_stable_json(data))
|