Files
neuron-tai/.scratch/distributed-inference-network/issues/13-coverage-first-shard-assignment.md
Dobromir Popov 344f432880 tasks
2026-06-30 21:19:47 +02:00

3.6 KiB

US-013 — Coverage-first tracker shard assignment

Upgrade the tracker's route selection to a coverage-first, speed-weighted bin-packing algorithm. The tracker maintains a live coverage map per model, issues rebalance directives when gaps appear, and assigns shard ranges to nodes based on declared VRAM, quantization, and benchmark throughput.

Context

Per ADR-0009:

  • Coverage gap = any layer range in a model with zero serving nodes → model is unroutable
  • Coverage-first: fill all gaps before adding redundancy to covered ranges
  • Speed-weighted bin-packing: fill fastest nodes to their VRAM limit first; cascade to next-fastest
  • Continuous rebalancing: triggered every 30s and on every node join/leave event
  • Rebalance directive: LOAD_SHARD or DROP_SHARD instruction sent from tracker to node

Node registration now includes capability data:

{
  "wallet": "...",
  "endpoint": "http://10.0.0.5:7000",
  "vram_bytes": 25769803776,
  "ram_bytes": 137438953472,
  "quantizations": ["nf4", "int8", "bfloat16"],
  "benchmark_tokens_per_sec": 12.4,
  "benchmark_model": "openai-community/gpt2"
}

Model preset metadata (stored in tracker config) includes bytes_per_layer at each quantization level and total_layers. The tracker can compute max assignable layers for any node.

Acceptance Criteria

  • Tracker accepts vram_bytes, ram_bytes, quantizations, benchmark_tokens_per_sec in the node registration payload
  • GET /v1/coverage/<model_preset> returns the coverage map: list of {start_layer, end_layer, node_count} for each assigned range
  • A model is only routable when all layer ranges have node_count >= 1
  • When a new node registers, the tracker assigns it to the highest-priority uncovered range (or expands the most-loaded range if fully covered)
  • When a node disconnects (heartbeat timeout), the tracker marks affected ranges as reduced-coverage and, if any reach 0, issues LOAD_SHARD directives to idle nodes within 30 seconds
  • The shard assignment respects VRAM: assigned_layers <= floor((vram_bytes * 0.8) / bytes_per_layer_at_quant)
  • Speed weighting: given two idle nodes both capable of covering a gap, the tracker assigns the wider sub-range to the faster node
  • An integration test: start tracker with a model preset, register three nodes with different VRAM, assert the coverage map shows 100% coverage and the largest VRAM node received the widest shard range
  • A second integration test: kill the node covering the middle layer range; assert tracker issues a LOAD_SHARD directive to an idle node and coverage recovers
  • python -m pytest passes
  • Commit only this story's changes

Implementation Notes

  • Coverage map data structure: dict[str, list[tuple[int, int, int]]] keyed by model preset
  • Bin-packing is run in-process on the tracker (not a separate service); it's triggered by registration/heartbeat events and a 30s periodic timer
  • Rebalance directives are delivered as responses to node heartbeat POSTs (node polls tracker every 10s anyway) — no new push channel needed
  • bytes_per_layer is read from a model_presets.json config file; add GPT-2 (12 layers, ~30MB/layer bfloat16) as the test preset
  • The current tracker _select_route function is extended, not replaced — backward compat with stub nodes that don't send VRAM data (default to 8GB / bfloat16 if omitted)

Comments

  • 2026-06-30: Follow-up capacity hardening lives in 20-memory-budget-shard-slots-and-dropout-relocation.md. US-013 remains the base coverage-first assignment and dropout rebalance story; US-020 owns operator --memory, --max-shards, shard-slot enforcement, and relocation limit hardening so the two scopes do not conflict.