Wire node memory and shard slot capabilities

This commit is contained in:
Dobromir Popov
2026-06-30 21:17:07 +02:00
parent 61074a8fe8
commit f1e4ed6a32
5 changed files with 53 additions and 4 deletions

View File

@@ -259,7 +259,7 @@ class _NodeEntry:
__slots__ = (
"node_id", "endpoint", "shard_start", "shard_end",
"model", "hf_repo", "num_layers", "shard_checksum", "hardware_profile", "wallet_address",
"score", "vram_bytes", "ram_bytes", "quantizations",
"score", "vram_bytes", "ram_bytes", "quantizations", "max_loaded_shards",
"benchmark_tokens_per_sec", "quantization", "managed_assignment",
"pending_directives", "last_heartbeat", "tracker_mode",
"relay_addr", "cert_fingerprint", "peer_id",
@@ -285,6 +285,7 @@ class _NodeEntry:
vram_bytes: int = DEFAULT_VRAM_BYTES,
ram_bytes: int = DEFAULT_RAM_BYTES,
quantizations: list[str] | None = None,
max_loaded_shards: int = 1,
benchmark_tokens_per_sec: float = DEFAULT_BENCHMARK_TOKENS_PER_SEC,
quantization: str | None = None,
managed_assignment: bool = False,
@@ -307,6 +308,7 @@ class _NodeEntry:
self.vram_bytes = vram_bytes
self.ram_bytes = ram_bytes
self.quantizations = quantizations or list(DEFAULT_QUANTIZATIONS)
self.max_loaded_shards = max_loaded_shards
self.benchmark_tokens_per_sec = benchmark_tokens_per_sec
self.quantization = quantization
self.managed_assignment = managed_assignment
@@ -1431,13 +1433,14 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
try:
vram_bytes = int(body.get("vram_bytes", DEFAULT_VRAM_BYTES))
ram_bytes = int(body.get("ram_bytes", DEFAULT_RAM_BYTES))
max_loaded_shards = int(body.get("max_loaded_shards", 1))
benchmark_tokens_per_sec = float(
body.get("benchmark_tokens_per_sec", DEFAULT_BENCHMARK_TOKENS_PER_SEC)
)
except (TypeError, ValueError):
self._send_json(400, {"error": "vram_bytes, ram_bytes, and benchmark_tokens_per_sec must be numeric"})
self._send_json(400, {"error": "vram_bytes, ram_bytes, max_loaded_shards, and benchmark_tokens_per_sec must be numeric"})
return
if vram_bytes < 0 or ram_bytes < 0 or benchmark_tokens_per_sec <= 0:
if vram_bytes < 0 or ram_bytes < 0 or max_loaded_shards < 1 or benchmark_tokens_per_sec <= 0:
self._send_json(400, {"error": "capability values must be positive"})
return
quantizations_body = body.get("quantizations", DEFAULT_QUANTIZATIONS)
@@ -1501,6 +1504,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
vram_bytes=vram_bytes,
ram_bytes=ram_bytes,
quantizations=quantizations,
max_loaded_shards=max_loaded_shards,
benchmark_tokens_per_sec=benchmark_tokens_per_sec,
quantization=quantization,
managed_assignment=managed_assignment or not explicit_shard,