File size: 4,279 Bytes
ec24fb0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | """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
|