blitzwing 0.2.1 → 0.2.3

package/package.json

@@ -1,32 +1,32 @@
-{
-  "name": "blitzwing",
-  "version": "0.2.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",
-    "runtime",
-    "README.md"
-  ],
-  "keywords": [
-    "blitzwing",
-    "petals",
-    "llm",
-    "distributed-inference"
-  ],
-  "license": "MIT",
-  "config": {
-    "discoveryUrl": "http://35.226.124.189:9000"
-  },
-  "dependencies": {
-    "@clack/prompts": "^0.9.1",
-    "picocolors": "^1.1.1"
-  }
-}
+{
+  "name": "blitzwing",
+  "version": "0.2.3",
+  "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",
+    "runtime",
+    "README.md"
+  ],
+  "keywords": [
+    "blitzwing",
+    "petals",
+    "llm",
+    "distributed-inference"
+  ],
+  "license": "MIT",
+  "config": {
+    "discoveryUrl": "http://35.226.124.189:9000"
+  },
+  "dependencies": {
+    "@clack/prompts": "^0.9.1",
+    "picocolors": "^1.1.1"
+  }
+}

package/runtime/app.py

@@ -18,7 +18,7 @@ from pydantic import BaseModel, Field
 
 from http_chain import HttpChainInference
 from local_runner import LocalShardRunner
-from tensor_codec import tensor_to_payload
+from tensor_codec import tensor_from_payload, tensor_to_payload
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -65,6 +65,14 @@ class PrefixResponse(BaseModel):
     hidden: dict
 
 
+class ContinueRequest(BaseModel):
+    hidden: dict
+
+
+class ContinueResponse(BaseModel):
+    hidden: dict
+
+
 class ShardManager:
     def __init__(self) -> None:
         self.model = os.getenv("MODEL_NAME", "bigscience/bloom-560m")
@@ -293,6 +301,7 @@ def _prewarm_local_runner() -> None:
 
 @app.on_event("startup")
 def on_startup() -> None:
+    manager.last_exit_code = None
     if manager._auto_start:
         try:
             manager.start(bootstrap=manager.new_swarm)
@@ -364,6 +373,26 @@ async def chain_prefix(body: PrefixRequest) -> PrefixResponse:
         raise HTTPException(status_code=503, detail=str(exc)) from exc
 
 
+@app.post("/v1/chain/continue", response_model=ContinueResponse)
+async def chain_continue(body: ContinueRequest) -> ContinueResponse:
+    st = manager.status()
+    if not st.running:
+        raise HTTPException(status_code=503, detail="Petals server not running")
+
+    def _run() -> ContinueResponse:
+        runner = _get_local_runner()
+        hidden = tensor_from_payload(body.hidden)
+        hidden = runner.forward_tail(hidden)
+        return ContinueResponse(hidden=tensor_to_payload(hidden))
+
+    try:
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(_inference_executor, _run)
+    except Exception as exc:  # noqa: BLE001
+        logger.exception("HTTP chain continue failed")
+        raise HTTPException(status_code=503, detail=str(exc)) from exc
+
+
 @app.post("/v1/chat/completions", response_model=ChatInferenceResponse)
 async def chat_completions(body: ChatInferenceRequest) -> ChatInferenceResponse:
     st = manager.status()

package/runtime/http_chain.py

@@ -1,11 +1,11 @@
-"""HTTP-chained distributed inference — no libp2p between nodes."""
+"""HTTP-chained distributed inference — every online host runs its layer slice."""
 
 from __future__ import annotations
 
 import logging
 import os
 from dataclasses import dataclass
-from typing import List, Optional, Sequence
+from typing import List, Optional, Sequence, Tuple
 
 import httpx
 import torch
@@ -68,11 +68,39 @@ def resolve_mother_shard_url(manifest: dict) -> str:
     )
 
 
-def mother_prefix_blocks(manifest: dict) -> str:
-    for host in manifest.get("hosts", []):
-        if host.get("role") == "mother" and host.get("block_indices"):
-            return str(host["block_indices"])
-    return "0:18"
+def ordered_hosts(manifest: dict) -> List[dict]:
+    hosts = [h for h in manifest.get("hosts", []) if h.get("block_indices")]
+    hosts.sort(key=lambda h: parse_range(str(h["block_indices"]))[0])
+    return hosts
+
+
+def validate_chain_coverage(hosts: List[dict], total_layers: int) -> None:
+    if not hosts:
+        raise RuntimeError("swarm_manifest has no hosts")
+    expected = 0
+    for host in hosts:
+        start, end = parse_range(str(host["block_indices"]))
+        if start != expected:
+            raise RuntimeError(
+                f"layer gap before {host.get('host_id')}: expected start={expected}, got {start}"
+            )
+        expected = end
+    if total_layers and expected != int(total_layers):
+        raise RuntimeError(
+            f"layer map incomplete: covered 0:{expected}, need 0:{total_layers}"
+        )
+
+
+def host_http_url(host: dict, manifest: dict) -> str:
+    if host.get("role") == "mother":
+        return resolve_mother_shard_url(manifest)
+    url = (host.get("shard_manager_url") or "").rstrip("/")
+    if not _is_public_http_url(url):
+        raise RuntimeError(
+            f"host {host.get('host_id')} has no public shard_manager_url "
+            f"(got {url or 'empty'}) — cannot include it in HTTP chain"
+        )
+    return url
 
 
 @dataclass
@@ -80,6 +108,22 @@ class HttpChainInference:
     runner: LocalShardRunner
     local_block_indices: str
 
+    def _remote_prefix(self, client: httpx.Client, url: str, input_ids: List[int]) -> torch.Tensor:
+        resp = client.post(
+            f"{url.rstrip('/')}/v1/chain/prefix",
+            json={"input_ids": input_ids},
+        )
+        resp.raise_for_status()
+        return tensor_from_payload(resp.json()["hidden"])
+
+    def _remote_continue(self, client: httpx.Client, url: str, hidden: torch.Tensor) -> torch.Tensor:
+        resp = client.post(
+            f"{url.rstrip('/')}/v1/chain/continue",
+            json={"hidden": tensor_to_payload(hidden)},
+        )
+        resp.raise_for_status()
+        return tensor_from_payload(resp.json()["hidden"])
+
     def generate(
         self,
         messages: Sequence[dict],
@@ -92,19 +136,22 @@ class HttpChainInference:
         if not swarm_manifest:
             raise RuntimeError("swarm_manifest required for HTTP-chained inference")
 
-        mother_url = resolve_mother_shard_url(swarm_manifest)
-        mother_blocks = mother_prefix_blocks(swarm_manifest)
-        _, mother_end = parse_range(mother_blocks)
-        local_start, _local_end = parse_range(self.local_block_indices)
-        prefix_timeout = float(os.getenv("HTTP_CHAIN_PREFIX_TIMEOUT", "120"))
-
-        if local_start != mother_end:
-            logger.warning(
-                "Layer gap between mother end=%s and local start=%s",
-                mother_end,
-                local_start,
+        hosts = ordered_hosts(swarm_manifest)
+        total_layers = int(swarm_manifest.get("total_layers") or 0)
+        validate_chain_coverage(hosts, total_layers)
+
+        local_start, local_end = parse_range(self.local_block_indices)
+        tail = hosts[-1]
+        tail_start, tail_end = parse_range(str(tail["block_indices"]))
+        if (tail_start, tail_end) != (local_start, local_end):
+            raise RuntimeError(
+                f"local blocks {self.local_block_indices} are not the tail "
+                f"(tail is {tail.get('host_id')} {tail['block_indices']})"
             )
 
+        upstream = hosts[:-1]
+        hop_timeout = float(os.getenv("HTTP_CHAIN_PREFIX_TIMEOUT", "120"))
+
         tokenizer = self.runner.tokenizer
         try:
             prompt = tokenizer.apply_chat_template(
@@ -124,17 +171,19 @@ class HttpChainInference:
         input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"]
         prompt_tokens = int(input_ids.shape[-1])
         max_new_tokens = max(1, min(int(max_tokens), 512))
+        hop_desc = " -> ".join(
+            f"{h.get('host_id')}[{h.get('block_indices')}]" for h in hosts
+        )
         logger.info(
-            "HTTP chain: mother=%s blocks=%s local=%s max_tokens=%s",
-            mother_url,
-            mother_blocks,
+            "HTTP chain hops=%s local=%s max_tokens=%s",
+            hop_desc,
             self.local_block_indices,
             max_new_tokens,
         )
         generated: List[int] = []
         finish_reason = "stop"
 
-        with httpx.Client(timeout=prefix_timeout) as client:
+        with httpx.Client(timeout=hop_timeout) as client:
             for step in range(max_new_tokens):
                 full_ids = (
                     torch.cat(
@@ -144,16 +193,32 @@ class HttpChainInference:
                     if generated
                     else input_ids
                 )
+                token_ids = full_ids[0].tolist()
 
-                prefix_resp = client.post(
-                    f"{mother_url.rstrip('/')}/v1/chain/prefix",
-                    json={"input_ids": full_ids[0].tolist()},
-                )
-                prefix_resp.raise_for_status()
-                hidden = tensor_from_payload(prefix_resp.json()["hidden"])
-                logger.debug("HTTP chain step %s: mother prefix ok", step)
+                if not upstream:
+                    # Mother-only / single-node: run everything locally.
+                    hidden = self.runner.forward_prefix(
+                        torch.tensor([token_ids], dtype=torch.long)
+                    )
+                else:
+                    first = upstream[0]
+                    first_url = host_http_url(first, swarm_manifest)
+                    hidden = self._remote_prefix(client, first_url, token_ids)
+                    logger.debug(
+                        "HTTP chain step %s: prefix via %s ok",
+                        step,
+                        first.get("host_id"),
+                    )
+                    for mid in upstream[1:]:
+                        mid_url = host_http_url(mid, swarm_manifest)
+                        hidden = self._remote_continue(client, mid_url, hidden)
+                        logger.debug(
+                            "HTTP chain step %s: continue via %s ok",
+                            step,
+                            mid.get("host_id"),
+                        )
+                    hidden = self.runner.forward_tail(hidden)
 
-                hidden = self.runner.forward_tail(hidden)
                 logits = self.runner.logits_from_hidden(hidden)
                 next_logits = logits[0, -1, :]
 

package/src/install.js

@@ -10,6 +10,72 @@ const PACKAGE_ROOT = path.join(__dirname, "..");
 const RUNTIME_SRC = path.join(PACKAGE_ROOT, "runtime");
 const RUNTIME_DEST = path.join(HOME_DIR, "runtime");
 
+function sleep(ms) {
+  return new Promise((r) => setTimeout(r, ms));
+}
+
+/**
+ * Kill leftover contributor shard/Petals processes so a new join does not talk to a
+ * stale uvicorn that still reports last_exit_code=-15 (SIGTERM from a prior leave).
+ */
+export function stopLocalContributorStack(extraPids = []) {
+  for (const pid of extraPids) {
+    if (!pid || !Number.isFinite(Number(pid))) continue;
+    try {
+      process.kill(Number(pid), "SIGTERM");
+    } catch {
+      /* already gone */
+    }
+  }
+
+  if (process.platform === "win32") {
+    for (const port of [SHARD_PORT, PETALS_PORT]) {
+      try {
+        const out = execFileSync("netstat", ["-ano"], { encoding: "utf8" });
+        const re = new RegExp(`:${port}\\s+.*LISTENING\\s+(\\d+)`, "i");
+        const m = out.match(re);
+        if (m) {
+          execFileSync("taskkill", ["/PID", m[1], "/T", "/F"], { stdio: "ignore" });
+        }
+      } catch {
+        /* ignore */
+      }
+    }
+    return;
+  }
+
+  try {
+    execFileSync(
+      "bash",
+      [
+        "-lc",
+        [
+          `fuser -k ${SHARD_PORT}/tcp ${PETALS_PORT}/tcp 2>/dev/null || true`,
+          `pkill -f 'uvicorn app:app --host 0.0.0.0 --port ${SHARD_PORT}' 2>/dev/null || true`,
+          `pkill -f 'petals.cli.run_server' 2>/dev/null || true`,
+          "sleep 1",
+        ].join("; "),
+      ],
+      { stdio: "ignore" }
+    );
+  } catch {
+    /* ignore */
+  }
+}
+
+function petalsLogShowsSessionStart(logPath) {
+  try {
+    if (!fs.existsSync(logPath)) return false;
+    const text = fs.readFileSync(logPath, "utf8");
+    return (
+      /Starting Petals:/.test(text) ||
+      /\[INFO\] Started\b/.test(text) ||
+      /Running a server on/.test(text)
+    );
+  } catch {
+    return false;
+  }
+}
 function venvPython() {
   if (process.platform === "win32") {
     return path.join(VENV_PATH, "Scripts", "python.exe");
@@ -226,7 +292,16 @@ export function writeLocalShardManager() {
 }
 
 export function startShardManagerProcess({ python, env, logPath }) {
+  // Always clear stale listeners before binding — leave() used to stop Petals only.
+  stopLocalContributorStack();
   const appPath = syncRuntimeFiles();
+  fs.mkdirSync(HOME_DIR, { recursive: true });
+  // Fresh logs so wait/error tails are from this session.
+  try {
+    fs.writeFileSync(path.join(HOME_DIR, "petals.log"), "");
+  } catch {
+    /* ignore */
+  }
   const out = fs.openSync(logPath, "a");
   const petalsPy = petalsImportable(venvPython()) ? venvPython() : python;
   const child = spawn(
@@ -253,28 +328,38 @@ 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();
+  let sawRunning = false;
+
   while (Date.now() - start < timeoutMs) {
     try {
       const res = await fetch(`http://${host}:${SHARD_PORT}/status`);
       if (res.ok) {
         const body = await res.json();
-        if (body.running && body.pid) return body;
+        if (body.running && body.pid) {
+          sawRunning = true;
+          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}`
-          );
+          // A prior leave/stop leaves uvicorn up with last_exit_code=-15. Ignore that
+          // until this session's Petals has actually started (or we already saw running).
+          const sessionStarted = petalsLogShowsSessionStart(petalsLog);
+          if (sawRunning || sessionStarted) {
+            const tail = fs.existsSync(petalsLog)
+              ? fs.readFileSync(petalsLog, "utf8").trim().split("\n").slice(-8).join("\n")
+              : "";
+            throw new Error(
+              `Petals exited (code ${body.last_exit_code}). ${tail || "See " + petalsLog}`
+            );
+          }
         }
       }
     } catch (err) {
       if (err instanceof Error && err.message.startsWith("Petals exited")) throw err;
     }
-    await new Promise((r) => setTimeout(r, 2000));
+    await sleep(2000);
   }
   const tail = fs.existsSync(petalsLog)
-    ? fs.readFileSync(petalsLog, "utf8").trim().split("\n").slice(-5).join("\n")
+    ? fs.readFileSync(petalsLog, "utf8").trim().split("\n").slice(-8).join("\n")
     : "";
   throw new Error(
     `Timed out waiting for Petals on http://${host}:${SHARD_PORT}/status. ${tail || "See " + petalsLog}`

package/src/wizard.js

@@ -8,6 +8,7 @@ import {
   ensureVenvAndPetals,
   startShardManagerProcess,
   waitForShardRunning,
+  stopLocalContributorStack,
   syncRuntimeFiles,
   venvPython,
 } from "./install.js";
@@ -370,6 +371,7 @@ async function doLeave() {
   } catch {
     /* ignore */
   }
+  stopLocalContributorStack([state.shard_pid, state.tunnel_pid].filter(Boolean));
   if (state.tunnel_pid) {
     stopTunnelProcess(state.tunnel_pid);
   }