Imahn Shekhzadeh

import argparse
import logging
import os
import random
from typing import Generator, List

import numpy as np
import torch
from torch import Tensor
from torch import distributed as dist
from torch.utils.data import DataLoader, DistributedSampler, TensorDataset


def infiniteloop(dataloader) -> Generator[List[Tensor], None, None]:
    while True:
        for data in iter(dataloader):
            yield data


def setup(
    rank: int,
    world_size: int,
    master_addr: str = "localhost",
    master_port: str = "12355",
    backend: str = "nccl",
) -> None:
    """
    Initialize the distributed environment.

    Args:
        rank: Rank of the current process.
        world_size: Number of processes participating in the job.
        master_addr: IP address of the master node.
        master_port: Port number of the master node.
        backend: Backend to use.
    """

    os.environ["MASTER_ADDR"] = master_addr
    os.environ["MASTER_PORT"] = master_port

    # initialize the process group
    dist.init_process_group(
        backend=backend,
        rank=rank,
        world_size=world_size,
    )


def get_args() -> argparse.Namespace:
    """Get arguments passed via CLI."""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--master_addr",
        type=str,
        default=None,
        help="IP address of the master node.",
    )
    parser.add_argument(
        "--master_port",
        type=str,
        default=None,
        help="Port of the master node.",
    )
    return parser.parse_args()


def seed_worker(worker_id: int) -> None:
    """
    Seed the worker for the dataloader. Function copy-pasted from [1].

    Args:
        worker_id: Worker ID.

    References:
        [1] https://pytorch.org/docs/stable/notes/randomness.html#dataloader
    """
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)


def run(
    rank: int | torch.device, world_size: int, args: argparse.Namespace
) -> None:
    """
    Run test.

    Args:
        rank: Rank of the current process. Can be `torch.device("cpu")` if no
            GPU is available.
        world_size: Number of processes participating in distributed training.
            If `world_size` is 1, no distributed training is used.
    """
    if world_size > 1:
        setup(
            rank=rank,
            world_size=world_size,
            master_addr=args.master_addr,
            master_port=args.master_port,
        )

    num_samples = 12
    batch_size = 4
    # When using a single GPU per process and per
    # DistributedDataParallel, we need to divide the batch size
    # ourselves based on the total number of GPUs of the current node.
    batch_size = int(batch_size / world_size)
    num_epochs = 2

    tensor = torch.randn(
        num_samples, 2, generator=torch.Generator().manual_seed(2)
    )
    dataset = TensorDataset(tensor)
    sampler = DistributedSampler(dataset) if world_size > 1 else None
    dataloader = DataLoader(
        dataset=dataset,
        sampler=sampler,
        shuffle=False if world_size > 1 else True,
        worker_init_fn=seed_worker,
        generator=torch.Generator().manual_seed(0),
        batch_size=batch_size,
        num_workers=4,
    )
    datalooper = infiniteloop(dataloader)
    num_batches__per_epoch = len(dataloader)
    logging.info(f"# Batches/epoch: {num_batches__per_epoch}")

    for epoch in range(num_epochs):
        if sampler is not None:
            # necessary to ensure shuffling of the data
            # https://pytorch.org/docs/stable/data.html
            sampler.set_epoch(epoch)

        for batch_idx in range(num_batches__per_epoch):
            data = next(datalooper)
            logging.info(
                f"\n\nRank: {rank}, Epoch: {epoch}, Batch: {batch_idx}, Data:\n{data}"
            )


if __name__ == "__main__":
    logging.basicConfig(
        level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s"
    )

    args = get_args()
    world_size = int(os.getenv("WORLD_SIZE", 1))
    logging.info(
        f"{args}\nWorld_size: {world_size}\nPyTorch version: "
        f"{torch.__version__}"
    )

    run(rank=int(os.getenv("RANK", 0)), world_size=world_size, args=args)