blitzwing 0.1.0
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.
- package/README.md +29 -0
- package/bin/blitzwing.js +7 -0
- package/package.json +31 -0
- package/src/api.js +68 -0
- package/src/config.js +9 -0
- package/src/index.js +1 -0
- package/src/install.js +218 -0
- package/src/net.js +36 -0
- package/src/state.js +32 -0
- package/src/wizard.js +318 -0
package/README.md
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
# blitzwing
|
|
2
|
+
|
|
3
|
+
Interactive setup wizard to contribute layers to a Blitzwing Petals mother swarm.
|
|
4
|
+
|
|
5
|
+
## Install
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
npm i -g blitzwing
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
## Use
|
|
12
|
+
|
|
13
|
+
```bash
|
|
14
|
+
blitzwing
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
The wizard talks to the **Discovery Service** (no mother URL to type), asks which model and how many layers to host, installs Petals if needed, and joins the swarm.
|
|
18
|
+
|
|
19
|
+
```bash
|
|
20
|
+
blitzwing status
|
|
21
|
+
blitzwing leave
|
|
22
|
+
```
|
|
23
|
+
|
|
24
|
+
Override discovery URL (dev/staging only):
|
|
25
|
+
|
|
26
|
+
```bash
|
|
27
|
+
export BLITZWING_DISCOVERY_URL=http://127.0.0.1:9000
|
|
28
|
+
blitzwing
|
|
29
|
+
```
|
package/bin/blitzwing.js
ADDED
package/package.json
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
{
|
|
2
|
+
"name": "blitzwing",
|
|
3
|
+
"version": "0.1.0",
|
|
4
|
+
"description": "Interactive setup wizard to join a Blitzwing Petals mother swarm as a compute contributor",
|
|
5
|
+
"bin": {
|
|
6
|
+
"blitzwing": "bin/blitzwing.js"
|
|
7
|
+
},
|
|
8
|
+
"type": "module",
|
|
9
|
+
"engines": {
|
|
10
|
+
"node": ">=18"
|
|
11
|
+
},
|
|
12
|
+
"files": [
|
|
13
|
+
"bin",
|
|
14
|
+
"src",
|
|
15
|
+
"README.md"
|
|
16
|
+
],
|
|
17
|
+
"keywords": [
|
|
18
|
+
"blitzwing",
|
|
19
|
+
"petals",
|
|
20
|
+
"llm",
|
|
21
|
+
"distributed-inference"
|
|
22
|
+
],
|
|
23
|
+
"license": "MIT",
|
|
24
|
+
"config": {
|
|
25
|
+
"discoveryUrl": "http://127.0.0.1:9000"
|
|
26
|
+
},
|
|
27
|
+
"dependencies": {
|
|
28
|
+
"@clack/prompts": "^0.9.1",
|
|
29
|
+
"picocolors": "^1.1.1"
|
|
30
|
+
}
|
|
31
|
+
}
|
package/src/api.js
ADDED
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+
export async function httpJson(url, options = {}) {
|
|
2
|
+
const res = await fetch(url, {
|
|
3
|
+
...options,
|
|
4
|
+
headers: {
|
|
5
|
+
"Content-Type": "application/json",
|
|
6
|
+
Accept: "application/json",
|
|
7
|
+
...(options.headers || {}),
|
|
8
|
+
},
|
|
9
|
+
});
|
|
10
|
+
const text = await res.text();
|
|
11
|
+
let body = null;
|
|
12
|
+
try {
|
|
13
|
+
body = text ? JSON.parse(text) : null;
|
|
14
|
+
} catch {
|
|
15
|
+
body = { raw: text };
|
|
16
|
+
}
|
|
17
|
+
if (!res.ok) {
|
|
18
|
+
const err = new Error(
|
|
19
|
+
typeof body === "object"
|
|
20
|
+
? body.detail?.message || body.detail || body.message || text || res.statusText
|
|
21
|
+
: text || res.statusText
|
|
22
|
+
);
|
|
23
|
+
err.status = res.status;
|
|
24
|
+
err.body = body;
|
|
25
|
+
throw err;
|
|
26
|
+
}
|
|
27
|
+
return body;
|
|
28
|
+
}
|
|
29
|
+
|
|
30
|
+
export async function listMothers(discoveryUrl) {
|
|
31
|
+
return httpJson(`${discoveryUrl.replace(/\/$/, "")}/v1/mothers`);
|
|
32
|
+
}
|
|
33
|
+
|
|
34
|
+
export async function motherHealth(motherUrl) {
|
|
35
|
+
return httpJson(`${motherUrl.replace(/\/$/, "")}/health`);
|
|
36
|
+
}
|
|
37
|
+
|
|
38
|
+
export async function motherHosts(motherUrl) {
|
|
39
|
+
return httpJson(`${motherUrl.replace(/\/$/, "")}/v1/hosts`);
|
|
40
|
+
}
|
|
41
|
+
|
|
42
|
+
export async function joinHost(motherUrl, payload) {
|
|
43
|
+
return httpJson(`${motherUrl.replace(/\/$/, "")}/v1/hosts/join`, {
|
|
44
|
+
method: "POST",
|
|
45
|
+
body: JSON.stringify(payload),
|
|
46
|
+
});
|
|
47
|
+
}
|
|
48
|
+
|
|
49
|
+
export async function readyHost(motherUrl, payload) {
|
|
50
|
+
return httpJson(`${motherUrl.replace(/\/$/, "")}/v1/hosts/ready`, {
|
|
51
|
+
method: "POST",
|
|
52
|
+
body: JSON.stringify(payload),
|
|
53
|
+
});
|
|
54
|
+
}
|
|
55
|
+
|
|
56
|
+
export async function heartbeatHost(motherUrl, payload) {
|
|
57
|
+
return httpJson(`${motherUrl.replace(/\/$/, "")}/v1/hosts/heartbeat`, {
|
|
58
|
+
method: "POST",
|
|
59
|
+
body: JSON.stringify(payload),
|
|
60
|
+
});
|
|
61
|
+
}
|
|
62
|
+
|
|
63
|
+
export async function leaveHost(motherUrl, payload) {
|
|
64
|
+
return httpJson(`${motherUrl.replace(/\/$/, "")}/v1/hosts/leave`, {
|
|
65
|
+
method: "POST",
|
|
66
|
+
body: JSON.stringify(payload),
|
|
67
|
+
});
|
|
68
|
+
}
|
package/src/config.js
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
/** Default Discovery Service URL — override with BLITZWING_DISCOVERY_URL or --discovery-url */
|
|
2
|
+
export const DEFAULT_DISCOVERY_URL =
|
|
3
|
+
process.env.BLITZWING_DISCOVERY_URL || "http://127.0.0.1:9000";
|
|
4
|
+
|
|
5
|
+
export const HOME_DIR = process.env.BLITZWING_HOME || `${process.env.HOME || process.env.USERPROFILE}/.blitzwing`;
|
|
6
|
+
export const STATE_PATH = `${HOME_DIR}/contributor.json`;
|
|
7
|
+
export const VENV_PATH = `${HOME_DIR}/venv`;
|
|
8
|
+
export const SHARD_PORT = Number(process.env.BLITZWING_SHARD_PORT || 8001);
|
|
9
|
+
export const PETALS_PORT = Number(process.env.BLITZWING_PETALS_PORT || 31337);
|
package/src/index.js
ADDED
|
@@ -0,0 +1 @@
|
|
|
1
|
+
export { run } from "./wizard.js";
|
package/src/install.js
ADDED
|
@@ -0,0 +1,218 @@
|
|
|
1
|
+
import { spawn, execFileSync } from "node:child_process";
|
|
2
|
+
import fs from "node:fs";
|
|
3
|
+
import path from "node:path";
|
|
4
|
+
import { HOME_DIR, VENV_PATH, SHARD_PORT, PETALS_PORT } from "./config.js";
|
|
5
|
+
import { which } from "./net.js";
|
|
6
|
+
|
|
7
|
+
function venvPython() {
|
|
8
|
+
if (process.platform === "win32") {
|
|
9
|
+
return path.join(VENV_PATH, "Scripts", "python.exe");
|
|
10
|
+
}
|
|
11
|
+
return path.join(VENV_PATH, "bin", "python");
|
|
12
|
+
}
|
|
13
|
+
|
|
14
|
+
function venvPip() {
|
|
15
|
+
if (process.platform === "win32") {
|
|
16
|
+
return path.join(VENV_PATH, "Scripts", "pip.exe");
|
|
17
|
+
}
|
|
18
|
+
return path.join(VENV_PATH, "bin", "pip");
|
|
19
|
+
}
|
|
20
|
+
|
|
21
|
+
export function ensurePython() {
|
|
22
|
+
const py = which("python3") || which("python");
|
|
23
|
+
if (!py) {
|
|
24
|
+
throw new Error("Python 3 is required. Install Python 3.10+ and re-run blitzwing.");
|
|
25
|
+
}
|
|
26
|
+
const ver = execFileSync(py, ["--version"], { encoding: "utf8" });
|
|
27
|
+
return { py, ver: ver.trim() };
|
|
28
|
+
}
|
|
29
|
+
|
|
30
|
+
export function ensureVenvAndPetals({ onLog }) {
|
|
31
|
+
const { py } = ensurePython();
|
|
32
|
+
fs.mkdirSync(HOME_DIR, { recursive: true });
|
|
33
|
+
|
|
34
|
+
if (!fs.existsSync(venvPython())) {
|
|
35
|
+
onLog?.(`Creating venv at ${VENV_PATH}`);
|
|
36
|
+
execFileSync(py, ["-m", "venv", VENV_PATH], { stdio: "inherit" });
|
|
37
|
+
}
|
|
38
|
+
|
|
39
|
+
const pip = venvPip();
|
|
40
|
+
const python = venvPython();
|
|
41
|
+
|
|
42
|
+
onLog?.("Installing CPU PyTorch + Petals (may take a few minutes on first run)...");
|
|
43
|
+
execFileSync(pip, ["install", "-U", "pip"], { stdio: "inherit" });
|
|
44
|
+
execFileSync(
|
|
45
|
+
pip,
|
|
46
|
+
["install", "torch", "--index-url", "https://download.pytorch.org/whl/cpu"],
|
|
47
|
+
{ stdio: "inherit" }
|
|
48
|
+
);
|
|
49
|
+
// Prefer git install for contributors without the monorepo
|
|
50
|
+
execFileSync(
|
|
51
|
+
pip,
|
|
52
|
+
["install", "git+https://github.com/bigscience-workshop/petals.git"],
|
|
53
|
+
{ stdio: "inherit" }
|
|
54
|
+
);
|
|
55
|
+
execFileSync(pip, ["install", "fastapi", "uvicorn[standard]", "pydantic", "httpx"], {
|
|
56
|
+
stdio: "inherit",
|
|
57
|
+
});
|
|
58
|
+
|
|
59
|
+
return { python, pip };
|
|
60
|
+
}
|
|
61
|
+
|
|
62
|
+
/**
|
|
63
|
+
* Write a tiny local shard-manager runner script so contributors don't need the monorepo.
|
|
64
|
+
* Embeds a minimal supervisor compatible with mother /reload API.
|
|
65
|
+
*/
|
|
66
|
+
export function writeLocalShardManager() {
|
|
67
|
+
const dir = path.join(HOME_DIR, "runtime");
|
|
68
|
+
fs.mkdirSync(dir, { recursive: true });
|
|
69
|
+
const appPath = path.join(dir, "shard_manager_app.py");
|
|
70
|
+
const code = `
|
|
71
|
+
import os, signal, subprocess, threading, time
|
|
72
|
+
from typing import Optional, List
|
|
73
|
+
from fastapi import FastAPI, HTTPException
|
|
74
|
+
from pydantic import BaseModel, Field
|
|
75
|
+
|
|
76
|
+
class ReloadRequest(BaseModel):
|
|
77
|
+
block_indices: str = Field(..., pattern=r"^\\d+:\\d+$")
|
|
78
|
+
|
|
79
|
+
class StatusResponse(BaseModel):
|
|
80
|
+
running: bool
|
|
81
|
+
pid: Optional[int] = None
|
|
82
|
+
block_indices: Optional[str] = None
|
|
83
|
+
model: str
|
|
84
|
+
public_ip: Optional[str] = None
|
|
85
|
+
port: int
|
|
86
|
+
last_exit_code: Optional[int] = None
|
|
87
|
+
|
|
88
|
+
class Mgr:
|
|
89
|
+
def __init__(self):
|
|
90
|
+
self.model = os.environ["MODEL_NAME"]
|
|
91
|
+
self.public_ip = os.environ.get("PUBLIC_IP")
|
|
92
|
+
self.port = int(os.environ.get("PETALS_PORT", "31337"))
|
|
93
|
+
self.identity_path = os.environ.get("IDENTITY_PATH", os.path.expanduser("~/.blitzwing/petals-identity"))
|
|
94
|
+
self.block_indices = os.environ["BLOCK_INDICES"]
|
|
95
|
+
peers = [p.strip() for p in os.environ.get("INITIAL_PEERS", "").split(",") if p.strip()]
|
|
96
|
+
self.initial_peers = peers
|
|
97
|
+
self.python = os.environ.get("PETALS_PYTHON", "python")
|
|
98
|
+
self._proc = None
|
|
99
|
+
self._lock = threading.Lock()
|
|
100
|
+
self.last_exit_code = None
|
|
101
|
+
|
|
102
|
+
def cmd(self, bi):
|
|
103
|
+
c = [self.python, "-m", "petals.cli.run_server", self.model,
|
|
104
|
+
"--device", "cpu", "--quant_type", "none",
|
|
105
|
+
"--block_indices", bi, "--port", str(self.port),
|
|
106
|
+
"--identity_path", self.identity_path, "--num_handlers", "1"]
|
|
107
|
+
if self.public_ip:
|
|
108
|
+
c += ["--public_ip", self.public_ip]
|
|
109
|
+
if self.initial_peers:
|
|
110
|
+
c += ["--initial_peers", *self.initial_peers]
|
|
111
|
+
return c
|
|
112
|
+
|
|
113
|
+
def start(self, bi=None):
|
|
114
|
+
with self._lock:
|
|
115
|
+
if bi: self.block_indices = bi
|
|
116
|
+
if self._proc and self._proc.poll() is None:
|
|
117
|
+
return
|
|
118
|
+
self._proc = subprocess.Popen(self.cmd(self.block_indices))
|
|
119
|
+
self.last_exit_code = None
|
|
120
|
+
|
|
121
|
+
def stop(self):
|
|
122
|
+
with self._lock:
|
|
123
|
+
if not self._proc or self._proc.poll() is not None:
|
|
124
|
+
self._proc = None
|
|
125
|
+
return
|
|
126
|
+
p = self._proc
|
|
127
|
+
p.send_signal(signal.SIGTERM)
|
|
128
|
+
try: p.wait(timeout=30)
|
|
129
|
+
except Exception:
|
|
130
|
+
p.kill(); p.wait(timeout=10)
|
|
131
|
+
self.last_exit_code = p.returncode
|
|
132
|
+
self._proc = None
|
|
133
|
+
|
|
134
|
+
def reload(self, bi):
|
|
135
|
+
self.stop(); time.sleep(1); self.start(bi)
|
|
136
|
+
|
|
137
|
+
def status(self):
|
|
138
|
+
running = self._proc is not None and self._proc.poll() is None
|
|
139
|
+
return StatusResponse(
|
|
140
|
+
running=running,
|
|
141
|
+
pid=(self._proc.pid if running else None),
|
|
142
|
+
block_indices=self.block_indices,
|
|
143
|
+
model=self.model,
|
|
144
|
+
public_ip=self.public_ip,
|
|
145
|
+
port=self.port,
|
|
146
|
+
last_exit_code=self.last_exit_code,
|
|
147
|
+
)
|
|
148
|
+
|
|
149
|
+
mgr = Mgr()
|
|
150
|
+
app = FastAPI()
|
|
151
|
+
|
|
152
|
+
@app.on_event("startup")
|
|
153
|
+
def _up():
|
|
154
|
+
mgr.start()
|
|
155
|
+
|
|
156
|
+
@app.on_event("shutdown")
|
|
157
|
+
def _down():
|
|
158
|
+
mgr.stop()
|
|
159
|
+
|
|
160
|
+
@app.get("/status")
|
|
161
|
+
def status():
|
|
162
|
+
return mgr.status()
|
|
163
|
+
|
|
164
|
+
@app.post("/reload")
|
|
165
|
+
def reload(body: ReloadRequest):
|
|
166
|
+
a,b = map(int, body.block_indices.split(":"))
|
|
167
|
+
if b <= a: raise HTTPException(400, "bad range")
|
|
168
|
+
mgr.reload(body.block_indices)
|
|
169
|
+
time.sleep(0.5)
|
|
170
|
+
return mgr.status()
|
|
171
|
+
|
|
172
|
+
@app.post("/stop")
|
|
173
|
+
def stop():
|
|
174
|
+
mgr.stop(); return mgr.status()
|
|
175
|
+
`;
|
|
176
|
+
fs.writeFileSync(appPath, code);
|
|
177
|
+
return appPath;
|
|
178
|
+
}
|
|
179
|
+
|
|
180
|
+
export function startShardManagerProcess({
|
|
181
|
+
python,
|
|
182
|
+
env,
|
|
183
|
+
logPath,
|
|
184
|
+
}) {
|
|
185
|
+
const appPath = writeLocalShardManager();
|
|
186
|
+
const out = fs.openSync(logPath, "a");
|
|
187
|
+
const child = spawn(
|
|
188
|
+
python,
|
|
189
|
+
["-m", "uvicorn", `shard_manager_app:app`, "--host", "0.0.0.0", "--port", String(SHARD_PORT)],
|
|
190
|
+
{
|
|
191
|
+
cwd: path.dirname(appPath),
|
|
192
|
+
env: { ...process.env, ...env, PETALS_PYTHON: python },
|
|
193
|
+
detached: true,
|
|
194
|
+
stdio: ["ignore", out, out],
|
|
195
|
+
}
|
|
196
|
+
);
|
|
197
|
+
child.unref();
|
|
198
|
+
return { pid: child.pid, logPath, appPath };
|
|
199
|
+
}
|
|
200
|
+
|
|
201
|
+
export async function waitForShardRunning({ timeoutMs = 300000 } = {}) {
|
|
202
|
+
const start = Date.now();
|
|
203
|
+
while (Date.now() - start < timeoutMs) {
|
|
204
|
+
try {
|
|
205
|
+
const res = await fetch(`http://127.0.0.1:${SHARD_PORT}/status`);
|
|
206
|
+
if (res.ok) {
|
|
207
|
+
const body = await res.json();
|
|
208
|
+
if (body.running) return body;
|
|
209
|
+
}
|
|
210
|
+
} catch {
|
|
211
|
+
/* retry */
|
|
212
|
+
}
|
|
213
|
+
await new Promise((r) => setTimeout(r, 2000));
|
|
214
|
+
}
|
|
215
|
+
throw new Error("Timed out waiting for local Petals shard manager to become ready");
|
|
216
|
+
}
|
|
217
|
+
|
|
218
|
+
export { SHARD_PORT, PETALS_PORT, venvPython };
|
package/src/net.js
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
import { execFileSync } from "node:child_process";
|
|
2
|
+
|
|
3
|
+
export async function detectPublicIp() {
|
|
4
|
+
const tryUrls = [
|
|
5
|
+
"http://metadata.google.internal/computeMetadata/v1/instance/network-interfaces/0/access-configs/0/external-ip",
|
|
6
|
+
"http://169.254.169.254/latest/meta-data/public-ipv4",
|
|
7
|
+
"https://api.ipify.org",
|
|
8
|
+
];
|
|
9
|
+
|
|
10
|
+
for (const url of tryUrls) {
|
|
11
|
+
try {
|
|
12
|
+
const headers = url.includes("metadata.google.internal")
|
|
13
|
+
? { "Metadata-Flavor": "Google" }
|
|
14
|
+
: {};
|
|
15
|
+
const ctrl = AbortSignal.timeout(3000);
|
|
16
|
+
const res = await fetch(url, { headers, signal: ctrl });
|
|
17
|
+
if (!res.ok) continue;
|
|
18
|
+
const text = (await res.text()).trim();
|
|
19
|
+
if (/^\d+\.\d+\.\d+\.\d+$/.test(text)) return text;
|
|
20
|
+
} catch {
|
|
21
|
+
/* try next */
|
|
22
|
+
}
|
|
23
|
+
}
|
|
24
|
+
return "";
|
|
25
|
+
}
|
|
26
|
+
|
|
27
|
+
export function which(cmd) {
|
|
28
|
+
try {
|
|
29
|
+
const out = execFileSync(process.platform === "win32" ? "where" : "which", [cmd], {
|
|
30
|
+
encoding: "utf8",
|
|
31
|
+
});
|
|
32
|
+
return out.split(/\r?\n/)[0].trim();
|
|
33
|
+
} catch {
|
|
34
|
+
return null;
|
|
35
|
+
}
|
|
36
|
+
}
|
package/src/state.js
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
import fs from "node:fs";
|
|
2
|
+
import path from "node:path";
|
|
3
|
+
import { STATE_PATH, HOME_DIR } from "./config.js";
|
|
4
|
+
|
|
5
|
+
export function ensureHome() {
|
|
6
|
+
fs.mkdirSync(HOME_DIR, { recursive: true });
|
|
7
|
+
}
|
|
8
|
+
|
|
9
|
+
export function loadState() {
|
|
10
|
+
try {
|
|
11
|
+
return JSON.parse(fs.readFileSync(STATE_PATH, "utf8"));
|
|
12
|
+
} catch {
|
|
13
|
+
return null;
|
|
14
|
+
}
|
|
15
|
+
}
|
|
16
|
+
|
|
17
|
+
export function saveState(state) {
|
|
18
|
+
ensureHome();
|
|
19
|
+
fs.writeFileSync(STATE_PATH, JSON.stringify(state, null, 2));
|
|
20
|
+
}
|
|
21
|
+
|
|
22
|
+
export function clearState() {
|
|
23
|
+
try {
|
|
24
|
+
fs.unlinkSync(STATE_PATH);
|
|
25
|
+
} catch {
|
|
26
|
+
/* ignore */
|
|
27
|
+
}
|
|
28
|
+
}
|
|
29
|
+
|
|
30
|
+
export function stateDir() {
|
|
31
|
+
return path.dirname(STATE_PATH);
|
|
32
|
+
}
|
package/src/wizard.js
ADDED
|
@@ -0,0 +1,318 @@
|
|
|
1
|
+
import * as p from "@clack/prompts";
|
|
2
|
+
import color from "picocolors";
|
|
3
|
+
import fs from "node:fs";
|
|
4
|
+
import path from "node:path";
|
|
5
|
+
import { spawn } from "node:child_process";
|
|
6
|
+
import { DEFAULT_DISCOVERY_URL, HOME_DIR, SHARD_PORT, PETALS_PORT } from "./config.js";
|
|
7
|
+
import { listMothers, motherHosts, joinHost, readyHost, leaveHost } from "./api.js";
|
|
8
|
+
import { detectPublicIp } from "./net.js";
|
|
9
|
+
import {
|
|
10
|
+
ensureVenvAndPetals,
|
|
11
|
+
startShardManagerProcess,
|
|
12
|
+
waitForShardRunning,
|
|
13
|
+
venvPython,
|
|
14
|
+
} from "./install.js";
|
|
15
|
+
import { saveState, loadState, clearState, ensureHome } from "./state.js";
|
|
16
|
+
|
|
17
|
+
function parseArgs(argv) {
|
|
18
|
+
const out = { cmd: null, discoveryUrl: process.env.BLITZWING_DISCOVERY_URL || DEFAULT_DISCOVERY_URL };
|
|
19
|
+
const rest = [...argv];
|
|
20
|
+
if (rest[0] === "status" || rest[0] === "leave" || rest[0] === "help") {
|
|
21
|
+
out.cmd = rest.shift();
|
|
22
|
+
}
|
|
23
|
+
while (rest.length) {
|
|
24
|
+
const a = rest.shift();
|
|
25
|
+
if (a === "--discovery-url") out.discoveryUrl = rest.shift();
|
|
26
|
+
else if (a === "--help" || a === "-h") out.cmd = "help";
|
|
27
|
+
}
|
|
28
|
+
return out;
|
|
29
|
+
}
|
|
30
|
+
|
|
31
|
+
export async function run(argv) {
|
|
32
|
+
const args = parseArgs(argv);
|
|
33
|
+
if (args.cmd === "help") {
|
|
34
|
+
printHelp();
|
|
35
|
+
return;
|
|
36
|
+
}
|
|
37
|
+
if (args.cmd === "status") {
|
|
38
|
+
await showStatus();
|
|
39
|
+
return;
|
|
40
|
+
}
|
|
41
|
+
if (args.cmd === "leave") {
|
|
42
|
+
await doLeave();
|
|
43
|
+
return;
|
|
44
|
+
}
|
|
45
|
+
await wizard(args);
|
|
46
|
+
}
|
|
47
|
+
|
|
48
|
+
function printHelp() {
|
|
49
|
+
console.log(`
|
|
50
|
+
${color.bold("blitzwing")} — join a Blitzwing mother swarm as a compute node
|
|
51
|
+
|
|
52
|
+
blitzwing Interactive setup wizard (recommended)
|
|
53
|
+
blitzwing status Show local contributor status
|
|
54
|
+
blitzwing leave Leave the swarm and reclaim your layers
|
|
55
|
+
|
|
56
|
+
Env:
|
|
57
|
+
BLITZWING_DISCOVERY_URL Override discovery service (default ${DEFAULT_DISCOVERY_URL})
|
|
58
|
+
`);
|
|
59
|
+
}
|
|
60
|
+
|
|
61
|
+
async function wizard(args) {
|
|
62
|
+
p.intro(color.bgCyan(color.black(" blitzwing ")));
|
|
63
|
+
ensureHome();
|
|
64
|
+
|
|
65
|
+
const spin = p.spinner();
|
|
66
|
+
spin.start("Loading network from Discovery Service…");
|
|
67
|
+
let mothers;
|
|
68
|
+
try {
|
|
69
|
+
mothers = await listMothers(args.discoveryUrl);
|
|
70
|
+
} catch (err) {
|
|
71
|
+
spin.stop("Discovery failed");
|
|
72
|
+
p.cancel(`Could not reach discovery at ${args.discoveryUrl}: ${err.message}`);
|
|
73
|
+
process.exit(1);
|
|
74
|
+
}
|
|
75
|
+
spin.stop(`Found ${mothers.length} model(s)`);
|
|
76
|
+
|
|
77
|
+
if (!mothers.length) {
|
|
78
|
+
p.cancel("No mothers registered yet. Ask the operator to bootstrap a mother node.");
|
|
79
|
+
process.exit(1);
|
|
80
|
+
}
|
|
81
|
+
|
|
82
|
+
let selected = mothers[0];
|
|
83
|
+
if (mothers.length > 1) {
|
|
84
|
+
const choice = await p.select({
|
|
85
|
+
message: "Which model do you want to help serve?",
|
|
86
|
+
options: mothers.map((m) => ({
|
|
87
|
+
value: m.model,
|
|
88
|
+
label: m.model,
|
|
89
|
+
hint: m.mother_url,
|
|
90
|
+
})),
|
|
91
|
+
});
|
|
92
|
+
if (p.isCancel(choice)) {
|
|
93
|
+
p.cancel("Setup cancelled");
|
|
94
|
+
process.exit(0);
|
|
95
|
+
}
|
|
96
|
+
selected = mothers.find((m) => m.model === choice);
|
|
97
|
+
} else {
|
|
98
|
+
p.log.info(`Model: ${color.cyan(selected.model)}`);
|
|
99
|
+
}
|
|
100
|
+
|
|
101
|
+
let maxLayers = selected.total_layers - 1;
|
|
102
|
+
try {
|
|
103
|
+
const hosts = await motherHosts(selected.mother_url);
|
|
104
|
+
if (typeof hosts.max_layers_available === "number") {
|
|
105
|
+
maxLayers = hosts.max_layers_available;
|
|
106
|
+
}
|
|
107
|
+
} catch (err) {
|
|
108
|
+
p.log.warn(`Could not read live host map (${err.message}); using total_layers-1=${maxLayers}`);
|
|
109
|
+
}
|
|
110
|
+
|
|
111
|
+
if (maxLayers < 1) {
|
|
112
|
+
p.cancel("No spare layers available on this swarm right now. Try again later.");
|
|
113
|
+
process.exit(1);
|
|
114
|
+
}
|
|
115
|
+
|
|
116
|
+
const layers = await p.text({
|
|
117
|
+
message: `How many layers can this machine host? (1–${maxLayers})`,
|
|
118
|
+
initialValue: String(Math.min(8, maxLayers)),
|
|
119
|
+
validate(v) {
|
|
120
|
+
const n = Number(v);
|
|
121
|
+
if (!Number.isInteger(n) || n < 1 || n > maxLayers) {
|
|
122
|
+
return `Enter an integer between 1 and ${maxLayers}`;
|
|
123
|
+
}
|
|
124
|
+
},
|
|
125
|
+
});
|
|
126
|
+
if (p.isCancel(layers)) {
|
|
127
|
+
p.cancel("Setup cancelled");
|
|
128
|
+
process.exit(0);
|
|
129
|
+
}
|
|
130
|
+
const layersN = Number(layers);
|
|
131
|
+
|
|
132
|
+
const detected = await detectPublicIp();
|
|
133
|
+
const publicIp = await p.text({
|
|
134
|
+
message: "Public IP for this machine (other peers must reach you)",
|
|
135
|
+
initialValue: detected || "",
|
|
136
|
+
validate(v) {
|
|
137
|
+
if (!v || !v.trim()) return "Public IP is required";
|
|
138
|
+
},
|
|
139
|
+
});
|
|
140
|
+
if (p.isCancel(publicIp)) {
|
|
141
|
+
p.cancel("Setup cancelled");
|
|
142
|
+
process.exit(0);
|
|
143
|
+
}
|
|
144
|
+
|
|
145
|
+
const confirm = await p.confirm({
|
|
146
|
+
message: `Join ${selected.model} hosting ${layersN} layers from ${String(publicIp).trim()}?`,
|
|
147
|
+
initialValue: true,
|
|
148
|
+
});
|
|
149
|
+
if (p.isCancel(confirm) || !confirm) {
|
|
150
|
+
p.cancel("Setup cancelled");
|
|
151
|
+
process.exit(0);
|
|
152
|
+
}
|
|
153
|
+
|
|
154
|
+
spin.start("Preparing Python environment + Petals");
|
|
155
|
+
try {
|
|
156
|
+
ensureVenvAndPetals({
|
|
157
|
+
onLog: (msg) => {
|
|
158
|
+
spin.message(msg);
|
|
159
|
+
},
|
|
160
|
+
});
|
|
161
|
+
} catch (err) {
|
|
162
|
+
spin.stop("Install failed");
|
|
163
|
+
p.cancel(err.message);
|
|
164
|
+
process.exit(1);
|
|
165
|
+
}
|
|
166
|
+
const python = venvPython();
|
|
167
|
+
spin.stop("Environment ready");
|
|
168
|
+
|
|
169
|
+
const shardManagerUrl = `http://${String(publicIp).trim()}:${SHARD_PORT}`;
|
|
170
|
+
|
|
171
|
+
spin.start("Requesting layer assignment from mother…");
|
|
172
|
+
let assignment;
|
|
173
|
+
try {
|
|
174
|
+
assignment = await joinHost(selected.mother_url, {
|
|
175
|
+
model: selected.model,
|
|
176
|
+
layers: layersN,
|
|
177
|
+
public_ip: String(publicIp).trim(),
|
|
178
|
+
shard_manager_url: shardManagerUrl,
|
|
179
|
+
});
|
|
180
|
+
} catch (err) {
|
|
181
|
+
spin.stop("Join rejected");
|
|
182
|
+
const max = err.body?.detail?.max_layers;
|
|
183
|
+
p.cancel(`${err.message}${max != null ? ` (max available: ${max})` : ""}`);
|
|
184
|
+
process.exit(1);
|
|
185
|
+
}
|
|
186
|
+
spin.stop(`Assigned blocks ${assignment.block_indices} (${assignment.layers_hosted} layers)`);
|
|
187
|
+
|
|
188
|
+
const logPath = path.join(HOME_DIR, "shard_manager.log");
|
|
189
|
+
spin.start("Starting local Petals server…");
|
|
190
|
+
const { pid } = startShardManagerProcess({
|
|
191
|
+
python,
|
|
192
|
+
logPath,
|
|
193
|
+
env: {
|
|
194
|
+
MODEL_NAME: selected.model,
|
|
195
|
+
PUBLIC_IP: String(publicIp).trim(),
|
|
196
|
+
BLOCK_INDICES: assignment.block_indices,
|
|
197
|
+
INITIAL_PEERS: (assignment.initial_peers || []).join(","),
|
|
198
|
+
PETALS_PORT: String(PETALS_PORT),
|
|
199
|
+
IDENTITY_PATH: path.join(HOME_DIR, "petals-identity"),
|
|
200
|
+
SHARD_AUTO_START: "1",
|
|
201
|
+
},
|
|
202
|
+
});
|
|
203
|
+
|
|
204
|
+
try {
|
|
205
|
+
await waitForShardRunning({ timeoutMs: 600000 });
|
|
206
|
+
} catch (err) {
|
|
207
|
+
spin.stop("Petals did not become ready");
|
|
208
|
+
p.cancel(`${err.message}. See ${logPath}`);
|
|
209
|
+
process.exit(1);
|
|
210
|
+
}
|
|
211
|
+
spin.stop("Petals is serving your layers");
|
|
212
|
+
|
|
213
|
+
spin.start("Finalizing handoff with mother…");
|
|
214
|
+
try {
|
|
215
|
+
await readyHost(selected.mother_url, { host_id: assignment.host_id });
|
|
216
|
+
} catch (err) {
|
|
217
|
+
spin.stop("Handoff failed");
|
|
218
|
+
p.cancel(err.message);
|
|
219
|
+
process.exit(1);
|
|
220
|
+
}
|
|
221
|
+
spin.stop("Handoff complete");
|
|
222
|
+
|
|
223
|
+
const state = {
|
|
224
|
+
host_id: assignment.host_id,
|
|
225
|
+
model: selected.model,
|
|
226
|
+
mother_url: selected.mother_url,
|
|
227
|
+
block_indices: assignment.block_indices,
|
|
228
|
+
layers_hosted: assignment.layers_hosted,
|
|
229
|
+
public_ip: String(publicIp).trim(),
|
|
230
|
+
shard_manager_url: shardManagerUrl,
|
|
231
|
+
shard_pid: pid,
|
|
232
|
+
discovery_url: args.discoveryUrl,
|
|
233
|
+
joined_at: new Date().toISOString(),
|
|
234
|
+
};
|
|
235
|
+
saveState(state);
|
|
236
|
+
startHeartbeatDaemon(state);
|
|
237
|
+
|
|
238
|
+
p.outro(
|
|
239
|
+
`${color.green("You are online.")} Hosting ${color.cyan(String(state.layers_hosted))} layers of ${color.cyan(state.model)} at ${state.block_indices}\n` +
|
|
240
|
+
`Run ${color.bold("blitzwing status")} anytime, or ${color.bold("blitzwing leave")} to exit.`
|
|
241
|
+
);
|
|
242
|
+
}
|
|
243
|
+
|
|
244
|
+
function startHeartbeatDaemon(state) {
|
|
245
|
+
const hbPath = path.join(HOME_DIR, "heartbeat.mjs");
|
|
246
|
+
fs.writeFileSync(
|
|
247
|
+
hbPath,
|
|
248
|
+
`const mother = ${JSON.stringify(state.mother_url)};
|
|
249
|
+
const hostId = ${JSON.stringify(state.host_id)};
|
|
250
|
+
async function beat() {
|
|
251
|
+
try {
|
|
252
|
+
await fetch(mother.replace(/\\/$/, "") + "/v1/hosts/heartbeat", {
|
|
253
|
+
method: "POST",
|
|
254
|
+
headers: { "Content-Type": "application/json" },
|
|
255
|
+
body: JSON.stringify({ host_id: hostId }),
|
|
256
|
+
});
|
|
257
|
+
} catch {}
|
|
258
|
+
}
|
|
259
|
+
setInterval(beat, 60000);
|
|
260
|
+
beat();
|
|
261
|
+
`
|
|
262
|
+
);
|
|
263
|
+
const child = spawn(process.execPath, [hbPath], {
|
|
264
|
+
detached: true,
|
|
265
|
+
stdio: "ignore",
|
|
266
|
+
});
|
|
267
|
+
child.unref();
|
|
268
|
+
}
|
|
269
|
+
|
|
270
|
+
async function showStatus() {
|
|
271
|
+
const state = loadState();
|
|
272
|
+
if (!state) {
|
|
273
|
+
console.log("Not joined. Run blitzwing to set up this machine as a node.");
|
|
274
|
+
return;
|
|
275
|
+
}
|
|
276
|
+
console.log(color.bold("Contributor status"));
|
|
277
|
+
console.log(` model: ${state.model}`);
|
|
278
|
+
console.log(` host_id: ${state.host_id}`);
|
|
279
|
+
console.log(` layers_hosted: ${state.layers_hosted}`);
|
|
280
|
+
console.log(` block_indices: ${state.block_indices}`);
|
|
281
|
+
console.log(` mother: ${state.mother_url}`);
|
|
282
|
+
console.log(` public_ip: ${state.public_ip}`);
|
|
283
|
+
try {
|
|
284
|
+
const res = await fetch(`http://127.0.0.1:${SHARD_PORT}/status`);
|
|
285
|
+
if (res.ok) {
|
|
286
|
+
const st = await res.json();
|
|
287
|
+
console.log(` petals_running: ${st.running}`);
|
|
288
|
+
console.log(` petals_pid: ${st.pid}`);
|
|
289
|
+
} else {
|
|
290
|
+
console.log(" petals_running: unknown");
|
|
291
|
+
}
|
|
292
|
+
} catch {
|
|
293
|
+
console.log(` petals_running: no (shard manager not reachable on :${SHARD_PORT})`);
|
|
294
|
+
}
|
|
295
|
+
}
|
|
296
|
+
|
|
297
|
+
async function doLeave() {
|
|
298
|
+
const state = loadState();
|
|
299
|
+
if (!state) {
|
|
300
|
+
p.cancel("Not joined.");
|
|
301
|
+
process.exit(1);
|
|
302
|
+
}
|
|
303
|
+
const spin = p.spinner();
|
|
304
|
+
spin.start("Leaving swarm…");
|
|
305
|
+
try {
|
|
306
|
+
await leaveHost(state.mother_url, { host_id: state.host_id });
|
|
307
|
+
} catch (err) {
|
|
308
|
+
p.log.warn(`Mother leave call failed: ${err.message}`);
|
|
309
|
+
}
|
|
310
|
+
try {
|
|
311
|
+
await fetch(`http://127.0.0.1:${SHARD_PORT}/stop`, { method: "POST" });
|
|
312
|
+
} catch {
|
|
313
|
+
/* ignore */
|
|
314
|
+
}
|
|
315
|
+
clearState();
|
|
316
|
+
spin.stop("Left swarm");
|
|
317
|
+
p.outro("Your layers were reclaimed by the mother (when reachable).");
|
|
318
|
+
}
|