feat: let admins manage model placement

This commit is contained in:
Dobromir Popov
2026-07-14 15:16:23 +02:00
parent 29351d6217
commit e2f3ae32b8
5 changed files with 166 additions and 3 deletions

View File

@@ -1545,6 +1545,26 @@ class TorchNodeServer:
def apply_tracker_directives(self, directives: list[dict]) -> dict | None:
"""Apply tracker shard directives (LOAD_SHARD replace, ADD_SHARD load-more)."""
drop_directive = next(
(directive for directive in reversed(directives) if directive.get("action") == "DROP_SHARD"),
None,
)
if drop_directive is not None:
model_id = str(drop_directive.get("model") or "")
removed = self._backends.pop(model_id, None)
if removed is None:
return None
if self._backends:
self._backend = next(iter(self._backends.values()))
self._tracker_mode = self._backend.shard_start == 0
else:
self._backend = None
self._tracker_mode = False
if self._server is not None:
self._server.backends = dict(self._backends)
self._server.backend = self._backend
self._server.tracker_mode = self._tracker_mode
return {"action": "DROP_SHARD", "model": model_id}
add_directive = next(
(directive for directive in reversed(directives) if directive.get("action") == "ADD_SHARD"),
None,
@@ -1574,6 +1594,8 @@ class TorchNodeServer:
flush=True,
)
try:
if replacing:
self._backends.clear()
new_backend = _load_backend(model_id, shard_start, shard_end, quantization, self._cache_dir)
except TypeError:
new_backend = _load_backend(model_id, shard_start, shard_end, quantization)