jazari 0.1.0

The orchestration layer for modern Slurm clusters.

🐘 Jazari

The orchestration layer for modern Slurm clusters.

Jazari is a command-line tool that makes launching distributed AI/ML training jobs on high-performance computing (HPC) clusters as easy as running a script on your laptop.

It abstracts away the complexity of writing Slurm (#SBATCH) scripts, handling multi-node networking, and managing environment variables.

---

Why Jazari?

If you are a researcher using a university or national supercomputer (like Compute Canada/Digital Alliance), you know the pain:

• The Script Nightmare: Copy-pasting old bash scripts, accidentally leaving in wrong parameters, and debugging obscure sbatch errors.
• Networking Headaches: Manually figuring out how to get PyTorch DDP to find the master node's IP address across multiple machines.
• Experiment Fragmentation: Having a 4-node training run spawn 4 separate experiments in Weights & Biases.
• Slurm Arcana: Remembering obscure flags for accounts, time formats, and memory allocation.

Jazari solves this. You write your Python training script, and Jazari handles the rest.

---

✨ Key Features

• 🚀 One-Command Launch: Go from Python script to running distributed job with a single CLI command.
• 🤖 Automatic Slurm Generation: Uses robust Jinja2 templates to generate correct, safe #SBATCH scripts on the fly.
• 🧠 Zero-Config DDP Networking: Automatically resolves master IP, ports, and world size for PyTorch Distributed Data Parallel.
• 📈 Seamless Weights & Biases: Securely propagates your local API key and ensures multi-node jobs log as a single, clean run.
• ⚙️ User "Init" Profiles: Save your default account, time limits, and preferences once and never type them again.
• 🎨 Beautiful CLI: Modern, color-coded output with rich error reporting.

---

🛠️ Installation

Jazari is designed to be installed in a virtual environment on your cluster's login node.

# 1. Clone the repository
git clone [https://github.com/levoz92/jazari.git](https://github.com/levoz92/jazari.git)
cd jazari

# 2. Create and activate a virtual environment (recommended)
python3 -m venv venv
source venv/bin/activate

# 3. Install in editable mode
pip install -e .

Verify the installation:

jazari --help

---

⚡ Quick Start

1. One-Time Setup

On your cluster login node, run the init command to save your defaults (like your allocation account).

jazari init

2. Create a Python Script

Here is a minimal PyTorch DDP example (train.py):

import os
import torch
import torch.distributed as dist

def main():
    # 1. Initialize Process Group (Jazari sets all necessary env vars)
    dist.init_process_group(backend="nccl")
    
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    print(f"👋 Hello from rank {rank} of {world_size} on node {os.uname().nodename}!")

    # Your training loop here...

    # 3. Clean up
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

3. Launch It!

Run your script on 2 nodes with 4 GPUs per node (8 GPUs total).

jazari run -N 2 -G 4 --name "my-big-run" --track-wandb python train.py --batch-size 128

That's it. Jazari will generate the script, submit it to Slurm, and stream the output.

---

📖 Usage Reference

jazari init

Interactively configure your default settings. These are saved to ~/.jazari/config.yaml.

jazari run [OPTIONS] COMMAND

Option	Shorthand	Description	Default
--nodes	-N	Number of compute nodes to request.	1 (or config default)
--gpus	-G	Number of GPUs per node.	1 (or config default)
--cpus	-c	Number of CPU cores per task.	1 (or config default)
--time	-t	Time limit in D-HH:MM format (e.g., 0-02:30 for 2.5 hours).	01:00:00 (or config default)
--account	-A	Slurm account to charge (e.g., def-user).	Config default
--name	-n	Name for the job in Slurm and W&B.	jazari_run
--track-wandb		Auto-configure Weights & Biases environment.	False (or config default)
--dry-run		Print the generated generated #SBATCH script to stdout without submitting.	False

---

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.