blitzwing 0.1.4 → 0.1.8

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "blitzwing",
-  "version": "0.1.4",
+  "version": "0.1.8",
   "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

@@ -33,11 +33,25 @@ function isPetalsCompatible(info) {
   return info && info.major === 3 && info.minor >= 10 && info.minor <= 11;
 }
 
+function pinnedPythonPath() {
+  try {
+    const p = path.join(HOME_DIR, "python");
+    if (!fs.existsSync(p)) return null;
+    const bin = fs.readFileSync(p, "utf8").trim();
+    return bin || null;
+  } catch {
+    return null;
+  }
+}
+
 export function ensurePython() {
   // Petals/hivemind break on 3.12+. Never fall back to a too-new interpreter.
+  const home = process.env.HOME || process.env.USERPROFILE || "";
   const candidates = [];
   const envBin = process.env.BLITZWING_PYTHON || process.env.PETALS_PYTHON;
   if (envBin) candidates.push(envBin);
+  const pinned = pinnedPythonPath();
+  if (pinned) candidates.push(pinned);
   if (process.env.CONDA_PREFIX) {
     candidates.push(
       path.join(process.env.CONDA_PREFIX, "bin", "python3.11"),
@@ -45,43 +59,77 @@ export function ensurePython() {
       path.join(process.env.CONDA_PREFIX, "bin", "python")
     );
   }
-  candidates.push(
-    path.join(process.env.HOME || "", "miniconda", "envs", "blitzwing", "bin", "python3.11"),
-    path.join(process.env.HOME || "", "miniconda", "envs", "blitzwing", "bin", "python"),
-    path.join(process.env.HOME || "", "miniconda3", "envs", "blitzwing", "bin", "python3.11"),
-    path.join(process.env.HOME || "", "miniconda3", "envs", "blitzwing", "bin", "python"),
-    "python3.11",
-    "python3.10"
-  );
+  for (const root of ["miniforge3", "mambaforge", "miniconda", "miniconda3", "anaconda3"]) {
+    candidates.push(
+      path.join(home, root, "envs", "blitzwing", "bin", "python3.11"),
+      path.join(home, root, "envs", "blitzwing", "bin", "python")
+    );
+  }
+  candidates.push("python3.11", "python3.10");
 
+  const tried = [];
   const seen = new Set();
   for (const cmd of candidates) {
     if (!cmd || seen.has(cmd)) continue;
     seen.add(cmd);
     const found = cmd.includes("/") || cmd.includes("\\") ? (fs.existsSync(cmd) ? cmd : null) : which(cmd);
-    if (!found) continue;
+    if (!found) {
+      tried.push(`${cmd} (missing)`);
+      continue;
+    }
     const info = pythonVersion(found);
     if (isPetalsCompatible(info)) {
+      fs.mkdirSync(HOME_DIR, { recursive: true });
+      fs.writeFileSync(path.join(HOME_DIR, "python"), found + "\n");
       return { py: found, ver: info.ver };
     }
+    tried.push(`${found} (${info ? info.ver : "unreadable"})`);
   }
 
   throw new Error(
-    "Python 3.10 or 3.11 is required (3.12+ will not work). " +
-      "Install with: conda create -n blitzwing python=3.11 && conda activate blitzwing  " +
-      "or set BLITZWING_PYTHON=/path/to/python3.11"
+    "Python 3.10 or 3.11 is required (3.12+ will not work).\n" +
+      "Checked:\n  - " +
+      tried.slice(0, 12).join("\n  - ") +
+      "\nFix:\n" +
+      "  curl -LsSf https://astral.sh/uv/install.sh | sh && source $HOME/.local/bin/env\n" +
+      "  uv python install 3.11\n" +
+      "  export BLITZWING_PYTHON=\"$(uv python find 3.11)\"\n" +
+      "  mkdir -p ~/.blitzwing && echo \"$BLITZWING_PYTHON\" > ~/.blitzwing/python\n" +
+      "  blitzwing"
   );
 }
 
 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 });
@@ -99,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",
@@ -129,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 };
 }
 
@@ -176,8 +239,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
 
@@ -186,7 +253,9 @@ 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]
@@ -197,7 +266,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):
@@ -206,7 +283,13 @@ 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)
@@ -217,6 +300,9 @@ class Mgr:
         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,
@@ -266,12 +352,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],
     }
@@ -280,21 +367,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/wizard.js

@@ -142,8 +142,22 @@ async function wizard(args) {
     process.exit(0);
   }
 
+  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) {
+      const s = String(v || "").trim();
+      if (!/^0\.0\.\d+$/.test(s)) return "Enter a Hedera account like 0.0.123456";
+    },
+  });
+  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 from ${String(publicIp).trim()} paying to ${hederaAccountId}?`,
     initialValue: true,
   });
   if (p.isCancel(confirm) || !confirm) {
@@ -176,6 +190,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");
@@ -202,7 +217,7 @@ async function wizard(args) {
   });
 
   try {
-    await waitForShardRunning({ timeoutMs: 600000 });
+    await waitForShardRunning({ timeoutMs: 600000, statusHost: String(publicIp).trim() });
   } catch (err) {
     spin.stop("Petals did not become ready");
     p.cancel(`${err.message}. See ${logPath}`);
@@ -230,6 +245,7 @@ async function wizard(args) {
     shard_manager_url: shardManagerUrl,
     shard_pid: pid,
     discovery_url: args.discoveryUrl,
+    hedera_account_id: hederaAccountId,
     joined_at: new Date().toISOString(),
   };
   saveState(state);
@@ -281,7 +297,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 = state.public_ip || "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 +325,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 = state.public_ip || "127.0.0.1";
+    await fetch(`http://${statusHost}:${SHARD_PORT}/stop`, { method: "POST" });
   } catch {
     /* ignore */
   }