feat: let admins manage model placement
This commit is contained in:
@@ -1528,6 +1528,49 @@ def _request_model_load_locked(server: "_TrackerHTTPServer", model_key: str) ->
|
||||
return None
|
||||
|
||||
|
||||
def _force_model_load_locked(server: "_TrackerHTTPServer", model_key: str) -> dict | None:
|
||||
"""Replace the fastest ready assignment after an explicit admin eviction."""
|
||||
resolved_name, preset = _resolve_model_preset(server.model_presets, model_key)
|
||||
if preset is None or not preset.get("hf_repo"):
|
||||
return None
|
||||
start, end = _preset_layer_bounds(preset)
|
||||
candidates = [node for node in server.registry.values()
|
||||
if node.status == "ready" and node.pending_new_assignment is None
|
||||
and _has_usable_quantization(node)]
|
||||
if not candidates:
|
||||
return None
|
||||
node = max(candidates, key=lambda item: item.benchmark_tokens_per_sec)
|
||||
shard_end = min(end, start + min(_node_layer_capacity(node, preset), end - start + 1) - 1)
|
||||
if shard_end < start:
|
||||
return None
|
||||
quantization = _node_quantization(node, preset)
|
||||
directive = _load_directive(node, str(preset["hf_repo"]), start, shard_end, quantization)
|
||||
replaced = node.hf_repo or node.model
|
||||
node.model, node.hf_repo = resolved_name, str(preset["hf_repo"])
|
||||
node.shard_start, node.shard_end, node.quantization = start, shard_end, quantization
|
||||
node.managed_assignment, node.pending_new_assignment = True, directive
|
||||
node.pending_directives.append(directive)
|
||||
_tracker_log(server, "warn", "model load forced", node_id=node.node_id,
|
||||
model=resolved_name, replaced_model=replaced, shard=f"{start}-{shard_end}")
|
||||
return {"node_id": node.node_id, "model": resolved_name, "hf_repo": preset["hf_repo"],
|
||||
"shard_start": start, "shard_end": shard_end, "replaced_model": replaced}
|
||||
|
||||
|
||||
def _release_model_locked(server: "_TrackerHTTPServer", model_key: str) -> int:
|
||||
"""Queue DROP_SHARD for every served shard and remove it from routing immediately."""
|
||||
resolved_name, preset = _resolve_model_preset(server.model_presets, model_key)
|
||||
if preset is None:
|
||||
return 0
|
||||
released = 0
|
||||
for node in server.registry.values():
|
||||
if not _node_matches_preset(node, resolved_name, preset) or node.shard_start is None or node.shard_end is None:
|
||||
continue
|
||||
node.pending_directives.append(_drop_directive(node, str(preset.get("hf_repo") or resolved_name), node.shard_start, node.shard_end, node.quantization or "bfloat16"))
|
||||
node.status = "loading"
|
||||
released += 1
|
||||
return released
|
||||
|
||||
|
||||
def _preferred_node_quantization(
|
||||
node: _NodeEntry,
|
||||
preset: dict,
|
||||
@@ -3043,6 +3086,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
if self.path == "/v1/models/load":
|
||||
self._handle_model_load_request()
|
||||
return
|
||||
if self.path == "/v1/models/release":
|
||||
self._handle_model_release_request()
|
||||
return
|
||||
if self.path == "/v1/models/vote":
|
||||
self._handle_model_coverage_vote()
|
||||
return
|
||||
@@ -3170,8 +3216,6 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
seen_ids: set[str] = set()
|
||||
for name, preset in server.model_presets.items():
|
||||
model_nodes = [node for node in alive if _node_matches_preset(node, name, preset)]
|
||||
if not model_nodes and not preset.get("recommended"):
|
||||
continue
|
||||
required_start, required_end = _preset_layer_bounds(preset)
|
||||
coverage = _coverage_percentage(
|
||||
model_nodes,
|
||||
@@ -4834,11 +4878,32 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
with server.lock:
|
||||
self._purge_expired_nodes()
|
||||
assignment = _request_model_load_locked(server, model)
|
||||
if assignment is None and body.get("force") is True:
|
||||
assignment = _force_model_load_locked(server, model)
|
||||
if assignment is None:
|
||||
self._send_json(409, {"error": "no ready joined node has an available model slot and sufficient capacity"})
|
||||
return
|
||||
self._send_json(202, {"status": "queued", "assignment": assignment})
|
||||
|
||||
def _handle_model_release_request(self):
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
if not self._require_role("admin", "validator"):
|
||||
return
|
||||
body = self._read_json_body()
|
||||
if body is None:
|
||||
return
|
||||
model = body.get("model")
|
||||
if not isinstance(model, str) or not model.strip():
|
||||
self._send_json(400, {"error": "model is required"})
|
||||
return
|
||||
with server.lock:
|
||||
self._purge_expired_nodes()
|
||||
released = _release_model_locked(server, model)
|
||||
if not released:
|
||||
self._send_json(404, {"error": "no served shards found for model"})
|
||||
return
|
||||
self._send_json(202, {"status": "release_queued", "released": released})
|
||||
|
||||
def _handle_model_coverage_vote(self):
|
||||
"""Record a rolling wish-list signal for an unavailable precision."""
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
|
||||
Reference in New Issue
Block a user