| """Logging, atomic writes, stage state, GPU discovery.""" |
| import json |
| import logging |
| import os |
| 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 process may use.""" |
| 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("M4_GPUS") |
| allowed = {int(x) for x in allowed.split(",")} if allowed else None |
| res = [] |
| for line in out.strip().splitlines(): |
| i, free, util = [int(x.strip()) for x in line.split(",")] |
| 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(): |
| return max(1, min(CFG["n_workers"], (os.cpu_count() or 8) // 2)) |
|
|
|
|
| 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) |
|
|
