distai 0.0.1.dev0

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

distai

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

Install

From this directory (shared/):

pip install .

For local development:

pip install -e .[dev]

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

Build And Publish

1. Update version in distai/__about__.py.
2. Build distributions:

python -m build

3. Validate metadata:

twine check dist/*

4. Upload to TestPyPI:

twine upload --repository testpypi dist/*

5. Upload to PyPI:

twine upload dist/*