blitzwing 0.1.5 → 0.1.9

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "blitzwing",
-  "version": "0.1.5",
+  "version": "0.1.9",
   "description": "Interactive setup wizard to join a Blitzwing Petals mother swarm as a compute contributor",
   "bin": {
     "blitzwing": "bin/blitzwing.js"
@@ -22,7 +22,7 @@
   ],
   "license": "MIT",
   "config": {
-    "discoveryUrl": "https://f44c-103-98-63-33.ngrok-free.app"
+    "discoveryUrl": "http://35.238.86.1:9000"
   },
   "dependencies": {
     "@clack/prompts": "^0.9.1",

package/src/config.js

@@ -2,10 +2,13 @@
 export const DEFAULT_DISCOVERY_URL =
   process.env.BLITZWING_DISCOVERY_URL ||
   process.env.npm_package_config_discoveryUrl ||
-  "https://f44c-103-98-63-33.ngrok-free.app";
+  "http://35.238.86.1:9000";
 
 export const HOME_DIR = process.env.BLITZWING_HOME || `${process.env.HOME || process.env.USERPROFILE}/.blitzwing`;
 export const STATE_PATH = `${HOME_DIR}/contributor.json`;
 export const VENV_PATH = `${HOME_DIR}/venv`;
-export const SHARD_PORT = Number(process.env.BLITZWING_SHARD_PORT || 8001);
-export const PETALS_PORT = Number(process.env.BLITZWING_PETALS_PORT || 31337);
+// Windows + WSL on same PC: avoid 8001/31337 — WSL forwards those on 127.0.0.1.
+const defaultShardPort = process.platform === "win32" ? 8011 : 8001;
+const defaultPetalsPort = process.platform === "win32" ? 31338 : 31337;
+export const SHARD_PORT = Number(process.env.BLITZWING_SHARD_PORT || defaultShardPort);
+export const PETALS_PORT = Number(process.env.BLITZWING_PETALS_PORT || defaultPetalsPort);

package/src/install.js

@@ -101,13 +101,35 @@ export function ensurePython() {
 
 function petalsImportable(python) {
   try {
-    execFileSync(python, ["-c", "import petals"], { stdio: "ignore" });
+    execFileSync(
+      python,
+      ["-c", "import petals.cli.run_server"],
+      { stdio: "ignore" }
+    );
     return true;
   } catch {
     return false;
   }
 }
 
+function uvicornImportable(python) {
+  try {
+    execFileSync(python, ["-c", "import uvicorn, fastapi"], { stdio: "ignore" });
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+function ensureShardManagerDeps(python, pip, env, onLog) {
+  if (uvicornImportable(python)) return;
+  onLog?.("Installing shard manager deps (uvicorn, fastapi)...");
+  execFileSync(python, ["-m", "pip", "install", "fastapi", "uvicorn[standard]", "pydantic"], {
+    stdio: "inherit",
+    env,
+  });
+}
+
 export function ensureVenvAndPetals({ onLog }) {
   const { py } = ensurePython();
   fs.mkdirSync(HOME_DIR, { recursive: true });
@@ -125,29 +147,43 @@ export function ensureVenvAndPetals({ onLog }) {
   };
 
   if (petalsImportable(python)) {
+    ensureShardManagerDeps(python, pip, env, onLog);
     onLog?.("Petals already installed — skipping dependency install");
+    fs.writeFileSync(path.join(HOME_DIR, "python"), python + "\n");
     return { python, pip };
   }
 
+  if (process.platform === "win32") {
+    throw new Error(
+      "Petals cannot be installed on native Windows (hivemind requires uvloop/Linux).\n" +
+        "Run the contributor inside WSL instead:\n" +
+        "  wsl bash /mnt/c/Users/MSI/Desktop/blitzwing/scripts/local_wsl_contributor_join.sh\n" +
+        "Or run the full E2E script:\n" +
+        "  powershell -File scripts/run_local_e2e.ps1"
+    );
+  }
+
   onLog?.("Installing CPU PyTorch + Petals (may take a few minutes on first run)...");
-  execFileSync(pip, ["install", "-U", "pip", "wheel", "setuptools<81"], {
+  execFileSync(python, ["-m", "pip", "install", "-U", "pip", "wheel", "setuptools<81"], {
     stdio: "inherit",
     env,
   });
   execFileSync(
-    pip,
-    ["install", "torch", "--index-url", "https://download.pytorch.org/whl/cpu"],
+    python,
+    ["-m", "pip", "install", "torch", "--index-url", "https://download.pytorch.org/whl/cpu"],
     { stdio: "inherit", env }
   );
   // Needed to compile hivemind protobufs during --no-build-isolation install
-  execFileSync(pip, ["install", "grpcio", "grpcio-tools", "protobuf"], {
+  execFileSync(python, ["-m", "pip", "install", "grpcio", "grpcio-tools", "protobuf"], {
     stdio: "inherit",
     env,
   });
   // Hivemind first with no build isolation (needs pkg_resources from setuptools<81)
   execFileSync(
-    pip,
+    python,
     [
+      "-m",
+      "pip",
       "install",
       "--no-build-isolation",
       "git+https://github.com/learning-at-home/hivemind.git@213bff98a62accb91f254e2afdccbf1d69ebdea9",
@@ -155,15 +191,16 @@ export function ensureVenvAndPetals({ onLog }) {
     { stdio: "inherit", env }
   );
   execFileSync(
-    pip,
-    ["install", "--no-build-isolation", "git+https://github.com/bigscience-workshop/petals.git"],
+    python,
+    ["-m", "pip", "install", "--no-build-isolation", "git+https://github.com/bigscience-workshop/petals.git"],
     { stdio: "inherit", env }
   );
-  execFileSync(pip, ["install", "fastapi", "uvicorn[standard]", "pydantic", "httpx"], {
+  execFileSync(python, ["-m", "pip", "install", "fastapi", "uvicorn[standard]", "pydantic", "httpx"], {
     stdio: "inherit",
     env,
   });
 
+  fs.writeFileSync(path.join(HOME_DIR, "python"), python + "\n");
   return { python, pip };
 }
 
@@ -183,6 +220,7 @@ from pydantic import BaseModel, Field
 
 class ReloadRequest(BaseModel):
     block_indices: str = Field(..., pattern=r"^\\d+:\\d+$")
+    initial_peers: Optional[List[str]] = None
 
 class StatusResponse(BaseModel):
     running: bool
@@ -202,8 +240,12 @@ class Mgr:
         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
+        announce_raw = os.environ.get("ANNOUNCE_MADDRS", "")
+        self.announce_maddrs = [a.strip() for a in announce_raw.split(",") if a.strip()]
         self.python = os.environ.get("PETALS_PYTHON", "python")
+        self.petals_log = os.environ.get("PETALS_LOG", os.path.expanduser("~/.blitzwing/petals.log"))
         self._proc = None
+        self._log_fp = None
         self._lock = threading.Lock()
         self.last_exit_code = None
 
@@ -212,10 +254,19 @@ class Mgr:
              "--device", "cpu", "--quant_type", "none",
              "--block_indices", bi, "--port", str(self.port),
              "--identity_path", self.identity_path, "--num_handlers", "1"]
-        if self.public_ip:
+        if self.announce_maddrs:
+            c += ["--announce_maddrs", *self.announce_maddrs]
+        elif self.public_ip:
             c += ["--public_ip", self.public_ip]
         if self.initial_peers:
             c += ["--initial_peers", *self.initial_peers]
+        use_auto_relay = os.environ.get("PETALS_USE_AUTO_RELAY", "1") not in ("0", "false", "False")
+        if not self.announce_maddrs and use_auto_relay:
+            pass
+        elif not use_auto_relay:
+            c += ["--no_auto_relay"]
+        if os.environ.get("PETALS_SKIP_REACHABILITY_CHECK", "1") in ("1", "true", "True"):
+            c += ["--skip_reachability_check"]
         return c
 
     def start(self, bi=None):
@@ -223,7 +274,15 @@ class Mgr:
             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))
+            if self._log_fp is None:
+                os.makedirs(os.path.dirname(self.petals_log) or ".", exist_ok=True)
+                self._log_fp = open(self.petals_log, "a", encoding="utf-8")
+            popen_kwargs = {"stdout": self._log_fp, "stderr": subprocess.STDOUT}
+            if os.name == "nt":
+                popen_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
+            else:
+                popen_kwargs["preexec_fn"] = os.setsid
+            self._proc = subprocess.Popen(self.cmd(self.block_indices), **popen_kwargs)
             self.last_exit_code = None
 
     def stop(self):
@@ -232,17 +291,28 @@ class Mgr:
                 self._proc = None
                 return
             p = self._proc
-            p.send_signal(signal.SIGTERM)
+            try:
+                if os.name == "nt":
+                    p.send_signal(signal.CTRL_BREAK_EVENT)
+                else:
+                    os.killpg(os.getpgid(p.pid), signal.SIGTERM)
+            except Exception:
+                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):
+    def reload(self, bi, initial_peers=None):
+        if initial_peers is not None:
+            self.initial_peers = initial_peers
         self.stop(); time.sleep(1); self.start(bi)
 
     def status(self):
+        if self._proc is not None and self._proc.poll() is not None:
+            self.last_exit_code = self._proc.returncode
+            self._proc = None
         running = self._proc is not None and self._proc.poll() is None
         return StatusResponse(
             running=running,
@@ -273,7 +343,7 @@ def status():
 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)
+    mgr.reload(body.block_indices, initial_peers=body.initial_peers)
     time.sleep(0.5)
     return mgr.status()
 
@@ -292,12 +362,13 @@ export function startShardManagerProcess({
 }) {
   const appPath = writeLocalShardManager();
   const out = fs.openSync(logPath, "a");
+  const petalsPy = petalsImportable(venvPython()) ? venvPython() : python;
   const child = spawn(
-    python,
+    petalsPy,
     ["-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 },
+      env: { ...process.env, ...env, PETALS_PYTHON: petalsPy, PETALS_LOG: path.join(HOME_DIR, "petals.log") },
       detached: true,
       stdio: ["ignore", out, out],
     }
@@ -306,21 +377,37 @@ export function startShardManagerProcess({
   return { pid: child.pid, logPath, appPath };
 }
 
-export async function waitForShardRunning({ timeoutMs = 300000 } = {}) {
+export async function waitForShardRunning({ timeoutMs = 300000, statusHost } = {}) {
+  const host = statusHost || process.env.BLITZWING_STATUS_HOST || "127.0.0.1";
+  const petalsLog = path.join(HOME_DIR, "petals.log");
   const start = Date.now();
   while (Date.now() - start < timeoutMs) {
     try {
-      const res = await fetch(`http://127.0.0.1:${SHARD_PORT}/status`);
+      const res = await fetch(`http://${host}:${SHARD_PORT}/status`);
       if (res.ok) {
         const body = await res.json();
-        if (body.running) return body;
+        if (body.running && body.pid) return body;
+        if (body.last_exit_code != null) {
+          const tail = fs.existsSync(petalsLog)
+            ? fs.readFileSync(petalsLog, "utf8").trim().split("\n").slice(-5).join("\n")
+            : "";
+          throw new Error(
+            `Petals exited (code ${body.last_exit_code}). ${tail || "See " + petalsLog}`
+          );
+        }
       }
-    } catch {
+    } catch (err) {
+      if (err instanceof Error && err.message.startsWith("Petals exited")) throw err;
       /* retry */
     }
     await new Promise((r) => setTimeout(r, 2000));
   }
-  throw new Error("Timed out waiting for local Petals shard manager to become ready");
+  const tail = fs.existsSync(petalsLog)
+    ? fs.readFileSync(petalsLog, "utf8").trim().split("\n").slice(-5).join("\n")
+    : "";
+  throw new Error(
+    `Timed out waiting for Petals on http://${host}:${SHARD_PORT}/status. ${tail || "See " + petalsLog}`
+  );
 }
 
 export { SHARD_PORT, PETALS_PORT, venvPython };

package/src/net.js

@@ -24,6 +24,59 @@ export async function detectPublicIp() {
   return "";
 }
 
+/** Local LAN/WSL IP for shard_manager_url metadata (not Petals announce). */
+export async function detectLocalIp() {
+  if (process.env.BLITZWING_LOCAL_IP) {
+    return process.env.BLITZWING_LOCAL_IP.trim();
+  }
+  try {
+    const { networkInterfaces } = await import("node:os");
+    const nets = networkInterfaces();
+    for (const name of Object.keys(nets)) {
+      for (const net of nets[name] || []) {
+        if (net.family === "IPv4" && !net.internal) {
+          return net.address;
+        }
+      }
+    }
+  } catch {
+    /* ignore */
+  }
+  return "127.0.0.1";
+}
+
+/** Optional ngrok TCP tunnel for dev; production should use VM public IP or auto-relay. */
+export async function discoverNgrokTcpAnnounce(petalsPort) {
+  const override = process.env.BLITZWING_ANNOUNCE_MADDRS?.trim();
+  if (override) return override;
+
+  try {
+    const res = await fetch("http://127.0.0.1:4040/api/tunnels", {
+      signal: AbortSignal.timeout(2000),
+    });
+    if (!res.ok) return null;
+    const data = await res.json();
+    const tunnel = (data.tunnels || []).find((t) => t.proto === "tcp");
+    if (!tunnel?.public_url) return null;
+    const url = new URL(tunnel.public_url);
+    return `/dns4/${url.hostname}/tcp/${url.port || petalsPort}`;
+  } catch {
+    return null;
+  }
+}
+
+export function extractPeerMultiaddrFromLog(logPath, fs) {
+  if (!fs.existsSync(logPath)) return null;
+  const text = fs.readFileSync(logPath, "utf8");
+  const line = text
+    .split("\n")
+    .reverse()
+    .find((l) => l.includes("Running a server on"));
+  if (!line) return null;
+  const m = line.match(/Running a server on \['([^']+)'\]/);
+  return m ? m[1] : null;
+}
+
 export function which(cmd) {
   try {
     const out = execFileSync(process.platform === "win32" ? "where" : "which", [cmd], {

package/src/wizard.js

@@ -5,7 +5,12 @@ import path from "node:path";
 import { spawn } from "node:child_process";
 import { DEFAULT_DISCOVERY_URL, HOME_DIR, SHARD_PORT, PETALS_PORT } from "./config.js";
 import { listMothers, motherHosts, joinHost, readyHost, leaveHost } from "./api.js";
-import { detectPublicIp } from "./net.js";
+import {
+  detectPublicIp,
+  detectLocalIp,
+  discoverNgrokTcpAnnounce,
+  extractPeerMultiaddrFromLog,
+} from "./net.js";
 import {
   ensureVenvAndPetals,
   startShardManagerProcess,
@@ -55,6 +60,9 @@ ${color.bold("blitzwing")} — join a Blitzwing mother swarm as a compute node
 
 Env:
   BLITZWING_DISCOVERY_URL   Override discovery service (default ${DEFAULT_DISCOVERY_URL})
+  BLITZWING_ANNOUNCE_MADDRS Explicit Petals announce multiaddr (VM with public IP)
+  BLITZWING_LOCAL_IP        Local IP for shard manager metadata
+  PETALS_USE_AUTO_RELAY=0   Disable libp2p auto-relay (use with public IP / port-forward)
 `);
 }
 
@@ -130,20 +138,86 @@ async function wizard(args) {
   const layersN = Number(layers);
 
   const detected = await detectPublicIp();
-  const publicIp = await p.text({
-    message: "Public IP for this machine (other peers must reach you)",
-    initialValue: detected || "",
+  const networkMode = await p.select({
+    message: "How will other peers reach your Petals node?",
+    options: [
+      {
+        value: "relay",
+        label: "Auto (recommended for home / NAT)",
+        hint: "Uses Petals libp2p relay — no port forwarding",
+      },
+      {
+        value: "public",
+        label: "Public IP / cloud VM",
+        hint: "TCP 31337 (and 8001) open on the internet",
+      },
+      {
+        value: "ngrok",
+        label: "Dev tunnel (ngrok TCP on :4040)",
+        hint: "Only if ngrok is already running",
+      },
+    ],
+    initialValue: "relay",
+  });
+  if (p.isCancel(networkMode)) {
+    p.cancel("Setup cancelled");
+    process.exit(0);
+  }
+
+  let publicIp = "";
+  let announceMaddrs = "";
+  let useAutoRelay = "1";
+
+  if (networkMode === "public") {
+    publicIp = await p.text({
+      message: "Public IPv4 (must be reachable on TCP " + PETALS_PORT + ")",
+      initialValue: detected || "",
+      validate(v) {
+        if (!v || !/^\d+\.\d+\.\d+\.\d+$/.test(String(v).trim())) {
+          return "Enter a valid public IPv4 address";
+        }
+      },
+    });
+    if (p.isCancel(publicIp)) {
+      p.cancel("Setup cancelled");
+      process.exit(0);
+    }
+    publicIp = String(publicIp).trim();
+    announceMaddrs = `/ip4/${publicIp}/tcp/${PETALS_PORT}`;
+    useAutoRelay = "0";
+  } else if (networkMode === "ngrok") {
+    const ngrok = await discoverNgrokTcpAnnounce(PETALS_PORT);
+    if (!ngrok) {
+      p.cancel("No ngrok TCP tunnel on http://127.0.0.1:4040. Start: ngrok tcp " + PETALS_PORT);
+      process.exit(1);
+    }
+    announceMaddrs = ngrok;
+    publicIp = ngrok.match(/dns4\/([^/]+)/)?.[1] || detected || "relay";
+    useAutoRelay = "0";
+    p.log.info(`Using ngrok announce ${announceMaddrs}`);
+  } else {
+    publicIp = detected || "relay";
+    p.log.info("Using Petals libp2p auto-relay (no inbound port forward required).");
+  }
+
+  const localIp = await detectLocalIp();
+
+  const hederaAccount = await p.text({
+    message: "Hedera account ID for layer payouts (e.g. 0.0.123456)",
+    initialValue: process.env.BLITZWING_HEDERA_ACCOUNT_ID || "",
     validate(v) {
-      if (!v || !v.trim()) return "Public IP is required";
+      const s = String(v || "").trim();
+      if (!/^0\.0\.\d+$/.test(s)) return "Enter a Hedera account like 0.0.123456";
     },
   });
-  if (p.isCancel(publicIp)) {
+  if (p.isCancel(hederaAccount)) {
     p.cancel("Setup cancelled");
     process.exit(0);
   }
+  const hederaAccountId = String(hederaAccount).trim();
 
   const confirm = await p.confirm({
-    message: `Join ${selected.model} hosting ${layersN} layers from ${String(publicIp).trim()}?`,
+    message: `Join ${selected.model} hosting ${layersN} layers (${networkMode}) paying to ${hederaAccountId}?`,
     initialValue: true,
   });
   if (p.isCancel(confirm) || !confirm) {
@@ -166,7 +240,8 @@ async function wizard(args) {
   const python = venvPython();
   spin.stop("Environment ready");
 
-  const shardManagerUrl = `http://${String(publicIp).trim()}:${SHARD_PORT}`;
+  const shardManagerUrl = `http://${localIp}:${SHARD_PORT}`;
+  const petalsLogPath = path.join(HOME_DIR, "petals.log");
 
   spin.start("Requesting layer assignment from mother…");
   let assignment;
@@ -176,6 +251,7 @@ async function wizard(args) {
       layers: layersN,
       public_ip: String(publicIp).trim(),
       shard_manager_url: shardManagerUrl,
+      hedera_account_id: hederaAccountId,
     });
   } catch (err) {
     spin.stop("Join rejected");
@@ -187,22 +263,28 @@ async function wizard(args) {
 
   const logPath = path.join(HOME_DIR, "shard_manager.log");
   spin.start("Starting local Petals server…");
-  const { pid } = startShardManagerProcess({
-    python,
-    logPath,
-    env: {
+  const shardEnv = {
       MODEL_NAME: selected.model,
-      PUBLIC_IP: String(publicIp).trim(),
+      PUBLIC_IP: announceMaddrs ? "" : String(publicIp).trim(),
       BLOCK_INDICES: assignment.block_indices,
       INITIAL_PEERS: (assignment.initial_peers || []).join(","),
       PETALS_PORT: String(PETALS_PORT),
       IDENTITY_PATH: path.join(HOME_DIR, "petals-identity"),
       SHARD_AUTO_START: "1",
-    },
+      PETALS_USE_AUTO_RELAY: useAutoRelay,
+      PETALS_SKIP_REACHABILITY_CHECK: "1",
+    };
+  if (announceMaddrs) {
+    shardEnv.ANNOUNCE_MADDRS = announceMaddrs;
+  }
+  const { pid } = startShardManagerProcess({
+    python,
+    logPath,
+    env: shardEnv,
   });
 
   try {
-    await waitForShardRunning({ timeoutMs: 600000 });
+    await waitForShardRunning({ timeoutMs: 600000, statusHost: "127.0.0.1" });
   } catch (err) {
     spin.stop("Petals did not become ready");
     p.cancel(`${err.message}. See ${logPath}`);
@@ -211,8 +293,12 @@ async function wizard(args) {
   spin.stop("Petals is serving your layers");
 
   spin.start("Finalizing handoff with mother…");
+  const peerMultiaddr = extractPeerMultiaddrFromLog(petalsLogPath, fs);
   try {
-    await readyHost(selected.mother_url, { host_id: assignment.host_id });
+    await readyHost(selected.mother_url, {
+      host_id: assignment.host_id,
+      peer_multiaddr: peerMultiaddr || undefined,
+    });
   } catch (err) {
     spin.stop("Handoff failed");
     p.cancel(err.message);
@@ -227,9 +313,12 @@ async function wizard(args) {
     block_indices: assignment.block_indices,
     layers_hosted: assignment.layers_hosted,
     public_ip: String(publicIp).trim(),
+    network_mode: networkMode,
+    announce_maddrs: announceMaddrs || null,
     shard_manager_url: shardManagerUrl,
     shard_pid: pid,
     discovery_url: args.discoveryUrl,
+    hedera_account_id: hederaAccountId,
     joined_at: new Date().toISOString(),
   };
   saveState(state);
@@ -281,7 +370,8 @@ async function showStatus() {
   console.log(`  mother:         ${state.mother_url}`);
   console.log(`  public_ip:      ${state.public_ip}`);
   try {
-    const res = await fetch(`http://127.0.0.1:${SHARD_PORT}/status`);
+    const statusHost = "127.0.0.1";
+    const res = await fetch(`http://${statusHost}:${SHARD_PORT}/status`);
     if (res.ok) {
       const st = await res.json();
       console.log(`  petals_running: ${st.running}`);
@@ -308,7 +398,8 @@ async function doLeave() {
     p.log.warn(`Mother leave call failed: ${err.message}`);
   }
   try {
-    await fetch(`http://127.0.0.1:${SHARD_PORT}/stop`, { method: "POST" });
+    const statusHost = "127.0.0.1";
+    await fetch(`http://${statusHost}:${SHARD_PORT}/stop`, { method: "POST" });
   } catch {
     /* ignore */
   }