blitzwing 0.1.9 → 0.2.0

package/README.md

@@ -1,6 +1,6 @@
 # blitzwing
 
-Interactive setup wizard to contribute layers to a Blitzwing Petals mother swarm.
+Join a Blitzwing mother swarm as a compute contributor. One command — no repo clone, no ngrok, no port forwarding.
 
 ## Install
 
@@ -8,22 +8,31 @@ Interactive setup wizard to contribute layers to a Blitzwing Petals mother swarm
 npm i -g blitzwing
 ```
 
+**Windows:** run inside **WSL** (Petals needs Linux). Install Node in WSL, then `npm i -g blitzwing`.
+
 ## Use
 
 ```bash
 blitzwing
 ```
 
-The wizard talks to the **Discovery Service** (no mother URL to type), asks which model and how many layers to host, installs Petals if needed, and joins the swarm.
+The wizard:
+
+1. Discovers mothers (no mother URL to type)
+2. Asks which model and how many layers
+3. Asks for your **Hedera account ID** (`0.0.N`) for payouts
+4. Installs Petals if needed
+5. Opens a free **Cloudflare Quick Tunnel** (no account / token)
+6. Joins the swarm and starts heartbeats
 
 ```bash
 blitzwing status
 blitzwing leave
 ```
 
-Override discovery URL (dev/staging only):
+## Env (optional)
 
 ```bash
-export BLITZWING_DISCOVERY_URL=http://127.0.0.1:9000
-blitzwing
+export BLITZWING_DISCOVERY_URL=http://127.0.0.1:9000   # override discovery
+export BLITZWING_HEDERA_ACCOUNT_ID=0.0.123456          # prefill payout account
 ```

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "blitzwing",
-  "version": "0.1.9",
+  "version": "0.2.0",
   "description": "Interactive setup wizard to join a Blitzwing Petals mother swarm as a compute contributor",
   "bin": {
     "blitzwing": "bin/blitzwing.js"
@@ -12,6 +12,7 @@
   "files": [
     "bin",
     "src",
+    "runtime",
     "README.md"
   ],
   "keywords": [

package/runtime/app.py

@@ -0,0 +1,398 @@
+"""Contributor shard manager — Petals supervisor + HTTP-chain inference endpoints."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import signal
+import subprocess
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import List, Optional
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+
+from http_chain import HttpChainInference
+from local_runner import LocalShardRunner
+from tensor_codec import tensor_to_payload
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class ReloadRequest(BaseModel):
+    block_indices: str = Field(..., pattern=r"^\d+:\d+$")
+    initial_peers: Optional[List[str]] = None
+
+
+class StatusResponse(BaseModel):
+    running: bool
+    pid: Optional[int] = None
+    block_indices: Optional[str] = None
+    model: str
+    public_ip: Optional[str] = None
+    port: int
+    last_exit_code: Optional[int] = None
+    identity_path: str
+    initial_peers: List[str] = []
+    new_swarm: bool = False
+
+
+class ChatInferenceRequest(BaseModel):
+    messages: List[dict]
+    max_tokens: int = 64
+    temperature: float = 0.7
+    top_p: float = 0.9
+    swarm_manifest: Optional[dict] = None
+
+
+class ChatInferenceResponse(BaseModel):
+    text: str
+    prompt_tokens: int
+    completion_tokens: int
+    finish_reason: str
+
+
+class PrefixRequest(BaseModel):
+    input_ids: List[int]
+
+
+class PrefixResponse(BaseModel):
+    hidden: dict
+
+
+class ShardManager:
+    def __init__(self) -> None:
+        self.model = os.getenv("MODEL_NAME", "bigscience/bloom-560m")
+        self.public_ip = os.getenv("PUBLIC_IP")
+        self.port = int(os.getenv("PETALS_PORT", "31337"))
+        self.identity_path = os.getenv(
+            "IDENTITY_PATH", str(Path.home() / ".blitzwing" / "petals-identity")
+        )
+        self.device = os.getenv("PETALS_DEVICE", "cpu")
+        self.quant_type = os.getenv("PETALS_QUANT_TYPE", "none")
+        self.num_handlers = int(os.getenv("PETALS_NUM_HANDLERS", "1"))
+        self.python = os.getenv("PETALS_PYTHON", "python")
+        peers_raw = os.getenv("INITIAL_PEERS", "")
+        self.initial_peers = [p.strip() for p in peers_raw.split(",") if p.strip()]
+        announce_raw = os.getenv("ANNOUNCE_MADDRS", "")
+        self.announce_maddrs = [a.strip() for a in announce_raw.split(",") if a.strip()]
+        self.new_swarm = os.getenv("NEW_SWARM", "0") in ("1", "true", "True")
+        self.use_auto_relay = self._resolve_use_auto_relay()
+        self.skip_reachability_check = os.getenv(
+            "PETALS_SKIP_REACHABILITY_CHECK", "1"
+        ) in ("1", "true", "True")
+        self.block_indices = os.getenv("BLOCK_INDICES", "0:22")
+        self._proc: Optional[subprocess.Popen] = None
+        self._lock = threading.Lock()
+        self.last_exit_code: Optional[int] = None
+        self._auto_start = os.getenv("SHARD_AUTO_START", "1") not in ("0", "false", "False")
+        self._bootstrapped = False
+        self._log_fp = None
+        self.petals_log = os.getenv(
+            "PETALS_LOG", str(Path.home() / ".blitzwing" / "petals.log")
+        )
+
+    def _resolve_use_auto_relay(self) -> bool:
+        raw = os.getenv("PETALS_USE_AUTO_RELAY")
+        if raw is not None:
+            return raw not in ("0", "false", "False", "no")
+        if self.announce_maddrs:
+            return False
+        return not self.new_swarm
+
+    def build_cmd(self, block_indices: str, *, bootstrap: bool = False) -> List[str]:
+        cmd = [
+            self.python,
+            "-m",
+            "petals.cli.run_server",
+            self.model,
+            "--device",
+            self.device,
+            "--quant_type",
+            self.quant_type,
+            "--block_indices",
+            block_indices,
+            "--port",
+            str(self.port),
+            "--identity_path",
+            self.identity_path,
+            "--num_handlers",
+            str(self.num_handlers),
+        ]
+        if self.announce_maddrs:
+            cmd.append("--announce_maddrs")
+            cmd.extend(self.announce_maddrs)
+        elif self.public_ip:
+            cmd.extend(["--public_ip", self.public_ip])
+        if bootstrap and self.new_swarm:
+            cmd.append("--new_swarm")
+        elif self.initial_peers and not os.getenv("BLITZWING_HTTP_ONLY"):
+            cmd.append("--initial_peers")
+            cmd.extend(self.initial_peers)
+        if not self.use_auto_relay:
+            cmd.append("--no_auto_relay")
+        if self.skip_reachability_check:
+            cmd.append("--skip_reachability_check")
+        return cmd
+
+    def start(self, block_indices: Optional[str] = None, *, bootstrap: bool = False) -> None:
+        with self._lock:
+            if block_indices:
+                self.block_indices = block_indices
+            if self._proc and self._proc.poll() is None:
+                raise RuntimeError("Petals server already running; call /reload instead")
+            use_bootstrap = bootstrap or (self.new_swarm and not self._bootstrapped)
+            cmd = self.build_cmd(self.block_indices, bootstrap=use_bootstrap)
+            logger.info("Starting Petals: %s", " ".join(cmd))
+            if self._log_fp is None:
+                Path(self.petals_log).parent.mkdir(parents=True, exist_ok=True)
+                self._log_fp = open(self.petals_log, "a", encoding="utf-8")
+            popen_kwargs: dict = {
+                "stdout": self._log_fp,
+                "stderr": subprocess.STDOUT,
+            }
+            if os.name != "nt":
+                popen_kwargs["preexec_fn"] = os.setsid
+            self._proc = subprocess.Popen(cmd, **popen_kwargs)
+            self.last_exit_code = None
+            self._bootstrapped = True
+
+    def stop(self, timeout: float = 30.0) -> None:
+        with self._lock:
+            if not self._proc or self._proc.poll() is not None:
+                self._proc = None
+                return
+            proc = self._proc
+            logger.info("Stopping Petals pid=%s", proc.pid)
+            try:
+                if os.name == "nt":
+                    proc.terminate()
+                else:
+                    try:
+                        os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
+                    except ProcessLookupError:
+                        proc.send_signal(signal.SIGTERM)
+                try:
+                    proc.wait(timeout=timeout)
+                except subprocess.TimeoutExpired:
+                    if os.name != "nt":
+                        try:
+                            os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
+                        except ProcessLookupError:
+                            proc.kill()
+                    else:
+                        proc.kill()
+                    proc.wait(timeout=10)
+            finally:
+                self.last_exit_code = proc.returncode
+                self._proc = None
+
+    def reload(self, block_indices: str, initial_peers: Optional[List[str]] = None) -> None:
+        if initial_peers is not None:
+            self.initial_peers = initial_peers
+        rebootstrap = self.new_swarm
+        logger.info(
+            "Reloading Petals with block_indices=%s initial_peers=%s rebootstrap=%s",
+            block_indices,
+            self.initial_peers,
+            rebootstrap,
+        )
+        self.stop()
+        time.sleep(1.0)
+        if rebootstrap:
+            self._bootstrapped = False
+        self.start(block_indices, bootstrap=rebootstrap)
+
+    def status(self) -> StatusResponse:
+        running = self._proc is not None and self._proc.poll() is None
+        pid = self._proc.pid if running and self._proc else None
+        if self._proc and not running:
+            self.last_exit_code = self._proc.returncode
+        return StatusResponse(
+            running=running,
+            pid=pid,
+            block_indices=self.block_indices,
+            model=self.model,
+            public_ip=self.public_ip,
+            port=self.port,
+            last_exit_code=self.last_exit_code,
+            identity_path=self.identity_path,
+            initial_peers=list(self.initial_peers),
+            new_swarm=self.new_swarm,
+        )
+
+
+manager = ShardManager()
+app = FastAPI(title="Blitzwing Contributor Shard Manager", version="0.2.0")
+
+_local_runner: Optional[LocalShardRunner] = None
+_http_chain: Optional[HttpChainInference] = None
+_inference_executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="shard-infer")
+
+
+def _petals_log_path() -> str:
+    return manager.petals_log
+
+
+def _reset_inference_caches() -> None:
+    global _local_runner, _http_chain
+    if _local_runner is not None or _http_chain is not None:
+        logger.info("Resetting inference caches (blocks=%s)", manager.block_indices)
+    _local_runner = None
+    _http_chain = None
+
+
+def _get_local_runner() -> LocalShardRunner:
+    global _local_runner
+    blocks = manager.block_indices
+    if _local_runner is not None:
+        current = f"{_local_runner.block_start}:{_local_runner.block_end}"
+        if current != blocks:
+            logger.warning(
+                "LocalShardRunner stale (%s != %s); recreating",
+                current,
+                blocks,
+            )
+            _reset_inference_caches()
+    if _local_runner is None:
+        _local_runner = LocalShardRunner(
+            manager.model,
+            blocks,
+            local_port=manager.port,
+            log_path=_petals_log_path(),
+        )
+    return _local_runner
+
+
+def _get_http_chain() -> HttpChainInference:
+    global _http_chain
+    if _http_chain is None:
+        _http_chain = HttpChainInference(
+            runner=_get_local_runner(),
+            local_block_indices=manager.block_indices,
+        )
+    return _http_chain
+
+
+def _prewarm_local_runner() -> None:
+    for _ in range(120):
+        if manager.status().running:
+            break
+        time.sleep(2)
+    try:
+        _get_local_runner()
+        logger.info("LocalShardRunner pre-warmed for HTTP chain")
+    except Exception:  # noqa: BLE001
+        logger.exception("LocalShardRunner pre-warm failed")
+
+
+@app.on_event("startup")
+def on_startup() -> None:
+    if manager._auto_start:
+        try:
+            manager.start(bootstrap=manager.new_swarm)
+        except Exception:  # noqa: BLE001
+            logger.exception("Failed to auto-start Petals server")
+    threading.Thread(
+        target=_prewarm_local_runner, name="prefix-prewarm", daemon=True
+    ).start()
+
+
+@app.on_event("shutdown")
+def on_shutdown() -> None:
+    manager.stop()
+
+
+@app.get("/health")
+def health() -> dict:
+    st = manager.status()
+    return {"status": "ok" if st.running else "degraded", "running": st.running}
+
+
+@app.get("/status", response_model=StatusResponse)
+def status() -> StatusResponse:
+    return manager.status()
+
+
+@app.post("/reload", response_model=StatusResponse)
+def reload(body: ReloadRequest) -> StatusResponse:
+    start, end = map(int, body.block_indices.split(":"))
+    if end <= start:
+        raise HTTPException(status_code=400, detail="block_indices end must be > start")
+    try:
+        manager.reload(body.block_indices, initial_peers=body.initial_peers)
+    except Exception as exc:  # noqa: BLE001
+        raise HTTPException(status_code=500, detail=str(exc)) from exc
+    _reset_inference_caches()
+    threading.Thread(
+        target=_prewarm_local_runner, name="prefix-prewarm", daemon=True
+    ).start()
+    time.sleep(0.5)
+    return manager.status()
+
+
+@app.post("/stop", response_model=StatusResponse)
+def stop() -> StatusResponse:
+    manager.stop()
+    return manager.status()
+
+
+@app.post("/v1/chain/prefix", response_model=PrefixResponse)
+async def chain_prefix(body: PrefixRequest) -> PrefixResponse:
+    st = manager.status()
+    if not st.running:
+        raise HTTPException(status_code=503, detail="Petals server not running")
+
+    def _run() -> PrefixResponse:
+        import torch
+
+        runner = _get_local_runner()
+        input_ids = torch.tensor([body.input_ids], dtype=torch.long)
+        hidden = runner.forward_prefix(input_ids)
+        return PrefixResponse(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 prefix forward 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()
+    if not st.running:
+        raise HTTPException(status_code=503, detail="Petals server not running")
+    if not body.swarm_manifest:
+        raise HTTPException(
+            status_code=400,
+            detail="swarm_manifest required for distributed HTTP inference",
+        )
+
+    def _run() -> ChatInferenceResponse:
+        result = _get_http_chain().generate(
+            body.messages,
+            max_tokens=body.max_tokens,
+            temperature=body.temperature,
+            top_p=body.top_p,
+            swarm_manifest=body.swarm_manifest,
+        )
+        return ChatInferenceResponse(
+            text=result.text,
+            prompt_tokens=result.prompt_tokens,
+            completion_tokens=result.completion_tokens,
+            finish_reason=result.finish_reason,
+        )
+
+    try:
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(_inference_executor, _run)
+    except Exception as exc:  # noqa: BLE001
+        logger.exception("Contributor HTTP-chained inference failed")
+        raise HTTPException(status_code=503, detail=str(exc)) from exc

package/runtime/http_chain.py

@@ -0,0 +1,189 @@
+"""HTTP-chained distributed inference — no libp2p between nodes."""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass
+from typing import List, Optional, Sequence
+
+import httpx
+import torch
+
+from inference import InferenceResult
+from local_runner import LocalShardRunner, parse_range
+from tensor_codec import tensor_from_payload, tensor_to_payload
+
+logger = logging.getLogger(__name__)
+
+
+def _is_public_http_url(url: str) -> bool:
+    if not url:
+        return False
+    lower = url.lower()
+    if lower.startswith("http://127.") or lower.startswith("http://localhost"):
+        return False
+    if url.startswith("http://172.") or url.startswith("http://192.168."):
+        return False
+    if url.startswith("http://10."):
+        return False
+    return url.startswith("http://") or url.startswith("https://")
+
+
+def _is_reachable_mother_url(url: str) -> bool:
+    if not _is_public_http_url(url):
+        return False
+    if url.rstrip("/").endswith(":8001"):
+        return False
+    return True
+
+
+def resolve_mother_shard_url(manifest: dict) -> str:
+    for env_key in ("MOTHER_PUBLIC_SHARD_URL",):
+        override = (os.getenv(env_key) or "").strip().rstrip("/")
+        if override:
+            return override
+
+    gateway = (
+        os.getenv("MOTHER_URL")
+        or os.getenv("BLITZWING_MOTHER_URL")
+        or os.getenv("MOTHER_PUBLIC_GATEWAY_URL")
+        or ""
+    ).strip().rstrip("/")
+    if gateway:
+        return gateway
+
+    for host in manifest.get("hosts", []):
+        if host.get("role") != "mother":
+            continue
+        url = (host.get("shard_manager_url") or "").rstrip("/")
+        if _is_reachable_mother_url(url):
+            return url
+        public_ip = host.get("public_ip")
+        if public_ip:
+            return f"http://{public_ip}:8000"
+    raise RuntimeError(
+        "Cannot resolve mother prefix URL from swarm manifest "
+        "(set BLITZWING_MOTHER_URL or MOTHER_PUBLIC_SHARD_URL)"
+    )
+
+
+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"
+
+
+@dataclass
+class HttpChainInference:
+    runner: LocalShardRunner
+    local_block_indices: str
+
+    def generate(
+        self,
+        messages: Sequence[dict],
+        *,
+        max_tokens: int = 64,
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        swarm_manifest: Optional[dict] = None,
+    ) -> InferenceResult:
+        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,
+            )
+
+        tokenizer = self.runner.tokenizer
+        try:
+            prompt = tokenizer.apply_chat_template(
+                list(messages),
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+        except Exception:  # noqa: BLE001
+            parts: List[str] = []
+            for msg in messages:
+                role = msg.get("role", "user")
+                content = msg.get("content", "")
+                parts.append(f"<|{role}|>\n{content}</s>")
+            parts.append("<|assistant|>\n")
+            prompt = "".join(parts)
+
+        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))
+        logger.info(
+            "HTTP chain: mother=%s blocks=%s local=%s max_tokens=%s",
+            mother_url,
+            mother_blocks,
+            self.local_block_indices,
+            max_new_tokens,
+        )
+        generated: List[int] = []
+        finish_reason = "stop"
+
+        with httpx.Client(timeout=prefix_timeout) as client:
+            for step in range(max_new_tokens):
+                full_ids = (
+                    torch.cat(
+                        [input_ids, torch.tensor([generated], dtype=input_ids.dtype)],
+                        dim=-1,
+                    )
+                    if generated
+                    else input_ids
+                )
+
+                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)
+
+                hidden = self.runner.forward_tail(hidden)
+                logits = self.runner.logits_from_hidden(hidden)
+                next_logits = logits[0, -1, :]
+
+                if temperature <= 0:
+                    next_id = int(torch.argmax(next_logits).item())
+                else:
+                    probs = torch.softmax(next_logits / float(temperature), dim=-1)
+                    if top_p < 1.0:
+                        sorted_probs, sorted_idx = torch.sort(probs, descending=True)
+                        cumulative = torch.cumsum(sorted_probs, dim=-1)
+                        mask = cumulative > float(top_p)
+                        mask[..., 1:] = mask[..., :-1].clone()
+                        mask[..., 0] = False
+                        sorted_probs[mask] = 0
+                        sorted_probs = sorted_probs / sorted_probs.sum()
+                        pick = torch.multinomial(sorted_probs, 1).item()
+                        next_id = int(sorted_idx[pick].item())
+                    else:
+                        next_id = int(torch.multinomial(probs, 1).item())
+
+                generated.append(next_id)
+                if next_id == tokenizer.eos_token_id:
+                    break
+            else:
+                finish_reason = "length"
+
+        text = tokenizer.decode(generated, skip_special_tokens=True)
+        return InferenceResult(
+            text=text,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=len(generated),
+            finish_reason=finish_reason,
+        )

package/runtime/inference.py

@@ -0,0 +1,13 @@
+"""Minimal inference result type for HTTP-chain contributors."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class InferenceResult:
+    text: str
+    prompt_tokens: int
+    completion_tokens: int
+    finish_reason: str