Metadata-Version: 2.4
Name: trainq
Version: 0.1.0
Summary: Redis-based GPU job queue for ML training/eval: VRAM-aware scheduling, priority aging, backfill.
Author: Korshort
License: MIT
Project-URL: Homepage, https://github.com/Korshort/trainq
Project-URL: Repository, https://github.com/Korshort/trainq
Project-URL: Issues, https://github.com/Korshort/trainq/issues
Keywords: gpu,job-queue,scheduler,machine-learning,redis
Classifier: Development Status :: 4 - Beta
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
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 :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: System :: Distributed Computing
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: redis>=5.0
Provides-Extra: hf
Requires-Dist: huggingface_hub>=0.20; extra == "hf"
Provides-Extra: dev
Requires-Dist: pytest>=7; extra == "dev"
Dynamic: license-file

# trainq — a lightweight GPU job queue

Submit training/eval jobs and `trainq` runs them as soon as the GPUs are free.
Redis-backed, framework-agnostic, single-file-simple. Your command can be
anything: PyTorch, JAX, `torchrun`, or a plain shell script.

```
trainq submit --command "python train.py" --ckpt-dir /data/exp1 \
  --log-path /data/exp1.log --gpu-count 1 --description "my first job"
```

**Why trainq?** On a shared GPU box, people step on each other, big jobs starve
small ones, and you never know when your run will actually start. trainq gives you
a queue that packs jobs onto GPUs by VRAM, fills idle gaps with short jobs
(backfill), ages up long-waiting jobs so nothing starves, and shows a live ETA
for everything in flight.

## Features

- **VRAM-aware sharing** — co-locate several jobs on one GPU; reservations prevent the startup-time OOM race.
- **Priority + aging** — explicit priorities, plus automatic promotion the longer a job waits (no starvation).
- **Backfill** — while a big job waits for GPUs, short jobs that finish before those GPUs free up run first.
- **Live ETA** — progress is parsed from your logs (several formats auto-detected) or a tiny universal convention.
- **Checkpoint hooks (optional)** — run any command and/or auto-upload to the Hugging Face Hub on each new checkpoint.
- **Crash-tolerant daemon** — if the daemon dies, running jobs keep going; on restart it re-tracks them by PID.

## Requirements

- Python 3.9+
- Redis (`sudo apt-get install -y redis-server`, or `brew install redis`)
- `nvidia-smi` for GPU scheduling (only the queue daemon needs it; the checkpoint watcher does not)

## Quick start

```bash
# 1) install  — from PyPI
pip install trainq
#    …or from source:
#    git clone https://github.com/Korshort/trainq && cd trainq && pip install -e .

# 2) start the queue daemon on your GPU box
python -m trainq.server >> trainq_server.log 2>&1 &

# 3) submit a job
trainq submit --command "python train.py --lr 1e-4" \
  --ckpt-dir /data/ckpt/exp1 --log-path /data/logs/exp1.log \
  --gpu-count 1 --estimated-vram-mb 15000 \
  --description "resnet50 finetune lr1e-4"

# 4) watch it
trainq list          # every job, with status, GPU, and live ETA
```

That's the whole loop. Everything below is optional depth.

## Try it in 2 minutes (no GPU needed)

You can see the checkpoint-hook system end to end on a laptop, no GPU or model
required. [`examples/dummy_train.py`](examples/dummy_train.py) is a fake trainer
that just writes progress and drops checkpoint folders on a timer.

```bash
# Terminal A — a "training" run that saves a checkpoint every 5 steps
python examples/dummy_train.py --ckpt-dir /tmp/demo --steps 20 --save-every 5 --sleep 1

# Terminal B — watch that directory and run a hook on each new checkpoint
python -m trainq.checkpoint_watcher --ckpt-dir /tmp/demo \
  --on-checkpoint 'echo ">> new checkpoint: {name} (step {step})"'
```

You'll see the hook fire for `checkpoint-5`, `-10`, `-15`, `-20`. See
[examples/](examples/) for the full walkthrough, including the GPU queue flow.

## Contributing

Contributions are welcome, and the no-GPU example above is the fastest way to
see the moving parts. Good first areas: more log-format parsers in
`trainq/progress.py`, more recipes in `examples/`, and tests. Please keep the
core dependency-light (`redis` only; `huggingface_hub` stays an optional extra).
Open an issue for bugs or ideas, or a PR for changes — see
[CONTRIBUTING.md](CONTRIBUTING.md) for the dev setup and how to run tests.

## Cookbook

Common recipes. Mix and match the flags.

**Pack two jobs onto one GPU** — give each job its VRAM estimate; trainq shares the GPU when it fits.
```bash
trainq submit --command "python train.py"  --ckpt-dir /d/a --log-path /d/a.log --estimated-vram-mb 18000 --description a
trainq submit --command "python eval.py"   --ckpt-dir /d/b --log-path /d/b.log --estimated-vram-mb 12000 --description b
```

**Restrict a job to certain GPUs** — `--gpu-ids` is an *allow-set*, not a pin; trainq picks `--gpu-count` from it.
```bash
trainq submit ... --gpu-ids 2,3 --gpu-count 1 --estimated-vram-mb 24000   # GPU 2 or 3, whichever has room
```

**Multi-GPU job** — include a launcher; trainq sets `CUDA_VISIBLE_DEVICES` for you.
```bash
trainq submit --command "torchrun --nproc_per_node=2 train.py" --gpu-count 2 ...
```

**Evaluate every checkpoint on a spare GPU** — the watcher starts automatically when a hook is set.
```bash
trainq submit --command "python train.py" --ckpt-dir /d/exp --log-path /d/exp.log \
  --checkpoint-glob 'checkpoint-*' \
  --on-checkpoint 'python eval.py --ckpt {checkpoint}' --watch-gpu 3 \
  --description "train + per-checkpoint eval"
```

**Auto-back up checkpoints to the Hugging Face Hub** — needs `pip install -e '.[hf]'` and `HF_TOKEN`.
```bash
trainq submit ... --hf-repo-id <user>/<repo>
```

**Run cleanup / send a notification after a job** — `TRAINQ_JOB_STATUS` is `completed` or `failed`.
```bash
trainq submit ... --on-complete 'echo "$TRAINQ_JOB_STATUS: {ckpt_dir}" | mail -s "job done" me@example.com'
```

**Run an eval the instant training finishes** — negative priority beats aged-up training jobs.
```bash
trainq submit --command "python train.py" --priority 5  --description train ...
trainq submit --command "python eval.py"  --priority -10 --gpu-count 1 --estimated-vram-mb 20000 --description eval
```

**Get an exact ETA regardless of framework** — write a tiny file from your training loop; trainq prefers it over log parsing.
```python
import json, os
json.dump({"step": step, "total": total}, open(os.path.join(ckpt_dir, "progress.json"), "w"))
```

## Command reference

```
trainq list                    # all jobs + status + live ETA
trainq status <job_id>         # full detail of one job
trainq submit --command ...    # queue a job
trainq cancel <job_id>         # cancel a pending job
trainq priority <job_id> <n>   # re-prioritize a pending job (float ok; lower runs first)
trainq clear                   # drop completed/failed/cancelled records
```

### `trainq submit` options

| Option | Default | Meaning |
|--------|---------|---------|
| `--command` | required | Command to run (evaluated by bash; may include env + `cd`) |
| `--ckpt-dir` | required | Checkpoint / output root |
| `--log-path` | required | Log output path (must be unique per job) |
| `--description` | required | Shown in `trainq list` |
| `--gpu-count` | 1 | GPUs needed. For 2+, use a launcher (`torchrun`, etc.) |
| `--gpu-ids` | auto | Allow-set (not a pin); trainq picks `--gpu-count` from it, e.g. `2,3` |
| `--priority` | 5 | 0 (highest) … 9 (lowest). Negative allowed (stays ahead even after aging) |
| `--estimated-hours` | 0 | Enables backfill scheduling. 0 = disabled |
| `--estimated-vram-mb` | 0 | Enables VRAM-based sharing. 0 = require a fully idle GPU |
| `--max-retries` | 1 | (informational — a failed job is currently marked FAILED immediately) |
| `--checkpoint-glob` | none | Glob for new checkpoints under `--ckpt-dir` (see Checkpoint hooks) |
| `--on-checkpoint` | none | Command per new checkpoint (`{checkpoint} {name} {step} {ckpt_dir}`) |
| `--hf-repo-id` | none | Auto-upload each new checkpoint to this HF Hub repo |
| `--watch-gpu` | none | Dedicated GPU for the `--on-checkpoint` command |
| `--on-complete` | none | Command after the job ends (`{ckpt_dir} {log_path}`, env `TRAINQ_JOB_STATUS`) |

## How scheduling works

**VRAM-based sharing (`--estimated-vram-mb`).** With an estimate, a GPU is
eligible when `free_vram − estimate − margin ≥ 0`, so several jobs can share
one card (largest-free GPU first). Without an estimate, only a fully idle GPU
(no processes) is used — safe, but you wait longer. A freshly launched process
takes time to actually claim VRAM, so trainq *reserves* the estimate on launch and
releases it once the process is seen holding memory, preventing a
double-booking OOM.

**Backfill (`--estimated-hours`).** When the top job can't get enough GPUs, trainq
looks for a lower job that will finish before those GPUs free up. If the
candidate uses GPUs the top job isn't waiting on, it can't delay the top job,
so it runs unconditionally.

**Priority aging.** Effective priority drops (runs sooner) the longer a job
waits — 1 point per 30h by default — so a low-priority job can't wait forever.
Aging is deliberately gentle so explicit priorities still mean something.

**ETA / progress.** `trainq list` re-parses logs on every call. Auto-detected
formats: fairseq2, HuggingFace Trainer, PaddleOCR (`epoch: [M/N]`),
`metrics.jsonl` (`iter`), and generic `step X/Y`. If none match, it falls back
to `--estimated-hours` (`~Xh`). The most robust option is the `progress.json`
convention shown in the Cookbook — it's format- and buffering-independent and
takes priority over log parsing.

## Checkpoint hooks (detail)

If a job sets `--on-checkpoint` or `--hf-repo-id`, the daemon launches a watcher
alongside it. The watcher scans `--ckpt-dir` for paths matching
`--checkpoint-glob` and, for each new one:

1. uploads it to `--hf-repo-id` if set (file or directory), then
2. runs `--on-checkpoint` if set, with `{checkpoint} {name} {step} {ckpt_dir}` substituted.

When the training process exits, the watcher does a final sweep and stops.

`--checkpoint-glob` examples: `checkpoint-*` (default, HF Trainer convention),
`ws_*/checkpoints/step_*`, `epoch_*`, `*.pt`.

> To avoid grabbing a checkpoint that's still being written, trainq only processes
> a checkpoint once it hasn't changed for `--stabilize-sec` seconds (default 10;
> raise it, or set 0 to disable). This is a debounce for frameworks that don't
> save atomically.

`--on-complete` runs once after the whole job ends, with `{ckpt_dir} {log_path}`
substituted and `TRAINQ_JOB_STATUS` set to `completed` or `failed`.

## Configuration

All configuration is via environment variables.

| Variable | Default | Description |
|----------|---------|-------------|
| `TRAINQ_REDIS_HOST` | localhost | Redis host |
| `TRAINQ_REDIS_PORT` | 6379 | Redis port |
| `TRAINQ_REDIS_DB` | 0 | Redis DB number |
| `TRAINQ_REDIS_PASSWORD` | (none) | Redis password |
| `TRAINQ_POLL_SEC` | 10 | Daemon poll interval (seconds) |
| `TRAINQ_VRAM_THRESHOLD_MB` | 1000 | Min free VRAM to call a GPU idle (MB) |
| `HF_TOKEN` | (none) | For HF upload; falls back to `./.env`, `~/.env`, `~/.cache/huggingface/token` |

## Deploying to a server

`scripts/deploy_trainq.sh` rsyncs the package to a remote GPU host, sets up a venv,
installs trainq, and (re)starts the daemon. Running jobs are re-tracked by PID, so
a redeploy doesn't interrupt them.

```bash
bash scripts/deploy_trainq.sh <host> [remote_user] [remote_dir]
```

Requires: `ssh <host>` access, and `python3` + `sudo` (apt) on the remote.

## Multi-user notes

The executor runs each job as the submitting user (`sudo -u` when that differs
from the daemon user). To use this on a shared host, the daemon user needs
passwordless `sudo -u` for the target users. In a single-user setup, the daemon
and jobs run as the same user and no sudo is involved.

## Architecture

| File | Role |
|------|------|
| `trainq/schema.py` | `Job` data model (serialize / deserialize) |
| `trainq/config.py` | Redis connection, key names, poll interval (`TRAINQ_*` env) |
| `trainq/gpu_monitor.py` | nvidia-smi parsing, VRAM reservation, availability |
| `trainq/scheduler.py` | Aging + backfill scheduling (effective priority) |
| `trainq/executor.py` | Job launch, process management, completion hook |
| `trainq/server.py` | Queue daemon main loop |
| `trainq/client.py` | CLI (submit/list/status/cancel/clear/priority) |
| `trainq/progress.py` | Log/file-based progress & ETA estimation |
| `trainq/checkpoint_watcher.py` | Checkpoint watcher CLI entry point |
| `trainq/watcher/daemon.py` | Detect new checkpoints, run hook, optional HF upload |
| `trainq/watcher/checkpoint_tracker.py` | Checkpoint detection (glob) + processed-state tracking |
| `trainq/watcher/hf_uploader.py` | Hugging Face Hub upload (optional) |

**Flow:** `trainq submit` → Redis sorted set `trainq:jobs` → the daemon polls and calls
`scheduler.pick_next_job` (run the top job if GPUs fit, else look for a backfill
candidate) → launch the job (and a watcher if hooks are set) → on finish, run
`on_complete`; on unresolved failure, leave it FAILED without running hooks.

## License

MIT — see [LICENSE](LICENSE).
