Metadata-Version: 2.4
Name: rayllm-orchestrator
Version: 0.1.2
Summary: Train and serve LLMs on any hardware with one command each
Author-email: Agastya Choudhary <agastya.r.choudhary@gmail.com>
License: MIT
Project-URL: Homepage, https://github.com/agastya-choudhary123/rayllm-orchestrator
Project-URL: Repository, https://github.com/agastya-choudhary123/rayllm-orchestrator
Requires-Python: >=3.9
Description-Content-Type: text/markdown
Requires-Dist: torch>=2.4
Requires-Dist: transformers>=4.44
Requires-Dist: datasets>=2.20
Requires-Dist: peft>=0.12
Requires-Dist: accelerate>=0.33
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Provides-Extra: gpu
Requires-Dist: ray[default]>=2.35; extra == "gpu"
Requires-Dist: vllm>=0.6.0; extra == "gpu"
Requires-Dist: bitsandbytes>=0.43; extra == "gpu"
Provides-Extra: mlx
Requires-Dist: mlx>=0.18; extra == "mlx"

# RayLLM-Orchestrator

**Fine-tune an LLM on your data and serve it behind an OpenAI-compatible API — one command each.**

```bash
# Train
rayllm train --model gpt2 --dataset my-data.jsonl --epochs 3

# Serve (in another terminal)
rayllm serve --model ./checkpoints/gpt2 --port 8000

# Test it
curl -X POST http://localhost:8000/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{"messages":[{"role":"user","content":"Hello"}]}'
```

Both `rayllm <cmd>` (installed console script) and `python orchestrator.py <cmd>`
(from a source checkout) work identically.

---

## What It Does

**Train** — real PyTorch fine-tuning with:
- Model loading from HuggingFace IDs or local paths
- Dataset tokenization (JSONL or HuggingFace datasets)
- LoRA / QLoRA (4-bit / 8-bit base weights)
- Checkpointing and resume from a saved epoch
- Optional multi-GPU via Ray + PyTorch FSDP (needs >1 CUDA GPU)
- On Apple Silicon, an MLX-LoRA path when `mlx-lm` is installed and the model
  is MLX-loadable; otherwise it falls back to the PyTorch path

**Serve** — one OpenAI-compatible HTTP API over several backends. You pick the
backend (a recommended default is offered); it is not chosen silently. See
[Choosing a backend](#choosing-a-backend).

---

## Installation

```bash
pip install rayllm-orchestrator
```

This installs the core stack: `torch`, `transformers`, `datasets`, `peft`,
`accelerate`. The serving backends below are **optional** and installed
separately — you only need the one(s) you want:

| Backend | Extra install | Where it runs |
|---------|---------------|---------------|
| transformers | (included) | CPU / CUDA / Apple MPS — works everywhere |
| MLX | `pip install mlx mlx-lm` | Apple Silicon (arm64 macOS) |
| llama.cpp | `pip install llama-cpp-python` | CPU, Metal on macOS — GGUF models |
| vLLM | `pip install vllm` | NVIDIA GPU (Linux) |
| Ollama | install [Ollama](https://ollama.com) + `pip install ollama` | uses the local Ollama daemon |

### From source
```bash
git clone https://github.com/agastya-choudhary123/rayllm-orchestrator
cd rayllm-orchestrator
pip install -e .
python orchestrator.py train --model gpt2 --dataset my-data.jsonl
```

### Docker
The repo ships a `Dockerfile` (and `Dockerfile.cpu`) you build locally — there is
no prebuilt image published to a registry:
```bash
docker build -t rayllm .
docker run -it rayllm train --model gpt2 --dataset my-data.jsonl --epochs 3
```

---

## Quick Start

### 1. Prepare your data

JSONL, one object per line. Supported shapes: `{"text": ...}`,
`{"prompt": ..., "completion": ...}`, or `{"messages": [...]}`.
```json
{"text": "Your training example here."}
{"text": "Another example."}
```

Or point at a HuggingFace dataset id:
```bash
rayllm train --model gpt2 --dataset wikitext --epochs 1
```

### 2. Train
```bash
rayllm train --model gpt2 --dataset my-data.jsonl --epochs 3 --quant 4bit
```
Downloads the model, tokenizes your data, trains (saving a checkpoint after each
epoch), and prints the checkpoint path. With `--quant`, base weights load in
4/8-bit and only LoRA adapters train (QLoRA); the PyTorch path merges the
adapters into a plain HuggingFace checkpoint the serving layer loads directly.

### 3. Serve
```bash
rayllm serve --model ./checkpoints/gpt2 --port 8000
```
Prompts you to choose a backend from what's installed (or pass `--backend`),
then starts an OpenAI-compatible API on port 8000.

### 4. Use it
```bash
curl -X POST http://localhost:8000/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{"messages":[{"role":"user","content":"Write a haiku about coding"}],"max_tokens":50}'
```

---

## Choosing a backend

List what's actually usable on your machine:
```bash
rayllm backends
```
Example output — availability is probed live (installed **and** runnable here),
and a format-aware recommendation is marked:
```
1. [✗] ollama         needs the ollama daemon running
2. [✓] mlx            Apple-Silicon-native            <- recommended for this model
3. [✗] llama_cpp      `pip install llama-cpp-python`
4. [✗] vllm           installed but no CUDA GPU found
5. [✓] transformers   universal fallback
```

Then either let `serve`/`run` prompt you (it offers the recommended one as the
Enter-default), or force a choice:
```bash
rayllm serve --model ./checkpoints/gpt2 --backend transformers
```

---

## Common Use Cases

### Apple Silicon (MLX)
```bash
pip install mlx mlx-lm
rayllm serve --model mlx-community/Qwen2.5-0.5B-Instruct-4bit --backend mlx
```
MLX is Apple's Metal-native framework and keeps 4-bit models memory-safe in
unified memory (a guard caps its cache so training can't starve the OS). See
[measured numbers](#measured-performance).

### Fine-tune with quantization (QLoRA)
```bash
rayllm train --model meta-llama/Llama-3.2-3B --dataset my-data.jsonl --quant 4bit
```
Loads base weights in 4-bit and trains only the LoRA adapters — far fewer
trainable parameters and lower memory than full fine-tuning.

### Multi-GPU training (Ray + FSDP)
```bash
rayllm train --model meta-llama/Llama-3.2-3B --dataset my-data.jsonl \
  --strategy fsdp-ray --num-workers 4
```
Requires Ray installed and more than one CUDA GPU; otherwise it logs a note and
runs the single-device loop. (`deepspeed-ray` is accepted but currently runs the
same Ray+FSDP path — there is no separate DeepSpeed backend.)

### Speculative decoding (MLX)
```bash
rayllm serve --model <large-mlx-model> --draft-model <small-mlx-model> --backend mlx
```
A small draft model proposes tokens the large model verifies — faster generation
with identical output. MLX backend only.

---

## Command Reference

### `train`
```
rayllm train --model <id-or-path> --dataset <id-or-path.jsonl>
  [--epochs N] [--strategy single|fsdp-ray|deepspeed-ray]
  [--quant none|8bit|4bit] [--out DIR] [--num-workers N]
```
Defaults: `--epochs 1`, `--strategy fsdp-ray` (falls back to single without Ray +
multiple GPUs), `--quant none`, `--out ./checkpoints`, `--num-workers 0` (auto).

### `serve`
```
rayllm serve --model <checkpoint-or-id>
  [--backend ollama|mlx|llama_cpp|vllm|transformers]
  [--quant ...] [--port 8000] [--host 0.0.0.0]
  [--max-model-len 4096] [--draft-model <id>]
```
If `--backend` is omitted you're prompted to choose (interactive terminals only).

### `run`
```
rayllm run --model <id-or-path> --data <id-or-path.jsonl>
  [--epochs N] [--port 8000] [--backend ...] [--no-serve] [--no-fast]
```
Fine-tune, then serve in one shot.

### `backends`
```
rayllm backends
```
Scan this machine and list which serving backends are available.

---

## Hardware Support

| Hardware | Train | Serve | Notes |
|----------|-------|-------|-------|
| CPU (any OS) | ✓ (torch) | ✓ (transformers / llama.cpp) | Works everywhere; slow on large models |
| Apple Silicon | ✓ (MLX or torch) | ✓ (MLX / transformers) | MLX needs `mlx-lm`; native Metal GPU |
| NVIDIA GPU (Linux) | ✓ (FSDP w/ Ray) | ✓ (vLLM / transformers) | vLLM/Ray installed separately |
| Multi-GPU | ✓ (Ray + FSDP) | ✓ (vLLM) | Needs >1 GPU |

The vLLM and multi-GPU/FSDP paths require real NVIDIA hardware and are not
exercised in this project's own testing — treat them as supported-but-unverified
here.

---

## Measured Performance

Real serving throughput measured with the repo's `benchmark_serving` helper
(`orchestrator/backends_mlx.py`) on **one Apple Silicon Mac (~17 GB unified
memory), MLX backend, 4-bit weights, 120-token generations**:

| Model | Throughput | Peak memory |
|-------|-----------:|------------:|
| Qwen2.5-0.5B-Instruct-4bit | ~276 tok/s | 0.33 GB |
| Qwen3-0.6B-4bit | ~240 tok/s | 0.44 GB |
| Qwen3-8B-4bit | ~22 tok/s | 4.72 GB |

These are single-machine snapshots, not guarantees — throughput and memory
depend heavily on model, quantization, prompt length, and hardware. Benchmark on
your own setup before relying on any number. No cross-backend speed comparisons
(e.g. MLX vs. PyTorch MPS) are claimed here because none have been measured.

---

## Troubleshooting

**`No module named 'torch'`** — install the core stack: `pip install rayllm-orchestrator` (or `pip install torch transformers datasets peft accelerate`).

**Out of memory during training** — use quantization: add `--quant 4bit`.

**"No serving backend is available"** — run `rayllm backends` to see what's
installed, then install one (e.g. `pip install mlx-lm` on a Mac) or pass
`--backend transformers`.

**Serve doesn't respond** — check it's up: `curl http://localhost:8000/v1/models`.

---

## Architecture

```
orchestrator.py            ← thin shim (python orchestrator.py ...)
orchestrator/
├── cli.py                ← the `rayllm` command (train / serve / run / backends)
├── train.py              ← training loop + MLX/torch dispatch, fallback
├── serve.py              ← serving + backend dispatch (OpenAI-compatible API)
├── models.py             ← model loading + backend detection/selection
├── data.py               ← tokenization, packing/bucketing
├── fast.py               ← bf16, prefetch, packing, LoRA optimizations
├── backends_mlx.py       ← Apple Silicon native train + serve (MLX)
└── util.py               ← logging, capability detection
```

Core dependencies: `torch`, `transformers`, `datasets`, `peft`, `accelerate`.
Backends (MLX, vLLM, llama.cpp, Ollama) and Ray are optional and installed as
needed.

---

## License

MIT. Contributions welcome — see [CONTRIBUTING.md](CONTRIBUTING.md).

## See Also

- [vLLM](https://github.com/vllm-project/vllm) — GPU serving with continuous batching
- [MLX](https://github.com/ml-explore/mlx) — Apple Silicon ML framework
- [Ray Train](https://docs.ray.io/en/latest/train/train.html) — distributed training
- [Hugging Face](https://huggingface.co) — models & datasets
