354 lines
14 KiB
Python
354 lines
14 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Versioned immutable manifests for runtime-ready host weights."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from pathlib import PurePosixPath
|
|
from typing import cast
|
|
|
|
import regex as re
|
|
|
|
from .identity import CanonicalJson, JsonValue, WeightArtifactIdentity, canonical_json
|
|
|
|
_SHA256_RE = re.compile(r"[0-9a-f]{64}")
|
|
|
|
|
|
class ManifestValidationError(ValueError):
|
|
"""Raised when a published manifest violates its schema."""
|
|
|
|
|
|
class TensorKind(str, Enum):
|
|
PARAMETER = "parameter"
|
|
BUFFER = "buffer"
|
|
|
|
|
|
def _require_text(name: str, value: object) -> str:
|
|
if not isinstance(value, str) or not value:
|
|
raise ManifestValidationError(f"{name} must be a non-empty string")
|
|
return value
|
|
|
|
|
|
def _require_int(name: str, value: object) -> int:
|
|
if isinstance(value, bool) or not isinstance(value, int):
|
|
raise ManifestValidationError(f"{name} must be an integer")
|
|
return value
|
|
|
|
|
|
def _require_sha256(name: str, value: object) -> str:
|
|
if not isinstance(value, str) or _SHA256_RE.fullmatch(value) is None:
|
|
raise ManifestValidationError(f"{name} must be a lowercase SHA-256 digest")
|
|
return value
|
|
|
|
|
|
def _require_safe_file_name(value: object) -> str:
|
|
if not isinstance(value, str):
|
|
raise ManifestValidationError("artifact file name must be a string")
|
|
path = PurePosixPath(value)
|
|
if not value or path.is_absolute() or len(path.parts) != 1 or path.name != value or value in {".", ".."}:
|
|
raise ManifestValidationError(f"artifact file name must be one safe relative basename: {value!r}")
|
|
return value
|
|
|
|
|
|
def _require_int_tuple(name: str, value: object) -> tuple[int, ...]:
|
|
if not isinstance(value, list):
|
|
raise ManifestValidationError(f"{name} must be an array")
|
|
return tuple(_require_int(f"{name}[{index}]", item) for index, item in enumerate(value))
|
|
|
|
|
|
def _exact_keys(value: dict[str, object], expected: set[str], kind: str) -> None:
|
|
keys = set(value)
|
|
if keys != expected:
|
|
raise ManifestValidationError(
|
|
f"{kind} fields differ: missing={sorted(expected - keys)}, unexpected={sorted(keys - expected)}"
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TensorManifestEntry:
|
|
name: str
|
|
kind: TensorKind
|
|
role: str
|
|
file_name: str
|
|
storage_key: str
|
|
shape: tuple[int, ...]
|
|
dtype: str
|
|
stride: tuple[int, ...]
|
|
nbytes: int
|
|
sha256: str
|
|
|
|
def __post_init__(self) -> None:
|
|
_require_text("tensor.name", self.name)
|
|
_require_text("tensor.storage_key", self.storage_key)
|
|
_require_text("tensor.role", self.role)
|
|
_require_text("tensor.dtype", self.dtype)
|
|
if not isinstance(self.kind, TensorKind):
|
|
raise ManifestValidationError("tensor kind must use TensorKind")
|
|
_require_safe_file_name(self.file_name)
|
|
if not isinstance(self.shape, tuple) or any(
|
|
isinstance(dimension, bool) or not isinstance(dimension, int) or dimension < 0 for dimension in self.shape
|
|
):
|
|
raise ManifestValidationError(f"tensor {self.name!r} has a negative shape dimension")
|
|
if (
|
|
not isinstance(self.stride, tuple)
|
|
or len(self.stride) != len(self.shape)
|
|
or any(isinstance(value, bool) or not isinstance(value, int) or value < 0 for value in self.stride)
|
|
):
|
|
raise ManifestValidationError(f"tensor {self.name!r} has an invalid stride")
|
|
if _require_int(f"tensors[{self.name!r}].nbytes", self.nbytes) < 0:
|
|
raise ManifestValidationError(f"tensor {self.name!r} has negative nbytes")
|
|
_require_sha256(f"tensors[{self.name!r}].sha256", self.sha256)
|
|
|
|
def semantic_dict(self) -> dict[str, JsonValue]:
|
|
"""Return content identity independent of physical file sharding."""
|
|
return {
|
|
"name": self.name,
|
|
"kind": self.kind.value,
|
|
"role": self.role,
|
|
"shape": list(self.shape),
|
|
"dtype": self.dtype,
|
|
"stride": list(self.stride),
|
|
"nbytes": self.nbytes,
|
|
"sha256": self.sha256,
|
|
}
|
|
|
|
def to_dict(self) -> dict[str, JsonValue]:
|
|
return {
|
|
"name": self.name,
|
|
"kind": self.kind.value,
|
|
"role": self.role,
|
|
"file_name": self.file_name,
|
|
"storage_key": self.storage_key,
|
|
"shape": list(self.shape),
|
|
"dtype": self.dtype,
|
|
"stride": list(self.stride),
|
|
"nbytes": self.nbytes,
|
|
"sha256": self.sha256,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, value: dict[str, object]) -> TensorManifestEntry:
|
|
expected = {
|
|
"name",
|
|
"kind",
|
|
"role",
|
|
"file_name",
|
|
"storage_key",
|
|
"shape",
|
|
"dtype",
|
|
"stride",
|
|
"nbytes",
|
|
"sha256",
|
|
}
|
|
_exact_keys(value, expected, "tensor manifest entry")
|
|
kind_value = _require_text("tensor.kind", value["kind"])
|
|
try:
|
|
kind = TensorKind(kind_value)
|
|
except ValueError as exc:
|
|
raise ManifestValidationError(f"unknown tensor kind {value['kind']!r}") from exc
|
|
return cls(
|
|
name=_require_text("tensor.name", value["name"]),
|
|
kind=kind,
|
|
role=_require_text("tensor.role", value["role"]),
|
|
file_name=_require_safe_file_name(value["file_name"]),
|
|
storage_key=_require_text("tensor.storage_key", value["storage_key"]),
|
|
shape=_require_int_tuple("tensor.shape", value["shape"]),
|
|
dtype=_require_text("tensor.dtype", value["dtype"]),
|
|
stride=_require_int_tuple("tensor.stride", value["stride"]),
|
|
nbytes=_require_int("tensor.nbytes", value["nbytes"]),
|
|
sha256=_require_sha256("tensor.sha256", value["sha256"]),
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FileManifestEntry:
|
|
file_name: str
|
|
size: int
|
|
sha256: str
|
|
|
|
def __post_init__(self) -> None:
|
|
_require_safe_file_name(self.file_name)
|
|
if _require_int(f"files[{self.file_name!r}].size", self.size) < 0:
|
|
raise ManifestValidationError(f"artifact file {self.file_name!r} has negative size")
|
|
_require_sha256(f"files[{self.file_name!r}].sha256", self.sha256)
|
|
|
|
def to_dict(self) -> dict[str, JsonValue]:
|
|
return {"file_name": self.file_name, "size": self.size, "sha256": self.sha256}
|
|
|
|
@classmethod
|
|
def from_dict(cls, value: dict[str, object]) -> FileManifestEntry:
|
|
_exact_keys(value, {"file_name", "size", "sha256"}, "file manifest entry")
|
|
return cls(
|
|
file_name=_require_safe_file_name(value["file_name"]),
|
|
size=_require_int("file.size", value["size"]),
|
|
sha256=_require_sha256("file.sha256", value["sha256"]),
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class WeightManifest:
|
|
schema_version: int
|
|
artifact_key: str
|
|
identity: WeightArtifactIdentity
|
|
producer_schema: str
|
|
restorer_schema: str
|
|
tensors: tuple[TensorManifestEntry, ...]
|
|
files: tuple[FileManifestEntry, ...]
|
|
artifact_content_sha256: str
|
|
format_metadata: CanonicalJson = field(default_factory=CanonicalJson.empty)
|
|
|
|
def __post_init__(self) -> None:
|
|
if _require_int("manifest.schema_version", self.schema_version) < 1:
|
|
raise ManifestValidationError("manifest schema_version must be positive")
|
|
_require_sha256("artifact_key", self.artifact_key)
|
|
_require_sha256("artifact_content_sha256", self.artifact_content_sha256)
|
|
if not isinstance(self.identity, WeightArtifactIdentity):
|
|
raise ManifestValidationError("manifest identity must use WeightArtifactIdentity")
|
|
if self.artifact_key != self.identity.key:
|
|
raise ManifestValidationError("manifest artifact_key does not match its semantic identity")
|
|
_require_text("manifest.producer_schema", self.producer_schema)
|
|
_require_text("manifest.restorer_schema", self.restorer_schema)
|
|
if (
|
|
not isinstance(self.tensors, tuple)
|
|
or not isinstance(self.files, tuple)
|
|
or not self.tensors
|
|
or not self.files
|
|
or any(not isinstance(entry, TensorManifestEntry) for entry in self.tensors)
|
|
or any(not isinstance(entry, FileManifestEntry) for entry in self.files)
|
|
):
|
|
raise ManifestValidationError("an artifact must contain at least one tensor and one file")
|
|
if not isinstance(self.format_metadata, CanonicalJson):
|
|
raise ManifestValidationError("manifest format_metadata must use CanonicalJson")
|
|
tensor_names = [entry.name for entry in self.tensors]
|
|
if tensor_names != sorted(tensor_names) or len(set(tensor_names)) != len(tensor_names):
|
|
raise ManifestValidationError("manifest tensors must have unique names in canonical order")
|
|
file_names = [entry.file_name for entry in self.files]
|
|
if file_names != sorted(file_names) or len(set(file_names)) != len(file_names):
|
|
raise ManifestValidationError("manifest files must have unique names in canonical order")
|
|
available_files = set(file_names)
|
|
if any(entry.file_name not in available_files for entry in self.tensors):
|
|
raise ManifestValidationError("tensor manifest references an undeclared artifact file")
|
|
storage_bindings = [(entry.file_name, entry.storage_key) for entry in self.tensors]
|
|
if len(set(storage_bindings)) != len(storage_bindings):
|
|
raise ManifestValidationError("manifest tensor storage bindings must be unique")
|
|
|
|
@property
|
|
def canonical_bytes(self) -> bytes:
|
|
return canonical_json(self.to_dict())
|
|
|
|
@property
|
|
def total_tensor_bytes(self) -> int:
|
|
return sum(entry.nbytes for entry in self.tensors)
|
|
|
|
def to_dict(self) -> dict[str, JsonValue]:
|
|
return {
|
|
"schema_version": self.schema_version,
|
|
"artifact_key": self.artifact_key,
|
|
"identity": self.identity.to_dict(),
|
|
"producer_schema": self.producer_schema,
|
|
"restorer_schema": self.restorer_schema,
|
|
"tensors": [entry.to_dict() for entry in self.tensors],
|
|
"files": [entry.to_dict() for entry in self.files],
|
|
"artifact_content_sha256": self.artifact_content_sha256,
|
|
"format_metadata": self.format_metadata.to_value(),
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, value: dict[str, object]) -> WeightManifest:
|
|
expected = {
|
|
"schema_version",
|
|
"artifact_key",
|
|
"identity",
|
|
"producer_schema",
|
|
"restorer_schema",
|
|
"tensors",
|
|
"files",
|
|
"artifact_content_sha256",
|
|
"format_metadata",
|
|
}
|
|
_exact_keys(value, expected, "weight manifest")
|
|
identity = value["identity"]
|
|
if not isinstance(identity, dict):
|
|
raise ManifestValidationError("manifest identity must be an object")
|
|
tensors = value["tensors"]
|
|
files = value["files"]
|
|
if not isinstance(tensors, list) or not all(isinstance(item, dict) for item in tensors):
|
|
raise ManifestValidationError("manifest tensors must be an array of objects")
|
|
if not isinstance(files, list) or not all(isinstance(item, dict) for item in files):
|
|
raise ManifestValidationError("manifest files must be an array of objects")
|
|
return cls(
|
|
schema_version=_require_int("manifest.schema_version", value["schema_version"]),
|
|
artifact_key=_require_sha256("manifest.artifact_key", value["artifact_key"]),
|
|
identity=WeightArtifactIdentity.from_dict(cast(dict[str, object], identity)),
|
|
producer_schema=_require_text("manifest.producer_schema", value["producer_schema"]),
|
|
restorer_schema=_require_text("manifest.restorer_schema", value["restorer_schema"]),
|
|
tensors=tuple(TensorManifestEntry.from_dict(cast(dict[str, object], item)) for item in tensors),
|
|
files=tuple(FileManifestEntry.from_dict(cast(dict[str, object], item)) for item in files),
|
|
artifact_content_sha256=_require_sha256(
|
|
"manifest.artifact_content_sha256", value["artifact_content_sha256"]
|
|
),
|
|
format_metadata=CanonicalJson.from_value(value["format_metadata"]),
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PublicationMarker:
|
|
schema_version: int
|
|
publication_id: str
|
|
manifest_sha256: str
|
|
artifact_content_sha256: str
|
|
integrity_mode: str
|
|
|
|
def __post_init__(self) -> None:
|
|
if _require_int("publication.schema_version", self.schema_version) < 1:
|
|
raise ManifestValidationError("publication schema_version must be positive")
|
|
_require_text("publication.publication_id", self.publication_id)
|
|
_require_text("publication.integrity_mode", self.integrity_mode)
|
|
_require_sha256("publication.manifest_sha256", self.manifest_sha256)
|
|
_require_sha256("publication.artifact_content_sha256", self.artifact_content_sha256)
|
|
|
|
@property
|
|
def canonical_bytes(self) -> bytes:
|
|
return canonical_json(self.to_dict())
|
|
|
|
def to_dict(self) -> dict[str, JsonValue]:
|
|
return {
|
|
"schema_version": self.schema_version,
|
|
"publication_id": self.publication_id,
|
|
"manifest_sha256": self.manifest_sha256,
|
|
"artifact_content_sha256": self.artifact_content_sha256,
|
|
"integrity_mode": self.integrity_mode,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, value: dict[str, object]) -> PublicationMarker:
|
|
expected = {
|
|
"schema_version",
|
|
"publication_id",
|
|
"manifest_sha256",
|
|
"artifact_content_sha256",
|
|
"integrity_mode",
|
|
}
|
|
_exact_keys(value, expected, "publication marker")
|
|
return cls(
|
|
schema_version=_require_int("publication.schema_version", value["schema_version"]),
|
|
publication_id=_require_text("publication.publication_id", value["publication_id"]),
|
|
manifest_sha256=_require_sha256("publication.manifest_sha256", value["manifest_sha256"]),
|
|
artifact_content_sha256=_require_sha256(
|
|
"publication.artifact_content_sha256", value["artifact_content_sha256"]
|
|
),
|
|
integrity_mode=_require_text("publication.integrity_mode", value["integrity_mode"]),
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"FileManifestEntry",
|
|
"ManifestValidationError",
|
|
"PublicationMarker",
|
|
"TensorKind",
|
|
"TensorManifestEntry",
|
|
"WeightManifest",
|
|
]
|