| """Shared helpers: secrets, seeding, JSONL I/O, environment fingerprint.""" |
| from __future__ import annotations |
|
|
| import json |
| import os |
| import platform |
| import random |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def ensure_python_headers() -> None: |
| """Triton (torch.compile, Liger, flash-linear-attention) compiles helpers against Python.h. This machine has |
| no python3.12-dev and no sudo, so the headers are unpacked to ~/.local/pydev (README, setup). Point gcc at |
| them when the system copy is missing; a no-op on machines that have the headers.""" |
| import sysconfig |
|
|
| if (Path(sysconfig.get_paths()["include"]) / "Python.h").exists(): |
| return |
| local = Path.home() / ".local" / "pydev" / "root" / "usr" / "include" |
| if (local / f"python{sys.version_info.major}.{sys.version_info.minor}" / "Python.h").exists(): |
| inc = [str(local / f"python{sys.version_info.major}.{sys.version_info.minor}"), str(local)] |
| os.environ["CPATH"] = os.pathsep.join(inc + [p for p in os.environ.get("CPATH", "").split(os.pathsep) if p]) |
|
|
|
|
| def load_env(path: Path = ROOT / ".env") -> None: |
| """Read KEY=VALUE lines from .env; map hf_token -> HF_TOKEN for huggingface_hub.""" |
| ensure_python_headers() |
| if not path.exists(): |
| return |
| for line in path.read_text().splitlines(): |
| line = line.strip() |
| if not line or line.startswith("#") or "=" not in line: |
| continue |
| key, val = (s.strip().strip('"').strip("'") for s in line.split("=", 1)) |
| if key.lower() in ("hf_token", "hugging_face_hub_token"): |
| key = "HF_TOKEN" |
| os.environ.setdefault(key, val) |
|
|
|
|
| def set_seed(seed: int) -> None: |
| random.seed(seed) |
| os.environ["PYTHONHASHSEED"] = str(seed) |
| try: |
| import numpy as np |
|
|
| np.random.seed(seed) |
| except ImportError: |
| pass |
| try: |
| import torch |
|
|
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| torch.backends.cuda.matmul.allow_tf32 = False |
| torch.backends.cudnn.allow_tf32 = False |
| except ImportError: |
| pass |
|
|
|
|
| def cap_gpu_memory(gb: float | None = None) -> None: |
| """Limit PyTorch's CUDA allocator (default 100 GB, env SMMLU_GPU_MEM_GB). |
| |
| The DGX Spark shares one 128 GB pool between CPU and GPU: an unbounded GPU job does not get a CUDA |
| out-of-memory error, it starves the system and the kernel OOM killer takes down the desktop session and |
| tmux (10 Oct, 11:15). With the cap, the job fails with torch.OutOfMemoryError and everything else survives. |
| """ |
| import torch |
|
|
| if not torch.cuda.is_available(): |
| return |
| gb = float(gb or os.environ.get("SMMLU_GPU_MEM_GB", 100)) |
| total = torch.cuda.get_device_properties(0).total_memory |
| torch.cuda.set_per_process_memory_fraction(min(gb * 1024**3 / total, 1.0), 0) |
|
|
|
|
| def read_jsonl(path: str | Path) -> list[dict]: |
| with open(path, encoding="utf-8") as f: |
| return [json.loads(line) for line in f if line.strip()] |
|
|
|
|
| def write_jsonl(path: str | Path, rows) -> None: |
| Path(path).parent.mkdir(parents=True, exist_ok=True) |
| with open(path, "w", encoding="utf-8") as f: |
| for r in rows: |
| f.write(json.dumps(r, ensure_ascii=False) + "\n") |
|
|
|
|
| def env_fingerprint() -> dict: |
| """Versions and code revision, stored with every run for reproducibility.""" |
| info = {"python": platform.python_version(), "machine": platform.machine()} |
| for mod in ("torch", "transformers", "datasets", "peft", "accelerate", "huggingface_hub"): |
| try: |
| info[mod] = __import__(mod).__version__ |
| except ImportError: |
| pass |
| try: |
| import torch |
|
|
| if torch.cuda.is_available(): |
| info["gpu"] = torch.cuda.get_device_name(0) |
| info["cuda"] = torch.version.cuda |
| except ImportError: |
| pass |
| try: |
| rev = subprocess.run(["git", "rev-parse", "HEAD"], cwd=ROOT, capture_output=True, text=True) |
| dirty = subprocess.run(["git", "status", "--porcelain"], cwd=ROOT, capture_output=True, text=True) |
| info["git_commit"] = rev.stdout.strip() or None |
| info["git_dirty"] = bool(dirty.stdout.strip()) |
| except FileNotFoundError: |
| pass |
| return info |
|
|
