Compare commits
45 Commits
cursor/fix
...
5feb5b96f8
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5feb5b96f8 | ||
|
|
daddbaa4a3 | ||
|
|
d648da3344 | ||
|
|
4a10eb6013 | ||
|
|
436e872abe | ||
|
|
1e44e8e578 | ||
|
|
52629d7762 | ||
|
|
0ffd195fec | ||
|
|
0b39d80375 | ||
|
|
aa7f4eb13b | ||
|
|
42d6fe2b15 | ||
|
|
1b9f62f82f | ||
|
|
a224644247 | ||
|
|
1ecc599f7f | ||
|
|
91e4bcf2c9 | ||
|
|
e44abc910d | ||
|
|
29db25108f | ||
|
|
e06969fcb5 | ||
|
|
194fa1d926 | ||
|
|
7419ace926 | ||
|
|
560de08edd | ||
|
|
9c73db0ef2 | ||
|
|
3d82188dc1 | ||
|
|
518c259cd3 | ||
|
|
f0dc3bd93f | ||
|
|
a0b37ad1b9 | ||
|
|
dae0719a32 | ||
|
|
e2b20883ca | ||
|
|
481ce6c6f5 | ||
|
|
7ba87051f5 | ||
|
|
ac0ca20b56 | ||
|
|
38355eba25 | ||
|
|
471893c9d5 | ||
|
|
a0dcbfbfd0 | ||
|
|
0d8162dcd3 | ||
|
|
3fc8228590 | ||
|
|
6374082b1b | ||
|
|
16614855bc | ||
|
|
cdd2699e63 | ||
|
|
912ee4f1fd | ||
|
|
f1eea5b6d4 | ||
|
|
456c43ea1d | ||
|
|
aba5fb12fa | ||
|
|
1eb1e0baa2 | ||
|
|
b1f08c45cd |
BIN
.billing.sqlite
BIN
.billing.sqlite
Binary file not shown.
@@ -30,3 +30,7 @@ Both are already migrated into `.scratch/alpha-hardening/prd.json` (AH-021 updat
|
||||
|
||||
**Why:** three audits agreed the alpha blockers are unauthenticated gossip (anyone can inject billing events), the free-credit faucet, and ephemeral bans.
|
||||
**How to apply:** work test-first per issue acceptance criteria; use `.venv`; `cryptography` belongs in node deps (wallet.py imports it — causes many of the 24 "failures" in a fresh env). See [[project-status]] and [[autonomous-work-style]].
|
||||
|
||||
## Routing telemetry resume (2026-07-07)
|
||||
|
||||
`.scratch/alpha-hardening/issues/24-routing-telemetry-resume.md` / AH-024 captures the interrupted Claude handoff. Learned routing is already committed at `518c259`; the dirty tree contains live-progress/current-request heartbeat/dashboard telemetry. First known blocker: `packages/tracker/meshnet_tracker/server.py:1490` uses `threading.Lock | None`, which crashes import because `threading.Lock` is a factory function at runtime. Fix that before running the targeted telemetry tests. Keep `.claude/settings.local.json` uncommitted unless explicitly approved.
|
||||
|
||||
@@ -42,7 +42,8 @@ Historical handoff note: `/mnt/c/Users/popov/Downloads/neuron-tai-alpha-handoff-
|
||||
- Verification: downloader/startup targeted subset passes (`pytest tests/test_node_startup.py -k "download_shard or same_shard"`). Full `tests/test_node_startup.py` has 46 passed and 4 unrelated Windows chmod/path separator failures.
|
||||
- Live Windows confirmation: `meshnet-node start --tracker http://192.168.0.179:8080 --model Qwen3.6-35B-A3B` reuses `F:\_STORAGE\models\qwen3.6-35b-a3b`, prints `Cached at`, registers, and reaches ready as node `5gMLrmyB-26b1f8a4204a`.
|
||||
- Follow-up fix: preset-model startup now starts the heartbeat thread after registration; without this, the node appeared briefly on the dashboard and was purged on first inference/route after heartbeat expiry. Tracker dashboard now has a "Console output" panel backed by `/v1/console` for node register/expiry, routing failures, and proxy events.
|
||||
- Qwen3.6-35B-A3B reserve-based split is expected: an 79 GB CPU node may be assigned layers 0-36, and a second node fills 37-39. Do not "fix" this by bypassing the 20% assignment reserve unless the shard-planning policy changes.
|
||||
- Qwen3.6-35B-A3B CPU runtime cap (2026-07-08): the old reserve-based split could assign an 79 GB CPU node layers 0-36, but real partial loading can exceed that budget and die without a Python traceback. Node startup now clips oversized CPU auto-assignments before loading, and tracker CPU assignment uses a stricter runtime headroom factor; do not revert this to the old 20% reserve-only policy.
|
||||
- Route hardening: tracker chat proxy and `/v1/route` diagnostics now use alias-aware preset node matching for split Qwen3.6 routes; dashboard derives grouped inference history from proxy route/complete console events and shows observed TPS after completion.
|
||||
- Live proxy hardening: model lookup trims outer whitespace before alias matching (`qwen3.6-35b-a3b ` resolves), and tracker route logs/dashboard queue depth combine heartbeat queue with tracker-local proxy in-flight counts so Postman-style bursts no longer show every selected route as queue `0`.
|
||||
- Split-shard streaming hardening: Qwen3.6-style distributed generation now emits SSE chunks token-by-token from the head node instead of buffering all generated text until completion. Tracker direct/relay stream proxy logs `proxy progress` with live tokens/TPS, dashboard Inference history shows currently processing requests with live TPS/tokens/queue, and relay stream completion no longer references an undefined `session_id`.
|
||||
- Native Windows Qwen3.6-MoE import fix: `flash-linear-attention` imports `triton`; without `triton-windows`, startup fails with misleading `Could not import module 'Qwen3_5MoeForCausalLM'`. Installed `triton-windows` in `C:\Users\popov\miniforge3` and added it as a Windows-only node dependency.
|
||||
|
||||
28
.gitattributes
vendored
Normal file
28
.gitattributes
vendored
Normal file
@@ -0,0 +1,28 @@
|
||||
# Normalize line endings across Windows/Linux checkouts.
|
||||
# All text files are stored as LF in the repo and checked out as LF
|
||||
# on every OS. Git auto-detects text vs binary.
|
||||
* text=auto eol=lf
|
||||
|
||||
# Explicitly binary — never touch these bytes.
|
||||
*.png binary
|
||||
*.jpg binary
|
||||
*.jpeg binary
|
||||
*.gif binary
|
||||
*.ico binary
|
||||
*.pdf binary
|
||||
*.zip binary
|
||||
*.gz binary
|
||||
*.tar binary
|
||||
*.wasm binary
|
||||
*.sqlite binary
|
||||
*.sqlite3 binary
|
||||
*.safetensors binary
|
||||
*.gguf binary
|
||||
|
||||
# Scripts that must stay LF even if someone forces CRLF locally.
|
||||
*.sh text eol=lf
|
||||
*.py text eol=lf
|
||||
|
||||
# Windows batch files genuinely need CRLF.
|
||||
*.bat text eol=crlf
|
||||
*.cmd text eol=crlf
|
||||
10
.gitignore
vendored
10
.gitignore
vendored
@@ -18,5 +18,11 @@ dist/
|
||||
!.env.example
|
||||
!.env.testnet
|
||||
.rocm-local/*
|
||||
billing.sqlite
|
||||
.pytest-tmp/*
|
||||
.pytest-tmp/*
|
||||
|
||||
# Local tracker/node sqlite databases (never commit runtime state)
|
||||
*.sqlite
|
||||
*.sqlite3
|
||||
logs/tracker/error.log
|
||||
logs/tracker/info.log
|
||||
logs/tracker/warning.log
|
||||
|
||||
@@ -9,6 +9,10 @@ Pre-release alpha audit + grilling (2026-07-04). Bucket 1 trust-boundary blocker
|
||||
|
||||
Locked scope: one settlement tracker, open node join, devnet mock-USDT, reputation carries forward → fraud must be bounded. See [ADR-0016](../../docs/adr/0016-alpha-scope-and-known-limitations.md).
|
||||
|
||||
**Resume task (2026-07-07):** [24 - Routing telemetry resume](./issues/24-routing-telemetry-resume.md) is `ready-for-agent`. Learned-routing commit `518c259` is already present; dirty tree contains current-request heartbeat/dashboard telemetry and a known import-time annotation crash in `server.py:1490`.
|
||||
|
||||
**Perf follow-up (2026-07-08):** [25 — Sharded per-node KV cache for distributed generation](./issues/25-per-node-kv-cache-distributed.md) is **implemented** ([ADR-0022](../../docs/adr/0022-sharded-per-node-kv-cache.md)): per-generation session ids, prefill/decode wire protocol (`X-Meshnet-Cache`/`X-Meshnet-Past-Len`), per-node sharded `DynamicCache(config=…)` (hybrid-attention-aware), TTL+LRU eviction with 409 cache-miss → full re-prefill fallback. Golden test proves token-identical output vs the stateless path; CPU two-shard measurement: 7.05 tps decaying 32% → 18.93 tps flat (2.68×). Remaining: re-measure on the live 2-node GPU topology and the Qwen3.6-35B-A3B mixed topology.
|
||||
|
||||
## Artifacts
|
||||
|
||||
| Path | Status |
|
||||
@@ -16,7 +20,7 @@ Locked scope: one settlement tracker, open node join, devnet mock-USDT, reputati
|
||||
| [research-verifiable-inference.md](./research-verifiable-inference.md) | Complete — SOTA research, §8 layered scheme, TOPLOC adopt |
|
||||
| [handoff.md](./handoff.md) | Session handoff — locked decisions, env notes |
|
||||
| [docs/adr/0016–0019](../../docs/adr/) | Alpha scope, auth, fraud, multi-tracker design |
|
||||
| [issues/](./issues/) | 22 work items (Buckets 1–3) |
|
||||
| [issues/](./issues/) | 25 work items (Buckets 1–3 + perf follow-ups) |
|
||||
|
||||
## ADRs (this feature)
|
||||
|
||||
@@ -76,6 +80,12 @@ Locked scope: one settlement tracker, open node join, devnet mock-USDT, reputati
|
||||
| [22 MEMORY + project-status index](./issues/22-doc-memory-project-status.md) (done) |
|
||||
| [21 Honest-noise calibration corpus](./issues/21-honest-noise-calibration-corpus.md) (ops; prod gate for audits) |
|
||||
|
||||
### Phase 5 — Distributed-inference performance (post-routing-fix)
|
||||
|
||||
| Issue | Depends on |
|
||||
|---|---|
|
||||
| [25 Sharded per-node KV cache](./issues/25-per-node-kv-cache-distributed.md) | ADR-0020 routing fix (done), [24 routing telemetry resume](./issues/24-routing-telemetry-resume.md) |
|
||||
|
||||
## First 3 to implement
|
||||
|
||||
1. **02 + 20** — Unified auth boundary + validator service token (shared helper and roles)
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
Status: ready-for-agent
|
||||
|
||||
Scoped 2026-07-07 from an interrupted Claude session. This is a resume/cleanup task for routing and live-progress work that is partly committed and partly left dirty in the working tree.
|
||||
|
||||
# 24 - Finish learned-routing telemetry and live-progress cleanup
|
||||
|
||||
## Current state
|
||||
|
||||
The main dynamic routing feature is already committed at `518c259` (`routing improvements - dynamic (wip)`):
|
||||
|
||||
- `packages/tracker/meshnet_tracker/routing_stats.py` - decayed-EWMA route stats store, epsilon-greedy route selection, diagnostics.
|
||||
- `packages/tracker/meshnet_tracker/server.py` - route enumeration per head, bandit selection in the chat proxy, epoch bumps on node join/leave, `/v1/routing`, route sample recording with 8-token hygiene.
|
||||
- `packages/tracker/meshnet_tracker/cli.py` - `--route-explore-share`, `--route-weight-alpha`, `--route-stats-half-life` and env vars.
|
||||
- `packages/tracker/meshnet_tracker/dashboard.html` - "Routing (learned)" panel.
|
||||
- `docs/adr/0021-dynamic-statistical-routing.md` - design record.
|
||||
- `tests/test_dynamic_routing.py` - includes the exact GPU(0-21)+CPU(0-39) topology, hybrid downstream `start_layer=22`, 0.6/0.4 traffic split for a 1.5 TPS ratio, and scout-rate behavior.
|
||||
|
||||
The current working tree still has uncommitted follow-up work:
|
||||
|
||||
- `packages/node/meshnet_node/torch_server.py` - tracks in-flight chat requests, exposes `TorchNodeServer.current_requests`, prints generation progress with TPS.
|
||||
- `packages/node/meshnet_node/startup.py` - sends `current_requests` in heartbeat payloads and increases heartbeat cadence while busy.
|
||||
- `packages/tracker/meshnet_tracker/server.py` - accepts heartbeat `current_requests`, includes them in `/v1/network/map`, and logs `proxy connecting` before upstream connection.
|
||||
- `packages/tracker/meshnet_tracker/dashboard.html` - enriches the call wall from heartbeat `current_requests` so active requests remain visible even before terminal proxy events.
|
||||
- `tests/test_real_model_backend.py` and `tests/test_tracker_routing.py` - targeted coverage for current-request snapshots, heartbeat sanitization/storage, and TPS progress logging.
|
||||
- `QUICKSTART.md` - documents optional linear-attention fast-path packages for Qwen3.5/3.6 GPU nodes.
|
||||
|
||||
There is also an untracked local file, `.claude/settings.local.json`, which should not be included unless the owner explicitly wants local Claude settings committed.
|
||||
|
||||
## Known blocker found during resume
|
||||
|
||||
Targeted pytest currently fails during import before reaching the new tests:
|
||||
|
||||
```text
|
||||
TypeError: unsupported operand type(s) for |: 'builtin_function_or_method' and 'NoneType'
|
||||
```
|
||||
|
||||
Immediate cause: `packages/tracker/meshnet_tracker/server.py:1490` annotates `ws_lock: threading.Lock | None = None`. `threading.Lock` is a factory function at runtime, not a type, so `| None` evaluates eagerly and crashes. This exists on `HEAD` too, not just in the dirty telemetry changes.
|
||||
|
||||
Fix options:
|
||||
|
||||
- Add `from __future__ import annotations` at the top of `server.py`, then run enough tests to catch any annotation side effects.
|
||||
- Or change that annotation to a safe runtime type such as `Any | None` / remove the union annotation. Keep the change minimal.
|
||||
|
||||
## What to do next
|
||||
|
||||
1. Fix the import-time `threading.Lock | None` crash.
|
||||
2. Re-run the targeted tests:
|
||||
|
||||
```bash
|
||||
.\.venv\Scripts\python.exe -m pytest tests/test_tracker_routing.py::test_tracker_heartbeat_stores_current_requests tests/test_tracker_routing.py::test_normalize_current_requests_sanitizes_payload tests/test_real_model_backend.py::test_current_requests_snapshot_while_generating tests/test_real_model_backend.py::test_distributed_generating_log_includes_tps -q
|
||||
```
|
||||
|
||||
3. Run the relevant routing regression tests:
|
||||
|
||||
```bash
|
||||
.\.venv\Scripts\python.exe -m pytest tests/test_dynamic_routing.py tests/test_tracker_routing.py -q
|
||||
```
|
||||
|
||||
4. If practical, run the non-integration suite:
|
||||
|
||||
```bash
|
||||
.\.venv\Scripts\python.exe -m pytest tests/ -q -m "not integration"
|
||||
```
|
||||
|
||||
5. Confirm or document the pre-existing failure from the interrupted session: `test_proxy_chat_splits_payout_by_tracker_assigned_route_span` reportedly failed on `HEAD` too and was unrelated.
|
||||
6. Commit the intentional work in two commits if it remains naturally split:
|
||||
- learned routing is already committed in `518c259`; leave it alone unless fixing regressions there.
|
||||
- commit the live-progress/current-request telemetry cleanup separately after tests pass.
|
||||
|
||||
## Acceptance criteria
|
||||
|
||||
- [ ] Importing `meshnet_tracker.server` no longer crashes on the lock annotation.
|
||||
- [ ] Current-request heartbeat payloads are sanitized and surfaced in `/v1/network/map`.
|
||||
- [ ] Node-side in-flight chat snapshots report request id, model, token count, elapsed seconds, tokens/sec, and routing completion.
|
||||
- [ ] Dashboard call wall can show active requests from heartbeat data, not only tracker console terminal events.
|
||||
- [ ] Targeted telemetry tests pass.
|
||||
- [ ] Dynamic routing tests still pass, including GPU(0-21)+CPU(0-39) hybrid-route enumeration and traffic split behavior.
|
||||
- [ ] Full or non-integration suite result is recorded; unrelated pre-existing failures are named explicitly.
|
||||
- [ ] `.claude/settings.local.json` remains uncommitted unless intentionally approved.
|
||||
|
||||
## ADR links
|
||||
|
||||
- [ADR-0020](../../docs/adr/0020-chat-streaming-live-progress-and-mixed-topology-routing.md)
|
||||
- [ADR-0021](../../docs/adr/0021-dynamic-statistical-routing.md)
|
||||
|
||||
## Blocked by
|
||||
|
||||
None. The import-time annotation crash is the first fix.
|
||||
|
||||
## Blocks
|
||||
|
||||
Clean handoff/commit of the interrupted live routing progress work.
|
||||
@@ -0,0 +1,58 @@
|
||||
Status: implemented 2026-07-08 — pending live 2-node GPU verification
|
||||
|
||||
Implemented in `packages/node/meshnet_node/model_backend.py` + `torch_server.py`; design in
|
||||
[ADR-0022](../../../docs/adr/0022-sharded-per-node-kv-cache.md); tests in
|
||||
`tests/test_kv_cache_distributed.py` (11 fast tests + env-gated golden test,
|
||||
`MESHNET_REAL_MODEL_TESTS=1`).
|
||||
|
||||
**Measured (two-shard Qwen2.5-0.5B 0-11/12-23, CPU, 44-token prompt, 40 steps):**
|
||||
stateless 7.05 tps decaying 32% (8.09 → 5.50 first-10 vs last-10); cached 18.93 tps and
|
||||
FLAT (17.21 → 19.28) — 2.68× overall, gap grows quadratically with length. Remaining
|
||||
acceptance item: re-measure on the live 2-node GPU topology (needs both machines).
|
||||
|
||||
Scoped 2026-07-08 from a live two-machine distributed-inference debugging session (Qwen2.5-0.5B GPU+GPU pipeline, and Qwen3.6-35B-A3B mixed GPU/CPU). The ADR-0020 mixed-topology `start_layer` bug is fixed (`518c259`, `e44abc9`, `1ecc599`); this issue is the next performance blocker in the same code path.
|
||||
|
||||
# 25 — Sharded per-node KV cache for distributed generation (MoE/hybrid-attention aware)
|
||||
|
||||
## What to build
|
||||
|
||||
The distributed generation loop (`torch_server.py:515-612`, `_do_chat_completions` distributed path) currently has **no KV cache at all**: `model_backend.py` passes `use_cache: False` in every layer-forward call (lines 763, 768, 770-771), and each autoregressive step re-encodes the *entire* prompt-so-far from scratch (`backend.encode_prompt(current_text)`), re-running every layer on every node in the route for every generated token.
|
||||
|
||||
Observed cost of this on a live 2-node Qwen2.5-0.5B GPU pipeline (layers 0-20 / 21-23): tps decayed from 22.3 (at 235 output tokens) to 12.6 (at 449 tokens) within a single generation — the expected quadratic-cost signature. On the Qwen3.6-35B-A3B mixed-topology case this collapses to ~0.07 tps even after the routing fix, partly for this reason.
|
||||
|
||||
`X-Meshnet-Session` already exists on the wire (`torch_server.py:707`, minted fresh **per token**, not per generation) but today only labels one activation transfer for chunk reassembly/logging — it is not used to key any cached state.
|
||||
|
||||
| Subtask | Owner package | Deliverable |
|
||||
|---|---|---|
|
||||
| Session lifecycle | `packages/node/meshnet_node/torch_server.py` | Mint session ID once per chat request (not per token); reuse across all steps of that generation; add `X-Meshnet-Seq-Len` / position header so a node can tell prefill from decode steps |
|
||||
| Per-node sharded cache | `packages/node/meshnet_node/model_backend.py` | `TorchModelShard` holds a `session_id → cache_state` map scoped to *its own* layer range only (naturally sharded — no node stores another node's KV); `forward_bytes` takes `use_cache=True` and returns/reuses `past_key_values` (or `use_cache=False` for the prefill token to keep failure/eviction simple) |
|
||||
| Prefill vs. decode split | `packages/node/meshnet_node/torch_server.py` | Step 0 sends the full prompt activation (current behavior); steps 1+ send only the newest token's hidden state (`[1, 1, hidden]`) with correct `position_ids`, cutting per-step payload from O(seq_len) to O(1) |
|
||||
| MoE / hybrid-attention state | `packages/node/meshnet_node/model_backend.py` | Cache abstraction must hold "whatever `use_cache=True` returns for this layer range," not assume standard K/V tensors — Qwen3.6's linear-attention/hybrid layers (see `[transformers] The fast path is not available...` warning already logged at startup) cache **recurrent conv/delta state**, not K/V pairs. MoE expert routing itself is layer-local and needs no cross-token cache, but confirm no expert-choice state leaks across the stateless-vs-cached boundary when `use_cache` toggles between prefill and decode |
|
||||
| Cache lifecycle | `packages/node/meshnet_node/torch_server.py` | TTL + LRU eviction per node (bounded by `max_loaded_shards`/memory budget); explicit "cache miss" response so a restarted/evicted node causes the head to fall back to a full re-prefill instead of a hard error — keep today's fully-stateless path as the recovery mode |
|
||||
| Correctness parity | `tests/` | Golden-output test: distributed multi-token output with caching enabled must match the existing stateless path token-for-token (or within sampling tolerance) for a fixed prompt/seed |
|
||||
|
||||
**Non-goals for first landing:** cross-node cache migration/rebalancing on route change (evict + re-prefill is acceptable initially); speculative decoding; batching multiple concurrent sessions' KV within one node beyond what eviction already requires.
|
||||
|
||||
**Code refs:**
|
||||
|
||||
- `packages/node/meshnet_node/torch_server.py:515-612` — distributed generation loop (`current_text = current_text + token_str`, full re-encode every step)
|
||||
- `packages/node/meshnet_node/torch_server.py:690-789` — `_run_downstream_pipeline`, session minting, `X-Meshnet-Session`/`X-Meshnet-Hop-Index`/`X-Meshnet-Start-Layer` headers
|
||||
- `packages/node/meshnet_node/model_backend.py:189-201, 330-351, 763-771` — `use_cache: False` call sites, `effective_start` layer-slicing logic that any cache keying must respect
|
||||
- `docs/adr/0020-chat-streaming-live-progress-and-mixed-topology-routing.md` — prerequisite routing fix this issue builds on
|
||||
- `docs/adr/0021-dynamic-statistical-routing.md` — route selection this cache must stay compatible with (a route change mid-generation should trigger cache-miss fallback, not corruption)
|
||||
|
||||
## Acceptance criteria
|
||||
|
||||
- [x] A session ID is stable across all steps of one chat generation (not re-minted per token) — minted once in `_do_chat_completions`, asserted in `test_session_is_stable_and_decode_payloads_are_single_token`
|
||||
- [x] Steps after the first prefill send only the new token's activation (`[1, 1, hidden]` via `encode_next_token`) with `X-Meshnet-Cache: decode` + `X-Meshnet-Past-Len`
|
||||
- [x] Each node caches state only for its own shard's layer range (`TorchModelShard.kv_sessions`; sharding falls out of per-node layer execution)
|
||||
- [x] Cache abstraction is not K/V-shaped-only: `DynamicCache(config=model.config)` — the same construction Qwen3.6-Next's own forward uses for hybrid linear-attention conv/delta state; store treats it as opaque; `TypeError` fallback disables caching per-backend
|
||||
- [x] Bounded memory: TTL (600 s, `MESHNET_KV_TTL_SECONDS`) + LRU (8, `MESHNET_KV_MAX_SESSIONS`); miss → HTTP 409 `{"error": "cache_miss"}` → head re-prefills (tested)
|
||||
- [x] Golden-output test: cached and stateless produce identical token ids on real two-shard Qwen2.5-0.5B (`test_cached_distributed_generation_matches_stateless_golden`, passed)
|
||||
- [x] Measured (CPU two-shard proxy, 40 steps): stateless 7.05 tps w/ 32% decay → cached 18.93 tps flat, 2.68×. ⚠️ still to run on the live 2-node GPU topology
|
||||
- [x] `tests/test_two_node_pipeline.py` and `tests/test_dynamic_routing.py` pass (30 passed; 6 tmp-dir fixture errors are a pre-existing Windows temp-permission env issue, identical on clean tree)
|
||||
- [x] Design captured in [ADR-0022](../../../docs/adr/0022-sharded-per-node-kv-cache.md) incl. cache-miss/route-change interaction with ADR-0021
|
||||
|
||||
## Notes
|
||||
|
||||
MoE routing (router + expert FFN) is layer-local per token and does not itself need a cross-token cache — it was ruled out as the cause of the earlier Qwen3.6 garbage-output bug (that was the ADR-0020 `start_layer` double-execution). The MoE angle that *does* matter here is architecture-awareness in the cache design: don't hardcode a K/V tensor shape assumption that breaks on Qwen3.6's hybrid attention layers.
|
||||
@@ -483,9 +483,50 @@
|
||||
"notes": "Source issue: .scratch/alpha-hardening/issues/23-dynamic-hf-pricing.md. High priority, ship-soon for launch, NOT an alpha-release blocker (unlike AH-021).",
|
||||
"dependsOn": [],
|
||||
"completionNotes": "Completed by agent"
|
||||
},
|
||||
{
|
||||
"id": "AH-024",
|
||||
"title": "24 - Finish learned-routing telemetry and live-progress cleanup",
|
||||
"description": "Status: ready-for-agent\n\nScoped 2026-07-07 from an interrupted Claude session. The learned-routing feature is already committed at 518c259 (`routing improvements - dynamic (wip)`): routing_stats.py, tracker route enumeration and bandit selection, CLI routing flags, `/v1/routing`, dashboard Routing (learned), ADR-0021, and tests/test_dynamic_routing.py including the GPU(0-21)+CPU(0-39) hybrid topology. The dirty working tree contains follow-up live-progress/current-request telemetry in torch_server.py, startup.py, tracker server/dashboard, tests, and QUICKSTART. Known blocker found during resume: importing meshnet_tracker.server currently crashes at `server.py:1490` because `ws_lock: threading.Lock | None = None` evaluates `threading.Lock` as a factory function, not a type. Fix that first, then verify and commit the telemetry cleanup separately from the already-committed dynamic-routing work. Leave `.claude/settings.local.json` uncommitted unless explicitly approved.\n\nSource issue has exact file list, commands, and the reported pre-existing unrelated failure (`test_proxy_chat_splits_payout_by_tracker_assigned_route_span`).",
|
||||
"acceptanceCriteria": [
|
||||
"Importing `meshnet_tracker.server` no longer crashes on the lock annotation",
|
||||
"Current-request heartbeat payloads are sanitized and surfaced in `/v1/network/map`",
|
||||
"Node-side in-flight chat snapshots report request id, model, token count, elapsed seconds, tokens/sec, and routing completion",
|
||||
"Dashboard call wall can show active requests from heartbeat data, not only tracker console terminal events",
|
||||
"Targeted telemetry tests pass",
|
||||
"Dynamic routing tests still pass, including GPU(0-21)+CPU(0-39) hybrid-route enumeration and traffic split behavior",
|
||||
"Full or non-integration suite result is recorded; unrelated pre-existing failures are named explicitly",
|
||||
"`.claude/settings.local.json` remains uncommitted unless intentionally approved"
|
||||
],
|
||||
"priority": 24,
|
||||
"passes": true,
|
||||
"notes": "Source issue: .scratch/alpha-hardening/issues/24-routing-telemetry-resume.md. Resume task for interrupted 2026-07-07 Claude session; first known fix is server.py:1490 annotation crash.",
|
||||
"dependsOn": [],
|
||||
"completionNotes": ""
|
||||
},
|
||||
{
|
||||
"id": "AH-025",
|
||||
"title": "25 — Sharded per-node KV cache for distributed generation (MoE/hybrid-attention aware)",
|
||||
"description": "Status: implemented 2026-07-08 — pending live 2-node GPU verification\n\nScoped 2026-07-08 from a live two-machine distributed-inference debugging session. The ADR-0020 mixed-topology start_layer bug is fixed (518c259, e44abc9, 1ecc599); this is the next performance blocker in the same path. The distributed generation loop has NO KV cache at all: model_backend.py passes use_cache: False in every layer-forward call, and each autoregressive step re-encodes the entire prompt-so-far from scratch, re-running every layer on every node in the route for every generated token. Observed on a live 2-node Qwen2.5-0.5B GPU pipeline: tps decayed from 22.3 (at 235 output tokens) to 12.6 (at 449 tokens) within a single generation, the expected quadratic-cost signature. X-Meshnet-Session already exists on the wire but is minted fresh per token and only labels one activation transfer for chunk reassembly/logging, not keyed to any cached state. Build: (1) stable per-request session lifecycle instead of per-token, (2) per-node sharded cache keyed by session scoped to that node's own layer range only, (3) prefill-vs-decode split so post-prefill steps send only the newest token's activation, (4) cache abstraction that holds whatever use_cache=True returns per layer range (not K/V-shaped-only) because Qwen3.6's hybrid linear-attention layers cache recurrent conv/delta state, not standard K/V, (5) TTL+LRU eviction with an explicit cache-miss fallback to full re-prefill so restarts/route-changes degrade gracefully instead of corrupting output. MoE expert routing itself is layer-local and was already ruled out as the cause of the earlier Qwen3.6 garbage-output bug (that was the start_layer double-execution); the MoE angle that matters here is architecture-awareness so the cache design does not hardcode a K/V shape assumption that breaks on Qwen3.6's hybrid attention layers.\n\nSource issue has full subtask table and code refs.",
|
||||
"acceptanceCriteria": [
|
||||
"A session ID is stable across all steps of one chat generation (not re-minted per token)",
|
||||
"Steps after the first prefill send only the new token's activation, not the full sequence, over the wire between nodes",
|
||||
"Each node caches past_key_values/recurrent state only for its own shard's layer range; no node ever holds another node's cache",
|
||||
"Cache works correctly for both standard-attention shards and Qwen3.6-style hybrid linear-attention/recurrent shards",
|
||||
"Bounded memory: TTL + LRU eviction; eviction/restart triggers a documented cache-miss response, not silent corruption",
|
||||
"Golden-output regression test proves cached and uncached distributed generation produce equivalent output for a fixed prompt",
|
||||
"Measured tps improvement recorded on the same 2-node Qwen2.5-0.5B topology used to observe the regression (target: flat tps across generation length)",
|
||||
"tests/test_two_node_pipeline.py and tests/test_dynamic_routing.py still pass",
|
||||
"Design captured in a new ADR (or an amendment to ADR-0020/0021) covering the cache-miss/route-change interaction"
|
||||
],
|
||||
"priority": 25,
|
||||
"passes": true,
|
||||
"notes": "Source issue: .scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md. Perf follow-up to the ADR-0020 routing fix; no prior story covered KV caching or MoE-specific caching needs.",
|
||||
"dependsOn": [],
|
||||
"completionNotes": "Implemented 2026-07-08 (ADR-0022, docs/adr/0022-sharded-per-node-kv-cache.md). Per-generation session id; X-Meshnet-Cache prefill/decode + X-Meshnet-Past-Len wire headers; decode steps send [1,1,hidden] via encode_next_token (tail now returns token_id so the head never re-tokenizes); per-node SessionCacheStore holds DynamicCache(config=model.config) — hybrid-attention/recurrent-state aware, sharded naturally by each node's own layer range; TTL (600s) + LRU (8) eviction; 409 {\"error\":\"cache_miss\"} -> head re-prefills full sequence under the same session (stateless path kept as recovery mode; legacy nodes without the protocol degrade to per-step prefill). Tests: tests/test_kv_cache_distributed.py — 11 fast tests + env-gated golden test (MESHNET_REAL_MODEL_TESTS=1) proving token-identical cached vs stateless output on a real two-shard Qwen2.5-0.5B split. Measured (CPU two-shard, 40 steps): stateless 7.05 tps decaying 32% -> cached 18.93 tps flat, 2.68x overall. Remaining: re-measure on the live 2-node GPU topology and Qwen3.6-35B-A3B mixed topology (needs both machines)."
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"updatedAt": "2026-07-06T06:01:25.474Z"
|
||||
"updatedAt": "2026-07-08T23:30:00.000Z"
|
||||
}
|
||||
}
|
||||
258
CONTEXT.md
258
CONTEXT.md
@@ -1,23 +1,23 @@
|
||||
# Distributed Inference Network
|
||||
|
||||
A volunteer GPU network where nodes independently load model shards, a tracker routes inference through optimal node chains, and contributors earn tokens for serving compute.
|
||||
|
||||
## Language
|
||||
|
||||
### Nodes & compute
|
||||
|
||||
**Node**:
|
||||
A volunteer machine that runs the node client, holds one or more shards on disk, and serves inference requests for those shards.
|
||||
_Avoid_: worker, peer, miner, server
|
||||
|
||||
**Shard**:
|
||||
A contiguous range of transformer layers from a model that a node loads and serves. Shards are the unit of storage, assignment, and reward.
|
||||
_Avoid_: partition, slice, chunk, segment
|
||||
|
||||
**Shard Swarm**:
|
||||
The P2P group of nodes that collectively seed and download a specific shard. One swarm exists per shard.
|
||||
_Avoid_: torrent, cluster, pool
|
||||
|
||||
# Distributed Inference Network
|
||||
|
||||
A volunteer GPU network where nodes independently load model shards, a tracker routes inference through optimal node chains, and contributors earn tokens for serving compute.
|
||||
|
||||
## Language
|
||||
|
||||
### Nodes & compute
|
||||
|
||||
**Node**:
|
||||
A volunteer machine that runs the node client, holds one or more shards on disk, and serves inference requests for those shards.
|
||||
_Avoid_: worker, peer, miner, server
|
||||
|
||||
**Shard**:
|
||||
A contiguous range of transformer layers from a model that a node loads and serves. Shards are the unit of storage, assignment, and reward.
|
||||
_Avoid_: partition, slice, chunk, segment
|
||||
|
||||
**Shard Swarm**:
|
||||
The P2P group of nodes that collectively seed and download a specific shard. One swarm exists per shard.
|
||||
_Avoid_: torrent, cluster, pool
|
||||
|
||||
**Inference Route**:
|
||||
An ordered sequence of nodes whose shards together cover all layers of a model. The tracker selects the optimal route per request.
|
||||
_Avoid_: pipeline, chain, path
|
||||
@@ -55,115 +55,115 @@ Realtime progress information for an active Route Session, including phase, gene
|
||||
_Avoid_: logs, debug output
|
||||
|
||||
### Tracker
|
||||
|
||||
**Tracker**:
|
||||
The coordinator service that maintains the node registry, scores nodes by throughput/latency, and assigns inference routes. Runs as a centralized service with a P2P gossip fallback.
|
||||
_Avoid_: coordinator, scheduler, director
|
||||
|
||||
|
||||
**Tracker**:
|
||||
The coordinator service that maintains the node registry, scores nodes by throughput/latency, and assigns inference routes. Runs as a centralized service with a P2P gossip fallback.
|
||||
_Avoid_: coordinator, scheduler, director
|
||||
|
||||
**Tracker Node**:
|
||||
A node that serves at least the first-layer shard (`layers[0..k]`) for a model and acts as the inference entry point for that model. Tracker nodes own the tokenizer and `embed_tokens`, receive client requests directly, select the onward route from the coverage map, and stream results and progress when possible. Any node advertising a new model to the network becomes its tracker node.
|
||||
_Avoid_: primary node, master node, gateway node
|
||||
|
||||
**Coverage Map**:
|
||||
The tracker's per-model mapping of layer ranges to node counts: `[(start_layer, end_layer, node_count), ...]`. A layer range with `node_count=0` is a coverage gap — the model is unroutable until the gap is filled. Coverage-first bin-packing fills all gaps before adding redundancy.
|
||||
_Avoid_: shard map, assignment table, coverage report
|
||||
|
||||
**Rebalance Directive**:
|
||||
A `LOAD_SHARD` or `DROP_SHARD` instruction the tracker issues to a node when the coverage map changes (node joins, node leaves, or load-balance reoptimization). Delivered as part of the node's heartbeat response.
|
||||
_Avoid_: rebalance command, shard instruction, migration order
|
||||
|
||||
**Node Score**:
|
||||
A throughput/latency rating the tracker maintains per node, used for route selection. Updated continuously from inference telemetry.
|
||||
_Avoid_: reputation, rating, rank
|
||||
|
||||
### Payments & fraud
|
||||
|
||||
**Stake**:
|
||||
Collateral a node stands to lose for fraud. In the current design the node's Pending Balance serves as stake — no upfront deposit is required. An optional USDT/TAI deposit may return later for routing priority.
|
||||
_Avoid_: deposit, bond, escrow
|
||||
|
||||
**Treasury**:
|
||||
The single project-owned Solana wallet that custodially holds client deposits, pays node payouts, and accumulates the Protocol Cut. Its keypair is loaded only on settlement-capable trackers.
|
||||
_Avoid_: escrow, vault, hot wallet
|
||||
|
||||
**Pending Balance**:
|
||||
A node's accrued, not-yet-paid USDT earnings on the tracker ledger. Doubles as the node's fraud collateral: it is forfeited in full when a validator catches a divergent output.
|
||||
_Avoid_: unpaid rewards, accrual, balance due
|
||||
|
||||
**Settlement Period**:
|
||||
The dynamic interval driving on-chain payouts: a node is paid when its Pending Balance exceeds the Payout Threshold or the period elapses, whichever comes first. Short in development (seconds), long in production (daily), configurable to grow with volume.
|
||||
_Avoid_: epoch, payout cycle, billing cycle
|
||||
|
||||
**Payout Threshold**:
|
||||
The minimum Pending Balance that triggers an immediate payout before the Settlement Period elapses. Includes a dust floor so payouts are never smaller than they are worth.
|
||||
_Avoid_: minimum payout, dust limit
|
||||
|
||||
**Protocol Cut**:
|
||||
The 10% of inference fees retained by the project for infrastructure; the remaining 90% goes to the nodes that served the request. Accumulates in the Treasury as the future TAI liquidity reserve.
|
||||
_Avoid_: spread, commission, house fee
|
||||
|
||||
**Deposit Watcher**:
|
||||
The tracker component that observes the Treasury's on-chain USDT deposits and credits the sending client's API-key ledger balance.
|
||||
_Avoid_: payment listener, chain scanner
|
||||
|
||||
**Mock USDT**:
|
||||
The self-created 6-decimal SPL mint that stands in for USDT on devnet, where real USDT does not exist. The mint address is configuration, so mainnet cutover is a config change.
|
||||
_Avoid_: test token, fake USDT, devnet dollar
|
||||
|
||||
**Tax**:
|
||||
The share of caller payments distributed to compute nodes as rewards. Taxes are weighted by completed work and historical node speed so faster, larger nodes earn proportionally more.
|
||||
_Avoid_: fee, toll, commission
|
||||
|
||||
**Caller Credit**:
|
||||
Free starting balance granted to a new caller/API key so they can try the network before topping up.
|
||||
_Avoid_: signup bonus, faucet, airdrop
|
||||
|
||||
**Free Compute Job**:
|
||||
Work a compute node performs without earning immediate rewards, usually during probation or bootstrap phases.
|
||||
_Avoid_: unpaid labor, warmup request
|
||||
|
||||
**Slash**:
|
||||
The penalty for a proven fraud incident: the node's entire Pending Balance is forfeited to the Treasury and a Strike is recorded.
|
||||
_Avoid_: penalize, burn, fine, forfeit
|
||||
|
||||
**Strike**:
|
||||
A fraud incident recorded on-chain against a node. Enough strikes result in a ban.
|
||||
_Avoid_: infraction, violation, flag
|
||||
|
||||
**Ban**:
|
||||
Permanent exclusion of a wallet from the network after exceeding the strike threshold. Recorded on-chain.
|
||||
_Avoid_: blacklist, block, suspension
|
||||
|
||||
**Probationary Period**:
|
||||
The first N jobs a new wallet must complete without earning, to raise the cost of re-entering after a ban.
|
||||
_Avoid_: trial period, warmup, grace period
|
||||
|
||||
**Token**:
|
||||
TAI, our native Solana SPL token. Deferred (ADR-0015): nodes are currently paid directly in USDT; TAI returns as the reward/upside layer once volume exists, funded by the accumulated Protocol Cut. Clients never need to hold it.
|
||||
_Avoid_: coin, reward token, native token
|
||||
|
||||
**Contract Boundary**:
|
||||
The Python interface in `packages/contracts` that represents registry, payment, and settlement behavior. During the prototype it is implemented by deterministic local wrappers; later the same boundary is backed by real Solana programs.
|
||||
_Avoid_: mock contract, fake chain, temporary hack
|
||||
|
||||
**Validator**:
|
||||
A trusted node (or the tracker itself) that re-runs a sample of inference requests to detect fraud.
|
||||
_Avoid_: auditor, checker, referee
|
||||
|
||||
**Validation Event**:
|
||||
A completed inference record that contains enough information for a validator to decide whether to sample and re-run the request: session id, model preset, messages, inference route, node wallets, and observed output.
|
||||
_Avoid_: audit log, trace, receipt
|
||||
|
||||
**Slash Proof**:
|
||||
The record submitted by a validator when a sampled re-run diverges from the observed output beyond tolerance. In the prototype this is deterministic local contract state; later it maps to an on-chain proof transaction.
|
||||
_Avoid_: accusation, report, claim
|
||||
|
||||
### Client-facing
|
||||
|
||||
**Client**:
|
||||
Any application or user that sends inference requests to the gateway. Prepays USDT into the Treasury; each request is metered against the resulting ledger balance at a per-1K-tokens price set per model.
|
||||
_Avoid_: user, caller, consumer
|
||||
|
||||
**Model Preset**:
|
||||
A named, versioned model available on the network (e.g. `llama-3-70b`). The tracker knows which nodes hold which shards for each preset.
|
||||
_Avoid_: model, checkpoint, version
|
||||
|
||||
**Coverage Map**:
|
||||
The tracker's per-model mapping of layer ranges to node counts: `[(start_layer, end_layer, node_count), ...]`. A layer range with `node_count=0` is a coverage gap — the model is unroutable until the gap is filled. Coverage-first bin-packing fills all gaps before adding redundancy.
|
||||
_Avoid_: shard map, assignment table, coverage report
|
||||
|
||||
**Rebalance Directive**:
|
||||
A `LOAD_SHARD` or `DROP_SHARD` instruction the tracker issues to a node when the coverage map changes (node joins, node leaves, or load-balance reoptimization). Delivered as part of the node's heartbeat response.
|
||||
_Avoid_: rebalance command, shard instruction, migration order
|
||||
|
||||
**Node Score**:
|
||||
A throughput/latency rating the tracker maintains per node, used for route selection. Updated continuously from inference telemetry.
|
||||
_Avoid_: reputation, rating, rank
|
||||
|
||||
### Payments & fraud
|
||||
|
||||
**Stake**:
|
||||
Collateral a node stands to lose for fraud. In the current design the node's Pending Balance serves as stake — no upfront deposit is required. An optional USDT/TAI deposit may return later for routing priority.
|
||||
_Avoid_: deposit, bond, escrow
|
||||
|
||||
**Treasury**:
|
||||
The single project-owned Solana wallet that custodially holds client deposits, pays node payouts, and accumulates the Protocol Cut. Its keypair is loaded only on settlement-capable trackers.
|
||||
_Avoid_: escrow, vault, hot wallet
|
||||
|
||||
**Pending Balance**:
|
||||
A node's accrued, not-yet-paid USDT earnings on the tracker ledger. Doubles as the node's fraud collateral: it is forfeited in full when a validator catches a divergent output.
|
||||
_Avoid_: unpaid rewards, accrual, balance due
|
||||
|
||||
**Settlement Period**:
|
||||
The dynamic interval driving on-chain payouts: a node is paid when its Pending Balance exceeds the Payout Threshold or the period elapses, whichever comes first. Short in development (seconds), long in production (daily), configurable to grow with volume.
|
||||
_Avoid_: epoch, payout cycle, billing cycle
|
||||
|
||||
**Payout Threshold**:
|
||||
The minimum Pending Balance that triggers an immediate payout before the Settlement Period elapses. Includes a dust floor so payouts are never smaller than they are worth.
|
||||
_Avoid_: minimum payout, dust limit
|
||||
|
||||
**Protocol Cut**:
|
||||
The 10% of inference fees retained by the project for infrastructure; the remaining 90% goes to the nodes that served the request. Accumulates in the Treasury as the future TAI liquidity reserve.
|
||||
_Avoid_: spread, commission, house fee
|
||||
|
||||
**Deposit Watcher**:
|
||||
The tracker component that observes the Treasury's on-chain USDT deposits and credits the sending client's API-key ledger balance.
|
||||
_Avoid_: payment listener, chain scanner
|
||||
|
||||
**Mock USDT**:
|
||||
The self-created 6-decimal SPL mint that stands in for USDT on devnet, where real USDT does not exist. The mint address is configuration, so mainnet cutover is a config change.
|
||||
_Avoid_: test token, fake USDT, devnet dollar
|
||||
|
||||
**Tax**:
|
||||
The share of caller payments distributed to compute nodes as rewards. Taxes are weighted by completed work and historical node speed so faster, larger nodes earn proportionally more.
|
||||
_Avoid_: fee, toll, commission
|
||||
|
||||
**Caller Credit**:
|
||||
Free starting balance granted to a new caller/API key so they can try the network before topping up.
|
||||
_Avoid_: signup bonus, faucet, airdrop
|
||||
|
||||
**Free Compute Job**:
|
||||
Work a compute node performs without earning immediate rewards, usually during probation or bootstrap phases.
|
||||
_Avoid_: unpaid labor, warmup request
|
||||
|
||||
**Slash**:
|
||||
The penalty for a proven fraud incident: the node's entire Pending Balance is forfeited to the Treasury and a Strike is recorded.
|
||||
_Avoid_: penalize, burn, fine, forfeit
|
||||
|
||||
**Strike**:
|
||||
A fraud incident recorded on-chain against a node. Enough strikes result in a ban.
|
||||
_Avoid_: infraction, violation, flag
|
||||
|
||||
**Ban**:
|
||||
Permanent exclusion of a wallet from the network after exceeding the strike threshold. Recorded on-chain.
|
||||
_Avoid_: blacklist, block, suspension
|
||||
|
||||
**Probationary Period**:
|
||||
The first N jobs a new wallet must complete without earning, to raise the cost of re-entering after a ban.
|
||||
_Avoid_: trial period, warmup, grace period
|
||||
|
||||
**Token**:
|
||||
TAI, our native Solana SPL token. Deferred (ADR-0015): nodes are currently paid directly in USDT; TAI returns as the reward/upside layer once volume exists, funded by the accumulated Protocol Cut. Clients never need to hold it.
|
||||
_Avoid_: coin, reward token, native token
|
||||
|
||||
**Contract Boundary**:
|
||||
The Python interface in `packages/contracts` that represents registry, payment, and settlement behavior. During the prototype it is implemented by deterministic local wrappers; later the same boundary is backed by real Solana programs.
|
||||
_Avoid_: mock contract, fake chain, temporary hack
|
||||
|
||||
**Validator**:
|
||||
A trusted node (or the tracker itself) that re-runs a sample of inference requests to detect fraud.
|
||||
_Avoid_: auditor, checker, referee
|
||||
|
||||
**Validation Event**:
|
||||
A completed inference record that contains enough information for a validator to decide whether to sample and re-run the request: session id, model preset, messages, inference route, node wallets, and observed output.
|
||||
_Avoid_: audit log, trace, receipt
|
||||
|
||||
**Slash Proof**:
|
||||
The record submitted by a validator when a sampled re-run diverges from the observed output beyond tolerance. In the prototype this is deterministic local contract state; later it maps to an on-chain proof transaction.
|
||||
_Avoid_: accusation, report, claim
|
||||
|
||||
### Client-facing
|
||||
|
||||
**Client**:
|
||||
Any application or user that sends inference requests to the gateway. Prepays USDT into the Treasury; each request is metered against the resulting ledger balance at a per-1K-tokens price set per model.
|
||||
_Avoid_: user, caller, consumer
|
||||
|
||||
**Model Preset**:
|
||||
A named, versioned model available on the network (e.g. `llama-3-70b`). The tracker knows which nodes hold which shards for each preset.
|
||||
_Avoid_: model, checkpoint, version
|
||||
|
||||
1009
QUICKSTART.md
1009
QUICKSTART.md
File diff suppressed because it is too large
Load Diff
BIN
accounts.sqlite
BIN
accounts.sqlite
Binary file not shown.
BIN
billing.sqlite
BIN
billing.sqlite
Binary file not shown.
@@ -103,8 +103,32 @@ Verify the install:
|
||||
|
||||
```bash
|
||||
meshnet-node --help
|
||||
python -c "import transformers; print(transformers.__version__)"
|
||||
```
|
||||
|
||||
`transformers` must be **≥ 5.12** for Qwen3.5/3.6-MoE models (older versions fail
|
||||
with `'Qwen3_5MoeConfig' object has no attribute 'vocab_size'`). If you install
|
||||
into an existing conda/miniforge env instead of a fresh venv, run
|
||||
`pip install -U transformers` there. The startup warning about
|
||||
`flash-linear-attention` / `causal-conv1d` ("fast path is not available") is
|
||||
harmless on CPU — those are optional GPU kernels.
|
||||
|
||||
If you run the node from native Windows instead of WSL2, install Triton for
|
||||
Windows in the same environment:
|
||||
|
||||
```powershell
|
||||
python -m pip install triton-windows
|
||||
```
|
||||
|
||||
Without it, Qwen3.5/3.6-MoE startup can fail with the misleading message
|
||||
`Could not import module 'Qwen3_5MoeForCausalLM'`.
|
||||
|
||||
**NVIDIA GPU on native Windows:** the CUDA fast path works — after
|
||||
`triton-windows`, install FLA with plain `pip install flash-linear-attention`
|
||||
(no `[cuda]` extra, no `causal-conv1d`; both are Linux-only packaging and fail
|
||||
on Windows). No CUDA toolkit / `nvcc` is needed. See the platform table in
|
||||
[QUICKSTART.md](../QUICKSTART.md#qwen3536-moe-notes) for details.
|
||||
|
||||
---
|
||||
|
||||
## Step 6 — Pre-download the model shard
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
# ADR-0020: Dashboard chat streaming, live request progress, and the mixed-topology routing flaw
|
||||
|
||||
## Status: Accepted (chat/streaming/styles implemented); routing flaw documented, fix pending
|
||||
|
||||
## Context
|
||||
|
||||
Live alpha testing (2026-07-07) with `Qwen3.6-35B-A3B` split across two LAN nodes surfaced
|
||||
three UX gaps and one routing correctness flaw:
|
||||
|
||||
1. **No visibility while a request is processing.** The Call wall showed
|
||||
"no in-flight requests" during a 52-second generation. Cause: the dashboard chat sent
|
||||
`stream: false`, and the tracker only emits `proxy progress` console events (the Call
|
||||
wall's live-status source, `_tracker_log_proxy_progress`, `server.py` ~2199) for
|
||||
**streamed** requests. Non-streamed proxying produces only
|
||||
`route selected → connected → complete`, and short requests complete inside the
|
||||
dashboard's 4-second poll window.
|
||||
2. **Chat did not stream.** The nodes support SSE token-by-token generation
|
||||
(`generate_text_streaming`, hardened earlier for split shards), and the tracker proxy
|
||||
passes `text/event-stream` through (`server.py` ~3256), but the chat panel blocked on
|
||||
full JSON and showed nothing until completion.
|
||||
3. **Chat panel styles drifted.** The "new chat layout" redesign left hardcoded one-off
|
||||
colors (`#1f4788`, `#2563b8`, `#10151d`, `#1a1012`, `#5c2020`, `#ffb4b4`) mixed with
|
||||
the CSS custom-property palette.
|
||||
|
||||
## Decisions
|
||||
|
||||
### 1. Chat streams by default (SSE)
|
||||
|
||||
`dashboard.html` `sendChat()` now sends `stream: true` and consumes the SSE body with a
|
||||
`ReadableStream` reader:
|
||||
|
||||
- Assistant tokens render incrementally into the last bubble (direct DOM update, full
|
||||
re-render only at boundaries), with a blinking `▍` cursor while streaming.
|
||||
- Chat status shows live progress: `generating… N tokens · X tok/s`.
|
||||
- The send button becomes a stop button (`■`) during generation, backed by an
|
||||
`AbortController`; a stopped generation keeps the partial text.
|
||||
- Non-SSE responses (JSON fallback, errors) are still handled; `data: {"error": ...}`
|
||||
stream events surface as error bubbles.
|
||||
- `streaming` flags are stripped when loading persisted sessions so an interrupted
|
||||
generation never leaves a stuck cursor.
|
||||
|
||||
### 2. Live in-flight visibility rides on streaming
|
||||
|
||||
No tracker change was needed: because chat now streams, the tracker emits `proxy progress`
|
||||
events (throttled to stdout, updated in place in the console ring via
|
||||
`update_console_key`), and the existing Call wall state machine
|
||||
(`buildCallWallStates`) renders processing rows with live tokens/TPS/queue.
|
||||
|
||||
**Known limitation (accepted):** non-streamed API requests still show no progress between
|
||||
`proxy connected` and `proxy complete` — there is nothing to report until the node
|
||||
returns. Callers wanting live visibility should use `stream: true`.
|
||||
|
||||
### 3. Chat style tokens
|
||||
|
||||
All chat colors route through `:root` custom properties (`--hover-bg`, `--chat-user-bg`
|
||||
`#1f6feb`, `--chat-user-border`, `--chat-error-bg/border/fg`). No hardcoded hex values
|
||||
remain in chat rules, so future palette changes are single-line edits.
|
||||
|
||||
## Documented flaw: mixed-topology routing (partial GPU head + full CPU node)
|
||||
|
||||
### Observed (2026-07-07, tracker 192.168.0.179:8080)
|
||||
|
||||
Two nodes registered for `qwen3.6-35b-a3b`:
|
||||
|
||||
| node | hardware | shard | benchmark |
|
||||
|---|---|---|---|
|
||||
| `5gMLrmyB-ec3afe6f1a03` (192.168.0.20) | RTX 4060, CUDA | 0–21 (partial, fast) | 11,164 |
|
||||
| `7j77FsPY-55249b0583e5` (192.168.0.179) | CPU | 0–39 (full, slow) | 425 |
|
||||
|
||||
When the tracker selected the GPU node as head, it injected:
|
||||
|
||||
```
|
||||
downstream=[{"endpoint": "http://192.168.0.179:7000", "start_layer": 0}]
|
||||
```
|
||||
|
||||
`start_layer: 0` — not 22. The downstream full node re-ran **all 40 layers from layer 0
|
||||
on hidden states that had already passed through the head's layers 0–21**, producing
|
||||
garbage logits. Evidence from the logs:
|
||||
|
||||
- GPU-headed requests: `generation complete tokens=1` and billed `out=0`/`out=1`/`out=3`
|
||||
— near-instant EOS from corrupt activations.
|
||||
- The same prompt routed directly to the CPU full node: 209 tokens over 52 s (healthy).
|
||||
- Observed TPS for GPU-headed requests was meaningless (2.5–19.0 "tok/s" on 0–3 token
|
||||
outputs), and those samples now pollute the rolling per-`(node, model)` throughput
|
||||
stats used for routing preference.
|
||||
- Clients were **billed** for these broken 1-token responses.
|
||||
|
||||
### Root cause
|
||||
|
||||
The route planner treats the full-coverage node as a standalone complete route
|
||||
(`route=7j77FsPY…[0-39]`) but still injects it as the head's downstream with the
|
||||
downstream node's own `shard_start` (0) instead of `head.shard_end + 1` (22). A partial
|
||||
head + full-model downstream is a topology the planner never had to handle before —
|
||||
prior split tests used disjoint shards (0–11 + 12–23) where `shard_start` happened to
|
||||
equal the correct continuation layer.
|
||||
|
||||
### Required fix (not yet implemented)
|
||||
|
||||
1. **Correct continuation layer:** when hop N ends at layer `e`, hop N+1 must execute
|
||||
from `start_layer = e + 1` regardless of the downstream node's own `shard_start`
|
||||
(the `X-Meshnet-Start-Layer` overlapping-shard mechanism from ADR-0012 exists for
|
||||
exactly this; the planner must set it for full-model downstream nodes too).
|
||||
2. **Route preference sanity:** with a healthy single-node full route available, prefer
|
||||
it over a multi-hop route unless the pipeline is estimated faster; a fast head that
|
||||
forces a slow full-model tail wins nothing (every token still crosses the CPU node).
|
||||
3. **Stat hygiene:** exclude or flag throughput samples from responses with ≤ a few
|
||||
output tokens, so broken routes don't skew routing preference.
|
||||
4. **Billing guard (consider):** suspiciously short completions from multi-hop routes
|
||||
during this window were billed; a minimum-viability check (or refund path) may be
|
||||
warranted once audits land.
|
||||
|
||||
### Verification for the fix
|
||||
|
||||
Reproduce with a partial GPU head (0–21) + full CPU node (0–39): a chat request routed
|
||||
through the GPU head must produce output equivalent to the direct CPU route, with
|
||||
`downstream start_layer=22` visible in `proxy route selected`, and multi-token streamed
|
||||
output on the Call wall.
|
||||
|
||||
## Verification of this ADR's implemented changes
|
||||
|
||||
- `pytest tests/test_dashboard.py` — 5 passed (stale "Chat / inference" panel assertion
|
||||
updated to the tabbed layout).
|
||||
- Embedded dashboard JS parses (`new Function(script)` under Node 22).
|
||||
- Live check: open `/dashboard` → Chat, send a prompt to `qwen3.6-35b-a3b` — tokens
|
||||
must appear incrementally with live tok/s in the status line, the Call wall must show
|
||||
the request as `processing` with live TPS, and the send button must stop generation
|
||||
mid-stream keeping partial text.
|
||||
119
docs/adr/0021-dynamic-statistical-routing.md
Normal file
119
docs/adr/0021-dynamic-statistical-routing.md
Normal file
@@ -0,0 +1,119 @@
|
||||
# ADR-0021: Dynamic statistical routing (bandit-style route selection)
|
||||
|
||||
## Status: Accepted, implemented
|
||||
|
||||
## Context
|
||||
|
||||
ADR-0020 documented the mixed-topology flaw: with a fast GPU node serving layers 0–21 and
|
||||
a slow CPU node serving 0–39 of `Qwen3.6-35B-A3B`, the tracker picked the GPU node as
|
||||
proxy head *independently* of route planning, injecting a downstream hop with the wrong
|
||||
`start_layer` (0 instead of 22) and corrupting generation.
|
||||
|
||||
Beyond the bug, the deeper issue is that the tracker **cannot know a priori** which route
|
||||
is faster. Is one CPU node running all 40 layers faster than a GPU running 0–21 plus a
|
||||
CPU hop for 22–39? Benchmarks don't answer that — network hops, MoE expert loading, and
|
||||
queue dynamics only show up in real end-to-end requests. The router must *measure*.
|
||||
|
||||
## Decision
|
||||
|
||||
Route selection is a **multi-armed bandit** over enumerated candidate routes, implemented
|
||||
in `packages/tracker/meshnet_tracker/routing_stats.py` and wired into the chat proxy in
|
||||
`server.py`.
|
||||
|
||||
### Arms: route signatures
|
||||
|
||||
A route's identity is `model_key | node_id[shard] -> node_id[shard] -> …`. Node ids embed
|
||||
wallet + shard, so a node re-registering with a different shard produces a new arm
|
||||
automatically. The proxy target is **always the route's own head** (`route_nodes[0]`),
|
||||
and each hop's `start_layer` is `previous_hop.shard_end + 1` — this fixes ADR-0020's flaw
|
||||
structurally: head choice and route planning can no longer disagree.
|
||||
|
||||
### Candidate enumeration (`_enumerate_routes`)
|
||||
|
||||
One candidate per distinct head (a node whose `shard_start` equals the model's first
|
||||
layer — it must tokenize/embed), greedily completed with longest-advancing hops. Each
|
||||
candidate carries a `prior_tps`: its bottleneck hop's queue-adjusted effective throughput
|
||||
× reputation. Capped at 8 candidates ranked by prior.
|
||||
|
||||
### Statistics: decayed EWMA + topology epochs
|
||||
|
||||
Per (model, signature), `RouteStatsStore` keeps an EWMA of observed end-to-end tokens/sec
|
||||
with **time-decayed sample mass** (half-life default 600 s). Two staleness mechanisms
|
||||
handle the morphing network:
|
||||
|
||||
- **Continuous**: sample mass decays; a route unproven for a while (mass < 0.5) drops out
|
||||
of the exploit pool and gets re-scouted.
|
||||
- **Abrupt**: any node join/leave/shard-change bumps the model's *topology epoch*. Stats
|
||||
from an older epoch keep their EWMA as a display prior but are demoted to the scout
|
||||
pool ("stale") until re-measured under the new topology.
|
||||
|
||||
Sample hygiene: completions below `min_sample_tokens` (default 8) are rejected — the
|
||||
1-token garbage responses from the ADR-0020 bug would otherwise poison arms with
|
||||
meaningless tps values. Routes with no samples for 24 h are pruned.
|
||||
|
||||
### Selection policy (`choose_route`)
|
||||
|
||||
1. **Scout** (probability `explore_share`, default 0.3): if any candidate is unproven /
|
||||
stale / decayed, route the request there — least-measured first, tiebreak on prior.
|
||||
These are the user's "discovery/scout routes". With *no* proven arms at all, selection
|
||||
is deterministic best-prior (matches the old benchmark-based behavior, keeps cold
|
||||
start sane and tests deterministic).
|
||||
2. **Exploit** (otherwise): weighted random among proven arms with
|
||||
`P(route) ∝ tps^alpha`, `alpha` default 1.0 — a 1.5×-faster route gets 1.5× the
|
||||
traffic. `alpha` is a config knob: >1 shifts toward winner-takes-most as the network
|
||||
matures, without redesign. (Proportional split is not throughput-optimal in queueing
|
||||
terms, but it keeps every arm warm with fresh samples; tune alpha up when traffic
|
||||
justifies it.)
|
||||
|
||||
Pinned routes (`"route": [...]` in the request body) bypass the bandit but still record
|
||||
samples.
|
||||
|
||||
### Configuration
|
||||
|
||||
| CLI flag | env var | default |
|
||||
|---|---|---|
|
||||
| `--route-explore-share` | `MESHNET_ROUTE_EXPLORE_SHARE` | 0.3 |
|
||||
| `--route-weight-alpha` | `MESHNET_ROUTE_WEIGHT_ALPHA` | 1.0 |
|
||||
| `--route-stats-half-life` | `MESHNET_ROUTE_STATS_HALF_LIFE` | 600 |
|
||||
| — | `MESHNET_ROUTE_MIN_SAMPLE_TOKENS` | 8 |
|
||||
|
||||
High explore share now (development, few requests); drop toward 0.05–0.1 once real
|
||||
traffic provides passive coverage.
|
||||
|
||||
### Visibility
|
||||
|
||||
- **`GET /v1/routing`** (optionally `?model=`): per model — topology epoch and the full
|
||||
candidate table: hops, learned tps, **coefficient** (tps ÷ best proven route's tps),
|
||||
**expected traffic share**, sample count, decayed weight, status
|
||||
(proven / unsampled / stale / decayed).
|
||||
- **Dashboard → Overview → "Routing (learned)"**: renders that table live (4 s poll),
|
||||
with the active config in the header line.
|
||||
- **Console/`proxy route selected`** events now include the routing decision
|
||||
(`{"mode": "scout"|"exploit"|"pinned"|"greedy-fallback", "signature": …}`), so the Call
|
||||
wall history shows which arm served each request.
|
||||
|
||||
## Storage considerations
|
||||
|
||||
Stats are **in-memory per tracker** for alpha: they are cheap to relearn (a few requests
|
||||
per route), and gossiping them would import ADR-0019's consistency questions for data
|
||||
that is intentionally ephemeral. If multi-tracker route learning is needed later, ship
|
||||
route samples over the existing stats gossip and merge EWMAs by decayed weight — the
|
||||
store's (value, mass, timestamp) representation merges cleanly.
|
||||
|
||||
## Consequences
|
||||
|
||||
- The GPU(0–21)+CPU(0–39) topology now works: both routes get measured, the coefficient
|
||||
is visible on the dashboard, and traffic shifts to whichever is actually faster.
|
||||
- Routing is no longer deterministic once samples exist. Tests needing determinism seed
|
||||
`server.route_rng` or rely on the cold-start deterministic path.
|
||||
- The billing-relevant fix: heads are always part of the planned route, so per-hop
|
||||
`start_layer` and work-unit spans are consistent.
|
||||
|
||||
## Verification
|
||||
|
||||
`tests/test_dynamic_routing.py` (11 tests): EWMA/decay/epoch semantics, near-empty sample
|
||||
rejection, traffic split ≈ tps ratio at alpha=1 (0.6/0.4 over 4000 seeded draws), scout
|
||||
rate ≈ explore share, mixed-topology enumeration (both routes, hybrid prior = bottleneck),
|
||||
head-is-route-head regression with `start_layer=22` on the hybrid route, and `/v1/routing`
|
||||
table shape. Live: start both nodes, run several chats, open the dashboard "Routing
|
||||
(learned)" panel and watch coefficients converge.
|
||||
102
docs/adr/0022-sharded-per-node-kv-cache.md
Normal file
102
docs/adr/0022-sharded-per-node-kv-cache.md
Normal file
@@ -0,0 +1,102 @@
|
||||
# ADR-0022: Sharded per-node KV cache for distributed generation
|
||||
|
||||
## Status: Accepted, implemented (alpha-hardening issue 25)
|
||||
|
||||
## Context
|
||||
|
||||
The distributed generation loop (`torch_server.py`, `_do_chat_completions` distributed
|
||||
path) had **no KV cache**: every layer-forward call passed `use_cache: False`, and each
|
||||
autoregressive step re-encoded the entire prompt-so-far from scratch, re-running every
|
||||
layer on every node in the route for every generated token. Measured on a live 2-node
|
||||
Qwen2.5-0.5B GPU pipeline: tps decayed from 22.3 to 12.6 within a single generation —
|
||||
the quadratic-cost signature. On Qwen3.6-35B-A3B mixed GPU/CPU topology this collapsed
|
||||
to ~0.07 tps even after the ADR-0020 routing fix.
|
||||
|
||||
`X-Meshnet-Session` existed on the wire but was minted fresh **per token** and keyed no
|
||||
state.
|
||||
|
||||
## Decision
|
||||
|
||||
### Session lifecycle
|
||||
|
||||
The head mints one session id per chat generation (not per token) and reuses it across
|
||||
every step. Two new request headers extend the `/forward` wire protocol:
|
||||
|
||||
- `X-Meshnet-Cache: prefill | decode` — absent means legacy stateless (unchanged
|
||||
behavior, and what old nodes send/understand).
|
||||
- `X-Meshnet-Past-Len: N` — decode only: the number of tokens the node's session cache
|
||||
must already hold. A mismatch is a cache miss, never silent corruption.
|
||||
|
||||
Step 0 (`prefill`) sends the full prompt activation as before; each node creates fresh
|
||||
session state for its own layer range. Steps 1+ (`decode`) send only the newest token's
|
||||
hidden state — `[1, 1, hidden]`, cutting per-step compute and wire payload from
|
||||
O(seq_len) to O(1). The head embeds the next token directly from the `token_id` the tail
|
||||
now returns alongside text (`{"text": …, "token_id": …}`), avoiding text
|
||||
re-tokenization drift; EOS is detected by id against tokenizer + generation-config eos
|
||||
sets.
|
||||
|
||||
### Per-node sharded cache
|
||||
|
||||
`TorchModelShard.kv_sessions` is a `SessionCacheStore`: `session_id → SessionCacheEntry`
|
||||
holding cache state **only for that shard's layer range** — sharding falls out naturally
|
||||
because each node only executes (and therefore only caches) its own layers. No node ever
|
||||
holds another node's state.
|
||||
|
||||
### MoE / hybrid-attention awareness
|
||||
|
||||
The cached object is whatever `use_cache=True` produces: a transformers
|
||||
`DynamicCache(config=model.config)` — the same construction the model's own `forward()`
|
||||
uses. With the config, transformers picks the right per-layer state: K/V tensors for
|
||||
standard attention, conv/recurrent delta state for Qwen3.6-style hybrid linear-attention
|
||||
layers, sliding-window variants, etc. The store treats it as opaque; nothing assumes a
|
||||
K/V tensor shape. Cache slots are indexed by absolute `layer_idx`, so a shard updating
|
||||
only layers 12–23 leaves 0–11 empty (verified: sparse `DynamicCache.update` works).
|
||||
MoE expert routing is layer-local per token and needs no cross-token state.
|
||||
|
||||
Layers are invoked with `past_key_values=<cache>, use_cache=True, cache_position=…`
|
||||
(transformers 5.x layer API; the cache is mutated in place). If a model's layers reject
|
||||
those kwargs, the backend logs once, sets `supports_kv_cache = False`, and stays on the
|
||||
stateless path permanently — exotic architectures degrade to today's behavior instead of
|
||||
failing.
|
||||
|
||||
### Cache miss and route-change interaction (ADR-0021)
|
||||
|
||||
Any decode-mode request that cannot be served — unknown session (evicted, node
|
||||
restarted), `past_len` mismatch, `start_layer` mismatch (the route or shard overlap
|
||||
changed mid-generation), or caching disabled — raises `KVCacheMiss`, answered as
|
||||
**HTTP 409 `{"error": "cache_miss"}`**. The head catches it and falls back to one full
|
||||
re-prefill of the accumulated sequence under the same session id, which atomically
|
||||
replaces every node's session state, then continues cached. The fully-stateless path is
|
||||
therefore still the recovery mode: eviction and restarts cost one prefill, never
|
||||
corruption or a failed generation. A decode request against a node whose caching is
|
||||
disabled is also a 409 — running a single-token payload statelessly would silently
|
||||
produce garbage.
|
||||
|
||||
Mixed fleets degrade the same way: if the tail predates the protocol and returns no
|
||||
`token_id`, the head simply prefills every step (exactly the old cost).
|
||||
|
||||
### Bounded memory
|
||||
|
||||
`SessionCacheStore` enforces TTL (default 600 s, `MESHNET_KV_TTL_SECONDS`) plus LRU cap
|
||||
(default 8 sessions, `MESHNET_KV_MAX_SESSIONS`), evaluated on every access. The head
|
||||
additionally drops its own session explicitly when a generation completes; downstream
|
||||
nodes rely on TTL/LRU (an explicit cross-node release RPC was judged not worth the
|
||||
failure modes — misses are cheap).
|
||||
|
||||
### Non-goals (first landing)
|
||||
|
||||
Cross-node cache migration on route change (evict + re-prefill is acceptable),
|
||||
speculative decoding, cross-session batching.
|
||||
|
||||
## Consequences
|
||||
|
||||
- Per-token cost drops from O(seq_len) layer re-execution + O(seq_len) wire transfer per
|
||||
hop to O(1) of both; tps stays flat across generation length instead of decaying.
|
||||
- Golden test (`tests/test_kv_cache_distributed.py`, env-gated by
|
||||
`MESHNET_REAL_MODEL_TESTS=1`) proves cached and stateless distributed generation emit
|
||||
identical token ids on a real two-shard Qwen2.5-0.5B split.
|
||||
- Nodes now hold per-session GPU/CPU memory between requests (bounded above); operators
|
||||
sizing `max_loaded_shards` should account for ~`sessions × seq_len × kv_bytes_per_token`
|
||||
per resident model.
|
||||
- The wire protocol is backward- and forward-compatible: headers are additive, absent
|
||||
headers mean stateless, and 409 is only sent in reply to explicit decode-mode requests.
|
||||
@@ -1,48 +1,48 @@
|
||||
# US-020 — Manual route selection + hop-penalty benchmarking
|
||||
|
||||
## Context
|
||||
|
||||
The tracker auto-selects inference routes based on synthetic benchmark scores. To measure
|
||||
the real cost of adding hops (latency per node boundary), we need:
|
||||
|
||||
1. A way to pin a request to a specific route so we control the variable.
|
||||
2. A benchmark endpoint that runs the same prompt through 1-node, 2-node, and 3-node
|
||||
routes and records per-hop latency.
|
||||
|
||||
Results are stored to disk. Routing algorithm is **not** changed in this story — this is
|
||||
data collection only. The data will inform a future routing optimisation story.
|
||||
|
||||
## Design decisions (grilled 2026-07-01)
|
||||
|
||||
| Decision | Choice |
|
||||
|---|---|
|
||||
| Route spec | Optional `route` field in JSON request body (list of node IDs) |
|
||||
| Trigger | Explicit only — `POST /v1/benchmark/hop-penalty` endpoint |
|
||||
| Auth | Header-presence stub (`Authorization` must be non-empty); real auth in future story |
|
||||
| Routing integration | Store data only; routing algorithm unchanged |
|
||||
| Persistence | Append to `benchmark_results.json` in tracker working dir; in-memory queryable |
|
||||
|
||||
## Acceptance criteria
|
||||
|
||||
- `POST /v1/chat/completions` accepts optional `"route": ["<node_id>", ...]` in the
|
||||
request body. If present, the tracker uses those nodes in order instead of auto-selecting.
|
||||
If absent, existing routing is unchanged (no breaking change for unaware clients).
|
||||
- Missing or invalid node IDs in `route` return HTTP 400 with a descriptive error.
|
||||
- `POST /v1/benchmark/hop-penalty` is auth-gated: requests without a non-empty
|
||||
`Authorization` header return HTTP 401. Body: `{"model": "...", "prompt": "...",
|
||||
"max_new_tokens": 64}`.
|
||||
- Benchmark fans out to up to three routes: 1-node (single node covering all layers),
|
||||
2-node (two consecutive shard nodes), 3-node (three nodes) — using whatever is
|
||||
currently registered. Routes with insufficient coverage are skipped, not errored.
|
||||
- Response includes per-route breakdown: `total_ms`, `per_hop_ms: [...]`,
|
||||
`tokens_generated`, `route: [node_id, ...]`.
|
||||
- Results are appended to `<tracker_working_dir>/benchmark_results.json` (created if
|
||||
absent) as a JSON array. Each entry includes timestamp, model, prompt hash, and the
|
||||
per-route breakdown.
|
||||
- `GET /v1/benchmark/results` returns the stored results array. Also auth-gated.
|
||||
- Clients that never send `route` or call `/v1/benchmark/*` are completely unaffected.
|
||||
- Integration test: send the same prompt via a pinned 1-node route and a pinned 2-node
|
||||
route; assert 2-node result has 2 entries in `per_hop_ms`; assert both records appear
|
||||
in `benchmark_results.json`.
|
||||
- `python -m pytest` passes from repo root.
|
||||
- Commit only this story's changes.
|
||||
# US-020 — Manual route selection + hop-penalty benchmarking
|
||||
|
||||
## Context
|
||||
|
||||
The tracker auto-selects inference routes based on synthetic benchmark scores. To measure
|
||||
the real cost of adding hops (latency per node boundary), we need:
|
||||
|
||||
1. A way to pin a request to a specific route so we control the variable.
|
||||
2. A benchmark endpoint that runs the same prompt through 1-node, 2-node, and 3-node
|
||||
routes and records per-hop latency.
|
||||
|
||||
Results are stored to disk. Routing algorithm is **not** changed in this story — this is
|
||||
data collection only. The data will inform a future routing optimisation story.
|
||||
|
||||
## Design decisions (grilled 2026-07-01)
|
||||
|
||||
| Decision | Choice |
|
||||
|---|---|
|
||||
| Route spec | Optional `route` field in JSON request body (list of node IDs) |
|
||||
| Trigger | Explicit only — `POST /v1/benchmark/hop-penalty` endpoint |
|
||||
| Auth | Header-presence stub (`Authorization` must be non-empty); real auth in future story |
|
||||
| Routing integration | Store data only; routing algorithm unchanged |
|
||||
| Persistence | Append to `benchmark_results.json` in tracker working dir; in-memory queryable |
|
||||
|
||||
## Acceptance criteria
|
||||
|
||||
- `POST /v1/chat/completions` accepts optional `"route": ["<node_id>", ...]` in the
|
||||
request body. If present, the tracker uses those nodes in order instead of auto-selecting.
|
||||
If absent, existing routing is unchanged (no breaking change for unaware clients).
|
||||
- Missing or invalid node IDs in `route` return HTTP 400 with a descriptive error.
|
||||
- `POST /v1/benchmark/hop-penalty` is auth-gated: requests without a non-empty
|
||||
`Authorization` header return HTTP 401. Body: `{"model": "...", "prompt": "...",
|
||||
"max_new_tokens": 64}`.
|
||||
- Benchmark fans out to up to three routes: 1-node (single node covering all layers),
|
||||
2-node (two consecutive shard nodes), 3-node (three nodes) — using whatever is
|
||||
currently registered. Routes with insufficient coverage are skipped, not errored.
|
||||
- Response includes per-route breakdown: `total_ms`, `per_hop_ms: [...]`,
|
||||
`tokens_generated`, `route: [node_id, ...]`.
|
||||
- Results are appended to `<tracker_working_dir>/benchmark_results.json` (created if
|
||||
absent) as a JSON array. Each entry includes timestamp, model, prompt hash, and the
|
||||
per-route breakdown.
|
||||
- `GET /v1/benchmark/results` returns the stored results array. Also auth-gated.
|
||||
- Clients that never send `route` or call `/v1/benchmark/*` are completely unaffected.
|
||||
- Integration test: send the same prompt via a pinned 1-node route and a pinned 2-node
|
||||
route; assert 2-node result has 2 entries in `per_hop_ms`; assert both records appear
|
||||
in `benchmark_results.json`.
|
||||
- `python -m pytest` passes from repo root.
|
||||
- Commit only this story's changes.
|
||||
|
||||
Binary file not shown.
@@ -286,7 +286,7 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
|
||||
self._send_json(200, completion)
|
||||
|
||||
def _proxy_to_head_worker(self, url: str, body_bytes: bytes) -> None:
|
||||
"""Forward a raw request body to a head worker and stream the response back."""
|
||||
"""Forward a raw request body to a head worker and relay SSE without buffering."""
|
||||
target_url = f"{url}/v1/chat/completions"
|
||||
req = urllib.request.Request(
|
||||
target_url,
|
||||
@@ -297,6 +297,19 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=30.0) as r:
|
||||
content_type = r.headers.get("Content-Type", "application/json")
|
||||
if "text/event-stream" in content_type:
|
||||
self.send_response(r.status)
|
||||
self.send_header("Content-Type", content_type)
|
||||
self.send_header("Cache-Control", "no-cache")
|
||||
self.send_header("X-Accel-Buffering", "no")
|
||||
self.end_headers()
|
||||
while True:
|
||||
line = r.readline()
|
||||
if not line:
|
||||
break
|
||||
self.wfile.write(line)
|
||||
self.wfile.flush()
|
||||
return
|
||||
resp_body = r.read()
|
||||
status = r.status
|
||||
except urllib.error.HTTPError as exc:
|
||||
|
||||
@@ -68,6 +68,7 @@ def _run_node(cfg: dict) -> None:
|
||||
tracker_source_disabled=bool(cfg.get("tracker_source_disabled", False)),
|
||||
torch_threads=cfg.get("torch_threads"),
|
||||
torch_interop_threads=cfg.get("torch_interop_threads"),
|
||||
node_name=cfg.get("node_name"),
|
||||
)
|
||||
except Exception as exc:
|
||||
print(f"\nERROR: {exc}", file=sys.stderr, flush=True)
|
||||
@@ -157,6 +158,8 @@ def _cmd_default(args) -> int:
|
||||
overrides["host"] = args.host
|
||||
if args.advertise_host:
|
||||
overrides["advertise_host"] = args.advertise_host
|
||||
if getattr(args, "node_name", None):
|
||||
overrides["node_name"] = args.node_name
|
||||
if args.route_timeout != 30.0:
|
||||
overrides["route_timeout"] = args.route_timeout
|
||||
if getattr(args, "memory", None) is not None:
|
||||
@@ -246,6 +249,8 @@ def _cmd_start(args) -> int:
|
||||
cfg["wallet_path"] = args.wallet
|
||||
if args.download_dir:
|
||||
cfg["download_dir"] = args.download_dir
|
||||
if getattr(args, "node_name", None):
|
||||
cfg["node_name"] = args.node_name
|
||||
|
||||
# Legacy start: just run without the dashboard (keep original blocking loop)
|
||||
from .startup import run_startup
|
||||
@@ -270,6 +275,7 @@ def _cmd_start(args) -> int:
|
||||
tracker_source_disabled=getattr(args, "tracker_source_disabled", False),
|
||||
torch_threads=getattr(args, "torch_threads", None),
|
||||
torch_interop_threads=getattr(args, "torch_interop_threads", None),
|
||||
node_name=cfg.get("node_name"),
|
||||
)
|
||||
except Exception as exc:
|
||||
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
|
||||
@@ -315,6 +321,7 @@ def main() -> None:
|
||||
parser.add_argument("--port", type=int, metavar="N", help="Port to listen on")
|
||||
parser.add_argument("--host", metavar="ADDR", help="Interface to bind (default 0.0.0.0)")
|
||||
parser.add_argument("--advertise-host", metavar="ADDR", help="Host/IP advertised to the tracker")
|
||||
parser.add_argument("--node-name", metavar="NAME", help="Friendly display name shown on the tracker dashboard")
|
||||
parser.add_argument("--route-timeout", type=float, metavar="SEC", default=30.0,
|
||||
help="Seconds to wait for tracker route lookup (default 30)")
|
||||
parser.add_argument("--memory", type=int, metavar="MB", default=None,
|
||||
@@ -350,6 +357,7 @@ def main() -> None:
|
||||
start_cmd.add_argument("--quantization", choices=["auto", "bfloat16", "int8", "nf4", "bf16"], default="auto")
|
||||
start_cmd.add_argument("--host", default="0.0.0.0")
|
||||
start_cmd.add_argument("--advertise-host")
|
||||
start_cmd.add_argument("--node-name", help="Friendly display name shown on the tracker dashboard")
|
||||
start_cmd.add_argument("--tracker-mode", action="store_true")
|
||||
start_cmd.add_argument("--tracker-url", default=None)
|
||||
start_cmd.add_argument("--wallet")
|
||||
|
||||
@@ -3,8 +3,12 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
@@ -27,12 +31,128 @@ class PartialModelLoadUnsupported(ModelBackendError):
|
||||
"""Raised when a shard cannot be materialized from a local snapshot subset."""
|
||||
|
||||
|
||||
class KVCacheMiss(ModelBackendError):
|
||||
"""Raised when a decode step references session state this node no longer holds.
|
||||
|
||||
The head recovers by re-prefilling the full sequence (the stateless path),
|
||||
so eviction or a node restart degrades throughput instead of corrupting output.
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TensorPayload:
|
||||
body: bytes
|
||||
shape: list[int]
|
||||
attention_mask_header: str | None
|
||||
position_ids_header: str | None
|
||||
# Number of tokens already cached before this payload's tokens (decode steps).
|
||||
past_len: int | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TailTokenResult:
|
||||
"""Tail-shard decode result: decoded text plus the raw token id.
|
||||
|
||||
The token id lets the head feed the next decode step (and detect EOS)
|
||||
without re-tokenizing text, which is not guaranteed to round-trip.
|
||||
"""
|
||||
|
||||
text: str
|
||||
token_id: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionCacheEntry:
|
||||
"""Per-session cached state for one shard's layer range.
|
||||
|
||||
`cache` is whatever `use_cache=True` produces for these layers — a
|
||||
transformers Cache holding K/V tensors for standard attention, or
|
||||
recurrent conv/delta state for hybrid linear-attention layers. The store
|
||||
treats it as opaque.
|
||||
"""
|
||||
|
||||
cache: Any
|
||||
seq_len: int
|
||||
effective_start: int
|
||||
last_used: float
|
||||
|
||||
|
||||
class SessionCacheStore:
|
||||
"""TTL + LRU bounded map of session_id → SessionCacheEntry.
|
||||
|
||||
Each node caches state only for its own layer range; no node ever holds
|
||||
another node's cache. Stale or mismatched entries raise KVCacheMiss so the
|
||||
head falls back to a full re-prefill instead of producing corrupt output.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_sessions: int = 8,
|
||||
ttl_seconds: float = 600.0,
|
||||
clock: Any = None,
|
||||
) -> None:
|
||||
self.max_sessions = max(1, int(max_sessions))
|
||||
self.ttl_seconds = float(ttl_seconds)
|
||||
self._clock = clock or time.monotonic
|
||||
self._entries: OrderedDict[str, SessionCacheEntry] = OrderedDict()
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def __len__(self) -> int:
|
||||
with self._lock:
|
||||
return len(self._entries)
|
||||
|
||||
def store(self, session_id: str, cache: Any, seq_len: int, effective_start: int) -> SessionCacheEntry:
|
||||
now = self._clock()
|
||||
with self._lock:
|
||||
self._entries.pop(session_id, None)
|
||||
entry = SessionCacheEntry(cache, seq_len, effective_start, now)
|
||||
self._entries[session_id] = entry
|
||||
self._evict_locked(now)
|
||||
return entry
|
||||
|
||||
def lookup(
|
||||
self,
|
||||
session_id: str,
|
||||
*,
|
||||
expected_seq_len: int | None = None,
|
||||
effective_start: int | None = None,
|
||||
) -> SessionCacheEntry:
|
||||
now = self._clock()
|
||||
with self._lock:
|
||||
self._evict_locked(now)
|
||||
entry = self._entries.get(session_id)
|
||||
if entry is None:
|
||||
raise KVCacheMiss(f"no cached state for session {session_id[:8]}")
|
||||
if expected_seq_len is not None and entry.seq_len != expected_seq_len:
|
||||
del self._entries[session_id]
|
||||
raise KVCacheMiss(
|
||||
f"session {session_id[:8]} cache holds {entry.seq_len} tokens, "
|
||||
f"expected {expected_seq_len}"
|
||||
)
|
||||
if effective_start is not None and entry.effective_start != effective_start:
|
||||
del self._entries[session_id]
|
||||
raise KVCacheMiss(
|
||||
f"session {session_id[:8]} cached with start_layer "
|
||||
f"{entry.effective_start}, requested {effective_start}"
|
||||
)
|
||||
entry.last_used = now
|
||||
self._entries.move_to_end(session_id)
|
||||
return entry
|
||||
|
||||
def drop(self, session_id: str) -> None:
|
||||
with self._lock:
|
||||
self._entries.pop(session_id, None)
|
||||
|
||||
def _evict_locked(self, now: float) -> None:
|
||||
if self.ttl_seconds > 0:
|
||||
expired = [
|
||||
sid for sid, entry in self._entries.items()
|
||||
if now - entry.last_used > self.ttl_seconds
|
||||
]
|
||||
for sid in expired:
|
||||
del self._entries[sid]
|
||||
while len(self._entries) > self.max_sessions:
|
||||
self._entries.popitem(last=False)
|
||||
|
||||
|
||||
def validate_quantization(value: str) -> Quantization:
|
||||
@@ -134,8 +254,9 @@ class TorchModelShard:
|
||||
self.model.to(self.device)
|
||||
except Exception as exc:
|
||||
if _looks_like_oom(exc):
|
||||
memory_kind = "VRAM" if self.device.type == "cuda" else "RAM"
|
||||
raise InsufficientVRAMError(
|
||||
f"insufficient VRAM to load {model_id} layers {shard_start}:{shard_end} "
|
||||
f"insufficient {memory_kind} to load {model_id} layers {shard_start}:{shard_end} "
|
||||
f"with {quantization} quantization; choose a smaller shard or lower quantization"
|
||||
) from exc
|
||||
raise
|
||||
@@ -162,8 +283,14 @@ class TorchModelShard:
|
||||
self._position_embeddings = _position_embeddings(self.model)
|
||||
self._norm = _final_norm(self.model) if self.is_tail else None
|
||||
self._lm_head = getattr(self.model, "lm_head", None) if self.is_tail else None
|
||||
# Per-session KV/recurrent-state cache for this shard's layer range.
|
||||
self.supports_kv_cache = True
|
||||
self.kv_sessions = SessionCacheStore(
|
||||
max_sessions=int(os.environ.get("MESHNET_KV_MAX_SESSIONS", "8")),
|
||||
ttl_seconds=float(os.environ.get("MESHNET_KV_TTL_SECONDS", "600")),
|
||||
)
|
||||
|
||||
def encode_prompt(self, prompt: str) -> TensorPayload:
|
||||
def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload:
|
||||
if not self.is_head or self._embed_tokens is None:
|
||||
raise ModelBackendError("text prompts can only be accepted by the head shard")
|
||||
encoded = self.tokenizer(prompt, return_tensors="pt")
|
||||
@@ -176,9 +303,44 @@ class TorchModelShard:
|
||||
hidden_states = self._embed_tokens(input_ids)
|
||||
if self._position_embeddings is not None:
|
||||
hidden_states = hidden_states + self._position_embeddings(position_ids)
|
||||
hidden_states = self._run_layers(hidden_states, attention_mask, position_ids)
|
||||
hidden_states = self._run_layers_session(
|
||||
hidden_states, attention_mask, position_ids,
|
||||
session_id=session_id, cache_mode="prefill" if session_id else None,
|
||||
)
|
||||
return self._payload(hidden_states, attention_mask, position_ids)
|
||||
|
||||
def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload:
|
||||
"""Decode step: embed one new token against this head's cached session.
|
||||
|
||||
Raises KVCacheMiss if the session was evicted — callers fall back to a
|
||||
full re-prefill via encode_prompt.
|
||||
"""
|
||||
if not self.is_head or self._embed_tokens is None:
|
||||
raise ModelBackendError("decode steps can only start at the head shard")
|
||||
if not self.supports_kv_cache:
|
||||
raise KVCacheMiss("kv cache disabled on this backend")
|
||||
entry = self.kv_sessions.lookup(
|
||||
session_id, effective_start=self._effective_start(None)
|
||||
)
|
||||
past_len = entry.seq_len
|
||||
input_ids = self.torch.tensor([[int(token_id)]], dtype=self.torch.long, device=self.device)
|
||||
position_ids = self.torch.tensor([[past_len]], dtype=self.torch.long, device=self.device)
|
||||
hidden_states = self._embed_tokens(input_ids)
|
||||
if self._position_embeddings is not None:
|
||||
hidden_states = hidden_states + self._position_embeddings(position_ids)
|
||||
hidden_states = self._run_layers(
|
||||
hidden_states, None, position_ids,
|
||||
cache=entry.cache, past_len=past_len,
|
||||
)
|
||||
entry.seq_len = past_len + 1
|
||||
return TensorPayload(
|
||||
body=_tensor_to_bytes(hidden_states.to(self.torch.bfloat16).contiguous()),
|
||||
shape=list(hidden_states.shape),
|
||||
attention_mask_header=None,
|
||||
position_ids_header=_int_tensor_header(position_ids),
|
||||
past_len=past_len,
|
||||
)
|
||||
|
||||
def forward_bytes(
|
||||
self,
|
||||
body: bytes,
|
||||
@@ -186,7 +348,10 @@ class TorchModelShard:
|
||||
attention_mask_header: str | None,
|
||||
position_ids_header: str | None,
|
||||
start_layer: int | None = None,
|
||||
) -> TensorPayload | str:
|
||||
session_id: str | None = None,
|
||||
cache_mode: str | None = None,
|
||||
past_len: int | None = None,
|
||||
) -> TensorPayload | TailTokenResult | str:
|
||||
hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to(
|
||||
self.device
|
||||
)
|
||||
@@ -196,26 +361,51 @@ class TorchModelShard:
|
||||
position_ids = _tensor_from_int64_header(
|
||||
position_ids_header, self.torch, self.device
|
||||
)
|
||||
hidden_states = self._run_layers(
|
||||
hidden_states, attention_mask, position_ids, start_layer=start_layer
|
||||
hidden_states = self._run_layers_session(
|
||||
hidden_states, attention_mask, position_ids, start_layer=start_layer,
|
||||
session_id=session_id, cache_mode=cache_mode, past_len=past_len,
|
||||
)
|
||||
if self.is_tail:
|
||||
return self.decode_tail(hidden_states)
|
||||
return self.decode_tail_token(hidden_states)
|
||||
return self._payload(hidden_states, attention_mask, position_ids)
|
||||
|
||||
def decode_tail(self, hidden_states: Any) -> str:
|
||||
return self.decode_tail_token(hidden_states).text
|
||||
|
||||
def decode_tail_token(self, hidden_states: Any) -> TailTokenResult:
|
||||
if self._norm is not None:
|
||||
hidden_states = self._norm(hidden_states)
|
||||
if self._lm_head is None:
|
||||
raise ModelBackendError("tail shard has no lm_head")
|
||||
logits = self._lm_head(hidden_states)
|
||||
token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item())
|
||||
return self.tokenizer.decode([token_id], skip_special_tokens=True)
|
||||
return TailTokenResult(
|
||||
text=self.tokenizer.decode([token_id], skip_special_tokens=True),
|
||||
token_id=token_id,
|
||||
)
|
||||
|
||||
def eos_token_ids(self) -> list[int]:
|
||||
"""All token ids that should terminate generation (tokenizer + generation config)."""
|
||||
ids: set[int] = set()
|
||||
tok_eos = getattr(self.tokenizer, "eos_token_id", None)
|
||||
gen_config = getattr(self.model, "generation_config", None)
|
||||
gen_eos = getattr(gen_config, "eos_token_id", None) if gen_config is not None else None
|
||||
for value in (tok_eos, gen_eos):
|
||||
if value is None:
|
||||
continue
|
||||
if isinstance(value, (list, tuple)):
|
||||
ids.update(int(v) for v in value)
|
||||
else:
|
||||
ids.add(int(value))
|
||||
return sorted(ids)
|
||||
|
||||
def release_session(self, session_id: str) -> None:
|
||||
self.kv_sessions.drop(session_id)
|
||||
|
||||
def generate_text(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 256,
|
||||
max_new_tokens: int = 5120,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
) -> str:
|
||||
@@ -245,7 +435,7 @@ class TorchModelShard:
|
||||
def generate_text_streaming(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 256,
|
||||
max_new_tokens: int = 5000,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
):
|
||||
@@ -321,21 +511,108 @@ class TorchModelShard:
|
||||
)
|
||||
return dict(self.tokenizer(prompt, return_tensors="pt"))
|
||||
|
||||
def _effective_start(self, start_layer: int | None) -> int:
|
||||
# start_layer overrides shard_start for overlapping-shard routing
|
||||
# (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent
|
||||
# indexing outside the loaded weights.
|
||||
return (
|
||||
max(self.shard_start, start_layer)
|
||||
if start_layer is not None
|
||||
else self.shard_start
|
||||
)
|
||||
|
||||
def _new_session_cache(self) -> Any | None:
|
||||
"""Build the model-appropriate cache object for one session.
|
||||
|
||||
DynamicCache(config=...) lets transformers pick the right per-layer
|
||||
state (K/V for standard attention, conv/recurrent state for hybrid
|
||||
linear-attention layers) — the same construction the model's own
|
||||
forward() uses when use_cache=True.
|
||||
"""
|
||||
try:
|
||||
from transformers import DynamicCache
|
||||
except ImportError:
|
||||
return None
|
||||
try:
|
||||
return DynamicCache(config=self.model.config)
|
||||
except TypeError:
|
||||
return DynamicCache()
|
||||
|
||||
def _run_layers_session(
|
||||
self,
|
||||
hidden_states: Any,
|
||||
attention_mask: Any,
|
||||
position_ids: Any,
|
||||
start_layer: int | None = None,
|
||||
session_id: str | None = None,
|
||||
cache_mode: str | None = None,
|
||||
past_len: int | None = None,
|
||||
) -> Any:
|
||||
"""Run this shard's layers, keying cached state by session when requested.
|
||||
|
||||
cache_mode "prefill" creates fresh session state; "decode" requires an
|
||||
existing entry (KVCacheMiss otherwise). None runs fully stateless —
|
||||
today's behavior, kept as the recovery path.
|
||||
"""
|
||||
effective_start = self._effective_start(start_layer)
|
||||
if not (session_id and cache_mode and self.supports_kv_cache):
|
||||
if cache_mode == "decode":
|
||||
# A decode payload is one token — running it stateless would
|
||||
# silently produce garbage. Force the head to re-prefill.
|
||||
raise KVCacheMiss("kv cache disabled on this backend")
|
||||
return self._run_layers(
|
||||
hidden_states, attention_mask, position_ids, start_layer=start_layer
|
||||
)
|
||||
if cache_mode == "decode":
|
||||
entry = self.kv_sessions.lookup(
|
||||
session_id,
|
||||
expected_seq_len=past_len,
|
||||
effective_start=effective_start,
|
||||
)
|
||||
seq_len = int(hidden_states.shape[1])
|
||||
# Decode attends over cache + new token; no padding, so no mask needed.
|
||||
hidden_states = self._run_layers(
|
||||
hidden_states, None, position_ids,
|
||||
start_layer=start_layer, cache=entry.cache, past_len=entry.seq_len,
|
||||
)
|
||||
entry.seq_len += seq_len
|
||||
return hidden_states
|
||||
# Prefill: fresh cache for this session (replaces any stale entry).
|
||||
cache = self._new_session_cache()
|
||||
if cache is None:
|
||||
return self._run_layers(
|
||||
hidden_states, attention_mask, position_ids, start_layer=start_layer
|
||||
)
|
||||
try:
|
||||
result = self._run_layers(
|
||||
hidden_states, attention_mask, position_ids,
|
||||
start_layer=start_layer, cache=cache, past_len=0,
|
||||
)
|
||||
except TypeError as exc:
|
||||
# Layers reject cache kwargs (exotic architecture) — disable caching
|
||||
# for this backend and stay on the stateless path.
|
||||
self.supports_kv_cache = False
|
||||
print(f" [node] kv cache unsupported by {self.model_id}: {exc}", flush=True)
|
||||
return self._run_layers(
|
||||
hidden_states, attention_mask, position_ids, start_layer=start_layer
|
||||
)
|
||||
self.kv_sessions.store(
|
||||
session_id, cache,
|
||||
seq_len=int(hidden_states.shape[1]),
|
||||
effective_start=effective_start,
|
||||
)
|
||||
return result
|
||||
|
||||
def _run_layers(
|
||||
self,
|
||||
hidden_states: Any,
|
||||
attention_mask: Any,
|
||||
position_ids: Any,
|
||||
start_layer: int | None = None,
|
||||
cache: Any = None,
|
||||
past_len: int = 0,
|
||||
) -> Any:
|
||||
# start_layer overrides shard_start for overlapping-shard routing
|
||||
# (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent
|
||||
# indexing outside the loaded weights.
|
||||
effective_start = (
|
||||
max(self.shard_start, start_layer)
|
||||
if start_layer is not None
|
||||
else self.shard_start
|
||||
)
|
||||
effective_start = self._effective_start(start_layer)
|
||||
position_embeddings = _rotary_position_embeddings(
|
||||
self.model,
|
||||
hidden_states,
|
||||
@@ -346,6 +623,12 @@ class TorchModelShard:
|
||||
hidden_states,
|
||||
self.torch,
|
||||
)
|
||||
cache_position = None
|
||||
if cache is not None:
|
||||
seq_len = int(hidden_states.shape[1])
|
||||
cache_position = self.torch.arange(
|
||||
past_len, past_len + seq_len, device=hidden_states.device
|
||||
)
|
||||
with self.torch.inference_mode():
|
||||
for layer in self.layers[effective_start:self.shard_end + 1]:
|
||||
hidden_states = _call_layer(
|
||||
@@ -354,6 +637,8 @@ class TorchModelShard:
|
||||
layer_attention_mask,
|
||||
position_ids,
|
||||
position_embeddings,
|
||||
cache=cache,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
return hidden_states.to(self.torch.bfloat16)
|
||||
|
||||
@@ -411,7 +696,7 @@ def _should_partial_materialize_shard(
|
||||
return False
|
||||
if total_layers_hint is None:
|
||||
return False
|
||||
return not (shard_start == 0 and shard_end >= total_layers_hint - 1)
|
||||
return True
|
||||
|
||||
|
||||
def _load_partial_model_from_snapshot(
|
||||
@@ -476,17 +761,41 @@ def _load_partial_model_from_snapshot(
|
||||
)
|
||||
|
||||
with init_empty_weights_fn():
|
||||
model = auto_model_for_causal_lm.from_config(cfg, torch_dtype=dtype)
|
||||
model = auto_model_for_causal_lm.from_config(_causal_lm_config(cfg), torch_dtype=dtype)
|
||||
tie_weights = getattr(model, "tie_weights", None)
|
||||
if callable(tie_weights):
|
||||
tie_weights()
|
||||
|
||||
# Multimodal/MTP checkpoints (e.g. Qwen3.5/3.6-MoE) carry vision and
|
||||
# multi-token-prediction tensors the text-only CausalLM never builds;
|
||||
# transformers' from_pretrained drops them via _keys_to_ignore_on_load_unexpected,
|
||||
# so the manual loader must skip them too.
|
||||
expected_keys = _model_state_dict_keys(model)
|
||||
tensors_by_file: dict[str, list[str]] = {}
|
||||
skipped: list[str] = []
|
||||
for tensor_name in sorted(tensor_names):
|
||||
rel_file = weight_map.get(tensor_name)
|
||||
if not isinstance(rel_file, str):
|
||||
continue
|
||||
if (
|
||||
expected_keys is not None
|
||||
and _checkpoint_tensor_name_for_model(model, tensor_name) not in expected_keys
|
||||
):
|
||||
skipped.append(tensor_name)
|
||||
continue
|
||||
tensors_by_file.setdefault(rel_file, []).append(tensor_name)
|
||||
if skipped:
|
||||
preview = ", ".join(skipped[:3])
|
||||
print(
|
||||
f" Skipping {len(skipped)} checkpoint tensors absent from the causal LM "
|
||||
f"(e.g. {preview})",
|
||||
flush=True,
|
||||
)
|
||||
if not tensors_by_file:
|
||||
raise PartialModelLoadUnsupported(
|
||||
f"no checkpoint tensors for layers {shard_start}-{shard_end} match the "
|
||||
f"causal LM built from {snapshot_dir}"
|
||||
)
|
||||
|
||||
for rel_file, names in tensors_by_file.items():
|
||||
checkpoint_file = snapshot_dir / rel_file
|
||||
@@ -498,7 +807,7 @@ def _load_partial_model_from_snapshot(
|
||||
for tensor_name in names:
|
||||
set_tensor_fn(
|
||||
model,
|
||||
tensor_name,
|
||||
_checkpoint_tensor_name_for_model(model, tensor_name),
|
||||
device,
|
||||
value=handle.get_tensor(tensor_name),
|
||||
dtype=dtype,
|
||||
@@ -569,38 +878,85 @@ def _native_torch_dtype(cfg: Any, torch: Any) -> Any:
|
||||
return torch.bfloat16
|
||||
|
||||
|
||||
def _causal_lm_config(cfg: Any) -> Any:
|
||||
"""Use the text decoder config for composite VLM/MoE presets."""
|
||||
get_text_config = getattr(cfg, "get_text_config", None)
|
||||
if callable(get_text_config):
|
||||
try:
|
||||
return get_text_config()
|
||||
except Exception:
|
||||
pass
|
||||
text_config = getattr(cfg, "text_config", None)
|
||||
if text_config is not None:
|
||||
return text_config
|
||||
return cfg
|
||||
|
||||
|
||||
def _model_state_dict_keys(model: Any) -> set[str] | None:
|
||||
"""Expected parameter/buffer names, or None when the model can't report them."""
|
||||
state_dict = getattr(model, "state_dict", None)
|
||||
if not callable(state_dict):
|
||||
return None
|
||||
try:
|
||||
return set(state_dict().keys())
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _checkpoint_tensor_name_for_model(model: Any, tensor_name: str) -> str:
|
||||
"""Map multimodal checkpoint keys onto text-only CausalLM modules when needed."""
|
||||
inner = getattr(model, "model", None)
|
||||
if inner is not None and hasattr(inner, "language_model"):
|
||||
return tensor_name
|
||||
if ".language_model." in tensor_name:
|
||||
return tensor_name.replace(".language_model.", ".")
|
||||
return tensor_name
|
||||
|
||||
|
||||
def _transformer_backbone(model: Any) -> Any:
|
||||
if hasattr(model, "model"):
|
||||
inner = model.model
|
||||
language_model = getattr(inner, "language_model", None)
|
||||
if language_model is not None:
|
||||
return language_model
|
||||
return inner
|
||||
if hasattr(model, "transformer"):
|
||||
return model.transformer
|
||||
raise ModelBackendError(
|
||||
"unsupported HuggingFace model architecture: no transformer backbone found"
|
||||
)
|
||||
|
||||
|
||||
def _model_layers(model: Any) -> Any:
|
||||
if hasattr(model, "model") and hasattr(model.model, "layers"):
|
||||
return model.model.layers
|
||||
if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
|
||||
return model.transformer.h
|
||||
backbone = _transformer_backbone(model)
|
||||
for attr in ("layers", "h", "blocks"):
|
||||
layers = getattr(backbone, attr, None)
|
||||
if layers is not None:
|
||||
return layers
|
||||
raise ModelBackendError(
|
||||
"unsupported HuggingFace model architecture: no transformer layers found"
|
||||
)
|
||||
|
||||
|
||||
def _embed_tokens(model: Any) -> Any:
|
||||
if hasattr(model, "model") and hasattr(model.model, "embed_tokens"):
|
||||
return model.model.embed_tokens
|
||||
if hasattr(model, "transformer") and hasattr(model.transformer, "wte"):
|
||||
return model.transformer.wte
|
||||
backbone = _transformer_backbone(model)
|
||||
for attr in ("embed_tokens", "wte"):
|
||||
embed = getattr(backbone, attr, None)
|
||||
if embed is not None:
|
||||
return embed
|
||||
raise ModelBackendError(
|
||||
"unsupported HuggingFace model architecture: no token embeddings found"
|
||||
)
|
||||
|
||||
|
||||
def _position_embeddings(model: Any) -> Any | None:
|
||||
if hasattr(model, "transformer") and hasattr(model.transformer, "wpe"):
|
||||
return model.transformer.wpe
|
||||
return None
|
||||
backbone = _transformer_backbone(model)
|
||||
return getattr(backbone, "wpe", None)
|
||||
|
||||
|
||||
def _rotary_embedding_module(model: Any) -> Any | None:
|
||||
if hasattr(model, "model") and hasattr(model.model, "rotary_emb"):
|
||||
return model.model.rotary_emb
|
||||
if hasattr(model, "transformer") and hasattr(model.transformer, "rotary_emb"):
|
||||
return model.transformer.rotary_emb
|
||||
return None
|
||||
backbone = _transformer_backbone(model)
|
||||
return getattr(backbone, "rotary_emb", None)
|
||||
|
||||
|
||||
def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> list[Any]:
|
||||
@@ -627,10 +983,11 @@ def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> l
|
||||
|
||||
|
||||
def _final_norm(model: Any) -> Any | None:
|
||||
if hasattr(model, "model") and hasattr(model.model, "norm"):
|
||||
return model.model.norm
|
||||
if hasattr(model, "transformer") and hasattr(model.transformer, "ln_f"):
|
||||
return model.transformer.ln_f
|
||||
backbone = _transformer_backbone(model)
|
||||
for attr in ("norm", "ln_f", "final_layer_norm"):
|
||||
norm = getattr(backbone, attr, None)
|
||||
if norm is not None:
|
||||
return norm
|
||||
return None
|
||||
|
||||
|
||||
@@ -681,6 +1038,8 @@ def _call_layer(
|
||||
attention_mask: Any,
|
||||
position_ids: Any,
|
||||
position_embeddings: Any | None = None,
|
||||
cache: Any = None,
|
||||
cache_position: Any = None,
|
||||
) -> Any:
|
||||
attempts = (
|
||||
{
|
||||
@@ -701,6 +1060,14 @@ def _call_layer(
|
||||
last_exc: Exception | None = None
|
||||
for kwargs in attempts:
|
||||
filtered = {key: value for key, value in kwargs.items() if value is not None}
|
||||
if cache is not None:
|
||||
# transformers 5.x layers take a Cache via past_key_values and
|
||||
# mutate it in place; cache_position is required by sliding-window
|
||||
# and hybrid recurrent layers.
|
||||
filtered["past_key_values"] = cache
|
||||
filtered["use_cache"] = True
|
||||
if cache_position is not None:
|
||||
filtered["cache_position"] = cache_position
|
||||
try:
|
||||
output = layer(hidden_states, **filtered)
|
||||
return output[0] if isinstance(output, tuple) else output
|
||||
@@ -743,7 +1110,12 @@ def _looks_like_oom(exc: BaseException) -> bool:
|
||||
current: BaseException | None = exc
|
||||
while current is not None:
|
||||
text = str(current).lower()
|
||||
if "out of memory" in text or "cuda error: out of memory" in text:
|
||||
if (
|
||||
"out of memory" in text
|
||||
or "cuda error: out of memory" in text
|
||||
or "paging file is too small" in text
|
||||
or "os error 1455" in text
|
||||
):
|
||||
return True
|
||||
current = current.__cause__ or current.__context__
|
||||
return False
|
||||
|
||||
@@ -164,6 +164,9 @@ class RelayHttpBridge:
|
||||
path = str(payload.get("path") or "/")
|
||||
headers = payload.get("headers") if isinstance(payload.get("headers"), dict) else {}
|
||||
|
||||
req_suffix = f" request_id={request_id}" if request_id else ""
|
||||
print(f" [node] relay {method} {path}{req_suffix}", flush=True)
|
||||
|
||||
# body_base64 carries binary data (e.g. bfloat16 activation tensors) safely.
|
||||
# Fallback to text "body" for backward-compat with non-binary requests.
|
||||
body_b64 = payload.get("body_base64")
|
||||
|
||||
@@ -197,12 +197,58 @@ def _max_assignable_layers(
|
||||
memory_mb: int,
|
||||
total_layers: int | None,
|
||||
bytes_per_layer: int | None = None,
|
||||
*,
|
||||
safety_fraction: float = 0.8,
|
||||
) -> int:
|
||||
if total_layers is None or total_layers <= 0 or memory_mb <= 0:
|
||||
return 0
|
||||
budget_bytes = memory_mb * 1024 * 1024
|
||||
layer_bytes = bytes_per_layer or _DEFAULT_BYTES_PER_LAYER
|
||||
return min(total_layers, int((budget_bytes * 0.8) // layer_bytes))
|
||||
return min(total_layers, int((budget_bytes * safety_fraction) // layer_bytes))
|
||||
|
||||
|
||||
def _runtime_shard_safety_fraction(device: str) -> float:
|
||||
"""CPU partial loads need room for model skeletons, tokenizer, and allocator peaks."""
|
||||
return 0.55 if device != "cuda" else 0.8
|
||||
|
||||
|
||||
def _cap_auto_assigned_shard(
|
||||
shard_start: int,
|
||||
shard_end: int,
|
||||
total_layers: int | None,
|
||||
memory_mb: int,
|
||||
bytes_per_layer: int | None,
|
||||
device: str,
|
||||
) -> tuple[int, bool, int]:
|
||||
if bytes_per_layer is None or total_layers is None:
|
||||
return shard_end, False, 0
|
||||
max_layers = _max_assignable_layers(
|
||||
memory_mb,
|
||||
total_layers,
|
||||
bytes_per_layer=bytes_per_layer,
|
||||
safety_fraction=_runtime_shard_safety_fraction(device),
|
||||
)
|
||||
if max_layers <= 0:
|
||||
return shard_end, False, max_layers
|
||||
assigned_layers = shard_end - shard_start + 1
|
||||
if assigned_layers <= max_layers:
|
||||
return shard_end, False, max_layers
|
||||
return min(total_layers - 1, shard_start + max_layers - 1), True, max_layers
|
||||
|
||||
|
||||
def _format_shard_label(
|
||||
shard_start: int,
|
||||
shard_end: int,
|
||||
total_layers: int | None = None,
|
||||
*,
|
||||
model_name: str | None = None,
|
||||
) -> str:
|
||||
layer_count = shard_end - shard_start + 1
|
||||
if isinstance(total_layers, int) and total_layers > 0:
|
||||
return f"layers {shard_start}–{shard_end} ({layer_count} of {total_layers})"
|
||||
if model_name:
|
||||
return f"layers {shard_start}–{shard_end} ({model_name})"
|
||||
return f"layers {shard_start}–{shard_end}"
|
||||
|
||||
|
||||
def _shard_budget_line(
|
||||
@@ -211,13 +257,19 @@ def _shard_budget_line(
|
||||
total_layers: int | None,
|
||||
quantization: str,
|
||||
bytes_per_layer: int | None = None,
|
||||
safety_fraction: float = 0.8,
|
||||
) -> str:
|
||||
memory_gb = memory_mb / 1024
|
||||
gb_str = f"{memory_gb:.1f} GB"
|
||||
budget_quantization = "bfloat16" if quantization == "auto" else quantization
|
||||
if total_layers is None or total_layers <= 0:
|
||||
return f"Memory budget: {gb_str} {memory_source}; shard budget: unknown model layer count"
|
||||
max_layers = _max_assignable_layers(memory_mb, total_layers, bytes_per_layer=bytes_per_layer)
|
||||
max_layers = _max_assignable_layers(
|
||||
memory_mb,
|
||||
total_layers,
|
||||
bytes_per_layer=bytes_per_layer,
|
||||
safety_fraction=safety_fraction,
|
||||
)
|
||||
# Remaining capacity after one full model load (rough estimate)
|
||||
shard_bytes = max_layers * (bytes_per_layer or _DEFAULT_BYTES_PER_LAYER)
|
||||
remaining_gb = (memory_mb * 1024 * 1024 - shard_bytes) / (1024 ** 3)
|
||||
@@ -334,25 +386,38 @@ def _attach_relay_bridge(node: StubNodeServer | TorchNodeServer, bridge: RelayHt
|
||||
_PENDING_NODE_ID = "pending"
|
||||
|
||||
|
||||
_HEARTBEAT_INTERVAL_IDLE = 20.0
|
||||
_HEARTBEAT_INTERVAL_BUSY = 3.0
|
||||
|
||||
|
||||
def _start_heartbeat(
|
||||
tracker_url: str,
|
||||
node_id: str,
|
||||
register_payload: dict,
|
||||
interval: float = 20.0,
|
||||
interval: float = _HEARTBEAT_INTERVAL_IDLE,
|
||||
node_ref: Any | None = None,
|
||||
start_time: float | None = None,
|
||||
) -> threading.Thread:
|
||||
"""Daemon thread: sends heartbeats and re-registers automatically after tracker restarts.
|
||||
|
||||
Heartbeat body carries cumulative stats (total_requests, failed_requests,
|
||||
queue_depth, uptime_seconds, status). Stats are buffered locally during
|
||||
outage and flushed on next successful heartbeat.
|
||||
queue_depth, current_requests, uptime_seconds, status). Stats are buffered
|
||||
locally during outage and flushed on next successful heartbeat.
|
||||
|
||||
Heartbeat response may include new_assignment: {model, shard_start, shard_end}
|
||||
which is logged for now (hot-reload implemented in US-026).
|
||||
"""
|
||||
_start_time = start_time or time.monotonic()
|
||||
|
||||
def _current_requests_snapshot() -> list[dict]:
|
||||
if node_ref is None:
|
||||
return []
|
||||
getter = getattr(node_ref, "current_requests", None)
|
||||
if getter is None:
|
||||
return []
|
||||
current = getter() if callable(getter) else getter
|
||||
return list(current) if isinstance(current, list) else []
|
||||
|
||||
def _get_stats() -> dict:
|
||||
uptime = time.monotonic() - _start_time
|
||||
stats: dict = {"uptime_seconds": round(uptime, 1), "status": "ready"}
|
||||
@@ -364,8 +429,16 @@ def _start_heartbeat(
|
||||
)
|
||||
stats["failed_requests"] = getattr(node_ref, "failed_requests", 0)
|
||||
stats["queue_depth"] = getattr(node_ref, "queue_depth", 0)
|
||||
current_requests = _current_requests_snapshot()
|
||||
if current_requests:
|
||||
stats["current_requests"] = current_requests
|
||||
return stats
|
||||
|
||||
def _sleep_interval() -> float:
|
||||
if _current_requests_snapshot() or (node_ref is not None and getattr(node_ref, "queue_depth", 0) > 0):
|
||||
return _HEARTBEAT_INTERVAL_BUSY
|
||||
return interval
|
||||
|
||||
def _reregister() -> bool:
|
||||
nonlocal node_id
|
||||
try:
|
||||
@@ -427,7 +500,7 @@ def _start_heartbeat(
|
||||
outage_streak = 1 if node_id == _PENDING_NODE_ID else 0
|
||||
|
||||
while True:
|
||||
time.sleep(interval)
|
||||
time.sleep(_sleep_interval())
|
||||
|
||||
if outage_streak > 0:
|
||||
# Tracker was down — attempt re-registration first (it may have restarted
|
||||
@@ -537,6 +610,15 @@ def _warn_virtual_network_ip(ip: str | None) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _registration_display_fields(node_name: str | None) -> dict[str, str]:
|
||||
if not node_name:
|
||||
return {}
|
||||
name = node_name.strip()
|
||||
if not name:
|
||||
return {}
|
||||
return {"friendly_name": name}
|
||||
|
||||
|
||||
def run_startup(
|
||||
tracker_url: str,
|
||||
port: int = 0,
|
||||
@@ -557,6 +639,7 @@ def run_startup(
|
||||
tracker_source_disabled: bool = False,
|
||||
torch_threads: int | None = None,
|
||||
torch_interop_threads: int | None = None,
|
||||
node_name: str | None = None,
|
||||
) -> StubNodeServer | TorchNodeServer:
|
||||
"""Execute the full startup sequence and return a running node server.
|
||||
|
||||
@@ -573,6 +656,7 @@ def run_startup(
|
||||
|
||||
tracker_url = tracker_url.rstrip("/")
|
||||
relay_url = _discover_relay_url(tracker_url)
|
||||
display_fields = _registration_display_fields(node_name)
|
||||
if max_loaded_shards < 1:
|
||||
raise ValueError("--max-shards must be at least 1")
|
||||
|
||||
@@ -693,6 +777,25 @@ def run_startup(
|
||||
if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"):
|
||||
shard_start = net_asgn["shard_start"]
|
||||
shard_end = net_asgn["shard_end"]
|
||||
asgn_total_layers = int(net_asgn.get("num_layers") or detected)
|
||||
asgn_bytes_per_layer = _assignment_bytes_per_layer(net_asgn, quantization)
|
||||
capped_shard_end, was_capped, max_runtime_layers = _cap_auto_assigned_shard(
|
||||
shard_start,
|
||||
shard_end,
|
||||
asgn_total_layers,
|
||||
memory_budget_mb,
|
||||
asgn_bytes_per_layer,
|
||||
device,
|
||||
)
|
||||
if was_capped:
|
||||
original_end = shard_end
|
||||
shard_end = capped_shard_end
|
||||
print(
|
||||
f" WARNING: tracker assigned layers {shard_start}-{original_end}, "
|
||||
f"but CPU-safe runtime budget fits {max_runtime_layers}/{asgn_total_layers} layers; "
|
||||
f"loading layers {shard_start}-{shard_end} instead.",
|
||||
flush=True,
|
||||
)
|
||||
full_sources = (
|
||||
[] if tracker_source_disabled
|
||||
else _full_model_sources(net_asgn.get("model_sources", []))
|
||||
@@ -708,7 +811,7 @@ def run_startup(
|
||||
)
|
||||
print(
|
||||
f" Tracker found uncovered shard: "
|
||||
f"layers {shard_start}–{shard_end} (of {detected})",
|
||||
f"layers {shard_start}-{shard_end} (of {detected})",
|
||||
flush=True,
|
||||
)
|
||||
except Exception:
|
||||
@@ -734,13 +837,11 @@ def run_startup(
|
||||
_node_start_time = time.monotonic()
|
||||
actual_port = node.start()
|
||||
total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
|
||||
if isinstance(total_layers, int) and total_layers > 0:
|
||||
layer_count = shard_end - shard_start + 1
|
||||
shard_label = f"layers {shard_start}–{shard_end}; {layer_count} of {total_layers}"
|
||||
else:
|
||||
shard_label = f"layers {shard_start}–{shard_end}"
|
||||
shard_label = _format_shard_label(shard_start, shard_end, total_layers)
|
||||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||
endpoint = f"http://{public_host}:{actual_port}"
|
||||
if hasattr(node, "set_advertised_endpoint"):
|
||||
node.set_advertised_endpoint(endpoint)
|
||||
local_base_url = f"http://127.0.0.1:{actual_port}"
|
||||
relay_bridge, relay_fields = _start_relay_bridge_if_available(
|
||||
tracker_url,
|
||||
@@ -781,6 +882,7 @@ def run_startup(
|
||||
),
|
||||
**registration_capabilities,
|
||||
**relay_fields,
|
||||
**display_fields,
|
||||
}
|
||||
tracker_node_id = _register_with_tracker(
|
||||
tracker_url, reg_payload, node, _node_start_time,
|
||||
@@ -827,18 +929,36 @@ def run_startup(
|
||||
assigned_shard_start: int = net_assignment["shard_start"]
|
||||
assigned_shard_end: int = net_assignment["shard_end"]
|
||||
assigned_num_layers: int = net_assignment["num_layers"]
|
||||
assigned_bytes_per_layer = _assignment_bytes_per_layer(net_assignment, quantization)
|
||||
capped_shard_end, was_capped, max_runtime_layers = _cap_auto_assigned_shard(
|
||||
assigned_shard_start,
|
||||
assigned_shard_end,
|
||||
assigned_num_layers,
|
||||
memory_budget_mb,
|
||||
assigned_bytes_per_layer,
|
||||
device,
|
||||
)
|
||||
if was_capped:
|
||||
original_end = assigned_shard_end
|
||||
assigned_shard_end = capped_shard_end
|
||||
print(
|
||||
f" WARNING: tracker assigned layers {assigned_shard_start}-{original_end}, "
|
||||
f"but CPU-safe runtime budget fits {max_runtime_layers}/{assigned_num_layers} layers; "
|
||||
f"loading layers {assigned_shard_start}-{assigned_shard_end} instead.",
|
||||
flush=True,
|
||||
)
|
||||
assigned_model_sources: list[dict] = net_assignment.get("model_sources", [])
|
||||
if _gap_found:
|
||||
print(
|
||||
f" Assigned gap: {assigned_hf_repo} "
|
||||
f"layers {assigned_shard_start}–{assigned_shard_end} "
|
||||
f"layers {assigned_shard_start}-{assigned_shard_end} "
|
||||
f"(of {assigned_num_layers})",
|
||||
flush=True,
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f" Assigned redundant copy: {assigned_hf_repo} "
|
||||
f"layers {assigned_shard_start}–{assigned_shard_end} "
|
||||
f"layers {assigned_shard_start}-{assigned_shard_end} "
|
||||
f"(of {assigned_num_layers})",
|
||||
flush=True,
|
||||
)
|
||||
@@ -871,6 +991,8 @@ def run_startup(
|
||||
actual_port = node.start()
|
||||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||
endpoint = f"http://{public_host}:{actual_port}"
|
||||
if hasattr(node, "set_advertised_endpoint"):
|
||||
node.set_advertised_endpoint(endpoint)
|
||||
local_base_url = f"http://127.0.0.1:{actual_port}"
|
||||
relay_bridge, relay_fields = _start_relay_bridge_if_available(
|
||||
tracker_url,
|
||||
@@ -909,19 +1031,23 @@ def run_startup(
|
||||
),
|
||||
**registration_capabilities,
|
||||
**relay_fields,
|
||||
**display_fields,
|
||||
}
|
||||
tracker_node_id = _register_with_tracker(
|
||||
tracker_url, auto_reg_payload, node, _node_start_time,
|
||||
)
|
||||
shard_count = assigned_shard_end - assigned_shard_start + 1
|
||||
shard_label = _format_shard_label(
|
||||
assigned_shard_start,
|
||||
assigned_shard_end,
|
||||
assigned_num_layers,
|
||||
)
|
||||
print(
|
||||
f"\n{'=' * 32}\n"
|
||||
f"meshnet-node ready (auto-joined)\n"
|
||||
f" Wallet: {address}\n"
|
||||
f" Model ID: {assigned_hf_repo}\n"
|
||||
f" Shard: layers {assigned_shard_start}–{assigned_shard_end} "
|
||||
f"({shard_count} of {assigned_num_layers})\n"
|
||||
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization)}\n"
|
||||
f" Shard: {shard_label}\n"
|
||||
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization, bytes_per_layer=assigned_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\n"
|
||||
f" Quantization: {quantization}\n"
|
||||
f" Endpoint: {endpoint}\n"
|
||||
f" Node ID: {tracker_node_id or 'unregistered'}\n"
|
||||
@@ -967,13 +1093,36 @@ def run_startup(
|
||||
peers: list[dict] = assignment.get("peers", [])
|
||||
model_sources: list[dict] = [] if tracker_source_disabled else assignment.get("model_sources", [])
|
||||
assignment_bytes_per_layer = _assignment_bytes_per_layer(assignment, quantization)
|
||||
model_layers_end = assignment.get("model_layers_end")
|
||||
assigned_total_layers = (
|
||||
int(model_layers_end) + 1
|
||||
if model_layers_end is not None
|
||||
else None
|
||||
)
|
||||
shard_label = _format_shard_label(
|
||||
shard_start,
|
||||
shard_end,
|
||||
assigned_total_layers,
|
||||
model_name=assigned_model,
|
||||
)
|
||||
if user_pinned_shard:
|
||||
print(
|
||||
f" Shard: layers {shard_start}-{shard_end} of {assigned_model} (pinned)",
|
||||
flush=True,
|
||||
shard_label = f"{shard_label} (pinned)"
|
||||
if user_pinned_shard and assigned_total_layers and assignment_bytes_per_layer:
|
||||
pinned_layers = shard_end - shard_start + 1
|
||||
max_layers = _max_assignable_layers(
|
||||
memory_budget_mb,
|
||||
assigned_total_layers,
|
||||
assignment_bytes_per_layer,
|
||||
safety_fraction=_runtime_shard_safety_fraction(device),
|
||||
)
|
||||
else:
|
||||
print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)
|
||||
if pinned_layers > max_layers:
|
||||
raise ValueError(
|
||||
f"Pinned shard layers {shard_start}–{shard_end} ({pinned_layers} layers) exceed "
|
||||
f"the {memory_budget_mb / 1024:.1f} GB {memory_budget_source} budget "
|
||||
f"(fits up to {max_layers}/{assigned_total_layers} layers at bfloat16). "
|
||||
"Drop --shard-start/--shard-end to let the tracker auto-assign, or pin a smaller range."
|
||||
)
|
||||
print(f" Shard: {shard_label}", flush=True)
|
||||
|
||||
# 4. Download shard
|
||||
print("Downloading shard...", flush=True)
|
||||
@@ -998,7 +1147,80 @@ def run_startup(
|
||||
)
|
||||
print(f" Cached at: {shard_path}", flush=True)
|
||||
|
||||
# 5. Start HTTP server
|
||||
# 5. Start HTTP server — real HF weights use TorchNodeServer; stub-model stays stub.
|
||||
_node_start_time = time.monotonic()
|
||||
if hf_repo and assigned_model != "stub-model":
|
||||
print("Loading real PyTorch model shard...", flush=True)
|
||||
node = TorchNodeServer(
|
||||
host=host,
|
||||
port=port,
|
||||
model_id=hf_repo,
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
quantization=quantization,
|
||||
tracker_url=tracker_url,
|
||||
route_timeout=route_timeout,
|
||||
cache_dir=shard_path,
|
||||
debug=debug,
|
||||
max_loaded_shards=max_loaded_shards,
|
||||
)
|
||||
actual_port = node.start()
|
||||
total_layers = getattr(getattr(node, "backend", None), "total_layers", None) or assigned_total_layers
|
||||
shard_label = _format_shard_label(shard_start, shard_end, total_layers, model_name=assigned_model)
|
||||
if user_pinned_shard:
|
||||
shard_label = f"{shard_label} (pinned)"
|
||||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||
endpoint = f"http://{public_host}:{actual_port}"
|
||||
if hasattr(node, "set_advertised_endpoint"):
|
||||
node.set_advertised_endpoint(endpoint)
|
||||
local_base_url = f"http://127.0.0.1:{actual_port}"
|
||||
relay_bridge, relay_fields = _start_relay_bridge_if_available(
|
||||
tracker_url,
|
||||
address,
|
||||
local_base_url,
|
||||
endpoint,
|
||||
relay_url=relay_url,
|
||||
)
|
||||
_attach_relay_bridge(node, relay_bridge)
|
||||
reg_payload = {
|
||||
"endpoint": endpoint,
|
||||
"model": assigned_model,
|
||||
"hf_repo": hf_repo,
|
||||
"num_layers": total_layers,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"downloaded_models": downloaded_models,
|
||||
"hardware_profile": hw,
|
||||
"wallet_address": address,
|
||||
"quantization": quantization,
|
||||
"score": 1.0,
|
||||
"tracker_mode": (shard_start == 0),
|
||||
"managed_assignment": not user_pinned_shard,
|
||||
"model_metadata": model_metadata_for(hf_repo, total_layers, cache_dir=shard_path),
|
||||
**registration_capabilities,
|
||||
**relay_fields,
|
||||
**display_fields,
|
||||
}
|
||||
tracker_node_id = _register_with_tracker(
|
||||
tracker_url, reg_payload, node, _node_start_time,
|
||||
)
|
||||
print(
|
||||
f"\n{'=' * 32}\n"
|
||||
f"meshnet-node ready\n"
|
||||
f" Wallet: {address}\n"
|
||||
f" Model ID: {hf_repo}\n"
|
||||
f" Shard: {shard_label}\n"
|
||||
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\n"
|
||||
f" Quantization: {quantization}\n"
|
||||
f" Endpoint: {endpoint}\n"
|
||||
f" Node ID: {tracker_node_id or 'unregistered'}\n"
|
||||
f" Hardware: {_hardware_label(device, gpu_name)}\n"
|
||||
f" Benchmark: {bench_tps:,.0f} (throughput index)\n"
|
||||
f"{'=' * 32}",
|
||||
flush=True,
|
||||
)
|
||||
return node
|
||||
|
||||
is_last = shard_end >= assignment.get("model_layers_end", shard_end)
|
||||
node = StubNodeServer(
|
||||
host=host,
|
||||
@@ -1009,7 +1231,6 @@ def run_startup(
|
||||
model=assigned_model,
|
||||
shard_path=shard_path,
|
||||
)
|
||||
_node_start_time = time.monotonic()
|
||||
actual_port = node.start()
|
||||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||
endpoint = f"http://{public_host}:{actual_port}"
|
||||
@@ -1038,6 +1259,7 @@ def run_startup(
|
||||
"managed_assignment": not user_pinned_shard,
|
||||
**registration_capabilities,
|
||||
**relay_fields,
|
||||
**display_fields,
|
||||
}
|
||||
try:
|
||||
reg_resp = _post_json(
|
||||
@@ -1055,12 +1277,18 @@ def run_startup(
|
||||
hw_str = device.upper()
|
||||
if gpu_name:
|
||||
hw_str += f" ({gpu_name}, {vram_mb / 1024:.1f} GB)"
|
||||
shard_label = _format_shard_label(
|
||||
shard_start,
|
||||
shard_end,
|
||||
assigned_total_layers,
|
||||
model_name=assigned_model,
|
||||
)
|
||||
print(
|
||||
f"\n{'=' * 32}\n"
|
||||
f"meshnet-node ready\n"
|
||||
f" Wallet: {address}\n"
|
||||
f" Shard: layers {shard_start}-{shard_end} ({assigned_model})\n"
|
||||
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assignment.get('model_layers_end', shard_end) + 1, quantization, bytes_per_layer=assignment_bytes_per_layer)}\n"
|
||||
f" Shard: {shard_label}\n"
|
||||
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\n"
|
||||
f" Endpoint: {endpoint}\n"
|
||||
f" Node ID: {node_id}\n"
|
||||
f" Hardware: {hw_str}\n"
|
||||
|
||||
@@ -17,11 +17,17 @@ from typing import Any
|
||||
|
||||
from .model_backend import (
|
||||
InsufficientVRAMError,
|
||||
KVCacheMiss,
|
||||
MissingModelDependencyError,
|
||||
Quantization,
|
||||
TailTokenResult,
|
||||
TorchModelShard,
|
||||
validate_quantization,
|
||||
)
|
||||
|
||||
|
||||
class _PipelineCacheMiss(Exception):
|
||||
"""A downstream hop reported 409 cache_miss — head must re-prefill."""
|
||||
from .server import (
|
||||
_WIRE_VERSION,
|
||||
_compress_body,
|
||||
@@ -31,6 +37,69 @@ from .server import (
|
||||
)
|
||||
|
||||
|
||||
def _endpoint_key(url: str) -> str:
|
||||
"""Normalize http(s) endpoints for host:port comparison."""
|
||||
parsed = urllib.parse.urlparse(url.rstrip("/"))
|
||||
host = (parsed.hostname or "").lower()
|
||||
if not host:
|
||||
return url.rstrip("/").lower()
|
||||
port = parsed.port
|
||||
if port is None:
|
||||
port = 443 if parsed.scheme == "https" else 80
|
||||
return f"{host}:{port}"
|
||||
|
||||
|
||||
def _own_endpoint_key(server: _TorchHTTPServer) -> str:
|
||||
advertised = getattr(server, "advertised_endpoint", None)
|
||||
if advertised:
|
||||
return _endpoint_key(advertised)
|
||||
host, port = server.server_address
|
||||
return _endpoint_key(f"http://{host}:{port}")
|
||||
|
||||
|
||||
def _clamp_downstream_hops(
|
||||
hops: list[dict],
|
||||
backend: TorchModelShard | None,
|
||||
) -> list[dict]:
|
||||
"""Ensure downstream start_layer continues after this shard's layers."""
|
||||
if not hops or backend is None:
|
||||
return hops
|
||||
shard_end = getattr(backend, "shard_end", None)
|
||||
if shard_end is None:
|
||||
return hops
|
||||
min_start = int(shard_end) + 1
|
||||
clamped: list[dict] = []
|
||||
for hop in hops:
|
||||
adjusted = dict(hop)
|
||||
if int(adjusted.get("start_layer", 0)) < min_start:
|
||||
adjusted["start_layer"] = min_start
|
||||
clamped.append(adjusted)
|
||||
return clamped
|
||||
|
||||
|
||||
def _format_downstream_route(hops: list[dict]) -> str:
|
||||
return ", ".join(
|
||||
f"{h['endpoint']}@{h.get('start_layer', 0)}" for h in hops
|
||||
)
|
||||
|
||||
|
||||
def _write_progress_line(state: list[bool], message: str, *, final: bool = False) -> None:
|
||||
"""Rewrite one in-place progress line (\\r) or finish with a newline."""
|
||||
if final:
|
||||
if state[0]:
|
||||
sys.stdout.write("\r" + message + "\n")
|
||||
state[0] = False
|
||||
else:
|
||||
print(message, flush=True)
|
||||
return
|
||||
if state[0]:
|
||||
sys.stdout.write("\r" + message)
|
||||
else:
|
||||
sys.stdout.write(message)
|
||||
state[0] = True
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
def _relay_hop(
|
||||
relay_addr: str,
|
||||
path: str,
|
||||
@@ -65,6 +134,13 @@ def _relay_hop(
|
||||
return status, resp_headers, resp_body
|
||||
|
||||
|
||||
def _is_cache_miss_body(body: bytes) -> bool:
|
||||
try:
|
||||
return json.loads(body).get("error") == "cache_miss"
|
||||
except (json.JSONDecodeError, AttributeError, UnicodeDecodeError):
|
||||
return False
|
||||
|
||||
|
||||
class _TorchHTTPServer(http.server.HTTPServer):
|
||||
def __init__(
|
||||
self,
|
||||
@@ -87,10 +163,31 @@ class _TorchHTTPServer(http.server.HTTPServer):
|
||||
self.route_timeout = route_timeout
|
||||
self.debug = debug
|
||||
self.max_loaded_shards = max(1, max_loaded_shards)
|
||||
self.advertised_endpoint: str | None = None
|
||||
self.total_requests: int = 0
|
||||
self.failed_requests: int = 0
|
||||
self.queue_depth: int = 0
|
||||
self._stats_lock = threading.Lock()
|
||||
self._active_requests: dict[str, dict[str, Any]] = {}
|
||||
|
||||
def snapshot_current_requests(self) -> list[dict[str, Any]]:
|
||||
"""In-flight request snapshots for tracker heartbeats."""
|
||||
now = time.monotonic()
|
||||
with self._stats_lock:
|
||||
out: list[dict[str, Any]] = []
|
||||
for rec in self._active_requests.values():
|
||||
elapsed = max(now - float(rec["started"]), 1e-6)
|
||||
tokens = int(rec.get("tokens") or 0)
|
||||
out.append({
|
||||
"request_id": str(rec["request_id"]),
|
||||
"model": str(rec.get("model") or ""),
|
||||
"kind": str(rec.get("kind") or "chat"),
|
||||
"tokens": tokens,
|
||||
"elapsed_seconds": round(elapsed, 1),
|
||||
"tokens_per_sec": round(tokens / elapsed, 2) if tokens > 0 else 0.0,
|
||||
"routing_complete": bool(rec.get("routing_complete")),
|
||||
})
|
||||
return out
|
||||
|
||||
def resolve_backend(self, model_name: str | None) -> TorchModelShard | None:
|
||||
if not model_name:
|
||||
@@ -113,6 +210,53 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
|
||||
pass
|
||||
|
||||
def _request_id(self) -> str:
|
||||
return (
|
||||
self.headers.get("X-Meshnet-Request-Id")
|
||||
or self.headers.get("X-Request-Id")
|
||||
or f"local-{time.time_ns():x}"
|
||||
)
|
||||
|
||||
def _request_log_suffix(self) -> str:
|
||||
req_id = self.headers.get("X-Meshnet-Request-Id") or self.headers.get("X-Request-Id")
|
||||
return f" request_id={req_id}" if req_id else ""
|
||||
|
||||
def _track_request_begin(
|
||||
self,
|
||||
server: "_TorchHTTPServer",
|
||||
request_id: str,
|
||||
model: str,
|
||||
) -> None:
|
||||
with server._stats_lock:
|
||||
server._active_requests[request_id] = {
|
||||
"request_id": request_id,
|
||||
"model": model,
|
||||
"kind": "chat",
|
||||
"started": time.monotonic(),
|
||||
"tokens": 0,
|
||||
"routing_complete": False,
|
||||
}
|
||||
|
||||
def _track_request_progress(
|
||||
self,
|
||||
server: "_TorchHTTPServer",
|
||||
request_id: str,
|
||||
*,
|
||||
tokens: int,
|
||||
routing_complete: bool = False,
|
||||
) -> None:
|
||||
with server._stats_lock:
|
||||
rec = server._active_requests.get(request_id)
|
||||
if rec is None:
|
||||
return
|
||||
rec["tokens"] = tokens
|
||||
if routing_complete:
|
||||
rec["routing_complete"] = True
|
||||
|
||||
def _track_request_end(self, server: "_TorchHTTPServer", request_id: str) -> None:
|
||||
with server._stats_lock:
|
||||
server._active_requests.pop(request_id, None)
|
||||
|
||||
def do_POST(self):
|
||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||
if self.path == "/forward":
|
||||
@@ -199,12 +343,35 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
return
|
||||
|
||||
server.forward_chunk_count += 1
|
||||
if int(self.headers.get("X-Meshnet-Hop-Index", "0")) > 0:
|
||||
hop_index = int(self.headers.get("X-Meshnet-Hop-Index", "0"))
|
||||
if hop_index > 0:
|
||||
server.received_activations = True
|
||||
if chunk_index_value == 0:
|
||||
shard_start = getattr(server.backend, "shard_start", "?")
|
||||
shard_end = getattr(server.backend, "shard_end", "?")
|
||||
print(
|
||||
f" [node] forward hop={hop_index} "
|
||||
f"layers={shard_start}-{shard_end} "
|
||||
f"session={session[:8]}{self._request_log_suffix()}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
start_layer_header = self.headers.get("X-Meshnet-Start-Layer")
|
||||
start_layer = int(start_layer_header) if start_layer_header else None
|
||||
|
||||
# Session KV-cache protocol: prefill establishes per-session state on
|
||||
# this node's layer range; decode reuses it. Absent header = legacy
|
||||
# stateless call (also the signature fake backends implement).
|
||||
cache_mode = self.headers.get("X-Meshnet-Cache")
|
||||
forward_kwargs: dict[str, object] = {}
|
||||
if cache_mode in ("prefill", "decode"):
|
||||
past_len_header = self.headers.get("X-Meshnet-Past-Len")
|
||||
forward_kwargs = {
|
||||
"session_id": session,
|
||||
"cache_mode": cache_mode,
|
||||
"past_len": int(past_len_header) if past_len_header else None,
|
||||
}
|
||||
|
||||
try:
|
||||
result = server.backend.forward_bytes(
|
||||
raw_body,
|
||||
@@ -212,11 +379,18 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
self.headers.get("X-Meshnet-Attn-Mask"),
|
||||
self.headers.get("X-Meshnet-Position-Ids"),
|
||||
start_layer=start_layer,
|
||||
**forward_kwargs,
|
||||
)
|
||||
except KVCacheMiss as exc:
|
||||
self._send_json(409, {"error": "cache_miss", "detail": str(exc)})
|
||||
return
|
||||
except Exception as exc:
|
||||
self._send_json(500, {"error": str(exc)})
|
||||
return
|
||||
|
||||
if isinstance(result, TailTokenResult):
|
||||
self._send_json(200, {"text": result.text, "token_id": result.token_id})
|
||||
return
|
||||
if isinstance(result, str):
|
||||
self._send_json(200, {"text": result})
|
||||
return
|
||||
@@ -280,12 +454,14 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
def _handle_chat_completions(self) -> None:
|
||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||
request_id = self._request_id()
|
||||
with server._stats_lock:
|
||||
server.total_requests += 1
|
||||
server.queue_depth += 1
|
||||
try:
|
||||
self._do_chat_completions(server)
|
||||
self._do_chat_completions(server, request_id)
|
||||
finally:
|
||||
self._track_request_end(server, request_id)
|
||||
with server._stats_lock:
|
||||
server.queue_depth -= 1
|
||||
|
||||
@@ -294,7 +470,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
with server._stats_lock:
|
||||
server.failed_requests += 1
|
||||
|
||||
def _do_chat_completions(self, server: "_TorchHTTPServer") -> None:
|
||||
def _do_chat_completions(self, server: "_TorchHTTPServer", request_id: str) -> None:
|
||||
body = self._read_json_body()
|
||||
if body is None:
|
||||
return
|
||||
@@ -307,31 +483,74 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
if backend is None or not backend.is_head:
|
||||
self._send_json(400, {"error": "model not loaded on this node"})
|
||||
return
|
||||
max_tokens = int(body.get("max_tokens") or body.get("max_new_tokens") or 256)
|
||||
max_tokens = int(body.get("max_tokens") or body.get("max_new_tokens") or 5120)
|
||||
temperature = float(body.get("temperature") or 1.0)
|
||||
top_p = float(body.get("top_p") or 1.0)
|
||||
|
||||
self._track_request_begin(server, request_id, model_name)
|
||||
print(
|
||||
f" [node] processing chat model={model_name!r} stream={stream} "
|
||||
f"max_tokens={max_tokens}{self._request_log_suffix()}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# Fast path: this node owns the complete model — use HF generate() with KV cache.
|
||||
# Avoids the single-token-per-forward-pass limitation of the distributed path.
|
||||
if backend.is_head and backend.is_tail:
|
||||
gen_started = time.monotonic()
|
||||
progress_line = [False]
|
||||
try:
|
||||
if stream:
|
||||
self._stream_openai_response(
|
||||
backend.generate_text_streaming(messages, max_tokens, temperature, top_p),
|
||||
model_name,
|
||||
token_count = 0
|
||||
|
||||
def _counting_stream():
|
||||
nonlocal token_count
|
||||
for token_text in backend.generate_text_streaming(
|
||||
messages, max_tokens, temperature, top_p,
|
||||
):
|
||||
if token_text:
|
||||
token_count += 1
|
||||
self._track_request_progress(
|
||||
server, request_id, tokens=token_count, routing_complete=True,
|
||||
)
|
||||
yield token_text
|
||||
|
||||
self._stream_openai_response(_counting_stream(), model_name)
|
||||
elapsed = time.monotonic() - gen_started
|
||||
tps = token_count / max(elapsed, 1e-6)
|
||||
_write_progress_line(
|
||||
progress_line,
|
||||
f" [node] chat complete (stream) tokens={token_count} "
|
||||
f"elapsed_s={elapsed:.1f} tps={tps:.2f}{self._request_log_suffix()}",
|
||||
final=True,
|
||||
)
|
||||
else:
|
||||
text = backend.generate_text(messages, max_tokens, temperature, top_p)
|
||||
completion_tokens = _backend_token_count(
|
||||
backend, "count_text_tokens", text, fallback=len(text.split()) or 1,
|
||||
)
|
||||
print(
|
||||
f" [node] chat complete tokens={completion_tokens} "
|
||||
f"elapsed_s={time.monotonic() - gen_started:.1f}{self._request_log_suffix()}",
|
||||
flush=True,
|
||||
)
|
||||
self._send_openai_response(text, model_name, False, messages, backend=backend)
|
||||
except Exception as exc:
|
||||
self._record_failed_request()
|
||||
print(
|
||||
f" [node] chat failed after {time.monotonic() - gen_started:.1f}s: {exc}"
|
||||
f"{self._request_log_suffix()}",
|
||||
flush=True,
|
||||
)
|
||||
self._send_json(500, {"error": f"generation failed: {exc}"})
|
||||
return
|
||||
|
||||
# Distributed path: autoregressive generation across shards.
|
||||
# We do N single-step forward passes (no cross-node KV cache), which is slow
|
||||
# but correct. Each step: head encodes current sequence → forwards through route
|
||||
# → tail returns the next token string → append → repeat.
|
||||
# Distributed path: autoregressive generation across shards with a
|
||||
# sharded per-node KV cache. Step 0 prefills the full prompt through the
|
||||
# route (each node caches state for its own layer range, keyed by a
|
||||
# per-generation session id); steps 1+ send only the newest token's
|
||||
# hidden state. A 409 cache_miss from any hop (eviction/restart/route
|
||||
# change) falls back to a full re-prefill — the old stateless behavior.
|
||||
remaining_route = self._get_remaining_route(model_name, backend=backend)
|
||||
print(
|
||||
f" [node] chat route model={model_name!r} max_tokens={max_tokens} "
|
||||
@@ -364,28 +583,115 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
generated: list[str] = []
|
||||
current_text = prompt_text
|
||||
|
||||
session_id = str(uuid.uuid4())
|
||||
use_kv = bool(getattr(backend, "supports_kv_cache", False))
|
||||
eos_ids: set[int] = set()
|
||||
if use_kv:
|
||||
try:
|
||||
eos_ids = set(backend.eos_token_ids())
|
||||
except Exception:
|
||||
eos_ids = set()
|
||||
|
||||
stream_emit = None
|
||||
if stream:
|
||||
stream_emit = self._start_openai_stream(model_name)
|
||||
self._track_request_progress(server, request_id, tokens=0, routing_complete=True)
|
||||
|
||||
for _ in range(max_tokens):
|
||||
_GENERATION_LOG_INTERVAL = 5.0
|
||||
gen_started = time.monotonic()
|
||||
last_gen_log = gen_started
|
||||
progress_line = [False]
|
||||
last_token_id: int | None = None
|
||||
|
||||
def _prefill_step() -> tuple[str, int | None]:
|
||||
"""Full-sequence prefill: initial step and cache-miss recovery."""
|
||||
payload = (
|
||||
backend.encode_prompt(current_text, session_id=session_id)
|
||||
if use_kv
|
||||
else backend.encode_prompt(current_text)
|
||||
)
|
||||
return self._run_downstream_pipeline(
|
||||
payload, remaining_route, backend=backend,
|
||||
session=session_id, cache_mode="prefill" if use_kv else None,
|
||||
)
|
||||
|
||||
for step in range(max_tokens):
|
||||
try:
|
||||
payload = backend.encode_prompt(current_text)
|
||||
if use_kv and step > 0 and last_token_id is not None:
|
||||
try:
|
||||
payload = backend.encode_next_token(last_token_id, session_id)
|
||||
token_str, token_id = self._run_downstream_pipeline(
|
||||
payload, remaining_route, backend=backend,
|
||||
session=session_id, cache_mode="decode",
|
||||
)
|
||||
except (KVCacheMiss, _PipelineCacheMiss) as miss:
|
||||
# Evicted/restarted node or head lost its own session:
|
||||
# re-prefill the whole sequence once and continue cached.
|
||||
print(
|
||||
f" [node] kv cache miss at step {step} ({miss}); "
|
||||
f"re-prefilling {len(current_text)} chars",
|
||||
flush=True,
|
||||
)
|
||||
token_str, token_id = _prefill_step()
|
||||
else:
|
||||
token_str, token_id = _prefill_step()
|
||||
except _PipelineCacheMiss as exc:
|
||||
print(f" [node] unexpected cache miss on prefill: {exc}", flush=True)
|
||||
break
|
||||
except Exception as exc:
|
||||
print(f" [node] distributed encode error: {exc}", flush=True)
|
||||
break
|
||||
token_str = self._run_downstream_pipeline(payload, remaining_route, backend=backend)
|
||||
if not token_str:
|
||||
break
|
||||
# Stop on error responses or EOS.
|
||||
if token_str.startswith(("pipeline error", "decode error", "no downstream", "error:")):
|
||||
break
|
||||
if token_id is not None and token_id in eos_ids:
|
||||
break
|
||||
if eos_token and token_str == eos_token:
|
||||
break
|
||||
generated.append(token_str)
|
||||
if stream_emit is not None:
|
||||
stream_emit(token_str)
|
||||
current_text = current_text + token_str
|
||||
if not token_str and token_id is None:
|
||||
break
|
||||
last_token_id = token_id
|
||||
# token_str can be empty for a skipped special token that is not
|
||||
# EOS — keep generating from its token_id without emitting text.
|
||||
if token_str:
|
||||
generated.append(token_str)
|
||||
if stream_emit is not None:
|
||||
stream_emit(token_str)
|
||||
current_text = current_text + token_str
|
||||
self._track_request_progress(
|
||||
server,
|
||||
request_id,
|
||||
tokens=len(generated),
|
||||
routing_complete=True,
|
||||
)
|
||||
now = time.monotonic()
|
||||
if step == 0 or now - last_gen_log >= _GENERATION_LOG_INTERVAL:
|
||||
elapsed = now - gen_started
|
||||
token_count = len(generated)
|
||||
tps = token_count / max(elapsed, 1e-6)
|
||||
_write_progress_line(
|
||||
progress_line,
|
||||
f" [node] generating step={step + 1}/{max_tokens} "
|
||||
f"tokens={token_count} elapsed_s={elapsed:.1f} tps={tps:.2f}",
|
||||
)
|
||||
last_gen_log = now
|
||||
|
||||
if use_kv:
|
||||
try:
|
||||
backend.release_session(session_id)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if generated:
|
||||
elapsed = time.monotonic() - gen_started
|
||||
token_count = len(generated)
|
||||
tps = token_count / max(elapsed, 1e-6)
|
||||
_write_progress_line(
|
||||
progress_line,
|
||||
f" [node] generation complete tokens={token_count} "
|
||||
f"elapsed_s={elapsed:.1f} tps={tps:.2f}",
|
||||
final=True,
|
||||
)
|
||||
|
||||
result_text = "".join(generated)
|
||||
if stream_emit is not None:
|
||||
@@ -401,6 +707,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
start_layer tells each downstream node which layer to begin from,
|
||||
enabling correct execution when shard ranges overlap.
|
||||
"""
|
||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||
active_backend = backend or server.backend
|
||||
|
||||
# Fast path: tracker pre-resolved the downstream route and injected it as a header.
|
||||
injected = self.headers.get("X-Meshnet-Route")
|
||||
if injected:
|
||||
@@ -419,14 +728,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
hops.append(hop)
|
||||
elif isinstance(item, str):
|
||||
hops.append({"endpoint": item, "start_layer": 0})
|
||||
print(f" [node] using injected downstream route: {[h['endpoint'] for h in hops]}", flush=True)
|
||||
hops = _clamp_downstream_hops(hops, active_backend)
|
||||
print(
|
||||
f" [node] using injected downstream route: {_format_downstream_route(hops)}",
|
||||
flush=True,
|
||||
)
|
||||
return hops
|
||||
except (json.JSONDecodeError, TypeError, KeyError):
|
||||
pass
|
||||
|
||||
# Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject).
|
||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||
active_backend = backend or server.backend
|
||||
if server.tracker_url is None:
|
||||
return []
|
||||
route_model = getattr(active_backend, "model_id", None) or model
|
||||
@@ -434,29 +745,51 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}"
|
||||
with urllib.request.urlopen(url, timeout=server.route_timeout) as r:
|
||||
route_resp = json.loads(r.read())
|
||||
own_port = server.server_address[1]
|
||||
own_key = _own_endpoint_key(server)
|
||||
nodes_info = route_resp.get("nodes", [])
|
||||
hops = []
|
||||
covered_up_to: int | None = None
|
||||
hops: list[dict] = []
|
||||
passed_self = False
|
||||
for node_info in nodes_info:
|
||||
ep = node_info.get("endpoint", "")
|
||||
if ep.rstrip("/").endswith(f":{own_port}"):
|
||||
covered_up_to = node_info.get("shard_end")
|
||||
if not ep:
|
||||
continue
|
||||
if covered_up_to is None:
|
||||
covered_up_to = (node_info.get("shard_start") or 1) - 1
|
||||
hop = {"endpoint": ep, "start_layer": covered_up_to + 1}
|
||||
if _endpoint_key(ep) == own_key:
|
||||
passed_self = True
|
||||
continue
|
||||
if not passed_self:
|
||||
continue
|
||||
hop = {
|
||||
"endpoint": ep,
|
||||
"start_layer": int(node_info.get("start_layer", 0)),
|
||||
}
|
||||
if node_info.get("relay_addr"):
|
||||
hop["relay_addr"] = str(node_info["relay_addr"])
|
||||
hops.append(hop)
|
||||
covered_up_to = node_info.get("shard_end", covered_up_to)
|
||||
print(f" [node] tracker downstream route: {[h['endpoint'] for h in hops]}", flush=True)
|
||||
hops = _clamp_downstream_hops(hops, active_backend)
|
||||
print(
|
||||
f" [node] tracker downstream route: {_format_downstream_route(hops)}",
|
||||
flush=True,
|
||||
)
|
||||
return hops
|
||||
except Exception as exc:
|
||||
print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
|
||||
return []
|
||||
|
||||
def _run_downstream_pipeline(self, payload: object, route: list[dict], *, backend: TorchModelShard | None = None) -> str:
|
||||
def _run_downstream_pipeline(
|
||||
self,
|
||||
payload: object,
|
||||
route: list[dict],
|
||||
*,
|
||||
backend: TorchModelShard | None = None,
|
||||
session: str | None = None,
|
||||
cache_mode: str | None = None,
|
||||
) -> tuple[str, int | None]:
|
||||
"""Forward an activation through the downstream route.
|
||||
|
||||
Returns (token_text, token_id) — token_id is None when a hop predates
|
||||
the KV-cache protocol. Raises _PipelineCacheMiss when a hop responds
|
||||
409 cache_miss (evicted/restarted node) so the caller can re-prefill.
|
||||
"""
|
||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||
active_backend = backend or server.backend
|
||||
if not route:
|
||||
@@ -468,12 +801,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
bytearray(payload.body), # type: ignore[union-attr]
|
||||
dtype=active_backend.torch.bfloat16,
|
||||
).reshape(payload.shape).to(active_backend.device) # type: ignore[union-attr]
|
||||
return active_backend.decode_tail(tensor)
|
||||
if hasattr(active_backend, "decode_tail_token"):
|
||||
tail = active_backend.decode_tail_token(tensor)
|
||||
return tail.text, tail.token_id
|
||||
return active_backend.decode_tail(tensor), None
|
||||
except Exception as exc:
|
||||
return f"decode error: {exc}"
|
||||
return "no downstream route available for non-tail shard"
|
||||
return f"decode error: {exc}", None
|
||||
return "no downstream route available for non-tail shard", None
|
||||
|
||||
session = str(uuid.uuid4())
|
||||
# Session is stable across all steps of one generation when the caller
|
||||
# provides it (KV-cache protocol); fresh per call otherwise (legacy).
|
||||
session = session or str(uuid.uuid4())
|
||||
shape = payload.shape # type: ignore[union-attr]
|
||||
attn_mask = payload.attention_mask_header # type: ignore[union-attr]
|
||||
pos_ids = payload.position_ids_header # type: ignore[union-attr]
|
||||
@@ -503,6 +841,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
"X-Meshnet-Hop-Index": str(hop_index),
|
||||
"X-Meshnet-Start-Layer": str(start_layer),
|
||||
}
|
||||
if cache_mode:
|
||||
headers["X-Meshnet-Cache"] = cache_mode
|
||||
past_len = getattr(payload, "past_len", None)
|
||||
if cache_mode == "decode" and past_len is not None:
|
||||
headers["X-Meshnet-Past-Len"] = str(past_len)
|
||||
if current_attn:
|
||||
headers["X-Meshnet-Attn-Mask"] = current_attn
|
||||
if current_pos:
|
||||
@@ -512,12 +855,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
status, resp_headers, resp_body = _relay_hop(
|
||||
relay_addr, "/forward", current_body, headers, timeout=120.0,
|
||||
)
|
||||
if status == 409 and _is_cache_miss_body(resp_body):
|
||||
raise _PipelineCacheMiss(node_url)
|
||||
if status >= 400:
|
||||
print(
|
||||
f" [node] relay hop {hop_index} returned {status} from {relay_addr}",
|
||||
flush=True,
|
||||
)
|
||||
return f"pipeline error at {node_url} via relay: status {status}"
|
||||
return f"pipeline error at {node_url} via relay: status {status}", None
|
||||
except _PipelineCacheMiss:
|
||||
raise
|
||||
except Exception as exc:
|
||||
print(
|
||||
f" [node] relay hop {hop_index} failed at {relay_addr}: {exc}; "
|
||||
@@ -536,26 +883,33 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
with urllib.request.urlopen(req, timeout=120.0) as r:
|
||||
resp_body = r.read()
|
||||
resp_headers = {k.lower(): v for k, v in r.headers.items()}
|
||||
except urllib.error.HTTPError as exc:
|
||||
body = exc.read()
|
||||
if exc.code == 409 and _is_cache_miss_body(body):
|
||||
raise _PipelineCacheMiss(node_url) from exc
|
||||
print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True)
|
||||
return f"pipeline error at {node_url}: {exc}", None
|
||||
except Exception as exc:
|
||||
print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True)
|
||||
return f"pipeline error at {node_url}: {exc}"
|
||||
return f"pipeline error at {node_url}: {exc}", None
|
||||
content_type = resp_headers.get("content-type", "")
|
||||
if "application/json" in content_type:
|
||||
try:
|
||||
data = json.loads(resp_body)
|
||||
text = str(data.get("text", ""))
|
||||
token_id = data.get("token_id")
|
||||
if server.debug:
|
||||
print(f" [node] pipeline hop {hop_index} returned text={text!r}", flush=True)
|
||||
return text
|
||||
return text, int(token_id) if token_id is not None else None
|
||||
except json.JSONDecodeError:
|
||||
return resp_body.decode("utf-8", errors="replace")
|
||||
return resp_body.decode("utf-8", errors="replace"), None
|
||||
# Binary activation — update and forward to next node
|
||||
shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape))
|
||||
current_shape = _parse_shape(shape_header)
|
||||
current_body = resp_body
|
||||
current_attn = resp_headers.get("x-meshnet-attn-mask")
|
||||
current_pos = resp_headers.get("x-meshnet-position-ids")
|
||||
return ""
|
||||
return "", None
|
||||
|
||||
def _stream_openai_response(self, token_iter, model: str) -> None:
|
||||
"""Stream tokens from an iterator as SSE chunks."""
|
||||
@@ -783,6 +1137,12 @@ class TorchNodeServer:
|
||||
def queue_depth(self) -> int:
|
||||
return self._server.queue_depth if self._server is not None else 0
|
||||
|
||||
@property
|
||||
def current_requests(self) -> list[dict[str, Any]]:
|
||||
if self._server is None:
|
||||
return []
|
||||
return self._server.snapshot_current_requests()
|
||||
|
||||
@property
|
||||
def loaded_model_ids(self) -> list[str]:
|
||||
return list(self._backends.keys())
|
||||
@@ -862,6 +1222,11 @@ class TorchNodeServer:
|
||||
self._thread.start()
|
||||
return self.port
|
||||
|
||||
def set_advertised_endpoint(self, endpoint: str) -> None:
|
||||
"""Set the LAN-facing endpoint used for route self-detection."""
|
||||
if self._server is not None:
|
||||
self._server.advertised_endpoint = endpoint
|
||||
|
||||
def stop(self) -> None:
|
||||
if self._server is None:
|
||||
return
|
||||
|
||||
@@ -16,7 +16,8 @@ dependencies = [
|
||||
"rich>=13",
|
||||
"safetensors>=0.4",
|
||||
"torch>=2.1",
|
||||
"transformers>=4.39",
|
||||
"transformers>=5.12",
|
||||
"triton-windows>=3.7; platform_system == 'Windows'",
|
||||
"websockets>=13",
|
||||
"zstandard>=0.22",
|
||||
"kernels>=0.11.1,<0.16",
|
||||
|
||||
@@ -8,7 +8,8 @@ regular user.
|
||||
Mutations are append-only events with unique ids — the same replication
|
||||
model as ``BillingLedger`` — so accounts and API keys converge across the
|
||||
tracker hive via gossip, and every dashboard can serve registration/login.
|
||||
Sessions are deliberately local to each tracker (bearer tokens in memory).
|
||||
Sessions are local to each tracker and persisted so dashboard cookies survive
|
||||
tracker restarts.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -26,9 +27,24 @@ DEFAULT_ACCOUNTS_DB_PATH = "accounts.sqlite"
|
||||
SESSION_TTL = 7 * 86400.0 # seconds
|
||||
PBKDF2_ITERATIONS = 200_000
|
||||
MIN_PASSWORD_LENGTH = 8
|
||||
_MAX_NICKNAME_LENGTH = 64
|
||||
API_KEY_PREFIX = "sk-mesh-"
|
||||
|
||||
_EMAIL_RE = re.compile(r"^[^@\s]+@[^@\s]+\.[^@\s]+$")
|
||||
_UNSET = object()
|
||||
|
||||
|
||||
def _normalize_nickname(value: object) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
if not isinstance(value, str):
|
||||
raise ValueError("nickname must be a string")
|
||||
nickname = value.strip()
|
||||
if not nickname:
|
||||
return None
|
||||
if len(nickname) > _MAX_NICKNAME_LENGTH:
|
||||
raise ValueError(f"nickname must be at most {_MAX_NICKNAME_LENGTH} characters")
|
||||
return nickname
|
||||
|
||||
|
||||
def _hash_password(password: str, salt: str) -> str:
|
||||
@@ -63,6 +79,7 @@ class AccountStore:
|
||||
email: str | None = None,
|
||||
wallet: str | None = None,
|
||||
password: str,
|
||||
nickname: str | None = None,
|
||||
) -> dict:
|
||||
"""Create an account. The first account becomes the admin.
|
||||
|
||||
@@ -76,6 +93,7 @@ class AccountStore:
|
||||
raise ValueError("invalid email address")
|
||||
if len(password or "") < MIN_PASSWORD_LENGTH:
|
||||
raise ValueError(f"password must be at least {MIN_PASSWORD_LENGTH} characters")
|
||||
nickname = _normalize_nickname(nickname)
|
||||
with self._lock:
|
||||
for identifier in filter(None, (email, wallet)):
|
||||
if identifier.lower() in self._by_identifier:
|
||||
@@ -90,11 +108,30 @@ class AccountStore:
|
||||
"role": "admin" if not self._accounts else "user",
|
||||
"password_hash": _hash_password(password, salt),
|
||||
"salt": salt,
|
||||
"nickname": nickname,
|
||||
"ts": time.time(),
|
||||
}
|
||||
self._apply_locked(event)
|
||||
return self._public_view(self._accounts[event["account_id"]])
|
||||
|
||||
def update_profile(self, account_id: str, *, nickname: str | None = _UNSET) -> dict:
|
||||
"""Update display fields for an account. Pass nickname=None to clear."""
|
||||
if nickname is not _UNSET:
|
||||
nickname = _normalize_nickname(nickname)
|
||||
with self._lock:
|
||||
if account_id not in self._accounts:
|
||||
raise ValueError("unknown account")
|
||||
event = {
|
||||
"id": f"profile-{uuid.uuid4().hex}",
|
||||
"type": "update_profile",
|
||||
"account_id": account_id,
|
||||
"ts": time.time(),
|
||||
}
|
||||
if nickname is not _UNSET:
|
||||
event["nickname"] = nickname
|
||||
self._apply_locked(event)
|
||||
return self._public_view(self._accounts[account_id])
|
||||
|
||||
def verify_login(self, identifier: str, password: str) -> dict | None:
|
||||
"""Return the public account view when credentials match, else None."""
|
||||
with self._lock:
|
||||
@@ -115,6 +152,8 @@ class AccountStore:
|
||||
"account_id": account_id,
|
||||
"expires": time.time() + SESSION_TTL,
|
||||
}
|
||||
self._dirty = True
|
||||
self.save_to_db()
|
||||
return token
|
||||
|
||||
def session_account(self, token: str | None) -> dict | None:
|
||||
@@ -134,7 +173,9 @@ class AccountStore:
|
||||
if not token:
|
||||
return
|
||||
with self._lock:
|
||||
self._sessions.pop(token, None)
|
||||
if self._sessions.pop(token, None) is not None:
|
||||
self._dirty = True
|
||||
self.save_to_db()
|
||||
|
||||
# ---- API keys ----
|
||||
|
||||
@@ -191,6 +232,7 @@ class AccountStore:
|
||||
"account_id": record["account_id"],
|
||||
"email": record.get("email"),
|
||||
"wallet": record.get("wallet"),
|
||||
"nickname": record.get("nickname"),
|
||||
"role": record["role"],
|
||||
"created_ts": record.get("ts", 0.0),
|
||||
}
|
||||
@@ -244,11 +286,19 @@ class AccountStore:
|
||||
"role": event.get("role", "user"),
|
||||
"password_hash": event["password_hash"],
|
||||
"salt": event["salt"],
|
||||
"nickname": event.get("nickname"),
|
||||
"ts": float(event.get("ts", 0.0)),
|
||||
}
|
||||
self._accounts[account_id] = record
|
||||
for identifier in filter(None, (record["email"], record["wallet"])):
|
||||
self._by_identifier.setdefault(identifier.lower(), account_id)
|
||||
elif etype == "update_profile":
|
||||
account_id = event["account_id"]
|
||||
record = self._accounts.get(account_id)
|
||||
if record is None:
|
||||
return
|
||||
if "nickname" in event:
|
||||
record["nickname"] = event.get("nickname")
|
||||
elif etype == "create_key":
|
||||
api_key = event["api_key"]
|
||||
if api_key not in self._revoked_keys:
|
||||
@@ -271,6 +321,10 @@ class AccountStore:
|
||||
"CREATE TABLE IF NOT EXISTS account_events "
|
||||
"(event_id TEXT PRIMARY KEY, payload TEXT NOT NULL, ts REAL NOT NULL)"
|
||||
)
|
||||
con.execute(
|
||||
"CREATE TABLE IF NOT EXISTS account_sessions "
|
||||
"(token TEXT PRIMARY KEY, account_id TEXT NOT NULL, expires REAL NOT NULL)"
|
||||
)
|
||||
con.commit()
|
||||
con.close()
|
||||
|
||||
@@ -279,6 +333,10 @@ class AccountStore:
|
||||
rows = con.execute(
|
||||
"SELECT payload FROM account_events ORDER BY ts, event_id"
|
||||
).fetchall()
|
||||
session_rows = con.execute(
|
||||
"SELECT token, account_id, expires FROM account_sessions WHERE expires >= ?",
|
||||
(time.time(),),
|
||||
).fetchall()
|
||||
con.close()
|
||||
with self._lock:
|
||||
for (payload,) in rows:
|
||||
@@ -288,6 +346,11 @@ class AccountStore:
|
||||
continue
|
||||
if event.get("id") not in self._seen_event_ids:
|
||||
self._apply_locked(event)
|
||||
self._sessions = {
|
||||
token: {"account_id": account_id, "expires": float(expires)}
|
||||
for token, account_id, expires in session_rows
|
||||
if account_id in self._accounts
|
||||
}
|
||||
self._dirty = False
|
||||
|
||||
def save_to_db(self) -> None:
|
||||
@@ -297,11 +360,21 @@ class AccountStore:
|
||||
if not self._dirty:
|
||||
return
|
||||
events = list(self._event_log)
|
||||
sessions = [
|
||||
(token, session["account_id"], float(session["expires"]))
|
||||
for token, session in self._sessions.items()
|
||||
if session["expires"] >= time.time()
|
||||
]
|
||||
self._dirty = False
|
||||
con = sqlite3.connect(self._db_path) # type: ignore[arg-type]
|
||||
con.executemany(
|
||||
"INSERT OR IGNORE INTO account_events (event_id, payload, ts) VALUES (?, ?, ?)",
|
||||
[(e["id"], json.dumps(e), float(e.get("ts", 0.0))) for e in events],
|
||||
)
|
||||
con.execute("DELETE FROM account_sessions")
|
||||
con.executemany(
|
||||
"INSERT INTO account_sessions (token, account_id, expires) VALUES (?, ?, ?)",
|
||||
sessions,
|
||||
)
|
||||
con.commit()
|
||||
con.close()
|
||||
|
||||
@@ -453,13 +453,12 @@ class BillingLedger:
|
||||
with self._lock:
|
||||
return self._node_pending.get(wallet, 0.0)
|
||||
|
||||
def usage_for(self, api_keys: list[str], *, recent_limit: int | None = None) -> dict:
|
||||
"""Aggregate charge history for a set of API keys (dashboard view)."""
|
||||
def usage_totals_for(self, api_keys: list[str]) -> dict:
|
||||
"""Aggregate charge totals without per-request records (dashboard summary)."""
|
||||
keys = set(api_keys)
|
||||
requests = 0
|
||||
total_tokens = 0
|
||||
total_cost = 0.0
|
||||
records: list[dict] = []
|
||||
with self._lock:
|
||||
for event in self._event_log:
|
||||
if event.get("type") != "charge" or event.get("api_key") not in keys:
|
||||
@@ -467,6 +466,20 @@ class BillingLedger:
|
||||
requests += 1
|
||||
total_tokens += int(event.get("total_tokens", 0))
|
||||
total_cost += float(event.get("cost", 0.0))
|
||||
return {
|
||||
"requests": requests,
|
||||
"total_tokens": total_tokens,
|
||||
"total_cost": total_cost,
|
||||
}
|
||||
|
||||
def usage_for(self, api_keys: list[str], *, recent_limit: int | None = None) -> dict:
|
||||
"""Aggregate charge history for a set of API keys (dashboard view)."""
|
||||
keys = set(api_keys)
|
||||
records: list[dict] = []
|
||||
with self._lock:
|
||||
for event in self._event_log:
|
||||
if event.get("type") != "charge" or event.get("api_key") not in keys:
|
||||
continue
|
||||
records.append({
|
||||
"api_key": event["api_key"],
|
||||
"model": event.get("model"),
|
||||
@@ -476,9 +489,9 @@ class BillingLedger:
|
||||
})
|
||||
recent = records[-recent_limit:] if recent_limit is not None else records
|
||||
return {
|
||||
"requests": requests,
|
||||
"total_tokens": total_tokens,
|
||||
"total_cost": total_cost,
|
||||
"requests": len(records),
|
||||
"total_tokens": sum(int(r.get("total_tokens", 0)) for r in records),
|
||||
"total_cost": sum(float(r.get("cost", 0.0)) for r in records),
|
||||
"records": records,
|
||||
"recent": recent,
|
||||
}
|
||||
|
||||
@@ -1,336 +1,419 @@
|
||||
"""meshnet-tracker CLI entry point."""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from .accounts import DEFAULT_ACCOUNTS_DB_PATH
|
||||
from .billing import DEFAULT_BILLING_DB_PATH
|
||||
from .hf_pricing import DEFAULT_HF_PRICING_LOG_DB_PATH
|
||||
from .server import (
|
||||
DEFAULT_CALLER_CREDIT_USDT,
|
||||
DEFAULT_DEVNET_TOPUP_USDT,
|
||||
TrackerServer,
|
||||
derive_relay_url_from_public_tracker_url,
|
||||
)
|
||||
|
||||
DEFAULT_REGISTRY_DB_PATH = "meshnet_registry.sqlite3"
|
||||
|
||||
|
||||
def _load_env_file(path: Path) -> None:
|
||||
"""Load simple KEY=VALUE pairs from an env file without overriding env vars."""
|
||||
if not path.exists():
|
||||
return
|
||||
try:
|
||||
lines = path.read_text().splitlines()
|
||||
except OSError:
|
||||
return
|
||||
for line in lines:
|
||||
text = line.strip()
|
||||
if not text or text.startswith("#"):
|
||||
continue
|
||||
if text.startswith("export "):
|
||||
text = text[len("export "):].strip()
|
||||
if "=" not in text:
|
||||
continue
|
||||
key, value = text.split("=", 1)
|
||||
key = key.strip()
|
||||
if not key or key in os.environ:
|
||||
continue
|
||||
value = value.strip()
|
||||
if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}:
|
||||
value = value[1:-1]
|
||||
os.environ[key] = value
|
||||
|
||||
|
||||
def _load_env_defaults() -> None:
|
||||
"""Load local and user-level tracker env defaults before parsing arguments."""
|
||||
_load_env_file(Path.cwd() / ".env")
|
||||
_load_env_file(Path.home() / ".config" / "meshnet" / "secrets.env")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
_load_env_defaults()
|
||||
common = argparse.ArgumentParser(add_help=False)
|
||||
common.add_argument("--host", default="0.0.0.0", help="Host interface to listen on")
|
||||
common.add_argument("--port", type=int, default=8080, help="Port to listen on")
|
||||
common.add_argument(
|
||||
"--heartbeat-timeout",
|
||||
type=float,
|
||||
default=30.0,
|
||||
help="Seconds before a node is removed from the registry after missed heartbeat",
|
||||
)
|
||||
common.add_argument(
|
||||
"--cluster-peers",
|
||||
default="",
|
||||
help="Comma-separated URLs of peer tracker nodes (enables Raft cluster mode)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--self-url",
|
||||
default=None,
|
||||
help="This tracker's own URL as seen by peers (auto-derived from --host/--port if omitted)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--stats-db",
|
||||
default=None,
|
||||
metavar="PATH",
|
||||
help="SQLite database path for persistent model usage statistics",
|
||||
)
|
||||
common.add_argument(
|
||||
"--relay-url",
|
||||
default=None,
|
||||
help="Public ws(s):// relay URL advertised to nodes, for example wss://ai.neuron.d-popov.com/ws",
|
||||
)
|
||||
common.add_argument(
|
||||
"--billing-db",
|
||||
default=DEFAULT_BILLING_DB_PATH,
|
||||
metavar="PATH",
|
||||
help=(
|
||||
"SQLite database path for the USDT billing ledger "
|
||||
f"(default: {DEFAULT_BILLING_DB_PATH}; ADR-0015)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--no-billing",
|
||||
action="store_true",
|
||||
help="Disable the USDT billing ledger",
|
||||
)
|
||||
common.add_argument(
|
||||
"--max-charge-per-request",
|
||||
type=float,
|
||||
default=None,
|
||||
help=(
|
||||
"Reject chat completion requests whose prompt plus requested completion "
|
||||
"token bound would cost more than this many USDT"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--starting-credit",
|
||||
type=float,
|
||||
default=DEFAULT_CALLER_CREDIT_USDT,
|
||||
metavar="USDT",
|
||||
help=(
|
||||
"One-time Caller Credit granted when an account creates its first "
|
||||
f"API key (default: {DEFAULT_CALLER_CREDIT_USDT}; set 0 to require "
|
||||
"deposits before inference)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--devnet-topup",
|
||||
type=float,
|
||||
default=DEFAULT_DEVNET_TOPUP_USDT,
|
||||
metavar="USDT",
|
||||
help=(
|
||||
"Dashboard devnet top-up faucet: each click credits this many USDT "
|
||||
f"to one of the account's keys (default: {DEFAULT_DEVNET_TOPUP_USDT}; "
|
||||
"MUST be 0 on mainnet deployments)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--registry-db",
|
||||
default=DEFAULT_REGISTRY_DB_PATH,
|
||||
metavar="PATH",
|
||||
help=(
|
||||
"SQLite database path for persisted strike/ban/reputation registry "
|
||||
f"state (default: {DEFAULT_REGISTRY_DB_PATH})"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--no-registry-contracts",
|
||||
action="store_true",
|
||||
help="Disable the local contract registry used for strike/ban/reputation enforcement",
|
||||
)
|
||||
common.add_argument(
|
||||
"--accounts-db",
|
||||
default=DEFAULT_ACCOUNTS_DB_PATH,
|
||||
metavar="PATH",
|
||||
help=(
|
||||
"SQLite database path for dashboard user accounts "
|
||||
f"(default: {DEFAULT_ACCOUNTS_DB_PATH})"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--no-accounts",
|
||||
action="store_true",
|
||||
help="Disable dashboard user accounts (registration/login)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--solana-rpc-url",
|
||||
default=None,
|
||||
help="Solana RPC URL (e.g. https://api.devnet.solana.com); enables the on-chain treasury",
|
||||
)
|
||||
common.add_argument(
|
||||
"--usdt-mint",
|
||||
default=None,
|
||||
help="SPL mint address of (mock) USDT — see scripts/devnet_setup.py",
|
||||
)
|
||||
common.add_argument(
|
||||
"--treasury-keypair",
|
||||
default=None,
|
||||
metavar="PATH",
|
||||
help="Treasury keypair JSON path (only on settlement-capable trackers)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--settle-period",
|
||||
type=float,
|
||||
default=86400.0,
|
||||
help="Max seconds between payouts to a node (dev: 60, prod: 86400)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--payout-threshold",
|
||||
type=float,
|
||||
default=5.0,
|
||||
help="Pending USDT that triggers an immediate payout (dev: 0)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--payout-dust-floor",
|
||||
type=float,
|
||||
default=0.01,
|
||||
help="Never pay out less than this many USDT",
|
||||
)
|
||||
common.add_argument(
|
||||
"--validator-service-token",
|
||||
default=None,
|
||||
help=(
|
||||
"Service token the validator uses on POST /v1/billing/forfeit "
|
||||
"(default: MESHNET_VALIDATOR_SERVICE_TOKEN env; ADR-0017)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--hive-secret",
|
||||
default=None,
|
||||
help=(
|
||||
"Shared secret authenticating gossip between tracker peers "
|
||||
"(default: MESHNET_HIVE_SECRET env; required for multi-tracker replication)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--toploc-calibration-db",
|
||||
default=None,
|
||||
metavar="PATH",
|
||||
help=(
|
||||
"SQLite path for the AH-021 honest-noise TOPLOC calibration corpus "
|
||||
"(enables POST /v1/calibration/toploc/run + GET /v1/calibration/toploc/results)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--toploc-reference-node-url",
|
||||
default=None,
|
||||
help="Reference node the calibration job teacher-forces claimed tokens against (see validator README)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--toploc-calibration-gate-min-hardware-profiles",
|
||||
type=int,
|
||||
default=1,
|
||||
help=(
|
||||
"Distinct (GPU model, dtype) profiles the corpus must cover before "
|
||||
"gate_status.ready is true (alpha exception: fleet size is acceptable)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--enable-hf-pricing",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Enable the daily dynamic pricing refresh (issue 23): for presets with a "
|
||||
"curated hf_aliases list, sets the client price to 80%% of the cheapest "
|
||||
"matching HuggingFace inference-marketplace rate. Presets without "
|
||||
"hf_aliases are unaffected and keep their static price."
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--hf-pricing-log-db",
|
||||
default=None,
|
||||
metavar="PATH",
|
||||
help=(
|
||||
"SQLite database path for the dynamic pricing change log "
|
||||
f"(default when --enable-hf-pricing is set: {DEFAULT_HF_PRICING_LOG_DB_PATH}; "
|
||||
"enables GET /v1/pricing/hf/history)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--hf-pricing-refresh-interval",
|
||||
type=float,
|
||||
default=86400.0,
|
||||
help="Seconds between dynamic pricing refresh passes (default: daily)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--models-dir",
|
||||
default=None,
|
||||
metavar="PATH",
|
||||
help="Local HuggingFace snapshot root advertised as tracker model-file source (default: MESHNET_MODELS_DIR)",
|
||||
)
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="meshnet-tracker",
|
||||
description="Distributed Inference Network node registry and route selection",
|
||||
parents=[common],
|
||||
)
|
||||
subparsers = parser.add_subparsers(dest="command")
|
||||
|
||||
subparsers.add_parser("start", help="Start the tracker server", parents=[common])
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.command in {None, "start"}:
|
||||
cluster_peers = [u.strip() for u in args.cluster_peers.split(",") if u.strip()]
|
||||
relay_url = args.relay_url or derive_relay_url_from_public_tracker_url(args.self_url)
|
||||
treasury = None
|
||||
if args.solana_rpc_url and args.usdt_mint and args.treasury_keypair:
|
||||
from meshnet_contracts.solana_adapter import SolanaCustodialTreasury
|
||||
|
||||
treasury = SolanaCustodialTreasury(
|
||||
args.solana_rpc_url, args.usdt_mint, args.treasury_keypair,
|
||||
)
|
||||
contracts = None
|
||||
if not args.no_registry_contracts:
|
||||
from meshnet_contracts import LocalSolanaContracts # type: ignore[import-not-found]
|
||||
|
||||
contracts = LocalSolanaContracts(registry_db=args.registry_db)
|
||||
server = TrackerServer(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
heartbeat_timeout=args.heartbeat_timeout,
|
||||
cluster_peers=cluster_peers or None,
|
||||
cluster_self_url=args.self_url,
|
||||
stats_db=getattr(args, "stats_db", None),
|
||||
relay_url=relay_url,
|
||||
enable_billing=not args.no_billing,
|
||||
billing_db=None if args.no_billing else args.billing_db,
|
||||
max_charge_per_request=args.max_charge_per_request,
|
||||
starting_credit=args.starting_credit,
|
||||
devnet_topup_amount=args.devnet_topup,
|
||||
contracts=contracts,
|
||||
accounts_db=None if args.no_accounts else args.accounts_db,
|
||||
treasury=treasury,
|
||||
settle_period=args.settle_period,
|
||||
payout_threshold=args.payout_threshold,
|
||||
payout_dust_floor=args.payout_dust_floor,
|
||||
validator_service_token=args.validator_service_token,
|
||||
hive_secret=args.hive_secret,
|
||||
toploc_calibration_db=args.toploc_calibration_db,
|
||||
toploc_reference_node_url=args.toploc_reference_node_url,
|
||||
toploc_calibration_gate_min_hardware_profiles=args.toploc_calibration_gate_min_hardware_profiles,
|
||||
enable_hf_pricing=args.enable_hf_pricing,
|
||||
hf_pricing_log_db=(
|
||||
args.hf_pricing_log_db
|
||||
or (DEFAULT_HF_PRICING_LOG_DB_PATH if args.enable_hf_pricing else None)
|
||||
),
|
||||
hf_pricing_refresh_interval=args.hf_pricing_refresh_interval,
|
||||
models_dir=args.models_dir,
|
||||
)
|
||||
port = server.start()
|
||||
print(f"meshnet-tracker listening on http://{args.host}:{port}", flush=True)
|
||||
try:
|
||||
while True:
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
server.stop()
|
||||
sys.exit(0)
|
||||
else:
|
||||
parser.print_help()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
"""meshnet-tracker CLI entry point."""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from .accounts import DEFAULT_ACCOUNTS_DB_PATH
|
||||
from .billing import DEFAULT_BILLING_DB_PATH
|
||||
from .hf_pricing import DEFAULT_HF_PRICING_LOG_DB_PATH
|
||||
from .logging_setup import (
|
||||
DEFAULT_LOG_BACKUP_COUNT,
|
||||
DEFAULT_LOG_DIR,
|
||||
DEFAULT_LOG_MAX_BYTES,
|
||||
configure_tracker_file_logging,
|
||||
)
|
||||
from .routing_stats import RoutingConfig
|
||||
from .server import (
|
||||
DEFAULT_CALLER_CREDIT_USDT,
|
||||
DEFAULT_DEVNET_TOPUP_USDT,
|
||||
TrackerServer,
|
||||
derive_relay_url_from_public_tracker_url,
|
||||
)
|
||||
|
||||
DEFAULT_REGISTRY_DB_PATH = "meshnet_registry.sqlite3"
|
||||
|
||||
|
||||
def _load_env_file(path: Path) -> None:
|
||||
"""Load simple KEY=VALUE pairs from an env file without overriding env vars."""
|
||||
if not path.exists():
|
||||
return
|
||||
try:
|
||||
lines = path.read_text().splitlines()
|
||||
except OSError:
|
||||
return
|
||||
for line in lines:
|
||||
text = line.strip()
|
||||
if not text or text.startswith("#"):
|
||||
continue
|
||||
if text.startswith("export "):
|
||||
text = text[len("export "):].strip()
|
||||
if "=" not in text:
|
||||
continue
|
||||
key, value = text.split("=", 1)
|
||||
key = key.strip()
|
||||
if not key or key in os.environ:
|
||||
continue
|
||||
value = value.strip()
|
||||
if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}:
|
||||
value = value[1:-1]
|
||||
os.environ[key] = value
|
||||
|
||||
|
||||
def _load_env_defaults() -> None:
|
||||
"""Load local and user-level tracker env defaults before parsing arguments."""
|
||||
_load_env_file(Path.cwd() / ".env")
|
||||
_load_env_file(Path.home() / ".config" / "meshnet" / "secrets.env")
|
||||
|
||||
|
||||
def _routing_config_from_args(args: argparse.Namespace) -> RoutingConfig | None:
|
||||
"""Build a RoutingConfig from CLI flags; None keeps env-var/server defaults."""
|
||||
overrides = {
|
||||
"explore_share": args.route_explore_share,
|
||||
"weight_alpha": args.route_weight_alpha,
|
||||
"stats_half_life_seconds": args.route_stats_half_life,
|
||||
}
|
||||
set_values = {key: value for key, value in overrides.items() if value is not None}
|
||||
if not set_values:
|
||||
return None
|
||||
return RoutingConfig(**set_values)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
_load_env_defaults()
|
||||
common = argparse.ArgumentParser(add_help=False)
|
||||
common.add_argument("--host", default="0.0.0.0", help="Host interface to listen on")
|
||||
common.add_argument("--port", type=int, default=8080, help="Port to listen on")
|
||||
common.add_argument(
|
||||
"--heartbeat-timeout",
|
||||
type=float,
|
||||
default=30.0,
|
||||
help="Seconds before a node is removed from the registry after missed heartbeat",
|
||||
)
|
||||
common.add_argument(
|
||||
"--cluster-peers",
|
||||
default="",
|
||||
help="Comma-separated URLs of peer tracker nodes (enables Raft cluster mode)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--self-url",
|
||||
default=None,
|
||||
help="This tracker's own URL as seen by peers (auto-derived from --host/--port if omitted)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--stats-db",
|
||||
default=None,
|
||||
metavar="PATH",
|
||||
help="SQLite database path for persistent model usage statistics",
|
||||
)
|
||||
common.add_argument(
|
||||
"--relay-url",
|
||||
default=None,
|
||||
help="Public ws(s):// relay URL advertised to nodes, for example wss://ai.neuron.d-popov.com/ws",
|
||||
)
|
||||
common.add_argument(
|
||||
"--billing-db",
|
||||
default=DEFAULT_BILLING_DB_PATH,
|
||||
metavar="PATH",
|
||||
help=(
|
||||
"SQLite database path for the USDT billing ledger "
|
||||
f"(default: {DEFAULT_BILLING_DB_PATH}; ADR-0015)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--no-billing",
|
||||
action="store_true",
|
||||
help="Disable the USDT billing ledger",
|
||||
)
|
||||
common.add_argument(
|
||||
"--max-charge-per-request",
|
||||
type=float,
|
||||
default=None,
|
||||
help=(
|
||||
"Reject chat completion requests whose prompt plus requested completion "
|
||||
"token bound would cost more than this many USDT"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--starting-credit",
|
||||
type=float,
|
||||
default=DEFAULT_CALLER_CREDIT_USDT,
|
||||
metavar="USDT",
|
||||
help=(
|
||||
"One-time Caller Credit granted when an account creates its first "
|
||||
f"API key (default: {DEFAULT_CALLER_CREDIT_USDT}; set 0 to require "
|
||||
"deposits before inference)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--devnet-topup",
|
||||
type=float,
|
||||
default=DEFAULT_DEVNET_TOPUP_USDT,
|
||||
metavar="USDT",
|
||||
help=(
|
||||
"Dashboard devnet top-up faucet: each click credits this many USDT "
|
||||
f"to one of the account's keys (default: {DEFAULT_DEVNET_TOPUP_USDT}; "
|
||||
"MUST be 0 on mainnet deployments)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--registry-db",
|
||||
default=DEFAULT_REGISTRY_DB_PATH,
|
||||
metavar="PATH",
|
||||
help=(
|
||||
"SQLite database path for persisted strike/ban/reputation registry "
|
||||
f"state (default: {DEFAULT_REGISTRY_DB_PATH})"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--no-registry-contracts",
|
||||
action="store_true",
|
||||
help="Disable the local contract registry used for strike/ban/reputation enforcement",
|
||||
)
|
||||
common.add_argument(
|
||||
"--accounts-db",
|
||||
default=DEFAULT_ACCOUNTS_DB_PATH,
|
||||
metavar="PATH",
|
||||
help=(
|
||||
"SQLite database path for dashboard user accounts "
|
||||
f"(default: {DEFAULT_ACCOUNTS_DB_PATH})"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--no-accounts",
|
||||
action="store_true",
|
||||
help="Disable dashboard user accounts (registration/login)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--solana-rpc-url",
|
||||
default=None,
|
||||
help="Solana RPC URL (e.g. https://api.devnet.solana.com); enables the on-chain treasury",
|
||||
)
|
||||
common.add_argument(
|
||||
"--usdt-mint",
|
||||
default=None,
|
||||
help="SPL mint address of (mock) USDT — see scripts/devnet_setup.py",
|
||||
)
|
||||
common.add_argument(
|
||||
"--treasury-keypair",
|
||||
default=None,
|
||||
metavar="PATH",
|
||||
help="Treasury keypair JSON path (only on settlement-capable trackers)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--settle-period",
|
||||
type=float,
|
||||
default=86400.0,
|
||||
help="Max seconds between payouts to a node (dev: 60, prod: 86400)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--payout-threshold",
|
||||
type=float,
|
||||
default=5.0,
|
||||
help="Pending USDT that triggers an immediate payout (dev: 0)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--payout-dust-floor",
|
||||
type=float,
|
||||
default=0.01,
|
||||
help="Never pay out less than this many USDT",
|
||||
)
|
||||
common.add_argument(
|
||||
"--validator-service-token",
|
||||
default=None,
|
||||
help=(
|
||||
"Service token the validator uses on POST /v1/billing/forfeit "
|
||||
"(default: MESHNET_VALIDATOR_SERVICE_TOKEN env; ADR-0017)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--hive-secret",
|
||||
default=None,
|
||||
help=(
|
||||
"Shared secret authenticating gossip between tracker peers "
|
||||
"(default: MESHNET_HIVE_SECRET env; required for multi-tracker replication)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--toploc-calibration-db",
|
||||
default=None,
|
||||
metavar="PATH",
|
||||
help=(
|
||||
"SQLite path for the AH-021 honest-noise TOPLOC calibration corpus "
|
||||
"(enables POST /v1/calibration/toploc/run + GET /v1/calibration/toploc/results)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--toploc-reference-node-url",
|
||||
default=None,
|
||||
help="Reference node the calibration job teacher-forces claimed tokens against (see validator README)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--toploc-calibration-gate-min-hardware-profiles",
|
||||
type=int,
|
||||
default=1,
|
||||
help=(
|
||||
"Distinct (GPU model, dtype) profiles the corpus must cover before "
|
||||
"gate_status.ready is true (alpha exception: fleet size is acceptable)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--enable-hf-pricing",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Enable the daily dynamic pricing refresh (issue 23): for presets with a "
|
||||
"curated hf_aliases list, sets the client price to 80%% of the cheapest "
|
||||
"matching HuggingFace inference-marketplace rate. Presets without "
|
||||
"hf_aliases are unaffected and keep their static price."
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--hf-pricing-log-db",
|
||||
default=None,
|
||||
metavar="PATH",
|
||||
help=(
|
||||
"SQLite database path for the dynamic pricing change log "
|
||||
f"(default when --enable-hf-pricing is set: {DEFAULT_HF_PRICING_LOG_DB_PATH}; "
|
||||
"enables GET /v1/pricing/hf/history)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--hf-pricing-refresh-interval",
|
||||
type=float,
|
||||
default=86400.0,
|
||||
help="Seconds between dynamic pricing refresh passes (default: daily)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--models-dir",
|
||||
default=None,
|
||||
metavar="PATH",
|
||||
help="Local HuggingFace snapshot root advertised as tracker model-file source (default: MESHNET_MODELS_DIR)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--route-explore-share",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="FRACTION",
|
||||
help=(
|
||||
"Fraction of requests routed down unproven/stale routes to gather "
|
||||
"throughput statistics (ADR-0021; default 0.3, lower once traffic grows)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--route-weight-alpha",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="ALPHA",
|
||||
help=(
|
||||
"Traffic weight exponent among proven routes: share ∝ tps^alpha "
|
||||
"(default 1.0 — a 1.5x-faster route gets 1.5x the traffic)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--route-stats-half-life",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="SECONDS",
|
||||
help="Half-life for decaying route throughput observations (default 600)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--log-dir",
|
||||
default=DEFAULT_LOG_DIR,
|
||||
metavar="PATH",
|
||||
help=(
|
||||
"Directory for rotating tracker logs "
|
||||
f"(default: {DEFAULT_LOG_DIR}; files: info.log, warning.log, error.log)"
|
||||
),
|
||||
)
|
||||
common.add_argument(
|
||||
"--log-max-bytes",
|
||||
type=int,
|
||||
default=DEFAULT_LOG_MAX_BYTES,
|
||||
metavar="BYTES",
|
||||
help=f"Rotate each tracker log file after this many bytes (default: {DEFAULT_LOG_MAX_BYTES})",
|
||||
)
|
||||
common.add_argument(
|
||||
"--log-backup-count",
|
||||
type=int,
|
||||
default=DEFAULT_LOG_BACKUP_COUNT,
|
||||
metavar="N",
|
||||
help=f"Number of rotated tracker log files to keep per level (default: {DEFAULT_LOG_BACKUP_COUNT})",
|
||||
)
|
||||
common.add_argument(
|
||||
"--no-file-logs",
|
||||
action="store_true",
|
||||
help="Disable rotating tracker log files and only write to the terminal",
|
||||
)
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="meshnet-tracker",
|
||||
description="Distributed Inference Network node registry and route selection",
|
||||
parents=[common],
|
||||
)
|
||||
subparsers = parser.add_subparsers(dest="command")
|
||||
|
||||
subparsers.add_parser("start", help="Start the tracker server", parents=[common])
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.command in {None, "start"}:
|
||||
if not args.no_file_logs:
|
||||
log_dir = configure_tracker_file_logging(
|
||||
args.log_dir,
|
||||
max_bytes=args.log_max_bytes,
|
||||
backup_count=args.log_backup_count,
|
||||
)
|
||||
print(f"meshnet-tracker logs: {log_dir}", flush=True)
|
||||
cluster_peers = [u.strip() for u in args.cluster_peers.split(",") if u.strip()]
|
||||
relay_url = args.relay_url or derive_relay_url_from_public_tracker_url(args.self_url)
|
||||
treasury = None
|
||||
if args.solana_rpc_url and args.usdt_mint and args.treasury_keypair:
|
||||
from meshnet_contracts.solana_adapter import SolanaCustodialTreasury
|
||||
|
||||
treasury = SolanaCustodialTreasury(
|
||||
args.solana_rpc_url, args.usdt_mint, args.treasury_keypair,
|
||||
)
|
||||
contracts = None
|
||||
if not args.no_registry_contracts:
|
||||
from meshnet_contracts import LocalSolanaContracts # type: ignore[import-not-found]
|
||||
|
||||
contracts = LocalSolanaContracts(registry_db=args.registry_db)
|
||||
server = TrackerServer(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
heartbeat_timeout=args.heartbeat_timeout,
|
||||
cluster_peers=cluster_peers or None,
|
||||
cluster_self_url=args.self_url,
|
||||
stats_db=getattr(args, "stats_db", None),
|
||||
relay_url=relay_url,
|
||||
enable_billing=not args.no_billing,
|
||||
billing_db=None if args.no_billing else args.billing_db,
|
||||
max_charge_per_request=args.max_charge_per_request,
|
||||
starting_credit=args.starting_credit,
|
||||
devnet_topup_amount=args.devnet_topup,
|
||||
contracts=contracts,
|
||||
accounts_db=None if args.no_accounts else args.accounts_db,
|
||||
treasury=treasury,
|
||||
settle_period=args.settle_period,
|
||||
payout_threshold=args.payout_threshold,
|
||||
payout_dust_floor=args.payout_dust_floor,
|
||||
validator_service_token=args.validator_service_token,
|
||||
hive_secret=args.hive_secret,
|
||||
toploc_calibration_db=args.toploc_calibration_db,
|
||||
toploc_reference_node_url=args.toploc_reference_node_url,
|
||||
toploc_calibration_gate_min_hardware_profiles=args.toploc_calibration_gate_min_hardware_profiles,
|
||||
enable_hf_pricing=args.enable_hf_pricing,
|
||||
hf_pricing_log_db=(
|
||||
args.hf_pricing_log_db
|
||||
or (DEFAULT_HF_PRICING_LOG_DB_PATH if args.enable_hf_pricing else None)
|
||||
),
|
||||
hf_pricing_refresh_interval=args.hf_pricing_refresh_interval,
|
||||
models_dir=args.models_dir,
|
||||
routing_config=_routing_config_from_args(args),
|
||||
)
|
||||
port = server.start()
|
||||
print(f"meshnet-tracker listening on http://{args.host}:{port}", flush=True)
|
||||
try:
|
||||
while True:
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
server.stop()
|
||||
sys.exit(0)
|
||||
else:
|
||||
parser.print_help()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
99
packages/tracker/meshnet_tracker/logging_setup.py
Normal file
99
packages/tracker/meshnet_tracker/logging_setup.py
Normal file
@@ -0,0 +1,99 @@
|
||||
"""Rotating file logging for the tracker CLI."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from logging.handlers import RotatingFileHandler
|
||||
from pathlib import Path
|
||||
from typing import TextIO
|
||||
|
||||
|
||||
DEFAULT_LOG_DIR = "logs/tracker"
|
||||
DEFAULT_LOG_MAX_BYTES = 10 * 1024 * 1024
|
||||
DEFAULT_LOG_BACKUP_COUNT = 5
|
||||
TRACKER_LOGGER_NAME = "meshnet.tracker"
|
||||
|
||||
|
||||
class _ExactLevelFilter(logging.Filter):
|
||||
def __init__(self, level: int) -> None:
|
||||
super().__init__()
|
||||
self._level = level
|
||||
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
return record.levelno == self._level
|
||||
|
||||
|
||||
class _TeeStream:
|
||||
def __init__(self, stream: TextIO, logger: logging.Logger, level: int) -> None:
|
||||
self._stream = stream
|
||||
self._logger = logger
|
||||
self._level = level
|
||||
self._buffer = ""
|
||||
|
||||
def write(self, text: str) -> int:
|
||||
self._stream.write(text)
|
||||
self._stream.flush()
|
||||
self._buffer += text
|
||||
while "\n" in self._buffer:
|
||||
line, self._buffer = self._buffer.split("\n", 1)
|
||||
line = line.rstrip()
|
||||
if line:
|
||||
self._logger.log(self._level, line)
|
||||
return len(text)
|
||||
|
||||
def flush(self) -> None:
|
||||
self._stream.flush()
|
||||
line = self._buffer.rstrip()
|
||||
if line:
|
||||
self._logger.log(self._level, line)
|
||||
self._buffer = ""
|
||||
|
||||
def isatty(self) -> bool:
|
||||
return self._stream.isatty()
|
||||
|
||||
|
||||
def _make_handler(path: Path, level: int, *, max_bytes: int, backup_count: int) -> RotatingFileHandler:
|
||||
handler = RotatingFileHandler(
|
||||
path,
|
||||
maxBytes=max_bytes,
|
||||
backupCount=backup_count,
|
||||
encoding="utf-8",
|
||||
)
|
||||
handler.setLevel(level)
|
||||
handler.addFilter(_ExactLevelFilter(level))
|
||||
handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
|
||||
return handler
|
||||
|
||||
|
||||
def configure_tracker_file_logging(
|
||||
log_dir: str | Path = DEFAULT_LOG_DIR,
|
||||
*,
|
||||
max_bytes: int = DEFAULT_LOG_MAX_BYTES,
|
||||
backup_count: int = DEFAULT_LOG_BACKUP_COUNT,
|
||||
tee_stdio: bool = True,
|
||||
) -> Path:
|
||||
"""Configure rotatable info/warning/error log files and return the directory."""
|
||||
|
||||
path = Path(log_dir).expanduser()
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
logger = logging.getLogger(TRACKER_LOGGER_NAME)
|
||||
logger.setLevel(logging.INFO)
|
||||
logger.propagate = False
|
||||
logger.handlers.clear()
|
||||
logger.addHandler(_make_handler(path / "info.log", logging.INFO, max_bytes=max_bytes, backup_count=backup_count))
|
||||
logger.addHandler(_make_handler(path / "warning.log", logging.WARNING, max_bytes=max_bytes, backup_count=backup_count))
|
||||
logger.addHandler(_make_handler(path / "error.log", logging.ERROR, max_bytes=max_bytes, backup_count=backup_count))
|
||||
|
||||
if tee_stdio:
|
||||
if not isinstance(sys.stdout, _TeeStream):
|
||||
sys.stdout = _TeeStream(sys.stdout, logger, logging.INFO) # type: ignore[assignment]
|
||||
if not isinstance(sys.stderr, _TeeStream):
|
||||
sys.stderr = _TeeStream(sys.stderr, logger, logging.ERROR) # type: ignore[assignment]
|
||||
|
||||
return path
|
||||
|
||||
|
||||
def tracker_logger() -> logging.Logger:
|
||||
return logging.getLogger(TRACKER_LOGGER_NAME)
|
||||
@@ -39,6 +39,38 @@
|
||||
]
|
||||
}
|
||||
},
|
||||
"qwen2.5-0.5b-instruct": {
|
||||
"layers_start": 0,
|
||||
"layers_end": 23,
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"aliases": [
|
||||
"qwen2.5-0.5b",
|
||||
"Qwen2.5-0.5B-Instruct",
|
||||
"Qwen/Qwen2.5-0.5B-Instruct"
|
||||
],
|
||||
"deployment_status": "supported",
|
||||
"price_per_1k_tokens": 0.002,
|
||||
"input_price_per_1k_tokens": 0.002,
|
||||
"output_price_per_1k_tokens": 0.002,
|
||||
"hf_aliases": [],
|
||||
"hf_verified_match_note": "Static 10× dev-funding markup over ~$0.20/1M commercial API reference (Qwen-class hosted rates). $0.002/1k = $2/1M blended input+output.",
|
||||
"required_model_bytes": 1056964608,
|
||||
"download_size_bytes": 1056964608,
|
||||
"native_quantization": "bfloat16",
|
||||
"canonical_audit_dtype": "bfloat16",
|
||||
"canonical_audit_quantization": "bfloat16",
|
||||
"bytes_per_layer": {
|
||||
"bfloat16": 44040192
|
||||
},
|
||||
"metadata": {
|
||||
"architecture": "Dense transformer (GQA)",
|
||||
"total_parameters": "0.5B",
|
||||
"num_layers": 24,
|
||||
"context_length": 32768,
|
||||
"native_quantization": "bfloat16",
|
||||
"download_size_gb": 1
|
||||
}
|
||||
},
|
||||
"qwen3.6-35b-a3b": {
|
||||
"layers_start": 0,
|
||||
"layers_end": 39,
|
||||
|
||||
257
packages/tracker/meshnet_tracker/routing_stats.py
Normal file
257
packages/tracker/meshnet_tracker/routing_stats.py
Normal file
@@ -0,0 +1,257 @@
|
||||
"""Learned route statistics for dynamic bandit-style route selection (ADR-0021).
|
||||
|
||||
The tracker treats each viable route (ordered chain of node shards covering a
|
||||
model) as a bandit arm. Observed end-to-end tokens/sec per route is kept as a
|
||||
time-decayed EWMA. Selection splits traffic between:
|
||||
|
||||
- **exploit**: weighted-random among *proven* routes, weight ∝ tps ** alpha
|
||||
(alpha=1.0 → a 1.5x-faster route gets 1.5x the traffic);
|
||||
- **scout**: with probability `explore_share`, the least-measured unproven or
|
||||
stale route is chosen so the tracker keeps learning as the network morphs.
|
||||
|
||||
Staleness has two mechanisms:
|
||||
- continuous: sample mass decays with `stats_half_life_seconds`, so old
|
||||
observations fade;
|
||||
- abrupt: every node join/leave bumps the model's *topology epoch*; stats from
|
||||
an older epoch keep their EWMA as a prior but drop back into the scout pool
|
||||
until re-measured.
|
||||
|
||||
Route signatures embed node ids and shard ranges, so a node re-registering
|
||||
with a different shard produces a new arm automatically.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Iterable
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RoutingConfig:
|
||||
explore_share: float = 0.3
|
||||
weight_alpha: float = 1.0
|
||||
stats_half_life_seconds: float = 600.0
|
||||
min_sample_tokens: int = 8
|
||||
# One fresh sample has mass 1.0 and decays from there; 0.5 keeps a single
|
||||
# observation "proven" for one half-life before demoting it to the scout pool.
|
||||
min_proven_weight: float = 0.5
|
||||
max_candidate_routes: int = 8
|
||||
prune_after_seconds: float = 86400.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class RouteStat:
|
||||
ewma_tps: float = 0.0
|
||||
weight: float = 0.0 # decayed effective sample mass
|
||||
last_sample_ts: float = 0.0
|
||||
epoch: int = 0
|
||||
samples: int = 0 # lifetime raw sample count (display only)
|
||||
|
||||
def decayed_weight(self, now: float, half_life: float) -> float:
|
||||
if self.weight <= 0.0:
|
||||
return 0.0
|
||||
age = max(0.0, now - self.last_sample_ts)
|
||||
return self.weight * 0.5 ** (age / half_life)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RouteCandidate:
|
||||
nodes: list[Any]
|
||||
signature: str
|
||||
prior_tps: float = 0.0
|
||||
|
||||
|
||||
def route_signature(model_key: str, nodes: Iterable[Any]) -> str:
|
||||
hops = "->".join(
|
||||
f"{getattr(n, 'node_id', '?')}[{getattr(n, 'shard_start', '?')}-{getattr(n, 'shard_end', '?')}]"
|
||||
for n in nodes
|
||||
)
|
||||
return f"{model_key}|{hops}"
|
||||
|
||||
|
||||
class RouteStatsStore:
|
||||
"""Thread-safe per-route decayed throughput statistics."""
|
||||
|
||||
def __init__(self, config: RoutingConfig | None = None) -> None:
|
||||
self.config = config or RoutingConfig()
|
||||
self._lock = threading.Lock()
|
||||
self._stats: dict[str, RouteStat] = {}
|
||||
self._epochs: dict[str, int] = {}
|
||||
|
||||
def epoch(self, model_key: str) -> int:
|
||||
with self._lock:
|
||||
return self._epochs.get(model_key, 0)
|
||||
|
||||
def bump_epoch(self, model_keys: Iterable[str | None]) -> None:
|
||||
"""Mark the topology changed for the given model keys (node join/leave)."""
|
||||
with self._lock:
|
||||
for key in model_keys:
|
||||
if key:
|
||||
self._epochs[key] = self._epochs.get(key, 0) + 1
|
||||
|
||||
def record_sample(
|
||||
self,
|
||||
model_key: str,
|
||||
signature: str,
|
||||
tokens: int,
|
||||
elapsed_seconds: float,
|
||||
now: float | None = None,
|
||||
) -> bool:
|
||||
"""Fold one completed request into the route's EWMA.
|
||||
|
||||
Returns False (and records nothing) for samples below
|
||||
`min_sample_tokens` — near-empty completions come from broken routes
|
||||
and would poison the arm with meaningless throughput values.
|
||||
"""
|
||||
cfg = self.config
|
||||
if tokens < cfg.min_sample_tokens or elapsed_seconds <= 0.0:
|
||||
return False
|
||||
tps = tokens / elapsed_seconds
|
||||
ts = time.time() if now is None else now
|
||||
with self._lock:
|
||||
stat = self._stats.get(signature)
|
||||
if stat is None:
|
||||
stat = RouteStat()
|
||||
self._stats[signature] = stat
|
||||
carried = stat.decayed_weight(ts, cfg.stats_half_life_seconds)
|
||||
total = carried + 1.0
|
||||
stat.ewma_tps = (stat.ewma_tps * carried + tps) / total
|
||||
stat.weight = total
|
||||
stat.last_sample_ts = ts
|
||||
stat.epoch = self._epochs.get(model_key, 0)
|
||||
stat.samples += 1
|
||||
return True
|
||||
|
||||
def snapshot(self, signature: str, model_key: str, now: float | None = None) -> dict:
|
||||
"""Point-in-time view of one route's learned state."""
|
||||
ts = time.time() if now is None else now
|
||||
cfg = self.config
|
||||
with self._lock:
|
||||
stat = self._stats.get(signature)
|
||||
current_epoch = self._epochs.get(model_key, 0)
|
||||
if stat is None:
|
||||
return {"tps": None, "weight": 0.0, "samples": 0, "status": "unsampled"}
|
||||
weight = stat.decayed_weight(ts, cfg.stats_half_life_seconds)
|
||||
if stat.epoch != current_epoch:
|
||||
status = "stale"
|
||||
elif weight < cfg.min_proven_weight:
|
||||
status = "decayed" if stat.samples else "unsampled"
|
||||
else:
|
||||
status = "proven"
|
||||
return {
|
||||
"tps": round(stat.ewma_tps, 4) if stat.samples else None,
|
||||
"weight": round(weight, 4),
|
||||
"samples": stat.samples,
|
||||
"status": status,
|
||||
}
|
||||
|
||||
def prune(self, now: float | None = None) -> int:
|
||||
"""Drop routes with no samples for `prune_after_seconds`."""
|
||||
ts = time.time() if now is None else now
|
||||
cutoff = ts - self.config.prune_after_seconds
|
||||
with self._lock:
|
||||
dead = [sig for sig, stat in self._stats.items() if stat.last_sample_ts < cutoff]
|
||||
for sig in dead:
|
||||
del self._stats[sig]
|
||||
return len(dead)
|
||||
|
||||
|
||||
def choose_route(
|
||||
candidates: list[RouteCandidate],
|
||||
store: RouteStatsStore,
|
||||
model_key: str,
|
||||
rng: random.Random | None = None,
|
||||
now: float | None = None,
|
||||
) -> tuple[RouteCandidate | None, dict]:
|
||||
"""Pick a route: ε-scout among unproven arms, else weighted ∝ tps**alpha.
|
||||
|
||||
Returns (candidate, decision) where decision explains the pick for logs
|
||||
and diagnostics: {"mode": "scout"|"exploit"|"prior", ...}.
|
||||
"""
|
||||
if not candidates:
|
||||
return None, {"mode": "none"}
|
||||
rng = rng or random
|
||||
cfg = store.config
|
||||
proven: list[tuple[RouteCandidate, float]] = []
|
||||
scouts: list[tuple[RouteCandidate, float]] = []
|
||||
for cand in candidates:
|
||||
snap = store.snapshot(cand.signature, model_key, now=now)
|
||||
if snap["status"] == "proven":
|
||||
proven.append((cand, max(float(snap["tps"] or 0.0), 1e-6)))
|
||||
else:
|
||||
scouts.append((cand, float(snap["weight"])))
|
||||
if scouts and (not proven or rng.random() < cfg.explore_share):
|
||||
# Least-measured first so new/stale arms accumulate samples fastest;
|
||||
# tiebreak on prior estimate so plausible routes get scouted first.
|
||||
scouts.sort(key=lambda item: (item[1], -item[0].prior_tps))
|
||||
pick = scouts[0][0]
|
||||
return pick, {"mode": "scout", "signature": pick.signature}
|
||||
if proven:
|
||||
weights = [tps ** cfg.weight_alpha for _, tps in proven]
|
||||
pick = rng.choices([cand for cand, _ in proven], weights=weights, k=1)[0]
|
||||
return pick, {
|
||||
"mode": "exploit",
|
||||
"signature": pick.signature,
|
||||
"candidates": len(proven),
|
||||
}
|
||||
# No stats anywhere yet — fall back to the prior (benchmark-derived) estimate.
|
||||
weights = [max(cand.prior_tps, 1e-6) ** cfg.weight_alpha for cand in candidates]
|
||||
pick = rng.choices(candidates, weights=weights, k=1)[0]
|
||||
return pick, {"mode": "prior", "signature": pick.signature}
|
||||
|
||||
|
||||
def route_table(
|
||||
candidates: list[RouteCandidate],
|
||||
store: RouteStatsStore,
|
||||
model_key: str,
|
||||
now: float | None = None,
|
||||
) -> list[dict]:
|
||||
"""Diagnostics rows: learned tps, coefficient vs best, expected traffic share."""
|
||||
cfg = store.config
|
||||
rows = []
|
||||
for cand in candidates:
|
||||
snap = store.snapshot(cand.signature, model_key, now=now)
|
||||
rows.append({"candidate": cand, **snap})
|
||||
proven = [r for r in rows if r["status"] == "proven"]
|
||||
scouts = [r for r in rows if r["status"] != "proven"]
|
||||
best_tps = max((float(r["tps"]) for r in proven), default=0.0)
|
||||
exploit_budget = 1.0 - (cfg.explore_share if scouts and proven else 0.0)
|
||||
if not proven:
|
||||
exploit_budget = 0.0
|
||||
weight_sum = sum(float(r["tps"]) ** cfg.weight_alpha for r in proven) or 1.0
|
||||
out = []
|
||||
for r in rows:
|
||||
cand: RouteCandidate = r["candidate"]
|
||||
if r["status"] == "proven":
|
||||
share = exploit_budget * (float(r["tps"]) ** cfg.weight_alpha) / weight_sum
|
||||
coefficient = round(float(r["tps"]) / best_tps, 3) if best_tps else None
|
||||
else:
|
||||
share = (
|
||||
(cfg.explore_share if proven else 1.0) / len(scouts)
|
||||
if scouts
|
||||
else 0.0
|
||||
)
|
||||
coefficient = None
|
||||
out.append({
|
||||
"signature": cand.signature,
|
||||
"hops": [
|
||||
{
|
||||
"node_id": getattr(n, "node_id", "?"),
|
||||
"shard": f"{getattr(n, 'shard_start', '?')}-{getattr(n, 'shard_end', '?')}",
|
||||
"endpoint": getattr(n, "endpoint", "?"),
|
||||
}
|
||||
for n in cand.nodes
|
||||
],
|
||||
"tps": r["tps"],
|
||||
"coefficient": coefficient,
|
||||
"expected_share": round(share, 4),
|
||||
"samples": r["samples"],
|
||||
"weight": r["weight"],
|
||||
"status": r["status"],
|
||||
"prior_tps": round(cand.prior_tps, 4),
|
||||
})
|
||||
out.sort(key=lambda r: (-(r["tps"] or 0.0), -r["prior_tps"]))
|
||||
return out
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,319 +1,449 @@
|
||||
"""Dashboard user accounts: registration, login, roles, API keys, usage.
|
||||
|
||||
Unit tests for AccountStore plus HTTP integration on the tracker:
|
||||
register/login/logout, per-account balance and usage, API-key lifecycle
|
||||
(revoked keys rejected by the OpenAI proxy), and the admin listing.
|
||||
"""
|
||||
|
||||
import json
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
import pytest
|
||||
|
||||
from meshnet_tracker.accounts import AccountStore
|
||||
from meshnet_tracker.auth import sign_hive_request
|
||||
from meshnet_tracker.billing import BillingLedger
|
||||
from meshnet_tracker.server import TrackerServer
|
||||
|
||||
HIVE_SECRET = "test-hive-secret"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- unit tests
|
||||
|
||||
|
||||
def test_first_account_is_admin_then_users():
|
||||
store = AccountStore()
|
||||
first = store.register(email="admin@example.com", password="secret-123")
|
||||
second = store.register(email="user@example.com", password="secret-123")
|
||||
assert first["role"] == "admin"
|
||||
assert second["role"] == "user"
|
||||
|
||||
|
||||
def test_register_requires_email_or_wallet_and_password_length():
|
||||
store = AccountStore()
|
||||
with pytest.raises(ValueError, match="email or a wallet"):
|
||||
store.register(password="secret-123")
|
||||
with pytest.raises(ValueError, match="invalid email"):
|
||||
store.register(email="not-an-email", password="secret-123")
|
||||
with pytest.raises(ValueError, match="at least 8"):
|
||||
store.register(email="a@b.co", password="short")
|
||||
|
||||
|
||||
def test_register_rejects_duplicate_identifiers():
|
||||
store = AccountStore()
|
||||
store.register(email="dup@example.com", password="secret-123")
|
||||
with pytest.raises(ValueError, match="already exists"):
|
||||
store.register(email="DUP@example.com", password="other-secret")
|
||||
|
||||
|
||||
def test_login_by_email_or_wallet():
|
||||
store = AccountStore()
|
||||
account = store.register(
|
||||
email="both@example.com", wallet="WalletXYZ", password="secret-123"
|
||||
)
|
||||
assert store.verify_login("both@example.com", "secret-123")["account_id"] == account["account_id"]
|
||||
assert store.verify_login("WalletXYZ", "secret-123")["account_id"] == account["account_id"]
|
||||
assert store.verify_login("both@example.com", "wrong-password") is None
|
||||
assert store.verify_login("nobody@example.com", "secret-123") is None
|
||||
|
||||
|
||||
def test_sessions_resolve_and_destroy():
|
||||
store = AccountStore()
|
||||
account = store.register(email="s@example.com", password="secret-123")
|
||||
token = store.create_session(account["account_id"])
|
||||
assert store.session_account(token)["account_id"] == account["account_id"]
|
||||
store.destroy_session(token)
|
||||
assert store.session_account(token) is None
|
||||
assert store.session_account("bogus") is None
|
||||
|
||||
|
||||
def test_api_key_lifecycle():
|
||||
store = AccountStore()
|
||||
account = store.register(email="k@example.com", password="secret-123")
|
||||
other = store.register(email="other@example.com", password="secret-123")
|
||||
key = store.create_api_key(account["account_id"])
|
||||
assert key.startswith("sk-mesh-")
|
||||
assert store.keys_for(account["account_id"]) == [key]
|
||||
# someone else's account cannot revoke it
|
||||
assert store.revoke_api_key(other["account_id"], key) is False
|
||||
assert store.revoke_api_key(account["account_id"], key) is True
|
||||
assert store.keys_for(account["account_id"]) == []
|
||||
assert store.is_key_revoked(key)
|
||||
|
||||
|
||||
def test_accounts_persist_across_restart(tmp_path):
|
||||
db = str(tmp_path / "accounts.db")
|
||||
store = AccountStore(db_path=db)
|
||||
account = store.register(email="p@example.com", password="secret-123")
|
||||
key = store.create_api_key(account["account_id"])
|
||||
store.save_to_db()
|
||||
|
||||
reloaded = AccountStore(db_path=db)
|
||||
assert reloaded.verify_login("p@example.com", "secret-123") is not None
|
||||
assert reloaded.keys_for(account["account_id"]) == [key]
|
||||
|
||||
|
||||
def test_account_events_replicate_and_dedupe():
|
||||
leader = AccountStore()
|
||||
follower = AccountStore()
|
||||
account = leader.register(email="r@example.com", password="secret-123")
|
||||
key = leader.create_api_key(account["account_id"])
|
||||
leader.revoke_api_key(account["account_id"], key)
|
||||
|
||||
events, cursor = leader.events_since(0)
|
||||
assert follower.apply_events(events) == len(events)
|
||||
assert follower.apply_events(events) == 0 # replay is a no-op
|
||||
assert follower.verify_login("r@example.com", "secret-123") is not None
|
||||
assert follower.is_key_revoked(key)
|
||||
more, _ = leader.events_since(cursor)
|
||||
assert more == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------- HTTP integration
|
||||
|
||||
|
||||
def _call(url, method="GET", body=None, token=None):
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if token:
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
data = json.dumps(body).encode() if body is not None else None
|
||||
req = urllib.request.Request(url, data=data, headers=headers, method=method)
|
||||
with urllib.request.urlopen(req) as r:
|
||||
return json.loads(r.read())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def account_tracker():
|
||||
"""Tracker with credit features pinned OFF (defaults are devnet-friendly 1.0)."""
|
||||
ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02)
|
||||
tracker = TrackerServer(
|
||||
billing=ledger,
|
||||
accounts=AccountStore(),
|
||||
hive_secret=HIVE_SECRET,
|
||||
starting_credit=0.0,
|
||||
devnet_topup_amount=0.0,
|
||||
)
|
||||
port = tracker.start()
|
||||
yield f"http://127.0.0.1:{port}", ledger
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_register_login_and_account_view(account_tracker):
|
||||
url, _ = account_tracker
|
||||
reg = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "admin@example.com", "password": "secret-123"})
|
||||
assert reg["account"]["role"] == "admin"
|
||||
assert reg["api_key"].startswith("sk-mesh-")
|
||||
assert reg["session_token"]
|
||||
|
||||
login = _call(f"{url}/v1/auth/login", "POST",
|
||||
{"identifier": "admin@example.com", "password": "secret-123"})
|
||||
me = _call(f"{url}/v1/account", token=login["session_token"])
|
||||
assert me["account"]["email"] == "admin@example.com"
|
||||
assert me["api_keys"] == [reg["api_key"]]
|
||||
assert me["total_balance"] == pytest.approx(0.0)
|
||||
assert me["usage"]["requests"] == 0
|
||||
|
||||
|
||||
def test_bad_credentials_and_missing_session_are_401(account_tracker):
|
||||
url, _ = account_tracker
|
||||
_call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "a@example.com", "password": "secret-123"})
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"{url}/v1/auth/login", "POST",
|
||||
{"identifier": "a@example.com", "password": "wrong-pass"})
|
||||
assert exc_info.value.code == 401
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"{url}/v1/account")
|
||||
assert exc_info.value.code == 401
|
||||
|
||||
|
||||
def test_key_create_revoke_and_revoked_key_rejected_by_proxy(account_tracker):
|
||||
url, _ = account_tracker
|
||||
reg = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "k@example.com", "password": "secret-123"})
|
||||
token = reg["session_token"]
|
||||
|
||||
new_key = _call(f"{url}/v1/account/keys", "POST", {}, token=token)["api_key"]
|
||||
me = _call(f"{url}/v1/account", token=token)
|
||||
assert sorted(me["api_keys"]) == sorted([reg["api_key"], new_key])
|
||||
|
||||
_call(f"{url}/v1/account/keys/revoke", "POST", {"api_key": new_key}, token=token)
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"{url}/v1/chat/completions", "POST",
|
||||
{"model": "any", "messages": []}, token=new_key)
|
||||
assert exc_info.value.code == 401
|
||||
assert "revoked" in exc_info.value.read().decode()
|
||||
|
||||
|
||||
def test_admin_listing_requires_admin_role(account_tracker):
|
||||
url, _ = account_tracker
|
||||
admin = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "admin@example.com", "password": "secret-123"})
|
||||
user = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"wallet": "WalletUser1", "password": "secret-123"})
|
||||
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"{url}/v1/admin/accounts", token=user["session_token"])
|
||||
assert exc_info.value.code == 403
|
||||
|
||||
listing = _call(f"{url}/v1/admin/accounts", token=admin["session_token"])
|
||||
accounts = listing["accounts"]
|
||||
assert len(accounts) == 2
|
||||
assert accounts[0]["role"] == "admin"
|
||||
assert accounts[1]["wallet"] == "WalletUser1"
|
||||
assert "balances" in accounts[0]
|
||||
|
||||
|
||||
def test_accounts_gossip_endpoint_applies_events(account_tracker):
|
||||
url, _ = account_tracker
|
||||
peer = AccountStore()
|
||||
peer.register(email="remote@example.com", password="secret-123")
|
||||
events, _ = peer.events_since(0)
|
||||
body = json.dumps({"events": events}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"{url}/v1/accounts/gossip", data=body,
|
||||
headers={"Content-Type": "application/json", **sign_hive_request(HIVE_SECRET, body)},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req) as r:
|
||||
result = json.loads(r.read())
|
||||
assert result["applied"] == len(events)
|
||||
login = _call(f"{url}/v1/auth/login", "POST",
|
||||
{"identifier": "remote@example.com", "password": "secret-123"})
|
||||
assert login["account"]["email"] == "remote@example.com"
|
||||
|
||||
|
||||
def test_accounts_endpoints_404_when_disabled():
|
||||
tracker = TrackerServer() # no accounts, no billing
|
||||
port = tracker.start()
|
||||
try:
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"http://127.0.0.1:{port}/v1/auth/register", "POST",
|
||||
{"email": "x@example.com", "password": "secret-123"})
|
||||
assert exc_info.value.code == 404
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
# ------------------------------------------- US-039/US-040: credit and top-up
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def funded_tracker():
|
||||
"""Tracker with Caller Credit and the devnet top-up faucet enabled."""
|
||||
ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02)
|
||||
tracker = TrackerServer(
|
||||
billing=ledger,
|
||||
accounts=AccountStore(),
|
||||
hive_secret=HIVE_SECRET,
|
||||
starting_credit=1.0,
|
||||
devnet_topup_amount=10.0,
|
||||
)
|
||||
port = tracker.start()
|
||||
yield f"http://127.0.0.1:{port}", ledger
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_caller_credit_granted_once_per_account(funded_tracker):
|
||||
url, ledger = funded_tracker
|
||||
reg = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "c@example.com", "password": "secret-123"})
|
||||
token = reg["session_token"]
|
||||
first_key = reg["api_key"]
|
||||
assert ledger.get_client_balance(first_key) == pytest.approx(1.0)
|
||||
|
||||
# A second key never re-grants — not even after revoking the first.
|
||||
second = _call(f"{url}/v1/account/keys", "POST", {}, token=token)
|
||||
assert second["caller_credit_granted"] is False
|
||||
assert ledger.get_client_balance(second["api_key"]) == pytest.approx(0.0)
|
||||
_call(f"{url}/v1/account/keys/revoke", "POST", {"api_key": first_key}, token=token)
|
||||
third = _call(f"{url}/v1/account/keys", "POST", {}, token=token)
|
||||
assert third["caller_credit_granted"] is False
|
||||
assert ledger.get_client_balance(third["api_key"]) == pytest.approx(0.0)
|
||||
|
||||
|
||||
def test_unknown_bearer_key_rejected_by_proxy(funded_tracker):
|
||||
url, ledger = funded_tracker
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"{url}/v1/chat/completions", "POST",
|
||||
{"model": "any", "messages": []}, token="sk-mesh-made-up-key")
|
||||
assert exc_info.value.code == 401
|
||||
assert "unknown API key" in exc_info.value.read().decode()
|
||||
# The invented key must not have become a billable client.
|
||||
assert ledger.get_client_balance("sk-mesh-made-up-key") == pytest.approx(0.0)
|
||||
|
||||
|
||||
def test_devnet_topup_credits_own_key_only(funded_tracker):
|
||||
url, ledger = funded_tracker
|
||||
owner = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "own@example.com", "password": "secret-123"})
|
||||
other = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "oth@example.com", "password": "secret-123"})
|
||||
|
||||
me = _call(f"{url}/v1/account", token=owner["session_token"])
|
||||
assert me["topup_amount"] == pytest.approx(10.0)
|
||||
|
||||
result = _call(f"{url}/v1/account/topup", "POST",
|
||||
{"api_key": owner["api_key"]}, token=owner["session_token"])
|
||||
assert result["credited"] == pytest.approx(10.0)
|
||||
assert result["balance"] == pytest.approx(11.0) # 1.0 caller credit + 10.0
|
||||
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"{url}/v1/account/topup", "POST",
|
||||
{"api_key": owner["api_key"]}, token=other["session_token"])
|
||||
assert exc_info.value.code == 403
|
||||
assert ledger.get_client_balance(owner["api_key"]) == pytest.approx(11.0)
|
||||
|
||||
|
||||
def test_topup_404_when_disabled(account_tracker):
|
||||
url, _ = account_tracker
|
||||
reg = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "t@example.com", "password": "secret-123"})
|
||||
me = _call(f"{url}/v1/account", token=reg["session_token"])
|
||||
assert me["topup_amount"] == pytest.approx(0.0)
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"{url}/v1/account/topup", "POST",
|
||||
{"api_key": reg["api_key"]}, token=reg["session_token"])
|
||||
assert exc_info.value.code == 404
|
||||
"""Dashboard user accounts: registration, login, roles, API keys, usage.
|
||||
|
||||
Unit tests for AccountStore plus HTTP integration on the tracker:
|
||||
register/login/logout, per-account balance and usage, API-key lifecycle
|
||||
(revoked keys rejected by the OpenAI proxy), and the admin listing.
|
||||
"""
|
||||
|
||||
import http.cookies
|
||||
import json
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
import pytest
|
||||
|
||||
from meshnet_tracker.accounts import AccountStore
|
||||
from meshnet_tracker.auth import sign_hive_request
|
||||
from meshnet_tracker.billing import BillingLedger
|
||||
from meshnet_tracker.server import TrackerServer
|
||||
|
||||
HIVE_SECRET = "test-hive-secret"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- unit tests
|
||||
|
||||
|
||||
def test_first_account_is_admin_then_users():
|
||||
store = AccountStore()
|
||||
first = store.register(email="admin@example.com", password="secret-123")
|
||||
second = store.register(email="user@example.com", password="secret-123")
|
||||
assert first["role"] == "admin"
|
||||
assert second["role"] == "user"
|
||||
|
||||
|
||||
def test_register_requires_email_or_wallet_and_password_length():
|
||||
store = AccountStore()
|
||||
with pytest.raises(ValueError, match="email or a wallet"):
|
||||
store.register(password="secret-123")
|
||||
with pytest.raises(ValueError, match="invalid email"):
|
||||
store.register(email="not-an-email", password="secret-123")
|
||||
with pytest.raises(ValueError, match="at least 8"):
|
||||
store.register(email="a@b.co", password="short")
|
||||
|
||||
|
||||
def test_register_rejects_duplicate_identifiers():
|
||||
store = AccountStore()
|
||||
store.register(email="dup@example.com", password="secret-123")
|
||||
with pytest.raises(ValueError, match="already exists"):
|
||||
store.register(email="DUP@example.com", password="other-secret")
|
||||
|
||||
|
||||
def test_register_and_update_nickname():
|
||||
store = AccountStore()
|
||||
account = store.register(
|
||||
email="nick@example.com",
|
||||
password="secret-123",
|
||||
nickname=" Alpha ",
|
||||
)
|
||||
assert account["nickname"] == "Alpha"
|
||||
updated = store.update_profile(account["account_id"], nickname="Beta")
|
||||
assert updated["nickname"] == "Beta"
|
||||
cleared = store.update_profile(account["account_id"], nickname=None)
|
||||
assert cleared["nickname"] is None
|
||||
|
||||
|
||||
def test_nickname_replicates_across_stores():
|
||||
leader = AccountStore()
|
||||
follower = AccountStore()
|
||||
account = leader.register(
|
||||
email="nick@example.com",
|
||||
password="secret-123",
|
||||
nickname="HiveNick",
|
||||
)
|
||||
leader.update_profile(account["account_id"], nickname="Renamed")
|
||||
events, _ = leader.events_since(0)
|
||||
follower.apply_events(events)
|
||||
view = follower.get_account(account["account_id"])
|
||||
assert view is not None
|
||||
assert view["nickname"] == "Renamed"
|
||||
|
||||
|
||||
def test_login_by_email_or_wallet():
|
||||
store = AccountStore()
|
||||
account = store.register(
|
||||
email="both@example.com", wallet="WalletXYZ", password="secret-123"
|
||||
)
|
||||
assert store.verify_login("both@example.com", "secret-123")["account_id"] == account["account_id"]
|
||||
assert store.verify_login("WalletXYZ", "secret-123")["account_id"] == account["account_id"]
|
||||
assert store.verify_login("both@example.com", "wrong-password") is None
|
||||
assert store.verify_login("nobody@example.com", "secret-123") is None
|
||||
|
||||
|
||||
def test_sessions_resolve_and_destroy():
|
||||
store = AccountStore()
|
||||
account = store.register(email="s@example.com", password="secret-123")
|
||||
token = store.create_session(account["account_id"])
|
||||
assert store.session_account(token)["account_id"] == account["account_id"]
|
||||
store.destroy_session(token)
|
||||
assert store.session_account(token) is None
|
||||
assert store.session_account("bogus") is None
|
||||
|
||||
|
||||
def test_sessions_persist_across_restart(tmp_path):
|
||||
db = str(tmp_path / "accounts.db")
|
||||
store = AccountStore(db_path=db)
|
||||
account = store.register(email="cookie@example.com", password="secret-123")
|
||||
token = store.create_session(account["account_id"])
|
||||
store.save_to_db()
|
||||
|
||||
reloaded = AccountStore(db_path=db)
|
||||
assert reloaded.session_account(token)["account_id"] == account["account_id"]
|
||||
|
||||
|
||||
def test_api_key_lifecycle():
|
||||
store = AccountStore()
|
||||
account = store.register(email="k@example.com", password="secret-123")
|
||||
other = store.register(email="other@example.com", password="secret-123")
|
||||
key = store.create_api_key(account["account_id"])
|
||||
assert key.startswith("sk-mesh-")
|
||||
assert store.keys_for(account["account_id"]) == [key]
|
||||
# someone else's account cannot revoke it
|
||||
assert store.revoke_api_key(other["account_id"], key) is False
|
||||
assert store.revoke_api_key(account["account_id"], key) is True
|
||||
assert store.keys_for(account["account_id"]) == []
|
||||
assert store.is_key_revoked(key)
|
||||
|
||||
|
||||
def test_accounts_persist_across_restart(tmp_path):
|
||||
db = str(tmp_path / "accounts.db")
|
||||
store = AccountStore(db_path=db)
|
||||
account = store.register(email="p@example.com", password="secret-123")
|
||||
key = store.create_api_key(account["account_id"])
|
||||
store.save_to_db()
|
||||
|
||||
reloaded = AccountStore(db_path=db)
|
||||
assert reloaded.verify_login("p@example.com", "secret-123") is not None
|
||||
assert reloaded.keys_for(account["account_id"]) == [key]
|
||||
|
||||
|
||||
def test_account_events_replicate_and_dedupe():
|
||||
leader = AccountStore()
|
||||
follower = AccountStore()
|
||||
account = leader.register(email="r@example.com", password="secret-123")
|
||||
key = leader.create_api_key(account["account_id"])
|
||||
leader.revoke_api_key(account["account_id"], key)
|
||||
|
||||
events, cursor = leader.events_since(0)
|
||||
assert follower.apply_events(events) == len(events)
|
||||
assert follower.apply_events(events) == 0 # replay is a no-op
|
||||
assert follower.verify_login("r@example.com", "secret-123") is not None
|
||||
assert follower.is_key_revoked(key)
|
||||
more, _ = leader.events_since(cursor)
|
||||
assert more == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------- HTTP integration
|
||||
|
||||
|
||||
def _call(url, method="GET", body=None, token=None):
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if token:
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
data = json.dumps(body).encode() if body is not None else None
|
||||
req = urllib.request.Request(url, data=data, headers=headers, method=method)
|
||||
with urllib.request.urlopen(req) as r:
|
||||
return json.loads(r.read())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def account_tracker():
|
||||
"""Tracker with credit features pinned OFF (defaults are devnet-friendly 1.0)."""
|
||||
ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02)
|
||||
tracker = TrackerServer(
|
||||
billing=ledger,
|
||||
accounts=AccountStore(),
|
||||
hive_secret=HIVE_SECRET,
|
||||
starting_credit=0.0,
|
||||
devnet_topup_amount=0.0,
|
||||
)
|
||||
port = tracker.start()
|
||||
yield f"http://127.0.0.1:{port}", ledger
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_register_login_and_account_view(account_tracker):
|
||||
url, _ = account_tracker
|
||||
reg = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "admin@example.com", "password": "secret-123"})
|
||||
assert reg["account"]["role"] == "admin"
|
||||
assert reg["api_key"].startswith("sk-mesh-")
|
||||
assert reg["session_token"]
|
||||
|
||||
login = _call(f"{url}/v1/auth/login", "POST",
|
||||
{"identifier": "admin@example.com", "password": "secret-123"})
|
||||
me = _call(f"{url}/v1/account", token=login["session_token"])
|
||||
assert me["account"]["email"] == "admin@example.com"
|
||||
assert me["api_keys"] == [reg["api_key"]]
|
||||
assert me["total_balance"] == pytest.approx(0.0)
|
||||
assert me["usage"]["requests"] == 0
|
||||
assert "records" not in me["usage"]
|
||||
assert "recent" not in me["usage"]
|
||||
|
||||
|
||||
def test_account_usage_endpoint_returns_records(account_tracker):
|
||||
url, ledger = account_tracker
|
||||
reg = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "usage@example.com", "password": "secret-123"})
|
||||
ledger.charge_request(reg["api_key"], "test-model", total_tokens=42, node_work=[("wallet-1", 1)])
|
||||
usage = _call(f"{url}/v1/account/usage", token=reg["session_token"])
|
||||
assert usage["requests"] == 1
|
||||
assert usage["total_tokens"] == 42
|
||||
assert len(usage["records"]) == 1
|
||||
assert usage["records"][0]["model"] == "test-model"
|
||||
|
||||
|
||||
def test_account_nickname_register_and_profile_update(account_tracker):
|
||||
url, _ = account_tracker
|
||||
reg = _call(f"{url}/v1/auth/register", "POST", {
|
||||
"email": "nick@example.com",
|
||||
"password": "secret-123",
|
||||
"nickname": "Operator",
|
||||
})
|
||||
assert reg["account"]["nickname"] == "Operator"
|
||||
|
||||
updated = _call(
|
||||
f"{url}/v1/account/profile",
|
||||
"POST",
|
||||
{"nickname": "Renamed"},
|
||||
token=reg["session_token"],
|
||||
)
|
||||
assert updated["account"]["nickname"] == "Renamed"
|
||||
|
||||
me = _call(f"{url}/v1/account", token=reg["session_token"])
|
||||
assert me["account"]["nickname"] == "Renamed"
|
||||
|
||||
|
||||
def test_login_sets_cookie_and_cookie_auth_survives_tracker_restart(tmp_path):
|
||||
accounts_db = str(tmp_path / "accounts.db")
|
||||
tracker = TrackerServer(
|
||||
billing=BillingLedger(starting_credit=0.0, default_price_per_1k=0.02),
|
||||
accounts_db=accounts_db,
|
||||
starting_credit=0.0,
|
||||
devnet_topup_amount=0.0,
|
||||
)
|
||||
port = tracker.start()
|
||||
url = f"http://127.0.0.1:{port}"
|
||||
try:
|
||||
_call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "cookie-http@example.com", "password": "secret-123"})
|
||||
req = urllib.request.Request(
|
||||
f"{url}/v1/auth/login",
|
||||
data=json.dumps({
|
||||
"identifier": "cookie-http@example.com",
|
||||
"password": "secret-123",
|
||||
}).encode(),
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req) as r:
|
||||
assert json.loads(r.read())["session_token"]
|
||||
cookie_header = r.headers["Set-Cookie"]
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
cookie = http.cookies.SimpleCookie(cookie_header)
|
||||
session_cookie = cookie["meshnet_session"].OutputString()
|
||||
|
||||
restarted = TrackerServer(
|
||||
billing=BillingLedger(starting_credit=0.0, default_price_per_1k=0.02),
|
||||
accounts_db=accounts_db,
|
||||
starting_credit=0.0,
|
||||
devnet_topup_amount=0.0,
|
||||
)
|
||||
restarted_port = restarted.start()
|
||||
restarted_url = f"http://127.0.0.1:{restarted_port}"
|
||||
try:
|
||||
req = urllib.request.Request(
|
||||
f"{restarted_url}/v1/account",
|
||||
headers={"Cookie": session_cookie},
|
||||
method="GET",
|
||||
)
|
||||
with urllib.request.urlopen(req) as r:
|
||||
me = json.loads(r.read())
|
||||
finally:
|
||||
restarted.stop()
|
||||
|
||||
assert me["account"]["email"] == "cookie-http@example.com"
|
||||
|
||||
|
||||
def test_bad_credentials_and_missing_session_are_401(account_tracker):
|
||||
url, _ = account_tracker
|
||||
_call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "a@example.com", "password": "secret-123"})
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"{url}/v1/auth/login", "POST",
|
||||
{"identifier": "a@example.com", "password": "wrong-pass"})
|
||||
assert exc_info.value.code == 401
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"{url}/v1/account")
|
||||
assert exc_info.value.code == 401
|
||||
|
||||
|
||||
def test_key_create_revoke_and_revoked_key_rejected_by_proxy(account_tracker):
|
||||
url, _ = account_tracker
|
||||
reg = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "k@example.com", "password": "secret-123"})
|
||||
token = reg["session_token"]
|
||||
|
||||
new_key = _call(f"{url}/v1/account/keys", "POST", {}, token=token)["api_key"]
|
||||
me = _call(f"{url}/v1/account", token=token)
|
||||
assert sorted(me["api_keys"]) == sorted([reg["api_key"], new_key])
|
||||
|
||||
_call(f"{url}/v1/account/keys/revoke", "POST", {"api_key": new_key}, token=token)
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"{url}/v1/chat/completions", "POST",
|
||||
{"model": "any", "messages": []}, token=new_key)
|
||||
assert exc_info.value.code == 401
|
||||
assert "revoked" in exc_info.value.read().decode()
|
||||
|
||||
|
||||
def test_admin_listing_requires_admin_role(account_tracker):
|
||||
url, _ = account_tracker
|
||||
admin = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "admin@example.com", "password": "secret-123"})
|
||||
user = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"wallet": "WalletUser1", "password": "secret-123"})
|
||||
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"{url}/v1/admin/accounts", token=user["session_token"])
|
||||
assert exc_info.value.code == 403
|
||||
|
||||
listing = _call(f"{url}/v1/admin/accounts", token=admin["session_token"])
|
||||
accounts = listing["accounts"]
|
||||
assert len(accounts) == 2
|
||||
assert accounts[0]["role"] == "admin"
|
||||
assert accounts[1]["wallet"] == "WalletUser1"
|
||||
assert "balances" in accounts[0]
|
||||
|
||||
|
||||
def test_accounts_gossip_endpoint_applies_events(account_tracker):
|
||||
url, _ = account_tracker
|
||||
peer = AccountStore()
|
||||
peer.register(email="remote@example.com", password="secret-123")
|
||||
events, _ = peer.events_since(0)
|
||||
body = json.dumps({"events": events}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"{url}/v1/accounts/gossip", data=body,
|
||||
headers={"Content-Type": "application/json", **sign_hive_request(HIVE_SECRET, body)},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req) as r:
|
||||
result = json.loads(r.read())
|
||||
assert result["applied"] == len(events)
|
||||
login = _call(f"{url}/v1/auth/login", "POST",
|
||||
{"identifier": "remote@example.com", "password": "secret-123"})
|
||||
assert login["account"]["email"] == "remote@example.com"
|
||||
|
||||
|
||||
def test_accounts_endpoints_404_when_disabled():
|
||||
tracker = TrackerServer() # no accounts, no billing
|
||||
port = tracker.start()
|
||||
try:
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"http://127.0.0.1:{port}/v1/auth/register", "POST",
|
||||
{"email": "x@example.com", "password": "secret-123"})
|
||||
assert exc_info.value.code == 404
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
# ------------------------------------------- US-039/US-040: credit and top-up
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def funded_tracker():
|
||||
"""Tracker with Caller Credit and the devnet top-up faucet enabled."""
|
||||
ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02)
|
||||
tracker = TrackerServer(
|
||||
billing=ledger,
|
||||
accounts=AccountStore(),
|
||||
hive_secret=HIVE_SECRET,
|
||||
starting_credit=1.0,
|
||||
devnet_topup_amount=10.0,
|
||||
)
|
||||
port = tracker.start()
|
||||
yield f"http://127.0.0.1:{port}", ledger
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_caller_credit_granted_once_per_account(funded_tracker):
|
||||
url, ledger = funded_tracker
|
||||
reg = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "c@example.com", "password": "secret-123"})
|
||||
token = reg["session_token"]
|
||||
first_key = reg["api_key"]
|
||||
assert ledger.get_client_balance(first_key) == pytest.approx(1.0)
|
||||
|
||||
# A second key never re-grants — not even after revoking the first.
|
||||
second = _call(f"{url}/v1/account/keys", "POST", {}, token=token)
|
||||
assert second["caller_credit_granted"] is False
|
||||
assert ledger.get_client_balance(second["api_key"]) == pytest.approx(0.0)
|
||||
_call(f"{url}/v1/account/keys/revoke", "POST", {"api_key": first_key}, token=token)
|
||||
third = _call(f"{url}/v1/account/keys", "POST", {}, token=token)
|
||||
assert third["caller_credit_granted"] is False
|
||||
assert ledger.get_client_balance(third["api_key"]) == pytest.approx(0.0)
|
||||
|
||||
|
||||
def test_unknown_bearer_key_rejected_by_proxy(funded_tracker):
|
||||
url, ledger = funded_tracker
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"{url}/v1/chat/completions", "POST",
|
||||
{"model": "any", "messages": []}, token="sk-mesh-made-up-key")
|
||||
assert exc_info.value.code == 401
|
||||
assert "unknown API key" in exc_info.value.read().decode()
|
||||
# The invented key must not have become a billable client.
|
||||
assert ledger.get_client_balance("sk-mesh-made-up-key") == pytest.approx(0.0)
|
||||
|
||||
|
||||
def test_devnet_topup_credits_own_key_only(funded_tracker):
|
||||
url, ledger = funded_tracker
|
||||
owner = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "own@example.com", "password": "secret-123"})
|
||||
other = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "oth@example.com", "password": "secret-123"})
|
||||
|
||||
me = _call(f"{url}/v1/account", token=owner["session_token"])
|
||||
assert me["topup_amount"] == pytest.approx(10.0)
|
||||
|
||||
result = _call(f"{url}/v1/account/topup", "POST",
|
||||
{"api_key": owner["api_key"]}, token=owner["session_token"])
|
||||
assert result["credited"] == pytest.approx(10.0)
|
||||
assert result["balance"] == pytest.approx(11.0) # 1.0 caller credit + 10.0
|
||||
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"{url}/v1/account/topup", "POST",
|
||||
{"api_key": owner["api_key"]}, token=other["session_token"])
|
||||
assert exc_info.value.code == 403
|
||||
assert ledger.get_client_balance(owner["api_key"]) == pytest.approx(11.0)
|
||||
|
||||
|
||||
def test_topup_404_when_disabled(account_tracker):
|
||||
url, _ = account_tracker
|
||||
reg = _call(f"{url}/v1/auth/register", "POST",
|
||||
{"email": "t@example.com", "password": "secret-123"})
|
||||
me = _call(f"{url}/v1/account", token=reg["session_token"])
|
||||
assert me["topup_amount"] == pytest.approx(0.0)
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_call(f"{url}/v1/account/topup", "POST",
|
||||
{"api_key": reg["api_key"]}, token=reg["session_token"])
|
||||
assert exc_info.value.code == 404
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -29,10 +29,126 @@ def test_dashboard_served_with_all_panels():
|
||||
for panel in PANELS:
|
||||
assert panel in html
|
||||
assert "<script>" in html # polling client embedded, no build step
|
||||
assert "resolveModelGroup" in html
|
||||
assert "buildModelAliasMap" in html
|
||||
assert "modelAliasKey(raw)" in html
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_dashboard_chat_uses_streaming_fetch():
|
||||
tracker = TrackerServer(billing=BillingLedger())
|
||||
port = tracker.start()
|
||||
try:
|
||||
html = urllib.request.urlopen(
|
||||
f"http://127.0.0.1:{port}/dashboard"
|
||||
).read().decode()
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
assert "stream: true" in html
|
||||
assert ".body.getReader()" in html
|
||||
assert '=== "[DONE]"' in html
|
||||
assert 'console.error("chat stream failed", err)' in html
|
||||
|
||||
|
||||
def test_network_map_includes_node_friendly_name():
|
||||
tracker = TrackerServer()
|
||||
port = tracker.start()
|
||||
try:
|
||||
body = json.dumps({
|
||||
"endpoint": "http://127.0.0.1:9010",
|
||||
"model": "stub-model",
|
||||
"shard_start": 0,
|
||||
"shard_end": 3,
|
||||
"hardware_profile": {},
|
||||
"friendly_name": "Kitchen GPU",
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{port}/v1/nodes/register",
|
||||
data=body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
urllib.request.urlopen(req).read()
|
||||
network = json.loads(
|
||||
urllib.request.urlopen(f"http://127.0.0.1:{port}/v1/network/map").read()
|
||||
)
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
assert network["nodes"][0]["friendly_name"] == "Kitchen GPU"
|
||||
|
||||
|
||||
def test_dashboard_chat_model_selector_shows_health_and_speed():
|
||||
tracker = TrackerServer()
|
||||
port = tracker.start()
|
||||
try:
|
||||
html = urllib.request.urlopen(
|
||||
f"http://127.0.0.1:{port}/dashboard"
|
||||
).read().decode()
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
assert "chatModelHealthHp" in html
|
||||
assert "modelServedCopiesFromMap" in html
|
||||
assert "nodeDisplayName" in html
|
||||
assert "accountDisplayName" in html
|
||||
assert "saveNickname" in html
|
||||
assert "chatModelTypicalTps" in html
|
||||
assert "chatModelOptionLabel" in html
|
||||
assert "findRoutingForModel" in html
|
||||
assert "tok/s" in html
|
||||
assert "toFixed(2)}HP" in html or '${copies(v)}HP' in html
|
||||
|
||||
|
||||
def test_dashboard_chat_sessions_use_delegated_handlers():
|
||||
tracker = TrackerServer()
|
||||
port = tracker.start()
|
||||
try:
|
||||
html = urllib.request.urlopen(
|
||||
f"http://127.0.0.1:{port}/dashboard"
|
||||
).read().decode()
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
assert "bindChatSessionList" in html
|
||||
assert "dataset.sessionId" in html
|
||||
assert "dataset.deleteSession" in html
|
||||
assert '[data-session-id]' in html
|
||||
assert 'onclick="selectChatSession(' not in html
|
||||
assert 'onclick="deleteChatSession(' not in html
|
||||
|
||||
|
||||
def test_dashboard_incremental_refresh_helpers():
|
||||
tracker = TrackerServer()
|
||||
port = tracker.start()
|
||||
try:
|
||||
html = urllib.request.urlopen(
|
||||
f"http://127.0.0.1:{port}/dashboard"
|
||||
).read().decode()
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
assert "renderIfChanged" in html
|
||||
assert "syncKeyedList" in html
|
||||
assert "refreshBlocked" in html
|
||||
assert "patchAccountPanelView" in html
|
||||
assert "buildAccountPanelShell" in html
|
||||
assert "refreshActiveTab" in html
|
||||
assert "TAB_FETCHERS" in html
|
||||
assert "loadAccountSummary" in html
|
||||
assert "loadAccountUsage" in html
|
||||
assert "fetchOverviewTab" in html
|
||||
assert "pollCallWallIfIdle" in html
|
||||
assert "pendingChatModelRefresh" in html
|
||||
assert 'renderChatHistory(true)' in html
|
||||
assert "renderChatHistory();" not in html
|
||||
assert "refreshIfIdle" not in html
|
||||
assert "refreshAccountIfIdle" not in html
|
||||
assert "setInterval(refreshIfIdle" not in html
|
||||
|
||||
|
||||
def test_dashboard_served_by_follower():
|
||||
"""A tracker that is not the leader (unreachable peers → never elected)
|
||||
still serves the dashboard from its own replicated state."""
|
||||
|
||||
306
tests/test_dynamic_routing.py
Normal file
306
tests/test_dynamic_routing.py
Normal file
@@ -0,0 +1,306 @@
|
||||
"""ADR-0021: dynamic bandit-style route selection with learned statistics."""
|
||||
|
||||
import http.server
|
||||
import json
|
||||
import random
|
||||
import threading
|
||||
import types
|
||||
import urllib.request
|
||||
|
||||
from meshnet_tracker.routing_stats import (
|
||||
RouteCandidate,
|
||||
RouteStatsStore,
|
||||
RoutingConfig,
|
||||
choose_route,
|
||||
route_signature,
|
||||
route_table,
|
||||
)
|
||||
from meshnet_tracker.server import TrackerServer, _enumerate_routes
|
||||
|
||||
|
||||
def _post_json(url: str, payload: dict) -> dict:
|
||||
req = urllib.request.Request(
|
||||
url,
|
||||
data=json.dumps(payload).encode(),
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=10.0) as resp:
|
||||
return json.loads(resp.read())
|
||||
|
||||
|
||||
def _get_json(url: str) -> dict:
|
||||
with urllib.request.urlopen(url, timeout=10.0) as resp:
|
||||
return json.loads(resp.read())
|
||||
|
||||
|
||||
def _fake_node(node_id, shard_start, shard_end, benchmark=100.0, endpoint=None):
|
||||
return types.SimpleNamespace(
|
||||
node_id=node_id,
|
||||
endpoint=endpoint or f"http://{node_id}:7000",
|
||||
model="qwen3.6-35b-a3b",
|
||||
hf_repo="unsloth/Qwen3.6-35B-A3B",
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
num_layers=40,
|
||||
benchmark_tokens_per_sec=benchmark,
|
||||
model_tokens_per_sec={},
|
||||
queue_depth=0,
|
||||
proxy_inflight=0,
|
||||
wallet_address=None,
|
||||
relay_addr=None,
|
||||
)
|
||||
|
||||
|
||||
# ---- RouteStatsStore ----------------------------------------------------
|
||||
|
||||
|
||||
def test_route_stats_sample_becomes_proven_and_decays():
|
||||
store = RouteStatsStore(RoutingConfig(stats_half_life_seconds=100.0))
|
||||
sig = "m|a[0-39]"
|
||||
assert store.snapshot(sig, "m", now=0.0)["status"] == "unsampled"
|
||||
assert store.record_sample("m", sig, tokens=100, elapsed_seconds=10.0, now=0.0)
|
||||
snap = store.snapshot(sig, "m", now=1.0)
|
||||
assert snap["status"] == "proven"
|
||||
assert snap["tps"] == 10.0
|
||||
# After many half-lives the sample mass decays below the proven threshold.
|
||||
assert store.snapshot(sig, "m", now=1000.0)["status"] == "decayed"
|
||||
|
||||
|
||||
def test_route_stats_rejects_near_empty_samples():
|
||||
store = RouteStatsStore(RoutingConfig(min_sample_tokens=8))
|
||||
assert not store.record_sample("m", "sig", tokens=3, elapsed_seconds=1.0)
|
||||
assert store.snapshot("sig", "m")["samples"] == 0
|
||||
|
||||
|
||||
def test_route_stats_epoch_bump_marks_stale():
|
||||
store = RouteStatsStore()
|
||||
sig = "m|a[0-39]"
|
||||
store.record_sample("m", sig, tokens=100, elapsed_seconds=10.0, now=0.0)
|
||||
assert store.snapshot(sig, "m", now=1.0)["status"] == "proven"
|
||||
store.bump_epoch(["m"])
|
||||
snap = store.snapshot(sig, "m", now=1.0)
|
||||
assert snap["status"] == "stale"
|
||||
assert snap["tps"] == 10.0 # EWMA kept as a prior for display
|
||||
# A fresh sample under the new epoch re-proves the route.
|
||||
store.record_sample("m", sig, tokens=100, elapsed_seconds=10.0, now=2.0)
|
||||
assert store.snapshot(sig, "m", now=3.0)["status"] == "proven"
|
||||
|
||||
|
||||
def test_route_stats_ewma_averages_samples():
|
||||
store = RouteStatsStore(RoutingConfig(stats_half_life_seconds=1e9))
|
||||
sig = "m|a"
|
||||
store.record_sample("m", sig, tokens=100, elapsed_seconds=10.0, now=0.0) # 10 tps
|
||||
store.record_sample("m", sig, tokens=200, elapsed_seconds=10.0, now=1.0) # 20 tps
|
||||
snap = store.snapshot(sig, "m", now=2.0)
|
||||
assert 14.9 < snap["tps"] < 15.1
|
||||
|
||||
|
||||
# ---- choose_route --------------------------------------------------------
|
||||
|
||||
|
||||
def _candidates_two_routes():
|
||||
fast = RouteCandidate(nodes=[], signature="m|fast", prior_tps=100.0)
|
||||
slow = RouteCandidate(nodes=[], signature="m|slow", prior_tps=50.0)
|
||||
return fast, slow
|
||||
|
||||
|
||||
def test_choose_route_without_samples_is_deterministic_best_prior():
|
||||
store = RouteStatsStore()
|
||||
fast, slow = _candidates_two_routes()
|
||||
for _ in range(20):
|
||||
picked, decision = choose_route([slow, fast], store, "m", rng=random.Random(7))
|
||||
assert picked is fast
|
||||
assert decision["mode"] == "scout"
|
||||
|
||||
|
||||
def test_choose_route_traffic_proportional_to_tps():
|
||||
store = RouteStatsStore(RoutingConfig(stats_half_life_seconds=1e9))
|
||||
fast, slow = _candidates_two_routes()
|
||||
now = 0.0
|
||||
for _ in range(5):
|
||||
now += 1.0
|
||||
store.record_sample("m", fast.signature, tokens=150, elapsed_seconds=10.0, now=now)
|
||||
store.record_sample("m", slow.signature, tokens=100, elapsed_seconds=10.0, now=now)
|
||||
rng = random.Random(42)
|
||||
picks = {"m|fast": 0, "m|slow": 0}
|
||||
for _ in range(4000):
|
||||
picked, decision = choose_route([fast, slow], store, "m", rng=rng, now=now)
|
||||
assert decision["mode"] == "exploit"
|
||||
picks[picked.signature] += 1
|
||||
share = picks["m|fast"] / 4000
|
||||
# 15 tps vs 10 tps at alpha=1 → expected fast share 0.6
|
||||
assert 0.55 < share < 0.65
|
||||
|
||||
|
||||
def test_choose_route_scouts_unproven_routes_at_explore_share():
|
||||
store = RouteStatsStore(RoutingConfig(explore_share=0.25, stats_half_life_seconds=1e9))
|
||||
fast, slow = _candidates_two_routes()
|
||||
now = 1.0
|
||||
store.record_sample("m", fast.signature, tokens=150, elapsed_seconds=10.0, now=now)
|
||||
rng = random.Random(11)
|
||||
scouted = 0
|
||||
for _ in range(4000):
|
||||
picked, decision = choose_route([fast, slow], store, "m", rng=rng, now=now)
|
||||
if decision["mode"] == "scout":
|
||||
scouted += 1
|
||||
assert picked is slow
|
||||
assert 0.20 < scouted / 4000 < 0.30
|
||||
|
||||
|
||||
# ---- _enumerate_routes ---------------------------------------------------
|
||||
|
||||
|
||||
def test_enumerate_routes_mixed_topology_yields_both_routes():
|
||||
gpu = _fake_node("gpu", 0, 21, benchmark=11000.0)
|
||||
cpu = _fake_node("cpu", 0, 39, benchmark=425.0)
|
||||
candidates = _enumerate_routes([gpu, cpu], 0, 39, model="qwen3.6-35b-a3b")
|
||||
signatures = {c.signature for c in candidates}
|
||||
assert signatures == {
|
||||
route_signature("qwen3.6-35b-a3b", [gpu, cpu]),
|
||||
route_signature("qwen3.6-35b-a3b", [cpu]),
|
||||
}
|
||||
hybrid = next(c for c in candidates if len(c.nodes) == 2)
|
||||
assert [n.node_id for n in hybrid.nodes] == ["gpu", "cpu"]
|
||||
# Hybrid route's prior is its bottleneck hop, not the fast head.
|
||||
assert hybrid.prior_tps == 425.0
|
||||
|
||||
|
||||
def test_enumerate_routes_requires_head_at_first_layer():
|
||||
tail_only = _fake_node("tail", 22, 39)
|
||||
assert _enumerate_routes([tail_only], 0, 39, model="m") == []
|
||||
|
||||
|
||||
def test_route_table_reports_coefficient_and_share():
|
||||
store = RouteStatsStore(RoutingConfig(explore_share=0.3, stats_half_life_seconds=1e9))
|
||||
fast, slow = _candidates_two_routes()
|
||||
now = 1.0
|
||||
for _ in range(3):
|
||||
store.record_sample("m", fast.signature, tokens=150, elapsed_seconds=10.0, now=now)
|
||||
store.record_sample("m", slow.signature, tokens=100, elapsed_seconds=10.0, now=now)
|
||||
now += 1.0
|
||||
rows = route_table([fast, slow], store, "m", now=now)
|
||||
by_sig = {r["signature"]: r for r in rows}
|
||||
assert by_sig["m|fast"]["coefficient"] == 1.0
|
||||
assert abs(by_sig["m|slow"]["coefficient"] - (10.0 / 15.0)) < 0.01
|
||||
# No scouts → full exploit budget split 0.6 / 0.4.
|
||||
assert abs(by_sig["m|fast"]["expected_share"] - 0.6) < 0.01
|
||||
assert abs(by_sig["m|slow"]["expected_share"] - 0.4) < 0.01
|
||||
|
||||
|
||||
# ---- integration: proxy uses route head + /v1/routing --------------------
|
||||
|
||||
|
||||
def test_proxy_head_is_route_head_and_routing_endpoint_lists_routes():
|
||||
"""Mixed topology (partial head 0-21 + full node 0-39): the proxy target
|
||||
must be the selected route's own head, downstream hops must continue at
|
||||
head.shard_end + 1 (the ADR-0020 flaw), and /v1/routing must list both
|
||||
candidate routes."""
|
||||
|
||||
class ChatHandler(http.server.BaseHTTPRequestHandler):
|
||||
def log_message(self, *args): # noqa: ARG002
|
||||
pass
|
||||
|
||||
def do_POST(self):
|
||||
length = int(self.headers.get("Content-Length", 0))
|
||||
self.rfile.read(length)
|
||||
route_header = self.headers.get("X-Meshnet-Route") or "[]"
|
||||
body = json.dumps({
|
||||
"choices": [{"message": {"role": "assistant", "content": route_header}}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 40},
|
||||
}).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
stubs = []
|
||||
threads = []
|
||||
for _ in range(2):
|
||||
stub = http.server.HTTPServer(("127.0.0.1", 0), ChatHandler)
|
||||
thread = threading.Thread(target=stub.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
stubs.append(stub)
|
||||
threads.append(thread)
|
||||
gpu_stub, cpu_stub = stubs
|
||||
|
||||
tracker = TrackerServer(model_presets={
|
||||
"qwen3.6-35b-a3b": {
|
||||
"layers_start": 0,
|
||||
"layers_end": 39,
|
||||
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
|
||||
"aliases": ["Qwen3.6-35B-A3B"],
|
||||
}
|
||||
})
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
tracker._server.route_rng = random.Random(3)
|
||||
for stub, shard_end, bench in ((gpu_stub, 21, 11000.0), (cpu_stub, 39, 425.0)):
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": f"http://127.0.0.1:{stub.server_address[1]}",
|
||||
"model": "qwen3.6-35b-a3b",
|
||||
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
|
||||
"num_layers": 40,
|
||||
"shard_start": 0,
|
||||
"shard_end": shard_end,
|
||||
"tracker_mode": True,
|
||||
"benchmark_tokens_per_sec": bench,
|
||||
"hardware_profile": {},
|
||||
"score": 1.0},
|
||||
)
|
||||
|
||||
for _ in range(8):
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
|
||||
{"model": "Qwen3.6-35B-A3B",
|
||||
"messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
|
||||
console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
|
||||
routing = _get_json(f"http://127.0.0.1:{tracker_port}/v1/routing")
|
||||
finally:
|
||||
tracker.stop()
|
||||
for stub, thread in zip(stubs, threads):
|
||||
stub.shutdown()
|
||||
stub.server_close()
|
||||
thread.join(timeout=1.0)
|
||||
|
||||
gpu_endpoint = f"http://127.0.0.1:{gpu_stub.server_address[1]}"
|
||||
cpu_endpoint = f"http://127.0.0.1:{cpu_stub.server_address[1]}"
|
||||
selected = [e for e in console["events"] if e["message"] == "proxy route selected"]
|
||||
assert selected
|
||||
for event in selected:
|
||||
fields = event["fields"]
|
||||
nodes = fields["nodes"]
|
||||
# The proxy head must be the route's first hop (ADR-0020 regression).
|
||||
assert fields["head_endpoint"] == nodes[0]["endpoint"]
|
||||
downstream = json.loads(fields["downstream"])
|
||||
if fields["head_endpoint"] == gpu_endpoint:
|
||||
# Partial head: downstream continues at layer 22, never 0.
|
||||
assert downstream == [{"endpoint": cpu_endpoint, "start_layer": 22}]
|
||||
else:
|
||||
assert fields["head_endpoint"] == cpu_endpoint
|
||||
assert downstream == []
|
||||
|
||||
table = routing["models"]["qwen3.6-35b-a3b"]
|
||||
assert len(table["routes"]) == 2
|
||||
sampled = [r for r in table["routes"] if r["samples"] > 0]
|
||||
assert sampled, "completed requests must produce route samples"
|
||||
|
||||
|
||||
def test_endpoint_key_distinguishes_same_port_different_hosts():
|
||||
from meshnet_node.torch_server import _clamp_downstream_hops, _endpoint_key
|
||||
|
||||
assert _endpoint_key("http://192.168.0.20:7000") == "192.168.0.20:7000"
|
||||
assert _endpoint_key("http://192.168.0.179:7000") == "192.168.0.179:7000"
|
||||
assert _endpoint_key("http://192.168.0.20:7000") != _endpoint_key("http://192.168.0.179:7000")
|
||||
|
||||
class Backend:
|
||||
shard_end = 21
|
||||
|
||||
hops = [{"endpoint": "http://192.168.0.179:7000", "start_layer": 0}]
|
||||
assert _clamp_downstream_hops(hops, Backend()) == [
|
||||
{"endpoint": "http://192.168.0.179:7000", "start_layer": 22},
|
||||
]
|
||||
@@ -186,3 +186,24 @@ def test_qwen_preset_prices_apply_to_all_aliases(tmp_path):
|
||||
assert billing.price_for("some/other-model") == pytest.approx(0.02)
|
||||
finally:
|
||||
pass
|
||||
|
||||
|
||||
def test_qwen25_preset_price_is_ten_x_commercial_reference(tmp_path):
|
||||
"""Qwen2.5-0.5B bills at 10× ~$0.20/1M reference ($0.002/1k), not the 0.02 default."""
|
||||
import pytest
|
||||
from meshnet_tracker.server import TrackerServer, _resolve_model_preset, DEFAULT_MODEL_PRESETS
|
||||
|
||||
name, preset = _resolve_model_preset(DEFAULT_MODEL_PRESETS, "Qwen/Qwen2.5-0.5B-Instruct")
|
||||
assert name == "qwen2.5-0.5b-instruct"
|
||||
assert preset["price_per_1k_tokens"] == pytest.approx(0.002)
|
||||
|
||||
tracker = TrackerServer(billing_db=str(tmp_path / "billing.sqlite"))
|
||||
billing = tracker._billing
|
||||
assert billing is not None
|
||||
for key in (
|
||||
"qwen2.5-0.5b",
|
||||
"Qwen2.5-0.5B-Instruct",
|
||||
"Qwen/Qwen2.5-0.5B-Instruct",
|
||||
):
|
||||
assert billing.price_for(key) == pytest.approx(0.002), key
|
||||
assert billing.price_for("unrelated-model") == pytest.approx(0.02)
|
||||
|
||||
384
tests/test_kv_cache_distributed.py
Normal file
384
tests/test_kv_cache_distributed.py
Normal file
@@ -0,0 +1,384 @@
|
||||
"""AH-25: sharded per-node KV cache for distributed generation.
|
||||
|
||||
Covers the SessionCacheStore (TTL + LRU + mismatch handling), the HTTP
|
||||
session protocol (stable session id, O(1) decode payloads, 409 cache-miss
|
||||
fallback, legacy stateless compatibility), and an env-gated golden test that
|
||||
proves cached and stateless distributed generation produce identical tokens
|
||||
on a real two-shard Qwen2.5-0.5B split.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import urllib.request
|
||||
|
||||
import pytest
|
||||
|
||||
from meshnet_node.model_backend import (
|
||||
KVCacheMiss,
|
||||
SessionCacheStore,
|
||||
TailTokenResult,
|
||||
TensorPayload,
|
||||
)
|
||||
from meshnet_node.torch_server import TorchNodeServer
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SessionCacheStore units
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _Clock:
|
||||
def __init__(self) -> None:
|
||||
self.now = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
|
||||
def test_store_lookup_roundtrip_advances_lru():
|
||||
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0, clock=_Clock())
|
||||
store.store("s1", cache=object(), seq_len=6, effective_start=12)
|
||||
entry = store.lookup("s1", expected_seq_len=6, effective_start=12)
|
||||
assert entry.seq_len == 6
|
||||
entry.seq_len += 1
|
||||
assert store.lookup("s1", expected_seq_len=7).seq_len == 7
|
||||
|
||||
|
||||
def test_lookup_unknown_session_raises_cache_miss():
|
||||
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
|
||||
with pytest.raises(KVCacheMiss):
|
||||
store.lookup("nope")
|
||||
|
||||
|
||||
def test_seq_len_mismatch_drops_entry_and_raises():
|
||||
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
|
||||
store.store("s1", cache=object(), seq_len=6, effective_start=0)
|
||||
with pytest.raises(KVCacheMiss):
|
||||
store.lookup("s1", expected_seq_len=9)
|
||||
# Entry must be gone — a poisoned cache is never reused.
|
||||
with pytest.raises(KVCacheMiss):
|
||||
store.lookup("s1")
|
||||
|
||||
|
||||
def test_effective_start_mismatch_raises():
|
||||
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
|
||||
store.store("s1", cache=object(), seq_len=6, effective_start=12)
|
||||
with pytest.raises(KVCacheMiss):
|
||||
store.lookup("s1", effective_start=21)
|
||||
|
||||
|
||||
def test_ttl_expiry_evicts_stale_sessions():
|
||||
clock = _Clock()
|
||||
store = SessionCacheStore(max_sessions=4, ttl_seconds=60.0, clock=clock)
|
||||
store.store("s1", cache=object(), seq_len=6, effective_start=0)
|
||||
clock.now = 61.0
|
||||
with pytest.raises(KVCacheMiss):
|
||||
store.lookup("s1")
|
||||
assert len(store) == 0
|
||||
|
||||
|
||||
def test_lru_eviction_bounds_session_count():
|
||||
clock = _Clock()
|
||||
store = SessionCacheStore(max_sessions=2, ttl_seconds=1000.0, clock=clock)
|
||||
store.store("s1", cache=object(), seq_len=1, effective_start=0)
|
||||
store.store("s2", cache=object(), seq_len=1, effective_start=0)
|
||||
store.lookup("s1") # s1 becomes most recent → s2 is LRU
|
||||
store.store("s3", cache=object(), seq_len=1, effective_start=0)
|
||||
assert len(store) == 2
|
||||
with pytest.raises(KVCacheMiss):
|
||||
store.lookup("s2")
|
||||
store.lookup("s1")
|
||||
store.lookup("s3")
|
||||
|
||||
|
||||
def test_drop_removes_session():
|
||||
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
|
||||
store.store("s1", cache=object(), seq_len=1, effective_start=0)
|
||||
store.drop("s1")
|
||||
with pytest.raises(KVCacheMiss):
|
||||
store.lookup("s1")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HTTP session protocol with fake cached backends
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _ChatTokenizer:
|
||||
eos_token = ""
|
||||
|
||||
def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False):
|
||||
return "debug prompt"
|
||||
|
||||
|
||||
class _CachedHeadBackend:
|
||||
model_id = "fake-model"
|
||||
total_layers = 12
|
||||
is_head = True
|
||||
is_tail = False
|
||||
supports_kv_cache = True
|
||||
tokenizer = _ChatTokenizer()
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.prefills: list[str | None] = []
|
||||
self.decode_calls: list[tuple[int, str]] = []
|
||||
self.released: list[str] = []
|
||||
self._seq: dict[str, int] = {}
|
||||
|
||||
def eos_token_ids(self) -> list[int]:
|
||||
return [99]
|
||||
|
||||
def release_session(self, session_id: str) -> None:
|
||||
self.released.append(session_id)
|
||||
|
||||
def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload:
|
||||
self.prefills.append(session_id)
|
||||
if session_id:
|
||||
self._seq[session_id] = 6
|
||||
return TensorPayload(
|
||||
body=b"\x00" * (1 * 6 * 8 * 2),
|
||||
shape=[1, 6, 8],
|
||||
attention_mask_header=None,
|
||||
position_ids_header=None,
|
||||
)
|
||||
|
||||
def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload:
|
||||
self.decode_calls.append((token_id, session_id))
|
||||
past = self._seq[session_id]
|
||||
self._seq[session_id] = past + 1
|
||||
return TensorPayload(
|
||||
body=b"\x00" * (1 * 1 * 8 * 2),
|
||||
shape=[1, 1, 8],
|
||||
attention_mask_header=None,
|
||||
position_ids_header=None,
|
||||
past_len=past,
|
||||
)
|
||||
|
||||
|
||||
class _CachedTailBackend:
|
||||
model_id = "fake-model"
|
||||
total_layers = 12
|
||||
is_head = False
|
||||
is_tail = True
|
||||
supports_kv_cache = True
|
||||
|
||||
def __init__(self, tokens, miss_on_call: int | None = None) -> None:
|
||||
self._tokens = list(tokens)
|
||||
self.miss_on_call = miss_on_call
|
||||
self.calls: list[dict] = []
|
||||
|
||||
def forward_bytes(
|
||||
self,
|
||||
body,
|
||||
shape,
|
||||
attention_mask_header,
|
||||
position_ids_header,
|
||||
start_layer=None,
|
||||
session_id=None,
|
||||
cache_mode=None,
|
||||
past_len=None,
|
||||
):
|
||||
call_index = len(self.calls)
|
||||
self.calls.append({
|
||||
"session": session_id,
|
||||
"mode": cache_mode,
|
||||
"past_len": past_len,
|
||||
"shape": list(shape),
|
||||
})
|
||||
if self.miss_on_call is not None and call_index == self.miss_on_call:
|
||||
raise KVCacheMiss("session evicted (test)")
|
||||
text, token_id = self._tokens.pop(0)
|
||||
return TailTokenResult(text=text, token_id=token_id)
|
||||
|
||||
|
||||
def _chat_once(head_port: int, tail_port: int, max_tokens: int) -> str:
|
||||
payload = json.dumps({
|
||||
"model": "fake-model",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": max_tokens,
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{head_port}/v1/chat/completions",
|
||||
data=payload,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"X-Meshnet-Route": json.dumps([
|
||||
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 6},
|
||||
]),
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
body = json.loads(resp.read())
|
||||
return body["choices"][0]["message"]["content"]
|
||||
|
||||
|
||||
def test_session_is_stable_and_decode_payloads_are_single_token():
|
||||
head_backend = _CachedHeadBackend()
|
||||
tail_backend = _CachedTailBackend([(" a", 1), (" b", 2), (" c", 3)])
|
||||
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
|
||||
tail = TorchNodeServer(backend=tail_backend)
|
||||
head_port = head.start()
|
||||
tail_port = tail.start()
|
||||
try:
|
||||
content = _chat_once(head_port, tail_port, max_tokens=3)
|
||||
finally:
|
||||
head.stop()
|
||||
tail.stop()
|
||||
|
||||
assert content == " a b c"
|
||||
assert len(tail_backend.calls) == 3
|
||||
# Step 0 is a full-prompt prefill; steps 1+ carry only the new token.
|
||||
assert tail_backend.calls[0]["mode"] == "prefill"
|
||||
assert tail_backend.calls[0]["shape"] == [1, 6, 8]
|
||||
for step, call in enumerate(tail_backend.calls[1:], start=1):
|
||||
assert call["mode"] == "decode"
|
||||
assert call["shape"] == [1, 1, 8]
|
||||
assert call["past_len"] == 6 + (step - 1)
|
||||
# One session id across every step of the generation.
|
||||
sessions = {call["session"] for call in tail_backend.calls}
|
||||
assert len(sessions) == 1
|
||||
session_id = sessions.pop()
|
||||
assert head_backend.prefills == [session_id]
|
||||
assert head_backend.decode_calls == [(1, session_id), (2, session_id)]
|
||||
# Head releases its own session state when the generation ends.
|
||||
assert head_backend.released == [session_id]
|
||||
|
||||
|
||||
def test_eos_token_id_stops_generation():
|
||||
head_backend = _CachedHeadBackend()
|
||||
tail_backend = _CachedTailBackend([(" a", 1), ("", 99)])
|
||||
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
|
||||
tail = TorchNodeServer(backend=tail_backend)
|
||||
head_port = head.start()
|
||||
tail_port = tail.start()
|
||||
try:
|
||||
content = _chat_once(head_port, tail_port, max_tokens=8)
|
||||
finally:
|
||||
head.stop()
|
||||
tail.stop()
|
||||
|
||||
assert content == " a"
|
||||
assert len(tail_backend.calls) == 2
|
||||
|
||||
|
||||
def test_downstream_cache_miss_falls_back_to_full_reprefill():
|
||||
head_backend = _CachedHeadBackend()
|
||||
# Call 1 (the first decode) raises KVCacheMiss → node answers 409 →
|
||||
# head re-prefills the full sequence and keeps generating.
|
||||
tail_backend = _CachedTailBackend(
|
||||
[(" a", 1), (" b", 2), (" c", 3)], miss_on_call=1,
|
||||
)
|
||||
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
|
||||
tail = TorchNodeServer(backend=tail_backend)
|
||||
head_port = head.start()
|
||||
tail_port = tail.start()
|
||||
try:
|
||||
content = _chat_once(head_port, tail_port, max_tokens=3)
|
||||
finally:
|
||||
head.stop()
|
||||
tail.stop()
|
||||
|
||||
assert content == " a b c"
|
||||
modes = [call["mode"] for call in tail_backend.calls]
|
||||
assert modes == ["prefill", "decode", "prefill", "decode"]
|
||||
# Head re-prefilled once, with the same stable session id.
|
||||
assert len(head_backend.prefills) == 2
|
||||
assert len(set(head_backend.prefills)) == 1
|
||||
|
||||
|
||||
def test_kv_head_with_legacy_tail_reprefills_every_step():
|
||||
"""Mixed fleet: tail predates the protocol and returns no token_id."""
|
||||
|
||||
class _LegacyTailBackend:
|
||||
model_id = "fake-model"
|
||||
total_layers = 12
|
||||
is_head = False
|
||||
is_tail = True
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
def forward_bytes(self, body, shape, attention_mask_header,
|
||||
position_ids_header, start_layer=None, **kwargs):
|
||||
self.calls += 1
|
||||
return " x" if self.calls < 3 else ""
|
||||
|
||||
head_backend = _CachedHeadBackend()
|
||||
tail_backend = _LegacyTailBackend()
|
||||
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
|
||||
tail = TorchNodeServer(backend=tail_backend)
|
||||
head_port = head.start()
|
||||
tail_port = tail.start()
|
||||
try:
|
||||
content = _chat_once(head_port, tail_port, max_tokens=5)
|
||||
finally:
|
||||
head.stop()
|
||||
tail.stop()
|
||||
|
||||
assert content == " x x"
|
||||
# No token_id from the tail → every step is a full prefill (legacy cost),
|
||||
# never a decode against a cache the tail doesn't keep.
|
||||
assert head_backend.decode_calls == []
|
||||
assert len(head_backend.prefills) == 3
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Golden test on a real two-shard split (env-gated: loads Qwen2.5-0.5B twice)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_GOLDEN_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
|
||||
requires_real_model = pytest.mark.skipif(
|
||||
os.environ.get("MESHNET_REAL_MODEL_TESTS") != "1",
|
||||
reason="set MESHNET_REAL_MODEL_TESTS=1 to run the real-model golden test",
|
||||
)
|
||||
|
||||
|
||||
@requires_real_model
|
||||
def test_cached_distributed_generation_matches_stateless_golden():
|
||||
pytest.importorskip("torch")
|
||||
from meshnet_node.model_backend import TorchModelShard
|
||||
|
||||
head = TorchModelShard(_GOLDEN_MODEL, 0, 11)
|
||||
tail = TorchModelShard(_GOLDEN_MODEL, 12, 23)
|
||||
steps = 12
|
||||
prompt = head.tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": "Count from 1 to 5."}],
|
||||
add_generation_prompt=True,
|
||||
tokenize=False,
|
||||
)
|
||||
|
||||
# Reference: today's stateless path — re-encode the full sequence each step.
|
||||
stateless_ids: list[int] = []
|
||||
text = prompt
|
||||
for _ in range(steps):
|
||||
payload = head.encode_prompt(text)
|
||||
result = tail.forward_bytes(
|
||||
payload.body, payload.shape,
|
||||
payload.attention_mask_header, payload.position_ids_header,
|
||||
start_layer=12,
|
||||
)
|
||||
stateless_ids.append(result.token_id)
|
||||
text += result.text
|
||||
|
||||
# Cached path: one prefill, then single-token decode steps.
|
||||
session = "golden-session"
|
||||
cached_ids: list[int] = []
|
||||
payload = head.encode_prompt(prompt, session_id=session)
|
||||
result = tail.forward_bytes(
|
||||
payload.body, payload.shape,
|
||||
payload.attention_mask_header, payload.position_ids_header,
|
||||
start_layer=12, session_id=session, cache_mode="prefill",
|
||||
)
|
||||
cached_ids.append(result.token_id)
|
||||
for _ in range(steps - 1):
|
||||
payload = head.encode_next_token(cached_ids[-1], session)
|
||||
assert payload.shape[1] == 1, "decode payload must be a single token"
|
||||
result = tail.forward_bytes(
|
||||
payload.body, payload.shape,
|
||||
None, payload.position_ids_header,
|
||||
start_layer=12, session_id=session, cache_mode="decode",
|
||||
past_len=payload.past_len,
|
||||
)
|
||||
cached_ids.append(result.token_id)
|
||||
|
||||
assert cached_ids == stateless_ids
|
||||
@@ -1118,7 +1118,7 @@ def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, ca
|
||||
assert captured_registration["vram_bytes"] == 6144 * 1024 * 1024
|
||||
assert captured_registration["max_loaded_shards"] == 2
|
||||
output = capsys.readouterr().out
|
||||
assert "Shard: layers 0–23; 24 of 24" in output
|
||||
assert "Shard: layers 0–23 (24 of 24)" in output
|
||||
assert "Node ID: node-test-123" in output
|
||||
|
||||
|
||||
@@ -1646,6 +1646,166 @@ def test_preset_model_startup_honors_pinned_shard_range(tmp_path, monkeypatch):
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_preset_startup_rejects_pinned_shard_above_memory_budget(tmp_path, monkeypatch):
|
||||
"""Pinned layer ranges that exceed the node memory budget fail before model load."""
|
||||
import meshnet_node.startup as startup_mod
|
||||
|
||||
monkeypatch.setattr(
|
||||
startup_mod,
|
||||
"detect_hardware",
|
||||
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 8 * 1024},
|
||||
)
|
||||
|
||||
tracker = TrackerServer(model_presets={
|
||||
"big-model": {
|
||||
"layers_start": 0,
|
||||
"layers_end": 39,
|
||||
"hf_repo": "org/big-model",
|
||||
"bytes_per_layer": {"bfloat16": 2 * 1024 * 1024 * 1024},
|
||||
},
|
||||
})
|
||||
tracker_port = tracker.start()
|
||||
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
||||
try:
|
||||
with pytest.raises(ValueError, match="Pinned shard layers 0–39"):
|
||||
run_startup(
|
||||
tracker_url=tracker_url,
|
||||
model="big-model",
|
||||
shard_start=0,
|
||||
shard_end=39,
|
||||
wallet_path=tmp_path / "wallet.json",
|
||||
cache_dir=tmp_path / "shards",
|
||||
)
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_network_auto_join_clips_oversized_cpu_assignment(tmp_path, monkeypatch, capsys):
|
||||
"""Old trackers may assign too many CPU layers; node clips before model load."""
|
||||
import meshnet_node.startup as startup_mod
|
||||
|
||||
torch_calls: list[dict] = []
|
||||
registrations: list[dict] = []
|
||||
|
||||
class FakeBackend:
|
||||
total_layers = 40
|
||||
|
||||
class FakeTorchNodeServer:
|
||||
def __init__(self, **kwargs):
|
||||
torch_calls.append(kwargs)
|
||||
self.backend = FakeBackend()
|
||||
self.tracker_node_id = None
|
||||
|
||||
def start(self):
|
||||
return 7000
|
||||
|
||||
def stop(self):
|
||||
pass
|
||||
|
||||
oversized_assignment = {
|
||||
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
|
||||
"model": "qwen3.6-35b-a3b",
|
||||
"shard_start": 0,
|
||||
"shard_end": 36,
|
||||
"num_layers": 40,
|
||||
"gap_found": False,
|
||||
"bytes_per_layer": {"bfloat16": 1_797_594_419},
|
||||
"model_sources": [],
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
startup_mod,
|
||||
"detect_hardware",
|
||||
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 79 * 1024},
|
||||
)
|
||||
monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
|
||||
monkeypatch.setattr(startup_mod, "_get_json", lambda *_args, **_kwargs: oversized_assignment)
|
||||
monkeypatch.setattr(startup_mod, "_post_json", lambda _url, payload: registrations.append(payload) or {"node_id": "n1"})
|
||||
monkeypatch.setattr(startup_mod, "_start_heartbeat", lambda *_args, **_kwargs: None)
|
||||
monkeypatch.setattr(startup_mod, "model_metadata_for", lambda *_args, **_kwargs: {"num_layers": 40})
|
||||
|
||||
node = run_startup(
|
||||
tracker_url="http://127.0.0.1:8080",
|
||||
wallet_path=tmp_path / "wallet.json",
|
||||
tracker_source_disabled=True,
|
||||
)
|
||||
try:
|
||||
assert torch_calls[0]["shard_start"] == 0
|
||||
assert torch_calls[0]["shard_end"] == 24
|
||||
assert registrations[0]["shard_end"] == 24
|
||||
output = capsys.readouterr().out
|
||||
assert "CPU-safe runtime budget fits 25/40 layers" in output
|
||||
assert "layers 0-24" in output
|
||||
finally:
|
||||
node.stop()
|
||||
|
||||
|
||||
def test_preset_model_with_hf_repo_loads_torch_backend(tmp_path, monkeypatch, capsys):
|
||||
"""Named presets that advertise hf_repo must load TorchNodeServer, not the stub server."""
|
||||
import meshnet_node.startup as startup_mod
|
||||
|
||||
class FakeBackend:
|
||||
total_layers = 16
|
||||
|
||||
torch_calls: list[dict] = []
|
||||
|
||||
class FakeTorchNodeServer:
|
||||
def __init__(self, **kwargs):
|
||||
torch_calls.append(kwargs)
|
||||
self.backend = FakeBackend()
|
||||
self.port = None
|
||||
self.chat_completion_count = 0
|
||||
self.tracker_node_id = None
|
||||
|
||||
def start(self):
|
||||
self.port = 7002
|
||||
return self.port
|
||||
|
||||
def stop(self):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(
|
||||
startup_mod,
|
||||
"detect_hardware",
|
||||
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16 * 1024},
|
||||
)
|
||||
monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
|
||||
monkeypatch.setattr(startup_mod, "StubNodeServer", lambda **_kw: (_ for _ in ()).throw(AssertionError("preset with hf_repo must not use StubNodeServer")))
|
||||
|
||||
model_dir = tmp_path / "node-shards" / "tiny-llama"
|
||||
model_dir.mkdir(parents=True)
|
||||
(model_dir / "config.json").write_text('{"num_hidden_layers": 16}')
|
||||
monkeypatch.setattr(startup_mod, "download_shard", lambda *_a, **_kw: model_dir)
|
||||
|
||||
tracker = TrackerServer(model_presets={
|
||||
"tiny-llama": {"layers_start": 0, "layers_end": 15, "hf_repo": "org/tiny-llama-shards"}
|
||||
})
|
||||
tracker_port = tracker.start()
|
||||
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
||||
try:
|
||||
node = run_startup(
|
||||
tracker_url=tracker_url,
|
||||
model="tiny-llama",
|
||||
wallet_path=tmp_path / "wallet.json",
|
||||
cache_dir=tmp_path / "node-shards",
|
||||
)
|
||||
try:
|
||||
assert len(torch_calls) == 1
|
||||
assert torch_calls[0]["model_id"] == "org/tiny-llama-shards"
|
||||
assert torch_calls[0]["cache_dir"] == model_dir
|
||||
output = capsys.readouterr().out
|
||||
assert "Loading real PyTorch model shard..." in output
|
||||
assert "Model ID: org/tiny-llama-shards" in output
|
||||
network_map = _get_json(f"{tracker_url}/v1/network/map")
|
||||
registered = network_map["nodes"][0]
|
||||
assert registered["hf_repo"] == "org/tiny-llama-shards"
|
||||
assert registered["num_layers"] == 16
|
||||
finally:
|
||||
node.stop()
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_torch_startup_retries_registration_when_tracker_unreachable(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
@@ -1896,6 +2056,55 @@ def test_network_assign_gap_found_field():
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_network_assign_uses_conservative_cpu_runtime_budget():
|
||||
"""CPU assignments leave headroom for partial-load overhead, not just raw weights."""
|
||||
import json as _json
|
||||
import urllib.request as _ur
|
||||
|
||||
tracker = TrackerServer(model_presets={
|
||||
"qwen3.6-35b-a3b": {
|
||||
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
|
||||
"aliases": ["unsloth/Qwen3.6-35B-A3B"],
|
||||
"layers_start": 0,
|
||||
"layers_end": 39,
|
||||
"recommended": True,
|
||||
"bytes_per_layer": {"bfloat16": 1_797_594_419},
|
||||
},
|
||||
})
|
||||
port = tracker.start()
|
||||
try:
|
||||
data = _json.dumps({
|
||||
"endpoint": "http://127.0.0.1:9200",
|
||||
"model": "qwen3.6-35b-a3b",
|
||||
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
|
||||
"num_layers": 40,
|
||||
"shard_start": 0,
|
||||
"shard_end": 39,
|
||||
"hardware_profile": {},
|
||||
"score": 1.0,
|
||||
}).encode()
|
||||
req = _ur.Request(
|
||||
f"http://127.0.0.1:{port}/v1/nodes/register",
|
||||
data=data,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with _ur.urlopen(req) as r:
|
||||
r.read()
|
||||
|
||||
resp = _get_json(
|
||||
f"http://127.0.0.1:{port}/v1/network/assign"
|
||||
"?device=cpu&vram_mb=0&ram_mb=80896"
|
||||
"&hf_repo=unsloth/Qwen3.6-35B-A3B"
|
||||
)
|
||||
|
||||
assert resp["gap_found"] is False
|
||||
assert resp["shard_start"] == 0
|
||||
assert resp["shard_end"] == 24
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_route_finds_hf_model_across_two_nodes():
|
||||
"""Tracker /v1/route returns ordered route for HF model even without a preset."""
|
||||
import json as _json
|
||||
|
||||
@@ -17,6 +17,7 @@ from meshnet_node.model_backend import (
|
||||
TensorPayload,
|
||||
TorchModelShard,
|
||||
_call_layer,
|
||||
_checkpoint_tensor_name_for_model,
|
||||
_load_partial_model_from_snapshot,
|
||||
_should_partial_materialize_shard,
|
||||
_decoder_attention_mask,
|
||||
@@ -225,7 +226,7 @@ def test_tail_forward_returns_text_completion_from_binary_activations():
|
||||
node.stop()
|
||||
|
||||
|
||||
def test_full_model_chat_completion_uses_generation_not_single_token_decode():
|
||||
def test_full_model_chat_completion_uses_generation_not_single_token_decode(capsys):
|
||||
node = TorchNodeServer(backend=_FakeFullBackend())
|
||||
port = node.start()
|
||||
try:
|
||||
@@ -237,7 +238,10 @@ def test_full_model_chat_completion_uses_generation_not_single_token_decode():
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{port}/v1/chat/completions",
|
||||
data=payload,
|
||||
headers={"Content-Type": "application/json"},
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"X-Meshnet-Request-Id": "req-test-123",
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||
@@ -248,6 +252,10 @@ def test_full_model_chat_completion_uses_generation_not_single_token_decode():
|
||||
finally:
|
||||
node.stop()
|
||||
|
||||
out = capsys.readouterr().out
|
||||
assert " [node] processing chat model='fake-model' stream=False max_tokens=7 request_id=req-test-123" in out
|
||||
assert " [node] chat complete tokens=1 elapsed_s=" in out
|
||||
|
||||
|
||||
def test_pipeline_hop_logs_are_suppressed_without_debug(capsys):
|
||||
tail_backend = _FakePipelineTailBackend()
|
||||
@@ -368,6 +376,92 @@ def test_split_shard_chat_streams_each_generated_token_incrementally():
|
||||
assert "data: [DONE]" in rest
|
||||
|
||||
|
||||
def test_current_requests_snapshot_while_generating():
|
||||
release_second = threading.Event()
|
||||
head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
|
||||
tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second))
|
||||
head_port = head.start()
|
||||
tail_port = tail.start()
|
||||
response = None
|
||||
try:
|
||||
payload = json.dumps({
|
||||
"model": "fake-model",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"stream": True,
|
||||
"max_tokens": 2,
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{head_port}/v1/chat/completions",
|
||||
data=payload,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"X-Meshnet-Request-Id": "req-live-1",
|
||||
"X-Meshnet-Route": json.dumps([
|
||||
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
|
||||
]),
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
response = urllib.request.urlopen(req, timeout=5)
|
||||
deadline = time.time() + 2.0
|
||||
while time.time() < deadline:
|
||||
live = head.current_requests
|
||||
if live and live[0]["request_id"] == "req-live-1" and live[0]["tokens"] >= 1:
|
||||
break
|
||||
time.sleep(0.02)
|
||||
assert head.current_requests
|
||||
snap = head.current_requests[0]
|
||||
assert snap["request_id"] == "req-live-1"
|
||||
assert snap["tokens"] >= 1
|
||||
assert snap["tokens_per_sec"] >= 0
|
||||
assert snap["routing_complete"] is True
|
||||
release_second.set()
|
||||
response.read()
|
||||
finally:
|
||||
release_second.set()
|
||||
if response is not None:
|
||||
response.close()
|
||||
head.stop()
|
||||
tail.stop()
|
||||
|
||||
assert head.current_requests == []
|
||||
|
||||
|
||||
def test_distributed_generating_log_includes_tps(capsys):
|
||||
head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
|
||||
tail = TorchNodeServer(backend=_FakePipelineTailBackend())
|
||||
head_port = head.start()
|
||||
tail_port = tail.start()
|
||||
try:
|
||||
payload = json.dumps({
|
||||
"model": "fake-model",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 1,
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{head_port}/v1/chat/completions",
|
||||
data=payload,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"X-Meshnet-Route": json.dumps([
|
||||
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
|
||||
]),
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||
json.loads(resp.read())
|
||||
finally:
|
||||
head.stop()
|
||||
tail.stop()
|
||||
|
||||
out = capsys.readouterr().out
|
||||
assert "generating step=1/1" in out
|
||||
assert " tps=" in out
|
||||
assert "generation complete tokens=1" in out
|
||||
assert out.count("generating step=1/1") == 1
|
||||
|
||||
|
||||
def test_int_tensor_header_serializes_torch_tensors():
|
||||
torch = pytest.importorskip("torch")
|
||||
|
||||
@@ -422,7 +516,7 @@ def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapsho
|
||||
39,
|
||||
total_layers_hint=40,
|
||||
uses_quantized_weights=False,
|
||||
) is False
|
||||
) is True
|
||||
assert _should_partial_materialize_shard(
|
||||
str(snapshot_dir),
|
||||
4,
|
||||
@@ -439,6 +533,208 @@ def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapsho
|
||||
) is False
|
||||
|
||||
|
||||
def test_checkpoint_tensor_name_remapped_for_text_only_causal_lm():
|
||||
class TextOnlyModel:
|
||||
def __init__(self):
|
||||
self.model = types.SimpleNamespace(layers=[])
|
||||
|
||||
model = TextOnlyModel()
|
||||
assert _checkpoint_tensor_name_for_model(
|
||||
model,
|
||||
"model.language_model.layers.0.mlp.gate.weight",
|
||||
) == "model.layers.0.mlp.gate.weight"
|
||||
assert _checkpoint_tensor_name_for_model(
|
||||
model,
|
||||
"model.language_model.embed_tokens.weight",
|
||||
) == "model.embed_tokens.weight"
|
||||
|
||||
|
||||
def test_checkpoint_tensor_name_kept_for_multimodal_backbone():
|
||||
class MultimodalModel:
|
||||
def __init__(self):
|
||||
self.model = types.SimpleNamespace(language_model=types.SimpleNamespace())
|
||||
|
||||
model = MultimodalModel()
|
||||
name = "model.language_model.layers.0.mlp.gate.weight"
|
||||
assert _checkpoint_tensor_name_for_model(model, name) == name
|
||||
|
||||
|
||||
def test_partial_snapshot_loader_remaps_language_model_checkpoint_keys(tmp_path):
|
||||
snapshot_dir = tmp_path / "snapshot"
|
||||
snapshot_dir.mkdir()
|
||||
(snapshot_dir / "config.json").write_text(json.dumps({
|
||||
"text_config": {"num_hidden_layers": 3},
|
||||
}))
|
||||
(snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
|
||||
"weight_map": {
|
||||
"model.language_model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors",
|
||||
}
|
||||
}))
|
||||
(snapshot_dir / "shard-2.safetensors").write_bytes(b"stub")
|
||||
|
||||
class FakeModule:
|
||||
def __init__(self):
|
||||
self.to_calls = []
|
||||
|
||||
def to(self, device):
|
||||
self.to_calls.append(device)
|
||||
return self
|
||||
|
||||
class FakeModel:
|
||||
def __init__(self):
|
||||
self.model = types.SimpleNamespace(
|
||||
layers=[FakeModule(), FakeModule(), FakeModule()],
|
||||
rotary_emb=FakeModule(),
|
||||
)
|
||||
|
||||
def tie_weights(self):
|
||||
pass
|
||||
|
||||
class AutoConfigStub:
|
||||
@staticmethod
|
||||
def from_pretrained(model_id):
|
||||
return types.SimpleNamespace(
|
||||
text_config=types.SimpleNamespace(num_hidden_layers=3),
|
||||
get_text_config=lambda: types.SimpleNamespace(num_hidden_layers=3),
|
||||
)
|
||||
|
||||
class AutoModelStub:
|
||||
@staticmethod
|
||||
def from_config(cfg, torch_dtype=None):
|
||||
return FakeModel()
|
||||
|
||||
set_calls = []
|
||||
|
||||
def fake_set_tensor(module, tensor_name, device, value=None, dtype=None):
|
||||
set_calls.append(tensor_name)
|
||||
|
||||
class FakeSafeOpen:
|
||||
def __init__(self, filename, framework, device):
|
||||
self.filename = Path(filename).name
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def get_tensor(self, tensor_name):
|
||||
return tensor_name
|
||||
|
||||
class UnusedContext:
|
||||
def __enter__(self):
|
||||
return None
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
_load_partial_model_from_snapshot(
|
||||
AutoConfigStub,
|
||||
AutoModelStub,
|
||||
types.SimpleNamespace(),
|
||||
str(snapshot_dir),
|
||||
1,
|
||||
1,
|
||||
"bf16",
|
||||
"cpu:0",
|
||||
init_empty_weights_fn=UnusedContext,
|
||||
set_tensor_fn=fake_set_tensor,
|
||||
safe_open_fn=FakeSafeOpen,
|
||||
)
|
||||
|
||||
assert set_calls == ["model.layers.1.self_attn.q_proj.weight"]
|
||||
|
||||
|
||||
def test_partial_snapshot_loader_skips_tensors_absent_from_causal_lm(tmp_path):
|
||||
# Multimodal/MTP checkpoints (Qwen3.5/3.6-MoE) carry mtp.* and model.visual.*
|
||||
# tensors that the text-only CausalLM never builds — they must be skipped,
|
||||
# not assigned (assignment raises AttributeError: 'mtp' / 'visual').
|
||||
snapshot_dir = tmp_path / "snapshot"
|
||||
snapshot_dir.mkdir()
|
||||
(snapshot_dir / "config.json").write_text(json.dumps({
|
||||
"text_config": {"num_hidden_layers": 3},
|
||||
}))
|
||||
(snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
|
||||
"weight_map": {
|
||||
"model.language_model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors",
|
||||
"mtp.layers.1.input_layernorm.weight": "shard-2.safetensors",
|
||||
"model.visual.blocks.1.attn.qkv.weight": "shard-2.safetensors",
|
||||
}
|
||||
}))
|
||||
(snapshot_dir / "shard-2.safetensors").write_bytes(b"stub")
|
||||
|
||||
class FakeModule:
|
||||
def to(self, device):
|
||||
return self
|
||||
|
||||
class FakeModel:
|
||||
def __init__(self):
|
||||
self.model = types.SimpleNamespace(
|
||||
layers=[FakeModule(), FakeModule(), FakeModule()],
|
||||
rotary_emb=FakeModule(),
|
||||
)
|
||||
|
||||
def tie_weights(self):
|
||||
pass
|
||||
|
||||
def state_dict(self):
|
||||
return {"model.layers.1.self_attn.q_proj.weight": None}
|
||||
|
||||
class AutoConfigStub:
|
||||
@staticmethod
|
||||
def from_pretrained(model_id):
|
||||
return types.SimpleNamespace(
|
||||
text_config=types.SimpleNamespace(num_hidden_layers=3),
|
||||
get_text_config=lambda: types.SimpleNamespace(num_hidden_layers=3),
|
||||
)
|
||||
|
||||
class AutoModelStub:
|
||||
@staticmethod
|
||||
def from_config(cfg, torch_dtype=None):
|
||||
return FakeModel()
|
||||
|
||||
set_calls = []
|
||||
|
||||
def fake_set_tensor(module, tensor_name, device, value=None, dtype=None):
|
||||
set_calls.append(tensor_name)
|
||||
|
||||
class FakeSafeOpen:
|
||||
def __init__(self, filename, framework, device):
|
||||
pass
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def get_tensor(self, tensor_name):
|
||||
return tensor_name
|
||||
|
||||
class UnusedContext:
|
||||
def __enter__(self):
|
||||
return None
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
_load_partial_model_from_snapshot(
|
||||
AutoConfigStub,
|
||||
AutoModelStub,
|
||||
types.SimpleNamespace(),
|
||||
str(snapshot_dir),
|
||||
1,
|
||||
1,
|
||||
"bf16",
|
||||
"cpu:0",
|
||||
init_empty_weights_fn=UnusedContext,
|
||||
set_tensor_fn=fake_set_tensor,
|
||||
safe_open_fn=FakeSafeOpen,
|
||||
)
|
||||
|
||||
assert set_calls == ["model.layers.1.self_attn.q_proj.weight"]
|
||||
|
||||
|
||||
def test_partial_snapshot_loader_materializes_only_assigned_tensors(tmp_path):
|
||||
snapshot_dir = tmp_path / "snapshot"
|
||||
snapshot_dir.mkdir()
|
||||
|
||||
@@ -142,6 +142,21 @@ def _send_chat_request(gateway_url: str, prompt: str) -> dict:
|
||||
return json.loads(r.read())
|
||||
|
||||
|
||||
def _send_streaming_chat_request(gateway_url: str, prompt: str):
|
||||
data = json.dumps({
|
||||
"model": GPT2_MODEL,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"stream": True,
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"{gateway_url}/v1/chat/completions",
|
||||
data=data,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
return urllib.request.urlopen(req)
|
||||
|
||||
|
||||
def test_all_responses_valid_openai_format(tracker_node_setup):
|
||||
"""Ten requests via gateway all return valid OpenAI chat completion format."""
|
||||
gateway_url, _, _ = tracker_node_setup
|
||||
@@ -155,6 +170,30 @@ def test_all_responses_valid_openai_format(tracker_node_setup):
|
||||
assert isinstance(message.get("content"), str), f"request {i}: content must be a string"
|
||||
|
||||
|
||||
def test_streaming_head_worker_response_is_not_buffered_with_content_length(tracker_node_setup):
|
||||
"""Gateway must relay head-worker SSE as a live stream, not a buffered JSON-sized body."""
|
||||
gateway_url, _, _ = tracker_node_setup
|
||||
|
||||
with _send_streaming_chat_request(gateway_url, "stream through head worker") as resp:
|
||||
assert resp.status == 200
|
||||
assert "text/event-stream" in resp.headers["Content-Type"]
|
||||
assert "Content-Length" not in resp.headers
|
||||
data_lines = []
|
||||
while len(data_lines) < 4:
|
||||
line = resp.readline().decode().strip()
|
||||
if line.startswith("data: "):
|
||||
data_lines.append(line)
|
||||
if line == "data: [DONE]":
|
||||
break
|
||||
|
||||
assert data_lines[-1] == "data: [DONE]"
|
||||
content = "".join(
|
||||
json.loads(line[6:])["choices"][0].get("delta", {}).get("content", "")
|
||||
for line in data_lines[:-1]
|
||||
)
|
||||
assert "head-worker" in content
|
||||
|
||||
|
||||
def test_both_tracker_nodes_receive_load(tracker_node_setup):
|
||||
"""Both head workers handle at least one request each out of ten."""
|
||||
gateway_url, tracker_node_a, tracker_node_b = tracker_node_setup
|
||||
|
||||
67
tests/test_tracker_logging.py
Normal file
67
tests/test_tracker_logging.py
Normal file
@@ -0,0 +1,67 @@
|
||||
import logging
|
||||
import sys
|
||||
|
||||
from meshnet_tracker.logging_setup import configure_tracker_file_logging, tracker_logger
|
||||
|
||||
|
||||
def test_tracker_file_logging_writes_separate_level_files(tmp_path):
|
||||
original_stdout = sys.stdout
|
||||
original_stderr = sys.stderr
|
||||
try:
|
||||
log_dir = configure_tracker_file_logging(tmp_path, tee_stdio=False)
|
||||
logger = tracker_logger()
|
||||
|
||||
logger.info("info-event")
|
||||
logger.warning("warning-event")
|
||||
logger.error("error-event")
|
||||
for handler in logger.handlers:
|
||||
handler.flush()
|
||||
|
||||
assert (log_dir / "info.log").read_text().count("info-event") == 1
|
||||
assert "warning-event" not in (log_dir / "info.log").read_text()
|
||||
assert "error-event" not in (log_dir / "info.log").read_text()
|
||||
|
||||
assert "warning-event" in (log_dir / "warning.log").read_text()
|
||||
assert "info-event" not in (log_dir / "warning.log").read_text()
|
||||
assert "error-event" not in (log_dir / "warning.log").read_text()
|
||||
|
||||
assert "error-event" in (log_dir / "error.log").read_text()
|
||||
assert "info-event" not in (log_dir / "error.log").read_text()
|
||||
assert "warning-event" not in (log_dir / "error.log").read_text()
|
||||
finally:
|
||||
sys.stdout = original_stdout
|
||||
sys.stderr = original_stderr
|
||||
|
||||
|
||||
def test_tracker_file_logging_tees_stdio_and_rotates(tmp_path):
|
||||
original_stdout = sys.stdout
|
||||
original_stderr = sys.stderr
|
||||
try:
|
||||
log_dir = configure_tracker_file_logging(
|
||||
tmp_path,
|
||||
max_bytes=120,
|
||||
backup_count=1,
|
||||
)
|
||||
|
||||
print("stdout goes to info", flush=True)
|
||||
print("stderr goes to error", file=sys.stderr, flush=True)
|
||||
for handler in tracker_logger().handlers:
|
||||
handler.flush()
|
||||
|
||||
assert "stdout goes to info" in (log_dir / "info.log").read_text()
|
||||
assert "stderr goes to error" in (log_dir / "error.log").read_text()
|
||||
|
||||
for index in range(12):
|
||||
tracker_logger().info("rotating-info-line-%02d", index)
|
||||
for handler in tracker_logger().handlers:
|
||||
handler.flush()
|
||||
|
||||
assert (log_dir / "info.log.1").exists()
|
||||
finally:
|
||||
sys.stdout = original_stdout
|
||||
sys.stderr = original_stderr
|
||||
logger = tracker_logger()
|
||||
for handler in logger.handlers:
|
||||
handler.close()
|
||||
logger.handlers.clear()
|
||||
logger.setLevel(logging.NOTSET)
|
||||
@@ -5,6 +5,7 @@ import json
|
||||
import threading
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
|
||||
import pytest
|
||||
@@ -466,6 +467,7 @@ def test_tracker_logs_stream_progress_before_request_completes():
|
||||
method="POST",
|
||||
)
|
||||
response = urllib.request.urlopen(req, timeout=3.0)
|
||||
assert response.headers["X-Accel-Buffering"] == "no"
|
||||
first_line = response.readline()
|
||||
assert first_line.startswith(b"data:")
|
||||
assert chunk_sent.wait(timeout=1.0)
|
||||
@@ -501,6 +503,169 @@ def test_tracker_logs_stream_progress_before_request_completes():
|
||||
node_thread.join(timeout=1.0)
|
||||
|
||||
|
||||
def test_tracker_stream_survives_idle_gap_between_sse_chunks():
|
||||
first_chunk_sent = threading.Event()
|
||||
|
||||
class IdleStreamingChatHandler(http.server.BaseHTTPRequestHandler):
|
||||
def log_message(self, format, *args):
|
||||
pass
|
||||
|
||||
def do_POST(self):
|
||||
if self.path != "/v1/chat/completions":
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
return
|
||||
self.rfile.read(int(self.headers.get("Content-Length", 0)))
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
|
||||
self.end_headers()
|
||||
first = json.dumps({
|
||||
"choices": [{"delta": {"content": "hello"}}],
|
||||
}).encode()
|
||||
second = json.dumps({
|
||||
"choices": [{"delta": {"content": " world"}}],
|
||||
}).encode()
|
||||
self.wfile.write(b"data: " + first + b"\n\n")
|
||||
self.wfile.flush()
|
||||
first_chunk_sent.set()
|
||||
time.sleep(1.0)
|
||||
self.wfile.write(b"data: " + second + b"\n\n")
|
||||
self.wfile.write(b"data: [DONE]\n\n")
|
||||
self.wfile.flush()
|
||||
|
||||
node = http.server.HTTPServer(("127.0.0.1", 0), IdleStreamingChatHandler)
|
||||
node_thread = threading.Thread(target=node.serve_forever, daemon=True)
|
||||
node_thread.start()
|
||||
tracker = TrackerServer(heartbeat_timeout=60.0)
|
||||
tracker_port = tracker.start()
|
||||
response = None
|
||||
try:
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": f"http://127.0.0.1:{node.server_address[1]}",
|
||||
"model": "idle-stream-model", "num_layers": 1,
|
||||
"shard_start": 0, "shard_end": 0,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
|
||||
data=json.dumps({
|
||||
"model": "idle-stream-model",
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}).encode(),
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
response = urllib.request.urlopen(req, timeout=3.0)
|
||||
assert response.readline().startswith(b"data:")
|
||||
assert first_chunk_sent.wait(timeout=1.0)
|
||||
|
||||
remaining = response.read().splitlines()
|
||||
assert b"data: [DONE]" in remaining
|
||||
finally:
|
||||
if response is not None:
|
||||
response.close()
|
||||
tracker.stop()
|
||||
node.shutdown()
|
||||
node.server_close()
|
||||
node_thread.join(timeout=1.0)
|
||||
|
||||
|
||||
def test_tracker_dashboard_can_cancel_inflight_proxy():
|
||||
chunk_sent = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
class StreamingChatHandler(http.server.BaseHTTPRequestHandler):
|
||||
def log_message(self, fmt, *args):
|
||||
pass
|
||||
|
||||
def do_POST(self):
|
||||
if self.path != "/v1/chat/completions":
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
return
|
||||
self.rfile.read(int(self.headers.get("Content-Length", 0)))
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
|
||||
self.end_headers()
|
||||
payload = json.dumps({
|
||||
"choices": [{"delta": {"content": "hello world"}}],
|
||||
}).encode()
|
||||
self.wfile.write(b"data: " + payload + b"\n\n")
|
||||
self.wfile.flush()
|
||||
chunk_sent.set()
|
||||
release.wait(timeout=3.0)
|
||||
self.wfile.write(b"data: [DONE]\n\n")
|
||||
self.wfile.flush()
|
||||
|
||||
node = http.server.HTTPServer(("127.0.0.1", 0), StreamingChatHandler)
|
||||
node_thread = threading.Thread(target=node.serve_forever, daemon=True)
|
||||
node_thread.start()
|
||||
tracker = TrackerServer(heartbeat_timeout=60.0)
|
||||
tracker_port = tracker.start()
|
||||
response = None
|
||||
request_id = None
|
||||
try:
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": f"http://127.0.0.1:{node.server_address[1]}",
|
||||
"model": "cancel-proxy-model", "num_layers": 1,
|
||||
"shard_start": 0, "shard_end": 0,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
|
||||
data=json.dumps({
|
||||
"model": "cancel-proxy-model",
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}).encode(),
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
response = urllib.request.urlopen(req, timeout=3.0)
|
||||
first_line = response.readline()
|
||||
assert first_line.startswith(b"data:")
|
||||
assert chunk_sent.wait(timeout=1.0)
|
||||
|
||||
console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
|
||||
selected = [
|
||||
event for event in console["events"]
|
||||
if event["message"] == "proxy route selected"
|
||||
]
|
||||
assert selected
|
||||
request_id = selected[-1]["fields"]["request_id"]
|
||||
|
||||
cancel = _post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/proxy/requests/{urllib.parse.quote(request_id, safe='')}/cancel",
|
||||
{},
|
||||
)
|
||||
assert cancel["status"] == "canceled"
|
||||
|
||||
deadline = time.time() + 5.0
|
||||
canceled_events = []
|
||||
while time.time() < deadline:
|
||||
console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
|
||||
canceled_events = [
|
||||
event for event in console["events"]
|
||||
if event["message"] == "proxy canceled"
|
||||
and event["fields"].get("request_id") == request_id
|
||||
]
|
||||
if canceled_events:
|
||||
break
|
||||
time.sleep(0.05)
|
||||
assert canceled_events
|
||||
finally:
|
||||
release.set()
|
||||
if response is not None:
|
||||
response.close()
|
||||
tracker.stop()
|
||||
node.shutdown()
|
||||
node.server_close()
|
||||
node_thread.join(timeout=1.0)
|
||||
|
||||
|
||||
def test_tracker_routes_hf_model_alias_from_quickstart():
|
||||
"""The documented qwen2.5-0.5b alias resolves a full HF repo registration."""
|
||||
tracker = TrackerServer()
|
||||
@@ -746,6 +911,57 @@ def test_tracker_route_endpoint_ignores_model_case_and_outer_whitespace():
|
||||
assert response["route"] == ["http://127.0.0.1:9101"]
|
||||
|
||||
|
||||
def test_tracker_route_prefers_distributed_over_single_full_shard():
|
||||
"""When a full 0-39 node and a partial 0-21 head coexist, /v1/route
|
||||
should return both hops — not the full shard alone."""
|
||||
tracker = TrackerServer(model_presets={
|
||||
"qwen3.6-35b-a3b": {
|
||||
"layers_start": 0,
|
||||
"layers_end": 39,
|
||||
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
|
||||
"aliases": ["Qwen3.6-35B-A3B"],
|
||||
}
|
||||
})
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://192.168.0.179:7000",
|
||||
"model": "qwen3.6-35b-a3b",
|
||||
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
|
||||
"num_layers": 40,
|
||||
"shard_start": 0,
|
||||
"shard_end": 39,
|
||||
"tracker_mode": True,
|
||||
"hardware_profile": {},
|
||||
"score": 1.0},
|
||||
)
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://192.168.0.20:7000",
|
||||
"model": "Qwen3.6-35B-A3B",
|
||||
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
|
||||
"num_layers": 40,
|
||||
"shard_start": 0,
|
||||
"shard_end": 21,
|
||||
"tracker_mode": True,
|
||||
"hardware_profile": {},
|
||||
"score": 1.0},
|
||||
)
|
||||
|
||||
response = _get_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/route?model=qwen3.6-35b-a3b"
|
||||
)
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
assert response["route"] == [
|
||||
"http://192.168.0.20:7000",
|
||||
"http://192.168.0.179:7000",
|
||||
]
|
||||
assert [node["start_layer"] for node in response["nodes"]] == [0, 22]
|
||||
|
||||
|
||||
def test_tracker_proxy_ignores_model_case_and_outer_whitespace():
|
||||
class ChatHandler(http.server.BaseHTTPRequestHandler):
|
||||
def log_message(self, fmt, *args):
|
||||
@@ -1403,6 +1619,70 @@ def test_tracker_heartbeat_updates_node():
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_heartbeat_stores_current_requests():
|
||||
"""Node-reported in-flight request snapshots appear on the network map."""
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
reg = _post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{
|
||||
"endpoint": "http://127.0.0.1:9001",
|
||||
"model": "progress-model",
|
||||
"shard_start": 0,
|
||||
"shard_end": 31,
|
||||
"hardware_profile": {},
|
||||
"score": 1.0,
|
||||
},
|
||||
)
|
||||
node_id = reg["node_id"]
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat",
|
||||
{
|
||||
"queue_depth": 1,
|
||||
"current_requests": [{
|
||||
"request_id": "req-abc123",
|
||||
"model": "progress-model",
|
||||
"kind": "chat",
|
||||
"tokens": 17,
|
||||
"elapsed_seconds": 42.5,
|
||||
"tokens_per_sec": 0.4,
|
||||
"routing_complete": True,
|
||||
}],
|
||||
},
|
||||
)
|
||||
network = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
|
||||
node = next(item for item in network["nodes"] if item["node_id"] == node_id)
|
||||
assert node["stats"]["queue_depth"] == 1
|
||||
assert node["stats"]["current_requests"] == [{
|
||||
"request_id": "req-abc123",
|
||||
"model": "progress-model",
|
||||
"kind": "chat",
|
||||
"tokens": 17,
|
||||
"elapsed_seconds": 42.5,
|
||||
"tokens_per_sec": 0.4,
|
||||
"routing_complete": True,
|
||||
}]
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_normalize_current_requests_sanitizes_payload():
|
||||
from meshnet_tracker.server import _normalize_current_requests
|
||||
|
||||
assert _normalize_current_requests(None) == []
|
||||
assert _normalize_current_requests([
|
||||
{"request_id": "req-1", "model": "m", "tokens": "9", "tokens_per_sec": "1.5"},
|
||||
{"model": "missing-id"},
|
||||
"bad",
|
||||
]) == [{
|
||||
"request_id": "req-1",
|
||||
"model": "m",
|
||||
"tokens": 9,
|
||||
"tokens_per_sec": 1.5,
|
||||
}]
|
||||
|
||||
|
||||
def test_tracker_heartbeat_expiry():
|
||||
"""Nodes that miss their heartbeat window are excluded from routes."""
|
||||
tracker = TrackerServer(heartbeat_timeout=0.05) # 50 ms
|
||||
|
||||
Reference in New Issue
Block a user