Buckets:
| """Run ownership, provenance, deterministic state, and adapter persistence.""" | |
| from __future__ import annotations | |
| import fcntl | |
| import hashlib | |
| import json | |
| import os | |
| import random | |
| import socket | |
| import sys | |
| import uuid | |
| from datetime import datetime, timedelta, timezone | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| SUBMISSION_DIRECTORY = "_submission" | |
| RUN_CLAIM = ".run-claim.json" | |
| def timestamped_run_name(name: str, *, now: datetime | None = None) -> str: | |
| if not name or Path(name).name != name or name in {".", ".."}: | |
| raise ValueError("run name must be one non-empty path component") | |
| if len(name) > 16 and name[8:9] == "-" and name[15:16] == "_": | |
| try: | |
| datetime.strptime(name[:15], "%Y%m%d-%H%M%S") | |
| return name | |
| except ValueError: | |
| pass | |
| offset = int(os.environ.get("SPATIAL_MMCOT_TZ_OFFSET", "-7")) | |
| if not -23 <= offset <= 23: | |
| raise ValueError("SPATIAL_MMCOT_TZ_OFFSET must be between -23 and 23") | |
| stamp_source = now or datetime.now(timezone(timedelta(hours=offset))) | |
| return f"{stamp_source.strftime('%Y%m%d-%H%M%S')}_{name}" | |
| def acquire_run_writer(path: str | Path): | |
| handle = (Path(path) / ".active-writer.lock").open("a+", encoding="utf-8") | |
| try: | |
| fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) | |
| except BlockingIOError as error: | |
| handle.close() | |
| raise RuntimeError(f"run already has an active writer: {path}") from error | |
| return handle | |
| def _validate_submission_attempt(path: Path) -> None: | |
| submission = path / SUBMISSION_DIRECTORY | |
| ready = submission / "READY" | |
| job_id = os.environ.get("SLURM_JOB_ID") | |
| if not ready.is_file() or not job_id: | |
| raise FileExistsError( | |
| f"existing run lacks a verified Slurm submission attempt: {path}" | |
| ) | |
| try: | |
| ready_payload = json.loads(ready.read_text(encoding="utf-8")) | |
| manifest = json.loads( | |
| (submission / "manifest.json").read_text(encoding="utf-8") | |
| ) | |
| attestation = json.loads( | |
| (submission / f"runtime-verification-{job_id}.json").read_text( | |
| encoding="utf-8" | |
| ) | |
| ) | |
| except (OSError, json.JSONDecodeError) as error: | |
| raise FileExistsError(f"invalid submission verification in {path}") from error | |
| if ( | |
| ready_payload.get("schema") != 1 | |
| or not isinstance(ready_payload.get("metadata_sha256"), dict) | |
| or attestation.get("schema") != "spatial_mmcot_runtime_verification_v1" | |
| or attestation.get("slurm_job_id") != job_id | |
| or attestation.get("source_hash") != manifest.get("copied_hash") | |
| or attestation.get("ready_sha256") != file_sha256(ready) | |
| ): | |
| raise FileExistsError(f"submission verification mismatch in {path}") | |
| def claim_run_directory(path: str | Path, *, resume: bool = False) -> Path: | |
| path = Path(path) | |
| if path.is_symlink(): | |
| raise ValueError(f"run directory cannot be a symlink: {path}") | |
| if os.environ.get("SLURM_JOB_ID") and not path.exists(): | |
| raise FileNotFoundError( | |
| "Slurm jobs must claim a submission-prepared run directory: " | |
| f"{path}" | |
| ) | |
| try: | |
| path.mkdir(parents=True, exist_ok=False) | |
| return path | |
| except FileExistsError: | |
| _validate_submission_attempt(path) | |
| job_id = os.environ.get("SLURM_JOB_ID") | |
| claim = path / RUN_CLAIM | |
| if claim.exists(): | |
| if resume: | |
| return path | |
| raise RuntimeError(f"run directory is already claimed: {path}") | |
| unexpected = [ | |
| child.name | |
| for child in path.iterdir() | |
| if child.name != SUBMISSION_DIRECTORY | |
| and not ( | |
| child.is_file() | |
| and child.name.startswith("slurm-") | |
| and child.name.endswith(".out") | |
| ) | |
| ] | |
| if unexpected: | |
| raise FileExistsError( | |
| f"submission-prepared run contains experiment artifacts: {unexpected}" | |
| ) | |
| try: | |
| descriptor = os.open( | |
| claim, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o640 | |
| ) | |
| except FileExistsError as error: | |
| raise RuntimeError(f"run directory is already claimed: {path}") from error | |
| with os.fdopen(descriptor, "w", encoding="utf-8") as handle: | |
| json.dump( | |
| { | |
| "schema": "spatial_mmcot_run_claim_v1", | |
| "slurm_job_id": job_id, | |
| "host": socket.gethostname(), | |
| "claimed_at": datetime.now(timezone.utc).isoformat(), | |
| }, | |
| handle, | |
| indent=2, | |
| ) | |
| handle.write("\n") | |
| return path | |
| def resolve_run_directory( | |
| config, task: str = "sft", *, category: str | None = None | |
| ) -> Path: | |
| if config.run.output_dir: | |
| return Path(config.run.output_dir) | |
| repository = Path( | |
| os.environ.get("SPATIAL_MMCOT_PROJECT_ROOT", Path(__file__).resolve().parents[3]) | |
| ).resolve() | |
| root = config.run.output_root or os.environ.get("SPATIAL_MMCOT_OUTPUT_ROOT") | |
| root = Path(root).expanduser() if root else repository / "experiments/output" | |
| if not root.is_absolute(): | |
| root = repository / root | |
| parent = root / "uniar-sft" / config.run.dataset | |
| parent /= getattr(config.run, "version", None) or "unversioned" | |
| return parent / (category or config.run.category) / task / timestamped_run_name( | |
| config.run.name | |
| ) | |
| def write_json(path: Path, value) -> None: | |
| temporary = path.parent / f".{path.name}.tmp-{uuid.uuid4().hex}" | |
| temporary.write_text( | |
| json.dumps(value, indent=2, ensure_ascii=False, default=str) + "\n", | |
| encoding="utf-8", | |
| ) | |
| temporary.replace(path) | |
| def file_sha256(path: Path) -> str: | |
| digest = hashlib.sha256() | |
| with path.open("rb") as handle: | |
| for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): | |
| digest.update(chunk) | |
| return digest.hexdigest() | |
| def _fingerprint_paths(root: Path, paths: list[Path], *, content_limit: int | None) -> str: | |
| root = root.resolve() | |
| digest = hashlib.sha256() | |
| for path in sorted({path.resolve() for path in paths}, key=str): | |
| info = path.stat() | |
| try: | |
| name = path.relative_to(root) | |
| except ValueError: | |
| name = path | |
| digest.update(f"{name}\0{info.st_size}\0{info.st_mtime_ns}\0".encode()) | |
| if content_limit is None or info.st_size <= content_limit: | |
| digest.update(file_sha256(path).encode()) | |
| return digest.hexdigest() | |
| def cache_fingerprint(root: str | Path, rows: tuple[dict, ...] | list[dict]) -> str: | |
| root = Path(root).resolve() | |
| metadata = [root / "READY", root / "manifest.jsonl"] | |
| shards = [root / name for name in {row["cache_shard"] for row in rows}] | |
| digest = hashlib.sha256() | |
| digest.update(_fingerprint_paths(root, metadata, content_limit=None).encode()) | |
| digest.update(_fingerprint_paths(root, shards, content_limit=0).encode()) | |
| return digest.hexdigest() | |
| def model_fingerprint(root: str | Path) -> str: | |
| root = Path(root).resolve() | |
| return _fingerprint_paths( | |
| root, | |
| [path for path in root.rglob("*") if path.is_file()], | |
| content_limit=64 * 1024 * 1024, | |
| ) | |
| def environment_record(device: torch.device, world_size: int) -> dict: | |
| return { | |
| "python": sys.version, | |
| "torch": torch.__version__, | |
| "cuda": torch.version.cuda, | |
| "world_size": world_size, | |
| "device": str(device), | |
| "gpu": torch.cuda.get_device_name(device) if device.type == "cuda" else None, | |
| "gpuq_assigned_id": os.environ.get("GPUQ_ASSIGNED_ID"), | |
| "visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), | |
| } | |
| def source_fingerprint(*packages: str) -> str: | |
| if not packages: | |
| raise ValueError("source_fingerprint requires explicit package names") | |
| ours = Path(__file__).resolve().parents[2] | |
| paths = [] | |
| for package in packages: | |
| root = ours / "src" / package | |
| package_paths = list(root.rglob("*.py")) if root.is_dir() else [] | |
| if not package_paths: | |
| raise ValueError(f"source package is missing or empty: {package}") | |
| paths.extend(package_paths) | |
| return _fingerprint_paths(ours, paths, content_limit=None) | |
| def seed_everything(seed: int, rank: int = 0) -> None: | |
| value = seed + rank | |
| random.seed(value) | |
| np.random.seed(value) | |
| torch.manual_seed(value) | |
| if torch.cuda.is_available(): | |
| torch.cuda.manual_seed_all(value) | |
| def rng_state() -> dict: | |
| return { | |
| "python": random.getstate(), | |
| "numpy": np.random.get_state(), | |
| "torch_cpu": torch.get_rng_state(), | |
| "torch_cuda": torch.cuda.get_rng_state() if torch.cuda.is_available() else None, | |
| } | |
| def restore_rng_state(state: dict) -> None: | |
| random.setstate(state["python"]) | |
| np.random.set_state(state["numpy"]) | |
| torch.set_rng_state(state["torch_cpu"]) | |
| if torch.cuda.is_available() and state["torch_cuda"] is not None: | |
| torch.cuda.set_rng_state(state["torch_cuda"]) | |
| def save_adapter(model, processor, path: Path) -> None: | |
| path.mkdir(parents=True, exist_ok=False) | |
| kwargs = {"safe_serialization": True} | |
| if hasattr(model, "peft_config"): | |
| kwargs["save_embedding_layers"] = False | |
| model.save_pretrained(path, **kwargs) | |
| processor.save_pretrained(path) | |
| def save_adapter_atomic(model, processor, path: Path) -> Path: | |
| """Publish a final adapter idempotently without exposing partial output.""" | |
| complete = path / "COMPLETE" | |
| if path.exists(): | |
| if complete.is_file(): | |
| return path | |
| raise FileExistsError(f"incomplete final adapter already exists: {path}") | |
| temporary = path.parent / f".{path.name}.tmp-{uuid.uuid4().hex}" | |
| save_adapter(model, processor, temporary) | |
| (temporary / "COMPLETE").write_text("complete\n", encoding="utf-8") | |
| temporary.rename(path) | |
| return path | |
| __all__ = [ | |
| "acquire_run_writer", | |
| "cache_fingerprint", | |
| "claim_run_directory", | |
| "environment_record", | |
| "file_sha256", | |
| "model_fingerprint", | |
| "resolve_run_directory", | |
| "restore_rng_state", | |
| "rng_state", | |
| "save_adapter", | |
| "save_adapter_atomic", | |
| "seed_everything", | |
| "source_fingerprint", | |
| "timestamped_run_name", | |
| "write_json", | |
| ] | |
Xet Storage Details
- Size:
- 10.4 kB
- Xet hash:
- 7b539f0459c7fd6a537ef90596e102031233c5b9499261a6fa808bddd6380aff
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.