blitzwing 0.1.8 → 0.1.10

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "blitzwing",
-  "version": "0.1.8",
+  "version": "0.1.10",
   "description": "Interactive setup wizard to join a Blitzwing Petals mother swarm as a compute contributor",
   "bin": {
     "blitzwing": "bin/blitzwing.js"

package/src/heartbeat.js

@@ -0,0 +1,64 @@
+import { spawn } from "node:child_process";
+import fs from "node:fs";
+import path from "node:path";
+import { heartbeatHost } from "./api.js";
+import { HOME_DIR } from "./config.js";
+
+/** Must be well below mother HEARTBEAT_TTL_SECONDS (default 60). */
+export const DEFAULT_HEARTBEAT_INTERVAL_MS = 20_000;
+
+/**
+ * Start a detached background process that heartbeats the mother orchestrator.
+ * @param {{ motherUrl: string, hostId: string, intervalMs?: number }} opts
+ */
+export function startContributorHeartbeatDaemon({
+  motherUrl,
+  hostId,
+  intervalMs = DEFAULT_HEARTBEAT_INTERVAL_MS,
+}) {
+  const hbPath = path.join(HOME_DIR, "heartbeat.mjs");
+  fs.mkdirSync(HOME_DIR, { recursive: true });
+  fs.writeFileSync(
+    hbPath,
+    `const mother = ${JSON.stringify(motherUrl)};
+const hostId = ${JSON.stringify(hostId)};
+const intervalMs = ${intervalMs};
+async function beat() {
+  try {
+    await fetch(mother.replace(/\\/$/, "") + "/v1/hosts/heartbeat", {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        "ngrok-skip-browser-warning": "true",
+      },
+      body: JSON.stringify({ host_id: hostId }),
+    });
+  } catch {}
+}
+setInterval(beat, intervalMs);
+beat();
+`
+  );
+  const child = spawn(process.execPath, [hbPath], {
+    detached: true,
+    stdio: "ignore",
+  });
+  child.unref();
+}
+
+/** In-process heartbeat loop (for tests or long-running CLI). */
+export async function startContributorHeartbeatLoop({
+  motherUrl,
+  hostId,
+  intervalMs = DEFAULT_HEARTBEAT_INTERVAL_MS,
+}) {
+  async function beat() {
+    try {
+      await heartbeatHost(motherUrl, { host_id: hostId });
+    } catch {
+      /* mother may be briefly unavailable */
+    }
+  }
+  setInterval(beat, intervalMs);
+  await beat();
+}

package/src/install.js

@@ -220,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
@@ -259,6 +260,13 @@ class Mgr:
             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):
@@ -296,7 +304,9 @@ class Mgr:
             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):
@@ -333,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()
 

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,
@@ -13,6 +18,7 @@ import {
   venvPython,
 } from "./install.js";
 import { saveState, loadState, clearState, ensureHome } from "./state.js";
+import { startContributorHeartbeatDaemon } from "./heartbeat.js";
 
 function parseArgs(argv) {
   const out = { cmd: null, discoveryUrl: process.env.BLITZWING_DISCOVERY_URL || DEFAULT_DISCOVERY_URL };
@@ -55,6 +61,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,18 +139,70 @@ 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 || "",
-    validate(v) {
-      if (!v || !v.trim()) return "Public IP is required";
-    },
+  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(publicIp)) {
+  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 || "",
@@ -157,7 +218,7 @@ async function wizard(args) {
   const hederaAccountId = String(hederaAccount).trim();
 
   const confirm = await p.confirm({
-    message: `Join ${selected.model} hosting ${layersN} layers from ${String(publicIp).trim()} paying to ${hederaAccountId}?`,
+    message: `Join ${selected.model} hosting ${layersN} layers (${networkMode}) paying to ${hederaAccountId}?`,
     initialValue: true,
   });
   if (p.isCancel(confirm) || !confirm) {
@@ -180,7 +241,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;
@@ -202,22 +264,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, statusHost: String(publicIp).trim() });
+    await waitForShardRunning({ timeoutMs: 600000, statusHost: "127.0.0.1" });
   } catch (err) {
     spin.stop("Petals did not become ready");
     p.cancel(`${err.message}. See ${logPath}`);
@@ -226,8 +294,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);
@@ -242,6 +314,8 @@ 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,
@@ -249,7 +323,10 @@ async function wizard(args) {
     joined_at: new Date().toISOString(),
   };
   saveState(state);
-  startHeartbeatDaemon(state);
+  startContributorHeartbeatDaemon({
+    motherUrl: state.mother_url,
+    hostId: state.host_id,
+  });
 
   p.outro(
     `${color.green("You are online.")} Hosting ${color.cyan(String(state.layers_hosted))} layers of ${color.cyan(state.model)} at ${state.block_indices}\n` +
@@ -257,32 +334,6 @@ async function wizard(args) {
   );
 }
 
-function startHeartbeatDaemon(state) {
-  const hbPath = path.join(HOME_DIR, "heartbeat.mjs");
-  fs.writeFileSync(
-    hbPath,
-    `const mother = ${JSON.stringify(state.mother_url)};
-const hostId = ${JSON.stringify(state.host_id)};
-async function beat() {
-  try {
-    await fetch(mother.replace(/\\/$/, "") + "/v1/hosts/heartbeat", {
-      method: "POST",
-      headers: { "Content-Type": "application/json" },
-      body: JSON.stringify({ host_id: hostId }),
-    });
-  } catch {}
-}
-setInterval(beat, 60000);
-beat();
-`
-  );
-  const child = spawn(process.execPath, [hbPath], {
-    detached: true,
-    stdio: "ignore",
-  });
-  child.unref();
-}
-
 async function showStatus() {
   const state = loadState();
   if (!state) {
@@ -297,7 +348,7 @@ async function showStatus() {
   console.log(`  mother:         ${state.mother_url}`);
   console.log(`  public_ip:      ${state.public_ip}`);
   try {
-    const statusHost = state.public_ip || "127.0.0.1";
+    const statusHost = "127.0.0.1";
     const res = await fetch(`http://${statusHost}:${SHARD_PORT}/status`);
     if (res.ok) {
       const st = await res.json();
@@ -325,7 +376,7 @@ async function doLeave() {
     p.log.warn(`Mother leave call failed: ${err.message}`);
   }
   try {
-    const statusHost = state.public_ip || "127.0.0.1";
+    const statusHost = "127.0.0.1";
     await fetch(`http://${statusHost}:${SHARD_PORT}/stop`, { method: "POST" });
   } catch {
     /* ignore */