gpumesh 0.1.0__py3-none-any.whl

gpumesh/__init__.py

@@ -0,0 +1,3 @@
+"""gpumesh - distributed compute sharing over a mesh of volunteer machines."""
+
+__version__ = "0.1.0"

gpumesh/__main__.py

@@ -0,0 +1,5 @@
+"""Allow running gpumesh as a module: python -m gpumesh"""
+
+from .cli import main
+
+main()

gpumesh/capability.py

@@ -0,0 +1,148 @@
+"""Hardware probing and benchmarking.
+
+Detects the best compute device available and produces a single numeric
+capability score used by the scheduler. torch and psutil are optional:
+without them the probe falls back to CPU info and a pure-Python benchmark.
+
+GPU detection uses nvidia-smi (no PyTorch required for detection).
+PyTorch with CUDA is auto-installed when a GPU is detected.
+"""
+
+import os
+import platform
+import shutil
+import socket
+import subprocess
+import sys
+import time
+
+
+def _detect_nvidia_gpu() -> dict:
+    """Detect NVIDIA GPU using nvidia-smi (no PyTorch required)."""
+    if shutil.which("nvidia-smi") is None:
+        return None
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader,nounits"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+        if result.returncode == 0 and result.stdout.strip():
+            parts = result.stdout.strip().split(",")
+            name = parts[0].strip()
+            memory = parts[1].strip() if len(parts) > 1 else "unknown"
+            return {"name": name, "memory_mb": memory}
+    except (subprocess.SubprocessError, FileNotFoundError):
+        pass
+    return None
+
+
+def _install_pytorch_cuda():
+    """Install PyTorch with CUDA support if NVIDIA GPU detected."""
+    print("[mesh] NVIDIA GPU detected, installing PyTorch with CUDA support...")
+    try:
+        subprocess.run(
+            [sys.executable, "-m", "pip", "install", "torch", "torchvision", "torchaudio",
+             "--index-url", "https://download.pytorch.org/whl/cu118"],
+            check=True,
+            timeout=300,
+        )
+        print("[mesh] PyTorch with CUDA installed successfully!")
+        return True
+    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
+        print(f"[mesh] Failed to install PyTorch: {e}")
+        print("[mesh] Continuing with CPU-only mode.")
+        return False
+
+
+def probe_device() -> dict:
+    info = {
+        "hostname": socket.gethostname(),
+        "platform": platform.platform(),
+        "cpu_count": os.cpu_count() or 1,
+        "device": "cpu",
+        "device_name": platform.processor() or "cpu",
+    }
+    
+    # First check for NVIDIA GPU using nvidia-smi (fast, no PyTorch needed)
+    nvidia_gpu = _detect_nvidia_gpu()
+    if nvidia_gpu:
+        info["nvidia_gpu"] = nvidia_gpu["name"]
+        info["nvidia_memory"] = nvidia_gpu["memory_mb"]
+    
+    # Try to use PyTorch for GPU detection and benchmarking
+    try:
+        import torch
+        if torch.cuda.is_available():
+            info["device"] = "cuda"
+            info["device_name"] = torch.cuda.get_device_name(0)
+        elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+            info["device"] = "mps"
+            info["device_name"] = "Apple Silicon GPU"
+    except ImportError:
+        # PyTorch not installed - try to auto-install if NVIDIA GPU detected
+        if nvidia_gpu:
+            if _install_pytorch_cuda():
+                try:
+                    import torch
+                    if torch.cuda.is_available():
+                        info["device"] = "cuda"
+                        info["device_name"] = torch.cuda.get_device_name(0)
+                except Exception:
+                    pass
+    
+    try:
+        import psutil
+        info["ram_gb"] = round(psutil.virtual_memory().total / 1e9, 1)
+    except ImportError:
+        pass
+    
+    return info
+
+
+def _bench_torch(device: str) -> float:
+    """Time a matmul on the given torch device; return ops/sec score."""
+    import torch
+
+    n = 1024
+    a = torch.rand(n, n, device=device)
+    b = torch.rand(n, n, device=device)
+    # warmup
+    (a @ b).sum().item()
+    start = time.perf_counter()
+    iters = 10
+    for _ in range(iters):
+        c = a @ b
+    c.sum().item()  # force sync on cuda/mps
+    elapsed = time.perf_counter() - start
+    flops = 2 * n**3 * iters
+    return flops / elapsed / 1e9  # GFLOP/s
+
+
+def _bench_python() -> float:
+    """Pure-Python matmul fallback. Slow by design; scores are comparable
+    across machines because everyone runs the same loop."""
+    n = 48
+    a = [[(i * j) % 7 / 7.0 for j in range(n)] for i in range(n)]
+    b = [[(i + j) % 5 / 5.0 for j in range(n)] for i in range(n)]
+    start = time.perf_counter()
+    iters = 5
+    for _ in range(iters):
+        [[sum(a[i][k] * b[k][j] for k in range(n)) for j in range(n)] for i in range(n)]
+    elapsed = time.perf_counter() - start
+    flops = 2 * n**3 * iters
+    return flops / elapsed / 1e9
+
+
+def benchmark(device: str) -> float:
+    try:
+        return round(_bench_torch(device), 3)
+    except ImportError:
+        return round(_bench_python(), 3)
+
+
+def full_probe() -> dict:
+    info = probe_device()
+    info["score"] = benchmark(info["device"])
+    return info

gpumesh/cli.py

@@ -0,0 +1,225 @@
+"""gpumesh command line interface.
+
+  gpumesh serve      [--port 8000] [--token SECRET] [--public]
+  gpumesh join       URL [--token SECRET]
+  gpumesh quickjoin  URL --token TOKEN    # One-click setup & join
+  gpumesh submit     SCRIPT --payloads FILE [--wait] --url URL [--token SECRET]
+  gpumesh status     JOB_ID --url URL [--token SECRET]
+  gpumesh cancel     JOB_ID --url URL [--token SECRET]
+  gpumesh workers    --url URL [--token SECRET]
+"""
+
+import argparse
+import os
+import secrets
+import socket
+import subprocess
+import sys
+
+from . import client, server, tunnel, worker
+
+
+def _lan_ip() -> str:
+    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        s.connect(("8.8.8.8", 80))
+        return s.getsockname()[0]
+    except OSError:
+        return "127.0.0.1"
+    finally:
+        s.close()
+
+
+def cmd_serve(args):
+    token = args.token or secrets.token_urlsafe(12)
+    httpd = server.serve("0.0.0.0", args.port, args.db, token)
+    ip = _lan_ip()
+    print(f"[mesh] coordinator listening on 0.0.0.0:{args.port}")
+    print(f"[mesh] token: {token}")
+    print(f"[mesh] LAN join command:")
+    print(f"       gpumesh join http://{ip}:{args.port} --token {token}")
+    if args.public:
+        tunnel.open_tunnel(args.port)
+    print("[mesh] Ctrl+C to stop")
+    try:
+        httpd.serve_forever()
+    except KeyboardInterrupt:
+        print("\n[mesh] shutting down")
+        httpd.gpumesh_stop.set()
+        httpd.shutdown()
+
+
+def cmd_join(args):
+    worker.run_worker(args.url, args.token, task_timeout=args.timeout)
+
+
+def cmd_submit(args):
+    job_id = client.submit_job(args.url, args.token, args.script,
+                               args.payloads, name=args.name)
+    print(f"[client] submitted job {job_id}")
+    if args.wait:
+        job = client.wait_for_job(args.url, args.token, job_id)
+        client.print_job(job)
+    else:
+        print(f"[client] check progress: gpumesh status {job_id}"
+              f" --url {args.url} --token {args.token}")
+
+
+def cmd_status(args):
+    job = client.get_status(args.url, args.token, args.job_id)
+    client.print_job(job)
+
+
+def cmd_cancel(args):
+    result = client.cancel_job(args.url, args.token, args.job_id)
+    if result is None:
+        print(f"[client] job {args.job_id} not found")
+    else:
+        print(f"[client] cancelled job {args.job_id}")
+        print(f"  pending tasks cancelled: {result['pending']}")
+        print(f"  running tasks cancelled: {result['running']}")
+        print(f"  already finished: {result['already_finished']}")
+
+
+def cmd_workers(args):
+    from .worker import MeshClient
+
+    resp = MeshClient(args.url, args.token).call("GET", "/api/workers")
+    for w in resp["workers"]:
+        state = "alive" if w["alive"] else "dead"
+        print(f"  {w['id']}  {w['hostname']:<20} {w['device']:<5} "
+              f"score={w['score']:<10} [{state}]")
+    if not resp["workers"]:
+        print("  (no workers yet)")
+
+
+def cmd_quickjoin(args):
+    """One-click setup: install dependencies, detect GPU, and join mesh."""
+    url = args.url
+    token = args.token
+    
+    print("[mesh] Quick Join - Setting up your machine as a worker...")
+    print()
+    
+    # Step 1: Check Python version
+    print("[1/4] Checking Python version...")
+    if sys.version_info < (3, 9):
+        print("[mesh] ERROR: Python 3.9 or higher is required.")
+        print("[mesh] Please install Python 3.9+ from https://www.python.org/downloads/")
+        return
+    print(f"[mesh] Python {sys.version_info.major}.{sys.version_info.minor} ✓")
+    
+    # Step 2: Install gpumesh if not already installed
+    print("[2/4] Checking gpumesh installation...")
+    try:
+        import gpumesh
+        print("[mesh] gpumesh already installed ✓")
+    except ImportError:
+        print("[mesh] Installing gpumesh...")
+        try:
+            subprocess.run(
+                [sys.executable, "-m", "pip", "install", "-e", "."],
+                check=True,
+                timeout=120,
+            )
+            print("[mesh] gpumesh installed ✓")
+        except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
+            print(f"[mesh] ERROR: Failed to install gpumesh: {e}")
+            print("[mesh] Please install manually: pip install gpumesh")
+            return
+    
+    # Step 3: Detect hardware and install GPU support if needed
+    print("[3/4] Detecting hardware...")
+    from . import capability
+    
+    # Check for NVIDIA GPU using nvidia-smi
+    nvidia_gpu = capability._detect_nvidia_gpu()
+    if nvidia_gpu:
+        print(f"[mesh] NVIDIA GPU detected: {nvidia_gpu['name']}")
+        
+        # Check if PyTorch with CUDA is available
+        try:
+            import torch
+            if torch.cuda.is_available():
+                print(f"[mesh] PyTorch with CUDA ready ✓")
+            else:
+                print("[mesh] PyTorch installed but CUDA not available")
+                capability._install_pytorch_cuda()
+        except ImportError:
+            print("[mesh] PyTorch not found, installing with CUDA support...")
+            capability._install_pytorch_cuda()
+    else:
+        print("[mesh] No NVIDIA GPU detected, using CPU")
+    
+    # Step 4: Join the mesh
+    print("[4/4] Joining the mesh...")
+    print(f"[mesh] Connecting to {url}...")
+    print()
+    
+    # Now run the normal join command
+    worker.run_worker(url, token)
+
+
+def _add_conn_args(p):
+    p.add_argument("--url", default=os.environ.get("GPUMESH_URL", ""),
+                   help="coordinator URL (or set GPUMESH_URL)")
+    p.add_argument("--token", default=os.environ.get("GPUMESH_TOKEN", ""),
+                   help="shared auth token (or set GPUMESH_TOKEN)")
+
+
+def main():
+    ap = argparse.ArgumentParser(prog="gpumesh",
+                                 description="borrow your friends' GPUs")
+    sub = ap.add_subparsers(dest="cmd", required=True)
+
+    p = sub.add_parser("serve", help="start a coordinator on this machine")
+    p.add_argument("--port", type=int, default=8000)
+    p.add_argument("--db", default="gpumesh.db")
+    p.add_argument("--token", default="", help="auth token (random if omitted)")
+    p.add_argument("--public", action="store_true",
+                   help="expose a public URL via ngrok")
+    p.set_defaults(func=cmd_serve)
+
+    p = sub.add_parser("join", help="offer this machine's compute to a mesh")
+    p.add_argument("url", help="coordinator URL")
+    p.add_argument("--token", default=os.environ.get("GPUMESH_TOKEN", ""))
+    p.add_argument("--timeout", type=float, default=240.0,
+                   help="per-task wall clock limit in seconds")
+    p.set_defaults(func=cmd_join)
+
+    p = sub.add_parser("submit", help="submit a job to the mesh")
+    p.add_argument("script", help="python script to run for every payload")
+    p.add_argument("--payloads", required=True,
+                   help="JSON file: list of payload objects (each may set 'cost')")
+    p.add_argument("--name", default="")
+    p.add_argument("--wait", action="store_true", help="block until finished")
+    _add_conn_args(p)
+    p.set_defaults(func=cmd_submit)
+
+    p = sub.add_parser("status", help="show job progress and results")
+    p.add_argument("job_id")
+    _add_conn_args(p)
+    p.set_defaults(func=cmd_status)
+
+    p = sub.add_parser("cancel", help="cancel a running job")
+    p.add_argument("job_id")
+    _add_conn_args(p)
+    p.set_defaults(func=cmd_cancel)
+
+    p = sub.add_parser("quickjoin", help="one-click setup: install, detect GPU, and join mesh")
+    p.add_argument("url", help="coordinator URL")
+    p.add_argument("--token", required=True, help="shared auth token")
+    p.set_defaults(func=cmd_quickjoin)
+
+    p = sub.add_parser("workers", help="list workers in the mesh")
+    _add_conn_args(p)
+    p.set_defaults(func=cmd_workers)
+
+    args = ap.parse_args()
+    if hasattr(args, "url") and not args.url and args.cmd != "serve":
+        ap.error("--url required (or set GPUMESH_URL)")
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

gpumesh/client.py

@@ -0,0 +1,62 @@
+"""Client-side helpers: submit jobs, poll status, pretty-print results."""
+
+import json
+import sys
+import time
+
+from .worker import MeshClient
+
+
+def cancel_job(url: str, token: str, job_id: str) -> dict:
+    """Cancel a running job. Returns a summary of cancelled tasks."""
+    return MeshClient(url, token).call("DELETE", f"/api/jobs/{job_id}")
+
+
+def submit_job(url: str, token: str, script_path: str, payloads_path: str,
+               name: str = "") -> str:
+    with open(script_path) as f:
+        script = f.read()
+    with open(payloads_path) as f:
+        payloads = json.load(f)
+    if not isinstance(payloads, list):
+        raise SystemExit("payloads file must contain a JSON list")
+
+    mesh = MeshClient(url, token)
+    resp = mesh.call("POST", "/api/jobs", {
+        "name": name or script_path,
+        "script": script,
+        "payloads": payloads,
+    })
+    return resp["job_id"]
+
+
+def get_status(url: str, token: str, job_id: str) -> dict:
+    return MeshClient(url, token).call("GET", f"/api/jobs/{job_id}")
+
+
+def wait_for_job(url: str, token: str, job_id: str, poll: float = 3.0) -> dict:
+    while True:
+        job = get_status(url, token, job_id)
+        counts = job["counts"]
+        total = sum(counts.values())
+        done = counts.get("done", 0) + counts.get("failed", 0)
+        sys.stdout.write(f"\r[job {job_id}] {done}/{total} tasks finished "
+                         f"({counts})   ")
+        sys.stdout.flush()
+        if job["finished"]:
+            print()
+            return job
+        time.sleep(poll)
+
+
+def print_job(job: dict):
+    print(f"job: {job['name']} ({job['id']})  finished={job['finished']}")
+    for t in job["tasks"]:
+        line = f"  task {t['id']}  [{t['status']}]  cost={t['cost']}"
+        if t["worker_id"]:
+            line += f"  worker={t['worker_id']}"
+        print(line)
+        if t["result"] is not None:
+            print(f"    result: {json.dumps(t['result'])}")
+        if t["error"]:
+            print(f"    error: {t['error']}")