gpumesh 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.
gpumesh-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,152 @@
1
+ Metadata-Version: 2.4
2
+ Name: gpumesh
3
+ Version: 0.1.0
4
+ Summary: Borrow your friends' GPUs: a terminal-based distributed compute mesh in pure Python
5
+ License: MIT
6
+ Requires-Python: >=3.9
7
+ Description-Content-Type: text/markdown
8
+ Provides-Extra: gpu
9
+ Requires-Dist: torch; extra == "gpu"
10
+ Provides-Extra: tunnel
11
+ Requires-Dist: pyngrok; extra == "tunnel"
12
+ Provides-Extra: sysinfo
13
+ Requires-Dist: psutil; extra == "sysinfo"
14
+
15
+ # gpumesh
16
+
17
+ **Borrow your friends' GPUs.** A terminal-based distributed compute mesh written in
18
+ pure Python (stdlib only — `torch`, `psutil`, `pyngrok` are optional extras).
19
+
20
+ You have ML work to run but a weak laptop. Your friend has a strong GPU sitting
21
+ idle. `gpumesh` turns any group of machines into a small compute cluster: one
22
+ machine coordinates, any number of machines join with a single command, and work
23
+ is split between them **proportionally to how fast each machine is**.
24
+
25
+ ```
26
+ your laptop (coordinator) friend's laptop (worker)
27
+ ┌───────────────────────┐ HTTP/JSON ┌────────────────────────┐
28
+ │ job queue │◄─────────────►│ hardware probe │
29
+ │ capability scheduler │ (optional │ matmul benchmark │
30
+ │ SQLite (WAL) │ ngrok │ sandboxed subprocess │
31
+ │ heartbeats + reaper │ tunnel) │ executor │
32
+ └──────────▲────────────┘ └────────────────────────┘
33
+ │ submit / status
34
+ client CLI
35
+ ```
36
+
37
+ ## Quick start
38
+
39
+ ```bash
40
+ pip install -e .
41
+ ```
42
+
43
+ **Machine 1 — start the coordinator:**
44
+
45
+ ```bash
46
+ gpumesh serve --port 8000
47
+ # prints a token and a ready-made join command, e.g.
48
+ # gpumesh join http://192.168.1.5:8000 --token Kv3xP9qL2mNa
49
+ ```
50
+
51
+ Add `--public` (with `pip install pyngrok`) to get a public URL friends can
52
+ reach over the internet.
53
+
54
+ **Machine 2..N — join as a worker (one command, as promised):**
55
+
56
+ ```bash
57
+ gpumesh join http://192.168.1.5:8000 --token Kv3xP9qL2mNa
58
+ ```
59
+
60
+ The worker probes its hardware (CUDA GPU → Apple Silicon → CPU), benchmarks
61
+ itself with a matmul, and starts pulling tasks. Your own laptop can join too —
62
+ then both machines compute in parallel.
63
+
64
+ **Submit work:**
65
+
66
+ ```bash
67
+ gpumesh submit examples/grid_search.py --payloads examples/payloads.json \
68
+ --url http://192.168.1.5:8000 --token Kv3xP9qL2mNa --wait
69
+ ```
70
+
71
+ A job is a Python script plus a JSON list of payloads (shards). Each payload
72
+ becomes one task; an optional `"cost"` field marks how heavy it is. Heavy
73
+ tasks go to fast machines, light tasks to slow ones.
74
+
75
+ **Watch the mesh:**
76
+
77
+ ```bash
78
+ gpumesh workers --url ... --token ...
79
+ gpumesh status JOB_ID --url ... --token ...
80
+ ```
81
+
82
+ ## Writing your own task script
83
+
84
+ The contract is two lines of glue:
85
+
86
+ ```python
87
+ import json, sys
88
+ payload = json.load(sys.stdin) # your shard's parameters
89
+ # ... do the work (os.environ["GPUMESH_DEVICE"] tells you cuda/mps/cpu) ...
90
+ print(json.dumps({"answer": 42})) # result = last stdout line, as JSON
91
+ ```
92
+
93
+ Anything data-parallel fits: hyperparameter search, batch inference,
94
+ cross-validation folds, rendering chunks, brute-force search shards.
95
+
96
+ ## How the scheduler splits work
97
+
98
+ 1. Every worker benchmarks itself at join time → a GFLOP/s **score**.
99
+ 2. Pending tasks are sorted by `cost`.
100
+ 3. When a worker asks for work, its score percentile among live workers picks
101
+ the matching percentile of task costs — the strongest machine gets the
102
+ heaviest remaining task, the weakest gets the lightest.
103
+ 4. Tasks are leased, not given away: if a worker dies (heartbeats stop) or the
104
+ lease expires, the reaper re-queues the task for someone else (max 3
105
+ attempts, then marked failed).
106
+
107
+ ## Design notes (the interview section)
108
+
109
+ - **Networking** — hand-rolled JSON-over-HTTP protocol on `http.server`
110
+ (no framework), worker heartbeats, poll-based task leasing that survives
111
+ flaky tunnels, shared-token auth, ngrok tunneling for NAT traversal.
112
+ - **DBMS** — SQLite in WAL mode; schema with foreign keys and indexes; atomic
113
+ lease acquisition via a conditional `UPDATE ... WHERE status='pending'` so
114
+ two workers can never grab the same task; transactions around multi-row
115
+ writes.
116
+ - **Operating systems** — every task runs in a fresh subprocess in its own
117
+ session (process group); timeouts kill the whole process tree with
118
+ `SIGKILL`; POSIX `RLIMIT_CPU` caps runaway loops; coordinator uses threads
119
+ with a lock around the shared DB connection.
120
+ - **Distributed systems** — capability-based scheduling, lease/heartbeat
121
+ failure detection, idempotent re-queue with bounded retries, work-stealing
122
+ style pull model (workers pull; the coordinator never needs to reach them
123
+ through NAT).
124
+ - **Task-parallel, not tensor-parallel** — splitting a single model's tensors
125
+ across machines over the internet dies on latency (that's what NCCL +
126
+ InfiniBand are for). Sharding independent work units is the model serverless
127
+ GPU platforms (e.g. RunPod) actually use.
128
+
129
+ ## Security warning
130
+
131
+ Workers execute code submitted to the coordinator. That is the point of the
132
+ tool, and it is also remote code execution by design. Only share your
133
+ coordinator URL + token with people you trust, treat the token as a password,
134
+ and don't leave a public tunnel running unattended. The subprocess sandbox
135
+ limits accidents, not attackers.
136
+
137
+ ## Layout
138
+
139
+ ```
140
+ gpumesh/
141
+ cli.py command line entry point (serve / join / submit / status / workers)
142
+ server.py coordinator: threaded HTTP JSON API + lease reaper
143
+ db.py SQLite layer: workers, jobs, tasks, atomic leasing
144
+ worker.py agent loop: register, heartbeat, lease, execute, report
145
+ sandbox.py subprocess isolation: timeouts, process-group kill, rlimits
146
+ capability.py hardware probe + matmul benchmark -> capability score
147
+ client.py job submission and status polling
148
+ tunnel.py optional ngrok public URL
149
+ examples/
150
+ grid_search.py hyperparameter-search demo task (pure Python)
151
+ payloads.json six shards with varying cost
152
+ ```
@@ -0,0 +1,138 @@
1
+ # gpumesh
2
+
3
+ **Borrow your friends' GPUs.** A terminal-based distributed compute mesh written in
4
+ pure Python (stdlib only — `torch`, `psutil`, `pyngrok` are optional extras).
5
+
6
+ You have ML work to run but a weak laptop. Your friend has a strong GPU sitting
7
+ idle. `gpumesh` turns any group of machines into a small compute cluster: one
8
+ machine coordinates, any number of machines join with a single command, and work
9
+ is split between them **proportionally to how fast each machine is**.
10
+
11
+ ```
12
+ your laptop (coordinator) friend's laptop (worker)
13
+ ┌───────────────────────┐ HTTP/JSON ┌────────────────────────┐
14
+ │ job queue │◄─────────────►│ hardware probe │
15
+ │ capability scheduler │ (optional │ matmul benchmark │
16
+ │ SQLite (WAL) │ ngrok │ sandboxed subprocess │
17
+ │ heartbeats + reaper │ tunnel) │ executor │
18
+ └──────────▲────────────┘ └────────────────────────┘
19
+ │ submit / status
20
+ client CLI
21
+ ```
22
+
23
+ ## Quick start
24
+
25
+ ```bash
26
+ pip install -e .
27
+ ```
28
+
29
+ **Machine 1 — start the coordinator:**
30
+
31
+ ```bash
32
+ gpumesh serve --port 8000
33
+ # prints a token and a ready-made join command, e.g.
34
+ # gpumesh join http://192.168.1.5:8000 --token Kv3xP9qL2mNa
35
+ ```
36
+
37
+ Add `--public` (with `pip install pyngrok`) to get a public URL friends can
38
+ reach over the internet.
39
+
40
+ **Machine 2..N — join as a worker (one command, as promised):**
41
+
42
+ ```bash
43
+ gpumesh join http://192.168.1.5:8000 --token Kv3xP9qL2mNa
44
+ ```
45
+
46
+ The worker probes its hardware (CUDA GPU → Apple Silicon → CPU), benchmarks
47
+ itself with a matmul, and starts pulling tasks. Your own laptop can join too —
48
+ then both machines compute in parallel.
49
+
50
+ **Submit work:**
51
+
52
+ ```bash
53
+ gpumesh submit examples/grid_search.py --payloads examples/payloads.json \
54
+ --url http://192.168.1.5:8000 --token Kv3xP9qL2mNa --wait
55
+ ```
56
+
57
+ A job is a Python script plus a JSON list of payloads (shards). Each payload
58
+ becomes one task; an optional `"cost"` field marks how heavy it is. Heavy
59
+ tasks go to fast machines, light tasks to slow ones.
60
+
61
+ **Watch the mesh:**
62
+
63
+ ```bash
64
+ gpumesh workers --url ... --token ...
65
+ gpumesh status JOB_ID --url ... --token ...
66
+ ```
67
+
68
+ ## Writing your own task script
69
+
70
+ The contract is two lines of glue:
71
+
72
+ ```python
73
+ import json, sys
74
+ payload = json.load(sys.stdin) # your shard's parameters
75
+ # ... do the work (os.environ["GPUMESH_DEVICE"] tells you cuda/mps/cpu) ...
76
+ print(json.dumps({"answer": 42})) # result = last stdout line, as JSON
77
+ ```
78
+
79
+ Anything data-parallel fits: hyperparameter search, batch inference,
80
+ cross-validation folds, rendering chunks, brute-force search shards.
81
+
82
+ ## How the scheduler splits work
83
+
84
+ 1. Every worker benchmarks itself at join time → a GFLOP/s **score**.
85
+ 2. Pending tasks are sorted by `cost`.
86
+ 3. When a worker asks for work, its score percentile among live workers picks
87
+ the matching percentile of task costs — the strongest machine gets the
88
+ heaviest remaining task, the weakest gets the lightest.
89
+ 4. Tasks are leased, not given away: if a worker dies (heartbeats stop) or the
90
+ lease expires, the reaper re-queues the task for someone else (max 3
91
+ attempts, then marked failed).
92
+
93
+ ## Design notes (the interview section)
94
+
95
+ - **Networking** — hand-rolled JSON-over-HTTP protocol on `http.server`
96
+ (no framework), worker heartbeats, poll-based task leasing that survives
97
+ flaky tunnels, shared-token auth, ngrok tunneling for NAT traversal.
98
+ - **DBMS** — SQLite in WAL mode; schema with foreign keys and indexes; atomic
99
+ lease acquisition via a conditional `UPDATE ... WHERE status='pending'` so
100
+ two workers can never grab the same task; transactions around multi-row
101
+ writes.
102
+ - **Operating systems** — every task runs in a fresh subprocess in its own
103
+ session (process group); timeouts kill the whole process tree with
104
+ `SIGKILL`; POSIX `RLIMIT_CPU` caps runaway loops; coordinator uses threads
105
+ with a lock around the shared DB connection.
106
+ - **Distributed systems** — capability-based scheduling, lease/heartbeat
107
+ failure detection, idempotent re-queue with bounded retries, work-stealing
108
+ style pull model (workers pull; the coordinator never needs to reach them
109
+ through NAT).
110
+ - **Task-parallel, not tensor-parallel** — splitting a single model's tensors
111
+ across machines over the internet dies on latency (that's what NCCL +
112
+ InfiniBand are for). Sharding independent work units is the model serverless
113
+ GPU platforms (e.g. RunPod) actually use.
114
+
115
+ ## Security warning
116
+
117
+ Workers execute code submitted to the coordinator. That is the point of the
118
+ tool, and it is also remote code execution by design. Only share your
119
+ coordinator URL + token with people you trust, treat the token as a password,
120
+ and don't leave a public tunnel running unattended. The subprocess sandbox
121
+ limits accidents, not attackers.
122
+
123
+ ## Layout
124
+
125
+ ```
126
+ gpumesh/
127
+ cli.py command line entry point (serve / join / submit / status / workers)
128
+ server.py coordinator: threaded HTTP JSON API + lease reaper
129
+ db.py SQLite layer: workers, jobs, tasks, atomic leasing
130
+ worker.py agent loop: register, heartbeat, lease, execute, report
131
+ sandbox.py subprocess isolation: timeouts, process-group kill, rlimits
132
+ capability.py hardware probe + matmul benchmark -> capability score
133
+ client.py job submission and status polling
134
+ tunnel.py optional ngrok public URL
135
+ examples/
136
+ grid_search.py hyperparameter-search demo task (pure Python)
137
+ payloads.json six shards with varying cost
138
+ ```
@@ -0,0 +1,3 @@
1
+ """gpumesh - distributed compute sharing over a mesh of volunteer machines."""
2
+
3
+ __version__ = "0.1.0"
@@ -0,0 +1,5 @@
1
+ """Allow running gpumesh as a module: python -m gpumesh"""
2
+
3
+ from .cli import main
4
+
5
+ main()
@@ -0,0 +1,148 @@
1
+ """Hardware probing and benchmarking.
2
+
3
+ Detects the best compute device available and produces a single numeric
4
+ capability score used by the scheduler. torch and psutil are optional:
5
+ without them the probe falls back to CPU info and a pure-Python benchmark.
6
+
7
+ GPU detection uses nvidia-smi (no PyTorch required for detection).
8
+ PyTorch with CUDA is auto-installed when a GPU is detected.
9
+ """
10
+
11
+ import os
12
+ import platform
13
+ import shutil
14
+ import socket
15
+ import subprocess
16
+ import sys
17
+ import time
18
+
19
+
20
+ def _detect_nvidia_gpu() -> dict:
21
+ """Detect NVIDIA GPU using nvidia-smi (no PyTorch required)."""
22
+ if shutil.which("nvidia-smi") is None:
23
+ return None
24
+ try:
25
+ result = subprocess.run(
26
+ ["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader,nounits"],
27
+ capture_output=True,
28
+ text=True,
29
+ timeout=5,
30
+ )
31
+ if result.returncode == 0 and result.stdout.strip():
32
+ parts = result.stdout.strip().split(",")
33
+ name = parts[0].strip()
34
+ memory = parts[1].strip() if len(parts) > 1 else "unknown"
35
+ return {"name": name, "memory_mb": memory}
36
+ except (subprocess.SubprocessError, FileNotFoundError):
37
+ pass
38
+ return None
39
+
40
+
41
+ def _install_pytorch_cuda():
42
+ """Install PyTorch with CUDA support if NVIDIA GPU detected."""
43
+ print("[mesh] NVIDIA GPU detected, installing PyTorch with CUDA support...")
44
+ try:
45
+ subprocess.run(
46
+ [sys.executable, "-m", "pip", "install", "torch", "torchvision", "torchaudio",
47
+ "--index-url", "https://download.pytorch.org/whl/cu118"],
48
+ check=True,
49
+ timeout=300,
50
+ )
51
+ print("[mesh] PyTorch with CUDA installed successfully!")
52
+ return True
53
+ except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
54
+ print(f"[mesh] Failed to install PyTorch: {e}")
55
+ print("[mesh] Continuing with CPU-only mode.")
56
+ return False
57
+
58
+
59
+ def probe_device() -> dict:
60
+ info = {
61
+ "hostname": socket.gethostname(),
62
+ "platform": platform.platform(),
63
+ "cpu_count": os.cpu_count() or 1,
64
+ "device": "cpu",
65
+ "device_name": platform.processor() or "cpu",
66
+ }
67
+
68
+ # First check for NVIDIA GPU using nvidia-smi (fast, no PyTorch needed)
69
+ nvidia_gpu = _detect_nvidia_gpu()
70
+ if nvidia_gpu:
71
+ info["nvidia_gpu"] = nvidia_gpu["name"]
72
+ info["nvidia_memory"] = nvidia_gpu["memory_mb"]
73
+
74
+ # Try to use PyTorch for GPU detection and benchmarking
75
+ try:
76
+ import torch
77
+ if torch.cuda.is_available():
78
+ info["device"] = "cuda"
79
+ info["device_name"] = torch.cuda.get_device_name(0)
80
+ elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
81
+ info["device"] = "mps"
82
+ info["device_name"] = "Apple Silicon GPU"
83
+ except ImportError:
84
+ # PyTorch not installed - try to auto-install if NVIDIA GPU detected
85
+ if nvidia_gpu:
86
+ if _install_pytorch_cuda():
87
+ try:
88
+ import torch
89
+ if torch.cuda.is_available():
90
+ info["device"] = "cuda"
91
+ info["device_name"] = torch.cuda.get_device_name(0)
92
+ except Exception:
93
+ pass
94
+
95
+ try:
96
+ import psutil
97
+ info["ram_gb"] = round(psutil.virtual_memory().total / 1e9, 1)
98
+ except ImportError:
99
+ pass
100
+
101
+ return info
102
+
103
+
104
+ def _bench_torch(device: str) -> float:
105
+ """Time a matmul on the given torch device; return ops/sec score."""
106
+ import torch
107
+
108
+ n = 1024
109
+ a = torch.rand(n, n, device=device)
110
+ b = torch.rand(n, n, device=device)
111
+ # warmup
112
+ (a @ b).sum().item()
113
+ start = time.perf_counter()
114
+ iters = 10
115
+ for _ in range(iters):
116
+ c = a @ b
117
+ c.sum().item() # force sync on cuda/mps
118
+ elapsed = time.perf_counter() - start
119
+ flops = 2 * n**3 * iters
120
+ return flops / elapsed / 1e9 # GFLOP/s
121
+
122
+
123
+ def _bench_python() -> float:
124
+ """Pure-Python matmul fallback. Slow by design; scores are comparable
125
+ across machines because everyone runs the same loop."""
126
+ n = 48
127
+ a = [[(i * j) % 7 / 7.0 for j in range(n)] for i in range(n)]
128
+ b = [[(i + j) % 5 / 5.0 for j in range(n)] for i in range(n)]
129
+ start = time.perf_counter()
130
+ iters = 5
131
+ for _ in range(iters):
132
+ [[sum(a[i][k] * b[k][j] for k in range(n)) for j in range(n)] for i in range(n)]
133
+ elapsed = time.perf_counter() - start
134
+ flops = 2 * n**3 * iters
135
+ return flops / elapsed / 1e9
136
+
137
+
138
+ def benchmark(device: str) -> float:
139
+ try:
140
+ return round(_bench_torch(device), 3)
141
+ except ImportError:
142
+ return round(_bench_python(), 3)
143
+
144
+
145
+ def full_probe() -> dict:
146
+ info = probe_device()
147
+ info["score"] = benchmark(info["device"])
148
+ return info
@@ -0,0 +1,225 @@
1
+ """gpumesh command line interface.
2
+
3
+ gpumesh serve [--port 8000] [--token SECRET] [--public]
4
+ gpumesh join URL [--token SECRET]
5
+ gpumesh quickjoin URL --token TOKEN # One-click setup & join
6
+ gpumesh submit SCRIPT --payloads FILE [--wait] --url URL [--token SECRET]
7
+ gpumesh status JOB_ID --url URL [--token SECRET]
8
+ gpumesh cancel JOB_ID --url URL [--token SECRET]
9
+ gpumesh workers --url URL [--token SECRET]
10
+ """
11
+
12
+ import argparse
13
+ import os
14
+ import secrets
15
+ import socket
16
+ import subprocess
17
+ import sys
18
+
19
+ from . import client, server, tunnel, worker
20
+
21
+
22
+ def _lan_ip() -> str:
23
+ s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
24
+ try:
25
+ s.connect(("8.8.8.8", 80))
26
+ return s.getsockname()[0]
27
+ except OSError:
28
+ return "127.0.0.1"
29
+ finally:
30
+ s.close()
31
+
32
+
33
+ def cmd_serve(args):
34
+ token = args.token or secrets.token_urlsafe(12)
35
+ httpd = server.serve("0.0.0.0", args.port, args.db, token)
36
+ ip = _lan_ip()
37
+ print(f"[mesh] coordinator listening on 0.0.0.0:{args.port}")
38
+ print(f"[mesh] token: {token}")
39
+ print(f"[mesh] LAN join command:")
40
+ print(f" gpumesh join http://{ip}:{args.port} --token {token}")
41
+ if args.public:
42
+ tunnel.open_tunnel(args.port)
43
+ print("[mesh] Ctrl+C to stop")
44
+ try:
45
+ httpd.serve_forever()
46
+ except KeyboardInterrupt:
47
+ print("\n[mesh] shutting down")
48
+ httpd.gpumesh_stop.set()
49
+ httpd.shutdown()
50
+
51
+
52
+ def cmd_join(args):
53
+ worker.run_worker(args.url, args.token, task_timeout=args.timeout)
54
+
55
+
56
+ def cmd_submit(args):
57
+ job_id = client.submit_job(args.url, args.token, args.script,
58
+ args.payloads, name=args.name)
59
+ print(f"[client] submitted job {job_id}")
60
+ if args.wait:
61
+ job = client.wait_for_job(args.url, args.token, job_id)
62
+ client.print_job(job)
63
+ else:
64
+ print(f"[client] check progress: gpumesh status {job_id}"
65
+ f" --url {args.url} --token {args.token}")
66
+
67
+
68
+ def cmd_status(args):
69
+ job = client.get_status(args.url, args.token, args.job_id)
70
+ client.print_job(job)
71
+
72
+
73
+ def cmd_cancel(args):
74
+ result = client.cancel_job(args.url, args.token, args.job_id)
75
+ if result is None:
76
+ print(f"[client] job {args.job_id} not found")
77
+ else:
78
+ print(f"[client] cancelled job {args.job_id}")
79
+ print(f" pending tasks cancelled: {result['pending']}")
80
+ print(f" running tasks cancelled: {result['running']}")
81
+ print(f" already finished: {result['already_finished']}")
82
+
83
+
84
+ def cmd_workers(args):
85
+ from .worker import MeshClient
86
+
87
+ resp = MeshClient(args.url, args.token).call("GET", "/api/workers")
88
+ for w in resp["workers"]:
89
+ state = "alive" if w["alive"] else "dead"
90
+ print(f" {w['id']} {w['hostname']:<20} {w['device']:<5} "
91
+ f"score={w['score']:<10} [{state}]")
92
+ if not resp["workers"]:
93
+ print(" (no workers yet)")
94
+
95
+
96
+ def cmd_quickjoin(args):
97
+ """One-click setup: install dependencies, detect GPU, and join mesh."""
98
+ url = args.url
99
+ token = args.token
100
+
101
+ print("[mesh] Quick Join - Setting up your machine as a worker...")
102
+ print()
103
+
104
+ # Step 1: Check Python version
105
+ print("[1/4] Checking Python version...")
106
+ if sys.version_info < (3, 9):
107
+ print("[mesh] ERROR: Python 3.9 or higher is required.")
108
+ print("[mesh] Please install Python 3.9+ from https://www.python.org/downloads/")
109
+ return
110
+ print(f"[mesh] Python {sys.version_info.major}.{sys.version_info.minor} ✓")
111
+
112
+ # Step 2: Install gpumesh if not already installed
113
+ print("[2/4] Checking gpumesh installation...")
114
+ try:
115
+ import gpumesh
116
+ print("[mesh] gpumesh already installed ✓")
117
+ except ImportError:
118
+ print("[mesh] Installing gpumesh...")
119
+ try:
120
+ subprocess.run(
121
+ [sys.executable, "-m", "pip", "install", "-e", "."],
122
+ check=True,
123
+ timeout=120,
124
+ )
125
+ print("[mesh] gpumesh installed ✓")
126
+ except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
127
+ print(f"[mesh] ERROR: Failed to install gpumesh: {e}")
128
+ print("[mesh] Please install manually: pip install gpumesh")
129
+ return
130
+
131
+ # Step 3: Detect hardware and install GPU support if needed
132
+ print("[3/4] Detecting hardware...")
133
+ from . import capability
134
+
135
+ # Check for NVIDIA GPU using nvidia-smi
136
+ nvidia_gpu = capability._detect_nvidia_gpu()
137
+ if nvidia_gpu:
138
+ print(f"[mesh] NVIDIA GPU detected: {nvidia_gpu['name']}")
139
+
140
+ # Check if PyTorch with CUDA is available
141
+ try:
142
+ import torch
143
+ if torch.cuda.is_available():
144
+ print(f"[mesh] PyTorch with CUDA ready ✓")
145
+ else:
146
+ print("[mesh] PyTorch installed but CUDA not available")
147
+ capability._install_pytorch_cuda()
148
+ except ImportError:
149
+ print("[mesh] PyTorch not found, installing with CUDA support...")
150
+ capability._install_pytorch_cuda()
151
+ else:
152
+ print("[mesh] No NVIDIA GPU detected, using CPU")
153
+
154
+ # Step 4: Join the mesh
155
+ print("[4/4] Joining the mesh...")
156
+ print(f"[mesh] Connecting to {url}...")
157
+ print()
158
+
159
+ # Now run the normal join command
160
+ worker.run_worker(url, token)
161
+
162
+
163
+ def _add_conn_args(p):
164
+ p.add_argument("--url", default=os.environ.get("GPUMESH_URL", ""),
165
+ help="coordinator URL (or set GPUMESH_URL)")
166
+ p.add_argument("--token", default=os.environ.get("GPUMESH_TOKEN", ""),
167
+ help="shared auth token (or set GPUMESH_TOKEN)")
168
+
169
+
170
+ def main():
171
+ ap = argparse.ArgumentParser(prog="gpumesh",
172
+ description="borrow your friends' GPUs")
173
+ sub = ap.add_subparsers(dest="cmd", required=True)
174
+
175
+ p = sub.add_parser("serve", help="start a coordinator on this machine")
176
+ p.add_argument("--port", type=int, default=8000)
177
+ p.add_argument("--db", default="gpumesh.db")
178
+ p.add_argument("--token", default="", help="auth token (random if omitted)")
179
+ p.add_argument("--public", action="store_true",
180
+ help="expose a public URL via ngrok")
181
+ p.set_defaults(func=cmd_serve)
182
+
183
+ p = sub.add_parser("join", help="offer this machine's compute to a mesh")
184
+ p.add_argument("url", help="coordinator URL")
185
+ p.add_argument("--token", default=os.environ.get("GPUMESH_TOKEN", ""))
186
+ p.add_argument("--timeout", type=float, default=240.0,
187
+ help="per-task wall clock limit in seconds")
188
+ p.set_defaults(func=cmd_join)
189
+
190
+ p = sub.add_parser("submit", help="submit a job to the mesh")
191
+ p.add_argument("script", help="python script to run for every payload")
192
+ p.add_argument("--payloads", required=True,
193
+ help="JSON file: list of payload objects (each may set 'cost')")
194
+ p.add_argument("--name", default="")
195
+ p.add_argument("--wait", action="store_true", help="block until finished")
196
+ _add_conn_args(p)
197
+ p.set_defaults(func=cmd_submit)
198
+
199
+ p = sub.add_parser("status", help="show job progress and results")
200
+ p.add_argument("job_id")
201
+ _add_conn_args(p)
202
+ p.set_defaults(func=cmd_status)
203
+
204
+ p = sub.add_parser("cancel", help="cancel a running job")
205
+ p.add_argument("job_id")
206
+ _add_conn_args(p)
207
+ p.set_defaults(func=cmd_cancel)
208
+
209
+ p = sub.add_parser("quickjoin", help="one-click setup: install, detect GPU, and join mesh")
210
+ p.add_argument("url", help="coordinator URL")
211
+ p.add_argument("--token", required=True, help="shared auth token")
212
+ p.set_defaults(func=cmd_quickjoin)
213
+
214
+ p = sub.add_parser("workers", help="list workers in the mesh")
215
+ _add_conn_args(p)
216
+ p.set_defaults(func=cmd_workers)
217
+
218
+ args = ap.parse_args()
219
+ if hasattr(args, "url") and not args.url and args.cmd != "serve":
220
+ ap.error("--url required (or set GPUMESH_URL)")
221
+ args.func(args)
222
+
223
+
224
+ if __name__ == "__main__":
225
+ main()