ezpz.examples.test⚓︎

@dataclass
class TrainConfig:
    """Runtime configuration for the `ezpz.examples.test` distributed smoke test."""

    warmup: int
    tp: int
    pp: int
    cp: int
    batch_size: int
    input_size: int
    output_size: int
    train_iters: int
    log_freq: int
    backend: str
    dtype: str
    print_freq: int
    pyinstrument_profiler: bool
    pytorch_profiler: bool
    pytorch_profiler_wait: int
    pytorch_profiler_warmup: int
    pytorch_profiler_active: int
    pytorch_profiler_repeat: int
    profile_memory: bool
    rank_zero_only: bool
    record_shapes: bool
    with_stack: bool
    with_flops: bool
    with_modules: bool
    acc_events: bool
    layer_sizes: list = field(default_factory=lambda: [512, 256, 128])
    dataset: str = "mnist"
    dataset_root: Optional[PathLike] = None
    num_workers: int = 0
    no_distributed_history: bool = False
    save_datasets: bool = False

    def __post_init__(self):
        """Initialise output paths and configure profiling context managers."""
        self._created_at = (
            ezpz.get_timestamp() if ezpz.get_rank() == 0 else None
        )
        self._created_at = ezpz.dist.broadcast(self._created_at, root=0)
        self.outdir = Path(os.getcwd()).joinpath(
            "outputs", "ezpz.examples.test", f"{self._created_at}"
        )
        self.outdir.mkdir(parents=True, exist_ok=True)
        dataset_root = (
            Path(self.dataset_root).expanduser()
            if self.dataset_root is not None
            else self.outdir.parent.joinpath("datasets", self.dataset)
        )
        dataset_root.mkdir(parents=True, exist_ok=True)
        self.dataset_root = dataset_root
        profiler_type = "torch" if self.pytorch_profiler else "pyinstrument"
        self.ctx = get_profiling_context(
            profiler_type=profiler_type,
            rank_zero_only=self.rank_zero_only,
            record_shapes=self.record_shapes,
            with_stack=self.with_stack,
            with_flops=self.with_flops,
            with_modules=self.with_modules,
            acc_events=self.acc_events,
            profile_memory=self.profile_memory,
            wait=self.pytorch_profiler_wait,
            warmup=self.pytorch_profiler_warmup,
            active=self.pytorch_profiler_active,
            repeat=self.pytorch_profiler_repeat,
            outdir=self.outdir,
        )
        logger.info(f"Outputs will be saved to {self.outdir}")

    def get_torch_dtype(self) -> torch.dtype:
        """Return the torch dtype requested by this configuration."""
        if self.dtype is None:
            return torch.get_default_dtype()
        if self.dtype in {
            "fp16",
            "half",
            "float16",
        }:
            return torch.float16
        if self.dtype in {
            "bfloat16",
            "bf16",
        }:
            return torch.bfloat16
        if self.dtype in {
            "float32",
            "fp32",
            "float",
            "single",
        }:
            return torch.float32
        logger.warning(f"Unknown dtype: {self.dtype=}, using float32")
        return torch.float32

@dataclass
class Trainer:
    """Co-ordinate training loops, logging, and profiling for the test model."""

    config: TrainConfig
    model: torch.nn.Module
    optimizer: torch.optim.Optimizer
    # history: ezpz.History = field(init=False)
    train_iter: int = 0
    rank: int = ezpz.get_rank()
    # device_type: str = ezpz.get_torch_device_type()
    device_type = os.environ.get("TORCH_DEVICE", ezpz.get_torch_device())
    world_size = ezpz.get_world_size()
    local_rank = ezpz.get_local_rank()
    device_id = f"{device_type}:{local_rank}"
    _train_loader: Optional[DataLoader] = field(init=False, default=None)
    _train_iterator: Optional[Iterator[tuple[torch.Tensor, torch.Tensor]]] = (
        field(init=False, default=None)
    )
    _feature_dim: int = field(init=False, default=0)

    def __post_init__(self):
        """Move the model to the target device and register logging hooks."""
        self.device_id = f"{self.device_type}:{self.local_rank}"
        self.dtype = self.config.get_torch_dtype()
        self.model.to(self.device_id)
        self.model.to(self.dtype)
        metrics_path = self.config.outdir.joinpath("metrics.jsonl")
        self.history: ezpz.history.History = ezpz.history.History(
            report_dir=self.config.outdir,
            report_enabled=True,
            jsonl_path=metrics_path,
            jsonl_overwrite=True,
            distributed_history=(
                1 < self.world_size <= 384 and not self.config.pytorch_profiler
            ),
        )

        if self.config.tp > 1 or self.config.pp > 1 or self.config.cp > 1:
            ezpz.dist.barrier(group=ezpz.tp.get_tensor_parallel_group())
            ezpz.dist.barrier(group=ezpz.tp.get_data_parallel_group())
            ezpz.dist.barrier(group=ezpz.tp.get_pipeline_parallel_group())
            ezpz.dist.barrier(group=ezpz.tp.get_context_parallel_group())

        if self.rank == 0 and not WANDB_DISABLED:
            logger.debug("Setting up wandb")
            wbconfig = {}
            wbconfig |= asdict(self.config)
            wbconfig |= ezpz.get_dist_info()
            _ = ezpz.setup_wandb(
                project_name="ezpz.examples.test",
                config=wbconfig,
            )
            if (wbrun := getattr(wandb, "run", None)) is not None and callable(
                wbrun.watch
            ):
                wbrun.watch(self.model, log="all")

        if self.world_size > 1:
            logger.debug("Hit torch.distributed.barrier()")
            ezpz.dist.barrier()
        self._train_loader, self._train_iterator = self._build_dataloader()
        self._feature_dim = self.config.input_size

    @ezpz.timeitlogit(rank=ezpz.get_rank())
    def _build_dataloader(
        self,
    ) -> tuple[DataLoader, Iterator[tuple[torch.Tensor, torch.Tensor]]]:
        """Construct a training dataloader for the requested dataset."""
        dataset_name = self.config.dataset.lower()
        if dataset_name == "mnist":
            try:
                from ezpz.data.vision import get_mnist
            except (
                ModuleNotFoundError
            ) as exc:  # pragma: no cover - optional dep
                msg = (
                    "torchvision is required to use the MNIST dataset. "
                    "Install it via `pip install torchvision`."
                )
                raise RuntimeError(msg) from exc
            assert self.config.dataset_root is not None
            dset_root = Path(self.config.dataset_root).expanduser().resolve()
            dset_root.mkdir(parents=True, exist_ok=True)
            bundle = get_mnist(
                train_batch_size=self.config.batch_size,
                test_batch_size=self.config.batch_size,
                outdir=dset_root,
                num_workers=self.config.num_workers,
                download=self.rank == 0,
                pin_memory=str(self.device_type).startswith(("cuda", "mps")),
            )
            train_loader = bundle["train"]["loader"]
            return train_loader, iter(train_loader)
        raise ValueError(f"Unknown dataset: {dataset_name!r}")

    @ezpz.timeitlogit(rank=ezpz.get_rank())
    def _next_batch(self) -> tuple[torch.Tensor, torch.Tensor]:
        """Return the next batch from the training dataloader."""
        assert self._train_loader is not None
        assert self._train_iterator is not None
        try:
            return next(self._train_iterator)
        except StopIteration:
            self._train_iterator = iter(self._train_loader)
            return next(self._train_iterator)

    @ezpz.timeitlogit(rank=ezpz.get_rank())
    def _prepare_inputs(self, inputs: torch.Tensor) -> torch.Tensor:
        """Move inputs to the configured device and coerce feature dimensions."""
        inputs = inputs.to(self.device_id)
        inputs = inputs.reshape(inputs.size(0), -1)
        if inputs.size(1) < self._feature_dim:
            pad = self._feature_dim - inputs.size(1)
            inputs = torch.nn.functional.pad(inputs, (0, pad))
        elif inputs.size(1) > self._feature_dim:
            inputs = inputs[:, : self._feature_dim]
        return inputs.to(self.dtype)

    @ezpz.timeitlogit(rank=ezpz.get_rank())
    def _forward_step(self) -> tuple[dict, torch.Tensor]:
        """Execute a forward pass returning metrics and the loss tensor."""
        t0 = time.perf_counter()
        batch_inputs, targets = self._next_batch()
        inputs = self._prepare_inputs(batch_inputs)
        targets = targets.to(self.device_id)
        logits = self.model(inputs)
        loss = calc_loss(logits, targets)
        accuracy = (logits.argmax(dim=1) == targets).float().mean()
        ezpz.dist.synchronize()
        metrics = {
            "loss": loss.detach(),
            "accuracy": accuracy.detach(),
            "dtf": time.perf_counter() - t0,
        }
        return metrics, loss

    @ezpz.timeitlogit(rank=ezpz.get_rank())
    def _backward_step(self, loss: torch.Tensor) -> float:
        """Perform the backwards/optimiser step and return elapsed seconds."""
        t0 = time.perf_counter()
        if self.config.backend == "deepspeed":
            self.model.backward(loss)
            self.model.step(loss)
        else:
            loss.backward()
            self.optimizer.step()
        ezpz.dist.synchronize()
        return time.perf_counter() - t0

    @ezpz.timeitlogit(rank=ezpz.get_rank())
    def train_step(self) -> dict:
        """Run one optimiser step, emitting periodic logs/metrics."""
        self.train_iter += 1
        metrics, loss = self._forward_step()
        metrics["dtb"] = self._backward_step(loss)
        self.optimizer.zero_grad()
        if self.train_iter == self.config.train_iters:
            return metrics
        if (
            self.train_iter % self.config.log_freq == 0
            or self.train_iter % self.config.print_freq == 0
        ):
            summary = self.history.update({"iter": self.train_iter, **metrics})
            if self.train_iter % self.config.print_freq == 0:
                logger.info(f"{summary}")
        return metrics

    @ezpz.timeitlogit(rank=ezpz.get_rank())
    def finalize(
        self, outdir: Optional[str | Path | os.PathLike] = None
    ) -> Dataset:
        """Flush profilers and return the aggregated training dataset."""
        import ambivalent
        import matplotlib.pyplot as plt

        plt.style.use(ambivalent.STYLES["ambivalent"])
        outdir = Path(outdir) if outdir is not None else self.config.outdir
        env_info = self._gather_environment_snapshot()
        dataset = self.history.finalize(
            run_name="ezpz.examples.test",
            dataset_fname="train",
            warmup=self.config.warmup,
            save=False,  # XXX: don't bother saving test data
            plot=(self.rank == 0),
            outdir=outdir,
            env_info=env_info,
        )
        logger.info(f"{dataset=}")
        if wandb is not None and ezpz.verify_wandb() and not WANDB_DISABLED:
            try:
                wandb.log(
                    {
                        "train_metrics": wandb.Table(
                            dataframe=dataset.to_dataframe()
                        )
                    }
                )
            except Exception:
                logger.warning("Failed to log final dataset to wandb")
        return dataset

    @ezpz.timeitlogit(rank=ezpz.get_rank())
    def train(self, profiler: Optional[torch.profiler.profile] = None) -> None:
        """Loop over all training iterations and return the final dataset."""
        for step in range(self.config.train_iters):
            if step == self.config.warmup:
                logger.info(f"Warmup complete at step {step}")
            _ = self.train_step()
            if profiler is not None:
                profiler.step()

    def _gather_environment_snapshot(self) -> dict[str, dict[str, str]]:
        """Collect key runtime environment details for reporting."""

        python_details = {
            "Version": (
                f"{sys.version_info.major}."
                f"{sys.version_info.minor}."
                f"{sys.version_info.micro}"
            ),
            "Implementation": sys.implementation.name,
            "Executable": sys.executable,
        }

        torch_details = {
            "Version": torch.__version__,
            "Device": str(self.device_id),
            "Backend": (
                ezpz.dist.get_backend()
                if hasattr(ezpz.dist, "get_backend")
                else "unknown"
            ),
        }

        host_name = (
            platform.uname().node
            if hasattr(platform, "uname")
            else os.environ.get("HOSTNAME", "unknown")
        )
        path_details = {
            "Working directory": str(Path.cwd()),
            "Output directory": str(self.config.outdir),
            "Dataset root": str(self.config.dataset_root),
            "Hostname": host_name,
        }

        dist_details = {
            "Rank": str(self.rank),
            "Local rank": str(self.local_rank),
            "World size": str(self.world_size),
        }

        env_vars: dict[str, str] = {}
        for key in (
            "MASTER_ADDR",
            "MASTER_PORT",
            "NODE_RANK",
            "LOCAL_RANK",
            "RANK",
            "WORLD_SIZE",
        ):
            value = os.environ.get(key)
            if value is not None:
                env_vars[key] = value

        snapshot: dict[str, dict[str, str]] = {
            "Paths": path_details,
            "Python": python_details,
            "Torch": torch_details,
            "Distributed": dist_details,
        }
        if env_vars:
            snapshot["Environment Variables"] = env_vars
        return snapshot