trainq 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- trainq-0.1.0/LICENSE +21 -0
- trainq-0.1.0/PKG-INFO +299 -0
- trainq-0.1.0/README.md +269 -0
- trainq-0.1.0/pyproject.toml +45 -0
- trainq-0.1.0/setup.cfg +4 -0
- trainq-0.1.0/tests/test_checkpoint_tracker.py +30 -0
- trainq-0.1.0/tests/test_progress.py +43 -0
- trainq-0.1.0/tests/test_schema.py +40 -0
- trainq-0.1.0/trainq/__init__.py +0 -0
- trainq-0.1.0/trainq/checkpoint_watcher.py +34 -0
- trainq-0.1.0/trainq/client.py +383 -0
- trainq-0.1.0/trainq/config.py +21 -0
- trainq-0.1.0/trainq/executor.py +379 -0
- trainq-0.1.0/trainq/gpu_monitor.py +231 -0
- trainq-0.1.0/trainq/progress.py +292 -0
- trainq-0.1.0/trainq/scheduler.py +219 -0
- trainq-0.1.0/trainq/schema.py +78 -0
- trainq-0.1.0/trainq/server.py +145 -0
- trainq-0.1.0/trainq/watcher/__init__.py +22 -0
- trainq-0.1.0/trainq/watcher/checkpoint_tracker.py +82 -0
- trainq-0.1.0/trainq/watcher/daemon.py +177 -0
- trainq-0.1.0/trainq/watcher/hf_uploader.py +132 -0
- trainq-0.1.0/trainq.egg-info/PKG-INFO +299 -0
- trainq-0.1.0/trainq.egg-info/SOURCES.txt +26 -0
- trainq-0.1.0/trainq.egg-info/dependency_links.txt +1 -0
- trainq-0.1.0/trainq.egg-info/entry_points.txt +2 -0
- trainq-0.1.0/trainq.egg-info/requires.txt +7 -0
- trainq-0.1.0/trainq.egg-info/top_level.txt +1 -0
trainq-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Korshort
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
trainq-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,299 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: trainq
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Redis-based GPU job queue for ML training/eval: VRAM-aware scheduling, priority aging, backfill.
|
|
5
|
+
Author: Korshort
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/Korshort/trainq
|
|
8
|
+
Project-URL: Repository, https://github.com/Korshort/trainq
|
|
9
|
+
Project-URL: Issues, https://github.com/Korshort/trainq/issues
|
|
10
|
+
Keywords: gpu,job-queue,scheduler,machine-learning,redis
|
|
11
|
+
Classifier: Development Status :: 4 - Beta
|
|
12
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
13
|
+
Classifier: Operating System :: POSIX :: Linux
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
19
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
20
|
+
Classifier: Topic :: System :: Distributed Computing
|
|
21
|
+
Requires-Python: >=3.9
|
|
22
|
+
Description-Content-Type: text/markdown
|
|
23
|
+
License-File: LICENSE
|
|
24
|
+
Requires-Dist: redis>=5.0
|
|
25
|
+
Provides-Extra: hf
|
|
26
|
+
Requires-Dist: huggingface_hub>=0.20; extra == "hf"
|
|
27
|
+
Provides-Extra: dev
|
|
28
|
+
Requires-Dist: pytest>=7; extra == "dev"
|
|
29
|
+
Dynamic: license-file
|
|
30
|
+
|
|
31
|
+
# trainq — a lightweight GPU job queue
|
|
32
|
+
|
|
33
|
+
Submit training/eval jobs and `trainq` runs them as soon as the GPUs are free.
|
|
34
|
+
Redis-backed, framework-agnostic, single-file-simple. Your command can be
|
|
35
|
+
anything: PyTorch, JAX, `torchrun`, or a plain shell script.
|
|
36
|
+
|
|
37
|
+
```
|
|
38
|
+
trainq submit --command "python train.py" --ckpt-dir /data/exp1 \
|
|
39
|
+
--log-path /data/exp1.log --gpu-count 1 --description "my first job"
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
**Why trainq?** On a shared GPU box, people step on each other, big jobs starve
|
|
43
|
+
small ones, and you never know when your run will actually start. trainq gives you
|
|
44
|
+
a queue that packs jobs onto GPUs by VRAM, fills idle gaps with short jobs
|
|
45
|
+
(backfill), ages up long-waiting jobs so nothing starves, and shows a live ETA
|
|
46
|
+
for everything in flight.
|
|
47
|
+
|
|
48
|
+
## Features
|
|
49
|
+
|
|
50
|
+
- **VRAM-aware sharing** — co-locate several jobs on one GPU; reservations prevent the startup-time OOM race.
|
|
51
|
+
- **Priority + aging** — explicit priorities, plus automatic promotion the longer a job waits (no starvation).
|
|
52
|
+
- **Backfill** — while a big job waits for GPUs, short jobs that finish before those GPUs free up run first.
|
|
53
|
+
- **Live ETA** — progress is parsed from your logs (several formats auto-detected) or a tiny universal convention.
|
|
54
|
+
- **Checkpoint hooks (optional)** — run any command and/or auto-upload to the Hugging Face Hub on each new checkpoint.
|
|
55
|
+
- **Crash-tolerant daemon** — if the daemon dies, running jobs keep going; on restart it re-tracks them by PID.
|
|
56
|
+
|
|
57
|
+
## Requirements
|
|
58
|
+
|
|
59
|
+
- Python 3.9+
|
|
60
|
+
- Redis (`sudo apt-get install -y redis-server`, or `brew install redis`)
|
|
61
|
+
- `nvidia-smi` for GPU scheduling (only the queue daemon needs it; the checkpoint watcher does not)
|
|
62
|
+
|
|
63
|
+
## Quick start
|
|
64
|
+
|
|
65
|
+
```bash
|
|
66
|
+
# 1) install — from PyPI
|
|
67
|
+
pip install trainq
|
|
68
|
+
# …or from source:
|
|
69
|
+
# git clone https://github.com/Korshort/trainq && cd trainq && pip install -e .
|
|
70
|
+
|
|
71
|
+
# 2) start the queue daemon on your GPU box
|
|
72
|
+
python -m trainq.server >> trainq_server.log 2>&1 &
|
|
73
|
+
|
|
74
|
+
# 3) submit a job
|
|
75
|
+
trainq submit --command "python train.py --lr 1e-4" \
|
|
76
|
+
--ckpt-dir /data/ckpt/exp1 --log-path /data/logs/exp1.log \
|
|
77
|
+
--gpu-count 1 --estimated-vram-mb 15000 \
|
|
78
|
+
--description "resnet50 finetune lr1e-4"
|
|
79
|
+
|
|
80
|
+
# 4) watch it
|
|
81
|
+
trainq list # every job, with status, GPU, and live ETA
|
|
82
|
+
```
|
|
83
|
+
|
|
84
|
+
That's the whole loop. Everything below is optional depth.
|
|
85
|
+
|
|
86
|
+
## Try it in 2 minutes (no GPU needed)
|
|
87
|
+
|
|
88
|
+
You can see the checkpoint-hook system end to end on a laptop, no GPU or model
|
|
89
|
+
required. [`examples/dummy_train.py`](examples/dummy_train.py) is a fake trainer
|
|
90
|
+
that just writes progress and drops checkpoint folders on a timer.
|
|
91
|
+
|
|
92
|
+
```bash
|
|
93
|
+
# Terminal A — a "training" run that saves a checkpoint every 5 steps
|
|
94
|
+
python examples/dummy_train.py --ckpt-dir /tmp/demo --steps 20 --save-every 5 --sleep 1
|
|
95
|
+
|
|
96
|
+
# Terminal B — watch that directory and run a hook on each new checkpoint
|
|
97
|
+
python -m trainq.checkpoint_watcher --ckpt-dir /tmp/demo \
|
|
98
|
+
--on-checkpoint 'echo ">> new checkpoint: {name} (step {step})"'
|
|
99
|
+
```
|
|
100
|
+
|
|
101
|
+
You'll see the hook fire for `checkpoint-5`, `-10`, `-15`, `-20`. See
|
|
102
|
+
[examples/](examples/) for the full walkthrough, including the GPU queue flow.
|
|
103
|
+
|
|
104
|
+
## Contributing
|
|
105
|
+
|
|
106
|
+
Contributions are welcome, and the no-GPU example above is the fastest way to
|
|
107
|
+
see the moving parts. Good first areas: more log-format parsers in
|
|
108
|
+
`trainq/progress.py`, more recipes in `examples/`, and tests. Please keep the
|
|
109
|
+
core dependency-light (`redis` only; `huggingface_hub` stays an optional extra).
|
|
110
|
+
Open an issue for bugs or ideas, or a PR for changes — see
|
|
111
|
+
[CONTRIBUTING.md](CONTRIBUTING.md) for the dev setup and how to run tests.
|
|
112
|
+
|
|
113
|
+
## Cookbook
|
|
114
|
+
|
|
115
|
+
Common recipes. Mix and match the flags.
|
|
116
|
+
|
|
117
|
+
**Pack two jobs onto one GPU** — give each job its VRAM estimate; trainq shares the GPU when it fits.
|
|
118
|
+
```bash
|
|
119
|
+
trainq submit --command "python train.py" --ckpt-dir /d/a --log-path /d/a.log --estimated-vram-mb 18000 --description a
|
|
120
|
+
trainq submit --command "python eval.py" --ckpt-dir /d/b --log-path /d/b.log --estimated-vram-mb 12000 --description b
|
|
121
|
+
```
|
|
122
|
+
|
|
123
|
+
**Restrict a job to certain GPUs** — `--gpu-ids` is an *allow-set*, not a pin; trainq picks `--gpu-count` from it.
|
|
124
|
+
```bash
|
|
125
|
+
trainq submit ... --gpu-ids 2,3 --gpu-count 1 --estimated-vram-mb 24000 # GPU 2 or 3, whichever has room
|
|
126
|
+
```
|
|
127
|
+
|
|
128
|
+
**Multi-GPU job** — include a launcher; trainq sets `CUDA_VISIBLE_DEVICES` for you.
|
|
129
|
+
```bash
|
|
130
|
+
trainq submit --command "torchrun --nproc_per_node=2 train.py" --gpu-count 2 ...
|
|
131
|
+
```
|
|
132
|
+
|
|
133
|
+
**Evaluate every checkpoint on a spare GPU** — the watcher starts automatically when a hook is set.
|
|
134
|
+
```bash
|
|
135
|
+
trainq submit --command "python train.py" --ckpt-dir /d/exp --log-path /d/exp.log \
|
|
136
|
+
--checkpoint-glob 'checkpoint-*' \
|
|
137
|
+
--on-checkpoint 'python eval.py --ckpt {checkpoint}' --watch-gpu 3 \
|
|
138
|
+
--description "train + per-checkpoint eval"
|
|
139
|
+
```
|
|
140
|
+
|
|
141
|
+
**Auto-back up checkpoints to the Hugging Face Hub** — needs `pip install -e '.[hf]'` and `HF_TOKEN`.
|
|
142
|
+
```bash
|
|
143
|
+
trainq submit ... --hf-repo-id <user>/<repo>
|
|
144
|
+
```
|
|
145
|
+
|
|
146
|
+
**Run cleanup / send a notification after a job** — `TRAINQ_JOB_STATUS` is `completed` or `failed`.
|
|
147
|
+
```bash
|
|
148
|
+
trainq submit ... --on-complete 'echo "$TRAINQ_JOB_STATUS: {ckpt_dir}" | mail -s "job done" me@example.com'
|
|
149
|
+
```
|
|
150
|
+
|
|
151
|
+
**Run an eval the instant training finishes** — negative priority beats aged-up training jobs.
|
|
152
|
+
```bash
|
|
153
|
+
trainq submit --command "python train.py" --priority 5 --description train ...
|
|
154
|
+
trainq submit --command "python eval.py" --priority -10 --gpu-count 1 --estimated-vram-mb 20000 --description eval
|
|
155
|
+
```
|
|
156
|
+
|
|
157
|
+
**Get an exact ETA regardless of framework** — write a tiny file from your training loop; trainq prefers it over log parsing.
|
|
158
|
+
```python
|
|
159
|
+
import json, os
|
|
160
|
+
json.dump({"step": step, "total": total}, open(os.path.join(ckpt_dir, "progress.json"), "w"))
|
|
161
|
+
```
|
|
162
|
+
|
|
163
|
+
## Command reference
|
|
164
|
+
|
|
165
|
+
```
|
|
166
|
+
trainq list # all jobs + status + live ETA
|
|
167
|
+
trainq status <job_id> # full detail of one job
|
|
168
|
+
trainq submit --command ... # queue a job
|
|
169
|
+
trainq cancel <job_id> # cancel a pending job
|
|
170
|
+
trainq priority <job_id> <n> # re-prioritize a pending job (float ok; lower runs first)
|
|
171
|
+
trainq clear # drop completed/failed/cancelled records
|
|
172
|
+
```
|
|
173
|
+
|
|
174
|
+
### `trainq submit` options
|
|
175
|
+
|
|
176
|
+
| Option | Default | Meaning |
|
|
177
|
+
|--------|---------|---------|
|
|
178
|
+
| `--command` | required | Command to run (evaluated by bash; may include env + `cd`) |
|
|
179
|
+
| `--ckpt-dir` | required | Checkpoint / output root |
|
|
180
|
+
| `--log-path` | required | Log output path (must be unique per job) |
|
|
181
|
+
| `--description` | required | Shown in `trainq list` |
|
|
182
|
+
| `--gpu-count` | 1 | GPUs needed. For 2+, use a launcher (`torchrun`, etc.) |
|
|
183
|
+
| `--gpu-ids` | auto | Allow-set (not a pin); trainq picks `--gpu-count` from it, e.g. `2,3` |
|
|
184
|
+
| `--priority` | 5 | 0 (highest) … 9 (lowest). Negative allowed (stays ahead even after aging) |
|
|
185
|
+
| `--estimated-hours` | 0 | Enables backfill scheduling. 0 = disabled |
|
|
186
|
+
| `--estimated-vram-mb` | 0 | Enables VRAM-based sharing. 0 = require a fully idle GPU |
|
|
187
|
+
| `--max-retries` | 1 | (informational — a failed job is currently marked FAILED immediately) |
|
|
188
|
+
| `--checkpoint-glob` | none | Glob for new checkpoints under `--ckpt-dir` (see Checkpoint hooks) |
|
|
189
|
+
| `--on-checkpoint` | none | Command per new checkpoint (`{checkpoint} {name} {step} {ckpt_dir}`) |
|
|
190
|
+
| `--hf-repo-id` | none | Auto-upload each new checkpoint to this HF Hub repo |
|
|
191
|
+
| `--watch-gpu` | none | Dedicated GPU for the `--on-checkpoint` command |
|
|
192
|
+
| `--on-complete` | none | Command after the job ends (`{ckpt_dir} {log_path}`, env `TRAINQ_JOB_STATUS`) |
|
|
193
|
+
|
|
194
|
+
## How scheduling works
|
|
195
|
+
|
|
196
|
+
**VRAM-based sharing (`--estimated-vram-mb`).** With an estimate, a GPU is
|
|
197
|
+
eligible when `free_vram − estimate − margin ≥ 0`, so several jobs can share
|
|
198
|
+
one card (largest-free GPU first). Without an estimate, only a fully idle GPU
|
|
199
|
+
(no processes) is used — safe, but you wait longer. A freshly launched process
|
|
200
|
+
takes time to actually claim VRAM, so trainq *reserves* the estimate on launch and
|
|
201
|
+
releases it once the process is seen holding memory, preventing a
|
|
202
|
+
double-booking OOM.
|
|
203
|
+
|
|
204
|
+
**Backfill (`--estimated-hours`).** When the top job can't get enough GPUs, trainq
|
|
205
|
+
looks for a lower job that will finish before those GPUs free up. If the
|
|
206
|
+
candidate uses GPUs the top job isn't waiting on, it can't delay the top job,
|
|
207
|
+
so it runs unconditionally.
|
|
208
|
+
|
|
209
|
+
**Priority aging.** Effective priority drops (runs sooner) the longer a job
|
|
210
|
+
waits — 1 point per 30h by default — so a low-priority job can't wait forever.
|
|
211
|
+
Aging is deliberately gentle so explicit priorities still mean something.
|
|
212
|
+
|
|
213
|
+
**ETA / progress.** `trainq list` re-parses logs on every call. Auto-detected
|
|
214
|
+
formats: fairseq2, HuggingFace Trainer, PaddleOCR (`epoch: [M/N]`),
|
|
215
|
+
`metrics.jsonl` (`iter`), and generic `step X/Y`. If none match, it falls back
|
|
216
|
+
to `--estimated-hours` (`~Xh`). The most robust option is the `progress.json`
|
|
217
|
+
convention shown in the Cookbook — it's format- and buffering-independent and
|
|
218
|
+
takes priority over log parsing.
|
|
219
|
+
|
|
220
|
+
## Checkpoint hooks (detail)
|
|
221
|
+
|
|
222
|
+
If a job sets `--on-checkpoint` or `--hf-repo-id`, the daemon launches a watcher
|
|
223
|
+
alongside it. The watcher scans `--ckpt-dir` for paths matching
|
|
224
|
+
`--checkpoint-glob` and, for each new one:
|
|
225
|
+
|
|
226
|
+
1. uploads it to `--hf-repo-id` if set (file or directory), then
|
|
227
|
+
2. runs `--on-checkpoint` if set, with `{checkpoint} {name} {step} {ckpt_dir}` substituted.
|
|
228
|
+
|
|
229
|
+
When the training process exits, the watcher does a final sweep and stops.
|
|
230
|
+
|
|
231
|
+
`--checkpoint-glob` examples: `checkpoint-*` (default, HF Trainer convention),
|
|
232
|
+
`ws_*/checkpoints/step_*`, `epoch_*`, `*.pt`.
|
|
233
|
+
|
|
234
|
+
> To avoid grabbing a checkpoint that's still being written, trainq only processes
|
|
235
|
+
> a checkpoint once it hasn't changed for `--stabilize-sec` seconds (default 10;
|
|
236
|
+
> raise it, or set 0 to disable). This is a debounce for frameworks that don't
|
|
237
|
+
> save atomically.
|
|
238
|
+
|
|
239
|
+
`--on-complete` runs once after the whole job ends, with `{ckpt_dir} {log_path}`
|
|
240
|
+
substituted and `TRAINQ_JOB_STATUS` set to `completed` or `failed`.
|
|
241
|
+
|
|
242
|
+
## Configuration
|
|
243
|
+
|
|
244
|
+
All configuration is via environment variables.
|
|
245
|
+
|
|
246
|
+
| Variable | Default | Description |
|
|
247
|
+
|----------|---------|-------------|
|
|
248
|
+
| `TRAINQ_REDIS_HOST` | localhost | Redis host |
|
|
249
|
+
| `TRAINQ_REDIS_PORT` | 6379 | Redis port |
|
|
250
|
+
| `TRAINQ_REDIS_DB` | 0 | Redis DB number |
|
|
251
|
+
| `TRAINQ_REDIS_PASSWORD` | (none) | Redis password |
|
|
252
|
+
| `TRAINQ_POLL_SEC` | 10 | Daemon poll interval (seconds) |
|
|
253
|
+
| `TRAINQ_VRAM_THRESHOLD_MB` | 1000 | Min free VRAM to call a GPU idle (MB) |
|
|
254
|
+
| `HF_TOKEN` | (none) | For HF upload; falls back to `./.env`, `~/.env`, `~/.cache/huggingface/token` |
|
|
255
|
+
|
|
256
|
+
## Deploying to a server
|
|
257
|
+
|
|
258
|
+
`scripts/deploy_trainq.sh` rsyncs the package to a remote GPU host, sets up a venv,
|
|
259
|
+
installs trainq, and (re)starts the daemon. Running jobs are re-tracked by PID, so
|
|
260
|
+
a redeploy doesn't interrupt them.
|
|
261
|
+
|
|
262
|
+
```bash
|
|
263
|
+
bash scripts/deploy_trainq.sh <host> [remote_user] [remote_dir]
|
|
264
|
+
```
|
|
265
|
+
|
|
266
|
+
Requires: `ssh <host>` access, and `python3` + `sudo` (apt) on the remote.
|
|
267
|
+
|
|
268
|
+
## Multi-user notes
|
|
269
|
+
|
|
270
|
+
The executor runs each job as the submitting user (`sudo -u` when that differs
|
|
271
|
+
from the daemon user). To use this on a shared host, the daemon user needs
|
|
272
|
+
passwordless `sudo -u` for the target users. In a single-user setup, the daemon
|
|
273
|
+
and jobs run as the same user and no sudo is involved.
|
|
274
|
+
|
|
275
|
+
## Architecture
|
|
276
|
+
|
|
277
|
+
| File | Role |
|
|
278
|
+
|------|------|
|
|
279
|
+
| `trainq/schema.py` | `Job` data model (serialize / deserialize) |
|
|
280
|
+
| `trainq/config.py` | Redis connection, key names, poll interval (`TRAINQ_*` env) |
|
|
281
|
+
| `trainq/gpu_monitor.py` | nvidia-smi parsing, VRAM reservation, availability |
|
|
282
|
+
| `trainq/scheduler.py` | Aging + backfill scheduling (effective priority) |
|
|
283
|
+
| `trainq/executor.py` | Job launch, process management, completion hook |
|
|
284
|
+
| `trainq/server.py` | Queue daemon main loop |
|
|
285
|
+
| `trainq/client.py` | CLI (submit/list/status/cancel/clear/priority) |
|
|
286
|
+
| `trainq/progress.py` | Log/file-based progress & ETA estimation |
|
|
287
|
+
| `trainq/checkpoint_watcher.py` | Checkpoint watcher CLI entry point |
|
|
288
|
+
| `trainq/watcher/daemon.py` | Detect new checkpoints, run hook, optional HF upload |
|
|
289
|
+
| `trainq/watcher/checkpoint_tracker.py` | Checkpoint detection (glob) + processed-state tracking |
|
|
290
|
+
| `trainq/watcher/hf_uploader.py` | Hugging Face Hub upload (optional) |
|
|
291
|
+
|
|
292
|
+
**Flow:** `trainq submit` → Redis sorted set `trainq:jobs` → the daemon polls and calls
|
|
293
|
+
`scheduler.pick_next_job` (run the top job if GPUs fit, else look for a backfill
|
|
294
|
+
candidate) → launch the job (and a watcher if hooks are set) → on finish, run
|
|
295
|
+
`on_complete`; on unresolved failure, leave it FAILED without running hooks.
|
|
296
|
+
|
|
297
|
+
## License
|
|
298
|
+
|
|
299
|
+
MIT — see [LICENSE](LICENSE).
|
trainq-0.1.0/README.md
ADDED
|
@@ -0,0 +1,269 @@
|
|
|
1
|
+
# trainq — a lightweight GPU job queue
|
|
2
|
+
|
|
3
|
+
Submit training/eval jobs and `trainq` runs them as soon as the GPUs are free.
|
|
4
|
+
Redis-backed, framework-agnostic, single-file-simple. Your command can be
|
|
5
|
+
anything: PyTorch, JAX, `torchrun`, or a plain shell script.
|
|
6
|
+
|
|
7
|
+
```
|
|
8
|
+
trainq submit --command "python train.py" --ckpt-dir /data/exp1 \
|
|
9
|
+
--log-path /data/exp1.log --gpu-count 1 --description "my first job"
|
|
10
|
+
```
|
|
11
|
+
|
|
12
|
+
**Why trainq?** On a shared GPU box, people step on each other, big jobs starve
|
|
13
|
+
small ones, and you never know when your run will actually start. trainq gives you
|
|
14
|
+
a queue that packs jobs onto GPUs by VRAM, fills idle gaps with short jobs
|
|
15
|
+
(backfill), ages up long-waiting jobs so nothing starves, and shows a live ETA
|
|
16
|
+
for everything in flight.
|
|
17
|
+
|
|
18
|
+
## Features
|
|
19
|
+
|
|
20
|
+
- **VRAM-aware sharing** — co-locate several jobs on one GPU; reservations prevent the startup-time OOM race.
|
|
21
|
+
- **Priority + aging** — explicit priorities, plus automatic promotion the longer a job waits (no starvation).
|
|
22
|
+
- **Backfill** — while a big job waits for GPUs, short jobs that finish before those GPUs free up run first.
|
|
23
|
+
- **Live ETA** — progress is parsed from your logs (several formats auto-detected) or a tiny universal convention.
|
|
24
|
+
- **Checkpoint hooks (optional)** — run any command and/or auto-upload to the Hugging Face Hub on each new checkpoint.
|
|
25
|
+
- **Crash-tolerant daemon** — if the daemon dies, running jobs keep going; on restart it re-tracks them by PID.
|
|
26
|
+
|
|
27
|
+
## Requirements
|
|
28
|
+
|
|
29
|
+
- Python 3.9+
|
|
30
|
+
- Redis (`sudo apt-get install -y redis-server`, or `brew install redis`)
|
|
31
|
+
- `nvidia-smi` for GPU scheduling (only the queue daemon needs it; the checkpoint watcher does not)
|
|
32
|
+
|
|
33
|
+
## Quick start
|
|
34
|
+
|
|
35
|
+
```bash
|
|
36
|
+
# 1) install — from PyPI
|
|
37
|
+
pip install trainq
|
|
38
|
+
# …or from source:
|
|
39
|
+
# git clone https://github.com/Korshort/trainq && cd trainq && pip install -e .
|
|
40
|
+
|
|
41
|
+
# 2) start the queue daemon on your GPU box
|
|
42
|
+
python -m trainq.server >> trainq_server.log 2>&1 &
|
|
43
|
+
|
|
44
|
+
# 3) submit a job
|
|
45
|
+
trainq submit --command "python train.py --lr 1e-4" \
|
|
46
|
+
--ckpt-dir /data/ckpt/exp1 --log-path /data/logs/exp1.log \
|
|
47
|
+
--gpu-count 1 --estimated-vram-mb 15000 \
|
|
48
|
+
--description "resnet50 finetune lr1e-4"
|
|
49
|
+
|
|
50
|
+
# 4) watch it
|
|
51
|
+
trainq list # every job, with status, GPU, and live ETA
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
That's the whole loop. Everything below is optional depth.
|
|
55
|
+
|
|
56
|
+
## Try it in 2 minutes (no GPU needed)
|
|
57
|
+
|
|
58
|
+
You can see the checkpoint-hook system end to end on a laptop, no GPU or model
|
|
59
|
+
required. [`examples/dummy_train.py`](examples/dummy_train.py) is a fake trainer
|
|
60
|
+
that just writes progress and drops checkpoint folders on a timer.
|
|
61
|
+
|
|
62
|
+
```bash
|
|
63
|
+
# Terminal A — a "training" run that saves a checkpoint every 5 steps
|
|
64
|
+
python examples/dummy_train.py --ckpt-dir /tmp/demo --steps 20 --save-every 5 --sleep 1
|
|
65
|
+
|
|
66
|
+
# Terminal B — watch that directory and run a hook on each new checkpoint
|
|
67
|
+
python -m trainq.checkpoint_watcher --ckpt-dir /tmp/demo \
|
|
68
|
+
--on-checkpoint 'echo ">> new checkpoint: {name} (step {step})"'
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
You'll see the hook fire for `checkpoint-5`, `-10`, `-15`, `-20`. See
|
|
72
|
+
[examples/](examples/) for the full walkthrough, including the GPU queue flow.
|
|
73
|
+
|
|
74
|
+
## Contributing
|
|
75
|
+
|
|
76
|
+
Contributions are welcome, and the no-GPU example above is the fastest way to
|
|
77
|
+
see the moving parts. Good first areas: more log-format parsers in
|
|
78
|
+
`trainq/progress.py`, more recipes in `examples/`, and tests. Please keep the
|
|
79
|
+
core dependency-light (`redis` only; `huggingface_hub` stays an optional extra).
|
|
80
|
+
Open an issue for bugs or ideas, or a PR for changes — see
|
|
81
|
+
[CONTRIBUTING.md](CONTRIBUTING.md) for the dev setup and how to run tests.
|
|
82
|
+
|
|
83
|
+
## Cookbook
|
|
84
|
+
|
|
85
|
+
Common recipes. Mix and match the flags.
|
|
86
|
+
|
|
87
|
+
**Pack two jobs onto one GPU** — give each job its VRAM estimate; trainq shares the GPU when it fits.
|
|
88
|
+
```bash
|
|
89
|
+
trainq submit --command "python train.py" --ckpt-dir /d/a --log-path /d/a.log --estimated-vram-mb 18000 --description a
|
|
90
|
+
trainq submit --command "python eval.py" --ckpt-dir /d/b --log-path /d/b.log --estimated-vram-mb 12000 --description b
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
**Restrict a job to certain GPUs** — `--gpu-ids` is an *allow-set*, not a pin; trainq picks `--gpu-count` from it.
|
|
94
|
+
```bash
|
|
95
|
+
trainq submit ... --gpu-ids 2,3 --gpu-count 1 --estimated-vram-mb 24000 # GPU 2 or 3, whichever has room
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
**Multi-GPU job** — include a launcher; trainq sets `CUDA_VISIBLE_DEVICES` for you.
|
|
99
|
+
```bash
|
|
100
|
+
trainq submit --command "torchrun --nproc_per_node=2 train.py" --gpu-count 2 ...
|
|
101
|
+
```
|
|
102
|
+
|
|
103
|
+
**Evaluate every checkpoint on a spare GPU** — the watcher starts automatically when a hook is set.
|
|
104
|
+
```bash
|
|
105
|
+
trainq submit --command "python train.py" --ckpt-dir /d/exp --log-path /d/exp.log \
|
|
106
|
+
--checkpoint-glob 'checkpoint-*' \
|
|
107
|
+
--on-checkpoint 'python eval.py --ckpt {checkpoint}' --watch-gpu 3 \
|
|
108
|
+
--description "train + per-checkpoint eval"
|
|
109
|
+
```
|
|
110
|
+
|
|
111
|
+
**Auto-back up checkpoints to the Hugging Face Hub** — needs `pip install -e '.[hf]'` and `HF_TOKEN`.
|
|
112
|
+
```bash
|
|
113
|
+
trainq submit ... --hf-repo-id <user>/<repo>
|
|
114
|
+
```
|
|
115
|
+
|
|
116
|
+
**Run cleanup / send a notification after a job** — `TRAINQ_JOB_STATUS` is `completed` or `failed`.
|
|
117
|
+
```bash
|
|
118
|
+
trainq submit ... --on-complete 'echo "$TRAINQ_JOB_STATUS: {ckpt_dir}" | mail -s "job done" me@example.com'
|
|
119
|
+
```
|
|
120
|
+
|
|
121
|
+
**Run an eval the instant training finishes** — negative priority beats aged-up training jobs.
|
|
122
|
+
```bash
|
|
123
|
+
trainq submit --command "python train.py" --priority 5 --description train ...
|
|
124
|
+
trainq submit --command "python eval.py" --priority -10 --gpu-count 1 --estimated-vram-mb 20000 --description eval
|
|
125
|
+
```
|
|
126
|
+
|
|
127
|
+
**Get an exact ETA regardless of framework** — write a tiny file from your training loop; trainq prefers it over log parsing.
|
|
128
|
+
```python
|
|
129
|
+
import json, os
|
|
130
|
+
json.dump({"step": step, "total": total}, open(os.path.join(ckpt_dir, "progress.json"), "w"))
|
|
131
|
+
```
|
|
132
|
+
|
|
133
|
+
## Command reference
|
|
134
|
+
|
|
135
|
+
```
|
|
136
|
+
trainq list # all jobs + status + live ETA
|
|
137
|
+
trainq status <job_id> # full detail of one job
|
|
138
|
+
trainq submit --command ... # queue a job
|
|
139
|
+
trainq cancel <job_id> # cancel a pending job
|
|
140
|
+
trainq priority <job_id> <n> # re-prioritize a pending job (float ok; lower runs first)
|
|
141
|
+
trainq clear # drop completed/failed/cancelled records
|
|
142
|
+
```
|
|
143
|
+
|
|
144
|
+
### `trainq submit` options
|
|
145
|
+
|
|
146
|
+
| Option | Default | Meaning |
|
|
147
|
+
|--------|---------|---------|
|
|
148
|
+
| `--command` | required | Command to run (evaluated by bash; may include env + `cd`) |
|
|
149
|
+
| `--ckpt-dir` | required | Checkpoint / output root |
|
|
150
|
+
| `--log-path` | required | Log output path (must be unique per job) |
|
|
151
|
+
| `--description` | required | Shown in `trainq list` |
|
|
152
|
+
| `--gpu-count` | 1 | GPUs needed. For 2+, use a launcher (`torchrun`, etc.) |
|
|
153
|
+
| `--gpu-ids` | auto | Allow-set (not a pin); trainq picks `--gpu-count` from it, e.g. `2,3` |
|
|
154
|
+
| `--priority` | 5 | 0 (highest) … 9 (lowest). Negative allowed (stays ahead even after aging) |
|
|
155
|
+
| `--estimated-hours` | 0 | Enables backfill scheduling. 0 = disabled |
|
|
156
|
+
| `--estimated-vram-mb` | 0 | Enables VRAM-based sharing. 0 = require a fully idle GPU |
|
|
157
|
+
| `--max-retries` | 1 | (informational — a failed job is currently marked FAILED immediately) |
|
|
158
|
+
| `--checkpoint-glob` | none | Glob for new checkpoints under `--ckpt-dir` (see Checkpoint hooks) |
|
|
159
|
+
| `--on-checkpoint` | none | Command per new checkpoint (`{checkpoint} {name} {step} {ckpt_dir}`) |
|
|
160
|
+
| `--hf-repo-id` | none | Auto-upload each new checkpoint to this HF Hub repo |
|
|
161
|
+
| `--watch-gpu` | none | Dedicated GPU for the `--on-checkpoint` command |
|
|
162
|
+
| `--on-complete` | none | Command after the job ends (`{ckpt_dir} {log_path}`, env `TRAINQ_JOB_STATUS`) |
|
|
163
|
+
|
|
164
|
+
## How scheduling works
|
|
165
|
+
|
|
166
|
+
**VRAM-based sharing (`--estimated-vram-mb`).** With an estimate, a GPU is
|
|
167
|
+
eligible when `free_vram − estimate − margin ≥ 0`, so several jobs can share
|
|
168
|
+
one card (largest-free GPU first). Without an estimate, only a fully idle GPU
|
|
169
|
+
(no processes) is used — safe, but you wait longer. A freshly launched process
|
|
170
|
+
takes time to actually claim VRAM, so trainq *reserves* the estimate on launch and
|
|
171
|
+
releases it once the process is seen holding memory, preventing a
|
|
172
|
+
double-booking OOM.
|
|
173
|
+
|
|
174
|
+
**Backfill (`--estimated-hours`).** When the top job can't get enough GPUs, trainq
|
|
175
|
+
looks for a lower job that will finish before those GPUs free up. If the
|
|
176
|
+
candidate uses GPUs the top job isn't waiting on, it can't delay the top job,
|
|
177
|
+
so it runs unconditionally.
|
|
178
|
+
|
|
179
|
+
**Priority aging.** Effective priority drops (runs sooner) the longer a job
|
|
180
|
+
waits — 1 point per 30h by default — so a low-priority job can't wait forever.
|
|
181
|
+
Aging is deliberately gentle so explicit priorities still mean something.
|
|
182
|
+
|
|
183
|
+
**ETA / progress.** `trainq list` re-parses logs on every call. Auto-detected
|
|
184
|
+
formats: fairseq2, HuggingFace Trainer, PaddleOCR (`epoch: [M/N]`),
|
|
185
|
+
`metrics.jsonl` (`iter`), and generic `step X/Y`. If none match, it falls back
|
|
186
|
+
to `--estimated-hours` (`~Xh`). The most robust option is the `progress.json`
|
|
187
|
+
convention shown in the Cookbook — it's format- and buffering-independent and
|
|
188
|
+
takes priority over log parsing.
|
|
189
|
+
|
|
190
|
+
## Checkpoint hooks (detail)
|
|
191
|
+
|
|
192
|
+
If a job sets `--on-checkpoint` or `--hf-repo-id`, the daemon launches a watcher
|
|
193
|
+
alongside it. The watcher scans `--ckpt-dir` for paths matching
|
|
194
|
+
`--checkpoint-glob` and, for each new one:
|
|
195
|
+
|
|
196
|
+
1. uploads it to `--hf-repo-id` if set (file or directory), then
|
|
197
|
+
2. runs `--on-checkpoint` if set, with `{checkpoint} {name} {step} {ckpt_dir}` substituted.
|
|
198
|
+
|
|
199
|
+
When the training process exits, the watcher does a final sweep and stops.
|
|
200
|
+
|
|
201
|
+
`--checkpoint-glob` examples: `checkpoint-*` (default, HF Trainer convention),
|
|
202
|
+
`ws_*/checkpoints/step_*`, `epoch_*`, `*.pt`.
|
|
203
|
+
|
|
204
|
+
> To avoid grabbing a checkpoint that's still being written, trainq only processes
|
|
205
|
+
> a checkpoint once it hasn't changed for `--stabilize-sec` seconds (default 10;
|
|
206
|
+
> raise it, or set 0 to disable). This is a debounce for frameworks that don't
|
|
207
|
+
> save atomically.
|
|
208
|
+
|
|
209
|
+
`--on-complete` runs once after the whole job ends, with `{ckpt_dir} {log_path}`
|
|
210
|
+
substituted and `TRAINQ_JOB_STATUS` set to `completed` or `failed`.
|
|
211
|
+
|
|
212
|
+
## Configuration
|
|
213
|
+
|
|
214
|
+
All configuration is via environment variables.
|
|
215
|
+
|
|
216
|
+
| Variable | Default | Description |
|
|
217
|
+
|----------|---------|-------------|
|
|
218
|
+
| `TRAINQ_REDIS_HOST` | localhost | Redis host |
|
|
219
|
+
| `TRAINQ_REDIS_PORT` | 6379 | Redis port |
|
|
220
|
+
| `TRAINQ_REDIS_DB` | 0 | Redis DB number |
|
|
221
|
+
| `TRAINQ_REDIS_PASSWORD` | (none) | Redis password |
|
|
222
|
+
| `TRAINQ_POLL_SEC` | 10 | Daemon poll interval (seconds) |
|
|
223
|
+
| `TRAINQ_VRAM_THRESHOLD_MB` | 1000 | Min free VRAM to call a GPU idle (MB) |
|
|
224
|
+
| `HF_TOKEN` | (none) | For HF upload; falls back to `./.env`, `~/.env`, `~/.cache/huggingface/token` |
|
|
225
|
+
|
|
226
|
+
## Deploying to a server
|
|
227
|
+
|
|
228
|
+
`scripts/deploy_trainq.sh` rsyncs the package to a remote GPU host, sets up a venv,
|
|
229
|
+
installs trainq, and (re)starts the daemon. Running jobs are re-tracked by PID, so
|
|
230
|
+
a redeploy doesn't interrupt them.
|
|
231
|
+
|
|
232
|
+
```bash
|
|
233
|
+
bash scripts/deploy_trainq.sh <host> [remote_user] [remote_dir]
|
|
234
|
+
```
|
|
235
|
+
|
|
236
|
+
Requires: `ssh <host>` access, and `python3` + `sudo` (apt) on the remote.
|
|
237
|
+
|
|
238
|
+
## Multi-user notes
|
|
239
|
+
|
|
240
|
+
The executor runs each job as the submitting user (`sudo -u` when that differs
|
|
241
|
+
from the daemon user). To use this on a shared host, the daemon user needs
|
|
242
|
+
passwordless `sudo -u` for the target users. In a single-user setup, the daemon
|
|
243
|
+
and jobs run as the same user and no sudo is involved.
|
|
244
|
+
|
|
245
|
+
## Architecture
|
|
246
|
+
|
|
247
|
+
| File | Role |
|
|
248
|
+
|------|------|
|
|
249
|
+
| `trainq/schema.py` | `Job` data model (serialize / deserialize) |
|
|
250
|
+
| `trainq/config.py` | Redis connection, key names, poll interval (`TRAINQ_*` env) |
|
|
251
|
+
| `trainq/gpu_monitor.py` | nvidia-smi parsing, VRAM reservation, availability |
|
|
252
|
+
| `trainq/scheduler.py` | Aging + backfill scheduling (effective priority) |
|
|
253
|
+
| `trainq/executor.py` | Job launch, process management, completion hook |
|
|
254
|
+
| `trainq/server.py` | Queue daemon main loop |
|
|
255
|
+
| `trainq/client.py` | CLI (submit/list/status/cancel/clear/priority) |
|
|
256
|
+
| `trainq/progress.py` | Log/file-based progress & ETA estimation |
|
|
257
|
+
| `trainq/checkpoint_watcher.py` | Checkpoint watcher CLI entry point |
|
|
258
|
+
| `trainq/watcher/daemon.py` | Detect new checkpoints, run hook, optional HF upload |
|
|
259
|
+
| `trainq/watcher/checkpoint_tracker.py` | Checkpoint detection (glob) + processed-state tracking |
|
|
260
|
+
| `trainq/watcher/hf_uploader.py` | Hugging Face Hub upload (optional) |
|
|
261
|
+
|
|
262
|
+
**Flow:** `trainq submit` → Redis sorted set `trainq:jobs` → the daemon polls and calls
|
|
263
|
+
`scheduler.pick_next_job` (run the top job if GPUs fit, else look for a backfill
|
|
264
|
+
candidate) → launch the job (and a watcher if hooks are set) → on finish, run
|
|
265
|
+
`on_complete`; on unresolved failure, leave it FAILED without running hooks.
|
|
266
|
+
|
|
267
|
+
## License
|
|
268
|
+
|
|
269
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,45 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=61"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "trainq"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Redis-based GPU job queue for ML training/eval: VRAM-aware scheduling, priority aging, backfill."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
authors = [{ name = "Korshort" }]
|
|
13
|
+
keywords = ["gpu", "job-queue", "scheduler", "machine-learning", "redis"]
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Development Status :: 4 - Beta",
|
|
16
|
+
"License :: OSI Approved :: MIT License",
|
|
17
|
+
"Operating System :: POSIX :: Linux",
|
|
18
|
+
"Programming Language :: Python :: 3",
|
|
19
|
+
"Programming Language :: Python :: 3.9",
|
|
20
|
+
"Programming Language :: Python :: 3.10",
|
|
21
|
+
"Programming Language :: Python :: 3.11",
|
|
22
|
+
"Programming Language :: Python :: 3.12",
|
|
23
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
24
|
+
"Topic :: System :: Distributed Computing",
|
|
25
|
+
]
|
|
26
|
+
dependencies = [
|
|
27
|
+
"redis>=5.0",
|
|
28
|
+
]
|
|
29
|
+
|
|
30
|
+
[project.optional-dependencies]
|
|
31
|
+
# Only needed when using HF Hub auto-backup (--hf-repo-id)
|
|
32
|
+
hf = ["huggingface_hub>=0.20"]
|
|
33
|
+
# Dev / test
|
|
34
|
+
dev = ["pytest>=7"]
|
|
35
|
+
|
|
36
|
+
[project.urls]
|
|
37
|
+
Homepage = "https://github.com/Korshort/trainq"
|
|
38
|
+
Repository = "https://github.com/Korshort/trainq"
|
|
39
|
+
Issues = "https://github.com/Korshort/trainq/issues"
|
|
40
|
+
|
|
41
|
+
[project.scripts]
|
|
42
|
+
trainq = "trainq.client:main"
|
|
43
|
+
|
|
44
|
+
[tool.setuptools]
|
|
45
|
+
packages = ["trainq", "trainq.watcher"]
|
trainq-0.1.0/setup.cfg
ADDED
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
from trainq.watcher import checkpoint_tracker as ct
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
def test_find_checkpoints_and_names(tmp_path):
|
|
5
|
+
for n in ["checkpoint-1", "checkpoint-2", "checkpoint-10"]:
|
|
6
|
+
(tmp_path / n).mkdir()
|
|
7
|
+
(tmp_path / "notes.txt").write_text("not a checkpoint")
|
|
8
|
+
|
|
9
|
+
found = ct.find_checkpoints(str(tmp_path), "checkpoint-*")
|
|
10
|
+
names = {ct.checkpoint_name(p) for p in found}
|
|
11
|
+
assert names == {"checkpoint-1", "checkpoint-2", "checkpoint-10"}
|
|
12
|
+
assert "notes.txt" not in names
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def test_extract_step():
|
|
16
|
+
assert ct.extract_step("checkpoint-500") == 500
|
|
17
|
+
assert ct.extract_step("step_42") == 42
|
|
18
|
+
assert ct.extract_step("epoch10") == 10
|
|
19
|
+
assert ct.extract_step("best") is None
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def test_state_roundtrip(tmp_path):
|
|
23
|
+
assert ct.load_state(str(tmp_path)) == {"processed": []}
|
|
24
|
+
state = {"processed": ["checkpoint-5"]}
|
|
25
|
+
ct.save_state(str(tmp_path), state)
|
|
26
|
+
assert ct.load_state(str(tmp_path))["processed"] == ["checkpoint-5"]
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def test_recent_mtime_missing(tmp_path):
|
|
30
|
+
assert ct.recent_mtime(str(tmp_path / "does-not-exist")) == 0.0
|