danrh
07f3dc254f
Introduce model-specific target profiles, fail-closed grounding and validation, content-addressed Director caching, LTX-2.5 contracts, comparison workflows, benchmarks, schemas, documentation, and regression coverage.
340 lines
11 KiB
Python
340 lines
11 KiB
Python
"""Small, content-addressed disk cache for completed Director packets.
|
|
|
|
The cache deliberately stores only JSON/text outputs. It never retains model,
|
|
processor, tensor, or CUDA objects, so a hit can be served before the
|
|
DiffusionGemma runtime is loaded.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import tempfile
|
|
from dataclasses import asdict, is_dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Mapping, Sequence
|
|
|
|
|
|
DIRECTOR_CACHE_SCHEMA = "dg-director-cache/2"
|
|
DIRECTOR_CACHE_KEY_SCHEMA = "dg-director-cache-key/2"
|
|
_SMALL_FILE_HASH_LIMIT = 64 * 1024 * 1024
|
|
_FILE_HASH_CACHE: dict[tuple[str, int, int], str] = {}
|
|
|
|
|
|
def _json_safe(value: Any) -> Any:
|
|
"""Convert cache-key inputs to deterministic, JSON-safe values."""
|
|
|
|
if value is None or isinstance(value, (str, bool, int)):
|
|
return value
|
|
if isinstance(value, float):
|
|
return value if math.isfinite(value) else repr(value)
|
|
if isinstance(value, Path):
|
|
return str(value)
|
|
if is_dataclass(value):
|
|
return _json_safe(asdict(value))
|
|
if isinstance(value, Mapping):
|
|
return {
|
|
str(key): _json_safe(item)
|
|
for key, item in sorted(value.items(), key=lambda pair: str(pair[0]))
|
|
}
|
|
if isinstance(value, (list, tuple)):
|
|
return [_json_safe(item) for item in value]
|
|
if isinstance(value, (set, frozenset)):
|
|
normalized = [_json_safe(item) for item in value]
|
|
return sorted(
|
|
normalized,
|
|
key=lambda item: json.dumps(item, sort_keys=True, separators=(",", ":")),
|
|
)
|
|
item = getattr(value, "item", None)
|
|
if callable(item):
|
|
try:
|
|
return _json_safe(item())
|
|
except Exception:
|
|
pass
|
|
raise TypeError(
|
|
"Director cache identity contains an unsupported value of type "
|
|
f"{type(value).__module__}.{type(value).__qualname__}."
|
|
)
|
|
|
|
|
|
def canonical_json_bytes(value: Any) -> bytes:
|
|
return json.dumps(
|
|
_json_safe(value),
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
ensure_ascii=True,
|
|
).encode("utf-8")
|
|
|
|
|
|
def canonical_sha256(value: Any) -> str:
|
|
return hashlib.sha256(canonical_json_bytes(value)).hexdigest()
|
|
|
|
|
|
def _sha256_file(path: Path) -> str:
|
|
stat = path.stat()
|
|
identity = (str(path.resolve(strict=False)), int(stat.st_size), int(stat.st_mtime_ns))
|
|
cached = _FILE_HASH_CACHE.get(identity)
|
|
if cached:
|
|
return cached
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
value = digest.hexdigest()
|
|
_FILE_HASH_CACHE[identity] = value
|
|
return value
|
|
|
|
|
|
def _file_identity(path: Path, *, content_hash: bool) -> dict[str, Any]:
|
|
try:
|
|
stat = path.stat()
|
|
except OSError as exc:
|
|
return {"name": path.name, "missing": True, "error": str(exc)}
|
|
result: dict[str, Any] = {
|
|
"name": path.name,
|
|
"size": int(stat.st_size),
|
|
"mtime_ns": int(stat.st_mtime_ns),
|
|
}
|
|
if content_hash and stat.st_size <= _SMALL_FILE_HASH_LIMIT:
|
|
try:
|
|
result["sha256"] = _sha256_file(path)
|
|
except OSError as exc:
|
|
result["hash_error"] = str(exc)
|
|
return result
|
|
|
|
|
|
def checkpoint_identity(model_path: str, status: Mapping[str, Any] | None = None) -> dict[str, Any]:
|
|
"""Return a cheap checkpoint identity without hashing model-sized shards."""
|
|
|
|
path = Path(str(model_path or "")).expanduser()
|
|
try:
|
|
resolved = path.resolve(strict=False)
|
|
except OSError:
|
|
resolved = path.absolute()
|
|
result: dict[str, Any] = {
|
|
"resolved_path": str(resolved),
|
|
"revision": str((status or {}).get("_commit_hash", "") or ""),
|
|
}
|
|
if not resolved.exists():
|
|
result["missing"] = True
|
|
return result
|
|
if resolved.is_file():
|
|
result["kind"] = "file"
|
|
# Hash ordinary config/checkpoint sidecars, but identify multi-gigabyte
|
|
# GGUF/model files by resolved path plus high-resolution file metadata.
|
|
result["file"] = _file_identity(
|
|
resolved,
|
|
content_hash=resolved.stat().st_size <= _SMALL_FILE_HASH_LIMIT,
|
|
)
|
|
return result
|
|
|
|
result["kind"] = "directory"
|
|
identity_names = (
|
|
"config.json",
|
|
"chat_template.jinja",
|
|
"generation_config.json",
|
|
"hf_quant_config.json",
|
|
"added_tokens.json",
|
|
"preprocessor_config.json",
|
|
"processor_config.json",
|
|
"special_tokens_map.json",
|
|
"tokenizer.json",
|
|
"tokenizer_config.json",
|
|
"tokenizer.model",
|
|
"model.safetensors.index.json",
|
|
)
|
|
result["identity_files"] = [
|
|
_file_identity(resolved / name, content_hash=True)
|
|
for name in identity_names
|
|
if (resolved / name).is_file()
|
|
]
|
|
|
|
shard_names: set[str] = set()
|
|
index_path = resolved / "model.safetensors.index.json"
|
|
if index_path.is_file():
|
|
try:
|
|
index = json.loads(index_path.read_text(encoding="utf-8"))
|
|
weight_map = index.get("weight_map") if isinstance(index, dict) else None
|
|
if isinstance(weight_map, dict):
|
|
shard_names.update(
|
|
str(name)
|
|
for name in weight_map.values()
|
|
if isinstance(name, str) and name
|
|
)
|
|
except (OSError, ValueError):
|
|
pass
|
|
if not shard_names:
|
|
shard_names.update(path.name for path in resolved.glob("*.safetensors"))
|
|
result["weight_files"] = [
|
|
_file_identity(resolved / name, content_hash=False)
|
|
for name in sorted(shard_names)
|
|
]
|
|
return result
|
|
|
|
|
|
def implementation_identity(root: Path, relative_paths: Sequence[str]) -> dict[str, Any]:
|
|
files: list[dict[str, Any]] = []
|
|
for relative in relative_paths:
|
|
path = root / relative
|
|
record = {"path": relative, **_file_identity(path, content_hash=True)}
|
|
files.append(record)
|
|
return {"files": files, "sha256": canonical_sha256(files)}
|
|
|
|
|
|
def ordered_tensor_content_hashes(images: Any) -> list[dict[str, Any]]:
|
|
"""Hash the exact ordered tensor samples delivered to the processor."""
|
|
|
|
if images is None:
|
|
return []
|
|
whole_value_is_one_sample = False
|
|
try:
|
|
image_shape = tuple(int(value) for value in images.shape)
|
|
whole_value_is_one_sample = len(image_shape) == 3
|
|
sample_count = 1 if whole_value_is_one_sample else int(image_shape[0])
|
|
except Exception:
|
|
try:
|
|
sample_count = len(images)
|
|
except Exception:
|
|
sample_count = 1
|
|
|
|
hashes: list[dict[str, Any]] = []
|
|
for index in range(max(0, sample_count)):
|
|
digest = hashlib.sha256()
|
|
try:
|
|
sample = images if whole_value_is_one_sample else images[index]
|
|
if hasattr(sample, "detach"):
|
|
sample = sample.detach().cpu().contiguous()
|
|
shape = [int(value) for value in getattr(sample, "shape", ())]
|
|
dtype = str(getattr(sample, "dtype", type(sample).__name__))
|
|
if hasattr(sample, "numpy"):
|
|
array = sample.numpy()
|
|
digest.update(memoryview(array).cast("B"))
|
|
else:
|
|
digest.update(repr(sample).encode("utf-8"))
|
|
hashes.append(
|
|
{
|
|
"position": index,
|
|
"shape": shape,
|
|
"dtype": dtype,
|
|
"sha256": digest.hexdigest(),
|
|
}
|
|
)
|
|
except Exception as exc:
|
|
digest.update(f"hash-error:{type(exc).__name__}:{exc}".encode("utf-8"))
|
|
hashes.append(
|
|
{
|
|
"position": index,
|
|
"hash_error": str(exc),
|
|
"sha256": digest.hexdigest(),
|
|
}
|
|
)
|
|
return hashes
|
|
|
|
|
|
def cache_root() -> Path:
|
|
override = os.environ.get("DG_DIRECTOR_CACHE_DIR", "").strip()
|
|
if override:
|
|
return Path(override).expanduser().resolve(strict=False)
|
|
try:
|
|
import folder_paths
|
|
|
|
return (
|
|
Path(folder_paths.get_system_user_directory("cache")).resolve(strict=False)
|
|
/ "diffusiongemma_prompt_builder"
|
|
/ "director_cache"
|
|
)
|
|
except Exception:
|
|
return (
|
|
Path(__file__).resolve().parents[2]
|
|
/ "user"
|
|
/ "__cache"
|
|
/ "diffusiongemma_prompt_builder"
|
|
/ "director_cache"
|
|
)
|
|
|
|
|
|
def cache_path(cache_key: str) -> Path:
|
|
normalized = str(cache_key or "").strip().lower()
|
|
if len(normalized) != 64 or any(character not in "0123456789abcdef" for character in normalized):
|
|
raise ValueError("Director cache key must be a 64-character SHA-256 digest.")
|
|
return cache_root() / normalized[:2] / f"{normalized}.json"
|
|
|
|
|
|
def load_cache_entry(cache_key: str) -> tuple[dict[str, Any] | None, str]:
|
|
path = cache_path(cache_key)
|
|
if not path.is_file():
|
|
return None, "not_found"
|
|
try:
|
|
maximum_bytes = max(
|
|
1024,
|
|
int(os.environ.get("DG_DIRECTOR_CACHE_MAX_BYTES", str(16 * 1024 * 1024))),
|
|
)
|
|
if path.stat().st_size > maximum_bytes:
|
|
return None, "entry_too_large"
|
|
payload = json.loads(path.read_text(encoding="utf-8"))
|
|
except Exception as exc:
|
|
return None, f"read_error:{type(exc).__name__}:{exc}"
|
|
if not isinstance(payload, dict):
|
|
return None, "invalid_root"
|
|
if payload.get("schema") != DIRECTOR_CACHE_SCHEMA:
|
|
return None, "schema_mismatch"
|
|
if payload.get("cache_key") != cache_key:
|
|
return None, "key_mismatch"
|
|
if not isinstance(payload.get("outputs"), dict):
|
|
return None, "outputs_missing"
|
|
expected_checksum = str(payload.get("outputs_sha256", "") or "")
|
|
if len(expected_checksum) != 64 or expected_checksum != canonical_sha256(payload["outputs"]):
|
|
return None, "outputs_checksum_mismatch"
|
|
return payload, "hit"
|
|
|
|
|
|
def save_cache_entry(cache_key: str, payload: Mapping[str, Any]) -> Path:
|
|
path = cache_path(cache_key)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
entry = dict(payload)
|
|
entry["schema"] = DIRECTOR_CACHE_SCHEMA
|
|
entry["cache_key"] = cache_key
|
|
entry["outputs_sha256"] = canonical_sha256(entry.get("outputs", {}))
|
|
encoded = json.dumps(entry, indent=2, sort_keys=True, ensure_ascii=True)
|
|
temporary_name = ""
|
|
try:
|
|
with tempfile.NamedTemporaryFile(
|
|
mode="w",
|
|
encoding="utf-8",
|
|
newline="\n",
|
|
prefix=f".{cache_key}.",
|
|
suffix=".partial",
|
|
dir=path.parent,
|
|
delete=False,
|
|
) as handle:
|
|
temporary_name = handle.name
|
|
handle.write(encoded)
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
os.replace(temporary_name, path)
|
|
finally:
|
|
if temporary_name:
|
|
try:
|
|
partial = Path(temporary_name)
|
|
if partial.exists():
|
|
partial.unlink()
|
|
except OSError:
|
|
pass
|
|
return path
|
|
|
|
|
|
__all__ = [
|
|
"DIRECTOR_CACHE_KEY_SCHEMA",
|
|
"DIRECTOR_CACHE_SCHEMA",
|
|
"cache_path",
|
|
"cache_root",
|
|
"canonical_sha256",
|
|
"checkpoint_identity",
|
|
"implementation_identity",
|
|
"load_cache_entry",
|
|
"ordered_tensor_content_hashes",
|
|
"save_cache_entry",
|
|
]
|