71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
|
|
|
|
DEFAULT_REQUIRED = ("final_metrics.json", "metrics.jsonl", "checkpoint.pt", "run_manifest.json")
|
|
|
|
|
|
def verify_artifacts(artifact_dir: str | Path, required: Iterable[str] = DEFAULT_REQUIRED) -> dict[str, Any]:
|
|
root = Path(artifact_dir)
|
|
if not root.exists():
|
|
raise FileNotFoundError(f"Artifact directory not found: {root}")
|
|
if not root.is_dir():
|
|
raise ValueError(f"Artifact path is not a directory: {root}")
|
|
|
|
missing = [name for name in required if not (root / name).is_file()]
|
|
if missing:
|
|
raise ValueError(f"Artifact directory missing required files: {', '.join(missing)}")
|
|
|
|
_validate_json(root / "final_metrics.json")
|
|
_validate_json(root / "run_manifest.json")
|
|
_validate_jsonl(root / "metrics.jsonl")
|
|
|
|
files = sorted(path for path in root.rglob("*") if path.is_file())
|
|
manifest = {
|
|
"artifact_dir": str(root),
|
|
"file_count": len(files),
|
|
"files": [
|
|
{
|
|
"path": str(path.relative_to(root)),
|
|
"bytes": path.stat().st_size,
|
|
"sha256": sha256_file(path),
|
|
}
|
|
for path in files
|
|
],
|
|
}
|
|
(root / "artifact_manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
|
|
(root / "checksums.txt").write_text(
|
|
"".join(f"{item['sha256']} {item['path']}\n" for item in manifest["files"])
|
|
)
|
|
return manifest
|
|
|
|
|
|
def sha256_file(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as file:
|
|
for chunk in iter(lambda: file.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _validate_json(path: Path) -> None:
|
|
try:
|
|
json.loads(path.read_text())
|
|
except json.JSONDecodeError as exc:
|
|
raise ValueError(f"Invalid JSON artifact {path}: {exc}") from exc
|
|
|
|
|
|
def _validate_jsonl(path: Path) -> None:
|
|
with path.open() as file:
|
|
for line_number, line in enumerate(file, start=1):
|
|
stripped = line.strip()
|
|
if not stripped:
|
|
continue
|
|
try:
|
|
json.loads(stripped)
|
|
except json.JSONDecodeError as exc:
|
|
raise ValueError(f"Invalid JSONL artifact {path}:{line_number}: {exc}") from exc
|