restore worker models

Author: OliverBryant

xinference/core/supervisor.py

@@ -1860,6 +1860,77 @@ async def add_worker(self, worker_address: str):
         self._worker_address_to_worker[worker_address] = worker_ref
         logger.debug("Worker %s has been added successfully", worker_address)
 
+    @log_async(logger=logger)
+    async def ensure_worker(
+        self, worker_address: str
+    ) -> xo.ActorRefType["WorkerActor"]:
+        from .worker import WorkerActor
+
+        worker_ref = await xo.actor_ref(
+            address=worker_address, uid=WorkerActor.default_uid()
+        )
+        if worker_address in self._worker_address_to_worker:
+            self._worker_address_to_worker[worker_address] = worker_ref
+            logger.debug("Worker %s already registered, refreshed ref", worker_address)
+        else:
+            self._worker_address_to_worker[worker_address] = worker_ref
+            logger.debug("Worker %s has been added successfully", worker_address)
+        return worker_ref
+
+    @log_async(logger=logger)
+    async def restore_worker_models(
+        self, worker_address: str, models: Dict[str, Dict[str, Any]]
+    ):
+        if not models:
+            return
+        worker_ref = await self.ensure_worker(worker_address)
+        restored = 0
+        for replica_model_uid in models.keys():
+            model_uid, rep_id = parse_replica_model_uid(replica_model_uid)
+            if rep_id < 0:
+                rep_id = 0
+
+            replica_info = self._model_uid_to_replica_info.get(model_uid, None)
+            if replica_info is None:
+                replica_count = rep_id + 1
+                replica_info = ReplicaInfo(
+                    replica=replica_count,
+                    scheduler=itertools.cycle(range(replica_count)),
+                )
+                self._model_uid_to_replica_info[model_uid] = replica_info
+            elif rep_id + 1 > replica_info.replica:
+                replica_info.replica = rep_id + 1
+                replica_info.scheduler = itertools.cycle(range(replica_info.replica))
+
+            if all(
+                w.address != worker_ref.address
+                for w in replica_info.replica_to_worker_refs[rep_id]
+            ):
+                replica_info.replica_to_worker_refs[rep_id].append(worker_ref)
+
+            existing = self._replica_model_uid_to_worker.get(replica_model_uid, None)
+            if existing is None:
+                self._replica_model_uid_to_worker[replica_model_uid] = worker_ref
+            elif isinstance(existing, (list, tuple)):
+                if all(w.address != worker_ref.address for w in existing):
+                    if isinstance(existing, tuple):
+                        self._replica_model_uid_to_worker[replica_model_uid] = [
+                            *existing,
+                            worker_ref,
+                        ]
+                    else:
+                        existing.append(worker_ref)
+            else:
+                if existing.address != worker_ref.address:
+                    self._replica_model_uid_to_worker[replica_model_uid] = [
+                        existing,
+                        worker_ref,
+                    ]
+            restored += 1
+        logger.info(
+            "Restored %s model replicas for worker %s", restored, worker_address
+        )
+
     @log_async(logger=logger)
     async def remove_worker(self, worker_address: str):
         uids_to_remove = []

xinference/core/worker.py

@@ -456,10 +456,20 @@ async def get_supervisor_ref(self, add_worker: bool = True) -> xo.ActorRefType:
             if self._supervisor_ref is not None:
                 return self._supervisor_ref
             self._supervisor_ref = supervisor_ref
-            if add_worker and len(self._model_uid_to_model) == 0:
-                # Newly started (or restarted), has no model, notify supervisor
-                await self._supervisor_ref.add_worker(self.address)
-                logger.info("Connected to supervisor as a fresh worker")
+            if add_worker:
+                await self._supervisor_ref.ensure_worker(self.address)
+                if len(self._model_uid_to_model) == 0:
+                    logger.info("Connected to supervisor as a fresh worker")
+                else:
+                    try:
+                        models = await self.list_models()
+                        await self._supervisor_ref.restore_worker_models(
+                            self.address, models
+                        )
+                    except Exception:
+                        logger.exception(
+                            "Failed to restore worker models to supervisor"
+                        )
 
             self._status_guard_ref = await xo.actor_ref(
                 address=self._supervisor_address, uid=StatusGuardActor.default_uid()