Wire node memory and shard slot capabilities
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user