trainq 0.1.0__py3-none-any.whl

trainq/checkpoint_watcher.py

@@ -0,0 +1,34 @@
+"""Checkpoint watcher CLI entry point.
+
+    python -m trainq.checkpoint_watcher --ckpt-dir DIR [--on-checkpoint CMD] [--hf-repo-id REPO]
+
+The implementation lives in the trainq.watcher package.
+"""
+
+from trainq.watcher import (
+    WatchConfig,
+    checkpoint_name,
+    extract_step,
+    find_checkpoints,
+    load_state,
+    save_state,
+    upload_checkpoint,
+    watch,
+)
+from trainq.watcher.daemon import main
+
+__all__ = [
+    "WatchConfig",
+    "watch",
+    "find_checkpoints",
+    "checkpoint_name",
+    "extract_step",
+    "load_state",
+    "save_state",
+    "upload_checkpoint",
+    "main",
+]
+
+
+if __name__ == "__main__":
+    main()

trainq/client.py

@@ -0,0 +1,383 @@
+"""Job queue client CLI.
+
+Usage:
+    python -m trainq.client submit --command "python train.py ..." \
+        --ckpt-dir /data/ckpt --log-path /data/log.txt \
+        --gpu-count 2 --priority 3
+
+    python -m trainq.client list
+    python -m trainq.client status <job_id>
+    python -m trainq.client cancel <job_id>
+    python -m trainq.client clear
+"""
+
+from __future__ import annotations
+
+import argparse
+import getpass
+import sys
+
+import redis
+
+from trainq.config import (
+    REDIS_HOST, REDIS_PORT, REDIS_DB, REDIS_PASSWORD,
+    QUEUE_KEY, JOB_DETAIL_KEY, RUNNING_KEY,
+)
+from datetime import datetime, timezone
+
+from trainq.progress import parse_training_progress
+from trainq.schema import Job, JobStatus
+
+AGING_HOURS_PER_POINT = 30.0  # -1 per 30h waited (matches scheduler.py)
+
+
+def _redis_client() -> redis.Redis:
+    return redis.Redis(
+        host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB,
+        password=REDIS_PASSWORD, decode_responses=True,
+    )
+
+
+# ── Commands ──────────────────────────────────────────────
+
+def cmd_submit(args: argparse.Namespace) -> None:
+    r = _redis_client()
+
+    # warn about a missing distributed launcher for multi-GPU (torchrun/accelerate/deepspeed/etc.)
+    _LAUNCHERS = ("torchrun", "accelerate", "deepspeed", "mpirun", "torch.distributed")
+    if args.gpu_count >= 2 and not any(l in args.command for l in _LAUNCHERS):
+        print("  ⚠️  gpu-count >= 2 but no distributed launcher found in the command.")
+        print("     Multi-GPU usually needs torchrun --nproc_per_node=N (or accelerate/deepspeed).")
+        print("     e.g. torchrun --nproc_per_node=2 train.py ...")
+        print("     Add --force to submit anyway. Aborting.")
+        if not getattr(args, 'force', False):
+            sys.exit(1)
+
+    gpu_ids = None
+    if args.gpu_ids:
+        gpu_ids = [int(x.strip()) for x in args.gpu_ids.split(",")]
+
+    if not args.description.strip():
+        print("  ❌  --description is required. Briefly describe the job.")
+        sys.exit(1)
+
+    job = Job(
+        command=args.command,
+        ckpt_dir=args.ckpt_dir,
+        log_path=args.log_path,
+        description=args.description.strip(),
+        gpu_count=args.gpu_count,
+        gpu_ids=gpu_ids,
+        priority=args.priority,
+        max_retries=args.max_retries,
+        estimated_hours=args.estimated_hours,
+        estimated_vram_mb=args.estimated_vram_mb,
+        checkpoint_glob=args.checkpoint_glob,
+        on_checkpoint=args.on_checkpoint,
+        hf_repo_id=args.hf_repo_id,
+        watch_gpu=args.watch_gpu,
+        on_complete=args.on_complete,
+        user=getpass.getuser(),
+    )
+    r.set(f"{JOB_DETAIL_KEY}:{job.job_id}", job.to_json())
+    r.zadd(QUEUE_KEY, {job.job_id: job.priority})
+    parts = [f"Submitted: {job.job_id}  priority={job.priority}  gpu={args.gpu_count}"]
+    if args.estimated_hours > 0:
+        parts.append(f"est={args.estimated_hours}h")
+    if args.estimated_vram_mb > 0:
+        parts.append(f"vram={args.estimated_vram_mb}MB")
+    if args.hf_repo_id:
+        parts.append(f"hf={args.hf_repo_id}")
+    if args.on_checkpoint:
+        parts.append("hook=on_checkpoint")
+    print("  ".join(parts))
+    if args.estimated_vram_mb == 0:
+        print("  ⚠️  estimated-vram-mb not set → only a fully idle GPU is used (may wait longer)")
+
+
+def cmd_list(args: argparse.Namespace) -> None:
+    r = _redis_client()
+
+    # pending (sorted set)
+    pending_ids = r.zrange(QUEUE_KEY, 0, -1, withscores=True)
+    running_ids = r.smembers(RUNNING_KEY)
+
+    # collect all job keys
+    all_keys = r.keys(f"{JOB_DETAIL_KEY}:*")
+    jobs: list[Job] = []
+    for k in all_keys:
+        raw = r.get(k)
+        if raw:
+            jobs.append(Job.from_json(raw))
+
+    if not jobs:
+        print("No jobs.")
+        return
+
+    # effective priority — frozen at started_at/finished_at once set.
+    # while running: wait fixed by started_at. When finished (completed/failed/
+    # cancelled): allow finished_at as the snapshot end in case it ended without started_at.
+    def _wait_hours(j: Job) -> float:
+        if j.started_at is not None:
+            return (
+                datetime.fromisoformat(j.started_at) - datetime.fromisoformat(j.created_at)
+            ).total_seconds() / 3600
+        if j.finished_at is not None:
+            return (
+                datetime.fromisoformat(j.finished_at) - datetime.fromisoformat(j.created_at)
+            ).total_seconds() / 3600
+        return (datetime.now(timezone.utc) - datetime.fromisoformat(j.created_at)).total_seconds() / 3600
+
+    def eff_pri(j: Job) -> float:
+        return j.priority - (_wait_hours(j) / AGING_HOURS_PER_POINT)
+
+    def wait_h(j: Job) -> float:
+        return _wait_hours(j)
+
+    def elapsed_str(j: Job) -> str:
+        """Elapsed time. Up to now while running, frozen at finished_at once ended."""
+        if j.started_at is None:
+            return "-"
+        end = (
+            datetime.fromisoformat(j.finished_at)
+            if j.finished_at is not None
+            else datetime.now(timezone.utc)
+        )
+        elapsed = (end - datetime.fromisoformat(j.started_at)).total_seconds() / 3600
+        if elapsed < 1:
+            return f"{elapsed * 60:.0f}m"
+        return f"{elapsed:.1f}h"
+
+    def progress_str(j: Job) -> str:
+        """Log-parsed progress + ETA."""
+        if j.status != JobStatus.RUNNING or j.started_at is None:
+            return "-"
+        elapsed = (datetime.now(timezone.utc) - datetime.fromisoformat(j.started_at)).total_seconds() / 3600
+        prog = parse_training_progress(j.log_path, j.command, j.ckpt_dir)
+        if prog is not None and prog.progress > 0:
+            eta = prog.estimate_remaining_hours(elapsed)
+            return f"{prog.progress:.0%} {eta:.1f}h"
+        # fallback: estimated_hours based
+        if j.estimated_hours > 0:
+            rem = max(0, j.estimated_hours - elapsed)
+            return f"~{rem:.1f}h"
+        return "?"
+
+    # sort by status
+    order = {JobStatus.RUNNING: 0, JobStatus.PENDING: 1, JobStatus.FAILED: 2, JobStatus.COMPLETED: 3, JobStatus.CANCELLED: 4}
+    jobs.sort(key=lambda j: (order.get(j.status, 9), eff_pri(j)))
+
+    def gpu_idx(j):
+        # actually assigned GPUs (running) first, else the allow-set (gpu_ids), else all ("all")
+        if j.assigned_gpus:
+            return ",".join(str(g) for g in j.assigned_gpus)
+        if j.gpu_ids:
+            return ",".join(str(g) for g in j.gpu_ids)
+        return "all"
+
+    DESC_WIDTH = 60
+    fmt = "{:<14} {:<10} {:>5} {:>4} {:>8} {:>6} {:>5} {:>6} {:>8} {:<9} {}"
+    print(fmt.format("JOB_ID", "STATUS", "EFF", "GPUs", "GPUIDX", "VRAM", "WAIT", "ELAPS", "ETA", "USER", "DESC"))
+    print("-" * (87 + DESC_WIDTH))
+    for j in jobs:
+        ep = f"{eff_pri(j):.1f}"
+        wh = f"{wait_h(j):.1f}h"
+        vram = f"{j.estimated_vram_mb/1024:.0f}GB" if j.estimated_vram_mb > 0 else "-"
+        if j.status == JobStatus.RUNNING:
+            elaps = elapsed_str(j)
+            eta = progress_str(j)
+        elif j.status == JobStatus.COMPLETED:
+            elaps = elapsed_str(j)
+            eta = "-"
+        else:  # pending, failed, cancelled
+            elaps = "-"
+            eta = f"{j.estimated_hours:.1f}h" if j.estimated_hours > 0 else "-"
+        desc = j.description[:DESC_WIDTH]
+        print(fmt.format(j.job_id, j.status.value, ep, j.gpu_count, gpu_idx(j), vram, wh, elaps, eta, (j.user or "-"), desc))
+
+
+def cmd_status(args: argparse.Namespace) -> None:
+    r = _redis_client()
+    raw = r.get(f"{JOB_DETAIL_KEY}:{args.job_id}")
+    if raw is None:
+        print(f"Job {args.job_id} not found.")
+        sys.exit(1)
+    job = Job.from_json(raw)
+    for field_name in [
+        "job_id", "status", "description", "command", "gpu_count", "gpu_ids", "assigned_gpus",
+        "priority", "estimated_hours", "estimated_vram_mb", "max_retries", "retry_count",
+        "ckpt_dir", "log_path",
+        "created_at", "started_at", "finished_at", "error_msg",
+    ]:
+        val = getattr(job, field_name)
+        if isinstance(val, JobStatus):
+            val = val.value
+        print(f"  {field_name:16s}: {val}")
+
+
+def cmd_cancel(args: argparse.Namespace) -> None:
+    r = _redis_client()
+    raw = r.get(f"{JOB_DETAIL_KEY}:{args.job_id}")
+    if raw is None:
+        print(f"Job {args.job_id} not found.")
+        sys.exit(1)
+    job = Job.from_json(raw)
+    if job.status != JobStatus.PENDING:
+        print(f"Cannot cancel: status is {job.status.value}")
+        sys.exit(1)
+    job.status = JobStatus.CANCELLED
+    r.set(f"{JOB_DETAIL_KEY}:{job.job_id}", job.to_json())
+    r.zrem(QUEUE_KEY, job.job_id)
+    print(f"Cancelled: {job.job_id}")
+
+
+def cmd_clear(args: argparse.Namespace) -> None:
+    r = _redis_client()
+    all_keys = r.keys(f"{JOB_DETAIL_KEY}:*")
+    removed = 0
+    for k in all_keys:
+        raw = r.get(k)
+        if raw:
+            job = Job.from_json(raw)
+            if job.status in (JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED):
+                r.delete(k)
+                removed += 1
+    print(f"Cleared {removed} finished jobs.")
+
+
+def cmd_priority(args: argparse.Namespace) -> None:
+    r = _redis_client()
+    raw = r.get(f"{JOB_DETAIL_KEY}:{args.job_id}")
+    if raw is None:
+        print(f"Job {args.job_id} not found.")
+        sys.exit(1)
+    job = Job.from_json(raw)
+    if job.status != JobStatus.PENDING:
+        print(f"Cannot change priority: status is {job.status.value} (only pending can be changed)")
+        sys.exit(1)
+    old = job.priority
+    job.priority = args.priority
+    r.set(f"{JOB_DETAIL_KEY}:{job.job_id}", job.to_json())
+    r.zadd(QUEUE_KEY, {job.job_id: args.priority})
+    print(f"Priority updated: {job.job_id}  {old} -> {args.priority}")
+
+
+# ── Main ──────────────────────────────────────────────────
+
+USAGE = """\
+trainq — GPU job queue client
+
+Usage:
+  trainq list                       list all jobs
+  trainq status <job_id>            show job details
+  trainq submit --command "..." ... submit a job
+  trainq cancel <job_id>            cancel a pending job
+  trainq priority <job_id> <val>    change a pending job's priority (float ok; lower runs first)
+  trainq clear                      remove completed/failed/cancelled jobs
+
+Minimal submit:
+  trainq submit --command "..." --ckpt-dir <dir> --log-path <path> --gpu-count 1 \\
+    --description "resnet50 finetune lr1e-4"
+
+Recommended submit (enables backfill + GPU sharing):
+  trainq submit --command "..." --ckpt-dir <dir> --log-path <path> \\
+    --gpu-count 1 --estimated-hours 2.0 --estimated-vram-mb 15000 \\
+    --description "eval sweep seed=0"
+
+  ⚠️  estimated-vram-mb not set → only a fully idle GPU is used (may wait longer)
+  ⚠️  gpu-count >= 2 → the command needs a distributed launcher (torchrun, etc.)
+  ⚠️  --gpu-ids is an allow-set (not a pin) → trainq picks --gpu-count from it.
+      e.g. --gpu-ids 2,3 --gpu-count 1 = GPU 2 or 3, whichever has VRAM room (not both).
+      To restrict, pass it per submit. There is no global setting.
+
+Checkpoint hooks (checkpoint watcher):
+  trainq submit ... --checkpoint-glob 'checkpoint-*' \\
+    --on-checkpoint 'python eval.py {checkpoint}' --watch-gpu 2
+  trainq submit ... --hf-repo-id <user>/<repo>   # auto-back up each new checkpoint to the HF Hub
+
+Live ETA (the ETA column in `list`: progress + remaining time):
+  Progress is read from logs/files. Auto-detected formats: fairseq2, HF Trainer,
+  PaddleOCR(epoch), metrics.jsonl(iter), 'step X/Y'. If none match, it estimates from --estimated-hours ('~Xh').
+  ★For format-agnostic real-time ETA: have the training script periodically write
+    {"step": <current>, "total": <total>} to {--ckpt-dir}/progress.json (buffering-independent, most robust).
+"""
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        prog="trainq",
+        description="GPU job queue client",
+        usage=USAGE,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    sub = parser.add_subparsers(dest="action")
+
+    # print usage when run without arguments
+    if len(sys.argv) == 1:
+        print(USAGE)
+        sys.exit(0)
+
+    # submit
+    p_sub = sub.add_parser("submit", help="Submit a job")
+    p_sub.add_argument("--command", required=True, help="Command to execute")
+    p_sub.add_argument("--ckpt-dir", required=True, help="Checkpoint / output root directory")
+    p_sub.add_argument("--log-path", required=True, help="Log file path")
+    p_sub.add_argument("--description", required=True, help="Job description, shown in `trainq list`")
+    p_sub.add_argument("--gpu-count", type=int, default=1, help="Number of GPUs needed (default: 1)")
+    p_sub.add_argument("--gpu-ids", default=None, help="Allowed GPU set, NOT forced: --gpu-count picks that many from it (VRAM-aware). Comma-separated, e.g. '2,3' with --gpu-count 1 uses GPU 2 or 3, whichever has room. No global setting; pass per submit to restrict")
+    p_sub.add_argument(
+        "--priority", type=int, default=5,
+        help=(
+            "Priority — 0 (highest) .. 9 (lowest) at submit, default: 5. "
+            "The effective priority drops below 0 (negative) via aging "
+            "(-1 per 30h waited, starvation guard — gentle so explicit priorities stay meaningful)."
+        ),
+    )
+    p_sub.add_argument("--max-retries", type=int, default=1, help="Max retry count on failure (default: 1)")
+    p_sub.add_argument("--estimated-hours", type=float, default=0.0, help="Estimated time in hours (optional, 0=unknown → backfill disabled)")
+    p_sub.add_argument("--estimated-vram-mb", type=int, default=0, help="Estimated VRAM per GPU in MB (optional, 0=require fully free GPU)")
+    p_sub.add_argument("--checkpoint-glob", default=None,
+                        help="Glob for new checkpoints (under ckpt_dir), e.g. 'checkpoint-*', 'ws_*/checkpoints/step_*'")
+    p_sub.add_argument("--on-checkpoint", default=None,
+                        help="Shell command per new checkpoint ({checkpoint} {name} {step} {ckpt_dir} substituted)")
+    p_sub.add_argument("--hf-repo-id", default=None,
+                        help="If set, auto-upload each new checkpoint to this HF Hub repo (requires HF_TOKEN)")
+    p_sub.add_argument("--watch-gpu", type=int, default=None,
+                        help="Dedicated GPU for the on_checkpoint command (separate from job GPUs)")
+    p_sub.add_argument("--on-complete", default=None,
+                        help="Shell command after the job ends ({ckpt_dir} {log_path} substituted, env TRAINQ_JOB_STATUS)")
+    p_sub.add_argument("--force", action="store_true",
+                        help="Skip validation checks and submit anyway")
+
+    # list
+    sub.add_parser("list", help="List all jobs")
+
+    # status
+    p_st = sub.add_parser("status", help="Show job details")
+    p_st.add_argument("job_id")
+
+    # cancel
+    p_ca = sub.add_parser("cancel", help="Cancel a pending job")
+    p_ca.add_argument("job_id")
+
+    # clear
+    sub.add_parser("clear", help="Remove completed/failed/cancelled jobs")
+
+    # priority
+    p_pri = sub.add_parser("priority", help="Change a pending job's priority")
+    p_pri.add_argument("job_id")
+    p_pri.add_argument("priority", type=float, help="New priority (float ok, e.g. 5.5). Lower runs first")
+
+    args = parser.parse_args()
+    {
+        "submit": cmd_submit,
+        "list": cmd_list,
+        "status": cmd_status,
+        "cancel": cmd_cancel,
+        "clear": cmd_clear,
+        "priority": cmd_priority,
+    }[args.action](args)
+
+
+if __name__ == "__main__":
+    main()

trainq/config.py

@@ -0,0 +1,21 @@
+"""Queue system configuration constants."""
+
+import os
+
+# Redis
+REDIS_HOST = os.getenv("TRAINQ_REDIS_HOST", "localhost")
+REDIS_PORT = int(os.getenv("TRAINQ_REDIS_PORT", "6379"))
+REDIS_DB = int(os.getenv("TRAINQ_REDIS_DB", "0"))
+REDIS_PASSWORD = os.getenv("TRAINQ_REDIS_PASSWORD", None)
+
+# Redis key prefix
+KEY_PREFIX = "trainq"
+QUEUE_KEY = f"{KEY_PREFIX}:jobs"          # sorted set (priority)
+JOB_DETAIL_KEY = f"{KEY_PREFIX}:job"      # hash per job  → {KEY_PREFIX}:job:{job_id}
+RUNNING_KEY = f"{KEY_PREFIX}:running"     # set of running job_ids
+
+# Server polling
+POLL_INTERVAL_SEC = int(os.getenv("TRAINQ_POLL_SEC", "10"))
+
+# GPU
+GPU_VRAM_FREE_THRESHOLD_MB = int(os.getenv("TRAINQ_VRAM_THRESHOLD_MB", "1000"))