Metadata-Version: 2.4
Name: distai
Version: 0.0.1.dev1
Summary: Distributed PyTorch helpers for process setup, sync printing, and model utilities
Author-email: Fuheng Wu <wufuheng@gmail.com>
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/yourusername/distai
Project-URL: Repository, https://github.com/yourusername/distai
Project-URL: Issues, https://github.com/yourusername/distai/issues
Keywords: pytorch,distributed,ddp,torchrun,machine-learning
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0.0
Requires-Dist: torchvision>=0.15.0
Provides-Extra: dev
Requires-Dist: black>=23.0.0; extra == "dev"
Requires-Dist: build>=1.2.0; extra == "dev"
Requires-Dist: flake8>=6.0.0; extra == "dev"
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: twine>=5.0.0; extra == "dev"
Dynamic: license-file

# distai

Package for the book [Distributed AI Systems: A practical guide to building scalable training, inference, and serving systems for production AI](https://www.amazon.com/dp/1807301710/).

![Distributed AI Systems cover](https://m.media-amazon.com/images/I/41sd7YI23nL.jpg)

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

## Install

Install from PyPI:

```bash
pip install distai
```

## Quick Usage

### With `torchrun` (recommended)

```bash
torchrun --nproc_per_node=4 your_script.py
```

```python
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`

```python
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`


