The smallest promotion

scripts/namespace_model_artifacts.py

@@ -0,0 +1,164 @@
+#!/usr/bin/env python3
+"""Publish rebuilt artifacts under a manifest namespace without changing sources."""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+
+def load_manifest(path: Path) -> tuple[dict, list[dict]]:
+  payload = json.loads(path.read_text(encoding="utf-8"))
+  if not isinstance(payload, dict) or not isinstance(payload.get("models"), list):
+    raise ValueError(f"Expected an object containing a models list: {path}")
+  return payload, payload["models"]
+
+
+def renamed_filename(filename: str, old_id: str, new_id: str) -> str:
+  old_prefix = f"{old_id}_driving_tinygrad.pkl"
+  if filename == old_prefix or filename.startswith(f"{old_prefix}."):
+    return f"{new_id}_driving_tinygrad.pkl{filename[len(old_prefix):]}"
+  return filename
+
+
+def rename_model_files(directory: Path, old_id: str, new_id: str) -> None:
+  if not directory.is_dir():
+    return
+
+  candidates = [
+    path for path in directory.iterdir()
+    if renamed_filename(path.name, old_id, new_id) != path.name
+  ]
+  if not candidates:
+    return
+
+  staged: list[tuple[Path, Path]] = []
+  for source in candidates:
+    temporary = source.with_name(f".{source.name}.namespace-tmp")
+    if temporary.exists():
+      temporary.unlink()
+    source.rename(temporary)
+    staged.append((temporary, directory / renamed_filename(source.name, old_id, new_id)))
+
+  for temporary, destination in staged:
+    if destination.exists():
+      destination.unlink()
+    temporary.rename(destination)
+
+  checksum = directory / f"{new_id}_driving_tinygrad.pkl.sha256"
+  if checksum.is_file():
+    text = checksum.read_text(encoding="utf-8")
+    text = text.replace(f"{old_id}_driving_tinygrad.pkl", f"{new_id}_driving_tinygrad.pkl")
+    checksum.write_text(text, encoding="utf-8")
+
+
+def remap_multipart_handoff(path: Path, id_map: dict[str, str]) -> None:
+  if not path.is_file():
+    return
+  payload = json.loads(path.read_text(encoding="utf-8"))
+  if not isinstance(payload, list):
+    raise ValueError(f"Expected a multipart handoff list: {path}")
+  for entry in payload:
+    if not isinstance(entry, dict):
+      continue
+    old_id = str(entry.get("id") or "")
+    new_id = id_map.get(old_id)
+    if not new_id:
+      continue
+    entry["id"] = new_id
+    for key in ("filename",):
+      if entry.get(key):
+        entry[key] = renamed_filename(str(entry[key]), old_id, new_id)
+    entry["parts"] = [renamed_filename(str(part), old_id, new_id) for part in entry.get("parts", [])]
+  path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
+
+
+def build_id_map(models: list[dict], suffix: str, manifest_version: str) -> dict[str, str]:
+  """Append the namespace suffix without colliding with an existing source ID."""
+  source_ids = [str(model["id"]) for model in models]
+  source_id_set = set(source_ids)
+  id_map: dict[str, str] = {}
+  for old_id in source_ids:
+    candidate = f"{old_id}{suffix}"
+    if candidate in source_id_set:
+      candidate = f"{old_id}{manifest_version}"
+    while candidate in id_map.values():
+      candidate += "_"
+    id_map[old_id] = candidate
+  return id_map
+
+
+def namespace_workspace(
+  workspace: Path,
+  base_manifest: Path,
+  manifest_version: str,
+  suffix: str,
+  completed_ids: list[str] | None = None,
+) -> Path:
+  payload, models = load_manifest(base_manifest)
+  if not suffix or any(not str(model.get("id") or "") for model in models):
+    raise ValueError("Every manifest model must have an ID and the suffix must be non-empty")
+
+  id_map = build_id_map(models, suffix, manifest_version)
+  if len(set(id_map.values())) != len(id_map):
+    raise ValueError("The requested model namespace contains ID collisions")
+
+  compiled = workspace / "compiled"
+  ready = workspace / "ready-for-resources"
+  ids_to_namespace = set(completed_ids) if completed_ids is not None else set(id_map)
+  unknown_ids = sorted(ids_to_namespace - set(id_map))
+  if unknown_ids:
+    raise KeyError(f"IDs are not present in the base manifest: {', '.join(unknown_ids)}")
+  missing = [
+    old_id for old_id in ids_to_namespace
+    if not (compiled / f"{old_id}_driving_tinygrad.pkl").is_file()
+    and not (compiled / f"{id_map[old_id]}_driving_tinygrad.pkl").is_file()
+  ]
+  if missing:
+    raise FileNotFoundError(f"Missing compiled artifacts for: {', '.join(sorted(missing))}")
+
+  for old_id, new_id in id_map.items():
+    if old_id not in ids_to_namespace:
+      continue
+    rename_model_files(compiled, old_id, new_id)
+    rename_model_files(ready, old_id, new_id)
+
+  namespaced_models = []
+  for model in models:
+    namespaced = dict(model)
+    namespaced["id"] = id_map[str(model["id"])]
+    namespaced_models.append(namespaced)
+
+  output = dict(payload)
+  output["models"] = namespaced_models
+  output_path = workspace / "manifests" / f"model_names_{manifest_version}.json"
+  output_path.parent.mkdir(parents=True, exist_ok=True)
+  output_path.write_text(json.dumps(output, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
+
+  remap_multipart_handoff(ready / "multipart.json", id_map)
+  (workspace / "source-maps" / f"model_id_namespace_{manifest_version}.json").write_text(
+    json.dumps(id_map, indent=2) + "\n", encoding="utf-8",
+  )
+  return output_path
+
+
+def main() -> int:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--workspace", type=Path, required=True)
+  parser.add_argument("--base-manifest", type=Path, required=True)
+  parser.add_argument("--manifest-version", default="v23")
+  parser.add_argument("--suffix", default="3")
+  parser.add_argument(
+    "--completed-ids",
+    nargs="+",
+    help="Only namespace these confirmed rebuilt IDs; still write the full manifest.",
+  )
+  args = parser.parse_args()
+
+  output = namespace_workspace(args.workspace, args.base_manifest, args.manifest_version, args.suffix, args.completed_ids)
+  print(output)
+  return 0
+
+
+if __name__ == "__main__":
+  raise SystemExit(main())