blitzwing 0.2.1 → 0.2.2

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.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",
+    "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")
@@ -364,6 +372,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, :]