Zachery Aaron Shores-Chmielewski
ae9ca3bcf3
Promote pipeline-parallel-inference to a first-class app and consolidate observability on the datastream wire, decoupling the dashboard crate from `distribution`. - apps/pipeline-parallel-inference: move the example out of `examples/` into `apps/` as its own workspace, rename binaries to `pp-worker`/`pp-orchestrator`, and strip release binaries - cluster: add `ClusterNode`, a synchronous facade over the actorized distribution protocol (IrohDriver + per-node Runtime hosting Swim/Registry/Metadata/Directory actors with a `MembershipFanout`), replacing ad-hoc `driver.node()`/`tick()` call sites - fleet: add per-node fleet telemetry that ships identity/resource records as `DatastreamFrame`s over the cluster transport to the orchestrator's `DatastreamSink`, folded into a `FleetView` on a 3s tick - provision: add best-effort, opt-in SSH boot-phase telemetry (`PP_DEPLOY_KEY`) that streams rented-node boot logs onto the orchestrator's datastream as `proc.boot.<stage>.*` - dashboard: rewire the crate dependency from `distribution` to `datastream`, drop the standalone `swactor-datastream-dashboard` binary, and rewrite `datastream_source.rs` to demux per-node frames into Overview/Distribution/Fleet views with live-node TTL filtering - distribution: refresh dist/netmap plugin copy and README from "Kademlia routing" to gossip-directory terminology Signed-off-by: Zachery Aaron Shores-Chmielewski <zacheryasc@gmail.com>
1675 lines
66 KiB
Python
1675 lines
66 KiB
Python
#!/usr/bin/env python3
|
||
"""Pipeline-parallel tinygrad worker — stdin/stdout JSON protocol.
|
||
|
||
Each worker hosts one *stage* of a pipeline-parallel inference run. The
|
||
stage index, total stage count, and the model name are read from
|
||
environment variables at startup:
|
||
|
||
STAGE=0 NUM_STAGES=2 MODEL=llama3.2:1b
|
||
|
||
The worker runs in one of two modes:
|
||
|
||
* **stub** (``--stub`` flag or ``PP_WORKER_STUB=1`` env): skips loading
|
||
any real model. Stage-0 ops return deterministic pseudo-bf16 hidden
|
||
bytes derived from the inputs; the final-stage op returns a
|
||
deterministic ``token_id``. ``tokenize`` whitespace-splits and hashes;
|
||
``detokenize`` formats the token ids as text. This mode is used by
|
||
the Rust actor and integration tests so they never need a GPU.
|
||
|
||
* **real** (default): loads ``$MODEL`` via ``Transformer.from_gguf`` and
|
||
``SimpleTokenizer.from_gguf_kv`` and runs the per-stage forward pass
|
||
directly against ``model.blk[start:end]`` (bypassing
|
||
``Transformer.forward`` and ``forward_jit``). Hidden state is cast to
|
||
float16 on the wire (2 bytes per element); decoded back to float16
|
||
on the receiving stage. Stage-1's ``forward_and_sample`` follows the
|
||
block range with ``output_norm`` + ``output`` and argmax-samples the
|
||
last position's logits.
|
||
|
||
Protocol — one JSON line in, one JSON line out.
|
||
|
||
Each op is restricted to the role that owns it (First = stage 0,
|
||
Last = stage N-1, Middle = anything in between). Ops invoked on the
|
||
wrong role return ``{"error": ...}`` and the worker keeps serving.
|
||
|
||
First-stage ops (``STAGE == 0``)::
|
||
|
||
-> {"op": "embed_and_forward", "request_id": <int>,
|
||
"tokens": [<int>, ...], "position": <int>}
|
||
<- {"request_id": <int>, "hidden_b64": "<base64>", "seq_len": <int>}
|
||
|
||
-> {"op": "decode_step", "request_id": <int>,
|
||
"token_id": <int>, "position": <int>}
|
||
<- {"request_id": <int>, "hidden_b64": "<base64>", "seq_len": 1}
|
||
|
||
-> {"op": "tokenize", "request_id": <int>, "prompt": "<text>"}
|
||
<- {"request_id": <int>, "tokens": [<int>, ...]}
|
||
|
||
Middle-stage op (``0 < STAGE < NUM_STAGES - 1``)::
|
||
|
||
-> {"op": "forward_range", "request_id": <int>,
|
||
"hidden_b64": "<base64>", "position": <int>, "seq_len": <int>}
|
||
<- {"request_id": <int>, "hidden_b64": "<base64>", "seq_len": <int>}
|
||
|
||
Last-stage ops (``STAGE == NUM_STAGES - 1``)::
|
||
|
||
-> {"op": "forward_and_sample", "request_id": <int>,
|
||
"hidden_b64": "<base64>", "position": <int>, "seq_len": <int>}
|
||
<- {"request_id": <int>, "token_id": <int>}
|
||
|
||
-> {"op": "detokenize", "request_id": <int>, "tokens": [<int>, ...]}
|
||
<- {"request_id": <int>, "text": "<text>"}
|
||
|
||
Errors of any kind reply with::
|
||
|
||
<- {"request_id": <int>?, "error": "<message>"}
|
||
|
||
and the worker keeps serving. EOF on stdin causes a clean exit 0.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import base64
|
||
import functools
|
||
import hashlib
|
||
import io
|
||
import json
|
||
import os
|
||
import re
|
||
import signal
|
||
import struct
|
||
import sys
|
||
import threading
|
||
import time
|
||
import urllib.parse
|
||
import urllib.request
|
||
from pathlib import Path
|
||
from typing import Sequence
|
||
|
||
# Stub-mode constants — small so test payloads stay tiny. Both values
|
||
# must match what the Rust stub-mode actor tests assume.
|
||
STUB_HIDDEN_DIM = 16
|
||
STUB_VOCAB_SIZE = 32
|
||
BYTES_PER_ELEM = 2 # bf16
|
||
|
||
# ASCII-only integer; rejects unicode digits and surrounding whitespace.
|
||
_INT_RE = re.compile(r"-?[0-9]+")
|
||
|
||
|
||
def compute_layer_range(stage: int, num_stages: int, total_blocks: int) -> tuple[int, int]:
|
||
"""Return the half-open ``[start, end)`` block range owned by ``stage``.
|
||
|
||
The split is ``total_blocks // num_stages`` per stage; the final stage
|
||
absorbs any remainder. ``num_stages`` must be ``>= 2`` — the example
|
||
does not serve single-node configurations (see
|
||
``examples/single-gpu-inference`` for that). Raises ``ValueError``
|
||
on out-of-range inputs.
|
||
"""
|
||
if num_stages < 2:
|
||
raise ValueError(f"num_stages must be >= 2, got {num_stages}")
|
||
if not (0 <= stage < num_stages):
|
||
raise ValueError(f"stage {stage} out of range [0, {num_stages})")
|
||
if total_blocks < num_stages:
|
||
raise ValueError(
|
||
f"total_blocks {total_blocks} cannot be split into {num_stages} stages"
|
||
)
|
||
k = total_blocks // num_stages
|
||
start = stage * k
|
||
end = (stage + 1) * k if stage < num_stages - 1 else total_blocks
|
||
return start, end
|
||
|
||
|
||
# ggml type tables for the sharded loader: quantized types map to
|
||
# (elements_per_block, bytes_per_block); native types map to byte width. These
|
||
# mirror tinygrad 0.12.0's ggml_data_to_tensor and let us size each tensor's raw
|
||
# byte slice so only the kept weights are copied off disk.
|
||
_GGML_QUANT_BLOCK = {2: (32, 18), 3: (32, 20), 8: (32, 34), 12: (256, 144), 14: (256, 210), 39: (32, 17)}
|
||
_GGML_NATIVE_ITEMSIZE = {0: 4, 1: 2, 16: 1, 17: 2, 18: 4}
|
||
|
||
|
||
def _ggml_tensor_nbytes(n_elements: int, ggml_type: int) -> int:
|
||
"""Raw byte size of an ``n_elements`` ggml tensor of ``ggml_type``."""
|
||
if ggml_type in _GGML_NATIVE_ITEMSIZE:
|
||
return _GGML_NATIVE_ITEMSIZE[ggml_type] * n_elements
|
||
if ggml_type in _GGML_QUANT_BLOCK:
|
||
elems_per_block, bytes_per_block = _GGML_QUANT_BLOCK[ggml_type]
|
||
return (n_elements // elems_per_block) * bytes_per_block
|
||
raise ValueError(f"unsupported ggml type {ggml_type}")
|
||
|
||
|
||
# ─── Sharded download (spec §4.1 / §4.2 / §4.7) ───────────────────────────
|
||
#
|
||
# The worker fetches only the byte ranges its stage actually needs
|
||
# (§4.1), caches the partial GGUF idempotently (§4.2), and emits
|
||
# `pp_download_progress` events with bounded latency while downloading
|
||
# (§4.7). The cached file is a SPARSE file with the same apparent size
|
||
# as the source — kept tensors live at their original byte offsets, so
|
||
# `_load_sharded_transformer` opens it unchanged. Filesystem holes
|
||
# absorb the non-kept regions so the actual disk usage is
|
||
# O(per-stage shard size), not O(full model size).
|
||
|
||
_PP_DOWNLOAD_READ_CHUNK = 256 * 1024
|
||
_PP_HEADER_INITIAL = 1024 * 1024
|
||
_PP_HEADER_MAX = 64 * 1024 * 1024
|
||
|
||
|
||
def _pp_round_up(n: int, align: int) -> int:
|
||
if align <= 0:
|
||
return n
|
||
return ((n + align - 1) // align) * align
|
||
|
||
|
||
class _PpHeaderTooShort(Exception):
|
||
"""Raised mid-parse when the header buffer ran out — caller grows it."""
|
||
|
||
|
||
def _pp_parse_gguf_header(buf: bytes) -> "tuple[list[tuple[str, tuple, int, int]], int, dict]":
|
||
"""Parse a GGUF header from `buf`. Returns (t_infos, data_start, kv).
|
||
|
||
`t_infos` is a list of ``(name, dims, ggml_type, offset)`` tuples
|
||
matching what ``_load_sharded_transformer``'s in-file parser
|
||
produces. `data_start` is the absolute byte offset where tensor
|
||
data begins. Raises `_PpHeaderTooShort` if the header is larger
|
||
than `buf` — caller should re-fetch with a larger buffer.
|
||
|
||
Format reference: tinygrad 0.12.0's gguf reader. GGUF versions 2
|
||
and 3 share the parse shape; the file's u32 version is checked.
|
||
"""
|
||
bio = io.BytesIO(buf)
|
||
|
||
def _need(n: int) -> bytes:
|
||
start = bio.tell()
|
||
out = bio.read(n)
|
||
if len(out) != n:
|
||
raise _PpHeaderTooShort(f"need {n} bytes at {start}, got {len(out)}")
|
||
return out
|
||
|
||
def _unpack(fmt: str, nbytes: int):
|
||
return struct.unpack(fmt, _need(nbytes))[0]
|
||
|
||
def _read_u32() -> int:
|
||
return _unpack("<I", 4)
|
||
|
||
def _read_i32() -> int:
|
||
return _unpack("<i", 4)
|
||
|
||
def _read_u64() -> int:
|
||
return _unpack("<Q", 8)
|
||
|
||
def _read_str() -> str:
|
||
length = _read_u64()
|
||
return _need(length).decode("utf-8")
|
||
|
||
def _read_arr():
|
||
elem_type = _read_i32()
|
||
count = _read_u64()
|
||
return [_readers[elem_type]() for _ in range(count)]
|
||
|
||
_readers = {
|
||
0: lambda: _unpack("<b", 1),
|
||
1: lambda: _unpack("<B", 1),
|
||
2: lambda: _unpack("<h", 2),
|
||
3: lambda: _unpack("<H", 2),
|
||
4: _read_u32,
|
||
5: _read_i32,
|
||
6: lambda: _unpack("<f", 4),
|
||
7: lambda: _unpack("<?", 1),
|
||
8: _read_str,
|
||
9: _read_arr,
|
||
10: _read_u64,
|
||
11: lambda: _unpack("<q", 8),
|
||
12: lambda: _unpack("<d", 8),
|
||
}
|
||
|
||
magic = _need(4)
|
||
if magic != b"GGUF":
|
||
raise ValueError(f"not a GGUF artifact (magic={magic!r})")
|
||
version = _read_i32()
|
||
if version not in (2, 3):
|
||
raise ValueError(f"unsupported GGUF version {version}")
|
||
n_tensors = _read_u64()
|
||
n_kv = _read_u64()
|
||
|
||
kv: "dict[str, object]" = {}
|
||
for _ in range(n_kv):
|
||
key = _read_str()
|
||
typ = _read_i32()
|
||
kv[key] = _readers[typ]()
|
||
|
||
t_infos: "list[tuple[str, tuple, int, int]]" = []
|
||
for _ in range(n_tensors):
|
||
name = _read_str()
|
||
n_dims = _read_u32()
|
||
dims = tuple(_read_u64() for _ in range(n_dims))
|
||
ggml_type = _read_i32()
|
||
offset = _read_u64()
|
||
t_infos.append((name, dims, ggml_type, offset))
|
||
|
||
alignment = int(kv.get("general.alignment", 32))
|
||
data_start = _pp_round_up(bio.tell(), alignment)
|
||
return t_infos, data_start, kv
|
||
|
||
|
||
def _pp_kept_names(t_infos, stage: int, num_stages: int, kv: dict) -> "set[str]":
|
||
"""Return the set of tensor names this stage requires. Matches the
|
||
`_kept` predicate inside `_load_sharded_transformer` — both code
|
||
paths must agree on the kept set or the loader would try to realize
|
||
a tensor whose bytes were not fetched."""
|
||
arch = str(kv["general.architecture"])
|
||
total_blocks = int(kv[f"{arch}.block_count"])
|
||
start, end = compute_layer_range(stage, num_stages, total_blocks)
|
||
is_last = stage == num_stages - 1
|
||
names = {info[0] for info in t_infos}
|
||
tied_output = "output.weight" not in names
|
||
kept: "set[str]" = set()
|
||
for info in t_infos:
|
||
name = info[0]
|
||
keep = False
|
||
for i in range(start, end):
|
||
if name.startswith(f"blk.{i}."):
|
||
keep = True
|
||
break
|
||
if not keep:
|
||
if name == "token_embd.weight" and (stage == 0 or (is_last and tied_output)):
|
||
keep = True
|
||
elif name == "output_norm.weight" and is_last:
|
||
keep = True
|
||
elif name == "output.weight" and is_last and not tied_output:
|
||
keep = True
|
||
if keep:
|
||
kept.add(name)
|
||
return kept
|
||
|
||
|
||
def _pp_kept_byte_ranges(t_infos, data_start: int, kept_names: "set[str]") -> "list[tuple[int, int]]":
|
||
"""Return the ``[(absolute_offset, nbytes)]`` byte ranges this stage
|
||
keeps, ordered by offset (so a streamed download writes ascending
|
||
offsets and the filesystem allocates fewer fragmented holes)."""
|
||
out: list[tuple[int, int]] = []
|
||
for name, dims, ggml_type, offset in t_infos:
|
||
if name not in kept_names:
|
||
continue
|
||
n_elements = 1
|
||
for d in dims:
|
||
n_elements *= int(d)
|
||
nbytes = _ggml_tensor_nbytes(n_elements, ggml_type)
|
||
out.append((data_start + int(offset), nbytes))
|
||
out.sort()
|
||
return out
|
||
|
||
|
||
def _pp_cache_paths(url: str) -> "tuple[Path, Path, Path]":
|
||
"""Return (cache_path, meta_path, partial_path) for ``url``.
|
||
|
||
Cache root defaults to ``$PP_MODEL_CACHE_DIR`` then ``~/.cache/pp-pipeline``.
|
||
The filename is ``<short-url-hash>-<basename>`` so two URLs that share a
|
||
basename cannot collide.
|
||
"""
|
||
raw = os.environ.get("PP_MODEL_CACHE_DIR", "").strip()
|
||
cache_root = Path(raw).expanduser() if raw else Path.home() / ".cache" / "pp-pipeline"
|
||
cache_root.mkdir(parents=True, exist_ok=True)
|
||
url_hash = hashlib.sha256(url.encode("utf-8")).hexdigest()[:16]
|
||
basename = os.path.basename(urllib.parse.urlparse(url).path) or "model.gguf"
|
||
cache_path = cache_root / f"{url_hash}-{basename}"
|
||
meta_path = cache_path.with_name(cache_path.name + ".pp_meta")
|
||
partial_path = cache_path.with_name(cache_path.name + ".partial")
|
||
return cache_path, meta_path, partial_path
|
||
|
||
|
||
def _pp_head(url: str) -> "tuple[int, str]":
|
||
"""HEAD request; return (Content-Length, Accept-Ranges header lowercased)."""
|
||
req = urllib.request.Request(url, method="HEAD")
|
||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||
total = int(resp.headers.get("Content-Length", "0"))
|
||
accept = (resp.headers.get("Accept-Ranges") or "").lower()
|
||
return total, accept
|
||
|
||
|
||
def _pp_range_get(url: str, start: int, end_inclusive: int) -> bytes:
|
||
"""Issue a Range GET; return the body bytes."""
|
||
req = urllib.request.Request(
|
||
url, headers={"Range": f"bytes={start}-{end_inclusive}"}
|
||
)
|
||
with urllib.request.urlopen(req, timeout=120) as resp:
|
||
return resp.read()
|
||
|
||
|
||
def _pp_emit_progress(stage: int, bytes_done: int, bytes_total: int, started_at_mono: float) -> None:
|
||
"""Emit one `pp_download_progress` event with the spec's field set
|
||
(§4.7). `started_at_mono` is `time.monotonic()` captured before the
|
||
first event so `elapsed_ms` is monotonic across the fetch."""
|
||
elapsed_ms = int((time.monotonic() - started_at_mono) * 1000)
|
||
mbps = round(((bytes_done * 8) / 1_000_000) / max(elapsed_ms / 1000, 1e-3), 1)
|
||
_emit_event(
|
||
"pp_download_progress",
|
||
stage_index=stage,
|
||
bytes_done=bytes_done,
|
||
bytes_total=bytes_total,
|
||
elapsed_ms=elapsed_ms,
|
||
mbps=mbps,
|
||
)
|
||
|
||
|
||
def _pp_fingerprint(f, offset: int, nbytes: int) -> str:
|
||
"""Hash the first 4KB + last 4KB of a kept tensor (or the whole
|
||
tensor if shorter). The §4.2 content-derived check: matches the
|
||
fingerprint recorded in the sidecar at download time."""
|
||
sample = 4096
|
||
f.seek(offset)
|
||
head = f.read(min(sample, nbytes))
|
||
if nbytes > sample:
|
||
f.seek(offset + nbytes - sample)
|
||
tail = f.read(sample)
|
||
else:
|
||
tail = b""
|
||
h = hashlib.sha256()
|
||
h.update(head)
|
||
h.update(tail)
|
||
h.update(nbytes.to_bytes(8, "little"))
|
||
return h.hexdigest()
|
||
|
||
|
||
def _pp_write_meta(
|
||
meta_path: Path,
|
||
cache_path: Path,
|
||
stage: int,
|
||
num_stages: int,
|
||
url: str,
|
||
total_size: int,
|
||
kept_ranges: "list[tuple[int, int]]",
|
||
) -> None:
|
||
fingerprints = []
|
||
with open(cache_path, "rb") as f:
|
||
for offset, nbytes in kept_ranges:
|
||
fingerprints.append({
|
||
"offset": offset,
|
||
"nbytes": nbytes,
|
||
"fingerprint": _pp_fingerprint(f, offset, nbytes),
|
||
})
|
||
meta_path.write_text(json.dumps({
|
||
"schema": 1,
|
||
"url": url,
|
||
"stage": stage,
|
||
"num_stages": num_stages,
|
||
"total_size": total_size,
|
||
"kept": fingerprints,
|
||
}))
|
||
|
||
|
||
def _pp_verify_cache(
|
||
cache_path: Path,
|
||
meta_path: Path,
|
||
stage: int,
|
||
num_stages: int,
|
||
) -> bool:
|
||
"""Spec §4.2 integrity check: file present at the expected apparent
|
||
size AND every recorded fingerprint re-matches the cached bytes.
|
||
Returns False on any discrepancy (including missing files, missing
|
||
sidecar, mismatched stage / num_stages, size mismatch, or any
|
||
fingerprint mismatch). A passing cache is used as-is — no refetch."""
|
||
if not cache_path.exists() or not meta_path.exists():
|
||
return False
|
||
try:
|
||
meta = json.loads(meta_path.read_text())
|
||
except (OSError, ValueError):
|
||
return False
|
||
if meta.get("schema") != 1:
|
||
return False
|
||
if meta.get("stage") != stage or meta.get("num_stages") != num_stages:
|
||
return False
|
||
if cache_path.stat().st_size != meta.get("total_size"):
|
||
return False
|
||
kept = meta.get("kept") or []
|
||
if not kept:
|
||
return False
|
||
try:
|
||
with open(cache_path, "rb") as f:
|
||
for entry in kept:
|
||
offset = int(entry["offset"])
|
||
nbytes = int(entry["nbytes"])
|
||
expected = entry["fingerprint"]
|
||
if _pp_fingerprint(f, offset, nbytes) != expected:
|
||
return False
|
||
except (OSError, KeyError, ValueError):
|
||
return False
|
||
return True
|
||
|
||
|
||
def _pp_download_sharded(url: str, stage: int, num_stages: int) -> str:
|
||
"""Spec §4.1 + §4.2 + §4.7: fetch only this stage's tensor bytes,
|
||
cache idempotently, emit progress events.
|
||
|
||
Returns the path to the on-disk file (sparse — apparent size matches
|
||
the source; only the kept ranges occupy disk blocks). On cache hit,
|
||
NO `pp_download_progress` events are emitted (spec §4.7).
|
||
"""
|
||
cache_path, meta_path, partial_path = _pp_cache_paths(url)
|
||
|
||
# §4.2: orphan-cleanup any leftover .partial from a previous killed
|
||
# fetch BEFORE any new fetch is initiated. Emit an event so the
|
||
# bundle reader can see that a stale temp was reaped.
|
||
if partial_path.exists():
|
||
try:
|
||
partial_path.unlink()
|
||
_emit_event(
|
||
"pp_cache_orphan_cleaned",
|
||
stage_index=stage,
|
||
path=str(partial_path),
|
||
)
|
||
except OSError:
|
||
pass
|
||
|
||
# §4.2 cache hit — return the cached file as-is.
|
||
if _pp_verify_cache(cache_path, meta_path, stage, num_stages):
|
||
_emit_event(
|
||
"pp_cache_hit",
|
||
stage_index=stage,
|
||
path=str(cache_path),
|
||
)
|
||
return str(cache_path)
|
||
|
||
# §4.1 fail-fast on no-range support.
|
||
total_size, accept_ranges = _pp_head(url)
|
||
if "bytes" not in accept_ranges:
|
||
_emit_event(
|
||
"pp_download_failed",
|
||
stage_index=stage,
|
||
reason="no_byte_range_support",
|
||
url=url,
|
||
accept_ranges=accept_ranges,
|
||
)
|
||
_die(f"source does not support byte-range requests: {url}")
|
||
if total_size <= 0:
|
||
_emit_event(
|
||
"pp_download_failed",
|
||
stage_index=stage,
|
||
reason="no_content_length",
|
||
url=url,
|
||
)
|
||
_die(f"source did not advertise Content-Length: {url}")
|
||
|
||
# Fetch the header in growing increments until we can parse it.
|
||
header_size = min(_PP_HEADER_INITIAL, total_size)
|
||
while True:
|
||
try:
|
||
header_bytes = _pp_range_get(url, 0, header_size - 1)
|
||
t_infos, data_start, kv = _pp_parse_gguf_header(header_bytes)
|
||
if data_start <= len(header_bytes):
|
||
break
|
||
# Parser succeeded structurally but data_start sits past our
|
||
# buffer — re-fetch enough to include the tensor data start.
|
||
header_size = min(data_start + 1024, total_size)
|
||
except _PpHeaderTooShort:
|
||
new_size = min(header_size * 2, total_size)
|
||
if new_size == header_size or new_size > _PP_HEADER_MAX:
|
||
_emit_event(
|
||
"pp_download_failed",
|
||
stage_index=stage,
|
||
reason="header_too_large",
|
||
header_size=header_size,
|
||
)
|
||
_die(f"GGUF header exceeded {_PP_HEADER_MAX} bytes")
|
||
header_size = new_size
|
||
|
||
kept_names = _pp_kept_names(t_infos, stage, num_stages, kv)
|
||
kept_ranges = _pp_kept_byte_ranges(t_infos, data_start, kept_names)
|
||
|
||
# bytes_total is the bytes this stage will pull from the network:
|
||
# header + kept-tensor regions. Not the full file (spec §4.1 means
|
||
# we never fetch the rest).
|
||
header_keep_bytes = data_start
|
||
kept_bytes_total = sum(nb for _, nb in kept_ranges)
|
||
bytes_total = header_keep_bytes + kept_bytes_total
|
||
|
||
interval_s_raw = os.environ.get("PP_DOWNLOAD_PROGRESS_INTERVAL_SECS", "").strip()
|
||
try:
|
||
interval_s = float(interval_s_raw) if interval_s_raw else 10.0
|
||
except ValueError:
|
||
interval_s = 10.0
|
||
if interval_s <= 0:
|
||
interval_s = 10.0
|
||
|
||
# Write the sparse output to `.partial`; rename on success. Opening
|
||
# with "wb" then truncate(total_size) creates a sparse file on
|
||
# Linux: only blocks we actually `write()` allocate disk.
|
||
started_at = time.monotonic()
|
||
with open(partial_path, "wb") as f:
|
||
f.truncate(total_size)
|
||
f.seek(0)
|
||
f.write(header_bytes[:data_start])
|
||
bytes_done = data_start
|
||
|
||
# Spec §4.7: first event MUST be at start of fetch, AFTER we
|
||
# know bytes_total. We have that now.
|
||
_pp_emit_progress(stage, bytes_done, bytes_total, started_at)
|
||
last_emit = time.monotonic()
|
||
|
||
for offset, nbytes in kept_ranges:
|
||
req = urllib.request.Request(
|
||
url,
|
||
headers={"Range": f"bytes={offset}-{offset + nbytes - 1}"},
|
||
)
|
||
try:
|
||
resp = urllib.request.urlopen(req, timeout=300)
|
||
except Exception as e:
|
||
_emit_event(
|
||
"pp_download_failed",
|
||
stage_index=stage,
|
||
reason="range_get_failed",
|
||
offset=offset,
|
||
nbytes=nbytes,
|
||
error=str(e),
|
||
)
|
||
# Spec §4.7: final event MUST be emitted on fetch failure.
|
||
_pp_emit_progress(stage, bytes_done, bytes_total, started_at)
|
||
_die(f"range GET failed at offset {offset}: {e}")
|
||
try:
|
||
f.seek(offset)
|
||
while True:
|
||
chunk = resp.read(_PP_DOWNLOAD_READ_CHUNK)
|
||
if not chunk:
|
||
break
|
||
f.write(chunk)
|
||
bytes_done += len(chunk)
|
||
now = time.monotonic()
|
||
if now - last_emit >= interval_s:
|
||
_pp_emit_progress(stage, bytes_done, bytes_total, started_at)
|
||
last_emit = now
|
||
finally:
|
||
resp.close()
|
||
|
||
# Spec §4.7: final event at completion.
|
||
_pp_emit_progress(stage, bytes_done, bytes_total, started_at)
|
||
|
||
# Sidecar before rename so a crash between rename + meta-write does
|
||
# not leave a "valid file, no sidecar" → would fail _pp_verify and
|
||
# refetch. Writing the sidecar first means a crash here leaves
|
||
# cache_path absent and partial_path present (which orphan-cleanup
|
||
# reaps on next boot).
|
||
_pp_write_meta(
|
||
meta_path, partial_path, stage, num_stages, url, total_size, kept_ranges
|
||
)
|
||
os.replace(partial_path, cache_path)
|
||
return str(cache_path)
|
||
|
||
|
||
def _load_sharded_transformer(gguf_path, stage: int, num_stages: int, max_context: int = 512):
|
||
"""Load only this stage's slice of the model onto the compute device.
|
||
|
||
Stock ``Transformer.from_gguf`` copies the *entire* GGUF onto the compute
|
||
device before any layer runs (it does ``gguf.to(None)``), so an 18 GB model
|
||
OOMs a 12 GB GPU no matter how the layers are split. Instead we parse the
|
||
GGUF header on the DISK device and copy only the tensors this stage needs —
|
||
``blk[start:end]`` plus ``token_embd`` (stage 0) and ``output_norm`` /
|
||
``output`` (last stage) — dequantizing each on the compute device. Returns
|
||
``(model, kv, start, end)``. Vendored against tinygrad 0.12.0's gguf format.
|
||
"""
|
||
import io
|
||
import struct
|
||
import functools
|
||
|
||
from tinygrad import Tensor, Device, nn
|
||
from tinygrad.helpers import prod, round_up, getenv
|
||
from tinygrad.nn.state import TensorIO, ggml_data_to_tensor
|
||
from tinygrad.apps.llm import Transformer
|
||
|
||
_t0 = time.monotonic()
|
||
gguf = Tensor(Path(gguf_path)) # device is DISK:<path> — nothing is copied to the GPU yet
|
||
|
||
# --- parse the GGUF header (kv metadata + tensor directory) off disk ---
|
||
reader = io.BufferedReader(TensorIO(gguf), 1_000_000)
|
||
|
||
def _unpack(fmt, nbytes):
|
||
return struct.unpack(fmt, reader.read(nbytes))[0]
|
||
|
||
def _read_str():
|
||
return str(reader.read(_read_u64()), "utf-8")
|
||
|
||
def _read_arr():
|
||
elem_reader, count = _readers[_read_i32()], _read_u64()
|
||
return [elem_reader() for _ in range(count)]
|
||
|
||
_readers = {8: _read_str, 9: _read_arr, **{t: functools.partial(_unpack, "<" + f, nb) for t, f, nb in
|
||
[(0, "c", 1), (1, "b", 1), (2, "H", 2), (3, "h", 2), (4, "I", 4), (5, "i", 4),
|
||
(6, "f", 4), (7, "?", 1), (10, "Q", 8), (11, "q", 8), (12, "d", 8)]}}
|
||
_read_u32, _read_i32, _read_u64 = _readers[4], _readers[5], _readers[10]
|
||
|
||
magic, version = reader.read(4), _read_i32()
|
||
n_tensors, n_kv = _read_u64(), _read_u64()
|
||
if magic != b"GGUF" or version not in (2, 3):
|
||
raise ValueError(f"invalid GGUF (magic={magic!r} version={version})")
|
||
kv = {}
|
||
for _ in range(n_kv):
|
||
key, typ = _read_str(), _read_i32()
|
||
kv[key] = _readers[typ]()
|
||
t_infos = [(_read_str(), tuple(_read_u64() for _ in range(_read_u32())), _read_i32(), _read_u64())
|
||
for _ in range(n_tensors)]
|
||
data_start = round_up(reader.tell(), kv.get("general.alignment", 32))
|
||
_t_header = time.monotonic()
|
||
|
||
arch = kv["general.architecture"]
|
||
total_blocks = int(kv[f"{arch}.block_count"])
|
||
start, end = compute_layer_range(stage, num_stages, total_blocks)
|
||
is_last = stage == num_stages - 1
|
||
names = {info[0] for info in t_infos}
|
||
tied_output = "output.weight" not in names # small models tie output to token_embd
|
||
|
||
def _kept(name: str) -> bool:
|
||
for i in range(start, end):
|
||
if name.startswith(f"blk.{i}."):
|
||
return True
|
||
if name == "token_embd.weight" and (stage == 0 or (is_last and tied_output)):
|
||
return True
|
||
if name == "output_norm.weight" and is_last:
|
||
return True
|
||
if name == "output.weight" and is_last and not tied_output:
|
||
return True
|
||
return False
|
||
|
||
half, device = getenv("HALF", 1), Device.DEFAULT
|
||
state_dict = {}
|
||
bytes_copied = 0
|
||
kept_count = 0
|
||
for name, dims, ggml_type, offset in t_infos:
|
||
n_elements = prod(dims)
|
||
if _kept(name):
|
||
nbytes = _ggml_tensor_nbytes(n_elements, ggml_type)
|
||
bytes_copied += nbytes
|
||
kept_count += 1
|
||
raw = gguf[data_start + offset: data_start + offset + nbytes].to(device)
|
||
tensor = ggml_data_to_tensor(raw, n_elements, ggml_type).reshape(*reversed(dims))
|
||
if arch == "llama": # interleaved -> half-split RoPE layout (llama-style only)
|
||
n_heads, n_kv_heads = kv[f"{arch}.attention.head_count"], kv[f"{arch}.attention.head_count_kv"]
|
||
if "attn_q.weight" in name:
|
||
tensor = tensor.rearrange("(n h two) d -> (n two h) d", n=n_heads, two=2)
|
||
if "attn_k.weight" in name:
|
||
tensor = tensor.rearrange("(n h two) d -> (n two h) d", n=n_kv_heads, two=2)
|
||
state_dict[name] = tensor.cast("float16") if half else tensor
|
||
else:
|
||
# DISK-rooted lazy tensor: only its .shape is read (model construction); never realized.
|
||
state_dict[name] = ggml_data_to_tensor(gguf[data_start + offset:], n_elements, ggml_type).reshape(*reversed(dims))
|
||
if tied_output and is_last:
|
||
state_dict["output.weight"] = state_dict["token_embd.weight"]
|
||
_t_statedict = time.monotonic()
|
||
|
||
n_heads = kv[f"{arch}.attention.head_count"]
|
||
model = Transformer(
|
||
num_blocks=total_blocks, dim=kv[f"{arch}.embedding_length"],
|
||
hidden_dim=kv.get(f"{arch}.expert_feed_forward_length", kv[f"{arch}.feed_forward_length"]),
|
||
n_heads=n_heads, n_kv_heads=kv[f"{arch}.attention.head_count_kv"],
|
||
norm_eps=kv[f"{arch}.attention.layer_norm_rms_epsilon"], vocab_size=len(kv["tokenizer.ggml.tokens"]),
|
||
head_dim=kv.get(f"{arch}.attention.key_length", kv[f"{arch}.embedding_length"] // n_heads),
|
||
rope_theta=kv[f"{arch}.rope.freq_base"], max_context=max_context,
|
||
qk_norm=int(state_dict["blk.0.attn_q_norm.weight"].shape[0]) if "blk.0.attn_q_norm.weight" in state_dict else 0,
|
||
num_experts=kv.get(f"{arch}.expert_count", 0), num_experts_per_tok=kv.get(f"{arch}.expert_used_count", 0))
|
||
|
||
_t_construct = time.monotonic()
|
||
|
||
# Provide only the kept weights; strict=False leaves the other blocks at their
|
||
# (lazy, never-run) init so they never touch the compute device.
|
||
kept = {name: tensor for name, tensor in state_dict.items()
|
||
if _kept(name) or (tied_output and is_last and name == "output.weight")}
|
||
# This is where the kept tensors are actually copied off disk and
|
||
# dequantized on the compute device — the dominant load cost.
|
||
nn.state.load_state_dict(model, kept, strict=False, verbose=False, consume=True, realize=True)
|
||
_t_realize = time.monotonic()
|
||
|
||
realize_ms = (_t_realize - _t_construct) * 1000
|
||
_emit_event(
|
||
"model_load_breakdown",
|
||
stage=stage,
|
||
resident_blocks=end - start,
|
||
total_blocks=total_blocks,
|
||
kept_tensors=kept_count,
|
||
bytes_copied=bytes_copied,
|
||
mb_copied=round(bytes_copied / 1_000_000, 1),
|
||
header_ms=round((_t_header - _t0) * 1000, 1),
|
||
statedict_build_ms=round((_t_statedict - _t_header) * 1000, 1),
|
||
construct_ms=round((_t_construct - _t_statedict) * 1000, 1),
|
||
realize_ms=round(realize_ms, 1),
|
||
realize_mb_per_s=round((bytes_copied / 1_000_000) / max(realize_ms / 1000, 1e-3), 1),
|
||
rss_mb=_rss_mb(),
|
||
)
|
||
return model, kv, start, end
|
||
|
||
|
||
def argmax_sample(logits: Sequence[float]) -> int:
|
||
"""Return the index of the maximum element in ``logits``.
|
||
|
||
On ties, the lowest index wins. Raises ``ValueError`` on empty input.
|
||
"""
|
||
if len(logits) == 0:
|
||
raise ValueError("argmax over empty logits")
|
||
best_idx = 0
|
||
best_val = logits[0]
|
||
for i in range(1, len(logits)):
|
||
if logits[i] > best_val:
|
||
best_val = logits[i]
|
||
best_idx = i
|
||
return best_idx
|
||
|
||
|
||
def _stub_hidden_bytes(tokens: Sequence[int], position: int) -> bytes:
|
||
"""Deterministic pseudo-bf16 bytes of length ``len(tokens) * hidden_dim * 2``.
|
||
|
||
Same ``(tokens, position)`` always produce the same bytes; different
|
||
inputs almost always differ. Bytes are derived from SHA-256 so the
|
||
payload is non-trivial (not all zeros, not monotonic).
|
||
"""
|
||
seed = hashlib.sha256(
|
||
b"pp-stub:"
|
||
+ str(position).encode()
|
||
+ b"|"
|
||
+ b",".join(str(int(t)).encode() for t in tokens)
|
||
).digest()
|
||
n_bytes = len(tokens) * STUB_HIDDEN_DIM * BYTES_PER_ELEM
|
||
out = bytearray()
|
||
counter = 0
|
||
while len(out) < n_bytes:
|
||
out.extend(hashlib.sha256(seed + counter.to_bytes(8, "little")).digest())
|
||
counter += 1
|
||
return bytes(out[:n_bytes])
|
||
|
||
|
||
def _stub_forward_range_bytes(hidden: bytes, position: int, seq_len: int) -> bytes:
|
||
"""Deterministic pseudo-bf16 bytes for a middle stage's ``forward_range``.
|
||
|
||
Output length is ``seq_len * STUB_HIDDEN_DIM * 2`` (matches the input
|
||
hidden's expected length). Same ``(hidden, position, seq_len)`` always
|
||
produce the same bytes; differing inputs almost always differ. Used
|
||
only in stub mode.
|
||
"""
|
||
n_bytes = seq_len * STUB_HIDDEN_DIM * BYTES_PER_ELEM
|
||
seed = hashlib.sha256(
|
||
b"pp-stub-fr:" + position.to_bytes(8, "little", signed=False) + hidden
|
||
).digest()
|
||
out = bytearray()
|
||
counter = 0
|
||
while len(out) < n_bytes:
|
||
out.extend(hashlib.sha256(seed + counter.to_bytes(8, "little")).digest())
|
||
counter += 1
|
||
return bytes(out[:n_bytes])
|
||
|
||
|
||
def _stub_token_id(hidden: bytes, position: int) -> int:
|
||
"""Deterministic stub token id in ``[0, STUB_VOCAB_SIZE)``."""
|
||
digest = hashlib.sha256(
|
||
b"pp-stub-sample:" + position.to_bytes(8, "little", signed=False) + hidden
|
||
).digest()
|
||
return int.from_bytes(digest[:4], "little") % STUB_VOCAB_SIZE
|
||
|
||
|
||
def _parse_env_int(name: str) -> int:
|
||
raw = os.environ.get(name, "")
|
||
stripped = raw.strip()
|
||
if not stripped:
|
||
_die(f"env {name} is required")
|
||
if not _INT_RE.fullmatch(stripped):
|
||
_die(f"env {name}={raw!r} is not an ASCII integer")
|
||
return int(stripped)
|
||
|
||
|
||
def _die(msg: str) -> "None":
|
||
print(f"pp_tinygrad_worker: {msg}", file=sys.stderr, flush=True)
|
||
raise SystemExit(2)
|
||
|
||
|
||
def _write(obj) -> None:
|
||
sys.stdout.write(json.dumps(obj) + "\n")
|
||
sys.stdout.flush()
|
||
|
||
|
||
# ─── Lifecycle event emission ─────────────────────────────────────────────
|
||
#
|
||
# The Rust StageActor parses every stdout line as JSON; any line carrying
|
||
# `"event": "<kind>"` is re-emitted as `Custom("worker_<kind>")` into the
|
||
# diagnostic bundle. The worker subprocess is otherwise opaque to the
|
||
# Rust side, so these are the only diagnostic signal the bundle ever sees
|
||
# from the Python layer (apart from exit code + stderr tail). We do NOT
|
||
# emit a `request_id` on event lines so the actor never confuses an event
|
||
# with an op reply.
|
||
|
||
_WORKER_START_MONOTONIC = time.monotonic()
|
||
_REQUESTS_SERVED = 0
|
||
|
||
|
||
def _emit_event(kind: str, **fields) -> None:
|
||
"""Emit a structured lifecycle event on stdout. The Rust actor folds
|
||
these into the diag bundle as `Custom("worker_<kind>")`."""
|
||
payload = {"event": kind, **fields}
|
||
try:
|
||
sys.stdout.write(json.dumps(payload) + "\n")
|
||
sys.stdout.flush()
|
||
except Exception:
|
||
# Best-effort: never let a logging failure crash the worker.
|
||
pass
|
||
|
||
|
||
def _uptime_ms() -> int:
|
||
return int((time.monotonic() - _WORKER_START_MONOTONIC) * 1000)
|
||
|
||
|
||
def _wall_ms() -> int:
|
||
"""Epoch milliseconds. Lets the bundle align worker events across nodes
|
||
and against the orchestrator's vast.ai create/lease timestamps — e.g.
|
||
(worker `starting`.wall_ms − instance create_ms) is the image-pull +
|
||
container-boot + worker-spawn cost the node can't see itself."""
|
||
return int(time.time() * 1000)
|
||
|
||
|
||
def _rss_mb() -> "int | None":
|
||
"""Resident-set size in MB, read from /proc/self/status (Linux).
|
||
Returns None on non-Linux or when the read fails — the field is
|
||
informational, never required."""
|
||
try:
|
||
with open("/proc/self/status", "r") as f:
|
||
for line in f:
|
||
if line.startswith("VmRSS:"):
|
||
parts = line.split()
|
||
# VmRSS: 12345 kB
|
||
return int(parts[1]) // 1024
|
||
except (OSError, ValueError, IndexError):
|
||
pass
|
||
return None
|
||
|
||
|
||
def _install_excepthook() -> None:
|
||
"""Catch every uncaught exception and emit a structured event before
|
||
the interpreter prints the traceback to stderr (which the actor's
|
||
ring buffer will also capture)."""
|
||
|
||
def _hook(exc_type, exc_value, exc_tb):
|
||
import traceback
|
||
|
||
tb_text = "".join(traceback.format_exception(exc_type, exc_value, exc_tb))
|
||
_emit_event(
|
||
"uncaught_exception",
|
||
type=exc_type.__name__,
|
||
value=str(exc_value),
|
||
traceback=tb_text,
|
||
uptime_ms=_uptime_ms(),
|
||
)
|
||
# Preserve the default behaviour so stderr still shows the trace
|
||
# (the actor's stderr ring buffer is a belt-and-braces backup).
|
||
sys.__excepthook__(exc_type, exc_value, exc_tb)
|
||
|
||
sys.excepthook = _hook
|
||
|
||
|
||
def _install_signal_handlers() -> None:
|
||
"""Emit `signal_received` and exit cleanly on SIGTERM/SIGINT.
|
||
SIGKILL and SIGSEGV cannot be caught — the Rust side relies on the
|
||
exit code / signal field of the eventual `worker_exited` Custom
|
||
event for those."""
|
||
|
||
def _on_signal(signum, _frame):
|
||
try:
|
||
name = signal.Signals(signum).name
|
||
except ValueError:
|
||
name = f"signal-{signum}"
|
||
_emit_event(
|
||
"signal_received",
|
||
signum=signum,
|
||
name=name,
|
||
uptime_ms=_uptime_ms(),
|
||
)
|
||
# 128 + signum is the conventional exit code for signal-driven
|
||
# termination; matches what /bin/sh reports.
|
||
sys.exit(128 + signum)
|
||
|
||
for s in (signal.SIGTERM, signal.SIGINT):
|
||
try:
|
||
signal.signal(s, _on_signal)
|
||
except (ValueError, OSError):
|
||
# Some environments (e.g. non-main thread) don't allow
|
||
# signal install — silently skip rather than crash here.
|
||
pass
|
||
|
||
|
||
def _start_heartbeat(interval_s: float = 30.0) -> None:
|
||
"""Daemon thread emitting `heartbeat` events. Lets the post-processor
|
||
distinguish *hung* (heartbeats stop but process alive — no exit event)
|
||
from *dead* (no heartbeat AND no exit — likely SIGKILL/SIGSEGV)."""
|
||
|
||
def _loop():
|
||
while True:
|
||
time.sleep(interval_s)
|
||
_emit_event(
|
||
"heartbeat",
|
||
uptime_ms=_uptime_ms(),
|
||
rss_mb=_rss_mb(),
|
||
requests_served=_REQUESTS_SERVED,
|
||
)
|
||
|
||
t = threading.Thread(target=_loop, name="pp-worker-heartbeat", daemon=True)
|
||
t.start()
|
||
|
||
|
||
def _is_nonneg_int(x) -> bool:
|
||
# ``bool`` is a subclass of ``int`` in Python; reject it explicitly so
|
||
# ``{"position": true}`` doesn't sneak through.
|
||
return isinstance(x, int) and not isinstance(x, bool) and x >= 0
|
||
|
||
|
||
class _RealModelState:
|
||
"""Holds the loaded tinygrad model + tokenizer for real mode.
|
||
|
||
All tinygrad imports live inside the loader so stub-mode never pays
|
||
the import cost. Once loaded, the state is reused for every request.
|
||
"""
|
||
|
||
def __init__(self, model_name: str, stage: int, num_stages: int):
|
||
# Import tinygrad lazily so stub-mode never touches it. This is
|
||
# the most likely crash site in real mode — emit lifecycle
|
||
# events around the import so the bundle records exactly when
|
||
# the worker started loading and how long it took.
|
||
_emit_event("importing_tinygrad", stage=stage, wall_ms=_wall_ms())
|
||
_import_start = time.monotonic()
|
||
import numpy as np
|
||
from tinygrad import Tensor, Device
|
||
from tinygrad.helpers import fetch, getenv
|
||
from tinygrad.apps.llm import SimpleTokenizer, models
|
||
|
||
import_ms = int((time.monotonic() - _import_start) * 1000)
|
||
_emit_event("tinygrad_imported", stage=stage, elapsed_ms=import_ms)
|
||
# Device + precision context: which backend the shard lands on and the
|
||
# toggles (HALF/JIT/BEAM) that dominate load + inference cost. Correlate
|
||
# with model_load_breakdown / op events to attribute time to dequant vs
|
||
# kernel compile vs steady-state matmul.
|
||
_emit_event(
|
||
"device",
|
||
stage=stage,
|
||
default_device=str(Device.DEFAULT),
|
||
half=getenv("HALF", 1),
|
||
jit=getenv("JIT", 1),
|
||
beam=getenv("BEAM", 0),
|
||
cuda_visible=os.environ.get("CUDA_VISIBLE_DEVICES"),
|
||
)
|
||
|
||
if model_name not in models:
|
||
available = ", ".join(sorted(models.keys()))
|
||
_die(f"unknown MODEL {model_name!r}; available: {available}")
|
||
|
||
url = models[model_name]
|
||
_emit_event("fetching_model", stage=stage, model=model_name, url=url, wall_ms=_wall_ms())
|
||
print(
|
||
f"pp_tinygrad_worker: stage={stage}/{num_stages} fetching {model_name}",
|
||
file=sys.stderr,
|
||
flush=True,
|
||
)
|
||
_fetch_start = time.monotonic()
|
||
# Spec §4.1: download only this stage's tensor byte ranges.
|
||
# Spec §4.2: cache idempotently with an integrity check; a
|
||
# complete-and-valid cache MUST NOT trigger a network fetch.
|
||
# Spec §4.7: `pp_download_progress` events are emitted by
|
||
# `_pp_download_sharded` while the fetch is in progress and
|
||
# NEVER on a cache hit. The unused `fetch` import remains as
|
||
# documentation of the prior code path; the sharded fetcher
|
||
# replaces it.
|
||
_ = fetch # silence the linter; kept for the diff reader
|
||
gguf_path = _pp_download_sharded(url, stage, num_stages)
|
||
fetch_ms = int((time.monotonic() - _fetch_start) * 1000)
|
||
# `gguf_bytes` is the file's apparent size (matches the source's
|
||
# total_size); actual on-disk usage is O(per-stage shard). Bundle
|
||
# readers reading `model_fetched.gguf_bytes` see the same value
|
||
# they did before the §4.1 change — the per-stage usage shows up
|
||
# in `pp_download_progress.bytes_total` (header + kept ranges).
|
||
try:
|
||
gguf_bytes = os.path.getsize(gguf_path)
|
||
except OSError:
|
||
gguf_bytes = 0
|
||
fetch_mb = gguf_bytes / 1_000_000
|
||
# A near-instant return with a non-zero apparent size means the
|
||
# sharded cache hit short-circuited the fetch. Distinguishable
|
||
# in the bundle from a real download via the presence (or not)
|
||
# of `pp_download_progress` events.
|
||
cache_hit = gguf_bytes > 0 and fetch_ms < 2000
|
||
download_mb_per_s = None if cache_hit else round(fetch_mb / max(fetch_ms / 1000, 1e-3), 1)
|
||
_emit_event(
|
||
"model_fetched",
|
||
stage=stage,
|
||
elapsed_ms=fetch_ms,
|
||
gguf_path=str(gguf_path),
|
||
gguf_bytes=gguf_bytes,
|
||
gguf_mb=round(fetch_mb, 1),
|
||
download_mb_per_s=download_mb_per_s,
|
||
cache_hit=cache_hit,
|
||
)
|
||
print(
|
||
f"pp_tinygrad_worker: stage={stage} loading model from {gguf_path}",
|
||
file=sys.stderr,
|
||
flush=True,
|
||
)
|
||
_emit_event("loading_model", stage=stage, model=model_name, wall_ms=_wall_ms())
|
||
_load_start = time.monotonic()
|
||
# Shard at load time: only this stage's block range (+ embed/output on the
|
||
# end stages) is copied to the compute device, so an 18 GB model fits on a
|
||
# 12 GB GPU. See _load_sharded_transformer for why stock from_gguf can't.
|
||
# The loader emits its own `model_load_breakdown` (header/dequant/realize).
|
||
model, kv, start, end = _load_sharded_transformer(
|
||
gguf_path, stage, num_stages, max_context=512
|
||
)
|
||
load_ms = int((time.monotonic() - _load_start) * 1000)
|
||
tokenizer = SimpleTokenizer.from_gguf_kv(kv)
|
||
|
||
arch = kv["general.architecture"]
|
||
hidden_dim = int(kv[f"{arch}.embedding_length"])
|
||
total_blocks = int(kv[f"{arch}.block_count"])
|
||
vocab_size = len(kv["tokenizer.ggml.tokens"])
|
||
|
||
_emit_event(
|
||
"model_loaded",
|
||
stage=stage,
|
||
elapsed_ms=load_ms,
|
||
rss_mb=_rss_mb(),
|
||
blocks_resident=end - start,
|
||
total_blocks=total_blocks,
|
||
)
|
||
# Stash the cold-start breakdown so main() can emit one `boot_profile`
|
||
# summary once the worker is ready (import + fetch + load + total).
|
||
self.timing = {
|
||
"import_ms": import_ms,
|
||
"fetch_ms": fetch_ms,
|
||
"fetch_cache_hit": cache_hit,
|
||
"gguf_bytes": gguf_bytes,
|
||
"download_mb_per_s": download_mb_per_s,
|
||
"load_ms": load_ms,
|
||
}
|
||
|
||
# Pin EOS token ids (best-effort) for callers that want to detect
|
||
# end-of-text from the sampled stream. We don't enforce stop here
|
||
# — the Rust Stage-1 actor owns the EOS decision — but exposing
|
||
# them in the ready line lets the orchestrator configure itself.
|
||
tokens_list = kv.get("tokenizer.ggml.tokens", [])
|
||
eos_ids: list[int] = []
|
||
for i, tok in enumerate(tokens_list):
|
||
if tok in ("<|end_of_text|>", "<|eot_id|>", "</s>", "<|endoftext|>"):
|
||
eos_ids.append(i)
|
||
|
||
# Tinygrad and numpy are kept as instance attributes so the
|
||
# per-request handlers don't re-import them.
|
||
self._np = np
|
||
self._Tensor = Tensor
|
||
self.model = model
|
||
self.tokenizer = tokenizer
|
||
self.hidden_dim = hidden_dim
|
||
self.vocab_size = vocab_size
|
||
self.total_blocks = total_blocks
|
||
self.stage = stage
|
||
self.num_stages = num_stages
|
||
self.start = start
|
||
self.end = end
|
||
self.eos_ids = eos_ids
|
||
# Per-op compute breakdown (deserialize / compute / host-copy ms), set by
|
||
# the forward ops and folded into the serve loop's `op` timing event.
|
||
self._last_compute: "dict | None" = None
|
||
|
||
print(
|
||
f"pp_tinygrad_worker: stage={stage} ready "
|
||
f"(blocks={total_blocks}, range=[{start},{end}), "
|
||
f"hidden_dim={hidden_dim}, vocab={vocab_size}, eos={eos_ids})",
|
||
file=sys.stderr,
|
||
flush=True,
|
||
)
|
||
|
||
# --- per-stage forward ops -------------------------------------------
|
||
|
||
def embed_and_forward(self, tokens: Sequence[int], position: int) -> tuple[bytes, int]:
|
||
Tensor = self._Tensor
|
||
t0 = time.monotonic()
|
||
t = Tensor([list(tokens)], dtype="int32")
|
||
x = self.model.token_embd(t)
|
||
for block in self.model.blk[self.start : self.end]:
|
||
x = block(x, position)
|
||
# Cast to half (2 bytes/elem) for the wire format, matching the
|
||
# stub. The model's weights are float16; the op output may have
|
||
# promoted to float32, so an explicit cast normalises this. tinygrad is
|
||
# lazy: embed + blocks + cast all *execute* at .realize() below (incl.
|
||
# JIT kernel compile on the first call), so compute_ms captures them.
|
||
t1 = time.monotonic()
|
||
x = x.cast("half").realize()
|
||
t2 = time.monotonic()
|
||
arr = x.numpy() # shape (1, seq_len, hidden_dim), dtype float16
|
||
out = arr.tobytes()
|
||
t3 = time.monotonic()
|
||
self._last_compute = {
|
||
"build_ms": round((t1 - t0) * 1000, 2),
|
||
"compute_ms": round((t2 - t1) * 1000, 2),
|
||
"host_copy_ms": round((t3 - t2) * 1000, 2),
|
||
"out_bytes": len(out),
|
||
}
|
||
return out, int(arr.shape[1])
|
||
|
||
def forward_range(
|
||
self, hidden_bytes: bytes, position: int, seq_len: int
|
||
) -> tuple[bytes, int]:
|
||
"""Run this middle stage's block range over an incoming hidden state.
|
||
|
||
Input is a flat float16 buffer of shape ``(1, seq_len, hidden_dim)``;
|
||
output is the same shape after applying ``model.blk[start:end]``
|
||
with KV-cache ``position``. Returns ``(bytes, out_seq_len)``.
|
||
"""
|
||
np = self._np
|
||
Tensor = self._Tensor
|
||
expected = seq_len * self.hidden_dim * BYTES_PER_ELEM
|
||
if len(hidden_bytes) != expected:
|
||
raise ValueError(
|
||
f"hidden length {len(hidden_bytes)} does not match "
|
||
f"seq_len*hidden_dim*2 ({seq_len}*{self.hidden_dim}*{BYTES_PER_ELEM} "
|
||
f"= {expected})"
|
||
)
|
||
t0 = time.monotonic()
|
||
arr = (
|
||
np.frombuffer(hidden_bytes, dtype=np.float16)
|
||
.reshape((1, seq_len, self.hidden_dim))
|
||
.copy()
|
||
)
|
||
x = Tensor(arr)
|
||
t1 = time.monotonic()
|
||
for block in self.model.blk[self.start : self.end]:
|
||
x = block(x, position)
|
||
x = x.cast("half").realize()
|
||
t2 = time.monotonic()
|
||
out_arr = x.numpy()
|
||
out = out_arr.tobytes()
|
||
t3 = time.monotonic()
|
||
self._last_compute = {
|
||
"deserialize_ms": round((t1 - t0) * 1000, 2),
|
||
"compute_ms": round((t2 - t1) * 1000, 2),
|
||
"host_copy_ms": round((t3 - t2) * 1000, 2),
|
||
"in_bytes": len(hidden_bytes),
|
||
"out_bytes": len(out),
|
||
}
|
||
return out, int(out_arr.shape[1])
|
||
|
||
def forward_and_sample(
|
||
self, hidden_bytes: bytes, position: int, seq_len: int
|
||
) -> int:
|
||
np = self._np
|
||
Tensor = self._Tensor
|
||
expected = seq_len * self.hidden_dim * BYTES_PER_ELEM
|
||
if len(hidden_bytes) != expected:
|
||
raise ValueError(
|
||
f"hidden length {len(hidden_bytes)} does not match "
|
||
f"seq_len*hidden_dim*2 ({seq_len}*{self.hidden_dim}*{BYTES_PER_ELEM} "
|
||
f"= {expected})"
|
||
)
|
||
t0 = time.monotonic()
|
||
arr = (
|
||
np.frombuffer(hidden_bytes, dtype=np.float16)
|
||
.reshape((1, seq_len, self.hidden_dim))
|
||
.copy()
|
||
)
|
||
x = Tensor(arr)
|
||
t1 = time.monotonic()
|
||
for block in self.model.blk[self.start : self.end]:
|
||
x = block(x, position)
|
||
x = self.model.output_norm(x)
|
||
logits = self.model.output(x)
|
||
# Argmax on the last position's logits. Matches what
|
||
# ``Transformer.forward`` does at llm.py:178. The .item() forces the
|
||
# blocks + output projection (over the full vocab) to execute here.
|
||
token_id = int(logits[0, -1, :].argmax().item())
|
||
t2 = time.monotonic()
|
||
self._last_compute = {
|
||
"deserialize_ms": round((t1 - t0) * 1000, 2),
|
||
"compute_ms": round((t2 - t1) * 1000, 2),
|
||
"in_bytes": len(hidden_bytes),
|
||
"token_id": token_id,
|
||
}
|
||
return token_id
|
||
|
||
def generate_full(self, prompt: str, max_tokens: int) -> list[int]:
|
||
"""End-to-end inference over the full block range.
|
||
|
||
Used as the reference for the sliced-vs-full equivalence tests.
|
||
Runs the same unjitted block iteration that the per-stage ops use
|
||
(``embed_and_forward`` + ``forward_and_sample``), but on all
|
||
``model.blk`` blocks within a single process — so the result is
|
||
bit-identical to a correctly-sliced pipeline with argmax sampling.
|
||
|
||
Returns the ``max_tokens`` sampled token ids (excluding the
|
||
prompt). Each invocation starts a fresh prefill at ``position=0``;
|
||
the lazily-allocated per-block KV cache is overwritten in place
|
||
as positions are revisited, so the same worker can serve multiple
|
||
independent prompts.
|
||
"""
|
||
if self.start != 0 or self.end != self.total_blocks:
|
||
raise RuntimeError(
|
||
"generate_full needs the whole model resident, but this stage only "
|
||
f"holds blk[{self.start}:{self.end}) of {self.total_blocks}. Use the "
|
||
"per-stage ops (embed_and_forward / forward_range / forward_and_sample)."
|
||
)
|
||
Tensor = self._Tensor
|
||
if max_tokens <= 0:
|
||
return []
|
||
prompt_tokens = self.tokenizer.encode(prompt)
|
||
if len(prompt_tokens) == 0:
|
||
raise ValueError("prompt tokenized to an empty list")
|
||
|
||
def _forward_sample(token_ids: Sequence[int], position: int) -> int:
|
||
t = Tensor([list(token_ids)], dtype="int32")
|
||
x = self.model.token_embd(t)
|
||
for block in self.model.blk[0 : self.total_blocks]:
|
||
x = block(x, position)
|
||
x = self.model.output_norm(x)
|
||
logits = self.model.output(x)
|
||
return int(logits[0, -1, :].argmax().item())
|
||
|
||
out: list[int] = []
|
||
# Prefill: full prompt at position 0; sample at the last position.
|
||
next_id = _forward_sample(prompt_tokens, 0)
|
||
out.append(next_id)
|
||
# Autoregressive decode: feed the newly-sampled token at the
|
||
# next position. This mirrors what ``Stage0Actor`` does over the
|
||
# wire when ``NextToken { position, token_id }`` arrives.
|
||
pos = len(prompt_tokens)
|
||
for _ in range(max_tokens - 1):
|
||
next_id = _forward_sample([next_id], pos)
|
||
out.append(next_id)
|
||
pos += 1
|
||
return out
|
||
|
||
|
||
def _validate_tokens_list(tokens) -> str | None:
|
||
"""Return None if ``tokens`` is a list of non-bool ints, else an error msg."""
|
||
if not isinstance(tokens, list):
|
||
return "'tokens' must be a list of ints"
|
||
for t in tokens:
|
||
if not isinstance(t, int) or isinstance(t, bool) or t < 0:
|
||
return "'tokens' must be a list of non-negative ints"
|
||
return None
|
||
|
||
|
||
def _handle_request(
|
||
req: dict,
|
||
stage: int,
|
||
num_stages: int,
|
||
real_state: "_RealModelState | None" = None,
|
||
) -> dict:
|
||
rid = req.get("request_id")
|
||
if "op" not in req:
|
||
return {"request_id": rid, "error": "missing 'op' field"}
|
||
op = req["op"]
|
||
is_first = stage == 0
|
||
is_last = stage == num_stages - 1
|
||
is_middle = not is_first and not is_last
|
||
|
||
# Tokenize lives on the first stage; detokenize on the last. Routing
|
||
# both through their natural roles avoids ambiguity when an N-stage
|
||
# cluster has a tokenizer-bearing worker on every node (real mode).
|
||
if op == "tokenize":
|
||
if not is_first:
|
||
return {
|
||
"request_id": rid,
|
||
"error": f"op {op!r} is only valid on stage 0 "
|
||
f"(this worker is stage {stage} of {num_stages})",
|
||
}
|
||
prompt = req.get("prompt")
|
||
if not isinstance(prompt, str):
|
||
return {"request_id": rid, "error": "'prompt' must be a string"}
|
||
if real_state is not None:
|
||
tokens = real_state.tokenizer.encode(prompt)
|
||
else:
|
||
tokens = [
|
||
int(t) for t in _stub_tokenize(prompt)
|
||
]
|
||
return {"request_id": rid, "tokens": tokens}
|
||
|
||
if op == "detokenize":
|
||
if not is_last:
|
||
return {
|
||
"request_id": rid,
|
||
"error": f"op {op!r} is only valid on the final stage "
|
||
f"(this worker is stage {stage} of {num_stages})",
|
||
}
|
||
tokens = req.get("tokens")
|
||
err = _validate_tokens_list(tokens)
|
||
if err is not None:
|
||
return {"request_id": rid, "error": err}
|
||
if real_state is not None:
|
||
text = real_state.tokenizer.decode(tokens)
|
||
else:
|
||
text = " ".join(str(t) for t in tokens)
|
||
return {"request_id": rid, "text": text}
|
||
|
||
if op in ("embed_and_forward", "decode_step"):
|
||
if not is_first:
|
||
return {
|
||
"request_id": rid,
|
||
"error": f"op {op!r} is only valid on stage 0 "
|
||
f"(this worker is stage {stage} of {num_stages})",
|
||
}
|
||
if op == "embed_and_forward":
|
||
tokens = req.get("tokens")
|
||
position = req.get("position", 0)
|
||
err = _validate_tokens_list(tokens)
|
||
if err is not None:
|
||
return {"request_id": rid, "error": err}
|
||
if len(tokens) == 0:
|
||
return {"request_id": rid, "error": "'tokens' must be non-empty"}
|
||
if not _is_nonneg_int(position):
|
||
return {"request_id": rid, "error": "'position' must be a non-negative int"}
|
||
if real_state is not None:
|
||
hidden, out_seq_len = real_state.embed_and_forward(tokens, position)
|
||
else:
|
||
hidden = _stub_hidden_bytes(tokens, position)
|
||
out_seq_len = len(tokens)
|
||
return {
|
||
"request_id": rid,
|
||
"hidden_b64": base64.b64encode(hidden).decode("ascii"),
|
||
"seq_len": out_seq_len,
|
||
}
|
||
# decode_step
|
||
token_id = req.get("token_id")
|
||
position = req.get("position")
|
||
if not _is_nonneg_int(token_id):
|
||
return {"request_id": rid, "error": "'token_id' must be a non-negative int"}
|
||
if not _is_nonneg_int(position):
|
||
return {"request_id": rid, "error": "'position' must be a non-negative int"}
|
||
if real_state is not None:
|
||
hidden, out_seq_len = real_state.embed_and_forward([token_id], position)
|
||
else:
|
||
hidden = _stub_hidden_bytes([token_id], position)
|
||
out_seq_len = 1
|
||
return {
|
||
"request_id": rid,
|
||
"hidden_b64": base64.b64encode(hidden).decode("ascii"),
|
||
"seq_len": out_seq_len,
|
||
}
|
||
|
||
if op == "forward_range":
|
||
if not is_middle:
|
||
return {
|
||
"request_id": rid,
|
||
"error": f"op {op!r} is only valid on a middle stage "
|
||
f"(this worker is stage {stage} of {num_stages})",
|
||
}
|
||
hidden_b64 = req.get("hidden_b64")
|
||
position = req.get("position")
|
||
seq_len = req.get("seq_len")
|
||
if not isinstance(hidden_b64, str):
|
||
return {"request_id": rid, "error": "'hidden_b64' must be a string"}
|
||
if not _is_nonneg_int(position):
|
||
return {"request_id": rid, "error": "'position' must be a non-negative int"}
|
||
if not _is_nonneg_int(seq_len) or seq_len == 0:
|
||
return {"request_id": rid, "error": "'seq_len' must be a positive int"}
|
||
try:
|
||
hidden = base64.b64decode(hidden_b64, validate=True)
|
||
except (base64.binascii.Error, ValueError) as e:
|
||
return {"request_id": rid, "error": f"invalid base64 in hidden_b64: {e}"}
|
||
hidden_dim = real_state.hidden_dim if real_state is not None else STUB_HIDDEN_DIM
|
||
expected = seq_len * hidden_dim * BYTES_PER_ELEM
|
||
if len(hidden) != expected:
|
||
return {
|
||
"request_id": rid,
|
||
"error": (
|
||
f"hidden length {len(hidden)} does not match "
|
||
f"seq_len*hidden_dim*2 ({seq_len}*{hidden_dim}*{BYTES_PER_ELEM} "
|
||
f"= {expected})"
|
||
),
|
||
}
|
||
if real_state is not None:
|
||
hidden_out, out_seq_len = real_state.forward_range(hidden, position, seq_len)
|
||
else:
|
||
hidden_out = _stub_forward_range_bytes(hidden, position, seq_len)
|
||
out_seq_len = seq_len
|
||
return {
|
||
"request_id": rid,
|
||
"hidden_b64": base64.b64encode(hidden_out).decode("ascii"),
|
||
"seq_len": out_seq_len,
|
||
}
|
||
|
||
if op == "generate_full":
|
||
# Reference path used by the sliced-vs-full equivalence tests. Always
|
||
# runs over the full block range, so the worker's STAGE/NUM_STAGES
|
||
# are ignored here — any ``NUM_STAGES >= 2`` works since the op
|
||
# iterates ``model.blk`` directly. The example does not boot at
|
||
# ``NUM_STAGES=1`` (single-node configurations belong to
|
||
# ``examples/single-gpu-inference``).
|
||
if real_state is None:
|
||
return {
|
||
"request_id": rid,
|
||
"error": "op 'generate_full' requires real mode (stub mode has no full model)",
|
||
}
|
||
prompt = req.get("prompt")
|
||
max_tokens = req.get("max_tokens")
|
||
if not isinstance(prompt, str):
|
||
return {"request_id": rid, "error": "'prompt' must be a string"}
|
||
if not _is_nonneg_int(max_tokens) or max_tokens == 0:
|
||
return {"request_id": rid, "error": "'max_tokens' must be a positive int"}
|
||
try:
|
||
tokens = real_state.generate_full(prompt, max_tokens)
|
||
except Exception as e:
|
||
return {"request_id": rid, "error": f"generate_full: {e}"}
|
||
return {"request_id": rid, "tokens": tokens}
|
||
|
||
if op == "forward_and_sample":
|
||
if not is_last:
|
||
return {
|
||
"request_id": rid,
|
||
"error": f"op {op!r} is only valid on the final stage "
|
||
f"(this worker is stage {stage} of {num_stages})",
|
||
}
|
||
hidden_b64 = req.get("hidden_b64")
|
||
position = req.get("position")
|
||
seq_len = req.get("seq_len")
|
||
if not isinstance(hidden_b64, str):
|
||
return {"request_id": rid, "error": "'hidden_b64' must be a string"}
|
||
if not _is_nonneg_int(position):
|
||
return {"request_id": rid, "error": "'position' must be a non-negative int"}
|
||
if not _is_nonneg_int(seq_len) or seq_len == 0:
|
||
return {"request_id": rid, "error": "'seq_len' must be a positive int"}
|
||
try:
|
||
hidden = base64.b64decode(hidden_b64, validate=True)
|
||
except (base64.binascii.Error, ValueError) as e:
|
||
return {"request_id": rid, "error": f"invalid base64 in hidden_b64: {e}"}
|
||
if real_state is not None:
|
||
hidden_dim = real_state.hidden_dim
|
||
else:
|
||
hidden_dim = STUB_HIDDEN_DIM
|
||
expected = seq_len * hidden_dim * BYTES_PER_ELEM
|
||
if len(hidden) != expected:
|
||
return {
|
||
"request_id": rid,
|
||
"error": (
|
||
f"hidden length {len(hidden)} does not match "
|
||
f"seq_len*hidden_dim*2 ({seq_len}*{hidden_dim}*{BYTES_PER_ELEM} "
|
||
f"= {expected})"
|
||
),
|
||
}
|
||
if real_state is not None:
|
||
token_id = real_state.forward_and_sample(hidden, position, seq_len)
|
||
else:
|
||
token_id = _stub_token_id(hidden, position)
|
||
return {"request_id": rid, "token_id": token_id}
|
||
|
||
return {"request_id": rid, "error": f"unknown op {op!r}"}
|
||
|
||
|
||
def _stub_tokenize(prompt: str) -> list[int]:
|
||
"""Whitespace-split a prompt into deterministic small integer ids.
|
||
|
||
Only used in stub mode. The exact mapping is not part of the
|
||
worker's contract — callers just need a list of ints whose length
|
||
equals the number of whitespace-separated words.
|
||
"""
|
||
out: list[int] = []
|
||
for i, word in enumerate(prompt.split()):
|
||
s = sum(ord(c) for c in word)
|
||
out.append((s % 1024) + i)
|
||
return out
|
||
|
||
|
||
def main(argv: Sequence[str] | None = None) -> int:
|
||
# Install diagnostic hooks first thing so any failure during arg
|
||
# parsing or env validation still produces a structured event.
|
||
_install_excepthook()
|
||
_install_signal_handlers()
|
||
|
||
parser = argparse.ArgumentParser(description="pipeline-parallel tinygrad worker")
|
||
parser.add_argument(
|
||
"--stub",
|
||
action="store_true",
|
||
help="Stub mode: skip model loading, use deterministic in-memory ops",
|
||
)
|
||
parser.add_argument(
|
||
"--model",
|
||
default=None,
|
||
help="Model name (defaults to $MODEL or llama3.2:1b in real mode)",
|
||
)
|
||
args = parser.parse_args(argv)
|
||
|
||
stub_mode = args.stub or os.environ.get("PP_WORKER_STUB", "").strip() == "1"
|
||
|
||
stage = _parse_env_int("STAGE")
|
||
num_stages = _parse_env_int("NUM_STAGES")
|
||
|
||
_emit_event(
|
||
"starting",
|
||
pid=os.getpid(),
|
||
stage=stage,
|
||
num_stages=num_stages,
|
||
stub=stub_mode,
|
||
model=(args.model or os.environ.get("MODEL", "")).strip() or None,
|
||
python_version=sys.version.split()[0],
|
||
argv=list(sys.argv),
|
||
# wall_ms anchors this worker's boot against the orchestrator's vast.ai
|
||
# instance create/lease time — the only way to measure image-pull +
|
||
# container-boot latency, which the worker can't observe directly.
|
||
wall_ms=_wall_ms(),
|
||
host=os.uname().nodename,
|
||
)
|
||
_start_heartbeat()
|
||
|
||
if num_stages < 2:
|
||
_die(
|
||
f"NUM_STAGES must be >= 2 (single-node configurations are not "
|
||
f"served by this example), got {num_stages}"
|
||
)
|
||
if not (0 <= stage < num_stages):
|
||
_die(f"STAGE {stage} out of range [0, {num_stages})")
|
||
|
||
real_state: _RealModelState | None = None
|
||
if not stub_mode:
|
||
model_name = (args.model or os.environ.get("MODEL", "")).strip() or "llama3.2:1b"
|
||
try:
|
||
real_state = _RealModelState(model_name, stage, num_stages)
|
||
except SystemExit:
|
||
raise
|
||
except Exception as e:
|
||
import traceback
|
||
|
||
tb_text = traceback.format_exc()
|
||
_emit_event(
|
||
"model_load_failed",
|
||
stage=stage,
|
||
model=model_name,
|
||
type=type(e).__name__,
|
||
value=str(e),
|
||
traceback=tb_text,
|
||
)
|
||
print(tb_text, file=sys.stderr, flush=True)
|
||
_die(f"failed to load model {model_name!r}: {e}")
|
||
|
||
ready: dict = {"status": "ready", "pid": os.getpid(), "stage": stage}
|
||
if real_state is not None:
|
||
ready["hidden_dim"] = real_state.hidden_dim
|
||
ready["vocab_size"] = real_state.vocab_size
|
||
ready["total_blocks"] = real_state.total_blocks
|
||
ready["layer_range"] = [real_state.start, real_state.end]
|
||
ready["eos_token_ids"] = real_state.eos_ids
|
||
_write(ready)
|
||
# Mirror ready as a structured event under the spec's event name
|
||
# (§4.6 / §6.9: per-stage filter on `pp_worker_ready`). The
|
||
# `pp_*` event kind is passed through unchanged by the Rust actor
|
||
# so the bundle records `Custom("pp_worker_ready")`, matching the
|
||
# name the orchestrator-side wired check filters on. The protocol
|
||
# `status: "ready"` line above is unchanged for the actor's ready
|
||
# signal.
|
||
_emit_event(
|
||
"pp_worker_ready",
|
||
pid=os.getpid(),
|
||
stage_index=stage,
|
||
uptime_ms=_uptime_ms(),
|
||
rss_mb=_rss_mb(),
|
||
)
|
||
# One-stop cold-start breakdown so a single event answers "where did
|
||
# bring-up time go" per node: import + (fetch|cache) + load == time-to-ready.
|
||
if real_state is not None:
|
||
_emit_event(
|
||
"boot_profile",
|
||
stage=stage,
|
||
total_to_ready_ms=_uptime_ms(),
|
||
rss_mb=_rss_mb(),
|
||
blocks_resident=real_state.end - real_state.start,
|
||
**real_state.timing,
|
||
)
|
||
|
||
global _REQUESTS_SERVED
|
||
seen_ops: set = set()
|
||
last_op_end = time.monotonic()
|
||
for line in sys.stdin:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
req = json.loads(line)
|
||
except (json.JSONDecodeError, ValueError) as e:
|
||
_write({"error": f"invalid JSON: {e}"})
|
||
continue
|
||
if not isinstance(req, dict):
|
||
_write({"error": f"request must be a JSON object, got {type(req).__name__}"})
|
||
continue
|
||
op = req.get("op")
|
||
if real_state is not None:
|
||
real_state._last_compute = None
|
||
t_start = time.monotonic()
|
||
# idle_ms_before is the pipeline bubble: how long this worker sat
|
||
# blocked on its upstream stage between finishing the last op and
|
||
# receiving this one. High idle => the bottleneck is elsewhere.
|
||
idle_ms = round((t_start - last_op_end) * 1000, 2)
|
||
try:
|
||
reply = _handle_request(req, stage, num_stages, real_state=real_state)
|
||
except Exception as e: # last-ditch safety net so the worker stays up
|
||
import traceback
|
||
|
||
print(traceback.format_exc(), file=sys.stderr, flush=True)
|
||
reply = {"request_id": req.get("request_id"), "error": f"internal: {e}"}
|
||
_write(reply)
|
||
t_end = time.monotonic()
|
||
# Per-op trace — the granular signal for end-to-end latency. `first_call`
|
||
# flags the JIT-compile-bearing first invocation of each op (kernels are
|
||
# compiled once, then cached). `rid` (not `request_id`) keeps the Rust
|
||
# actor from ever mistaking this event line for an op reply.
|
||
is_first = op not in seen_ops
|
||
seen_ops.add(op)
|
||
_emit_event(
|
||
"op",
|
||
op=op,
|
||
rid=req.get("request_id"),
|
||
stage=stage,
|
||
duration_ms=round((t_end - t_start) * 1000, 2),
|
||
idle_ms_before=idle_ms,
|
||
first_call=is_first,
|
||
ok=isinstance(reply, dict) and "error" not in reply,
|
||
in_tokens=len(req["tokens"]) if isinstance(req.get("tokens"), list) else None,
|
||
in_seq_len=req.get("seq_len"),
|
||
out_seq_len=reply.get("seq_len") if isinstance(reply, dict) else None,
|
||
compute=(real_state._last_compute if real_state is not None else None),
|
||
uptime_ms=_uptime_ms(),
|
||
)
|
||
last_op_end = t_end
|
||
_REQUESTS_SERVED += 1
|
||
|
||
_emit_event(
|
||
"exiting",
|
||
reason="eof",
|
||
uptime_ms=_uptime_ms(),
|
||
requests_served=_REQUESTS_SERVED,
|
||
)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|