3.12 Checkpointing, Fault Tolerance & Long-Running Jobs¶

"""
checkpoint.py — Production-grade FSDP checkpointing utilities.

Requirements: PyTorch >= 2.1, torchdata, FSDP model and optimizer.
"""

import os
import json
import torch
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.fully_sharded_data_parallel import (
    StateDictType,
    FullStateDictConfig,
    ShardedStateDictConfig,
)
from torch.distributed.checkpoint import (
    save,
    load,
    FileSystemWriter,
    FileSystemReader,
)
from torch.distributed.checkpoint.metadata import BytesStorageMetadata
from pathlib import Path
import random
import numpy as np


# --------------------------------------------------------------------------
# RNG state helpers
# --------------------------------------------------------------------------

def get_rng_state() -> dict:
    """Capture all RNG states on the current rank."""
    return {
        "python": random.getstate(),
        "numpy": np.random.get_state(),
        "torch_cpu": torch.get_rng_state(),
        "torch_cuda": torch.cuda.get_rng_state(),  # current device only
    }


def restore_rng_state(state: dict) -> None:
    """Restore RNG states — must be called on the same rank."""
    random.setstate(state["python"])
    np.random.set_state(state["numpy"])
    torch.set_rng_state(state["torch_cpu"])
    torch.cuda.set_rng_state(state["torch_cuda"])


# --------------------------------------------------------------------------
# Sharded checkpoint save (topology-agnostic via DCP)
# --------------------------------------------------------------------------

def save_checkpoint(
    model: FSDP,
    optimizer: torch.optim.Optimizer,
    step: int,
    lr_scheduler,
    data_loader_state: dict,
    checkpoint_dir: str,
    rank: int,
) -> None:
    """
    Save a complete training checkpoint using PyTorch Distributed Checkpoint.
    All ranks participate; I/O is parallel across ranks.

    Args:
        model:            The FSDP-wrapped model.
        optimizer:        The optimizer (may also be FSDP-sharded).
        step:             Global training step.
        lr_scheduler:     LR scheduler object.
        data_loader_state: Dict with shard file index and byte offset.
        checkpoint_dir:   Root directory for all checkpoints.
        rank:             Current rank (for per-rank files).
    """
    ckpt_path = Path(checkpoint_dir) / f"step_{step:08d}"
    ckpt_path.mkdir(parents=True, exist_ok=True)

    # 1. Collect sharded model + optimizer state dict (stays sharded, no gather)
    with FSDP.state_dict_type(
        model,
        StateDictType.SHARDED_STATE_DICT,
        ShardedStateDictConfig(offload_to_cpu=True),
    ):
        model_state = model.state_dict()
        optim_state = FSDP.optim_state_dict(model, optimizer)

    # 2. Parallel distributed save — every rank writes its own shards
    save(
        {"model": model_state, "optimizer": optim_state},
        storage_writer=FileSystemWriter(ckpt_path / "model_optim"),
    )

    # 3. Per-rank RNG state (tiny, but critical for exact reproducibility)
    rng_path = ckpt_path / f"rng_state_rank{rank}.pt"
    torch.save(get_rng_state(), rng_path)

    # 4. Training metadata — only rank 0 writes the shared metadata file
    if rank == 0:
        metadata = {
            "step": step,
            "lr_scheduler": lr_scheduler.state_dict(),
            "data_loader": data_loader_state,
        }
        with open(ckpt_path / "metadata.json", "w") as f:
            json.dump(metadata, f, indent=2)

    # 5. Synchronise: all ranks must finish before we consider the checkpoint
    #    "complete". Write an atomic sentinel file last.
    dist.barrier()
    if rank == 0:
        (ckpt_path / "COMPLETE").touch()
        print(f"[rank 0] Checkpoint saved: {ckpt_path}")


# --------------------------------------------------------------------------
# Sharded checkpoint load
# --------------------------------------------------------------------------

def load_checkpoint(
    model: FSDP,
    optimizer: torch.optim.Optimizer,
    checkpoint_dir: str,
    rank: int,
) -> dict:
    """
    Load a checkpoint. Returns metadata dict (step, scheduler state, etc.).
    Works even if the current topology differs from the saved topology.
    """
    ckpt_path = Path(checkpoint_dir)
    assert (ckpt_path / "COMPLETE").exists(), \
        f"Checkpoint at {ckpt_path} is incomplete or corrupted!"

    # 1. Load model + optimizer via DCP (handles resharding automatically)
    with FSDP.state_dict_type(
        model,
        StateDictType.SHARDED_STATE_DICT,
        ShardedStateDictConfig(offload_to_cpu=True),
    ):
        state_dict = {"model": model.state_dict(), "optimizer": {}}
        load(
            state_dict,
            storage_reader=FileSystemReader(ckpt_path / "model_optim"),
        )
        model.load_state_dict(state_dict["model"])
        optim_state = FSDP.optim_state_dict_to_load(
            model, optimizer, state_dict["optimizer"]
        )
        optimizer.load_state_dict(optim_state)

    # 2. Restore per-rank RNG state.
    #    weights_only=False is required here: since PyTorch 2.6 `torch.load`
    #    defaults to weights_only=True, whose restricted unpickler rejects the
    #    numpy ndarray hidden inside np.random.get_state()'s tuple. Only ever
    #    pass weights_only=False for files you produced yourself.
    rng_path = ckpt_path / f"rng_state_rank{rank}.pt"
    if rng_path.exists():
        restore_rng_state(
            torch.load(rng_path, map_location="cpu", weights_only=False)
        )

    # 3. Load shared metadata (every rank reads it for the scheduler / step)
    with open(ckpt_path / "metadata.json") as f:
        metadata = json.load(f)

    dist.barrier()
    return metadata


# --------------------------------------------------------------------------
# Checkpoint manager: keeps last K checkpoints, rotates older ones
# --------------------------------------------------------------------------

class CheckpointManager:
    def __init__(self, root_dir: str, keep_last_k: int = 3):
        self.root_dir = Path(root_dir)
        self.root_dir.mkdir(parents=True, exist_ok=True)
        self.keep_last_k = keep_last_k
        self._saved: list[Path] = []

    def save(self, model, optimizer, step, lr_scheduler,
             data_loader_state, rank):
        save_checkpoint(
            model, optimizer, step, lr_scheduler,
            data_loader_state, str(self.root_dir), rank
        )
        ckpt_path = self.root_dir / f"step_{step:08d}"
        self._saved.append(ckpt_path)
        self._saved.sort()

        # Remove checkpoints older than keep_last_k
        while len(self._saved) > self.keep_last_k:
            old = self._saved.pop(0)
            if rank == 0 and old.exists():
                import shutil
                shutil.rmtree(old)
                print(f"[rank 0] Removed old checkpoint: {old}")

    def latest(self) -> Path | None:
        """Find the most recent COMPLETE checkpoint under root_dir."""
        candidates = sorted(self.root_dir.glob("step_*"))
        for ckpt in reversed(candidates):
            if (ckpt / "COMPLETE").exists():
                return ckpt
        return None

"""
dcp_modern.py — The parallelism-agnostic PyTorch DCP idiom (PyTorch >= 2.4).
Works unchanged for single-GPU, DDP, FSDP1 and FSDP2 (`fully_shard`).
"""

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import (
    get_state_dict, set_state_dict,
)
from torch.distributed.checkpoint.stateful import Stateful


class TrainState(Stateful):
    """Bundles everything that must resume together into one DCP entry.

    Because it implements the Stateful protocol, DCP calls .state_dict() at
    save time and .load_state_dict() at load time — so the model, optimizer,
    LR schedule, step counter and data cursor are written as ONE consistent
    checkpoint, with no separate metadata file to fall out of sync.
    """

    def __init__(self, model, optimizer, scheduler, dataloader):
        self.model, self.optimizer = model, optimizer
        self.scheduler, self.dataloader = scheduler, dataloader
        self.step = 0

    def state_dict(self) -> dict:
        # get_state_dict() returns *sharded* DTensors under FSDP and plain
        # tensors otherwise; DCP handles both without special-casing.
        model_sd, optim_sd = get_state_dict(self.model, self.optimizer)
        return {
            "model": model_sd,
            "optim": optim_sd,
            "sched": self.scheduler.state_dict(),
            "data": self.dataloader.state_dict(),   # StatefulDataLoader
            "step": self.step,
        }

    def load_state_dict(self, sd: dict) -> None:
        # set_state_dict() reshards the loaded tensors onto the CURRENT
        # topology — this is what makes a 512-GPU checkpoint loadable on 256.
        set_state_dict(
            self.model, self.optimizer,
            model_state_dict=sd["model"], optim_state_dict=sd["optim"],
        )
        self.scheduler.load_state_dict(sd["sched"])
        self.dataloader.load_state_dict(sd["data"])
        self.step = sd["step"]


def save(train_state, step: int, root: str, blocking: bool = False):
    """Synchronous or asynchronous save. Returns a future when async."""
    app_state = {"app": train_state}          # key must match on load
    ckpt_id = f"{root}/step_{step:08d}"
    if blocking:
        dcp.save(app_state, checkpoint_id=ckpt_id)
        return None
    # async_save stages tensors to CPU, then writes from a background
    # process; await the returned future before the NEXT async_save.
    return dcp.async_save(app_state, checkpoint_id=ckpt_id)


def restore(train_state, ckpt_id: str) -> None:
    """In-place load. The dict must be pre-populated with the same keys so
    DCP knows the destination shapes/shardings to resolve chunks into."""
    app_state = {"app": train_state}
    dcp.load(app_state, checkpoint_id=ckpt_id)

"""
async_checkpoint.py — Async checkpointing with background I/O thread.

The snapshot step (GPU -> CPU copy) is the only synchronisation point.
Disk I/O runs concurrently with the next training steps.
"""

import threading
import time
from copy import deepcopy
from pathlib import Path
import torch
import torch.distributed as dist


class AsyncCheckpointer:
    """
    Async checkpointing: copies state to CPU, then writes in a background
    thread so training can continue immediately.

    Usage:
        checkpointer = AsyncCheckpointer(save_fn=save_checkpoint, ...)
        # In training loop:
        checkpointer.save_if_due(step, model, optimizer, ...)
        # At the end of training:
        checkpointer.wait()
    """

    def __init__(self, save_fn, checkpoint_interval: int, root_dir: str):
        self.save_fn = save_fn
        self.interval = checkpoint_interval
        self.root_dir = root_dir
        self._thread: threading.Thread | None = None
        self._lock = threading.Lock()
        self._pending_error: Exception | None = None

    def _check_for_errors(self):
        """Propagate any exception from the background thread."""
        with self._lock:
            if self._pending_error is not None:
                raise self._pending_error

    @staticmethod
    def _to_cpu(obj):
        """Recursively copy every tensor in a nested state dict to CPU."""
        if torch.is_tensor(obj):
            return obj.detach().to("cpu", copy=True)
        if isinstance(obj, dict):
            return {k: AsyncCheckpointer._to_cpu(v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [AsyncCheckpointer._to_cpu(v) for v in obj]
        if isinstance(obj, tuple):
            return tuple(AsyncCheckpointer._to_cpu(v) for v in obj)
        return deepcopy(obj)

    def _snapshot_to_cpu(self, model, optimizer) -> tuple[dict, dict]:
        """
        Copy all tensors out to CPU RAM. This is the only step that blocks
        training, and it is fast because it is a device-to-host copy, not a
        disk write.

        Caveat: `.to("cpu")` allocates *pageable* memory. Production stagers
        (PyTorch DCP, DeepSpeed) copy into a pre-allocated *pinned* buffer,
        which both doubles achievable PCIe bandwidth and lets the copy run
        on a side CUDA stream, overlapping the next forward pass.
        """
        # model.state_dict() already returns CPU copies under
        # FSDP + SHARDED_STATE_DICT + offload_to_cpu=True; _to_cpu() makes
        # that unconditional so this class also works without that config.
        model_state = self._to_cpu(model.state_dict())
        # The optimizer state must ALSO leave the GPU. A plain
        # deepcopy(optimizer.state_dict()) clones Adam's moments *on device*,
        # transiently doubling optimizer VRAM at the worst possible moment.
        optim_state = self._to_cpu(optimizer.state_dict())
        return model_state, optim_state

    def _background_save(self, model_state, optim_state, step, metadata):
        """Runs in background thread: pure I/O, no GPU interaction."""
        try:
            start = time.time()
            ckpt_path = Path(self.root_dir) / f"step_{step:08d}"
            ckpt_path.mkdir(parents=True, exist_ok=True)
            torch.save(
                {"model": model_state, "optimizer": optim_state,
                 "metadata": metadata},
                ckpt_path / "checkpoint.pt"
            )
            (ckpt_path / "COMPLETE").touch()
            elapsed = time.time() - start
            print(f"[async ckpt] step {step} written in {elapsed:.1f}s")
        except Exception as e:
            with self._lock:
                self._pending_error = e

    def save_if_due(self, step, model, optimizer, metadata):
        """Call this every training step. Launches async save when due."""
        self._check_for_errors()
        if step % self.interval != 0:
            return

        # Wait for any previous background save to finish before starting new
        if self._thread is not None and self._thread.is_alive():
            self._thread.join()
        self._check_for_errors()

        # Snapshot — this blocks, but is fast (GPU->CPU copy)
        t0 = time.time()
        model_state, optim_state = self._snapshot_to_cpu(model, optimizer)
        dist.barrier()  # All ranks finish snapshot before any continues
        print(f"[async ckpt] snapshot took {time.time() - t0:.1f}s, "
              f"background write starting...")

        # Launch background write thread
        self._thread = threading.Thread(
            target=self._background_save,
            args=(model_state, optim_state, step, metadata),
            daemon=True,
        )
        self._thread.start()

    def wait(self):
        """Block until any in-flight async checkpoint is written."""
        if self._thread is not None:
            self._thread.join()
        self._check_for_errors()

"""
in_memory_storage.py — Minimal in-memory checkpoint storage backend.

Stores checkpoint bytes in a shared dict on each rank.
Production systems use more sophisticated shared-memory mechanisms
(e.g., mmap, Ray's plasma store).
"""

from torch.distributed.checkpoint.storage import StorageWriter, StorageReader
from torch.distributed.checkpoint.metadata import Metadata, StorageMeta
from io import BytesIO
import io


_IN_MEMORY_STORE: dict[str, bytes] = {}  # In a real system: cross-rank store


class InMemoryWriter(StorageWriter):
    def __init__(self):
        self._buffers: dict[str, bytes] = {}

    def set_up_storage_writer(self, is_coordinator: bool) -> None:
        pass  # no-op for in-memory

    def prepare_local_plan(self, plan):
        return plan

    def prepare_global_plan(self, global_plans):
        return global_plans

    def write_data(self, plan, planner):
        # Write each planned chunk to an in-memory buffer
        futures = []
        for bucket in plan.items:
            data = planner.resolve_data(bucket)
            buf = BytesIO()
            torch.save(data, buf)
            self._buffers[bucket.storage_index.fqn] = buf.getvalue()
        _IN_MEMORY_STORE.update(self._buffers)

    def finish(self, metadata, results):
        _IN_MEMORY_STORE["__metadata__"] = metadata


class InMemoryReader(StorageReader):
    def read_metadata(self) -> Metadata:
        return _IN_MEMORY_STORE["__metadata__"]

    def set_up_storage_reader(self, metadata, is_coordinator):
        pass

    def prepare_local_plan(self, plan):
        return plan

    def prepare_global_plan(self, global_plans):
        return global_plans

    def read_data(self, plan, planner):
        for req in plan.items:
            data = torch.load(BytesIO(_IN_MEMORY_STORE[req.storage_index.fqn]))
            planner.commit_tensor(req, data)

"""
stateful_dataloader.py — A minimal stateful data loader that saves/restores
its position within a sharded dataset.
"""

import json
from pathlib import Path
import torch
from torch.utils.data import DataLoader, IterableDataset


class ShardedTextDataset(IterableDataset):
    """
    Streams tokens from pre-tokenised shard files (.pt tensors).
    Saves and restores its cursor for exact reproducibility.
    """

    def __init__(
        self,
        shard_paths: list[str],
        seq_len: int,
        start_shard: int = 0,
        start_offset: int = 0,
    ):
        self.shard_paths = shard_paths
        self.seq_len = seq_len
        # Restored cursor positions
        self.start_shard = start_shard
        self.start_offset = start_offset

    def get_state(self) -> dict:
        """Call after every batch to get the current cursor state."""
        return {
            "current_shard": self._current_shard,
            "current_offset": self._current_offset,
        }

    def __iter__(self):
        self._current_shard = self.start_shard
        self._current_offset = self.start_offset

        for shard_idx in range(self.start_shard, len(self.shard_paths)):
            tokens = torch.load(self.shard_paths[shard_idx])  # 1D tensor
            start = self.start_offset if shard_idx == self.start_shard else 0
            self._current_shard = shard_idx

            pos = start
            while pos + self.seq_len + 1 <= len(tokens):
                x = tokens[pos : pos + self.seq_len]
                y = tokens[pos + 1 : pos + self.seq_len + 1]
                self._current_offset = pos
                yield x, y
                pos += self.seq_len

            # Move to next shard
            self.start_offset = 0  # only use start_offset for first shard

"""
stateful_loader_usage.py — Resumable data loading with torchdata.
`pip install torchdata`
"""

from torchdata.stateful_dataloader import StatefulDataLoader


class ResumableTokenStream(ShardedTextDataset):
    """Add the Stateful hooks StatefulDataLoader looks for on the dataset."""

    def state_dict(self) -> dict:
        return self.get_state()

    def load_state_dict(self, sd: dict) -> None:
        self.start_shard = sd["current_shard"]
        self.start_offset = sd["current_offset"]


def build_loader(shard_paths, seq_len, batch_size, num_workers=4):
    ds = ResumableTokenStream(shard_paths, seq_len)
    # Identical signature to torch.utils.data.DataLoader.
    return StatefulDataLoader(ds, batch_size=batch_size,
                              num_workers=num_workers)


# At checkpoint time: loader.state_dict() -> goes into the checkpoint.
# At resume time:     loader.load_state_dict(saved) BEFORE iterating.

def compute_rank_start_offset(
    global_token_offset: int,
    world_size: int,
    rank: int,
    tokens_per_rank_per_step: int,
) -> int:
    """
    Given the global number of tokens consumed so far, compute the
    starting offset for this rank in the new topology.

    global_token_offset: total tokens processed before the crash
    world_size: new number of data-parallel ranks
    rank: this rank's index in the new topology
    tokens_per_rank_per_step: seq_len * micro_batch_size
    """
    # Tokens consumed per global step
    tokens_per_step = tokens_per_rank_per_step * world_size
    # Completed steps so far
    steps_done = global_token_offset // tokens_per_step
    # This rank's start offset in the new topology
    return steps_done * tokens_per_rank_per_step + rank * tokens_per_rank_per_step

"""
determinism_setup.py — Configure determinism at various strictness levels.
"""

import torch
import os


def configure_determinism(level: str = "soft") -> None:
    """
    Configure PyTorch determinism.

    level="soft"   — restore RNG, no strict kernel determinism.
                     Fast. Resumable runs track closely but not bit-exactly.
    level="strict" — full determinism. Slower; use for debugging.
    """
    if level == "strict":
        # Force deterministic CUDA kernels
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        torch.use_deterministic_algorithms(True)
        # Required for some deterministic ops on CUDA
        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
        print("Strict determinism enabled (expect ~15% slowdown).")
    elif level == "soft":
        # Allow cuDNN to pick optimal (non-deterministic) algorithms
        torch.backends.cudnn.benchmark = True
        torch.use_deterministic_algorithms(False)
        print("Soft determinism: RNG states saved, kernels non-deterministic.")
    else:
        raise ValueError(f"Unknown level: {level}")

def seed_everything(base_seed: int, rank: int) -> None:
    """
    Set all RNG seeds deterministically for a given rank.
    The per-rank seed is derived from the base seed and rank index,
    ensuring different seeds across ranks but reproducibility from
    a given base seed.
    """
    import random
    import numpy as np

    rank_seed = base_seed + rank * 31337  # prime offset per rank
    random.seed(rank_seed)
    np.random.seed(rank_seed % (2**31))
    torch.manual_seed(rank_seed)
    torch.cuda.manual_seed_all(rank_seed)

import hashlib
from pathlib import Path


def sha256_file(path: Path) -> str:
    """Compute SHA-256 hash of a file."""
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):  # 1 MB chunks
            h.update(chunk)
    return h.hexdigest()


def write_checksums(ckpt_dir: Path) -> None:
    """Write SHA-256 checksums for all checkpoint files."""
    checksums = {}
    for p in ckpt_dir.rglob("*"):
        if p.is_file() and p.name != "checksums.json":
            checksums[str(p.relative_to(ckpt_dir))] = sha256_file(p)
    with open(ckpt_dir / "checksums.json", "w") as f:
        json.dump(checksums, f, indent=2)


def verify_checksums(ckpt_dir: Path) -> bool:
    """Return True if all files match their saved checksums."""
    checksum_file = ckpt_dir / "checksums.json"
    if not checksum_file.exists():
        return False  # No checksums — cannot verify
    with open(checksum_file) as f:
        expected = json.load(f)
    for rel_path, expected_hash in expected.items():
        actual = sha256_file(ckpt_dir / rel_path)
        if actual != expected_hash:
            print(f"CHECKSUM MISMATCH: {rel_path}")
            return False
    return True

# 1. Consolidate the sharded DCP directory into a single torch.save file.
#    (torch.distributed.checkpoint.format_utils also exposes the inverse,
#     torch_save_to_dcp, for importing a legacy single-file checkpoint.)
python -c "
from torch.distributed.checkpoint.format_utils import dcp_to_torch_save
dcp_to_torch_save('ckpt/step_00050000/model_optim', 'ckpt/consolidated.pt')
"

# 2. Strip optimizer state (~89% of the bytes) and re-serialise the weights
#    as safetensors — the format the HF ecosystem loads by default. It is a
#    zero-copy mmap-able layout with no pickle, so loading it cannot execute
#    code, and it is what you publish.
mkdir -p release && python export_safetensors.py

# export_safetensors.py — weights only, in the format the ecosystem expects.
import torch
from safetensors.torch import save_file

sd = torch.load("ckpt/consolidated.pt", map_location="cpu",
                weights_only=True)["model"]
# safetensors requires contiguous, non-shared storage per tensor.
sd = {k: v.contiguous().clone() for k, v in sd.items()}
save_file(sd, "release/model.safetensors",
          metadata={"format": "pt"})   # HF loaders check this key

#!/bin/bash
# watchdog_launch.sh — Auto-restart training on failure.
# Uses torchrun with fault-tolerant options.

MAX_RESTARTS=20
RESTART_COUNT=0
CHECKPOINT_DIR="/mnt/checkpoints/my_run"

while [ $RESTART_COUNT -lt $MAX_RESTARTS ]; do
    echo "Starting training attempt $((RESTART_COUNT + 1))..."

    torchrun \
        --nnodes="${NNODES}" \
        --nproc_per_node=8 \
        --rdzv_backend=c10d \
        --rdzv_endpoint="${MASTER_ADDR}:29500" \
        --max_restarts=0 \
        train.py \
            --checkpoint_dir "${CHECKPOINT_DIR}" \
            --resume_from_latest

    EXIT_CODE=$?
    if [ $EXIT_CODE -eq 0 ]; then
        echo "Training completed successfully."
        break
    else
        echo "Training failed (exit code $EXIT_CODE). Restarting..."
        RESTART_COUNT=$((RESTART_COUNT + 1))
        sleep 30  # Brief pause to allow failed nodes to be replaced
    fi
done

if [ $RESTART_COUNT -eq $MAX_RESTARTS ]; then
    echo "ERROR: Exceeded max restarts. Manual intervention required."
    exit 1
fi

class TrainingHeartbeat:
    """Write a heartbeat file periodically so external monitors can detect
    hung jobs (e.g., deadlocked collectives)."""

    def __init__(self, heartbeat_path: str, interval_steps: int = 10):
        self.path = Path(heartbeat_path)
        self.interval = interval_steps

    def beat(self, step: int) -> None:
        if step % self.interval == 0:
            self.path.write_text(
                json.dumps({"step": step, "ts": time.time()})
            )

# Enable the flight recorder on every rank before launching torchrun.
export TORCH_NCCL_TRACE_BUFFER_SIZE=2000     # entries retained per rank
export TORCH_NCCL_DUMP_ON_TIMEOUT=1          # dump the ring buffer on hang
export TORCH_NCCL_DEBUG_INFO_TEMP_FILE=/mnt/logs/nccl_trace_rank
# Fail fast instead of hanging forever on a dead peer:
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1

"""
train_loop.py — Fault-tolerant pretraining main loop sketch.

Assumes: FSDP model, AdamW optimizer, cosine LR schedule,
         sharded data loader, CheckpointManager, AsyncCheckpointer.
"""

import os
import torch
import torch.distributed as dist
from pathlib import Path
from torchdata.stateful_dataloader import StatefulDataLoader


def main():
    # --- Distributed init ---
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device = torch.device(f"cuda:{rank % 8}")
    torch.cuda.set_device(device)

    # --- Build model, optimizer, scheduler (not shown for brevity) ---
    model = build_fsdp_model(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95))
    scheduler = get_cosine_schedule(optimizer, warmup_steps=2000, total_steps=500_000)

    # --- Checkpoint manager ---
    ckpt_manager = CheckpointManager(
        root_dir="/mnt/checkpoints/my_run",
        keep_last_k=3,
    )
    heartbeat = TrainingHeartbeat(
        "/mnt/checkpoints/my_run/heartbeat.json",
        interval_steps=5,
    )

    # --- Resume from latest checkpoint if available ---
    start_step = 0
    latest_ckpt = ckpt_manager.latest()
    if latest_ckpt is not None:
        if rank == 0:
            print(f"Resuming from checkpoint: {latest_ckpt}")
        metadata = load_checkpoint(model, optimizer, str(latest_ckpt), rank)
        start_step = metadata["step"] + 1
        scheduler.load_state_dict(metadata["lr_scheduler"])
        # Restore data loader cursor
        data_loader_state = metadata.get("data_loader")
    else:
        data_loader_state = None

    # --- Data loader ---
    # StatefulDataLoader (torchdata) rather than DataLoader: with
    # num_workers > 0 only it can aggregate the per-worker cursors and the
    # in-flight prefetched batches. ResumableTokenStream is the Stateful
    # subclass from Section 3.12.6.
    dataset = ResumableTokenStream(shard_paths=get_shard_paths(), seq_len=2048)
    dl = StatefulDataLoader(dataset, batch_size=4, num_workers=4)
    if data_loader_state is not None:
        dl.load_state_dict(data_loader_state)   # must precede iteration
    loader = iter(dl)

    # --- Training loop ---
    model.train()
    for step in range(start_step, 500_000):
        x, y = next(loader)
        x, y = x.to(device), y.to(device)

        optimizer.zero_grad()
        loss = model(x, labels=y).loss
        loss.backward()

        # Gradient clipping — important for training stability
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        optimizer.step()
        scheduler.step()

        # Heartbeat so external watchdog knows we're alive
        heartbeat.beat(step)

        # Periodic logging
        if step % 100 == 0 and rank == 0:
            print(f"step={step:6d}  loss={loss.item():.4f}  "
                  f"lr={scheduler.get_last_lr()[0]:.2e}")

        # Save checkpoint
        if step % 1000 == 0:
            ckpt_manager.save(
                model, optimizer, step, scheduler,
                dl.state_dict(), rank
            )

    # Ensure final async checkpoint write completes
    dist.barrier()
    if rank == 0:
        print("Training complete.")


if __name__ == "__main__":
    main()