PufferLib — high-throughput RL (GPU)

super run . --script run.sh

#!/usr/bin/env bash
# Train an RL agent on Breakout with PufferLib — high-throughput GPU-vectorized
# RL (millions of env steps/sec). Proves the PufferLib stack trains on supercomputer.
# See README.md.
set -euo pipefail
mkdir -p /workspace/out

# Soak knob (stress tier): when SUPERCOMPUTER_STRESS_DURATION_S is set, keep training
# fresh rounds until that many seconds have passed since this script started;
# the artifact (written only after the loop) summarizes the last round's log.
# Unset/0 = one round, exactly the daily-smoke behavior.
STRESS_END=$(( $(date +%s) + ${SUPERCOMPUTER_STRESS_DURATION_S:-0} ))

# Fail loudly if the GPU isn't visible — we compile + train on it.
nvidia-smi -L

# PufferLib compiles CUDA extensions that link libnvidia-ml. The driver ships
# only libnvidia-ml.so.1, so expose an unversioned .so for the linker. We build
# here (on the GPU host) rather than at image-build time so nvcc's -arch=native
# sees a real GPU and the NVML/driver libs are present (torch is pre-cached).
NVML=$(ldconfig -p 2>/dev/null | awk '/libnvidia-ml\.so\.1/{print $NF; exit}')
[ -z "$NVML" ] && NVML=$(find / -name 'libnvidia-ml.so.1' 2>/dev/null | head -1)
[ -n "$NVML" ] && ln -sf "$NVML" /usr/lib/libnvidia-ml.so || echo "WARN: libnvidia-ml.so.1 not found"
export LIBRARY_PATH="/usr/lib:/usr/local/cuda/lib64/stubs:${LIBRARY_PATH:-}"

cd /opt/pufferlib
pip install -e . --break-system-packages   # torch cached; compiles ext on GPU
bash build.sh breakout                       # compile the Breakout Ocean (C) env

LOG=/workspace/out/train.log
# Run under a pty (script) so PufferLib's live dashboard renders into the log for
# the summarizer. A few million steps shows the score climbing in minutes; the
# soak loop just trains round after round (the log holds the last round).
while :; do
    script -qe -c "python -m pufferlib.pufferl train breakout --train.total-timesteps 10000000" "$LOG"
    if [ "$(date +%s)" -ge "$STRESS_END" ]; then break; fi
done

# Distill the run into the artifact (fails if no training evidence).
exec python "$(dirname "$0")/summarize.py" "$LOG"

"""Distill a PufferLib training run into the smoke artifact.

PufferLib's live dashboard is ANSI; we strip it and look for training evidence
(SPS / steps / epoch) plus any saved checkpoint. A PASS means training actually
ran on the GPU, not just that the process launched.
"""

import glob
import json
import os
import re
import sys
from pathlib import Path

OUT = Path("/workspace/out/pufferlib.json")
OUT.parent.mkdir(parents=True, exist_ok=True)

_ANSI = re.compile(r"\x1b\[[0-9;?]*[A-Za-z]")
_SPS = re.compile(r"SPS\s+([0-9][0-9.,]*\s*[KMBT]?)")


def main() -> None:
    raw = Path(sys.argv[1]).read_text(errors="replace") if len(sys.argv) > 1 else ""
    text = _ANSI.sub("", raw)

    # Dashboard fields that only appear once the training loop is running.
    evidence = [k for k in ("SPS", "Steps", "Epoch", "Uptime") if k in text]
    m = _SPS.search(text)
    sps = m.group(1).strip() if m else None

    # Checkpoints PufferLib saves during/after training (location varies by run).
    checkpoints = glob.glob("/opt/pufferlib/**/*.pt", recursive=True)
    checkpoints += glob.glob(os.path.expanduser("~/**/*.pt"), recursive=True)

    if not evidence and not checkpoints:
        print("ERROR: no training evidence (no dashboard metrics, no checkpoint)",
              file=sys.stderr)
        sys.exit(1)

    OUT.write_text(json.dumps({
        "env": "breakout",
        "framework": "pufferlib 4.0 (puffer train)",
        "trained": True,
        "evidence_fields": evidence,
        "sps": sps,
        "checkpoint_found": bool(checkpoints),
    }, indent=2) + "\n")
    print("wrote", OUT)


if __name__ == "__main__":
    main()

{
    "sku": "gpu_1x_l4",
    "image": {
        "base": "nvidia/cuda:12.6.2-cudnn-devel-ubuntu24.04",
        "apt": ["git", "build-essential", "clang", "ccache", "libomp-dev", "libglfw3", "libgl1-mesa-dev", "python3.12-dev", "python3-pip", "python-is-python3", "ca-certificates", "curl"],
        "pip": [],
        "run": [
            "git clone --depth 1 --branch 4.0 https://github.com/PufferAI/PufferLib.git /opt/pufferlib",
            "uv pip install --system torch"
        ],
        "env": {
            "PYTHONUNBUFFERED": "1",
            "WANDB_MODE": "disabled",
            "SSL_CERT_FILE": "/etc/ssl/certs/ca-certificates.crt",
            "PIP_CERT": "/etc/ssl/certs/ca-certificates.crt",
            "REQUESTS_CA_BUNDLE": "/etc/ssl/certs/ca-certificates.crt"
        }
    }
}