feat: checkpoint distributed gguf runtime stories
This commit is contained in:
Dobromir Popov
375
packages/node/meshnet_node/runtime_recipe.py
Normal file
375
packages/node/meshnet_node/runtime_recipe.py
Normal file
@@ -0,0 +1,375 @@
|
||||
"""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))
|
||||
Reference in New Issue
Block a user