feat: Working vastai single-node deployment for LLM inference
This commit is contained in:
parent
2c3c05ed85
commit
cb789c3ef1
24 changed files with 8750 additions and 1 deletions
|
|
@ -1,2 +1,4 @@
|
|||
*
|
||||
!target/x86_64-unknown-linux-musl/release/swactor
|
||||
!examples/single-gpu-inference/target/release/gpu-node
|
||||
!examples/single-gpu-inference/tinygrad_worker.py
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ members = [
|
|||
"tests/integration",
|
||||
"xtask",
|
||||
]
|
||||
exclude = ["crates/bindings/wasm-crypto"]
|
||||
exclude = ["crates/bindings/wasm-crypto", "examples"]
|
||||
|
||||
[package]
|
||||
name = "swactor"
|
||||
|
|
|
|||
4949
examples/single-gpu-inference/Cargo.lock
generated
Normal file
4949
examples/single-gpu-inference/Cargo.lock
generated
Normal file
File diff suppressed because it is too large
Load diff
29
examples/single-gpu-inference/Cargo.toml
Normal file
29
examples/single-gpu-inference/Cargo.toml
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
[workspace]
|
||||
|
||||
[package]
|
||||
name = "single-gpu-inference"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
publish = false
|
||||
|
||||
[dependencies]
|
||||
swactor = { path = "../..", features = ["transport", "serde", "std"] }
|
||||
swactor-process = { path = "../../crates/process" }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
reqwest = { version = "0.12", features = ["json"] }
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
distribution = { path = "../../crates/distribution", features = ["iroh"] }
|
||||
iroh = "0.96"
|
||||
urlencoding = "2"
|
||||
|
||||
[[bin]]
|
||||
name = "gpu-node"
|
||||
path = "src/bin/gpu_node.rs"
|
||||
|
||||
[[bin]]
|
||||
name = "single-gpu-inference"
|
||||
path = "src/bin/single_gpu_inference.rs"
|
||||
|
||||
[dev-dependencies]
|
||||
wiremock = "0.6"
|
||||
27
examples/single-gpu-inference/Dockerfile
Normal file
27
examples/single-gpu-inference/Dockerfile
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
FROM nvidia/cuda:12.6.3-devel-ubuntu24.04
|
||||
|
||||
# Install Python 3, pip, and CUDA runtime compiler (tinygrad compiles kernels via NVRTC)
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
python3 \
|
||||
python3-venv \
|
||||
python3-pip \
|
||||
ca-certificates && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install tinygrad and numpy
|
||||
RUN python3 -m pip install --no-cache-dir --break-system-packages \
|
||||
tinygrad==0.12.0 \
|
||||
numpy
|
||||
|
||||
# Copy the gpu-node binary and tinygrad worker
|
||||
# Build context should be the workspace root:
|
||||
# docker build -f examples/single-gpu-inference/Dockerfile -t <tag> .
|
||||
COPY examples/single-gpu-inference/target/release/gpu-node /usr/local/bin/gpu-node
|
||||
COPY examples/single-gpu-inference/tinygrad_worker.py /usr/local/share/tinygrad_worker.py
|
||||
|
||||
# Enable CUDA backend for tinygrad
|
||||
ENV CUDA=1
|
||||
ENV WORKER_SCRIPT=/usr/local/share/tinygrad_worker.py
|
||||
|
||||
CMD ["gpu-node"]
|
||||
78
examples/single-gpu-inference/INTEGRATION_TEST_PLAN.md
Normal file
78
examples/single-gpu-inference/INTEGRATION_TEST_PLAN.md
Normal file
|
|
@ -0,0 +1,78 @@
|
|||
# Plan: Local Mock Integration Test
|
||||
|
||||
## Context
|
||||
|
||||
The smoke-test crate has 5 component test groups (T-codec, T-worker, T-vastai, T-actor, T-cluster) that each test a piece of the distributed inference pipeline in isolation. What's missing is a single test that wires them together: two swactor nodes on localhost, one running the real `InferenceActor` (with `echo_worker.py`), communicating over iroh/QUIC. This proves the full local chain before spending money on vast.ai.
|
||||
|
||||
## Design Problem
|
||||
|
||||
`InferenceActor::Incoming` is `InferenceActorMsg` (a union of `Request` and `Process` variants). But the network codec delivers raw `InferenceRequest`. When `rt.deliver_raw()` delivers a deserialized `InferenceRequest` to the `InferenceActor`, the downcast to `InferenceActorMsg` fails silently.
|
||||
|
||||
**Solution:** Add a `RequestBridge` actor — same pattern as the existing `ProcessBridge`. It receives `InferenceRequest` from the network, wraps it as `InferenceActorMsg::Request(req)`, and forwards to the `InferenceActor`. ~10 lines.
|
||||
|
||||
## Changes
|
||||
|
||||
### 1. Add `RequestBridge` to `examples/single-gpu-inference/src/inference_actor.rs`
|
||||
|
||||
A public actor struct placed after the existing `ProcessBridge` (~line 55). Fields: `target: ActorAddress`. Implements `ActorInterface` with `Incoming = InferenceRequest`, wraps and forwards to target as `InferenceActorMsg::Request`.
|
||||
|
||||
### 2. Export it from `examples/single-gpu-inference/src/lib.rs`
|
||||
|
||||
Already exports `pub mod inference_actor` — `RequestBridge` just needs to be `pub`.
|
||||
|
||||
### 3. Create `examples/single-gpu-inference/tests/t_integration.rs`
|
||||
|
||||
One test function: `distributed_inference_through_echo_worker`.
|
||||
|
||||
**Setup (reuse patterns from t_cluster.rs and t_actor.rs):**
|
||||
- Copy iroh helpers: `make_driver`, `make_converged_pair`, `IrohActorTransport`, `encode_wire`/`decode_wire`, `drain_actor_messages` (with reduced internal sleep for localhost — 100ms instead of 500ms)
|
||||
- Copy process helpers: `echo_worker_spec`, `is_process_alive`
|
||||
|
||||
**Test flow:**
|
||||
1. Converge two iroh drivers via `make_converged_pair()`
|
||||
2. Create `rt_a` (local) and `rt_b` (remote) runtimes
|
||||
3. On `rt_b`: spawn `InferenceActor` (with `echo_worker_spec()`) + `RequestBridge` pointing at it
|
||||
4. On `rt_a`: create `response_inbox` for `InferenceResponse`
|
||||
5. Build `IrohActorTransport` in each direction, wire transport routers:
|
||||
- `rt_a`: `bridge_addr → transport_a_to_b`
|
||||
- `rt_b`: `inbox_addr → transport_b_to_a`
|
||||
6. Install codec registries and transport routers (`&mut self` — must happen after all spawns/inbox creation)
|
||||
7. Tick `rt_b` in a polling loop until `InferenceActorStatus::WorkerReady` (echo_worker.py started)
|
||||
8. `rt_a.send_to(bridge_addr, InferenceRequest { prompt: "Hello from node A", reply_to: inbox_addr, ... })`
|
||||
9. Pump loop (10s timeout): `drain_actor_messages` on both drivers → tick both runtimes → check `response_inbox`
|
||||
10. Assert response text contains `"Hello from node A"` (echo worker reflects prompt)
|
||||
11. Cleanup: stop inference actor, tick until echo_worker.py pid is dead, shutdown both drivers
|
||||
|
||||
**Message path through the system:**
|
||||
```
|
||||
rt_a.send_to(bridge_addr, InferenceRequest)
|
||||
→ transport router → IrohActorTransport (QUIC to node B)
|
||||
→ drain_actor_messages → rt_b.deliver_raw(bridge_addr, InferenceRequest)
|
||||
→ RequestBridge.handle() → ctx.send(inference_addr, InferenceActorMsg::Request(req))
|
||||
→ InferenceActor.handle() → writes JSON to echo_worker.py stdin
|
||||
→ echo_worker.py → writes JSON to stdout
|
||||
→ ProcessActor → ProcessNotification::Output → ProcessBridge → InferenceActor
|
||||
→ InferenceActor.process_output_line() → ctx.send(reply_to, InferenceResponse)
|
||||
→ transport router → IrohActorTransport (QUIC to node A)
|
||||
→ drain_actor_messages → rt_a.deliver_raw(inbox_addr, InferenceResponse)
|
||||
→ response_inbox.try_recv() ✓
|
||||
```
|
||||
|
||||
## Files Modified
|
||||
|
||||
| File | Change |
|
||||
|---|---|
|
||||
| `examples/single-gpu-inference/src/inference_actor.rs` | Add `pub struct RequestBridge` (~10 lines) |
|
||||
| `examples/single-gpu-inference/tests/t_integration.rs` | New file — one integration test (~200 lines) |
|
||||
|
||||
## Verification
|
||||
|
||||
```bash
|
||||
# Run just the new test
|
||||
cargo test --manifest-path examples/single-gpu-inference/Cargo.toml t_integration
|
||||
|
||||
# Confirm existing tests still pass
|
||||
cargo test --manifest-path examples/single-gpu-inference/Cargo.toml
|
||||
```
|
||||
|
||||
Expected: all 22 tests pass (21 existing + 1 new). Test runtime ~10-15s (dominated by iroh drain sleeps and echo_worker.py process I/O).
|
||||
13
examples/single-gpu-inference/Makefile
Normal file
13
examples/single-gpu-inference/Makefile
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
.PHONY: venv clean-venv
|
||||
|
||||
VENV_DIR := .venv
|
||||
PYTHON := $(VENV_DIR)/bin/python
|
||||
|
||||
venv: $(VENV_DIR)/bin/python
|
||||
|
||||
$(VENV_DIR)/bin/python:
|
||||
uv venv $(VENV_DIR) --python python3.11
|
||||
uv pip install --python $(VENV_DIR)/bin/python "tinygrad==0.12.0" numpy
|
||||
|
||||
clean-venv:
|
||||
rm -rf $(VENV_DIR)
|
||||
243
examples/single-gpu-inference/SPEC.md
Normal file
243
examples/single-gpu-inference/SPEC.md
Normal file
|
|
@ -0,0 +1,243 @@
|
|||
# Smoke Test Specification: Rent a GPU, Run Inference, Say Hello
|
||||
|
||||
## 1. Wanted Behavior
|
||||
|
||||
`smoke-run --vastai --api-key <key>` runs on the local machine. It:
|
||||
1. Finds the cheapest available GPU on vast.ai and rents it.
|
||||
2. Starts a local swactor node with iroh transport.
|
||||
3. The rented instance boots our pre-built Docker image containing `gpu-node` + tinygrad.
|
||||
4. The remote `gpu-node` reads `SEED_ADDR` from env and joins the local cluster via iroh/QUIC.
|
||||
5. The local node discovers the remote `InferenceActor` via SWIM name resolution (`"inference"`).
|
||||
6. The local node sends an `InferenceRequest` to the remote actor.
|
||||
7. The remote actor runs tinygrad inference, replies with `InferenceResponse`.
|
||||
8. Asserts the response is non-empty text and prints it.
|
||||
9. Destroys the vast.ai instance.
|
||||
10. Exits 0.
|
||||
|
||||
If anything fails, the script destroys the instance (if one was created) and exits 1.
|
||||
|
||||
Total wall-clock budget: 10 minutes.
|
||||
|
||||
```
|
||||
local machine (smoke-run) vast.ai GPU (gpu-node)
|
||||
| |
|
||||
| 1. search offers (REST) |
|
||||
| 2. create instance (REST) |
|
||||
| SEED_ADDR=<local node id> |
|
||||
| ------------------------------------> |
|
||||
| | [pulling docker image]
|
||||
| 3. poll vast.ai status every 10s | [gpu-node starting]
|
||||
| ------------------------------------> | [tinygrad loading model into VRAM]
|
||||
| "loading" |
|
||||
| <------------------------------------ |
|
||||
| "running" |
|
||||
| <------------------------------------ |
|
||||
| |
|
||||
| 4. SWIM cluster join (iroh/QUIC) |
|
||||
| <--------------------------------->>> |
|
||||
| [cluster converged] |
|
||||
| [name "inference" resolved] |
|
||||
| |
|
||||
| 5. InferenceRequest (swactor msg) |
|
||||
| "Say hello" |
|
||||
| ------------------------------------> |
|
||||
| | [tinygrad forward pass]
|
||||
| InferenceResponse |
|
||||
| "Hello! How can I help you?" |
|
||||
| <------------------------------------ |
|
||||
| |
|
||||
| 6. assert len(response) > 0 |
|
||||
| 7. DELETE instance (vast.ai REST) |
|
||||
| ------------------------------------> |
|
||||
| 8. exit 0 |
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. What Already Exists
|
||||
|
||||
The entire pipeline is implemented and tested locally. The remaining work is deploying it to vast.ai for real.
|
||||
|
||||
### Implemented components
|
||||
|
||||
| Component | File(s) | Status |
|
||||
|---|---|---|
|
||||
| Message types + codec | `src/messages.rs` | Done, tested (T-codec: 7 tests) |
|
||||
| InferenceActor + ProcessBridge | `src/inference_actor.rs` | Done, tested (T-actor: 4 tests) |
|
||||
| RequestBridge (network→actor type bridge) | `src/inference_actor.rs` | Done |
|
||||
| Iroh actor transport (shared) | `src/iroh_transport.rs` | Done, used by tests + binaries |
|
||||
| vast.ai REST client | `src/vastai.rs` | Done, tested (T-vastai: 7 tests) |
|
||||
| tinygrad worker (real GGUF model) | `tinygrad_worker.py` | Done, tested (T-worker: 5 test classes) |
|
||||
| Cluster transport tests | `tests/t_cluster.rs` | Done (3 tests) |
|
||||
| In-process integration test (echo) | `tests/t_integration.rs` | Done, passes |
|
||||
| In-process integration test (tinygrad) | `tests/t_integration.rs` | Done, passes (`#[ignore]`, needs .venv) |
|
||||
| `gpu-node` binary | `src/bin/gpu_node.rs` | Done, compiles |
|
||||
| `smoke-run` binary (localhost + vastai) | `src/bin/smoke_run.rs` | Done, compiles |
|
||||
| Binary e2e test (echo) | `tests/t_binary.rs` | Done, passes |
|
||||
| Binary e2e test (tinygrad) | `tests/t_binary.rs` | Done, passes (`#[ignore]`, needs .venv) |
|
||||
| Dockerfile | `Dockerfile` | Done |
|
||||
|
||||
### Test commands
|
||||
|
||||
```bash
|
||||
# Fast tests (no GPU, no downloads, no vast.ai) — 22 tests
|
||||
cargo test --manifest-path examples/single-gpu-inference/Cargo.toml
|
||||
|
||||
# Slow tests (downloads ~1GB GGUF model, runs tinygrad on CPU)
|
||||
cargo test --manifest-path examples/single-gpu-inference/Cargo.toml -- --ignored
|
||||
|
||||
# Binary e2e on localhost (spawns gpu-node + smoke-run as child processes)
|
||||
cargo test --manifest-path examples/single-gpu-inference/Cargo.toml binary_e2e_echo_worker
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. What Remains — Deployment to vast.ai
|
||||
|
||||
### 3.1 Fix relay mode for WAN
|
||||
|
||||
Both binaries currently use `RelayMode::Disabled`, which works on localhost but not over WAN. The remote `gpu-node` behind a vast.ai NAT cannot reach the local node without iroh relay servers.
|
||||
|
||||
**Changes needed:**
|
||||
- `gpu-node`: change `RelayMode::Disabled` to `RelayMode::Default` so iroh uses its public relay infrastructure for NAT traversal.
|
||||
- `smoke-run` (vastai path): same — use `RelayMode::Default`.
|
||||
- The localhost path (`smoke-run --seed`) can keep `RelayMode::Disabled`.
|
||||
|
||||
### 3.2 Fix worker script path in Dockerfile
|
||||
|
||||
The `Dockerfile` copies `tinygrad_worker.py` to `/usr/local/share/tinygrad_worker.py`, but `gpu-node` defaults `WORKER_SCRIPT` to `./tinygrad_worker.py`.
|
||||
|
||||
**Fix**: either change the Dockerfile `COPY` destination to `/app/tinygrad_worker.py` and set `WORKDIR /app`, or set `ENV WORKER_SCRIPT=/usr/local/share/tinygrad_worker.py` in the Dockerfile.
|
||||
|
||||
### 3.3 Add `pid` to tinygrad worker ready signal
|
||||
|
||||
`tinygrad_worker.py` emits `{"status": "ready"}` but `echo_worker.py` emits `{"status": "ready", "pid": <pid>}`. The `InferenceActor` parses the `pid` field for `WorkerReady { pid }` status reporting. Without it, `worker_pid` is `None` — not fatal, but cleanup assertions in tests rely on it.
|
||||
|
||||
**Fix**: change the ready signal to `_write({"status": "ready", "pid": os.getpid()})`.
|
||||
|
||||
### 3.4 Build and push the Docker image
|
||||
|
||||
```bash
|
||||
# Cross-compile gpu-node for linux/amd64 (if not already on linux/amd64)
|
||||
cargo build --release --bin gpu-node --manifest-path examples/single-gpu-inference/Cargo.toml
|
||||
|
||||
# Build image
|
||||
docker build -t <your-registry>/swactor-gpu:latest -f examples/single-gpu-inference/Dockerfile .
|
||||
|
||||
# Push to registry (vast.ai pulls from here)
|
||||
docker push <your-registry>/swactor-gpu:latest
|
||||
```
|
||||
|
||||
The image name in `vastai.rs::create_instance` is hardcoded to `"swactor-gpu:latest"`. Update this to match whatever registry you push to, or make it a parameter.
|
||||
|
||||
### 3.5 Run it for real
|
||||
|
||||
```bash
|
||||
VAST_API_KEY=<key> cargo run --manifest-path examples/single-gpu-inference/Cargo.toml --bin single-gpu-inference -- \
|
||||
--vastai --api-key <key> --gpu RTX_4090
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 4. Architecture
|
||||
|
||||
### Message flow (same for localhost and vast.ai)
|
||||
|
||||
```
|
||||
smoke-run gpu-node
|
||||
───────── ────────
|
||||
rt.send_to(bridge_addr, InferenceRequest)
|
||||
→ codec encodes → TransportRouter
|
||||
→ IrohActorTransport (QUIC to gpu-node)
|
||||
drain_and_collect_reply_addrs()
|
||||
→ decode_wire → codecs.receive()
|
||||
→ rt.deliver_raw(bridge_addr, InferenceRequest)
|
||||
→ extract reply_to, add return transport route
|
||||
RequestBridge.handle()
|
||||
→ ctx.send(inference_addr, InferenceActorMsg::Request)
|
||||
InferenceActor.handle()
|
||||
→ JSON to tinygrad_worker.py stdin
|
||||
tinygrad_worker.py
|
||||
→ Transformer forward pass
|
||||
→ JSON response to stdout
|
||||
InferenceActor.process_output_line()
|
||||
→ ctx.send(reply_to, InferenceResponse)
|
||||
→ TransportRouter → IrohActorTransport (QUIC back)
|
||||
drain_actor_messages()
|
||||
→ decode_wire → codecs.receive()
|
||||
→ rt.deliver_raw(inbox_addr, InferenceResponse)
|
||||
response_inbox.try_recv() ✓
|
||||
```
|
||||
|
||||
### Name discovery
|
||||
|
||||
The `gpu-node` registers its `RequestBridge` under the name `"inference"` via `DistributedNode::register_name`. This propagates through SWIM gossip piggyback. The `smoke-run` orchestrator calls `driver.node().resolve_name("inference")` to discover the bridge's `ActorAddress` without needing to know it ahead of time.
|
||||
|
||||
### Dynamic return routing
|
||||
|
||||
The `gpu-node`'s drain loop inspects incoming `InferenceRequest` payloads to extract the `reply_to` address. It then dynamically registers a transport route for that address pointing back to the only alive SWIM member (the orchestrator). This is necessary because the `gpu-node` doesn't know the orchestrator's inbox address at startup.
|
||||
|
||||
---
|
||||
|
||||
## 5. tinygrad Worker
|
||||
|
||||
`tinygrad_worker.py` — managed by swactor's process crate via stdin/stdout JSON.
|
||||
|
||||
**Model**: `llama3.2:1b` from tinygrad's built-in GGUF catalog (~1GB download, fits in any modern GPU's VRAM). Loaded via `tinygrad.apps.llm.Transformer.from_gguf()`.
|
||||
|
||||
**Protocol**:
|
||||
```
|
||||
← stdout: {"status": "ready", "pid": 12345}
|
||||
→ stdin: {"prompt": "Say hello", "max_tokens": 64, "temperature": 0.7}
|
||||
← stdout: {"response": "Hello! How can I help you today?"}
|
||||
→ stdin: {invalid json}
|
||||
← stdout: {"error": "invalid JSON: ..."}
|
||||
```
|
||||
|
||||
**Modes**:
|
||||
- Default: downloads and loads real GGUF model, runs real inference.
|
||||
- `--stub`: canned responses, no tinygrad import (for fast protocol tests).
|
||||
- `--model <name>`: override model (e.g., `qwen3:0.6b` for smaller download).
|
||||
|
||||
**Environment**:
|
||||
- `CUDA=1` → tinygrad uses CUDA backend (GPU).
|
||||
- No env var → tinygrad auto-detects (CUDA if available, else CPU).
|
||||
- `PYTHON=1` → forces tinygrad's pure-Python CPU backend (no clang needed).
|
||||
|
||||
---
|
||||
|
||||
## 6. Error Modes
|
||||
|
||||
| Failure | Detection | Response |
|
||||
|---------|-----------|----------|
|
||||
| No GPU offers available | vast.ai returns empty list | Print error, exit 1. No cleanup needed. |
|
||||
| Instance creation rejected | API returns non-success | Print error, exit 1. No cleanup needed. |
|
||||
| Instance never reaches `running` | 60 polls exhausted or terminal status | Destroy instance, exit 1. |
|
||||
| Cluster never converges | SWIM timeout (2 min) | Destroy instance, exit 1. |
|
||||
| Name `"inference"` never resolves | Timeout (1 min) | Destroy instance, exit 1. |
|
||||
| tinygrad worker never reports ready | Process timeout (10 min) | Destroy instance, exit 1. |
|
||||
| Empty/malformed response | Assertion on response text | Destroy instance, exit 1. |
|
||||
| Destroy fails | Catch around cleanup | Print warning, exit 1. |
|
||||
|
||||
Every path that allocates an instance also destroys it.
|
||||
|
||||
---
|
||||
|
||||
## 7. Environment Requirements
|
||||
|
||||
- Rust toolchain
|
||||
- Python 3.10+ with tinygrad 0.12.0 + numpy (in `.venv/`)
|
||||
- Docker (for building the GPU node image)
|
||||
- `VAST_API_KEY` (vast.ai deployment only — not needed for local tests)
|
||||
- Network access to iroh relay servers (vast.ai deployment only)
|
||||
|
||||
---
|
||||
|
||||
## 8. What This Proves
|
||||
|
||||
1. swactor nodes on rented GPUs can join a cluster with a local node over WAN via iroh/QUIC.
|
||||
2. swactor messages traverse the WAN transparently — no HTTP layer needed for node-to-node communication.
|
||||
3. tinygrad can load and run a LLaMA model on rented hardware via its CUDA backend.
|
||||
4. The process crate manages the Python child process lifecycle cleanly.
|
||||
5. SWIM gossip propagates actor name registrations, enabling dynamic service discovery.
|
||||
6. The full lifecycle (rent → cluster → discover → infer → teardown) is fully automated.
|
||||
16
examples/single-gpu-inference/echo_worker.py
Normal file
16
examples/single-gpu-inference/echo_worker.py
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Stub worker for testing. Same JSON protocol as tinygrad_worker.py."""
|
||||
import json, os, sys
|
||||
|
||||
print(json.dumps({"status": "ready", "pid": os.getpid()}), flush=True)
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
req = json.loads(line)
|
||||
if req.get("prompt") == "__crash__":
|
||||
os._exit(1)
|
||||
print(json.dumps({"response": f"echo: {req['prompt']}"}), flush=True)
|
||||
except Exception as e:
|
||||
print(json.dumps({"error": str(e)}), flush=True)
|
||||
310
examples/single-gpu-inference/src/bin/gpu_node.rs
Normal file
310
examples/single-gpu-inference/src/bin/gpu_node.rs
Normal file
|
|
@ -0,0 +1,310 @@
|
|||
//! gpu-node — GPU inference node for distributed tinygrad inference.
|
||||
//!
|
||||
//! Runs inside a Docker container (or locally for testing). Joins an existing
|
||||
//! cluster via `SEED_ADDR`, spawns an `InferenceActor` backed by
|
||||
//! `tinygrad_worker.py`, and registers the bridge under the name `"inference"`
|
||||
//! so the orchestrator can discover it.
|
||||
//!
|
||||
//! Environment variables:
|
||||
//! - `SEED_ADDR` (required): hex-encoded NodeId of the seed node to join
|
||||
//! - `SEED_DIRECT` (optional): comma-separated `ip:port` direct addresses for the seed
|
||||
//! - `WORKER_CMD` (optional): path to Python interpreter (default: `python3`)
|
||||
//! - `WORKER_SCRIPT` (optional): path to tinygrad_worker.py (default: `./tinygrad_worker.py`)
|
||||
//! - `PYTHON` (optional): set to `1` for CPU fallback backend (no GPU/clang)
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use distribution::iroh_driver::{IrohDriver, IrohDriverConfig};
|
||||
use distribution::node::DistributedNodeConfig;
|
||||
use distribution::registry::RegistryConfig;
|
||||
use distribution::swim::probe::SwimConfig;
|
||||
use iroh::RelayMode;
|
||||
|
||||
use swactor::actor::ActorAddress;
|
||||
use swactor::runtime::{Runtime, RuntimeConfig};
|
||||
use swactor::transport::{CodecRegistry, TransportRouter};
|
||||
|
||||
use single_gpu_inference::inference_actor::{InferenceActor, InferenceActorStatus, RequestBridge};
|
||||
use single_gpu_inference::iroh_transport::{decode_wire, IrohActorTransport, ACTOR_ALPN};
|
||||
use single_gpu_inference::messages::{inference_codec_registry, InferenceRequest};
|
||||
use swactor_process::{ProcessMode, ProcessSpec};
|
||||
|
||||
fn parse_hex_node_id(s: &str) -> [u8; 32] {
|
||||
assert!(
|
||||
s.len() == 64,
|
||||
"SEED_ADDR must be 64 hex characters (32 bytes)"
|
||||
);
|
||||
let mut bytes = [0u8; 32];
|
||||
for i in 0..32 {
|
||||
bytes[i] = u8::from_str_radix(&s[i * 2..i * 2 + 2], 16)
|
||||
.unwrap_or_else(|_| panic!("invalid hex at position {}", i * 2));
|
||||
}
|
||||
bytes
|
||||
}
|
||||
|
||||
fn node_config() -> DistributedNodeConfig {
|
||||
DistributedNodeConfig {
|
||||
swim: SwimConfig {
|
||||
probe_interval: 10,
|
||||
probe_timeout: 15,
|
||||
indirect_probes: 2,
|
||||
suspicion_timeout: 60,
|
||||
dead_reprobe_interval: 100,
|
||||
..SwimConfig::default()
|
||||
},
|
||||
cache_capacity: 100,
|
||||
republish_interval: 50,
|
||||
registry: RegistryConfig::default(),
|
||||
metadata_lambda: 3,
|
||||
}
|
||||
}
|
||||
|
||||
fn worker_spec() -> ProcessSpec {
|
||||
let cmd = std::env::var("WORKER_CMD").unwrap_or_else(|_| "python3".into());
|
||||
let script =
|
||||
std::env::var("WORKER_SCRIPT").unwrap_or_else(|_| "./tinygrad_worker.py".into());
|
||||
|
||||
let mut env = HashMap::new();
|
||||
if let Ok(val) = std::env::var("PYTHON") {
|
||||
env.insert("PYTHON".into(), val);
|
||||
}
|
||||
|
||||
ProcessSpec {
|
||||
command: cmd,
|
||||
args: vec![script],
|
||||
env,
|
||||
working_dir: None,
|
||||
mode: ProcessMode::Automated,
|
||||
initial_pty_size: None,
|
||||
kill_timeout: Some(Duration::from_secs(5)),
|
||||
stdin_buffer_limit: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Drain incoming actor messages from iroh into the swactor runtime.
|
||||
///
|
||||
/// Like `drain_actor_messages` from the shared module, but also inspects
|
||||
/// `InferenceRequest` payloads to extract `reply_to` addresses. Returns
|
||||
/// the set of reply-to addresses found so the caller can dynamically
|
||||
/// register transport routes for the response path.
|
||||
fn drain_and_collect_reply_addrs(
|
||||
driver: &IrohDriver,
|
||||
codecs: &CodecRegistry,
|
||||
rt: &Runtime,
|
||||
drain_sleep: Duration,
|
||||
) -> Vec<ActorAddress> {
|
||||
let conns = driver.drain_other_connections();
|
||||
if conns.is_empty() {
|
||||
return vec![];
|
||||
}
|
||||
let handle = driver.tokio_handle();
|
||||
let (tx, rx) = std::sync::mpsc::channel::<Vec<u8>>();
|
||||
|
||||
for (_node_id, conn) in conns {
|
||||
let tx = tx.clone();
|
||||
handle.spawn(async move {
|
||||
loop {
|
||||
match tokio::time::timeout(Duration::from_secs(1), conn.accept_uni()).await {
|
||||
Ok(Ok(mut recv)) => {
|
||||
if let Ok(data) = recv.read_to_end(256 * 1024).await {
|
||||
let _ = tx.send(data);
|
||||
}
|
||||
}
|
||||
_ => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
drop(tx);
|
||||
std::thread::sleep(drain_sleep);
|
||||
|
||||
let mut reply_addrs = Vec::new();
|
||||
for data in rx.try_iter() {
|
||||
if let Some(envelope) = decode_wire(&data) {
|
||||
// Inspect InferenceRequest payloads for reply_to addresses
|
||||
if envelope.type_tag == "smoke::InferenceRequest" {
|
||||
if let Ok(req) = serde_json::from_slice::<InferenceRequest>(&envelope.payload) {
|
||||
reply_addrs.push(req.reply_to);
|
||||
}
|
||||
}
|
||||
if let Ok((addr, msg)) = codecs.receive(envelope) {
|
||||
let _ = rt.deliver_raw(addr, msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
reply_addrs
|
||||
}
|
||||
|
||||
fn main() {
|
||||
// Create iroh driver
|
||||
let mut driver = IrohDriver::new(IrohDriverConfig {
|
||||
secret_key: None,
|
||||
relay_mode: RelayMode::Default,
|
||||
node: node_config(),
|
||||
peer_auth: None,
|
||||
additional_alpns: vec![ACTOR_ALPN.to_vec()],
|
||||
})
|
||||
.expect("failed to create iroh driver");
|
||||
|
||||
// Print node address info so tests/orchestrators can discover us
|
||||
let my_id = driver.node_id();
|
||||
let my_hex: String = my_id.0.iter().map(|b| format!("{:02x}", b)).collect();
|
||||
let direct_addrs: Vec<String> = driver
|
||||
.direct_addresses()
|
||||
.iter()
|
||||
.map(|sa| sa.to_string())
|
||||
.collect();
|
||||
let direct_str = direct_addrs.join(",");
|
||||
eprintln!("gpu-node started (node_id: {my_hex})");
|
||||
eprintln!("GPU_NODE_ADDR {my_hex} {direct_str}");
|
||||
|
||||
// Optionally join a seed node (required for remote deployments,
|
||||
// optional for localhost testing where the orchestrator joins us)
|
||||
if let Ok(seed_hex_raw) = std::env::var("SEED_ADDR") {
|
||||
let seed_hex = seed_hex_raw.trim().to_string();
|
||||
let seed_bytes = parse_hex_node_id(&seed_hex);
|
||||
let seed_key =
|
||||
iroh::PublicKey::from_bytes(&seed_bytes).expect("invalid seed public key");
|
||||
|
||||
let mut seed_addr = iroh::EndpointAddr::from(seed_key);
|
||||
// Add relay URL so we can find the seed over the internet
|
||||
if let Ok(relay) = std::env::var("SEED_RELAY") {
|
||||
let relay = relay.trim().to_string();
|
||||
if let Ok(relay_url) = relay.parse::<iroh::RelayUrl>() {
|
||||
eprintln!("using seed relay: {relay}");
|
||||
seed_addr = seed_addr.with_relay_url(relay_url);
|
||||
}
|
||||
}
|
||||
if let Ok(direct) = std::env::var("SEED_DIRECT") {
|
||||
for part in direct.split(',') {
|
||||
if let Ok(sa) = part.trim().parse::<SocketAddr>() {
|
||||
seed_addr = seed_addr.with_ip_addr(sa);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
eprintln!("joining seed: {seed_hex}");
|
||||
driver.join(&[seed_addr]);
|
||||
} else {
|
||||
eprintln!("no SEED_ADDR set — listening for incoming connections");
|
||||
}
|
||||
|
||||
// Create actor runtime
|
||||
let mut rt = Runtime::new(RuntimeConfig::default());
|
||||
let codecs = Arc::new(inference_codec_registry());
|
||||
|
||||
// Spawn InferenceActor + RequestBridge
|
||||
let status_inbox = rt.new_inbox::<InferenceActorStatus>().unwrap();
|
||||
let sender = rt.create_sender();
|
||||
let actor =
|
||||
InferenceActor::new(worker_spec(), sender).with_status_addr(*status_inbox.addr());
|
||||
let inference_addr = rt.spawn(actor).unwrap();
|
||||
|
||||
let bridge = RequestBridge { target: inference_addr };
|
||||
let bridge_addr = rt.spawn(bridge).unwrap();
|
||||
|
||||
// Register "inference" name in the cluster
|
||||
driver
|
||||
.node_mut()
|
||||
.register_name("inference".into(), bridge_addr);
|
||||
eprintln!("registered name 'inference' -> bridge {:?}", bridge_addr);
|
||||
|
||||
// Install codecs on the runtime
|
||||
rt.set_codec_registry(codecs.clone());
|
||||
|
||||
// Transport router — routes are added dynamically as requests arrive
|
||||
let router = Arc::new(TransportRouter::new());
|
||||
rt.set_transport_router(router.clone());
|
||||
|
||||
// Wait for worker to be ready (model download + load can take minutes)
|
||||
let start = Instant::now();
|
||||
let mut worker_ready = false;
|
||||
while start.elapsed() < Duration::from_secs(600) {
|
||||
rt.tick();
|
||||
if let Some(status) = status_inbox.try_recv() {
|
||||
match status {
|
||||
InferenceActorStatus::WorkerReady { pid } => {
|
||||
eprintln!("worker ready (pid: {:?})", pid);
|
||||
worker_ready = true;
|
||||
break;
|
||||
}
|
||||
InferenceActorStatus::ProcessStarted => {
|
||||
eprintln!("worker process started");
|
||||
}
|
||||
InferenceActorStatus::ProcessExited { status } => {
|
||||
eprintln!("worker exited during startup: {:?}", status);
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
driver.recv();
|
||||
driver.tick();
|
||||
std::thread::sleep(Duration::from_millis(50));
|
||||
}
|
||||
if !worker_ready {
|
||||
eprintln!("worker did not become ready within 600s");
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
// Main loop
|
||||
eprintln!("entering main loop");
|
||||
loop {
|
||||
driver.recv();
|
||||
driver.tick();
|
||||
|
||||
// Drain incoming actor messages and collect reply_to addresses
|
||||
let reply_addrs =
|
||||
drain_and_collect_reply_addrs(&driver, &codecs, &rt, Duration::from_millis(50));
|
||||
|
||||
// Dynamically register transport routes for reply_to addresses.
|
||||
// These addresses live on the remote orchestrator node — we need a
|
||||
// transport route so InferenceActor can send InferenceResponse back.
|
||||
for reply_addr in reply_addrs {
|
||||
// Find the peer to route back to. With a single orchestrator peer,
|
||||
// we route all reply addresses to the only alive member.
|
||||
let snap = driver.snapshot();
|
||||
for member in &snap.members {
|
||||
if member.state == "alive" && member.node_id.len() == 64 {
|
||||
let peer_bytes = parse_hex_node_id(&member.node_id);
|
||||
if let Ok(peer_key) = iroh::PublicKey::from_bytes(&peer_bytes) {
|
||||
let peer_addr = iroh::EndpointAddr::from(peer_key);
|
||||
let transport = Arc::new(IrohActorTransport::new(
|
||||
driver.endpoint().clone(),
|
||||
peer_addr,
|
||||
driver.tokio_handle(),
|
||||
));
|
||||
router.add_route(reply_addr, transport);
|
||||
eprintln!("added return route for {:?}", reply_addr);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
rt.tick();
|
||||
|
||||
// Check worker health — log but don't exit, so SWIM stays alive
|
||||
// for diagnostics when the worker crashes
|
||||
if let Some(status) = status_inbox.try_recv() {
|
||||
match status {
|
||||
InferenceActorStatus::ProcessExited { status } => {
|
||||
eprintln!("worker process exited: {:?}", status);
|
||||
eprintln!("keeping main loop alive for diagnostics");
|
||||
}
|
||||
other => {
|
||||
eprintln!("worker status: {:?}", other);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
|
||||
// Cleanup
|
||||
let _ = rt.stop_actor(inference_addr);
|
||||
rt.tick();
|
||||
driver.shutdown();
|
||||
eprintln!("gpu-node shut down");
|
||||
}
|
||||
573
examples/single-gpu-inference/src/bin/single_gpu_inference.rs
Normal file
573
examples/single-gpu-inference/src/bin/single_gpu_inference.rs
Normal file
|
|
@ -0,0 +1,573 @@
|
|||
//! single-gpu-inference — Local-side orchestrator for distributed tinygrad inference.
|
||||
//!
|
||||
//! Starts a local iroh node, waits for the remote gpu-node to join and
|
||||
//! register its `"inference"` name, then sends an `InferenceRequest` through
|
||||
//! the full distributed pipeline and prints the response.
|
||||
//!
|
||||
//! Usage:
|
||||
//! # Localhost mode (default): expects gpu-node already running with SEED_ADDR
|
||||
//! single-gpu-inference --seed <hex-node-id> [--seed-direct ip:port,...]
|
||||
//!
|
||||
//! # vast.ai mode: rents a GPU, deploys the Docker image, runs the test
|
||||
//! single-gpu-inference --vastai --api-key <key> [--gpu RTX_4090]
|
||||
//!
|
||||
//! Environment variables:
|
||||
//! PYTHON=1 — passed through to worker (CPU fallback, no GPU/clang)
|
||||
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use distribution::iroh_driver::{IrohDriver, IrohDriverConfig};
|
||||
use distribution::node::DistributedNodeConfig;
|
||||
use distribution::registry::RegistryConfig;
|
||||
use distribution::swim::probe::SwimConfig;
|
||||
use iroh::{PublicKey, RelayMode};
|
||||
|
||||
use swactor::runtime::{Runtime, RuntimeConfig};
|
||||
use swactor::transport::TransportRouter;
|
||||
|
||||
use single_gpu_inference::iroh_transport::{drain_actor_messages, IrohActorTransport, ACTOR_ALPN};
|
||||
use single_gpu_inference::messages::{inference_codec_registry, InferenceRequest, InferenceResponse};
|
||||
|
||||
fn parse_hex_node_id(s: &str) -> [u8; 32] {
|
||||
assert!(
|
||||
s.len() == 64,
|
||||
"node id must be 64 hex characters (32 bytes)"
|
||||
);
|
||||
let mut bytes = [0u8; 32];
|
||||
for i in 0..32 {
|
||||
bytes[i] = u8::from_str_radix(&s[i * 2..i * 2 + 2], 16)
|
||||
.unwrap_or_else(|_| panic!("invalid hex at position {}", i * 2));
|
||||
}
|
||||
bytes
|
||||
}
|
||||
|
||||
fn node_config() -> DistributedNodeConfig {
|
||||
DistributedNodeConfig {
|
||||
swim: SwimConfig {
|
||||
probe_interval: 10,
|
||||
probe_timeout: 15,
|
||||
indirect_probes: 2,
|
||||
suspicion_timeout: 60,
|
||||
dead_reprobe_interval: 100,
|
||||
..SwimConfig::default()
|
||||
},
|
||||
cache_capacity: 100,
|
||||
republish_interval: 50,
|
||||
registry: RegistryConfig::default(),
|
||||
metadata_lambda: 3,
|
||||
}
|
||||
}
|
||||
|
||||
fn print_usage() {
|
||||
eprintln!("Usage:");
|
||||
eprintln!(" single-gpu-inference --seed <hex-node-id> [--seed-direct ip:port,...]");
|
||||
eprintln!(" single-gpu-inference --vastai --api-key <key> [--gpu RTX_4090] [--image ghcr.io/user/swactor-gpu:latest]");
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
|
||||
let mut seed_hex: Option<String> = None;
|
||||
let mut seed_direct: Vec<SocketAddr> = Vec::new();
|
||||
let mut vastai = false;
|
||||
let mut api_key: Option<String> = None;
|
||||
let mut gpu_name = "RTX 4090".to_string();
|
||||
let mut image = "swactor-gpu:latest".to_string();
|
||||
|
||||
let mut i = 1;
|
||||
while i < args.len() {
|
||||
match args[i].as_str() {
|
||||
"--seed" => {
|
||||
i += 1;
|
||||
if i >= args.len() {
|
||||
eprintln!("--seed requires a value");
|
||||
std::process::exit(1);
|
||||
}
|
||||
seed_hex = Some(args[i].clone());
|
||||
}
|
||||
"--seed-direct" => {
|
||||
i += 1;
|
||||
if i >= args.len() {
|
||||
eprintln!("--seed-direct requires a value");
|
||||
std::process::exit(1);
|
||||
}
|
||||
for part in args[i].split(',') {
|
||||
if let Ok(sa) = part.trim().parse::<SocketAddr>() {
|
||||
seed_direct.push(sa);
|
||||
}
|
||||
}
|
||||
}
|
||||
"--vastai" => {
|
||||
vastai = true;
|
||||
}
|
||||
"--api-key" => {
|
||||
i += 1;
|
||||
if i >= args.len() {
|
||||
eprintln!("--api-key requires a value");
|
||||
std::process::exit(1);
|
||||
}
|
||||
api_key = Some(args[i].trim().to_string());
|
||||
}
|
||||
"--gpu" => {
|
||||
i += 1;
|
||||
if i >= args.len() {
|
||||
eprintln!("--gpu requires a value");
|
||||
std::process::exit(1);
|
||||
}
|
||||
gpu_name = args[i].clone();
|
||||
}
|
||||
"--image" => {
|
||||
i += 1;
|
||||
if i >= args.len() {
|
||||
eprintln!("--image requires a value");
|
||||
std::process::exit(1);
|
||||
}
|
||||
image = args[i].clone();
|
||||
}
|
||||
"--help" | "-h" => {
|
||||
print_usage();
|
||||
return;
|
||||
}
|
||||
other => {
|
||||
eprintln!("unknown argument: {other}");
|
||||
print_usage();
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
|
||||
if vastai {
|
||||
run_vastai(api_key.expect("--api-key required with --vastai"), &gpu_name, &image);
|
||||
} else {
|
||||
let seed = seed_hex.expect("--seed required in localhost mode");
|
||||
run_localhost(&seed, &seed_direct);
|
||||
}
|
||||
}
|
||||
|
||||
fn run_localhost(seed_hex: &str, seed_direct: &[SocketAddr]) {
|
||||
let seed_bytes = parse_hex_node_id(seed_hex);
|
||||
let seed_key =
|
||||
iroh::PublicKey::from_bytes(&seed_bytes).expect("invalid seed public key");
|
||||
|
||||
// Build seed endpoint address
|
||||
let mut seed_addr = iroh::EndpointAddr::from(seed_key);
|
||||
for &sa in seed_direct {
|
||||
seed_addr = seed_addr.with_ip_addr(sa);
|
||||
}
|
||||
|
||||
// Create local iroh driver
|
||||
let mut driver = IrohDriver::new(IrohDriverConfig {
|
||||
secret_key: None,
|
||||
relay_mode: RelayMode::Disabled,
|
||||
node: node_config(),
|
||||
peer_auth: None,
|
||||
additional_alpns: vec![ACTOR_ALPN.to_vec()],
|
||||
})
|
||||
.expect("failed to create iroh driver");
|
||||
|
||||
let my_id = driver.node_id();
|
||||
let my_hex: String = my_id.0.iter().map(|b| format!("{:02x}", b)).collect();
|
||||
eprintln!("single-gpu-inference started (node id: {my_hex})");
|
||||
eprintln!("joining seed: {seed_hex}");
|
||||
|
||||
driver.join(&[seed_addr.clone()]);
|
||||
|
||||
// Wait for SWIM convergence — the remote node must appear as alive
|
||||
eprintln!("waiting for cluster convergence...");
|
||||
let start = Instant::now();
|
||||
let mut converged = false;
|
||||
while start.elapsed() < Duration::from_secs(30) {
|
||||
driver.recv();
|
||||
driver.tick();
|
||||
|
||||
let peer_key = PublicKey::from_bytes(&seed_bytes).unwrap();
|
||||
let snap = driver.snapshot();
|
||||
let peer_hex: String = peer_key.as_bytes().iter().map(|b| format!("{:02x}", b)).collect();
|
||||
let alive = snap
|
||||
.members
|
||||
.iter()
|
||||
.any(|m| m.node_id == peer_hex && m.state == "alive");
|
||||
if alive {
|
||||
eprintln!("cluster converged — remote node is alive");
|
||||
converged = true;
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(50));
|
||||
}
|
||||
if !converged {
|
||||
eprintln!("cluster did not converge within 30s");
|
||||
driver.shutdown();
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
// Wait for name resolution — the remote node registers "inference"
|
||||
eprintln!("resolving name 'inference'...");
|
||||
let start = Instant::now();
|
||||
let mut bridge_addr = None;
|
||||
while start.elapsed() < Duration::from_secs(30) {
|
||||
driver.recv();
|
||||
driver.tick();
|
||||
|
||||
if let Some((addr, _node_id)) = driver.node().resolve_name("inference") {
|
||||
eprintln!("resolved 'inference' -> {:?}", addr);
|
||||
bridge_addr = Some(addr);
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(100));
|
||||
}
|
||||
let bridge_addr = bridge_addr.expect("failed to resolve 'inference' name within 30s");
|
||||
|
||||
// Create actor runtime and response inbox
|
||||
let mut rt = Runtime::new(RuntimeConfig::default());
|
||||
let codecs = Arc::new(inference_codec_registry());
|
||||
let response_inbox = rt.new_inbox::<InferenceResponse>().unwrap();
|
||||
let inbox_addr = *response_inbox.addr();
|
||||
|
||||
// Build transport to the remote node
|
||||
let transport_to_remote = Arc::new(IrohActorTransport::new(
|
||||
driver.endpoint().clone(),
|
||||
seed_addr,
|
||||
driver.tokio_handle(),
|
||||
));
|
||||
|
||||
// Wire routes: send to bridge_addr on remote
|
||||
let router = TransportRouter::new();
|
||||
router.add_route(bridge_addr, transport_to_remote);
|
||||
|
||||
rt.set_codec_registry(codecs.clone());
|
||||
rt.set_transport_router(Arc::new(router));
|
||||
|
||||
// Send InferenceRequest
|
||||
eprintln!("sending InferenceRequest...");
|
||||
rt.send_to(
|
||||
bridge_addr,
|
||||
InferenceRequest {
|
||||
prompt: "Say hello".into(),
|
||||
max_tokens: 32,
|
||||
temperature: 0.7,
|
||||
reply_to: inbox_addr,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Pump loop: drain messages, tick, check for response
|
||||
let start = Instant::now();
|
||||
let timeout = Duration::from_secs(300); // generous for model download + inference
|
||||
let mut got_response = false;
|
||||
|
||||
while start.elapsed() < timeout {
|
||||
std::thread::sleep(Duration::from_millis(100));
|
||||
|
||||
driver.recv();
|
||||
driver.tick();
|
||||
drain_actor_messages(&driver, &codecs, &rt, Duration::from_millis(100));
|
||||
rt.tick();
|
||||
|
||||
if let Some(response) = response_inbox.try_recv() {
|
||||
if response.text.is_empty() {
|
||||
eprintln!("received empty response (worker may not be ready), retrying...");
|
||||
continue;
|
||||
}
|
||||
println!("=== Inference Response ===");
|
||||
println!("{}", response.text);
|
||||
println!("==========================");
|
||||
got_response = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if !got_response {
|
||||
eprintln!("did not receive InferenceResponse within timeout");
|
||||
driver.shutdown();
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
driver.shutdown();
|
||||
eprintln!("single-gpu-inference complete");
|
||||
}
|
||||
|
||||
fn run_vastai(api_key_raw: String, gpu_name: &str, image: &str) {
|
||||
let api_key = api_key_raw.trim().to_string();
|
||||
// Create a tokio runtime for the async vast.ai client
|
||||
let tokio_rt = tokio::runtime::Runtime::new().expect("failed to create tokio runtime");
|
||||
|
||||
// Create local iroh driver first so we know our node id
|
||||
// Use RelayMode::Default for WAN NAT traversal to vast.ai instances
|
||||
let mut driver = IrohDriver::new(IrohDriverConfig {
|
||||
secret_key: None,
|
||||
relay_mode: RelayMode::Default,
|
||||
node: node_config(),
|
||||
peer_auth: None,
|
||||
additional_alpns: vec![ACTOR_ALPN.to_vec()],
|
||||
})
|
||||
.expect("failed to create iroh driver");
|
||||
|
||||
let my_id = driver.node_id();
|
||||
let my_hex: String = my_id.0.iter().map(|b| format!("{:02x}", b)).collect();
|
||||
eprintln!("single-gpu-inference started (node id: {my_hex})");
|
||||
|
||||
let base_url = "https://cloud.vast.ai";
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
// Wait for relay connection so the remote gpu-node can find us over the internet
|
||||
eprintln!("waiting for relay connection...");
|
||||
let relay_url = {
|
||||
let start = Instant::now();
|
||||
loop {
|
||||
driver.recv();
|
||||
driver.tick();
|
||||
if let Some(url) = driver.home_relay_url() {
|
||||
eprintln!("home relay: {url}");
|
||||
break Some(url.to_string());
|
||||
}
|
||||
if start.elapsed() > Duration::from_secs(15) {
|
||||
eprintln!("warning: no relay URL available after 15s");
|
||||
break None;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(100));
|
||||
}
|
||||
};
|
||||
|
||||
// Retry loop: find offer → create → wait for running → wait for convergence.
|
||||
// Broken hosts (CDI errors, GFW, networking) get excluded on retry.
|
||||
let max_retries = 3;
|
||||
let mut excluded_offers: Vec<u64> = Vec::new();
|
||||
let mut running_contract_id = 0u64;
|
||||
let mut converged = false;
|
||||
|
||||
for attempt in 1..=max_retries {
|
||||
// 1. Find cheapest offer (excluding previously failed ones)
|
||||
eprintln!("finding {gpu_name} offer on vast.ai (attempt {attempt}/{max_retries})...");
|
||||
let offer = match tokio_rt.block_on(single_gpu_inference::vastai::find_offer(
|
||||
&client, base_url, &api_key, gpu_name, &excluded_offers,
|
||||
)) {
|
||||
Ok(o) => o,
|
||||
Err(e) => {
|
||||
eprintln!("no offers available: {e}");
|
||||
driver.shutdown();
|
||||
std::process::exit(1);
|
||||
}
|
||||
};
|
||||
eprintln!(
|
||||
"found offer {} ({}, ${:.3}/hr, geo={:?})",
|
||||
offer.id, offer.gpu_name, offer.dph_total, offer.geolocation
|
||||
);
|
||||
let offer_id = offer.id;
|
||||
|
||||
// 2. Create instance with our node id as SEED_ADDR
|
||||
eprintln!("creating instance...");
|
||||
let instance = match tokio_rt.block_on(single_gpu_inference::vastai::create_instance(
|
||||
&client, base_url, &api_key, offer_id, &my_hex,
|
||||
relay_url.as_deref(), image,
|
||||
)) {
|
||||
Ok(i) => i,
|
||||
Err(e) => {
|
||||
eprintln!("failed to create instance: {e}");
|
||||
excluded_offers.push(offer_id);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let contract_id = instance.contract_id;
|
||||
eprintln!("instance created (contract: {contract_id})");
|
||||
|
||||
// 3. Wait for instance to be running
|
||||
eprintln!("waiting for instance to start...");
|
||||
match tokio_rt.block_on(single_gpu_inference::vastai::wait_for_running(
|
||||
&client, base_url, &api_key, contract_id,
|
||||
Duration::from_secs(10), 60,
|
||||
)) {
|
||||
Ok(r) => {
|
||||
eprintln!("instance running at {}:{}", r.ip, r.port);
|
||||
running_contract_id = contract_id;
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("instance failed to start: {e}");
|
||||
eprintln!("destroying instance and excluding offer {offer_id}...");
|
||||
let _ = tokio_rt.block_on(single_gpu_inference::vastai::destroy_instance(
|
||||
&client, base_url, &api_key, contract_id,
|
||||
));
|
||||
excluded_offers.push(offer_id);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Wait for SWIM convergence (gpu-node must connect back via iroh relay)
|
||||
eprintln!("waiting for cluster convergence...");
|
||||
let start = Instant::now();
|
||||
let mut last_log = Instant::now();
|
||||
while start.elapsed() < Duration::from_secs(120) {
|
||||
driver.recv();
|
||||
driver.tick();
|
||||
|
||||
let snap = driver.snapshot();
|
||||
if snap.members.iter().any(|m| m.state == "alive") {
|
||||
eprintln!("cluster converged");
|
||||
converged = true;
|
||||
break;
|
||||
}
|
||||
if last_log.elapsed() >= Duration::from_secs(15) {
|
||||
let elapsed = start.elapsed().as_secs();
|
||||
let members: Vec<_> = snap.members.iter()
|
||||
.map(|m| format!("{}={}", &m.node_id[..8], m.state))
|
||||
.collect();
|
||||
eprintln!(" convergence: {elapsed}s elapsed, members: {members:?}");
|
||||
last_log = Instant::now();
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(100));
|
||||
}
|
||||
|
||||
if converged {
|
||||
break;
|
||||
}
|
||||
|
||||
// Convergence failed — fetch logs for debugging, then destroy
|
||||
eprintln!("cluster did not converge, fetching instance logs...");
|
||||
if let Ok(log_url) = tokio_rt.block_on(single_gpu_inference::vastai::request_logs(
|
||||
&client, base_url, &api_key, contract_id,
|
||||
)) {
|
||||
if let Ok(logs) = tokio_rt.block_on(single_gpu_inference::vastai::fetch_logs(&client, &log_url)) {
|
||||
eprintln!("--- instance logs (contract {contract_id}) ---");
|
||||
// Print last 40 lines to avoid flooding
|
||||
let lines: Vec<&str> = logs.lines().collect();
|
||||
let start = if lines.len() > 40 { lines.len() - 40 } else { 0 };
|
||||
for line in &lines[start..] {
|
||||
eprintln!(" {line}");
|
||||
}
|
||||
eprintln!("--- end logs ---");
|
||||
}
|
||||
}
|
||||
eprintln!("destroying instance and excluding offer {offer_id}...");
|
||||
let _ = tokio_rt.block_on(single_gpu_inference::vastai::destroy_instance(
|
||||
&client, base_url, &api_key, contract_id,
|
||||
));
|
||||
excluded_offers.push(offer_id);
|
||||
}
|
||||
|
||||
if !converged {
|
||||
eprintln!("all {max_retries} attempts failed (no host converged), giving up");
|
||||
driver.shutdown();
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
// Resolve "inference" name
|
||||
eprintln!("resolving 'inference' name...");
|
||||
let start = Instant::now();
|
||||
let mut bridge_addr = None;
|
||||
let mut remote_node_id = None;
|
||||
while start.elapsed() < Duration::from_secs(60) {
|
||||
driver.recv();
|
||||
driver.tick();
|
||||
if let Some((addr, node_id)) = driver.node().resolve_name("inference") {
|
||||
bridge_addr = Some(addr);
|
||||
remote_node_id = Some(node_id);
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(200));
|
||||
}
|
||||
let bridge_addr = bridge_addr.expect("failed to resolve 'inference' name");
|
||||
let remote_node_id = remote_node_id.unwrap();
|
||||
|
||||
// Build transport to remote
|
||||
let remote_key = PublicKey::from_bytes(&remote_node_id.0).expect("invalid remote node key");
|
||||
let remote_endpoint_addr = iroh::EndpointAddr::from(remote_key);
|
||||
|
||||
let mut rt = Runtime::new(RuntimeConfig::default());
|
||||
let codecs = Arc::new(inference_codec_registry());
|
||||
let response_inbox = rt.new_inbox::<InferenceResponse>().unwrap();
|
||||
let inbox_addr = *response_inbox.addr();
|
||||
|
||||
let transport_to_remote = Arc::new(IrohActorTransport::new(
|
||||
driver.endpoint().clone(),
|
||||
remote_endpoint_addr,
|
||||
driver.tokio_handle(),
|
||||
));
|
||||
let router = TransportRouter::new();
|
||||
router.add_route(bridge_addr, transport_to_remote);
|
||||
rt.set_codec_registry(codecs.clone());
|
||||
rt.set_transport_router(Arc::new(router));
|
||||
|
||||
// Send request
|
||||
eprintln!("sending InferenceRequest...");
|
||||
rt.send_to(
|
||||
bridge_addr,
|
||||
InferenceRequest {
|
||||
prompt: "Say hello".into(),
|
||||
max_tokens: 32,
|
||||
temperature: 0.7,
|
||||
reply_to: inbox_addr,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Pump loop
|
||||
let start = Instant::now();
|
||||
let mut last_diag = Instant::now();
|
||||
let mut got_response = false;
|
||||
while start.elapsed() < Duration::from_secs(300) {
|
||||
std::thread::sleep(Duration::from_millis(100));
|
||||
driver.recv();
|
||||
driver.tick();
|
||||
drain_actor_messages(&driver, &codecs, &rt, Duration::from_millis(50));
|
||||
rt.tick();
|
||||
|
||||
// Periodic SWIM health diagnostics
|
||||
if last_diag.elapsed() >= Duration::from_secs(15) {
|
||||
let snap = driver.snapshot();
|
||||
let members: Vec<_> = snap.members.iter()
|
||||
.map(|m| format!("{}={}", &m.node_id[..8], m.state))
|
||||
.collect();
|
||||
let elapsed = start.elapsed().as_secs();
|
||||
eprintln!(" inference wait: {elapsed}s elapsed, members: {members:?}");
|
||||
last_diag = Instant::now();
|
||||
}
|
||||
|
||||
if let Some(response) = response_inbox.try_recv() {
|
||||
if response.text.is_empty() {
|
||||
continue;
|
||||
}
|
||||
println!("=== Inference Response ===");
|
||||
println!("{}", response.text);
|
||||
println!("==========================");
|
||||
got_response = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if !got_response {
|
||||
eprintln!("did not receive InferenceResponse within timeout");
|
||||
// Fetch logs for debugging before destroying
|
||||
eprintln!("fetching instance logs...");
|
||||
if let Ok(log_url) = tokio_rt.block_on(single_gpu_inference::vastai::request_logs(
|
||||
&client, base_url, &api_key, running_contract_id,
|
||||
)) {
|
||||
if let Ok(logs) = tokio_rt.block_on(single_gpu_inference::vastai::fetch_logs(&client, &log_url)) {
|
||||
eprintln!("--- instance logs (contract {running_contract_id}) ---");
|
||||
let lines: Vec<&str> = logs.lines().collect();
|
||||
let tail = if lines.len() > 60 { lines.len() - 60 } else { 0 };
|
||||
for line in &lines[tail..] {
|
||||
eprintln!(" {line}");
|
||||
}
|
||||
eprintln!("--- end logs ---");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 5. Destroy instance
|
||||
eprintln!("destroying instance...");
|
||||
let _ = tokio_rt.block_on(single_gpu_inference::vastai::destroy_instance(
|
||||
&client,
|
||||
base_url,
|
||||
&api_key,
|
||||
running_contract_id,
|
||||
));
|
||||
|
||||
driver.shutdown();
|
||||
if got_response {
|
||||
eprintln!("single-gpu-inference complete");
|
||||
} else {
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
220
examples/single-gpu-inference/src/inference_actor.rs
Normal file
220
examples/single-gpu-inference/src/inference_actor.rs
Normal file
|
|
@ -0,0 +1,220 @@
|
|||
//! InferenceActor — bridges swactor messaging to a Python child process.
|
||||
//!
|
||||
//! Spawns a Python worker (e.g. `echo_worker.py` or `tinygrad_worker.py`) via
|
||||
//! the process crate and translates `InferenceRequest` messages into stdin JSON,
|
||||
//! then parses stdout JSON into `InferenceResponse` replies.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
|
||||
use swactor::actor::{ActorAddress, ActorInterface, Ctx};
|
||||
use swactor::runtime::ExternalSender;
|
||||
use swactor_process::{
|
||||
spawn_local_process, ExitStatus, ProcessCommand, ProcessNotification, ProcessSpec,
|
||||
};
|
||||
|
||||
use crate::messages::{InferenceRequest, InferenceResponse};
|
||||
|
||||
// ── Messages ──────────────────────────────────────────────────────────────
|
||||
|
||||
/// Union type for messages the InferenceActor can receive.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum InferenceActorMsg {
|
||||
/// An inference request from a client.
|
||||
Request(InferenceRequest),
|
||||
/// A forwarded notification from the child process.
|
||||
Process(ProcessNotification),
|
||||
}
|
||||
|
||||
/// Status notifications emitted to an optional observer address.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum InferenceActorStatus {
|
||||
ProcessStarted,
|
||||
WorkerReady { pid: Option<u32> },
|
||||
ProcessExited { status: ExitStatus },
|
||||
}
|
||||
|
||||
// ── ProcessBridge ─────────────────────────────────────────────────────────
|
||||
|
||||
/// Receives `ProcessNotification` from the ProcessActor and forwards it
|
||||
/// wrapped as `InferenceActorMsg::Process` to the InferenceActor.
|
||||
///
|
||||
/// Necessary because swactor actors have a single `Incoming` type — the
|
||||
/// ProcessActor sends `ProcessNotification`, but InferenceActor expects
|
||||
/// `InferenceActorMsg`.
|
||||
struct ProcessBridge {
|
||||
target: ActorAddress,
|
||||
}
|
||||
|
||||
impl ActorInterface for ProcessBridge {
|
||||
type Incoming = ProcessNotification;
|
||||
type Response = ();
|
||||
|
||||
fn handle(&mut self, ctx: &Ctx, msg: ProcessNotification) {
|
||||
let _ = ctx.send(self.target, InferenceActorMsg::Process(msg));
|
||||
}
|
||||
}
|
||||
|
||||
// ── RequestBridge ────────────────────────────────────────────────────────
|
||||
|
||||
/// Receives `InferenceRequest` from the network and forwards it wrapped as
|
||||
/// `InferenceActorMsg::Request` to the InferenceActor.
|
||||
///
|
||||
/// Necessary because the network codec delivers raw `InferenceRequest`, but
|
||||
/// InferenceActor expects `InferenceActorMsg`.
|
||||
pub struct RequestBridge {
|
||||
pub target: ActorAddress,
|
||||
}
|
||||
|
||||
impl ActorInterface for RequestBridge {
|
||||
type Incoming = InferenceRequest;
|
||||
type Response = ();
|
||||
|
||||
fn handle(&mut self, ctx: &Ctx, msg: InferenceRequest) {
|
||||
let _ = ctx.send(self.target, InferenceActorMsg::Request(msg));
|
||||
}
|
||||
}
|
||||
|
||||
// ── InferenceActor ────────────────────────────────────────────────────────
|
||||
|
||||
pub struct InferenceActor {
|
||||
spec: ProcessSpec,
|
||||
sender: ExternalSender,
|
||||
process_addr: Option<ActorAddress>,
|
||||
bridge_addr: Option<ActorAddress>,
|
||||
pending_replies: VecDeque<ActorAddress>,
|
||||
ready: bool,
|
||||
process_alive: bool,
|
||||
worker_pid: Option<u32>,
|
||||
status_addr: Option<ActorAddress>,
|
||||
output_buffer: String,
|
||||
}
|
||||
|
||||
impl InferenceActor {
|
||||
pub fn new(spec: ProcessSpec, sender: ExternalSender) -> Self {
|
||||
Self {
|
||||
spec,
|
||||
sender,
|
||||
process_addr: None,
|
||||
bridge_addr: None,
|
||||
pending_replies: VecDeque::new(),
|
||||
ready: false,
|
||||
process_alive: false,
|
||||
worker_pid: None,
|
||||
status_addr: None,
|
||||
output_buffer: String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set an observer address that receives `InferenceActorStatus` updates.
|
||||
pub fn with_status_addr(mut self, addr: ActorAddress) -> Self {
|
||||
self.status_addr = Some(addr);
|
||||
self
|
||||
}
|
||||
|
||||
fn process_output_line(&mut self, ctx: &Ctx, line: &str) {
|
||||
let Ok(val) = serde_json::from_str::<serde_json::Value>(line) else {
|
||||
// Log non-JSON output (Python tracebacks, error messages, etc.)
|
||||
if !line.is_empty() {
|
||||
eprintln!("worker: {line}");
|
||||
}
|
||||
return;
|
||||
};
|
||||
|
||||
if val.get("status").and_then(|v| v.as_str()) == Some("ready") {
|
||||
self.ready = true;
|
||||
let pid = val.get("pid").and_then(|v| v.as_u64()).map(|p| p as u32);
|
||||
self.worker_pid = pid;
|
||||
if let Some(addr) = self.status_addr {
|
||||
let _ = ctx.send(addr, InferenceActorStatus::WorkerReady { pid });
|
||||
}
|
||||
} else if let Some(response) = val.get("response").and_then(|v| v.as_str()) {
|
||||
if let Some(reply_to) = self.pending_replies.pop_front() {
|
||||
let _ = ctx.send(reply_to, InferenceResponse { text: response.to_string() });
|
||||
}
|
||||
} else if let Some(err) = val.get("error").and_then(|v| v.as_str()) {
|
||||
eprintln!("worker error: {err}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ActorInterface for InferenceActor {
|
||||
type Incoming = InferenceActorMsg;
|
||||
type Response = ();
|
||||
|
||||
fn on_start(&mut self, ctx: &Ctx) {
|
||||
let proc_addr = spawn_local_process(ctx, &self.sender, self.spec.clone())
|
||||
.expect("failed to spawn worker process");
|
||||
|
||||
let bridge = ProcessBridge { target: ctx.self_addr() };
|
||||
let bridge_addr = ctx.spawn(bridge).expect("failed to spawn process bridge");
|
||||
|
||||
let _ = ctx.send(proc_addr, ProcessCommand::Subscribe { address: bridge_addr });
|
||||
|
||||
self.process_addr = Some(proc_addr);
|
||||
self.bridge_addr = Some(bridge_addr);
|
||||
}
|
||||
|
||||
fn handle(&mut self, ctx: &Ctx, msg: InferenceActorMsg) {
|
||||
match msg {
|
||||
InferenceActorMsg::Request(req) => {
|
||||
if !self.ready || !self.process_alive {
|
||||
let _ = ctx.send(req.reply_to, InferenceResponse { text: String::new() });
|
||||
return;
|
||||
}
|
||||
self.pending_replies.push_back(req.reply_to);
|
||||
let json = serde_json::json!({
|
||||
"prompt": req.prompt,
|
||||
"max_tokens": req.max_tokens,
|
||||
"temperature": req.temperature,
|
||||
});
|
||||
let mut data = serde_json::to_vec(&json).unwrap();
|
||||
data.push(b'\n');
|
||||
if let Some(proc_addr) = self.process_addr {
|
||||
let _ = ctx.send(proc_addr, ProcessCommand::WriteStdin { data });
|
||||
}
|
||||
}
|
||||
InferenceActorMsg::Process(notif) => match notif {
|
||||
ProcessNotification::Started { .. } => {
|
||||
self.process_alive = true;
|
||||
if let Some(addr) = self.status_addr {
|
||||
let _ = ctx.send(addr, InferenceActorStatus::ProcessStarted);
|
||||
}
|
||||
}
|
||||
ProcessNotification::Output { data, .. } => {
|
||||
let text = String::from_utf8_lossy(&data);
|
||||
self.output_buffer.push_str(&text);
|
||||
while let Some(pos) = self.output_buffer.find('\n') {
|
||||
let line = self.output_buffer[..pos].to_string();
|
||||
self.output_buffer = self.output_buffer[pos + 1..].to_string();
|
||||
self.process_output_line(ctx, line.trim());
|
||||
}
|
||||
}
|
||||
ProcessNotification::Exited { status, .. } => {
|
||||
self.process_alive = false;
|
||||
self.ready = false;
|
||||
// Drain pending requests with empty responses
|
||||
for reply_to in self.pending_replies.drain(..) {
|
||||
let _ = ctx.send(reply_to, InferenceResponse { text: String::new() });
|
||||
}
|
||||
if let Some(addr) = self.status_addr {
|
||||
let _ = ctx.send(addr, InferenceActorStatus::ProcessExited { status });
|
||||
}
|
||||
}
|
||||
ProcessNotification::Error { .. } => {
|
||||
self.process_alive = false;
|
||||
self.ready = false;
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn on_stop(&mut self, ctx: &Ctx) {
|
||||
if let Some(proc_addr) = self.process_addr {
|
||||
let _ = ctx.send(proc_addr, ProcessCommand::Close);
|
||||
let _ = ctx.stop_actor(proc_addr);
|
||||
}
|
||||
if let Some(bridge_addr) = self.bridge_addr {
|
||||
let _ = ctx.stop_actor(bridge_addr);
|
||||
}
|
||||
}
|
||||
}
|
||||
161
examples/single-gpu-inference/src/iroh_transport.rs
Normal file
161
examples/single-gpu-inference/src/iroh_transport.rs
Normal file
|
|
@ -0,0 +1,161 @@
|
|||
//! Shared iroh transport infrastructure for actor-level QUIC messaging.
|
||||
//!
|
||||
//! Provides `IrohActorTransport` (sends `WireEnvelope`s over iroh QUIC
|
||||
//! uni-directional streams), wire encoding/decoding, and a drain helper
|
||||
//! that feeds incoming messages into a swactor `Runtime`.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use distribution::iroh_driver::IrohDriver;
|
||||
use swactor::actor::ActorAddress;
|
||||
use swactor::runtime::Runtime;
|
||||
use swactor::transport::{CodecRegistry, Transport, WireEnvelope};
|
||||
use swactor::Error;
|
||||
|
||||
/// ALPN protocol for actor-level messages (distinct from SWIM protocol).
|
||||
pub const ACTOR_ALPN: &[u8] = b"swactor/actor/1";
|
||||
|
||||
/// Sends WireEnvelopes over iroh QUIC uni-directional streams.
|
||||
///
|
||||
/// Caches the QUIC connection so it stays alive between sends — dropping
|
||||
/// a connection before the receiver reads its streams causes "closed by
|
||||
/// peer" errors.
|
||||
///
|
||||
/// Wire format: [32B dest_addr][4B tag_len BE][tag bytes][payload bytes]
|
||||
pub struct IrohActorTransport {
|
||||
endpoint: iroh::Endpoint,
|
||||
target_addr: iroh::EndpointAddr,
|
||||
handle: tokio::runtime::Handle,
|
||||
conn: std::sync::Mutex<Option<iroh::endpoint::Connection>>,
|
||||
}
|
||||
|
||||
impl IrohActorTransport {
|
||||
pub fn new(
|
||||
endpoint: iroh::Endpoint,
|
||||
target_addr: iroh::EndpointAddr,
|
||||
handle: tokio::runtime::Handle,
|
||||
) -> Self {
|
||||
Self {
|
||||
endpoint,
|
||||
target_addr,
|
||||
handle,
|
||||
conn: std::sync::Mutex::new(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Transport for IrohActorTransport {
|
||||
fn send(&self, envelope: WireEnvelope) -> Result<(), Error> {
|
||||
let data = encode_wire(&envelope);
|
||||
let ep = self.endpoint.clone();
|
||||
let target = self.target_addr.clone();
|
||||
let cached = self.conn.lock().unwrap().clone();
|
||||
|
||||
let conn = self.handle.block_on(async move {
|
||||
if let Some(c) = cached {
|
||||
if c.close_reason().is_none() {
|
||||
return Ok(c);
|
||||
}
|
||||
}
|
||||
ep.connect(target, ACTOR_ALPN)
|
||||
.await
|
||||
.map_err(|e| Error::from(format!("iroh connect: {e}")))
|
||||
})?;
|
||||
|
||||
*self.conn.lock().unwrap() = Some(conn.clone());
|
||||
|
||||
self.handle.block_on(async move {
|
||||
let mut stream = conn
|
||||
.open_uni()
|
||||
.await
|
||||
.map_err(|e| Error::from(format!("iroh open_uni: {e}")))?;
|
||||
stream
|
||||
.write_all(&data)
|
||||
.await
|
||||
.map_err(|e| Error::from(format!("iroh write: {e}")))?;
|
||||
stream
|
||||
.finish()
|
||||
.map_err(|e| Error::from(format!("iroh finish: {e}")))?;
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn encode_wire(env: &WireEnvelope) -> Vec<u8> {
|
||||
let tag = env.type_tag.as_bytes();
|
||||
let mut buf = Vec::with_capacity(32 + 4 + tag.len() + env.payload.len());
|
||||
buf.extend_from_slice(&env.dest.0);
|
||||
buf.extend_from_slice(&(tag.len() as u32).to_be_bytes());
|
||||
buf.extend_from_slice(tag);
|
||||
buf.extend_from_slice(&env.payload);
|
||||
buf
|
||||
}
|
||||
|
||||
pub fn decode_wire(data: &[u8]) -> Option<WireEnvelope> {
|
||||
if data.len() < 36 {
|
||||
return None;
|
||||
}
|
||||
let mut addr = [0u8; 32];
|
||||
addr.copy_from_slice(&data[..32]);
|
||||
let tag_len = u32::from_be_bytes(data[32..36].try_into().ok()?) as usize;
|
||||
if data.len() < 36 + tag_len {
|
||||
return None;
|
||||
}
|
||||
let type_tag = String::from_utf8(data[36..36 + tag_len].to_vec()).ok()?;
|
||||
let payload = data[36 + tag_len..].to_vec();
|
||||
Some(WireEnvelope {
|
||||
dest: ActorAddress(addr),
|
||||
type_tag,
|
||||
payload,
|
||||
})
|
||||
}
|
||||
|
||||
/// Drain actor messages arriving via iroh into a swactor runtime.
|
||||
///
|
||||
/// Spawns tokio tasks (on the driver's runtime) to read uni-directional
|
||||
/// streams from pending actor-ALPN connections. Results are collected
|
||||
/// via a channel and delivered to the swactor runtime.
|
||||
///
|
||||
/// `drain_sleep` controls how long to wait for spawned tasks to read
|
||||
/// streams before collecting results. Use shorter durations (e.g. 100ms)
|
||||
/// for localhost tests, longer (e.g. 500ms) for cross-node scenarios.
|
||||
pub fn drain_actor_messages(
|
||||
driver: &IrohDriver,
|
||||
codecs: &CodecRegistry,
|
||||
rt: &Runtime,
|
||||
drain_sleep: Duration,
|
||||
) {
|
||||
let conns = driver.drain_other_connections();
|
||||
if conns.is_empty() {
|
||||
return;
|
||||
}
|
||||
let handle = driver.tokio_handle();
|
||||
let (tx, rx) = std::sync::mpsc::channel::<Vec<u8>>();
|
||||
|
||||
for (_node_id, conn) in conns {
|
||||
let tx = tx.clone();
|
||||
handle.spawn(async move {
|
||||
loop {
|
||||
match tokio::time::timeout(Duration::from_secs(1), conn.accept_uni()).await {
|
||||
Ok(Ok(mut recv)) => {
|
||||
if let Ok(data) = recv.read_to_end(256 * 1024).await {
|
||||
let _ = tx.send(data);
|
||||
}
|
||||
}
|
||||
_ => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
drop(tx);
|
||||
|
||||
std::thread::sleep(drain_sleep);
|
||||
|
||||
for data in rx.try_iter() {
|
||||
if let Some(envelope) = decode_wire(&data) {
|
||||
if let Ok((addr, msg)) = codecs.receive(envelope) {
|
||||
let _ = rt.deliver_raw(addr, msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
4
examples/single-gpu-inference/src/lib.rs
Normal file
4
examples/single-gpu-inference/src/lib.rs
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
pub mod inference_actor;
|
||||
pub mod iroh_transport;
|
||||
pub mod messages;
|
||||
pub mod vastai;
|
||||
67
examples/single-gpu-inference/src/messages.rs
Normal file
67
examples/single-gpu-inference/src/messages.rs
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
//! Inference message types, codec, and registry for the smoke test.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use swactor::actor::ActorAddress;
|
||||
use swactor::transport::{Codec, CodecRegistry, NetworkMessage};
|
||||
use swactor::Error;
|
||||
|
||||
// ─── Message Types ─────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct InferenceRequest {
|
||||
pub prompt: String,
|
||||
pub max_tokens: u32,
|
||||
pub temperature: f32,
|
||||
pub reply_to: ActorAddress,
|
||||
}
|
||||
|
||||
impl NetworkMessage for InferenceRequest {
|
||||
fn type_tag() -> &'static str {
|
||||
"smoke::InferenceRequest"
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct InferenceResponse {
|
||||
pub text: String,
|
||||
}
|
||||
|
||||
impl NetworkMessage for InferenceResponse {
|
||||
fn type_tag() -> &'static str {
|
||||
"smoke::InferenceResponse"
|
||||
}
|
||||
}
|
||||
|
||||
// ─── JSON Codec ────────────────────────────────────────────────────────────
|
||||
|
||||
pub struct InferenceCodec;
|
||||
|
||||
impl Codec<InferenceRequest> for InferenceCodec {
|
||||
fn encode(&self, msg: &InferenceRequest) -> Result<Vec<u8>, Error> {
|
||||
serde_json::to_vec(msg).map_err(|e| Error::from(format!("encode InferenceRequest: {e}")))
|
||||
}
|
||||
fn decode(&self, bytes: &[u8]) -> Result<InferenceRequest, Error> {
|
||||
serde_json::from_slice(bytes)
|
||||
.map_err(|e| Error::from(format!("decode InferenceRequest: {e}")))
|
||||
}
|
||||
}
|
||||
|
||||
impl Codec<InferenceResponse> for InferenceCodec {
|
||||
fn encode(&self, msg: &InferenceResponse) -> Result<Vec<u8>, Error> {
|
||||
serde_json::to_vec(msg).map_err(|e| Error::from(format!("encode InferenceResponse: {e}")))
|
||||
}
|
||||
fn decode(&self, bytes: &[u8]) -> Result<InferenceResponse, Error> {
|
||||
serde_json::from_slice(bytes)
|
||||
.map_err(|e| Error::from(format!("decode InferenceResponse: {e}")))
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Registry ──────────────────────────────────────────────────────────────
|
||||
|
||||
/// Build a `CodecRegistry` with inference message types registered.
|
||||
pub fn inference_codec_registry() -> CodecRegistry {
|
||||
let mut cr = CodecRegistry::new();
|
||||
cr.register::<InferenceRequest, _>(InferenceCodec);
|
||||
cr.register::<InferenceResponse, _>(InferenceCodec);
|
||||
cr
|
||||
}
|
||||
292
examples/single-gpu-inference/src/vastai.rs
Normal file
292
examples/single-gpu-inference/src/vastai.rs
Normal file
|
|
@ -0,0 +1,292 @@
|
|||
//! vast.ai REST API client for the smoke test orchestrator.
|
||||
//!
|
||||
//! Functions: find_offer, create_instance, wait_for_running, destroy_instance.
|
||||
//! All functions accept a `base_url` parameter so tests can point at a mock server.
|
||||
|
||||
use reqwest::Client;
|
||||
use serde::Deserialize;
|
||||
use std::time::Duration;
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct Offer {
|
||||
pub id: u64,
|
||||
pub gpu_name: String,
|
||||
pub dph_total: f64,
|
||||
#[serde(default)]
|
||||
pub geolocation: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct InstanceInfo {
|
||||
pub contract_id: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RunningInstance {
|
||||
pub ip: String,
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct SearchResponse {
|
||||
offers: Vec<Offer>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct CreateResponse {
|
||||
new_contract: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct InstanceResponse {
|
||||
instances: InstanceStatus,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct InstanceStatus {
|
||||
actual_status: Option<String>,
|
||||
intended_status: Option<String>,
|
||||
#[serde(default)]
|
||||
status_msg: Option<String>,
|
||||
#[serde(default)]
|
||||
public_ipaddr: Option<String>,
|
||||
#[serde(default)]
|
||||
ssh_port: Option<u16>,
|
||||
}
|
||||
|
||||
/// Find the cheapest offer matching a GPU type, excluding specific offer IDs.
|
||||
pub async fn find_offer(
|
||||
client: &Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
gpu_name: &str,
|
||||
exclude_ids: &[u64],
|
||||
) -> Result<Offer, String> {
|
||||
let query = serde_json::json!({
|
||||
"gpu_name": {"eq": gpu_name},
|
||||
"rentable": {"eq": true},
|
||||
"rented": {"eq": false},
|
||||
"reliability2": {"gte": 0.99},
|
||||
"cuda_max_good": {"gte": 12.0},
|
||||
"verified": {"eq": true},
|
||||
"direct_port_count": {"gte": 1},
|
||||
"inet_down": {"gte": 100.0},
|
||||
});
|
||||
let url = format!(
|
||||
"{base_url}/api/v0/bundles/?q={}",
|
||||
urlencoding::encode(&query.to_string())
|
||||
);
|
||||
let resp = client
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("find_offer request failed: {e}"))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return Err(format!("find_offer HTTP {status}: {body}"));
|
||||
}
|
||||
|
||||
let body: SearchResponse = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("find_offer parse failed: {e}"))?;
|
||||
|
||||
// Filter out hosts with unknown or Chinese geolocation — Docker Hub
|
||||
// and iroh relays are unreachable from behind the Great Firewall.
|
||||
let filtered: Vec<Offer> = body
|
||||
.offers
|
||||
.into_iter()
|
||||
.filter(|o| {
|
||||
o.geolocation
|
||||
.as_deref()
|
||||
.map_or(false, |g| !g.to_uppercase().contains("CN"))
|
||||
})
|
||||
.collect();
|
||||
|
||||
let candidates: Vec<Offer> = filtered
|
||||
.into_iter()
|
||||
.filter(|o| !exclude_ids.contains(&o.id))
|
||||
.collect();
|
||||
|
||||
candidates
|
||||
.into_iter()
|
||||
.min_by(|a, b| a.dph_total.partial_cmp(&b.dph_total).unwrap())
|
||||
.ok_or_else(|| "no offers available (after geo/exclusion filter)".to_string())
|
||||
}
|
||||
|
||||
/// Create a vast.ai instance from an offer, passing SEED_ADDR and SEED_RELAY in the env.
|
||||
pub async fn create_instance(
|
||||
client: &Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
offer_id: u64,
|
||||
seed_addr: &str,
|
||||
seed_relay: Option<&str>,
|
||||
image: &str,
|
||||
) -> Result<InstanceInfo, String> {
|
||||
let url = format!("{base_url}/api/v0/asks/{offer_id}/");
|
||||
let mut env = serde_json::json!({ "SEED_ADDR": seed_addr });
|
||||
if let Some(relay) = seed_relay {
|
||||
env["SEED_RELAY"] = serde_json::Value::String(relay.to_string());
|
||||
}
|
||||
let body = serde_json::json!({
|
||||
"image": image,
|
||||
"env": env,
|
||||
"onstart": "exec /usr/local/bin/gpu-node 2>&1",
|
||||
"disk": 20,
|
||||
});
|
||||
|
||||
let resp = client
|
||||
.put(&url)
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("create_instance request failed: {e}"))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return Err(format!("create_instance HTTP {status}: {body}"));
|
||||
}
|
||||
|
||||
let parsed: CreateResponse = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("create_instance parse failed: {e}"))?;
|
||||
|
||||
Ok(InstanceInfo {
|
||||
contract_id: parsed.new_contract,
|
||||
})
|
||||
}
|
||||
|
||||
/// Poll vast.ai until the instance reaches `running`, then extract IP + port.
|
||||
/// Returns error immediately on terminal statuses like `exited`.
|
||||
pub async fn wait_for_running(
|
||||
client: &Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
contract_id: u64,
|
||||
poll_interval: Duration,
|
||||
max_polls: u32,
|
||||
) -> Result<RunningInstance, String> {
|
||||
let url = format!("{base_url}/api/v0/instances/{contract_id}/");
|
||||
|
||||
for poll in 0..max_polls {
|
||||
let resp = client
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("wait_for_running request failed: {e}"))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return Err(format!("wait_for_running HTTP {status}: {body}"));
|
||||
}
|
||||
|
||||
let wrapper: InstanceResponse = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("wait_for_running parse failed: {e}"))?;
|
||||
let status = wrapper.instances;
|
||||
|
||||
let actual = status.actual_status.as_deref().unwrap_or("unknown");
|
||||
let intended = status.intended_status.as_deref().unwrap_or("unknown");
|
||||
eprintln!(" poll {}/{}: status={actual}", poll + 1, max_polls);
|
||||
|
||||
// Check for error in status_msg (host-side failures like OCI errors)
|
||||
if let Some(msg) = &status.status_msg {
|
||||
if msg.contains("Error") || msg.contains("failed") {
|
||||
return Err(format!("instance error: {msg}"));
|
||||
}
|
||||
}
|
||||
|
||||
// Check if intended_status has gone to stopped (instance gave up)
|
||||
if intended == "stopped" && actual != "running" {
|
||||
let msg = status.status_msg.unwrap_or_default();
|
||||
return Err(format!("instance stopped: {msg}"));
|
||||
}
|
||||
|
||||
match actual {
|
||||
"running" => {
|
||||
let ip = status
|
||||
.public_ipaddr
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
let port = status.ssh_port.unwrap_or(0);
|
||||
return Ok(RunningInstance { ip, port });
|
||||
}
|
||||
"exited" | "error" => {
|
||||
return Err(format!("instance reached terminal status: {actual}"));
|
||||
}
|
||||
_ => {
|
||||
tokio::time::sleep(poll_interval).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err("instance did not reach running within poll limit".to_string())
|
||||
}
|
||||
|
||||
/// Request instance logs and return the download URL.
|
||||
/// Logs take a few seconds to become available after this call.
|
||||
pub async fn request_logs(
|
||||
client: &Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
contract_id: u64,
|
||||
) -> Result<String, String> {
|
||||
let url = format!("{base_url}/api/v0/instances/request_logs/{contract_id}/");
|
||||
let resp = client
|
||||
.put(&url)
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("request_logs failed: {e}"))?;
|
||||
|
||||
let body: serde_json::Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("request_logs parse failed: {e}"))?;
|
||||
|
||||
body["result_url"]
|
||||
.as_str()
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| "no result_url in log response".to_string())
|
||||
}
|
||||
|
||||
/// Fetch instance logs from S3 URL. Returns the log text.
|
||||
pub async fn fetch_logs(client: &Client, log_url: &str) -> Result<String, String> {
|
||||
// Wait for the log to become available
|
||||
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
|
||||
let resp = client
|
||||
.get(log_url)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("fetch_logs failed: {e}"))?;
|
||||
resp.text()
|
||||
.await
|
||||
.map_err(|e| format!("fetch_logs read failed: {e}"))
|
||||
}
|
||||
|
||||
/// Destroy a vast.ai instance.
|
||||
pub async fn destroy_instance(
|
||||
client: &Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
contract_id: u64,
|
||||
) -> Result<(), String> {
|
||||
let url = format!("{base_url}/api/v0/instances/{contract_id}/");
|
||||
client
|
||||
.delete(&url)
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("destroy_instance request failed: {e}"))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
215
examples/single-gpu-inference/tests/t_actor.rs
Normal file
215
examples/single-gpu-inference/tests/t_actor.rs
Normal file
|
|
@ -0,0 +1,215 @@
|
|||
//! T-actor: InferenceActor + process bridge component tests.
|
||||
//!
|
||||
//! Spawns the InferenceActor on a single swactor runtime with `echo_worker.py`
|
||||
//! as the child process. No networking, no GPU.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use swactor::actor::ActorAddress;
|
||||
use swactor::runtime::{Inbox, Runtime, RuntimeConfig};
|
||||
|
||||
use single_gpu_inference::inference_actor::{InferenceActor, InferenceActorMsg, InferenceActorStatus};
|
||||
use single_gpu_inference::messages::{InferenceRequest, InferenceResponse};
|
||||
use swactor_process::{ExitStatus, ProcessMode, ProcessSpec};
|
||||
|
||||
// ── Helpers ───────────────────────────────────────────────────────────────
|
||||
|
||||
fn echo_worker_spec() -> ProcessSpec {
|
||||
ProcessSpec {
|
||||
command: "python3".into(),
|
||||
args: vec![format!("{}/echo_worker.py", env!("CARGO_MANIFEST_DIR"))],
|
||||
env: HashMap::new(),
|
||||
working_dir: None,
|
||||
mode: ProcessMode::Automated,
|
||||
initial_pty_size: None,
|
||||
kill_timeout: Some(Duration::from_secs(2)),
|
||||
stdin_buffer_limit: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Tick the runtime in a polling loop until the inbox has a message or timeout.
|
||||
fn tick_until_recv<M: swactor::actor::Message>(
|
||||
rt: &Runtime,
|
||||
inbox: &Inbox<M>,
|
||||
timeout: Duration,
|
||||
) -> Option<M> {
|
||||
let start = Instant::now();
|
||||
while start.elapsed() < timeout {
|
||||
rt.tick();
|
||||
if let Some(msg) = inbox.try_recv() {
|
||||
return Some(msg);
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(5));
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Spin up an InferenceActor and wait until it reports WorkerReady.
|
||||
/// Returns (actor address, worker PID).
|
||||
fn spawn_and_wait_ready(
|
||||
rt: &Runtime,
|
||||
status_inbox: &Inbox<InferenceActorStatus>,
|
||||
) -> (ActorAddress, u32) {
|
||||
let sender = rt.create_sender();
|
||||
let actor = InferenceActor::new(echo_worker_spec(), sender)
|
||||
.with_status_addr(*status_inbox.addr());
|
||||
let addr = rt.spawn(actor).unwrap();
|
||||
|
||||
let timeout = Duration::from_secs(5);
|
||||
let mut got_started = false;
|
||||
let mut worker_pid = None;
|
||||
|
||||
let start = Instant::now();
|
||||
while start.elapsed() < timeout {
|
||||
if let Some(status) = tick_until_recv(rt, status_inbox, Duration::from_millis(100)) {
|
||||
match status {
|
||||
InferenceActorStatus::ProcessStarted => got_started = true,
|
||||
InferenceActorStatus::WorkerReady { pid } => {
|
||||
assert!(got_started, "WorkerReady should come after ProcessStarted");
|
||||
worker_pid = pid;
|
||||
break;
|
||||
}
|
||||
other => panic!("unexpected status during startup: {:?}", other),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let pid = worker_pid.expect("worker should report PID within timeout");
|
||||
(addr, pid)
|
||||
}
|
||||
|
||||
fn is_process_alive(pid: u32) -> bool {
|
||||
std::fs::metadata(format!("/proc/{}", pid)).is_ok()
|
||||
}
|
||||
|
||||
// ── Tests ─────────────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn actor_spawns_process_and_receives_started() {
|
||||
let rt = Runtime::new(RuntimeConfig::default());
|
||||
let status_inbox = rt.new_inbox::<InferenceActorStatus>().unwrap();
|
||||
|
||||
let (_addr, pid) = spawn_and_wait_ready(&rt, &status_inbox);
|
||||
|
||||
// The process should be alive
|
||||
assert!(is_process_alive(pid), "worker process should be running");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inference_request_flows_through_process_and_reply_arrives() {
|
||||
let rt = Runtime::new(RuntimeConfig::default());
|
||||
let status_inbox = rt.new_inbox::<InferenceActorStatus>().unwrap();
|
||||
let response_inbox = rt.new_inbox::<InferenceResponse>().unwrap();
|
||||
|
||||
let (addr, _pid) = spawn_and_wait_ready(&rt, &status_inbox);
|
||||
|
||||
// Send an inference request
|
||||
rt.send_to(
|
||||
addr,
|
||||
InferenceActorMsg::Request(InferenceRequest {
|
||||
prompt: "Hello, world!".into(),
|
||||
max_tokens: 8,
|
||||
temperature: 0.7,
|
||||
reply_to: *response_inbox.addr(),
|
||||
}),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let response = tick_until_recv(&rt, &response_inbox, Duration::from_secs(5))
|
||||
.expect("should receive InferenceResponse");
|
||||
|
||||
assert!(
|
||||
response.text.contains("Hello, world!"),
|
||||
"echo worker should reflect the prompt, got: {:?}",
|
||||
response.text
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn worker_crash_is_handled_without_poisoning_runtime() {
|
||||
let rt = Runtime::new(RuntimeConfig::default());
|
||||
let status_inbox = rt.new_inbox::<InferenceActorStatus>().unwrap();
|
||||
let response_inbox = rt.new_inbox::<InferenceResponse>().unwrap();
|
||||
|
||||
let (addr, _pid) = spawn_and_wait_ready(&rt, &status_inbox);
|
||||
|
||||
// Send a request whose prompt triggers an os._exit(1) in the worker
|
||||
rt.send_to(
|
||||
addr,
|
||||
InferenceActorMsg::Request(InferenceRequest {
|
||||
prompt: "__crash__".into(),
|
||||
max_tokens: 1,
|
||||
temperature: 0.0,
|
||||
reply_to: *response_inbox.addr(),
|
||||
}),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// The actor should report ProcessExited
|
||||
let status = tick_until_recv(&rt, &status_inbox, Duration::from_secs(5))
|
||||
.expect("should receive ProcessExited status");
|
||||
match status {
|
||||
InferenceActorStatus::ProcessExited { status } => {
|
||||
assert_ne!(
|
||||
status,
|
||||
ExitStatus::Code(0),
|
||||
"crashed worker should not exit 0"
|
||||
);
|
||||
}
|
||||
other => panic!("expected ProcessExited, got: {:?}", other),
|
||||
}
|
||||
|
||||
// Prove the runtime is still alive: spawn a trivial actor and interact with it
|
||||
#[derive(Clone)]
|
||||
struct Ping {
|
||||
reply_to: ActorAddress,
|
||||
}
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
struct Pong;
|
||||
|
||||
struct PongActor;
|
||||
impl swactor::actor::ActorInterface for PongActor {
|
||||
type Incoming = Ping;
|
||||
type Response = Pong;
|
||||
fn handle(&mut self, ctx: &swactor::runtime::Ctx, msg: Ping) {
|
||||
let _ = ctx.send(msg.reply_to, Pong);
|
||||
}
|
||||
}
|
||||
|
||||
let pong_inbox = rt.new_inbox::<Pong>().unwrap();
|
||||
let pong_addr = rt.spawn(PongActor).unwrap();
|
||||
rt.send_to(pong_addr, Ping { reply_to: *pong_inbox.addr() }).unwrap();
|
||||
|
||||
let pong = tick_until_recv(&rt, &pong_inbox, Duration::from_secs(2));
|
||||
assert_eq!(pong, Some(Pong), "runtime should still be functional after worker crash");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stopping_actor_kills_child_process() {
|
||||
let rt = Runtime::new(RuntimeConfig::default());
|
||||
let status_inbox = rt.new_inbox::<InferenceActorStatus>().unwrap();
|
||||
|
||||
let (addr, pid) = spawn_and_wait_ready(&rt, &status_inbox);
|
||||
assert!(is_process_alive(pid), "worker should be alive before stop");
|
||||
|
||||
// Stop the InferenceActor
|
||||
rt.stop_actor(addr).unwrap();
|
||||
|
||||
// Tick until the child process is gone
|
||||
let start = Instant::now();
|
||||
let timeout = Duration::from_secs(5);
|
||||
while start.elapsed() < timeout {
|
||||
rt.tick();
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
if !is_process_alive(pid) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
assert!(
|
||||
!is_process_alive(pid),
|
||||
"worker process (pid {}) should be dead after actor stop",
|
||||
pid
|
||||
);
|
||||
}
|
||||
164
examples/single-gpu-inference/tests/t_binary.rs
Normal file
164
examples/single-gpu-inference/tests/t_binary.rs
Normal file
|
|
@ -0,0 +1,164 @@
|
|||
//! Binary integration test — validates gpu-node and smoke-run work end-to-end.
|
||||
//!
|
||||
//! Spawns both binaries as child processes on localhost. gpu-node runs an
|
||||
//! InferenceActor backed by a Python worker; smoke-run connects, sends an
|
||||
//! InferenceRequest through the distributed pipeline, and receives the response.
|
||||
//!
|
||||
//! This proves the binaries work through actual process boundaries — the same
|
||||
//! code path used in production, unlike the in-process tests in t_integration.rs.
|
||||
//!
|
||||
//! - `binary_e2e_echo_worker` — fast, uses echo_worker.py (canned echo responses).
|
||||
//! - `binary_e2e_tinygrad` — slow (#[ignore]), uses real tinygrad inference.
|
||||
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::process::{Child, Command, Stdio};
|
||||
use std::time::Duration;
|
||||
|
||||
/// Parse the GPU_NODE_ADDR line from gpu-node's stderr.
|
||||
/// Format: `GPU_NODE_ADDR <hex_node_id> <ip:port,ip:port,...>`
|
||||
fn parse_gpu_node_addr(line: &str) -> Option<(String, String)> {
|
||||
let line = line.trim();
|
||||
if !line.starts_with("GPU_NODE_ADDR ") {
|
||||
return None;
|
||||
}
|
||||
let parts: Vec<&str> = line.splitn(3, ' ').collect();
|
||||
if parts.len() < 3 {
|
||||
return None;
|
||||
}
|
||||
Some((parts[1].to_string(), parts[2].to_string()))
|
||||
}
|
||||
|
||||
/// Start gpu-node, wait for it to print its address, run smoke-run, verify response.
|
||||
fn run_binary_e2e(worker_cmd: &str, worker_script: &str, extra_env: Vec<(&str, &str)>) {
|
||||
let gpu_node_bin = env!("CARGO_BIN_EXE_gpu-node");
|
||||
let smoke_run_bin = env!("CARGO_BIN_EXE_single-gpu-inference");
|
||||
|
||||
// 1. Start gpu-node without SEED_ADDR (it just listens for connections)
|
||||
let mut cmd = Command::new(gpu_node_bin);
|
||||
cmd.env("WORKER_CMD", worker_cmd)
|
||||
.env("WORKER_SCRIPT", worker_script)
|
||||
.env_remove("SEED_ADDR")
|
||||
.stderr(Stdio::piped());
|
||||
for (k, v) in &extra_env {
|
||||
cmd.env(k, v);
|
||||
}
|
||||
let mut gpu_node = cmd.spawn().expect("failed to spawn gpu-node");
|
||||
|
||||
// 2. Read gpu-node stderr in a background thread to find its address
|
||||
// and keep draining so the pipe buffer doesn't fill up.
|
||||
let gpu_stderr = gpu_node.stderr.take().unwrap();
|
||||
let (addr_tx, addr_rx) = std::sync::mpsc::channel::<(String, String)>();
|
||||
let stderr_thread = std::thread::spawn(move || {
|
||||
let reader = BufReader::new(gpu_stderr);
|
||||
let mut sent = false;
|
||||
for line in reader.lines().flatten() {
|
||||
eprintln!("[gpu-node] {}", line);
|
||||
if !sent {
|
||||
if let Some(addr_info) = parse_gpu_node_addr(&line) {
|
||||
let _ = addr_tx.send(addr_info);
|
||||
sent = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// 3. Wait for gpu-node to print its address
|
||||
let (node_id, direct_addrs) = addr_rx
|
||||
.recv_timeout(Duration::from_secs(30))
|
||||
.expect("gpu-node did not print GPU_NODE_ADDR within 30s");
|
||||
eprintln!("gpu-node ready: node_id={node_id}, direct={direct_addrs}");
|
||||
|
||||
// 4. Spawn smoke-run pointing at gpu-node
|
||||
let mut smoke_run = Command::new(smoke_run_bin)
|
||||
.arg("--seed")
|
||||
.arg(&node_id)
|
||||
.arg("--seed-direct")
|
||||
.arg(&direct_addrs)
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()
|
||||
.expect("failed to spawn smoke-run");
|
||||
|
||||
// 5. Wait for smoke-run with timeout (poll every 200ms, bail after 120s)
|
||||
let start = std::time::Instant::now();
|
||||
let timeout = Duration::from_secs(120);
|
||||
loop {
|
||||
match smoke_run.try_wait() {
|
||||
Ok(Some(_)) => break,
|
||||
Ok(None) => {
|
||||
if start.elapsed() > timeout {
|
||||
let _ = smoke_run.kill();
|
||||
let _ = smoke_run.wait();
|
||||
cleanup(&mut gpu_node, stderr_thread);
|
||||
panic!("smoke-run did not exit within {timeout:?}");
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(200));
|
||||
}
|
||||
Err(e) => {
|
||||
cleanup(&mut gpu_node, stderr_thread);
|
||||
panic!("error waiting for smoke-run: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
let output = smoke_run.wait_with_output().expect("failed to read smoke-run output");
|
||||
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let stderr_out = String::from_utf8_lossy(&output.stderr);
|
||||
eprintln!("=== smoke-run stderr ===\n{stderr_out}");
|
||||
eprintln!("=== smoke-run stdout ===\n{stdout}");
|
||||
|
||||
// 6. Clean up gpu-node
|
||||
cleanup(&mut gpu_node, stderr_thread);
|
||||
|
||||
// 7. Verify
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"smoke-run exited with {:?}",
|
||||
output.status
|
||||
);
|
||||
assert!(
|
||||
stdout.contains("=== Inference Response ==="),
|
||||
"stdout should contain response header"
|
||||
);
|
||||
|
||||
// Extract and verify the response text
|
||||
let response_text: String = stdout
|
||||
.lines()
|
||||
.skip_while(|l| *l != "=== Inference Response ===")
|
||||
.skip(1) // skip the header itself
|
||||
.take_while(|l| *l != "==========================")
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
assert!(
|
||||
!response_text.is_empty(),
|
||||
"response text between markers should be non-empty"
|
||||
);
|
||||
eprintln!("response: {response_text:?}");
|
||||
}
|
||||
|
||||
fn cleanup(gpu_node: &mut Child, stderr_thread: std::thread::JoinHandle<()>) {
|
||||
let _ = gpu_node.kill();
|
||||
let _ = gpu_node.wait();
|
||||
let _ = stderr_thread.join();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn binary_e2e_echo_worker() {
|
||||
let manifest_dir = env!("CARGO_MANIFEST_DIR");
|
||||
let script = format!("{manifest_dir}/echo_worker.py");
|
||||
run_binary_e2e("python3", &script, vec![]);
|
||||
}
|
||||
|
||||
/// Full binary e2e with real tinygrad inference (~1B GGUF model).
|
||||
///
|
||||
/// Requires `.venv` with tinygrad installed and downloads a large model.
|
||||
/// Run explicitly with:
|
||||
/// cargo test --package smoke-test binary_e2e_tinygrad -- --ignored
|
||||
#[test]
|
||||
#[ignore]
|
||||
fn binary_e2e_tinygrad() {
|
||||
let manifest_dir = env!("CARGO_MANIFEST_DIR");
|
||||
let python = format!("{manifest_dir}/.venv/bin/python");
|
||||
let script = format!("{manifest_dir}/tinygrad_worker.py");
|
||||
run_binary_e2e(&python, &script, vec![]);
|
||||
}
|
||||
270
examples/single-gpu-inference/tests/t_cluster.rs
Normal file
270
examples/single-gpu-inference/tests/t_cluster.rs
Normal file
|
|
@ -0,0 +1,270 @@
|
|||
//! T-cluster: Cross-node messaging component tests.
|
||||
//!
|
||||
//! Two swactor nodes on localhost via iroh. Tests SWIM convergence,
|
||||
//! actor-level InferenceRequest/InferenceResponse exchange, and
|
||||
//! SWIM death detection after node shutdown.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use distribution::iroh_driver::{IrohDriver, IrohDriverConfig};
|
||||
use distribution::node::DistributedNodeConfig;
|
||||
use distribution::registry::RegistryConfig;
|
||||
use distribution::swim::probe::SwimConfig;
|
||||
use iroh::{PublicKey, RelayMode};
|
||||
|
||||
use swactor::actor::ActorInterface;
|
||||
use swactor::runtime::{Ctx, Runtime, RuntimeConfig};
|
||||
use swactor::transport::TransportRouter;
|
||||
|
||||
use single_gpu_inference::iroh_transport::{drain_actor_messages, IrohActorTransport, ACTOR_ALPN};
|
||||
use single_gpu_inference::messages::{inference_codec_registry, InferenceRequest, InferenceResponse};
|
||||
|
||||
// ── Test SWIM config ─────────────────────────────────────────────────────
|
||||
|
||||
fn test_node_config() -> DistributedNodeConfig {
|
||||
DistributedNodeConfig {
|
||||
swim: SwimConfig {
|
||||
probe_interval: 1,
|
||||
probe_timeout: 3,
|
||||
indirect_probes: 1,
|
||||
suspicion_timeout: 5,
|
||||
dead_reprobe_interval: 0,
|
||||
..SwimConfig::default()
|
||||
},
|
||||
cache_capacity: 100,
|
||||
republish_interval: 50,
|
||||
registry: RegistryConfig::default(),
|
||||
metadata_lambda: 3,
|
||||
}
|
||||
}
|
||||
|
||||
fn make_driver() -> IrohDriver {
|
||||
IrohDriver::new(IrohDriverConfig {
|
||||
secret_key: None,
|
||||
relay_mode: RelayMode::Disabled,
|
||||
node: test_node_config(),
|
||||
peer_auth: None,
|
||||
additional_alpns: vec![ACTOR_ALPN.to_vec()],
|
||||
})
|
||||
.expect("failed to create iroh driver")
|
||||
}
|
||||
|
||||
// ── Pump helpers ─────────────────────────────────────────────────────────
|
||||
|
||||
fn pump_one(driver: &mut IrohDriver) {
|
||||
driver.recv();
|
||||
driver.tick();
|
||||
}
|
||||
|
||||
fn pump_until_pair(
|
||||
a: &mut IrohDriver,
|
||||
b: &mut IrohDriver,
|
||||
timeout: Duration,
|
||||
check_fn: fn(&IrohDriver, &IrohDriver) -> bool,
|
||||
) -> bool {
|
||||
let start = Instant::now();
|
||||
while start.elapsed() < timeout {
|
||||
pump_one(a);
|
||||
pump_one(b);
|
||||
if check_fn(a, b) {
|
||||
return true;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
fn sees_alive(driver: &IrohDriver, peer_key: &PublicKey) -> bool {
|
||||
let snap = driver.snapshot();
|
||||
let peer_hex: String = peer_key
|
||||
.as_bytes()
|
||||
.iter()
|
||||
.map(|b| format!("{:02x}", b))
|
||||
.collect();
|
||||
snap.members
|
||||
.iter()
|
||||
.any(|m| m.node_id == peer_hex && m.state == "alive")
|
||||
}
|
||||
|
||||
fn both_alive(a: &IrohDriver, b: &IrohDriver) -> bool {
|
||||
let a_key = PublicKey::from_bytes(&a.node_id().0).unwrap();
|
||||
let b_key = PublicKey::from_bytes(&b.node_id().0).unwrap();
|
||||
sees_alive(a, &b_key) && sees_alive(b, &a_key)
|
||||
}
|
||||
|
||||
/// Create two IrohDrivers and converge them via seed join.
|
||||
fn make_converged_pair() -> (IrohDriver, IrohDriver) {
|
||||
let mut driver_a = make_driver();
|
||||
let mut driver_b = make_driver();
|
||||
|
||||
let a_addr = driver_a.endpoint_addr();
|
||||
driver_b.join(&[a_addr]);
|
||||
|
||||
let converged = pump_until_pair(
|
||||
&mut driver_a,
|
||||
&mut driver_b,
|
||||
Duration::from_secs(5),
|
||||
both_alive,
|
||||
);
|
||||
assert!(converged, "cluster setup: nodes did not converge within 5s");
|
||||
|
||||
(driver_a, driver_b)
|
||||
}
|
||||
|
||||
// ── Echo actor (replies InferenceResponse for any InferenceRequest) ─────
|
||||
|
||||
struct EchoInferenceActor;
|
||||
|
||||
impl ActorInterface for EchoInferenceActor {
|
||||
type Incoming = InferenceRequest;
|
||||
type Response = InferenceResponse;
|
||||
|
||||
fn handle(&mut self, ctx: &Ctx, msg: InferenceRequest) {
|
||||
let _ = ctx.send(
|
||||
msg.reply_to,
|
||||
InferenceResponse {
|
||||
text: format!("echo: {}", msg.prompt),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ── Tests ────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Node B joins node A via seed address. SWIM converges — both nodes see
|
||||
/// each other alive within 5 seconds.
|
||||
#[test]
|
||||
fn cluster_converges_via_iroh_seed_join() {
|
||||
let mut driver_a = make_driver();
|
||||
let mut driver_b = make_driver();
|
||||
|
||||
let a_addr = driver_a.endpoint_addr();
|
||||
driver_b.join(&[a_addr]);
|
||||
|
||||
let converged = pump_until_pair(
|
||||
&mut driver_a,
|
||||
&mut driver_b,
|
||||
Duration::from_secs(5),
|
||||
both_alive,
|
||||
);
|
||||
|
||||
assert!(converged, "nodes did not converge within 5s");
|
||||
assert_eq!(driver_a.snapshot().alive_count, 1);
|
||||
assert_eq!(driver_b.snapshot().alive_count, 1);
|
||||
|
||||
driver_a.shutdown();
|
||||
driver_b.shutdown();
|
||||
}
|
||||
|
||||
/// Actor on node A sends InferenceRequest to actor on node B via the
|
||||
/// transport router + codec. InferenceResponse arrives back at node A.
|
||||
#[test]
|
||||
fn inference_request_roundtrips_across_two_nodes() {
|
||||
let (mut driver_a, mut driver_b) = make_converged_pair();
|
||||
|
||||
let codecs = Arc::new(inference_codec_registry());
|
||||
|
||||
let mut rt_a = Runtime::new(RuntimeConfig::default());
|
||||
let mut rt_b = Runtime::new(RuntimeConfig::default());
|
||||
|
||||
// Spawn echo actor on node B
|
||||
let echo_addr = rt_b.spawn(EchoInferenceActor).unwrap();
|
||||
rt_b.tick();
|
||||
|
||||
// Inbox on node A for responses
|
||||
let response_inbox = rt_a.new_inbox::<InferenceResponse>().unwrap();
|
||||
let inbox_addr = *response_inbox.addr();
|
||||
|
||||
// Build iroh-backed transports for actor messages
|
||||
let transport_a_to_b = Arc::new(IrohActorTransport::new(
|
||||
driver_a.endpoint().clone(),
|
||||
driver_b.endpoint_addr(),
|
||||
driver_a.tokio_handle(),
|
||||
));
|
||||
let transport_b_to_a = Arc::new(IrohActorTransport::new(
|
||||
driver_b.endpoint().clone(),
|
||||
driver_a.endpoint_addr(),
|
||||
driver_b.tokio_handle(),
|
||||
));
|
||||
|
||||
// Wire routes: A knows echo_addr is on B, B knows inbox_addr is on A
|
||||
let router_a = TransportRouter::new();
|
||||
router_a.add_route(echo_addr, transport_a_to_b);
|
||||
let router_b = TransportRouter::new();
|
||||
router_b.add_route(inbox_addr, transport_b_to_a);
|
||||
|
||||
rt_a.set_codec_registry(codecs.clone());
|
||||
rt_a.set_transport_router(Arc::new(router_a));
|
||||
rt_b.set_codec_registry(codecs.clone());
|
||||
rt_b.set_transport_router(Arc::new(router_b));
|
||||
|
||||
// Send InferenceRequest from node A → actor on node B
|
||||
rt_a.send_to(
|
||||
echo_addr,
|
||||
InferenceRequest {
|
||||
prompt: "Hello from node A".into(),
|
||||
max_tokens: 8,
|
||||
temperature: 0.7,
|
||||
reply_to: inbox_addr,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Allow iroh transport to deliver, then drain into runtime B
|
||||
std::thread::sleep(Duration::from_millis(200));
|
||||
drain_actor_messages(&driver_b, &codecs, &rt_b, Duration::from_millis(500));
|
||||
rt_b.tick();
|
||||
|
||||
// Actor replied — allow transport to deliver, then drain into runtime A
|
||||
std::thread::sleep(Duration::from_millis(200));
|
||||
drain_actor_messages(&driver_a, &codecs, &rt_a, Duration::from_millis(500));
|
||||
|
||||
let response = response_inbox
|
||||
.try_recv()
|
||||
.expect("InferenceResponse should arrive at node A");
|
||||
assert!(
|
||||
response.text.contains("Hello from node A"),
|
||||
"expected echo of prompt, got: {:?}",
|
||||
response.text
|
||||
);
|
||||
|
||||
driver_a.shutdown();
|
||||
driver_b.shutdown();
|
||||
}
|
||||
|
||||
/// When node B shuts down, node A detects the death via SWIM within the
|
||||
/// configured suspicion window.
|
||||
#[test]
|
||||
fn node_death_detected_via_swim_after_shutdown() {
|
||||
let (mut driver_a, mut driver_b) = make_converged_pair();
|
||||
|
||||
assert_eq!(
|
||||
driver_a.snapshot().alive_count, 1,
|
||||
"precondition: A sees B alive"
|
||||
);
|
||||
|
||||
// Kill node B
|
||||
driver_b.shutdown();
|
||||
|
||||
// Pump node A until it sees zero alive peers
|
||||
let start = Instant::now();
|
||||
let timeout = Duration::from_secs(10);
|
||||
let mut detected = false;
|
||||
while start.elapsed() < timeout {
|
||||
driver_a.recv();
|
||||
driver_a.tick();
|
||||
if driver_a.snapshot().alive_count == 0 {
|
||||
detected = true;
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
|
||||
assert!(
|
||||
detected,
|
||||
"node A should detect node B's death via SWIM within the suspicion window"
|
||||
);
|
||||
|
||||
driver_a.shutdown();
|
||||
}
|
||||
153
examples/single-gpu-inference/tests/t_codec.rs
Normal file
153
examples/single-gpu-inference/tests/t_codec.rs
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
//! T-codec: InferenceRequest / InferenceResponse serialization tests.
|
||||
//!
|
||||
//! Validates the message contract that all other smoke test groups depend on:
|
||||
//! - Roundtrip fidelity for both message types
|
||||
//! - Graceful error handling on corrupted input
|
||||
|
||||
use single_gpu_inference::messages::{
|
||||
inference_codec_registry, InferenceCodec, InferenceRequest, InferenceResponse,
|
||||
};
|
||||
use swactor::actor::ActorAddress;
|
||||
use swactor::transport::Codec;
|
||||
|
||||
fn request_codec() -> &'static dyn Codec<InferenceRequest> {
|
||||
&InferenceCodec
|
||||
}
|
||||
|
||||
fn response_codec() -> &'static dyn Codec<InferenceResponse> {
|
||||
&InferenceCodec
|
||||
}
|
||||
|
||||
// ─── Roundtrip Tests ───────────────────────────────────────────────────────
|
||||
|
||||
/// A request with representative field values survives encode → decode
|
||||
/// through both the direct codec and the type-erased registry path.
|
||||
#[test]
|
||||
fn inference_request_roundtrips_through_codec_and_registry() {
|
||||
let original = InferenceRequest {
|
||||
prompt: "Tell me about distributed systems".into(),
|
||||
max_tokens: 128,
|
||||
temperature: 0.7,
|
||||
reply_to: ActorAddress::new_random(),
|
||||
};
|
||||
|
||||
// Direct codec path
|
||||
let codec = request_codec();
|
||||
let bytes = codec.encode(&original).expect("encode should succeed");
|
||||
let decoded = codec.decode(&bytes).expect("decode should succeed");
|
||||
assert_eq!(decoded, original);
|
||||
|
||||
// Registry (type-erased) path — encode via TypeId, decode via type_tag
|
||||
let registry = inference_codec_registry();
|
||||
let (tag, payload) = registry
|
||||
.encode(
|
||||
std::any::TypeId::of::<InferenceRequest>(),
|
||||
Box::new(original.clone()),
|
||||
)
|
||||
.expect("registry encode should succeed");
|
||||
assert_eq!(tag, "smoke::InferenceRequest");
|
||||
|
||||
let any_msg = registry
|
||||
.decode(&tag, &payload)
|
||||
.expect("registry decode should succeed");
|
||||
let decoded = any_msg
|
||||
.downcast::<InferenceRequest>()
|
||||
.expect("downcast should succeed");
|
||||
assert_eq!(*decoded, original);
|
||||
}
|
||||
|
||||
/// A response roundtrips through both the direct codec and the registry.
|
||||
#[test]
|
||||
fn inference_response_roundtrips_through_codec_and_registry() {
|
||||
let original = InferenceResponse {
|
||||
text: "Hello! I'd be happy to discuss distributed systems.".into(),
|
||||
};
|
||||
|
||||
let codec = response_codec();
|
||||
let bytes = codec.encode(&original).expect("encode should succeed");
|
||||
let decoded = codec.decode(&bytes).expect("decode should succeed");
|
||||
assert_eq!(decoded, original);
|
||||
|
||||
let registry = inference_codec_registry();
|
||||
let (tag, payload) = registry
|
||||
.encode(
|
||||
std::any::TypeId::of::<InferenceResponse>(),
|
||||
Box::new(original.clone()),
|
||||
)
|
||||
.expect("registry encode should succeed");
|
||||
assert_eq!(tag, "smoke::InferenceResponse");
|
||||
|
||||
let any_msg = registry
|
||||
.decode(&tag, &payload)
|
||||
.expect("registry decode should succeed");
|
||||
let decoded = any_msg
|
||||
.downcast::<InferenceResponse>()
|
||||
.expect("downcast should succeed");
|
||||
assert_eq!(*decoded, original);
|
||||
}
|
||||
|
||||
// ─── Corruption Tests ──────────────────────────────────────────────────────
|
||||
|
||||
/// Completely random bytes are not valid JSON — the decoder must return
|
||||
/// an error rather than panicking or producing a garbage message.
|
||||
#[test]
|
||||
fn corrupted_bytes_produce_error_for_request() {
|
||||
let codec = request_codec();
|
||||
let garbage: Vec<u8> = vec![0xFF, 0x00, 0xDE, 0xAD, 0xBE, 0xEF];
|
||||
let result = codec.decode(&garbage);
|
||||
assert!(result.is_err(), "garbage bytes must produce an error");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn corrupted_bytes_produce_error_for_response() {
|
||||
let codec = response_codec();
|
||||
let garbage: Vec<u8> = vec![0xFF, 0x00, 0xDE, 0xAD, 0xBE, 0xEF];
|
||||
let result = codec.decode(&garbage);
|
||||
assert!(result.is_err(), "garbage bytes must produce an error");
|
||||
}
|
||||
|
||||
/// Truncated payload — valid JSON prefix cut short mid-value.
|
||||
#[test]
|
||||
fn truncated_request_bytes_produce_error() {
|
||||
let codec = request_codec();
|
||||
let original = InferenceRequest {
|
||||
prompt: "hello".into(),
|
||||
max_tokens: 64,
|
||||
temperature: 0.5,
|
||||
reply_to: ActorAddress::new_random(),
|
||||
};
|
||||
let bytes = codec.encode(&original).unwrap();
|
||||
// Chop off the last half
|
||||
let truncated = &bytes[..bytes.len() / 2];
|
||||
let result = codec.decode(truncated);
|
||||
assert!(result.is_err(), "truncated bytes must produce an error");
|
||||
}
|
||||
|
||||
/// Empty input — zero bytes is not valid JSON.
|
||||
#[test]
|
||||
fn empty_bytes_produce_error() {
|
||||
let req_codec = request_codec();
|
||||
let res_codec = response_codec();
|
||||
assert!(req_codec.decode(&[]).is_err());
|
||||
assert!(res_codec.decode(&[]).is_err());
|
||||
}
|
||||
|
||||
/// Valid JSON but wrong schema — a response payload fed to the request
|
||||
/// decoder. The missing required fields must cause an error.
|
||||
#[test]
|
||||
fn wrong_message_type_produces_error() {
|
||||
let res_codec = response_codec();
|
||||
let req_codec = request_codec();
|
||||
|
||||
let response = InferenceResponse {
|
||||
text: "oops".into(),
|
||||
};
|
||||
let response_bytes = res_codec.encode(&response).unwrap();
|
||||
|
||||
// Decoding response bytes as a request should fail (missing prompt, max_tokens, etc.)
|
||||
let result = req_codec.decode(&response_bytes);
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"decoding response bytes as request must fail"
|
||||
);
|
||||
}
|
||||
415
examples/single-gpu-inference/tests/t_integration.rs
Normal file
415
examples/single-gpu-inference/tests/t_integration.rs
Normal file
|
|
@ -0,0 +1,415 @@
|
|||
//! T-integration: End-to-end distributed inference tests.
|
||||
//!
|
||||
//! Two swactor nodes on localhost via iroh/QUIC. Node B runs an
|
||||
//! `InferenceActor` backed by a Python worker. Node A sends an
|
||||
//! `InferenceRequest` across the network, through the `RequestBridge`,
|
||||
//! into the actor, through the child process, and back.
|
||||
//!
|
||||
//! - `distributed_inference_through_echo_worker` — fast, uses canned echo responses.
|
||||
//! - `distributed_inference_through_tinygrad` — slow (#[ignore]), downloads a real
|
||||
//! ~1B GGUF model and runs real tinygrad inference on CPU.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use distribution::iroh_driver::{IrohDriver, IrohDriverConfig};
|
||||
use distribution::node::DistributedNodeConfig;
|
||||
use distribution::registry::RegistryConfig;
|
||||
use distribution::swim::probe::SwimConfig;
|
||||
use iroh::{PublicKey, RelayMode};
|
||||
|
||||
use swactor::runtime::{Runtime, RuntimeConfig};
|
||||
use swactor::transport::TransportRouter;
|
||||
|
||||
use single_gpu_inference::inference_actor::{InferenceActor, InferenceActorStatus, RequestBridge};
|
||||
use single_gpu_inference::iroh_transport::{drain_actor_messages, IrohActorTransport, ACTOR_ALPN};
|
||||
use single_gpu_inference::messages::{inference_codec_registry, InferenceRequest, InferenceResponse};
|
||||
use swactor_process::{ProcessMode, ProcessSpec};
|
||||
|
||||
// ── Iroh helpers (from t_cluster.rs) ────────────────────────────────────
|
||||
|
||||
fn test_node_config() -> DistributedNodeConfig {
|
||||
DistributedNodeConfig {
|
||||
swim: SwimConfig {
|
||||
probe_interval: 1,
|
||||
probe_timeout: 3,
|
||||
indirect_probes: 1,
|
||||
suspicion_timeout: 5,
|
||||
dead_reprobe_interval: 0,
|
||||
..SwimConfig::default()
|
||||
},
|
||||
cache_capacity: 100,
|
||||
republish_interval: 50,
|
||||
registry: RegistryConfig::default(),
|
||||
metadata_lambda: 3,
|
||||
}
|
||||
}
|
||||
|
||||
fn make_driver() -> IrohDriver {
|
||||
IrohDriver::new(IrohDriverConfig {
|
||||
secret_key: None,
|
||||
relay_mode: RelayMode::Disabled,
|
||||
node: test_node_config(),
|
||||
peer_auth: None,
|
||||
additional_alpns: vec![ACTOR_ALPN.to_vec()],
|
||||
})
|
||||
.expect("failed to create iroh driver")
|
||||
}
|
||||
|
||||
fn pump_one(driver: &mut IrohDriver) {
|
||||
driver.recv();
|
||||
driver.tick();
|
||||
}
|
||||
|
||||
fn pump_until_pair(
|
||||
a: &mut IrohDriver,
|
||||
b: &mut IrohDriver,
|
||||
timeout: Duration,
|
||||
check_fn: fn(&IrohDriver, &IrohDriver) -> bool,
|
||||
) -> bool {
|
||||
let start = Instant::now();
|
||||
while start.elapsed() < timeout {
|
||||
pump_one(a);
|
||||
pump_one(b);
|
||||
if check_fn(a, b) {
|
||||
return true;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
fn sees_alive(driver: &IrohDriver, peer_key: &PublicKey) -> bool {
|
||||
let snap = driver.snapshot();
|
||||
let peer_hex: String = peer_key
|
||||
.as_bytes()
|
||||
.iter()
|
||||
.map(|b| format!("{:02x}", b))
|
||||
.collect();
|
||||
snap.members
|
||||
.iter()
|
||||
.any(|m| m.node_id == peer_hex && m.state == "alive")
|
||||
}
|
||||
|
||||
fn both_alive(a: &IrohDriver, b: &IrohDriver) -> bool {
|
||||
let a_key = PublicKey::from_bytes(&a.node_id().0).unwrap();
|
||||
let b_key = PublicKey::from_bytes(&b.node_id().0).unwrap();
|
||||
sees_alive(a, &b_key) && sees_alive(b, &a_key)
|
||||
}
|
||||
|
||||
fn make_converged_pair() -> (IrohDriver, IrohDriver) {
|
||||
let mut driver_a = make_driver();
|
||||
let mut driver_b = make_driver();
|
||||
|
||||
let a_addr = driver_a.endpoint_addr();
|
||||
driver_b.join(&[a_addr]);
|
||||
|
||||
let converged = pump_until_pair(
|
||||
&mut driver_a,
|
||||
&mut driver_b,
|
||||
Duration::from_secs(5),
|
||||
both_alive,
|
||||
);
|
||||
assert!(converged, "cluster setup: nodes did not converge within 5s");
|
||||
|
||||
(driver_a, driver_b)
|
||||
}
|
||||
|
||||
// ── Process helpers (from t_actor.rs) ────────────────────────────────────
|
||||
|
||||
fn echo_worker_spec() -> ProcessSpec {
|
||||
ProcessSpec {
|
||||
command: "python3".into(),
|
||||
args: vec![format!("{}/echo_worker.py", env!("CARGO_MANIFEST_DIR"))],
|
||||
env: HashMap::new(),
|
||||
working_dir: None,
|
||||
mode: ProcessMode::Automated,
|
||||
initial_pty_size: None,
|
||||
kill_timeout: Some(Duration::from_secs(2)),
|
||||
stdin_buffer_limit: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn tinygrad_worker_spec() -> ProcessSpec {
|
||||
let manifest = env!("CARGO_MANIFEST_DIR");
|
||||
ProcessSpec {
|
||||
command: format!("{manifest}/.venv/bin/python"),
|
||||
args: vec![format!("{manifest}/tinygrad_worker.py")],
|
||||
env: HashMap::new(),
|
||||
working_dir: None,
|
||||
mode: ProcessMode::Automated,
|
||||
initial_pty_size: None,
|
||||
kill_timeout: Some(Duration::from_secs(5)),
|
||||
stdin_buffer_limit: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_process_alive(pid: u32) -> bool {
|
||||
std::fs::metadata(format!("/proc/{}", pid)).is_ok()
|
||||
}
|
||||
|
||||
// ── Test ─────────────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn distributed_inference_through_echo_worker() {
|
||||
// 1. Converge two iroh drivers
|
||||
let (mut driver_a, mut driver_b) = make_converged_pair();
|
||||
|
||||
let codecs = Arc::new(inference_codec_registry());
|
||||
|
||||
// 2. Create runtimes
|
||||
let mut rt_a = Runtime::new(RuntimeConfig::default());
|
||||
let mut rt_b = Runtime::new(RuntimeConfig::default());
|
||||
|
||||
// 3. On rt_b: spawn InferenceActor with echo_worker + RequestBridge
|
||||
let status_inbox = rt_b.new_inbox::<InferenceActorStatus>().unwrap();
|
||||
let sender = rt_b.create_sender();
|
||||
let actor = InferenceActor::new(echo_worker_spec(), sender)
|
||||
.with_status_addr(*status_inbox.addr());
|
||||
let inference_addr = rt_b.spawn(actor).unwrap();
|
||||
|
||||
let bridge = RequestBridge { target: inference_addr };
|
||||
let bridge_addr = rt_b.spawn(bridge).unwrap();
|
||||
|
||||
// 4. On rt_a: create response inbox
|
||||
let response_inbox = rt_a.new_inbox::<InferenceResponse>().unwrap();
|
||||
let inbox_addr = *response_inbox.addr();
|
||||
|
||||
// 5. Build iroh transports
|
||||
let transport_a_to_b = Arc::new(IrohActorTransport::new(
|
||||
driver_a.endpoint().clone(),
|
||||
driver_b.endpoint_addr(),
|
||||
driver_a.tokio_handle(),
|
||||
));
|
||||
let transport_b_to_a = Arc::new(IrohActorTransport::new(
|
||||
driver_b.endpoint().clone(),
|
||||
driver_a.endpoint_addr(),
|
||||
driver_b.tokio_handle(),
|
||||
));
|
||||
|
||||
// Wire routes: A sends to bridge_addr on B, B sends to inbox_addr on A
|
||||
let router_a = TransportRouter::new();
|
||||
router_a.add_route(bridge_addr, transport_a_to_b);
|
||||
let router_b = TransportRouter::new();
|
||||
router_b.add_route(inbox_addr, transport_b_to_a);
|
||||
|
||||
// 6. Install codecs and transport routers
|
||||
rt_a.set_codec_registry(codecs.clone());
|
||||
rt_a.set_transport_router(Arc::new(router_a));
|
||||
rt_b.set_codec_registry(codecs.clone());
|
||||
rt_b.set_transport_router(Arc::new(router_b));
|
||||
|
||||
// 7. Tick rt_b until InferenceActor reports WorkerReady
|
||||
let start = Instant::now();
|
||||
let mut worker_pid = None;
|
||||
while start.elapsed() < Duration::from_secs(10) {
|
||||
rt_b.tick();
|
||||
if let Some(status) = status_inbox.try_recv() {
|
||||
match status {
|
||||
InferenceActorStatus::WorkerReady { pid } => {
|
||||
worker_pid = pid;
|
||||
break;
|
||||
}
|
||||
InferenceActorStatus::ProcessStarted => {}
|
||||
other => panic!("unexpected status during startup: {:?}", other),
|
||||
}
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(5));
|
||||
}
|
||||
let worker_pid = worker_pid.expect("echo_worker.py should report ready within 10s");
|
||||
|
||||
// 8. Send InferenceRequest from node A → bridge on node B
|
||||
rt_a.send_to(
|
||||
bridge_addr,
|
||||
InferenceRequest {
|
||||
prompt: "Hello from node A".into(),
|
||||
max_tokens: 8,
|
||||
temperature: 0.7,
|
||||
reply_to: inbox_addr,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// 9. Pump loop: drain messages on both sides, tick both runtimes
|
||||
let start = Instant::now();
|
||||
let mut got_response = false;
|
||||
while start.elapsed() < Duration::from_secs(10) {
|
||||
// Let transport deliver
|
||||
std::thread::sleep(Duration::from_millis(100));
|
||||
|
||||
// Drain incoming actor messages on both sides
|
||||
drain_actor_messages(&driver_b, &codecs, &rt_b, Duration::from_millis(100));
|
||||
drain_actor_messages(&driver_a, &codecs, &rt_a, Duration::from_millis(100));
|
||||
|
||||
// Tick both runtimes
|
||||
rt_b.tick();
|
||||
rt_a.tick();
|
||||
|
||||
// Check for response
|
||||
if let Some(response) = response_inbox.try_recv() {
|
||||
// 10. Assert
|
||||
assert!(
|
||||
response.text.contains("Hello from node A"),
|
||||
"expected echo of prompt, got: {:?}",
|
||||
response.text
|
||||
);
|
||||
got_response = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
assert!(got_response, "should receive InferenceResponse within 10s");
|
||||
|
||||
// 11. Cleanup: stop inference actor, wait for process death, shutdown drivers
|
||||
rt_b.stop_actor(inference_addr).unwrap();
|
||||
|
||||
let start = Instant::now();
|
||||
while start.elapsed() < Duration::from_secs(5) {
|
||||
rt_b.tick();
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
if !is_process_alive(worker_pid) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
!is_process_alive(worker_pid),
|
||||
"echo_worker.py (pid {}) should be dead after actor stop",
|
||||
worker_pid
|
||||
);
|
||||
|
||||
driver_a.shutdown();
|
||||
driver_b.shutdown();
|
||||
}
|
||||
|
||||
/// Full distributed inference through tinygrad with a real ~1B GGUF model.
|
||||
///
|
||||
/// Ignored by default — requires the `.venv` with tinygrad installed and
|
||||
/// downloads a ~1 GB model on first run. Run explicitly with:
|
||||
/// cargo test --package smoke-test distributed_inference -- --ignored
|
||||
#[test]
|
||||
#[ignore]
|
||||
fn distributed_inference_through_tinygrad() {
|
||||
// 1. Converge two iroh drivers
|
||||
let (mut driver_a, mut driver_b) = make_converged_pair();
|
||||
|
||||
let codecs = Arc::new(inference_codec_registry());
|
||||
|
||||
// 2. Create runtimes
|
||||
let mut rt_a = Runtime::new(RuntimeConfig::default());
|
||||
let mut rt_b = Runtime::new(RuntimeConfig::default());
|
||||
|
||||
// 3. On rt_b: spawn InferenceActor with tinygrad_worker + RequestBridge
|
||||
let status_inbox = rt_b.new_inbox::<InferenceActorStatus>().unwrap();
|
||||
let sender = rt_b.create_sender();
|
||||
let actor = InferenceActor::new(tinygrad_worker_spec(), sender)
|
||||
.with_status_addr(*status_inbox.addr());
|
||||
let inference_addr = rt_b.spawn(actor).unwrap();
|
||||
|
||||
let bridge = RequestBridge { target: inference_addr };
|
||||
let bridge_addr = rt_b.spawn(bridge).unwrap();
|
||||
|
||||
// 4. On rt_a: create response inbox
|
||||
let response_inbox = rt_a.new_inbox::<InferenceResponse>().unwrap();
|
||||
let inbox_addr = *response_inbox.addr();
|
||||
|
||||
// 5. Build iroh transports
|
||||
let transport_a_to_b = Arc::new(IrohActorTransport::new(
|
||||
driver_a.endpoint().clone(),
|
||||
driver_b.endpoint_addr(),
|
||||
driver_a.tokio_handle(),
|
||||
));
|
||||
let transport_b_to_a = Arc::new(IrohActorTransport::new(
|
||||
driver_b.endpoint().clone(),
|
||||
driver_a.endpoint_addr(),
|
||||
driver_b.tokio_handle(),
|
||||
));
|
||||
|
||||
let router_a = TransportRouter::new();
|
||||
router_a.add_route(bridge_addr, transport_a_to_b);
|
||||
let router_b = TransportRouter::new();
|
||||
router_b.add_route(inbox_addr, transport_b_to_a);
|
||||
|
||||
// 6. Install codecs and transport routers
|
||||
rt_a.set_codec_registry(codecs.clone());
|
||||
rt_a.set_transport_router(Arc::new(router_a));
|
||||
rt_b.set_codec_registry(codecs.clone());
|
||||
rt_b.set_transport_router(Arc::new(router_b));
|
||||
|
||||
// 7. Wait for tinygrad model download + load (generous timeout)
|
||||
let start = Instant::now();
|
||||
let mut worker_pid = None;
|
||||
while start.elapsed() < Duration::from_secs(600) {
|
||||
rt_b.tick();
|
||||
if let Some(status) = status_inbox.try_recv() {
|
||||
match status {
|
||||
InferenceActorStatus::WorkerReady { pid } => {
|
||||
worker_pid = pid;
|
||||
break;
|
||||
}
|
||||
InferenceActorStatus::ProcessStarted => {}
|
||||
other => panic!("unexpected status during startup: {:?}", other),
|
||||
}
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(50));
|
||||
}
|
||||
let worker_pid = worker_pid.expect("tinygrad_worker.py should report ready (model loaded)");
|
||||
|
||||
// 8. Send InferenceRequest from node A → bridge on node B
|
||||
rt_a.send_to(
|
||||
bridge_addr,
|
||||
InferenceRequest {
|
||||
prompt: "Say hello".into(),
|
||||
max_tokens: 32,
|
||||
temperature: 0.7,
|
||||
reply_to: inbox_addr,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// 9. Pump loop — generation on CPU can be slow
|
||||
let start = Instant::now();
|
||||
let mut got_response = false;
|
||||
while start.elapsed() < Duration::from_secs(300) {
|
||||
std::thread::sleep(Duration::from_millis(200));
|
||||
|
||||
drain_actor_messages(&driver_b, &codecs, &rt_b, Duration::from_millis(100));
|
||||
drain_actor_messages(&driver_a, &codecs, &rt_a, Duration::from_millis(100));
|
||||
|
||||
rt_b.tick();
|
||||
rt_a.tick();
|
||||
|
||||
if let Some(response) = response_inbox.try_recv() {
|
||||
assert!(
|
||||
!response.text.is_empty(),
|
||||
"expected non-empty generated text from tinygrad, got empty string"
|
||||
);
|
||||
eprintln!("tinygrad response: {:?}", response.text);
|
||||
got_response = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
got_response,
|
||||
"should receive InferenceResponse from tinygrad within timeout"
|
||||
);
|
||||
|
||||
// 10. Cleanup
|
||||
rt_b.stop_actor(inference_addr).unwrap();
|
||||
|
||||
let start = Instant::now();
|
||||
while start.elapsed() < Duration::from_secs(10) {
|
||||
rt_b.tick();
|
||||
std::thread::sleep(Duration::from_millis(50));
|
||||
if !is_process_alive(worker_pid) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
!is_process_alive(worker_pid),
|
||||
"tinygrad_worker.py (pid {}) should be dead after actor stop",
|
||||
worker_pid
|
||||
);
|
||||
|
||||
driver_a.shutdown();
|
||||
driver_b.shutdown();
|
||||
}
|
||||
235
examples/single-gpu-inference/tests/t_vastai.rs
Normal file
235
examples/single-gpu-inference/tests/t_vastai.rs
Normal file
|
|
@ -0,0 +1,235 @@
|
|||
//! T-vastai: vast.ai API client tests with mocked HTTP responses.
|
||||
//!
|
||||
//! Validates the orchestration layer that rents and manages GPU instances.
|
||||
//! All tests hit a local wiremock server — no real API calls, no VAST_API_KEY needed.
|
||||
|
||||
use reqwest::Client;
|
||||
use single_gpu_inference::vastai;
|
||||
use std::time::Duration;
|
||||
use wiremock::matchers::{header, method, path_regex};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
// ─── find_offer ───────────────────────────────────────────────────────────
|
||||
|
||||
/// Three offers at different prices — the client must pick the cheapest one.
|
||||
#[tokio::test]
|
||||
async fn find_offer_picks_cheapest_from_multiple_offers() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
Mock::given(method("GET"))
|
||||
.and(path_regex("/api/v0/bundles.*"))
|
||||
.and(header("Authorization", "Bearer test-key"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"offers": [
|
||||
{ "id": 101, "gpu_name": "RTX 3090", "dph_total": 0.30, "geolocation": "US" },
|
||||
{ "id": 102, "gpu_name": "RTX 3090", "dph_total": 0.15, "geolocation": "US" },
|
||||
{ "id": 103, "gpu_name": "RTX 3090", "dph_total": 0.25, "geolocation": "US" },
|
||||
]
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let client = Client::new();
|
||||
let offer = vastai::find_offer(&client, &server.uri(), "test-key", "RTX 3090", &[])
|
||||
.await
|
||||
.expect("should find an offer");
|
||||
|
||||
assert_eq!(offer.id, 102, "must pick the cheapest offer (id=102, $0.15/hr)");
|
||||
assert!((offer.dph_total - 0.15).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
/// Empty offer list — the client must return an error, not panic.
|
||||
#[tokio::test]
|
||||
async fn find_offer_returns_error_on_empty_list() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
Mock::given(method("GET"))
|
||||
.and(path_regex("/api/v0/bundles.*"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"offers": []
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let client = Client::new();
|
||||
let result = vastai::find_offer(&client, &server.uri(), "test-key", "RTX 3090", &[]).await;
|
||||
|
||||
assert!(result.is_err(), "empty offer list must produce an error");
|
||||
assert!(
|
||||
result.unwrap_err().contains("no offers"),
|
||||
"error should mention no offers"
|
||||
);
|
||||
}
|
||||
|
||||
// ─── create_instance ──────────────────────────────────────────────────────
|
||||
|
||||
/// Successful creation — parse the new_contract ID and verify SEED_ADDR is in the request.
|
||||
#[tokio::test]
|
||||
async fn create_instance_parses_contract_and_sends_seed_addr() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
Mock::given(method("PUT"))
|
||||
.and(path_regex("/api/v0/asks/102/"))
|
||||
.and(header("Authorization", "Bearer test-key"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"new_contract": 9999,
|
||||
"success": true
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let client = Client::new();
|
||||
let info = vastai::create_instance(
|
||||
&client,
|
||||
&server.uri(),
|
||||
"test-key",
|
||||
102,
|
||||
"iroh://node-abc123",
|
||||
None,
|
||||
"swactor-gpu:latest",
|
||||
)
|
||||
.await
|
||||
.expect("create should succeed");
|
||||
|
||||
assert_eq!(info.contract_id, 9999);
|
||||
}
|
||||
|
||||
/// Verify the request body contains SEED_ADDR by inspecting the recorded request.
|
||||
#[tokio::test]
|
||||
async fn create_instance_includes_seed_addr_in_env_payload() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
Mock::given(method("PUT"))
|
||||
.and(path_regex("/api/v0/asks/55/"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"new_contract": 7777,
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let client = Client::new();
|
||||
vastai::create_instance(
|
||||
&client,
|
||||
&server.uri(),
|
||||
"test-key",
|
||||
55,
|
||||
"iroh://seed-address-xyz",
|
||||
None,
|
||||
"swactor-gpu:latest",
|
||||
)
|
||||
.await
|
||||
.expect("create should succeed");
|
||||
|
||||
// Inspect the recorded request to verify the body contains SEED_ADDR.
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
assert_eq!(requests.len(), 1);
|
||||
let body: serde_json::Value = serde_json::from_slice(&requests[0].body).unwrap();
|
||||
assert_eq!(
|
||||
body["env"]["SEED_ADDR"],
|
||||
"iroh://seed-address-xyz",
|
||||
"request body must contain SEED_ADDR in the env payload"
|
||||
);
|
||||
}
|
||||
|
||||
// ─── wait_for_running ─────────────────────────────────────────────────────
|
||||
|
||||
/// Polling sequence: loading → loading → running. Must extract IP and port.
|
||||
#[tokio::test]
|
||||
async fn wait_for_running_handles_loading_then_running_sequence() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
// First two polls: loading
|
||||
Mock::given(method("GET"))
|
||||
.and(path_regex("/api/v0/instances/9999/"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"instances": { "actual_status": "loading" }
|
||||
})))
|
||||
.up_to_n_times(2)
|
||||
.expect(2)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
// Third poll: running with IP and port
|
||||
Mock::given(method("GET"))
|
||||
.and(path_regex("/api/v0/instances/9999/"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"instances": {
|
||||
"actual_status": "running",
|
||||
"public_ipaddr": "203.0.113.42",
|
||||
"ssh_port": 31337
|
||||
}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let client = Client::new();
|
||||
let running = vastai::wait_for_running(
|
||||
&client,
|
||||
&server.uri(),
|
||||
"test-key",
|
||||
9999,
|
||||
Duration::from_millis(10), // fast polling for tests
|
||||
10,
|
||||
)
|
||||
.await
|
||||
.expect("should eventually reach running");
|
||||
|
||||
assert_eq!(running.ip, "203.0.113.42");
|
||||
assert_eq!(running.port, 31337);
|
||||
}
|
||||
|
||||
/// Terminal status (exited) — must return error immediately without further polling.
|
||||
#[tokio::test]
|
||||
async fn wait_for_running_returns_error_on_exited_status() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
Mock::given(method("GET"))
|
||||
.and(path_regex("/api/v0/instances/1234/"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"instances": { "actual_status": "exited" }
|
||||
})))
|
||||
.expect(1) // Must only be called once — no further polling after terminal status.
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let client = Client::new();
|
||||
let result = vastai::wait_for_running(
|
||||
&client,
|
||||
&server.uri(),
|
||||
"test-key",
|
||||
1234,
|
||||
Duration::from_millis(10),
|
||||
10,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(result.is_err(), "terminal status must produce an error");
|
||||
assert!(
|
||||
result.unwrap_err().contains("terminal status"),
|
||||
"error should mention terminal status"
|
||||
);
|
||||
}
|
||||
|
||||
// ─── destroy_instance ─────────────────────────────────────────────────────
|
||||
|
||||
/// Verify the correct DELETE request is sent for instance teardown.
|
||||
#[tokio::test]
|
||||
async fn destroy_instance_sends_correct_delete_request() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
Mock::given(method("DELETE"))
|
||||
.and(path_regex("/api/v0/instances/9999/"))
|
||||
.and(header("Authorization", "Bearer test-key"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"success": true
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let client = Client::new();
|
||||
vastai::destroy_instance(&client, &server.uri(), "test-key", 9999)
|
||||
.await
|
||||
.expect("destroy should succeed");
|
||||
}
|
||||
154
examples/single-gpu-inference/tests/test_worker.py
Normal file
154
examples/single-gpu-inference/tests/test_worker.py
Normal file
|
|
@ -0,0 +1,154 @@
|
|||
"""T-worker: tinygrad worker protocol tests.
|
||||
|
||||
Spawns tinygrad_worker.py --stub as a subprocess and verifies the
|
||||
stdin/stdout JSON protocol defined in SPEC.md §2.3.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
WORKER_SCRIPT = os.path.join(os.path.dirname(__file__), "..", "tinygrad_worker.py")
|
||||
PYTHON = sys.executable
|
||||
|
||||
|
||||
def spawn_worker():
|
||||
"""Spawn the worker in --stub mode and return the Popen handle."""
|
||||
proc = subprocess.Popen(
|
||||
[PYTHON, WORKER_SCRIPT, "--stub"],
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
)
|
||||
return proc
|
||||
|
||||
|
||||
def read_line(proc, timeout=5):
|
||||
"""Read one line from the worker's stdout, with a timeout."""
|
||||
import selectors
|
||||
sel = selectors.DefaultSelector()
|
||||
sel.register(proc.stdout, selectors.EVENT_READ)
|
||||
events = sel.select(timeout=timeout)
|
||||
if not events:
|
||||
raise TimeoutError(f"No output from worker within {timeout}s")
|
||||
line = proc.stdout.readline()
|
||||
sel.close()
|
||||
if not line:
|
||||
raise EOFError("Worker closed stdout")
|
||||
return json.loads(line.strip())
|
||||
|
||||
|
||||
def send_line(proc, obj):
|
||||
"""Write a JSON line to the worker's stdin."""
|
||||
proc.stdin.write(json.dumps(obj) + "\n")
|
||||
proc.stdin.flush()
|
||||
|
||||
|
||||
def send_raw(proc, text):
|
||||
"""Write raw text to the worker's stdin."""
|
||||
proc.stdin.write(text + "\n")
|
||||
proc.stdin.flush()
|
||||
|
||||
|
||||
class TestWorkerStartup:
|
||||
"""Worker prints {"status": "ready"} on startup."""
|
||||
|
||||
def test_worker_emits_ready_on_startup(self):
|
||||
proc = spawn_worker()
|
||||
try:
|
||||
msg = read_line(proc)
|
||||
assert msg == {"status": "ready"}, f"Expected ready status, got: {msg}"
|
||||
finally:
|
||||
proc.terminate()
|
||||
proc.wait(timeout=5)
|
||||
|
||||
|
||||
class TestWorkerInference:
|
||||
"""Valid request produces a valid response with non-empty text."""
|
||||
|
||||
def test_valid_request_returns_non_empty_response(self):
|
||||
proc = spawn_worker()
|
||||
try:
|
||||
ready = read_line(proc)
|
||||
assert ready["status"] == "ready"
|
||||
|
||||
send_line(proc, {"prompt": "Say hello", "max_tokens": 8})
|
||||
response = read_line(proc)
|
||||
|
||||
assert "response" in response, f"Response missing 'response' field: {response}"
|
||||
assert isinstance(response["response"], str)
|
||||
assert len(response["response"]) > 0, "Response text must be non-empty"
|
||||
finally:
|
||||
proc.terminate()
|
||||
proc.wait(timeout=5)
|
||||
|
||||
def test_multiple_requests_in_sequence(self):
|
||||
"""Worker handles multiple requests without restarting."""
|
||||
proc = spawn_worker()
|
||||
try:
|
||||
read_line(proc) # ready
|
||||
|
||||
for prompt in ["Hello", "World", "Test"]:
|
||||
send_line(proc, {"prompt": prompt, "max_tokens": 8})
|
||||
response = read_line(proc)
|
||||
assert "response" in response
|
||||
assert len(response["response"]) > 0
|
||||
finally:
|
||||
proc.terminate()
|
||||
proc.wait(timeout=5)
|
||||
|
||||
|
||||
class TestWorkerMalformedInput:
|
||||
"""Malformed JSON produces {"error": "..."} and worker continues."""
|
||||
|
||||
def test_malformed_json_returns_error_and_continues(self):
|
||||
proc = spawn_worker()
|
||||
try:
|
||||
read_line(proc) # ready
|
||||
|
||||
# Send garbage
|
||||
send_raw(proc, "this is not json {{{")
|
||||
err = read_line(proc)
|
||||
assert "error" in err, f"Expected error response, got: {err}"
|
||||
|
||||
# Worker should still be alive and handle the next valid request
|
||||
send_line(proc, {"prompt": "Still alive?", "max_tokens": 8})
|
||||
response = read_line(proc)
|
||||
assert "response" in response, "Worker should continue after malformed input"
|
||||
assert len(response["response"]) > 0
|
||||
finally:
|
||||
proc.terminate()
|
||||
proc.wait(timeout=5)
|
||||
|
||||
def test_partial_json_returns_error(self):
|
||||
proc = spawn_worker()
|
||||
try:
|
||||
read_line(proc) # ready
|
||||
|
||||
send_raw(proc, '{"prompt": "incomplete')
|
||||
err = read_line(proc)
|
||||
assert "error" in err
|
||||
finally:
|
||||
proc.terminate()
|
||||
proc.wait(timeout=5)
|
||||
|
||||
|
||||
class TestWorkerEOFShutdown:
|
||||
"""Closing stdin (EOF) causes worker to exit cleanly with code 0."""
|
||||
|
||||
def test_eof_causes_clean_exit(self):
|
||||
proc = spawn_worker()
|
||||
read_line(proc) # ready
|
||||
|
||||
# Close stdin
|
||||
proc.stdin.close()
|
||||
|
||||
# Worker should exit cleanly
|
||||
exit_code = proc.wait(timeout=5)
|
||||
assert exit_code == 0, f"Worker exited with code {exit_code}, expected 0"
|
||||
159
examples/single-gpu-inference/tinygrad_worker.py
Normal file
159
examples/single-gpu-inference/tinygrad_worker.py
Normal file
|
|
@ -0,0 +1,159 @@
|
|||
#!/usr/bin/env python3
|
||||
"""tinygrad compute worker — stdin/stdout JSON protocol.
|
||||
|
||||
Startup: loads a model (or uses --stub for testing), prints {"status": "ready"}.
|
||||
|
||||
Protocol (newline-delimited JSON):
|
||||
→ stdin: {"prompt": "Say hello", "max_tokens": 64, "temperature": 0.7}
|
||||
← stdout: {"response": "Hello! How can I help you today?"}
|
||||
|
||||
Errors:
|
||||
← stdout: {"error": "description of what went wrong"}
|
||||
|
||||
Flags:
|
||||
--stub Skip model loading; return a canned response for every request.
|
||||
Used for component tests that exercise the protocol without a GPU.
|
||||
--model NAME Model from tinygrad's built-in catalog (default: llama3.2:1b).
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="tinygrad inference worker")
|
||||
parser.add_argument("--stub", action="store_true",
|
||||
help="Stub mode: skip model loading, return canned responses")
|
||||
parser.add_argument("--model", default="llama3.2:1b",
|
||||
help="Model name from tinygrad catalog (default: llama3.2:1b)")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.stub:
|
||||
model_data = None
|
||||
else:
|
||||
try:
|
||||
model_data = _load_model(args.model)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
_log(traceback.format_exc())
|
||||
_write({"error": f"model load failed: {e}"})
|
||||
sys.exit(1)
|
||||
|
||||
# Signal readiness
|
||||
_write({"status": "ready", "pid": os.getpid()})
|
||||
|
||||
# Request loop
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
request = json.loads(line)
|
||||
except (json.JSONDecodeError, ValueError) as e:
|
||||
_write({"error": f"invalid JSON: {e}"})
|
||||
continue
|
||||
|
||||
if "prompt" not in request:
|
||||
_write({"error": "missing 'prompt' field"})
|
||||
continue
|
||||
|
||||
prompt = request["prompt"]
|
||||
max_tokens = request.get("max_tokens", 64)
|
||||
temperature = request.get("temperature", 0.7)
|
||||
|
||||
try:
|
||||
text = _generate(model_data, prompt, max_tokens, temperature, stub=args.stub)
|
||||
_write({"response": text})
|
||||
except Exception as e:
|
||||
_write({"error": f"generation failed: {e}"})
|
||||
|
||||
|
||||
def _write(obj):
|
||||
"""Write a JSON object as a single line to stdout and flush."""
|
||||
print(json.dumps(obj), flush=True)
|
||||
|
||||
|
||||
def _log(msg):
|
||||
"""Write a log message to stderr (not part of the JSON protocol)."""
|
||||
print(msg, file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
def _load_model(model_name):
|
||||
"""Load a GGUF model via tinygrad 0.12.0's built-in catalog."""
|
||||
from tinygrad import Tensor
|
||||
from tinygrad.helpers import fetch
|
||||
from tinygrad.apps.llm import Transformer, SimpleTokenizer, models
|
||||
|
||||
if model_name not in models:
|
||||
available = ", ".join(models.keys())
|
||||
raise ValueError(f"Unknown model '{model_name}'. Available: {available}")
|
||||
|
||||
url = models[model_name]
|
||||
_log(f"Downloading {model_name} from {url}...")
|
||||
gguf_path = fetch(url)
|
||||
|
||||
_log(f"Loading model from {gguf_path}...")
|
||||
model, kv = Transformer.from_gguf(Tensor(gguf_path), max_context=512)
|
||||
tokenizer = SimpleTokenizer.from_gguf_kv(kv)
|
||||
|
||||
# Find stop token IDs for generation
|
||||
tokens_list = kv.get("tokenizer.ggml.tokens", [])
|
||||
stop_ids = set()
|
||||
for i, tok in enumerate(tokens_list):
|
||||
if tok in ("<|end_of_text|>", "<|eot_id|>", "</s>", "<|endoftext|>"):
|
||||
stop_ids.add(i)
|
||||
|
||||
# Find EOS token ID for chat template end-of-turn
|
||||
eot_id = None
|
||||
for i, tok in enumerate(tokens_list):
|
||||
if tok == "<|eot_id|>":
|
||||
eot_id = i
|
||||
break
|
||||
if eot_id is None:
|
||||
for i, tok in enumerate(tokens_list):
|
||||
if tok in ("</s>", "<|end_of_text|>"):
|
||||
eot_id = i
|
||||
break
|
||||
|
||||
_log(f"Model loaded. Stop IDs: {stop_ids}, EOT ID: {eot_id}")
|
||||
return {"model": model, "tokenizer": tokenizer, "stop_ids": stop_ids, "eot_id": eot_id}
|
||||
|
||||
|
||||
def _format_chat_tokens(tokenizer, prompt, eot_id):
|
||||
"""Format a prompt using Llama 3 instruct chat template."""
|
||||
tokens = tokenizer.role("user")
|
||||
tokens += tokenizer.encode(prompt)
|
||||
if eot_id is not None:
|
||||
tokens += tokenizer.end_turn(eot_id)
|
||||
tokens += tokenizer.role("assistant")
|
||||
return tokens
|
||||
|
||||
|
||||
def _generate(model_data, prompt, max_tokens, temperature, stub=False):
|
||||
"""Generate text from a prompt."""
|
||||
if stub:
|
||||
return f"stub response to: {prompt}"
|
||||
|
||||
model = model_data["model"]
|
||||
tokenizer = model_data["tokenizer"]
|
||||
stop_ids = model_data["stop_ids"]
|
||||
eot_id = model_data["eot_id"]
|
||||
|
||||
# Use chat template for instruction-tuned models
|
||||
tokens = _format_chat_tokens(tokenizer, prompt, eot_id)
|
||||
prompt_len = len(tokens)
|
||||
|
||||
for i, tok_id in enumerate(model.generate(tokens)):
|
||||
if tok_id in stop_ids:
|
||||
tokens.pop() # remove the stop token from output
|
||||
break
|
||||
if i + 1 >= max_tokens:
|
||||
break
|
||||
|
||||
return tokenizer.decode(tokens[prompt_len:])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Reference in a new issue