| """Logging, atomic writes, stage state, GPU discovery, and a crash-safe shard work queue.""" |
| import json |
| import logging |
| import os |
| import socket |
| import subprocess |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
|
|
| from config import CFG, LOGS, STATE, WORK |
|
|
|
|
| def get_logger(name): |
| LOGS.mkdir(parents=True, exist_ok=True) |
| log = logging.getLogger(name) |
| if not log.handlers: |
| log.setLevel(logging.INFO) |
| fmt = logging.Formatter("%(asctime)s %(name)s %(levelname)s %(message)s", "%m-%d %H:%M:%S") |
| h = logging.StreamHandler(sys.stdout) |
| h.setFormatter(fmt) |
| log.addHandler(h) |
| log.propagate = False |
| return log |
|
|
|
|
| def _tmp(path): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| return path.with_name(f".{path.name}.tmp{os.getpid()}") |
|
|
|
|
| def save_parquet(df, path): |
| tmp = _tmp(path) |
| df.to_parquet(tmp, index=False) |
| os.replace(tmp, path) |
|
|
|
|
| def save_npz(path, **arrays): |
| tmp = _tmp(path) |
| with open(tmp, "wb") as f: |
| np.savez(f, **arrays) |
| os.replace(tmp, path) |
|
|
|
|
| def save_npy(path, arr): |
| tmp = _tmp(path) |
| with open(tmp, "wb") as f: |
| np.save(f, arr) |
| os.replace(tmp, path) |
|
|
|
|
| def save_json(obj, path): |
| tmp = _tmp(path) |
| tmp.write_text(json.dumps(obj, indent=2, default=str)) |
| os.replace(tmp, path) |
|
|
|
|
| def load_json(path): |
| return json.loads(Path(path).read_text()) |
|
|
|
|
| def stage_done(name): |
| return (STATE / f"{name}.done").exists() |
|
|
|
|
| def mark_done(name, info=None): |
| save_json({"finished": time.strftime("%Y-%m-%d %H:%M:%S"), **(info or {})}, STATE / f"{name}.done") |
|
|
|
|
| class Progress: |
| def __init__(self, log, total, what, every=60): |
| self.log, self.total, self.what, self.every = log, total, what, every |
| self.t0 = self.last = time.time() |
| self.done = 0 |
|
|
| def step(self, n=1): |
| self.done += n |
| now = time.time() |
| if now - self.last >= self.every or self.done >= self.total: |
| self.last = now |
| rate = self.done / max(now - self.t0, 1e-9) |
| eta = (self.total - self.done) / rate if rate > 0 else float("nan") |
| self.log.info(f"{self.what}: {self.done}/{self.total} ({100 * self.done / max(self.total, 1):.1f}%) " |
| f"elapsed {(now - self.t0) / 60:.1f}m eta {eta / 60:.1f}m") |
|
|
|
|
| def gpu_status(): |
| """[(index, free_mb, util_pct)] for the GPUs this pipeline may use (M5_GPUS restricts).""" |
| try: |
| out = subprocess.run(["nvidia-smi", "--query-gpu=index,memory.free,utilization.gpu", |
| "--format=csv,noheader,nounits"], capture_output=True, text=True, timeout=30).stdout |
| except Exception: |
| return [] |
| allowed = os.environ.get("M5_GPUS") |
| allowed = {int(x) for x in allowed.split(",") if x != ""} if allowed else None |
| res = [] |
| for line in out.strip().splitlines(): |
| try: |
| i, free, util = [int(x.strip()) for x in line.split(",")] |
| except ValueError: |
| continue |
| if allowed is None or i in allowed: |
| res.append((i, free, util)) |
| return res |
|
|
|
|
| def gpu_patient(fn, log, what, wait_s=60, max_wait_s=3 * 3600): |
| """Run fn(); on CUDA OOM (other users filling the shared GPU) free cache, wait, retry.""" |
| import torch |
| t0 = time.time() |
| while True: |
| try: |
| return fn() |
| except torch.cuda.OutOfMemoryError: |
| torch.cuda.empty_cache() |
| if time.time() - t0 > max_wait_s: |
| raise |
| log.warning(f"{what}: GPU out of memory (shared card busy), retrying in {wait_s}s") |
| time.sleep(wait_s) |
|
|
|
|
| def n_workers(): |
| """Worker processes: the configured cap, reduced when other users already load the (shared) CPUs.""" |
| ncpu = os.cpu_count() or 8 |
| try: |
| busy = os.getloadavg()[0] |
| except OSError: |
| busy = 0 |
| return int(max(8, min(CFG["n_workers"], ncpu - busy - 4))) |
|
|
|
|
| def shard_path(kind, split, i, ext="parquet", sub=None): |
| base = WORK / kind / split |
| if sub: |
| base = base / sub |
| return base / f"shard_{i:05d}.{ext}" |
|
|
|
|
| def read_shards(kind, split, n, columns=None, sub=None, ext="parquet"): |
| return pd.concat([pd.read_parquet(shard_path(kind, split, i, ext, sub), columns=columns) for i in range(n)], |
| ignore_index=True) |
|
|
|
|
| |
|
|
| def _pid_alive(pid): |
| try: |
| os.kill(pid, 0) |
| return True |
| except ProcessLookupError: |
| return False |
| except PermissionError: |
| return True |
|
|
|
|
| REMOTE_STALE_S = 2.5 * 3600 |
|
|
|
|
| def claim(out_path): |
| """Claim the job that produces out_path. Returns True if this process owns it now. A claim is a file created |
| with O_EXCL holding host:pid; claims of dead processes (crash, kill, OOM) are broken and re-taken, so any |
| number of workers on any GPUs can drain one shard list without double work or lost shards.""" |
| out_path = Path(out_path) |
| if out_path.exists(): |
| return False |
| c = out_path.with_name(out_path.name + ".claim") |
| c.parent.mkdir(parents=True, exist_ok=True) |
| me = f"{socket.gethostname()}:{os.getpid()}" |
| for _ in range(2): |
| try: |
| fd = os.open(c, os.O_CREAT | os.O_EXCL | os.O_WRONLY) |
| os.write(fd, me.encode()) |
| os.close(fd) |
| if out_path.exists(): |
| c.unlink(missing_ok=True) |
| return False |
| return True |
| except FileExistsError: |
| try: |
| host, pid = c.read_text().split(":") |
| if host.startswith("remote"): |
| stale = time.time() - c.stat().st_mtime > REMOTE_STALE_S |
| else: |
| stale = host == socket.gethostname() and not _pid_alive(int(pid)) |
| except (ValueError, FileNotFoundError): |
| stale = True |
| if not stale: |
| return False |
| c.unlink(missing_ok=True) |
| return False |
|
|
|
|
| def release(out_path): |
| Path(str(out_path) + ".claim").unlink(missing_ok=True) |
|
|
|
|
| def drain(paths, fn, log, what): |
| """Run fn(i) for every index whose output does not exist yet, claiming each one first. Safe to run from |
| several processes at once. Returns the number of jobs this process completed.""" |
| todo = [i for i, p in enumerate(paths) if not Path(p).exists()] |
| prog = Progress(log, len(todo), what) |
| n = 0 |
| for i in todo: |
| if not claim(paths[i]): |
| continue |
| try: |
| fn(i) |
| n += 1 |
| finally: |
| release(paths[i]) |
| prog.step() |
| return n |
|
|
|
|
| def spawn_gpu_workers(stage_args, log, min_free_mb): |
| """Start one worker process per assigned GPU (python stages.py <stage_args...>), wait for all, fail if any |
| failed. Workers share the job list through claim().""" |
| gpus = [g for g in os.environ.get("M5_ASSIGNED_GPUS", "0").split(",") if g != ""] |
| src = Path(__file__).resolve().parent |
| procs = [] |
| for g in gpus: |
| env = dict(os.environ, CUDA_VISIBLE_DEVICES=g, M5_ASSIGNED_GPUS=g) |
| |
| |
| procs.append(subprocess.Popen([sys.executable, "-u", str(src / "stages.py"), *stage_args, |
| "--owner=model-5-pipeline-DO-NOT-KILL"], env=env, cwd=str(src))) |
| log.info(f"worker {stage_args} on GPU {g} pid={procs[-1].pid}") |
| codes = [p.wait() for p in procs] |
| if any(codes): |
| raise RuntimeError(f"GPU workers {stage_args} failed: {codes}") |
|
|
|
|
| def disk_free_gb(path): |
| st = os.statvfs(path) |
| return st.f_bavail * st.f_frsize / 1e9 |
|
|
