distai 0.0.1.dev1

Distributed PyTorch helpers for process setup, sync printing, and model utilities

distai

Package for the book Distributed AI Systems: A practical guide to building scalable training, inference, and serving systems for production AI.

Distributed PyTorch helpers for process setup, synchronized logging, and small model utilities.

Install

Install from PyPI:

pip install distai

Quick Usage

With torchrun (recommended)

torchrun --nproc_per_node=4 your_script.py

from distai import init_distributed, sync_print

rank, world_size, device, local_rank = init_distributed(use_cpu=False)
sync_print(f"Rank {rank} says hello", rank=rank, world_size=world_size)

Without torchrun

from distai import run_distributed


def my_worker(rank, world_size, device, local_rank):
    print(f"Rank {rank} running on {device}")


run_distributed(my_worker, world_size=4, use_cpu=False)

See example_no_torchrun.py for a complete runnable example.

Public API

• init_distributed
• sync_print
• run_distributed
• get_node_info
• setup_distributed
• cleanup_distributed
• get_resnet18_fashionmnist
• get_resnet18_cifar10