Wire node memory and shard slot capabilities
This commit is contained in:
@@ -265,6 +265,7 @@ def run_startup(
|
||||
contracts: Any | None = None,
|
||||
route_timeout: float = 30.0,
|
||||
vram_mb_override: int | None = None,
|
||||
max_loaded_shards: int = 1,
|
||||
debug: bool = False,
|
||||
) -> StubNodeServer | TorchNodeServer:
|
||||
"""Execute the full startup sequence and return a running node server.
|
||||
@@ -282,6 +283,8 @@ def run_startup(
|
||||
|
||||
tracker_url = tracker_url.rstrip("/")
|
||||
relay_url = _discover_relay_url(tracker_url)
|
||||
if max_loaded_shards < 1:
|
||||
raise ValueError("--max-shards must be at least 1")
|
||||
|
||||
# 1. Hardware detection
|
||||
if advertise_host is None and host == "0.0.0.0":
|
||||
@@ -315,6 +318,11 @@ def run_startup(
|
||||
print(" WARNING: No CUDA GPU detected — running in CPU mode", flush=True)
|
||||
else:
|
||||
print(f" GPU: {gpu_name} ({vram_mb} MB VRAM)", flush=True)
|
||||
registration_capabilities = {
|
||||
"max_loaded_shards": max_loaded_shards,
|
||||
}
|
||||
if vram_mb_override is not None or vram_mb > 0:
|
||||
registration_capabilities["vram_bytes"] = max(0, int(vram_mb)) * 1024 * 1024
|
||||
|
||||
# 2. Wallet
|
||||
print("Loading wallet...", flush=True)
|
||||
@@ -404,6 +412,7 @@ def run_startup(
|
||||
"score": 1.0,
|
||||
"tracker_mode": (shard_start == 0),
|
||||
"managed_assignment": not user_pinned_shard,
|
||||
**registration_capabilities,
|
||||
**relay_fields,
|
||||
}
|
||||
tracker_node_id: str | None = None
|
||||
@@ -493,6 +502,7 @@ def run_startup(
|
||||
"score": 1.0,
|
||||
"tracker_mode": (assigned_shard_start == 0),
|
||||
"managed_assignment": True,
|
||||
**registration_capabilities,
|
||||
**relay_fields,
|
||||
}
|
||||
tracker_node_id = None
|
||||
@@ -593,6 +603,7 @@ def run_startup(
|
||||
"hardware_profile": hw,
|
||||
"wallet_address": address,
|
||||
"score": 1.0,
|
||||
**registration_capabilities,
|
||||
**relay_fields,
|
||||
},
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user