fix: reconcile merged runtime behavior

This commit is contained in:
Dobromir Popov
2026-07-17 14:29:09 +03:00
parent 6aced6a005
commit f0197bfa83
3 changed files with 45 additions and 58 deletions

View File

@@ -1576,19 +1576,24 @@ class TorchNodeServer:
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}
if removed is not 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
# A DROP_SHARD may be paired with a subsequent LOAD_SHARD in the
# same heartbeat; continue so the replacement is applied atomically.
if not any(
directive.get("action") in {"ADD_SHARD", "LOAD_SHARD"}
for directive in directives
):
return {"action": "DROP_SHARD", "model": model_id} if removed is not None else None
add_directive = next(
(directive for directive in reversed(directives) if directive.get("action") == "ADD_SHARD"),
None,
@@ -1602,8 +1607,11 @@ class TorchNodeServer:
return None
shard_start = int(directive["shard_start"])
shard_end = int(directive["shard_end"])
quantization = str(directive.get("quantization") or self._backend.quantization)
model_id = str(directive.get("model") or self._backend.model_id)
quantization = str(
directive.get("quantization")
or (self._backend.quantization if self._backend is not None else "auto")
)
model_id = str(directive.get("model") or (self._backend.model_id if self._backend is not None else ""))
replacing = directive.get("action") == "LOAD_SHARD"
if not replacing and len(self._backends) >= self._max_loaded_shards:
print(