From b0b04034d7c58cb937e6e33481944aa6c0acca37 Mon Sep 17 00:00:00 2001
From: Zachery Aaron Shores-Chmielewski <zachanon@gmail.com>
Date: Wed, 4 Feb 2026 16:07:32 +0700
Subject: [PATCH] init

Toolkit for interaction with vastai instances. Eventual framework for low price spot distributed training.
---
 README.md |  57 ++++++++
 dtrain    | 405 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 462 insertions(+)
 create mode 100644 README.md
 create mode 100755 dtrain

diff --git a/README.md b/README.md
new file mode 100644
index 0000000..f52d5f8
--- /dev/null
+++ b/README.md
@@ -0,0 +1,57 @@
+# dtrain
+
+CLI for distributed training on vast.ai. Provision GPU nodes, deploy scripts, and monitor training with simple commands.
+
+## Why
+
+Training on multiple GPUs across multiple machines shouldn't require clicking through web UIs or juggling SSH sessions. dtrain wraps the vast.ai CLI to provide a streamlined workflow: rent nodes, deploy code, monitor progress, tear down.
+
+## Requirements
+
+- Python 3
+- [vast.ai CLI](https://vast.ai/docs/cli/overview) configured with your API key
+- SSH key registered with vast.ai
+
+## Usage
+
+```bash
+# Find available GPUs (sorted by price)
+dtrain search --vram 8
+
+# Rent 2 cheapest instances with 8GB+ VRAM
+dtrain rent -n 2 --vram 8
+
+# Check instance status (wait for SSH to be ready)
+dtrain status
+
+# Deploy and run your training script
+dtrain deploy --script train.py
+
+# Monitor training progress
+dtrain ps
+
+# Run arbitrary commands on all nodes
+dtrain run "nvidia-smi"
+dtrain run "tail -20 /workspace/train.log"
+
+# Tear down when done
+dtrain destroy
+```
+
+## Commands
+
+| Command | Description |
+|---------|-------------|
+| `search` | Find GPU offers, sorted by price |
+| `rent` | Rent instances by ID or auto-select N cheapest |
+| `status` | Show instance info and SSH commands |
+| `ps` | Show training status (log age + last output) |
+| `deploy` | Copy script to nodes and run in background |
+| `run` | Run command on all nodes |
+| `destroy` | Tear down all instances |
+
+## How it works
+
+- Scripts are deployed to `/workspace/` and run via `nohup` with output to `<script>.log`
+- `dtrain ps` reads the most recent log file to show what's happening
+- All state comes from vast.ai - dtrain doesn't store anything locally
diff --git a/dtrain b/dtrain
new file mode 100755
index 0000000..c86f13f
--- /dev/null
+++ b/dtrain
@@ -0,0 +1,405 @@
+#!/usr/bin/env python3
+"""
+dtrain - CLI for distributed training on vast.ai
+
+Commands:
+    search   - Find cheap GPU offers, sorted by price
+    rent     - Rent instances by ID or auto-select N cheapest
+    status   - Show instance info (ID, GPU, status, SSH command)
+    ps       - Show training status (log age + last output line)
+    deploy   - Copy script to all instances and run in background
+    run      - Run arbitrary command on all instances
+    destroy  - Tear down all instances
+
+Requires: vastai CLI configured with API key, SSH key registered with vast.ai
+"""
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+
+
+def run_vast(*args):
+    """Run a vastai CLI command and return the result."""
+    cmd = ["vastai"] + list(args)
+    result = subprocess.run(cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        print(f"Error: {result.stderr}", file=sys.stderr)
+        return None
+    return result.stdout
+
+
+# ============ COMMANDS ============
+
+def cmd_search(args):
+    """Find cheap GPU offers."""
+    query = f"gpu_ram>={args.vram} inet_down>=100"
+    output = run_vast("search", "offers", query, "-o", "dph_total", "--limit", str(args.limit), "--raw")
+    if not output:
+        return 1
+
+    offers = json.loads(output)
+    print(f"{'ID':<12} {'GPU':<16} {'$/hr':<8} {'VRAM':<6} {'Location':<20}")
+    print("-" * 70)
+    for o in offers:
+        print(f"{o['id']:<12} {o['gpu_name']:<16} ${o['dph_total']:<7.4f} {o['gpu_ram']/1024:.0f}GB   {o['geolocation'][:20]:<20}")
+    return 0
+
+
+def cmd_rent(args):
+    """Rent instances by ID or pick n cheapest."""
+    # If --cheap N is set, find the N cheapest offers
+    if args.cheap:
+        query = f"gpu_ram>={args.vram} inet_down>=100"
+        output = run_vast("search", "offers", query, "-o", "dph_total", "--limit", str(args.cheap), "--raw")
+        if not output:
+            return 1
+        offers = json.loads(output)
+        offer_ids = [o['id'] for o in offers]
+        total = sum(o['dph_total'] for o in offers)
+        print(f"Selected {len(offers)} cheapest:")
+        for o in offers:
+            print(f"  {o['id']} - {o['gpu_name']:<16} ${o['dph_total']:.4f}/hr  {o['geolocation'][:20]}")
+        print(f"Total: ${total:.4f}/hr")
+    else:
+        offer_ids = args.ids
+
+    if not offer_ids:
+        print("Error: provide IDs or use --cheap N")
+        return 1
+
+    print(f"\nRenting {len(offer_ids)} instance(s)...")
+    for oid in offer_ids:
+        output = run_vast("create", "instance", str(oid),
+                          "--image", args.image,
+                          "--disk", str(args.disk))
+        if output:
+            print(output.strip())
+
+    print("\nRun 'dtrain status' to check instance status.")
+    return 0
+
+
+def cmd_status(args):
+    """Show running instances."""
+    output = run_vast("show", "instances", "--raw")
+    if not output:
+        print("No instances or error fetching.")
+        return 1
+
+    instances = json.loads(output)
+    if not instances:
+        print("No running instances.")
+        return 0
+
+    print(f"{'ID':<10} {'GPU':<16} {'Status':<12} {'$/hr':<8} {'SSH'}")
+    print("-" * 80)
+    for inst in instances:
+        status = inst.get('actual_status', 'unknown')
+        gpu = inst.get('gpu_name', 'N/A')
+        cost = inst.get('dph_total') or 0
+        ssh_host = inst.get('ssh_host', '')
+        ssh_port = inst.get('ssh_port', '')
+        ssh_cmd = f"ssh -p {ssh_port} root@{ssh_host}" if ssh_host else "(not ready)"
+        print(f"{inst['id']:<10} {gpu:<16} {status:<12} ${cost:<7.4f} {ssh_cmd}")
+    return 0
+
+
+def get_instances():
+    """Get list of running instances."""
+    output = run_vast("show", "instances", "--raw")
+    if not output:
+        return []
+    return json.loads(output)
+
+
+def run_on_instance(inst, command, timeout=10):
+    """Run a command on an instance via SSH. Returns (returncode, output) or (None, msg) if not ready."""
+    ssh_host = inst.get('ssh_host', '')
+    ssh_port = inst.get('ssh_port', '')
+    if not ssh_host:
+        return None, "not ready"
+
+    # BatchMode=yes prevents password prompts from hanging
+    cmd = ["ssh", "-p", str(ssh_port),
+           "-o", "StrictHostKeyChecking=no",
+           "-o", "ConnectTimeout=5",
+           "-o", "BatchMode=yes",
+           f"root@{ssh_host}", command]
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
+        return result.returncode, result.stdout + result.stderr
+    except subprocess.TimeoutExpired:
+        return 1, "timeout"
+
+
+def cmd_deploy(args):
+    """Deploy a script to all instances, verify, and run it."""
+    instances = get_instances()
+    if not instances:
+        print("No running instances.")
+        return 1
+
+    filename = os.path.basename(args.script)
+    dest_path = args.dest.rstrip('/') + '/' + filename
+
+    ready_instances = []
+
+    # Phase 1: Copy script to all instances
+    print("Copying script to instances...")
+    for inst in instances:
+        ssh_host = inst.get('ssh_host', '')
+        ssh_port = inst.get('ssh_port', '')
+        if not ssh_host:
+            print(f"  {inst['id']} - not ready, skipping")
+            continue
+
+        # Copy file
+        dest = f"root@{ssh_host}:{args.dest}"
+        cmd = ["scp", "-P", str(ssh_port), "-o", "StrictHostKeyChecking=no", args.script, dest]
+        print(f"  {inst['id']} - copying {args.script}...", end=" ")
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        if result.returncode != 0:
+            print(f"FAILED: {result.stderr.strip()}")
+            continue
+
+        # Verify
+        ret, out = run_on_instance(inst, f"ls -la {dest_path}")
+        if ret == 0:
+            print("OK")
+            ready_instances.append((inst, dest_path))
+        else:
+            print(f"FAILED (verify): {out.strip()}")
+
+    if not ready_instances:
+        print("No instances ready for execution.")
+        return 1
+
+    if args.no_run:
+        return 0
+
+    # Phase 2: Run script on all instances (in background)
+    print(f"\nStarting {filename} on {len(ready_instances)} instance(s)...")
+    for inst, path in ready_instances:
+        logfile = filename.replace('.py', '.log')
+        run_cmd = f"cd {args.dest} && nohup python {filename} > {logfile} 2>&1 &"
+        ret, out = run_on_instance(inst, run_cmd)
+        if ret is None:
+            print(f"  {inst['id']} - not ready")
+        elif ret == 0:
+            print(f"  {inst['id']} - started (log: {args.dest}{logfile})")
+        else:
+            print(f"  {inst['id']} - FAILED: {out.strip()}")
+
+    # Summary
+    total = len(instances)
+    skipped = total - len(ready_instances)
+    print(f"\nSummary: {len(ready_instances)}/{total} instances deployed")
+    if skipped:
+        print(f"         {skipped} skipped/failed")
+
+    return 0
+
+
+def cmd_ps(args):
+    """Show training status by reading log files on each node."""
+    instances = get_instances()
+    if not instances:
+        print("No running instances.")
+        return 1
+
+    # Print header immediately
+    print(f"{'ID':<10} {'GPU':<14} {'Status':<10} {'Log age':<10} {'Last output'}")
+    print("-" * 100)
+
+    for inst in instances:
+        iid = inst['id']
+        gpu = inst.get('gpu_name', 'N/A')[:14]
+        status = inst.get('actual_status', 'unknown')
+
+        if not inst.get('ssh_host'):
+            print(f"{iid:<10} {gpu:<14} {status:<10} {'-':<10} (no ssh)")
+            continue
+
+        # Show we're checking this node
+        print(f"{iid:<10} {gpu:<14} {status:<10} ", end="", flush=True)
+
+        # Find most recent .log file, output format: "seconds_since_modified|last_line"
+        log_cmd = 'log=$(ls -t /workspace/*.log 2>/dev/null | head -1); if [ -n "$log" ]; then age=$(($(date +%s) - $(stat -c %Y "$log"))); echo "$age|$(tail -1 "$log" | head -c 60)"; else echo "NO_LOG|"; fi'
+        ret, out = run_on_instance(inst, log_cmd)
+
+        # Check for SSH/connection errors
+        error_phrases = ['connection refused', 'no route to host', 'connection timed out',
+                        'permission denied', 'host key verification failed', 'network is unreachable', 'timeout']
+        out_lower = (out or '').lower()
+        is_conn_error = any(err in out_lower for err in error_phrases)
+
+        if is_conn_error:
+            err_msg = out.strip().split('\n')[0][:40] if out else 'SSH error'
+            print(f"{'-':<10} ({err_msg})")
+            continue
+
+        # Parse "age|last_line" - skip vast.ai SSH banner lines by looking for our format
+        result_line = None
+        for line in (out or '').strip().split('\n'):
+            if '|' in line and (line[0].isdigit() or line.startswith('NO_LOG')):
+                result_line = line
+                break
+
+        if result_line:
+            age_str, last_line = result_line.split('|', 1)
+            if age_str == 'NO_LOG':
+                print(f"{'-':<10} (no log file)")
+            else:
+                # Convert raw seconds to friendly format
+                if age_str.endswith('s'):
+                    age_str = age_str  # already formatted
+                elif age_str.isdigit():
+                    secs = int(age_str)
+                    if secs < 60:
+                        age_str = f"{secs}s ago"
+                    elif secs < 3600:
+                        age_str = f"{secs//60}m ago"
+                    else:
+                        age_str = f"{secs//3600}h ago"
+                print(f"{age_str:<10} {last_line}")
+        else:
+            print(f"{'-':<10} (no output)")
+
+    return 0
+
+
+def cmd_run(args):
+    """Run a command on all instances."""
+    instances = get_instances()
+    if not instances:
+        print("No running instances.")
+        return 1
+
+    command = " ".join(args.command)
+
+    # Track results for summary
+    results = {'ok': 0, 'failed': 0, 'not_ready': 0}
+
+    for inst in instances:
+        iid = inst['id']
+        gpu = inst.get('gpu_name', 'N/A')
+        status = inst.get('actual_status', 'unknown')
+        ssh_host = inst.get('ssh_host', '')
+
+        # Header with instance info
+        header = f"[{iid}] {gpu} ({status})"
+
+        if not ssh_host:
+            print(f"{header} -- SKIPPED: {status}, waiting for SSH")
+            results['not_ready'] += 1
+            continue
+
+        ret, out = run_on_instance(inst, command)
+
+        if ret is None:
+            print(f"{header} -- SKIPPED: SSH not available")
+            results['not_ready'] += 1
+        elif ret == 0:
+            print(f"{header} -- OK")
+            if out.strip():
+                # Indent output for readability
+                for line in out.strip().split('\n'):
+                    print(f"    {line}")
+            results['ok'] += 1
+        else:
+            print(f"{header} -- FAILED (exit {ret})")
+            if out.strip():
+                for line in out.strip().split('\n'):
+                    print(f"    {line}")
+            results['failed'] += 1
+
+    # Summary
+    total = len(instances)
+    print(f"\n--- {results['ok']}/{total} ok", end="")
+    if results['failed']:
+        print(f", {results['failed']} failed", end="")
+    if results['not_ready']:
+        print(f", {results['not_ready']} not ready", end="")
+    print(" ---")
+
+    return 0 if results['failed'] == 0 else 1
+
+
+def cmd_destroy(args):
+    """Destroy instances."""
+    output = run_vast("show", "instances", "--raw")
+    if not output:
+        return 1
+
+    instances = json.loads(output)
+    if not instances:
+        print("No instances to destroy.")
+        return 0
+
+    print(f"Destroying {len(instances)} instance(s)...")
+    for inst in instances:
+        print(f"  {inst['id']} - {inst.get('gpu_name', 'N/A')}")
+        run_vast("destroy", "instance", str(inst['id']))
+
+    print("Done.")
+    return 0
+
+
+# ============ MAIN ============
+
+def main():
+    parser = argparse.ArgumentParser(prog="dtrain", description="CLI for distributed training on vast.ai")
+    sub = parser.add_subparsers(dest="cmd")
+
+    # search
+    p = sub.add_parser("search", help="Find cheap GPU offers")
+    p.add_argument("--vram", type=int, default=8, help="Min VRAM in GB (default: 8)")
+    p.add_argument("--limit", type=int, default=10, help="Max results (default: 10)")
+    p.set_defaults(func=cmd_search)
+
+    # rent
+    p = sub.add_parser("rent", help="Rent instances by ID or pick cheapest")
+    p.add_argument("ids", type=int, nargs="*", help="Offer ID(s) to rent")
+    p.add_argument("-n", "--cheap", type=int, help="Rent N cheapest offers instead")
+    p.add_argument("--vram", type=int, default=8, help="Min VRAM for --cheap (default: 8)")
+    p.add_argument("--image", default="pytorch/pytorch:latest", help="Docker image")
+    p.add_argument("--disk", type=int, default=20, help="Disk GB (default: 20)")
+    p.set_defaults(func=cmd_rent)
+
+    # status
+    p = sub.add_parser("status", help="Show running instances")
+    p.set_defaults(func=cmd_status)
+
+    # ps
+    p = sub.add_parser("ps", help="Show training status (log age + last line)")
+    p.set_defaults(func=cmd_ps)
+
+    # deploy
+    p = sub.add_parser("deploy", help="Deploy script to all instances and run it")
+    p.add_argument("--script", required=True, help="Script to deploy")
+    p.add_argument("--dest", default="/workspace/", help="Destination path (default: /workspace/)")
+    p.add_argument("--no-run", action="store_true", help="Only copy, don't run the script")
+    p.set_defaults(func=cmd_deploy)
+
+    # run
+    p = sub.add_parser("run", help="Run command on all instances")
+    p.add_argument("command", nargs="+", help="Command to run")
+    p.set_defaults(func=cmd_run)
+
+    # destroy
+    p = sub.add_parser("destroy", help="Destroy all instances")
+    p.set_defaults(func=cmd_destroy)
+
+    args = parser.parse_args()
+    if not args.cmd:
+        parser.print_help()
+        return 0
+
+    return args.func(args)
+
+
+if __name__ == "__main__":
+    sys.exit(main())