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 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,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -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