98 lines
3.3 KiB
Python
98 lines
3.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Deterministic identity helpers for the DLO runtime-weight cache."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
import json
|
|
import os
|
|
from collections.abc import Mapping
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import TypeAlias
|
|
|
|
import torch
|
|
|
|
JsonScalar: TypeAlias = None | bool | int | float | str
|
|
JsonValue: TypeAlias = JsonScalar | list["JsonValue"] | dict[str, "JsonValue"]
|
|
PathIdentity: TypeAlias = str | os.PathLike[str] | os.PathLike[bytes]
|
|
|
|
|
|
class IdentityNormalizationError(ValueError):
|
|
"""Raised when a cache identity contains a process-unstable value."""
|
|
|
|
|
|
def _type_identity(value: object) -> str:
|
|
value_type = value if isinstance(value, type) else type(value)
|
|
return f"{value_type.__module__}.{value_type.__qualname__}"
|
|
|
|
|
|
def canonicalize_existing_local_path(value: PathIdentity) -> str:
|
|
"""Resolve equivalent local paths while leaving model repository IDs alone."""
|
|
original = os.fsdecode(os.fspath(value))
|
|
candidate = Path(original).expanduser()
|
|
try:
|
|
if candidate.exists():
|
|
return str(candidate.resolve())
|
|
except OSError:
|
|
pass
|
|
return original
|
|
|
|
|
|
def normalize_identity(value: object) -> JsonValue:
|
|
"""Convert supported loader inputs into deterministic JSON-compatible data."""
|
|
if value is None or isinstance(value, (bool, int, float, str)):
|
|
return value
|
|
if isinstance(value, Enum):
|
|
return {"enum": _type_identity(value), "name": value.name}
|
|
if isinstance(value, torch.dtype):
|
|
return str(value)
|
|
if isinstance(value, os.PathLike):
|
|
return canonicalize_existing_local_path(value)
|
|
if isinstance(value, type):
|
|
return _type_identity(value)
|
|
if dataclasses.is_dataclass(value):
|
|
return normalize_identity(dataclasses.asdict(value))
|
|
|
|
model_dump = getattr(value, "model_dump", None)
|
|
if callable(model_dump):
|
|
try:
|
|
dumped = model_dump(mode="json")
|
|
except (TypeError, ValueError):
|
|
dumped = model_dump()
|
|
return normalize_identity(dumped)
|
|
|
|
if isinstance(value, Mapping):
|
|
normalized: dict[str, JsonValue] = {}
|
|
for key, item in value.items():
|
|
if not isinstance(key, str):
|
|
raise IdentityNormalizationError("runtime-cache identity mappings require string keys")
|
|
normalized[key] = normalize_identity(item)
|
|
return {key: normalized[key] for key in sorted(normalized)}
|
|
if isinstance(value, (set, frozenset)):
|
|
return sorted((normalize_identity(item) for item in value), key=canonical_json)
|
|
if isinstance(value, (list, tuple)):
|
|
return [normalize_identity(item) for item in value]
|
|
|
|
raise IdentityNormalizationError(f"runtime-cache identity does not support values of type {_type_identity(value)}")
|
|
|
|
|
|
def canonical_json(value: object) -> bytes:
|
|
"""Serialize a supported identity with stable ordering and separators."""
|
|
return json.dumps(
|
|
normalize_identity(value),
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
ensure_ascii=True,
|
|
).encode()
|
|
|
|
|
|
__all__ = [
|
|
"IdentityNormalizationError",
|
|
"JsonValue",
|
|
"canonical_json",
|
|
"canonicalize_existing_local_path",
|
|
"normalize_identity",
|
|
]
|