feat: add p2p shard swarm
This commit is contained in:
@@ -4,6 +4,7 @@ HTTP API contract:
|
||||
- POST /v1/nodes/register
|
||||
Request JSON: {
|
||||
"endpoint": "http://host:port", "shard_start": int, "shard_end": int,
|
||||
"model": str optional, "shard_checksum": str optional,
|
||||
"hardware_profile": object, "wallet_address": str optional,
|
||||
"score": number optional
|
||||
}
|
||||
@@ -34,7 +35,8 @@ DEFAULT_MODEL_PRESETS: dict[str, dict] = {
|
||||
class _NodeEntry:
|
||||
__slots__ = (
|
||||
"node_id", "endpoint", "shard_start", "shard_end",
|
||||
"hardware_profile", "wallet_address", "score", "last_heartbeat",
|
||||
"model", "shard_checksum", "hardware_profile", "wallet_address",
|
||||
"score", "last_heartbeat",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
@@ -43,6 +45,8 @@ class _NodeEntry:
|
||||
endpoint: str,
|
||||
shard_start: int,
|
||||
shard_end: int,
|
||||
model: str | None,
|
||||
shard_checksum: str | None,
|
||||
hardware_profile: dict,
|
||||
wallet_address: str | None,
|
||||
score: float,
|
||||
@@ -51,6 +55,8 @@ class _NodeEntry:
|
||||
self.endpoint = endpoint
|
||||
self.shard_start = shard_start
|
||||
self.shard_end = shard_end
|
||||
self.model = model
|
||||
self.shard_checksum = shard_checksum
|
||||
self.hardware_profile = hardware_profile
|
||||
self.wallet_address = wallet_address
|
||||
self.score = score
|
||||
@@ -209,6 +215,16 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
if not isinstance(hardware_profile, dict):
|
||||
self._send_json(400, {"error": "hardware_profile must be an object"})
|
||||
return
|
||||
model = body.get("model")
|
||||
if model is None:
|
||||
model = "stub-model"
|
||||
if not isinstance(model, str):
|
||||
self._send_json(400, {"error": "model must be a string"})
|
||||
return
|
||||
shard_checksum = body.get("shard_checksum")
|
||||
if shard_checksum is not None and not isinstance(shard_checksum, str):
|
||||
self._send_json(400, {"error": "shard_checksum must be a string"})
|
||||
return
|
||||
wallet_address = body.get("wallet_address")
|
||||
if wallet_address is not None and not isinstance(wallet_address, str):
|
||||
self._send_json(400, {"error": "wallet_address must be a string"})
|
||||
@@ -224,6 +240,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
endpoint=endpoint.rstrip("/"),
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
model=model,
|
||||
shard_checksum=shard_checksum,
|
||||
hardware_profile=hardware_profile,
|
||||
wallet_address=wallet_address,
|
||||
score=score,
|
||||
@@ -279,7 +297,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
with server.lock:
|
||||
self._purge_expired_nodes()
|
||||
alive = list(server.registry.values())
|
||||
alive = [node for node in server.registry.values() if node.model == model]
|
||||
if server.contracts is not None:
|
||||
alive = [
|
||||
node for node in alive
|
||||
@@ -317,11 +335,21 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
shard_start = gap_start
|
||||
shard_end = min(required_end, shard_start + max_layers - 1)
|
||||
|
||||
peers = [
|
||||
{"endpoint": node.endpoint, "checksum": node.shard_checksum}
|
||||
for node in alive
|
||||
if node.model == model
|
||||
and node.shard_start == shard_start
|
||||
and node.shard_end == shard_end
|
||||
and node.shard_checksum
|
||||
]
|
||||
|
||||
self._send_json(200, {
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"model": model,
|
||||
"model_layers_end": required_end,
|
||||
"peers": peers,
|
||||
**({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}),
|
||||
})
|
||||
|
||||
@@ -344,7 +372,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
with server.lock:
|
||||
self._purge_expired_nodes()
|
||||
alive = list(server.registry.values())
|
||||
alive = [node for node in server.registry.values() if node.model == model]
|
||||
if server.contracts is not None:
|
||||
alive = [
|
||||
node for node in alive
|
||||
@@ -364,6 +392,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
"wallet_address": e.wallet_address,
|
||||
"shard_start": e.shard_start,
|
||||
"shard_end": e.shard_end,
|
||||
"model": e.model,
|
||||
"shard_checksum": e.shard_checksum,
|
||||
"score": e.score,
|
||||
}
|
||||
for e in route
|
||||
|
||||
Reference in New Issue
Block a user