blitzwing 0.1.1 → 0.1.2

package/package.json

@@ -1,31 +1,31 @@
-{
-  "name": "blitzwing",
-  "version": "0.1.1",
-  "description": "Interactive setup wizard to join a Blitzwing Petals mother swarm as a compute contributor",
-  "bin": {
-    "blitzwing": "bin/blitzwing.js"
-  },
-  "type": "module",
-  "engines": {
-    "node": ">=18"
-  },
-  "files": [
-    "bin",
-    "src",
-    "README.md"
-  ],
-  "keywords": [
-    "blitzwing",
-    "petals",
-    "llm",
-    "distributed-inference"
-  ],
-  "license": "MIT",
-  "config": {
-    "discoveryUrl": "https://f44c-103-98-63-33.ngrok-free.app"
-  },
-  "dependencies": {
-    "@clack/prompts": "^0.9.1",
-    "picocolors": "^1.1.1"
-  }
-}
+{
+  "name": "blitzwing",
+  "version": "0.1.2",
+  "description": "Interactive setup wizard to join a Blitzwing Petals mother swarm as a compute contributor",
+  "bin": {
+    "blitzwing": "bin/blitzwing.js"
+  },
+  "type": "module",
+  "engines": {
+    "node": ">=18"
+  },
+  "files": [
+    "bin",
+    "src",
+    "README.md"
+  ],
+  "keywords": [
+    "blitzwing",
+    "petals",
+    "llm",
+    "distributed-inference"
+  ],
+  "license": "MIT",
+  "config": {
+    "discoveryUrl": "https://f44c-103-98-63-33.ngrok-free.app"
+  },
+  "dependencies": {
+    "@clack/prompts": "^0.9.1",
+    "picocolors": "^1.1.1"
+  }
+}

package/src/install.js

@@ -1,218 +1,277 @@
-import { spawn, execFileSync } from "node:child_process";
-import fs from "node:fs";
-import path from "node:path";
-import { HOME_DIR, VENV_PATH, SHARD_PORT, PETALS_PORT } from "./config.js";
-import { which } from "./net.js";
-
-function venvPython() {
-  if (process.platform === "win32") {
-    return path.join(VENV_PATH, "Scripts", "python.exe");
-  }
-  return path.join(VENV_PATH, "bin", "python");
-}
-
-function venvPip() {
-  if (process.platform === "win32") {
-    return path.join(VENV_PATH, "Scripts", "pip.exe");
-  }
-  return path.join(VENV_PATH, "bin", "pip");
-}
-
-export function ensurePython() {
-  const py = which("python3") || which("python");
-  if (!py) {
-    throw new Error("Python 3 is required. Install Python 3.10+ and re-run blitzwing.");
-  }
-  const ver = execFileSync(py, ["--version"], { encoding: "utf8" });
-  return { py, ver: ver.trim() };
-}
-
-export function ensureVenvAndPetals({ onLog }) {
-  const { py } = ensurePython();
-  fs.mkdirSync(HOME_DIR, { recursive: true });
-
-  if (!fs.existsSync(venvPython())) {
-    onLog?.(`Creating venv at ${VENV_PATH}`);
-    execFileSync(py, ["-m", "venv", VENV_PATH], { stdio: "inherit" });
-  }
-
-  const pip = venvPip();
-  const python = venvPython();
-
-  onLog?.("Installing CPU PyTorch + Petals (may take a few minutes on first run)...");
-  execFileSync(pip, ["install", "-U", "pip"], { stdio: "inherit" });
-  execFileSync(
-    pip,
-    ["install", "torch", "--index-url", "https://download.pytorch.org/whl/cpu"],
-    { stdio: "inherit" }
-  );
-  // Prefer git install for contributors without the monorepo
-  execFileSync(
-    pip,
-    ["install", "git+https://github.com/bigscience-workshop/petals.git"],
-    { stdio: "inherit" }
-  );
-  execFileSync(pip, ["install", "fastapi", "uvicorn[standard]", "pydantic", "httpx"], {
-    stdio: "inherit",
-  });
-
-  return { python, pip };
-}
-
-/**
- * Write a tiny local shard-manager runner script so contributors don't need the monorepo.
- * Embeds a minimal supervisor compatible with mother /reload API.
- */
-export function writeLocalShardManager() {
-  const dir = path.join(HOME_DIR, "runtime");
-  fs.mkdirSync(dir, { recursive: true });
-  const appPath = path.join(dir, "shard_manager_app.py");
-  const code = `
-import os, signal, subprocess, threading, time
-from typing import Optional, List
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel, Field
-
-class ReloadRequest(BaseModel):
-    block_indices: str = Field(..., pattern=r"^\\d+:\\d+$")
-
-class StatusResponse(BaseModel):
-    running: bool
-    pid: Optional[int] = None
-    block_indices: Optional[str] = None
-    model: str
-    public_ip: Optional[str] = None
-    port: int
-    last_exit_code: Optional[int] = None
-
-class Mgr:
-    def __init__(self):
-        self.model = os.environ["MODEL_NAME"]
-        self.public_ip = os.environ.get("PUBLIC_IP")
-        self.port = int(os.environ.get("PETALS_PORT", "31337"))
-        self.identity_path = os.environ.get("IDENTITY_PATH", os.path.expanduser("~/.blitzwing/petals-identity"))
-        self.block_indices = os.environ["BLOCK_INDICES"]
-        peers = [p.strip() for p in os.environ.get("INITIAL_PEERS", "").split(",") if p.strip()]
-        self.initial_peers = peers
-        self.python = os.environ.get("PETALS_PYTHON", "python")
-        self._proc = None
-        self._lock = threading.Lock()
-        self.last_exit_code = None
-
-    def cmd(self, bi):
-        c = [self.python, "-m", "petals.cli.run_server", self.model,
-             "--device", "cpu", "--quant_type", "none",
-             "--block_indices", bi, "--port", str(self.port),
-             "--identity_path", self.identity_path, "--num_handlers", "1"]
-        if self.public_ip:
-            c += ["--public_ip", self.public_ip]
-        if self.initial_peers:
-            c += ["--initial_peers", *self.initial_peers]
-        return c
-
-    def start(self, bi=None):
-        with self._lock:
-            if bi: self.block_indices = bi
-            if self._proc and self._proc.poll() is None:
-                return
-            self._proc = subprocess.Popen(self.cmd(self.block_indices))
-            self.last_exit_code = None
-
-    def stop(self):
-        with self._lock:
-            if not self._proc or self._proc.poll() is not None:
-                self._proc = None
-                return
-            p = self._proc
-            p.send_signal(signal.SIGTERM)
-            try: p.wait(timeout=30)
-            except Exception:
-                p.kill(); p.wait(timeout=10)
-            self.last_exit_code = p.returncode
-            self._proc = None
-
-    def reload(self, bi):
-        self.stop(); time.sleep(1); self.start(bi)
-
-    def status(self):
-        running = self._proc is not None and self._proc.poll() is None
-        return StatusResponse(
-            running=running,
-            pid=(self._proc.pid if running else None),
-            block_indices=self.block_indices,
-            model=self.model,
-            public_ip=self.public_ip,
-            port=self.port,
-            last_exit_code=self.last_exit_code,
-        )
-
-mgr = Mgr()
-app = FastAPI()
-
-@app.on_event("startup")
-def _up():
-    mgr.start()
-
-@app.on_event("shutdown")
-def _down():
-    mgr.stop()
-
-@app.get("/status")
-def status():
-    return mgr.status()
-
-@app.post("/reload")
-def reload(body: ReloadRequest):
-    a,b = map(int, body.block_indices.split(":"))
-    if b <= a: raise HTTPException(400, "bad range")
-    mgr.reload(body.block_indices)
-    time.sleep(0.5)
-    return mgr.status()
-
-@app.post("/stop")
-def stop():
-    mgr.stop(); return mgr.status()
-`;
-  fs.writeFileSync(appPath, code);
-  return appPath;
-}
-
-export function startShardManagerProcess({
-  python,
-  env,
-  logPath,
-}) {
-  const appPath = writeLocalShardManager();
-  const out = fs.openSync(logPath, "a");
-  const child = spawn(
-    python,
-    ["-m", "uvicorn", `shard_manager_app:app`, "--host", "0.0.0.0", "--port", String(SHARD_PORT)],
-    {
-      cwd: path.dirname(appPath),
-      env: { ...process.env, ...env, PETALS_PYTHON: python },
-      detached: true,
-      stdio: ["ignore", out, out],
-    }
-  );
-  child.unref();
-  return { pid: child.pid, logPath, appPath };
-}
-
-export async function waitForShardRunning({ timeoutMs = 300000 } = {}) {
-  const start = Date.now();
-  while (Date.now() - start < timeoutMs) {
-    try {
-      const res = await fetch(`http://127.0.0.1:${SHARD_PORT}/status`);
-      if (res.ok) {
-        const body = await res.json();
-        if (body.running) return body;
-      }
-    } catch {
-      /* retry */
-    }
-    await new Promise((r) => setTimeout(r, 2000));
-  }
-  throw new Error("Timed out waiting for local Petals shard manager to become ready");
-}
-
-export { SHARD_PORT, PETALS_PORT, venvPython };
+import { spawn, execFileSync } from "node:child_process";
+import fs from "node:fs";
+import path from "node:path";
+import { HOME_DIR, VENV_PATH, SHARD_PORT, PETALS_PORT } from "./config.js";
+import { which } from "./net.js";
+
+function venvPython() {
+  if (process.platform === "win32") {
+    return path.join(VENV_PATH, "Scripts", "python.exe");
+  }
+  return path.join(VENV_PATH, "bin", "python");
+}
+
+function venvPip() {
+  if (process.platform === "win32") {
+    return path.join(VENV_PATH, "Scripts", "pip.exe");
+  }
+  return path.join(VENV_PATH, "bin", "pip");
+}
+
+export function ensurePython() {
+  // Prefer 3.10/3.11 — Petals/hivemind break on 3.12+ / 3.13+
+  const candidates = ["python3.11", "python3.10", "python3", "python"];
+  let py = null;
+  let ver = "";
+  for (const cmd of candidates) {
+    const found = which(cmd);
+    if (!found) continue;
+    try {
+      ver = execFileSync(found, ["--version"], { encoding: "utf8" }).trim();
+      const m = ver.match(/Python (\d+)\.(\d+)/);
+      if (!m) continue;
+      const major = Number(m[1]);
+      const minor = Number(m[2]);
+      if (major === 3 && minor >= 10 && minor <= 11) {
+        py = found;
+        break;
+      }
+      if (!py) {
+        py = found; // fallback, may fail later
+      }
+    } catch {
+      /* try next */
+    }
+  }
+  if (!py) {
+    throw new Error("Python 3.10 or 3.11 is required. Install via conda/miniconda and re-run.");
+  }
+  const m = ver.match(/Python (\d+)\.(\d+)/);
+  if (m && (Number(m[1]) !== 3 || Number(m[2]) > 11)) {
+    throw new Error(
+      `${ver} is too new for Petals. Use Python 3.11 (e.g. conda create -n blitzwing python=3.11).`
+    );
+  }
+  return { py, ver };
+}
+
+function petalsImportable(python) {
+  try {
+    execFileSync(python, ["-c", "import petals"], { stdio: "ignore" });
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+export function ensureVenvAndPetals({ onLog }) {
+  const { py } = ensurePython();
+  fs.mkdirSync(HOME_DIR, { recursive: true });
+
+  if (!fs.existsSync(venvPython())) {
+    onLog?.(`Creating venv at ${VENV_PATH}`);
+    execFileSync(py, ["-m", "venv", VENV_PATH], { stdio: "inherit" });
+  }
+
+  const pip = venvPip();
+  const python = venvPython();
+  const env = {
+    ...process.env,
+    PIP_NO_BUILD_ISOLATION: "1",
+  };
+
+  if (petalsImportable(python)) {
+    onLog?.("Petals already installed — skipping dependency install");
+    return { python, pip };
+  }
+
+  onLog?.("Installing CPU PyTorch + Petals (may take a few minutes on first run)...");
+  execFileSync(pip, ["install", "-U", "pip", "wheel", "setuptools<81"], {
+    stdio: "inherit",
+    env,
+  });
+  execFileSync(
+    pip,
+    ["install", "torch", "--index-url", "https://download.pytorch.org/whl/cpu"],
+    { stdio: "inherit", env }
+  );
+  // Hivemind first with no build isolation (needs pkg_resources from setuptools<81)
+  execFileSync(
+    pip,
+    [
+      "install",
+      "--no-build-isolation",
+      "git+https://github.com/learning-at-home/hivemind.git@213bff98a62accb91f254e2afdccbf1d69ebdea9",
+    ],
+    { stdio: "inherit", env }
+  );
+  execFileSync(
+    pip,
+    ["install", "--no-build-isolation", "git+https://github.com/bigscience-workshop/petals.git"],
+    { stdio: "inherit", env }
+  );
+  execFileSync(pip, ["install", "fastapi", "uvicorn[standard]", "pydantic", "httpx"], {
+    stdio: "inherit",
+    env,
+  });
+
+  return { python, pip };
+}
+
+/**
+ * Write a tiny local shard-manager runner script so contributors don't need the monorepo.
+ * Embeds a minimal supervisor compatible with mother /reload API.
+ */
+export function writeLocalShardManager() {
+  const dir = path.join(HOME_DIR, "runtime");
+  fs.mkdirSync(dir, { recursive: true });
+  const appPath = path.join(dir, "shard_manager_app.py");
+  const code = `
+import os, signal, subprocess, threading, time
+from typing import Optional, List
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+
+class ReloadRequest(BaseModel):
+    block_indices: str = Field(..., pattern=r"^\\d+:\\d+$")
+
+class StatusResponse(BaseModel):
+    running: bool
+    pid: Optional[int] = None
+    block_indices: Optional[str] = None
+    model: str
+    public_ip: Optional[str] = None
+    port: int
+    last_exit_code: Optional[int] = None
+
+class Mgr:
+    def __init__(self):
+        self.model = os.environ["MODEL_NAME"]
+        self.public_ip = os.environ.get("PUBLIC_IP")
+        self.port = int(os.environ.get("PETALS_PORT", "31337"))
+        self.identity_path = os.environ.get("IDENTITY_PATH", os.path.expanduser("~/.blitzwing/petals-identity"))
+        self.block_indices = os.environ["BLOCK_INDICES"]
+        peers = [p.strip() for p in os.environ.get("INITIAL_PEERS", "").split(",") if p.strip()]
+        self.initial_peers = peers
+        self.python = os.environ.get("PETALS_PYTHON", "python")
+        self._proc = None
+        self._lock = threading.Lock()
+        self.last_exit_code = None
+
+    def cmd(self, bi):
+        c = [self.python, "-m", "petals.cli.run_server", self.model,
+             "--device", "cpu", "--quant_type", "none",
+             "--block_indices", bi, "--port", str(self.port),
+             "--identity_path", self.identity_path, "--num_handlers", "1"]
+        if self.public_ip:
+            c += ["--public_ip", self.public_ip]
+        if self.initial_peers:
+            c += ["--initial_peers", *self.initial_peers]
+        return c
+
+    def start(self, bi=None):
+        with self._lock:
+            if bi: self.block_indices = bi
+            if self._proc and self._proc.poll() is None:
+                return
+            self._proc = subprocess.Popen(self.cmd(self.block_indices))
+            self.last_exit_code = None
+
+    def stop(self):
+        with self._lock:
+            if not self._proc or self._proc.poll() is not None:
+                self._proc = None
+                return
+            p = self._proc
+            p.send_signal(signal.SIGTERM)
+            try: p.wait(timeout=30)
+            except Exception:
+                p.kill(); p.wait(timeout=10)
+            self.last_exit_code = p.returncode
+            self._proc = None
+
+    def reload(self, bi):
+        self.stop(); time.sleep(1); self.start(bi)
+
+    def status(self):
+        running = self._proc is not None and self._proc.poll() is None
+        return StatusResponse(
+            running=running,
+            pid=(self._proc.pid if running else None),
+            block_indices=self.block_indices,
+            model=self.model,
+            public_ip=self.public_ip,
+            port=self.port,
+            last_exit_code=self.last_exit_code,
+        )
+
+mgr = Mgr()
+app = FastAPI()
+
+@app.on_event("startup")
+def _up():
+    mgr.start()
+
+@app.on_event("shutdown")
+def _down():
+    mgr.stop()
+
+@app.get("/status")
+def status():
+    return mgr.status()
+
+@app.post("/reload")
+def reload(body: ReloadRequest):
+    a,b = map(int, body.block_indices.split(":"))
+    if b <= a: raise HTTPException(400, "bad range")
+    mgr.reload(body.block_indices)
+    time.sleep(0.5)
+    return mgr.status()
+
+@app.post("/stop")
+def stop():
+    mgr.stop(); return mgr.status()
+`;
+  fs.writeFileSync(appPath, code);
+  return appPath;
+}
+
+export function startShardManagerProcess({
+  python,
+  env,
+  logPath,
+}) {
+  const appPath = writeLocalShardManager();
+  const out = fs.openSync(logPath, "a");
+  const child = spawn(
+    python,
+    ["-m", "uvicorn", `shard_manager_app:app`, "--host", "0.0.0.0", "--port", String(SHARD_PORT)],
+    {
+      cwd: path.dirname(appPath),
+      env: { ...process.env, ...env, PETALS_PYTHON: python },
+      detached: true,
+      stdio: ["ignore", out, out],
+    }
+  );
+  child.unref();
+  return { pid: child.pid, logPath, appPath };
+}
+
+export async function waitForShardRunning({ timeoutMs = 300000 } = {}) {
+  const start = Date.now();
+  while (Date.now() - start < timeoutMs) {
+    try {
+      const res = await fetch(`http://127.0.0.1:${SHARD_PORT}/status`);
+      if (res.ok) {
+        const body = await res.json();
+        if (body.running) return body;
+      }
+    } catch {
+      /* retry */
+    }
+    await new Promise((r) => setTimeout(r, 2000));
+  }
+  throw new Error("Timed out waiting for local Petals shard manager to become ready");
+}
+
+export { SHARD_PORT, PETALS_PORT, venvPython };