We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
BUG: worker timeout when downloading a model


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `xinference/core/service.py`
Content:
```
1 # Copyright 2022-2023 XProbe Inc.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #      http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 import asyncio
16 import time
17 from dataclasses import dataclass
18 from logging import getLogger
19 from typing import Callable, Dict, List, Optional, Set, Tuple
20 
21 import xoscar as xo
22 
23 from ..core import ModelActor
24 from ..model import ModelSpec
25 from .resource import ResourceStatus, gather_node_info
26 
27 logger = getLogger(__name__)
28 
29 
30 DEFAULT_NODE_DEAD_TIMEOUT = 30
31 DEFAULT_NODE_CHECK_INTERVAL = 1
32 
33 
34 def log(func: Callable):
35     # TODO: support non-async function
36     import time
37     from functools import wraps
38 
39     @wraps(func)
40     async def wrapped(*args, **kwargs):
41         logger.debug(f"Enter {func.__name__}, args: {args}, kwargs: {kwargs}")
42         start = time.time()
43         ret = await func(*args, **kwargs)
44         logger.debug(
45             f"Leave {func.__name__}, elapsed time: {int(time.time() - start)} ms"
46         )
47         return ret
48 
49     return wrapped
50 
51 
52 @dataclass
53 class WorkerStatus:
54     update_time: float
55     status: Dict[str, ResourceStatus]
56 
57 
58 class SupervisorActor(xo.Actor):
59     def __init__(self):
60         super().__init__()
61         self._worker_address_to_worker: Dict[str, xo.ActorRefType[WorkerActor]] = {}
62         self._model_uid_to_worker: Dict[str, xo.ActorRefType[WorkerActor]] = {}
63         self._worker_status: Dict[str, WorkerStatus] = {}
64 
65     @classmethod
66     def uid(cls) -> str:
67         return "supervisor"
68 
69     async def __post_create__(self):
70         self._check_dead_nodes_task = asyncio.create_task(self._check_dead_nodes())
71 
72     async def __pre_destroy__(self):
73         self._check_dead_nodes_task.cancel()
74 
75     async def _choose_worker(self) -> xo.ActorRefType["WorkerActor"]:
76         # TODO: better allocation strategy.
77         min_running_model_count = None
78         target_worker = None
79         for worker in self._worker_address_to_worker.values():
80             running_model_count = await worker.get_model_count()
81             if (
82                 min_running_model_count is None
83                 or running_model_count < min_running_model_count
84             ):
85                 min_running_model_count = running_model_count
86                 target_worker = worker
87 
88         if target_worker:
89             return target_worker
90 
91         raise RuntimeError("No available worker found")
92 
93     @log
94     async def launch_builtin_model(
95         self,
96         model_uid: str,
97         model_name: str,
98         model_size_in_billions: Optional[int],
99         model_format: Optional[str],
100         quantization: Optional[str],
101         **kwargs,
102     ) -> xo.ActorRefType["ModelActor"]:
103         assert model_uid not in self._model_uid_to_worker
104 
105         worker_ref = await self._choose_worker()
106         model_ref = await worker_ref.launch_builtin_model(
107             model_uid=model_uid,
108             model_name=model_name,
109             model_size_in_billions=model_size_in_billions,
110             model_format=model_format,
111             quantization=quantization,
112             **kwargs,
113         )
114         self._model_uid_to_worker[model_uid] = worker_ref
115 
116         return model_ref
117 
118     async def _check_dead_nodes(self):
119         while True:
120             for address, status in self._worker_status.items():
121                 if time.time() - status.update_time > DEFAULT_NODE_DEAD_TIMEOUT:
122                     dead_models = []
123                     for model_uid in self._model_uid_to_worker:
124                         if self._model_uid_to_worker[model_uid].address == address:
125                             dead_models.append(model_uid)
126                     logger.error(
127                         "Worker timeout. address: %s, influenced models: %s",
128                         address,
129                         dead_models,
130                     )
131                     self._worker_status.pop(address)
132                     self._worker_address_to_worker.pop(address)
133             await asyncio.sleep(5)
134 
135     @log
136     async def terminate_model(self, model_uid: str):
137         assert model_uid in self._model_uid_to_worker
138 
139         worker_ref = self._model_uid_to_worker[model_uid]
140         await worker_ref.terminate_model(model_uid=model_uid)
141         del self._model_uid_to_worker[model_uid]
142 
143     @log
144     async def get_model(self, model_uid: str) -> xo.ActorRefType["ModelActor"]:
145         worker_ref = self._model_uid_to_worker[model_uid]
146         return await worker_ref.get_model(model_uid=model_uid)
147 
148     @log
149     async def describe_model(self, model_uid: str):
150         worker_ref = self._model_uid_to_worker[model_uid]
151         return await worker_ref.describe_model(model_uid=model_uid)
152 
153     @log
154     async def list_models(self) -> List[Tuple[str, ModelSpec]]:
155         ret = []
156         for worker in self._worker_address_to_worker.values():
157             ret.extend(await worker.list_models())
158         return ret
159 
160     @log
161     async def add_worker(self, worker_address: str):
162         assert worker_address not in self._worker_address_to_worker
163 
164         worker_ref = await xo.actor_ref(address=worker_address, uid=WorkerActor.uid())
165         self._worker_address_to_worker[worker_address] = worker_ref
166 
167     async def report_worker_status(
168         self, worker_address: str, status: Dict[str, ResourceStatus]
169     ):
170         self._worker_status[worker_address] = WorkerStatus(
171             update_time=time.time(), status=status
172         )
173 
174 
175 class WorkerActor(xo.Actor):
176     def __init__(self, supervisor_address: str, subpool_addresses: List[str]):
177         super().__init__()
178         self._supervisor_address = supervisor_address
179         self._supervisor_ref = None
180         self._model_uid_to_model: Dict[str, xo.ActorRefType["ModelActor"]] = {}
181         self._model_uid_to_model_spec: Dict[str, ModelSpec] = {}
182         self._subpool_address_to_model_uids: Dict[str, Set[str]] = dict(
183             [(subpool_address, set()) for subpool_address in subpool_addresses]
184         )
185         logger.debug(f"Worker actor initialized with subpools: {subpool_addresses}")
186 
187     @classmethod
188     def uid(cls) -> str:
189         return "worker"
190 
191     async def __post_create__(self):
192         self._supervisor_ref: xo.ActorRefType["SupervisorActor"] = await xo.actor_ref(
193             address=self._supervisor_address, uid=SupervisorActor.uid()
194         )
195         await self._supervisor_ref.add_worker(self.address)
196         self._upload_task = asyncio.create_task(self._periodical_report_status())
197 
198     async def __pre_destroy__(self):
199         self._upload_task.cancel()
200 
201     async def get_model_count(self) -> int:
202         return len(self._model_uid_to_model)
203 
204     def _choose_subpool(self) -> str:
205         min_running_model_count = None
206         target_subpool_address = None
207         for subpool_address in self._subpool_address_to_model_uids:
208             running_model_count = len(
209                 self._subpool_address_to_model_uids[subpool_address]
210             )
211             if (
212                 min_running_model_count is None
213                 or running_model_count < min_running_model_count
214             ):
215                 min_running_model_count = running_model_count
216                 target_subpool_address = subpool_address
217 
218         if target_subpool_address:
219             logger.debug(
220                 "Subpool selected: %s, running model count: %d",
221                 target_subpool_address,
222                 min_running_model_count,
223             )
224             return target_subpool_address
225 
226         raise RuntimeError("No available slot found")
227 
228     @log
229     async def launch_builtin_model(
230         self,
231         model_uid: str,
232         model_name: str,
233         model_size_in_billions: Optional[int],
234         model_format: Optional[str],
235         quantization: Optional[str],
236         **kwargs,
237     ) -> xo.ActorRefType["ModelActor"]:
238         assert model_uid not in self._model_uid_to_model
239 
240         from ..model import MODEL_FAMILIES
241 
242         for model_family in MODEL_FAMILIES:
243             model_spec = model_family.match(
244                 model_name=model_name,
245                 model_format=model_format,
246                 model_size_in_billions=model_size_in_billions,
247                 quantization=quantization,
248             )
249 
250             if model_spec is None:
251                 continue
252 
253             cls = model_family.cls
254             save_path = model_family.cache(
255                 model_spec.model_size_in_billions, model_spec.quantization
256             )
257             model = cls(model_uid, model_spec, save_path, kwargs)
258             subpool_address = self._choose_subpool()
259             model_ref = await xo.create_actor(
260                 ModelActor, address=subpool_address, uid=model_uid, model=model
261             )
262             await model_ref.load()
263             self._model_uid_to_model[model_uid] = model_ref
264             self._model_uid_to_model_spec[model_uid] = model_spec
265             self._subpool_address_to_model_uids[subpool_address].add(model_uid)
266             return model_ref
267 
268         raise ValueError(
269             f"Model not found, name: {model_name}, format: {model_format},"
270             f" size: {model_size_in_billions}, quantization: {quantization}"
271         )
272 
273     @log
274     async def terminate_model(self, model_uid: str):
275         assert model_uid in self._model_uid_to_model
276 
277         model_ref = self._model_uid_to_model[model_uid]
278         await xo.destroy_actor(model_ref)
279         del self._model_uid_to_model[model_uid]
280         del self._model_uid_to_model_spec[model_uid]
281         for subpool_address in self._subpool_address_to_model_uids:
282             if model_uid in self._subpool_address_to_model_uids[subpool_address]:
283                 self._subpool_address_to_model_uids[subpool_address].remove(model_uid)
284 
285     @log
286     async def list_models(self) -> List[Tuple[str, ModelSpec]]:
287         ret = []
288         for k, v in self._model_uid_to_model_spec.items():
289             ret.append((k, v))
290         return ret
291 
292     @log
293     async def get_model(self, model_uid: str) -> xo.ActorRefType["ModelActor"]:
294         return self._model_uid_to_model[model_uid]
295 
296     @log
297     async def describe_model(self, model_uid: str) -> ModelSpec:
298         return self._model_uid_to_model_spec[model_uid]
299 
300     async def report_status(self):
301         status = await asyncio.to_thread(gather_node_info)
302         await self._supervisor_ref.report_worker_status(self.address, status)
303 
304     async def _periodical_report_status(self):
305         while True:
306             try:
307                 await self.report_status()
308             except asyncio.CancelledError:  # pragma: no cover
309                 break
310             except RuntimeError as ex:  # pragma: no cover
311                 if "cannot schedule new futures" not in str(ex):
312                     # when atexit is triggered, the default pool might be shutdown
313                     # and to_thread will fail
314                     break
315             except (
316                 Exception
317             ) as ex:  # pragma: no cover  # noqa: E722  # nosec  # pylint: disable=bare-except
318                 logger.error(f"Failed to upload node info: {ex}")
319             try:
320                 await asyncio.sleep(DEFAULT_NODE_CHECK_INTERVAL)
321             except asyncio.CancelledError:  # pragma: no cover
322                 break
323
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/xinference/core/service.py b/xinference/core/service.py
--- a/xinference/core/service.py
+++ b/xinference/core/service.py
@@ -90,7 +90,6 @@
 
         raise RuntimeError("No available worker found")
 
-    @log
     async def launch_builtin_model(
         self,
         model_uid: str,
@@ -100,10 +99,22 @@
         quantization: Optional[str],
         **kwargs,
     ) -> xo.ActorRefType["ModelActor"]:
+        logger.debug(
+            (
+                f"Enter launch_builtin_model, model_uid: %s, model_name: %s, model_size: %s, "
+                f"model_format: %s, quantization: %s"
+            ),
+            model_uid,
+            model_name,
+            str(model_size_in_billions) if model_size_in_billions else "",
+            model_format,
+            quantization,
+        )
+
         assert model_uid not in self._model_uid_to_worker
 
         worker_ref = await self._choose_worker()
-        model_ref = await worker_ref.launch_builtin_model(
+        model_ref = yield worker_ref.launch_builtin_model(
             model_uid=model_uid,
             model_name=model_name,
             model_size_in_billions=model_size_in_billions,
@@ -111,9 +122,10 @@
             quantization=quantization,
             **kwargs,
         )
+        # TODO: not protected.
         self._model_uid_to_worker[model_uid] = worker_ref
 
-        return model_ref
+        raise xo.Return(model_ref)
 
     async def _check_dead_nodes(self):
         while True:
@@ -250,10 +262,13 @@
             if model_spec is None:
                 continue
 
-            cls = model_family.cls
-            save_path = model_family.cache(
-                model_spec.model_size_in_billions, model_spec.quantization
+            save_path = await asyncio.to_thread(
+                model_family.cache,
+                model_spec.model_size_in_billions,
+                model_spec.quantization,
             )
+
+            cls = model_family.cls
             model = cls(model_uid, model_spec, save_path, kwargs)
             subpool_address = self._choose_subpool()
             model_ref = await xo.create_actor(

{"golden_diff": "diff --git a/xinference/core/service.py b/xinference/core/service.py\n--- a/xinference/core/service.py\n+++ b/xinference/core/service.py\n@@ -90,7 +90,6 @@\n \n         raise RuntimeError(\"No available worker found\")\n \n-    @log\n     async def launch_builtin_model(\n         self,\n         model_uid: str,\n@@ -100,10 +99,22 @@\n         quantization: Optional[str],\n         **kwargs,\n     ) -> xo.ActorRefType[\"ModelActor\"]:\n+        logger.debug(\n+            (\n+                f\"Enter launch_builtin_model, model_uid: %s, model_name: %s, model_size: %s, \"\n+                f\"model_format: %s, quantization: %s\"\n+            ),\n+            model_uid,\n+            model_name,\n+            str(model_size_in_billions) if model_size_in_billions else \"\",\n+            model_format,\n+            quantization,\n+        )\n+\n         assert model_uid not in self._model_uid_to_worker\n \n         worker_ref = await self._choose_worker()\n-        model_ref = await worker_ref.launch_builtin_model(\n+        model_ref = yield worker_ref.launch_builtin_model(\n             model_uid=model_uid,\n             model_name=model_name,\n             model_size_in_billions=model_size_in_billions,\n@@ -111,9 +122,10 @@\n             quantization=quantization,\n             **kwargs,\n         )\n+        # TODO: not protected.\n         self._model_uid_to_worker[model_uid] = worker_ref\n \n-        return model_ref\n+        raise xo.Return(model_ref)\n \n     async def _check_dead_nodes(self):\n         while True:\n@@ -250,10 +262,13 @@\n             if model_spec is None:\n                 continue\n \n-            cls = model_family.cls\n-            save_path = model_family.cache(\n-                model_spec.model_size_in_billions, model_spec.quantization\n+            save_path = await asyncio.to_thread(\n+                model_family.cache,\n+                model_spec.model_size_in_billions,\n+                model_spec.quantization,\n             )\n+\n+            cls = model_family.cls\n             model = cls(model_uid, model_spec, save_path, kwargs)\n             subpool_address = self._choose_subpool()\n             model_ref = await xo.create_actor(\n", "issue": "BUG: worker timeout when downloading a model\n\n", "before_files": [{"content": "# Copyright 2022-2023 XProbe Inc.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport time\nfrom dataclasses import dataclass\nfrom logging import getLogger\nfrom typing import Callable, Dict, List, Optional, Set, Tuple\n\nimport xoscar as xo\n\nfrom ..core import ModelActor\nfrom ..model import ModelSpec\nfrom .resource import ResourceStatus, gather_node_info\n\nlogger = getLogger(__name__)\n\n\nDEFAULT_NODE_DEAD_TIMEOUT = 30\nDEFAULT_NODE_CHECK_INTERVAL = 1\n\n\ndef log(func: Callable):\n    # TODO: support non-async function\n    import time\n    from functools import wraps\n\n    @wraps(func)\n    async def wrapped(*args, **kwargs):\n        logger.debug(f\"Enter {func.__name__}, args: {args}, kwargs: {kwargs}\")\n        start = time.time()\n        ret = await func(*args, **kwargs)\n        logger.debug(\n            f\"Leave {func.__name__}, elapsed time: {int(time.time() - start)} ms\"\n        )\n        return ret\n\n    return wrapped\n\n\n@dataclass\nclass WorkerStatus:\n    update_time: float\n    status: Dict[str, ResourceStatus]\n\n\nclass SupervisorActor(xo.Actor):\n    def __init__(self):\n        super().__init__()\n        self._worker_address_to_worker: Dict[str, xo.ActorRefType[WorkerActor]] = {}\n        self._model_uid_to_worker: Dict[str, xo.ActorRefType[WorkerActor]] = {}\n        self._worker_status: Dict[str, WorkerStatus] = {}\n\n    @classmethod\n    def uid(cls) -> str:\n        return \"supervisor\"\n\n    async def __post_create__(self):\n        self._check_dead_nodes_task = asyncio.create_task(self._check_dead_nodes())\n\n    async def __pre_destroy__(self):\n        self._check_dead_nodes_task.cancel()\n\n    async def _choose_worker(self) -> xo.ActorRefType[\"WorkerActor\"]:\n        # TODO: better allocation strategy.\n        min_running_model_count = None\n        target_worker = None\n        for worker in self._worker_address_to_worker.values():\n            running_model_count = await worker.get_model_count()\n            if (\n                min_running_model_count is None\n                or running_model_count < min_running_model_count\n            ):\n                min_running_model_count = running_model_count\n                target_worker = worker\n\n        if target_worker:\n            return target_worker\n\n        raise RuntimeError(\"No available worker found\")\n\n    @log\n    async def launch_builtin_model(\n        self,\n        model_uid: str,\n        model_name: str,\n        model_size_in_billions: Optional[int],\n        model_format: Optional[str],\n        quantization: Optional[str],\n        **kwargs,\n    ) -> xo.ActorRefType[\"ModelActor\"]:\n        assert model_uid not in self._model_uid_to_worker\n\n        worker_ref = await self._choose_worker()\n        model_ref = await worker_ref.launch_builtin_model(\n            model_uid=model_uid,\n            model_name=model_name,\n            model_size_in_billions=model_size_in_billions,\n            model_format=model_format,\n            quantization=quantization,\n            **kwargs,\n        )\n        self._model_uid_to_worker[model_uid] = worker_ref\n\n        return model_ref\n\n    async def _check_dead_nodes(self):\n        while True:\n            for address, status in self._worker_status.items():\n                if time.time() - status.update_time > DEFAULT_NODE_DEAD_TIMEOUT:\n                    dead_models = []\n                    for model_uid in self._model_uid_to_worker:\n                        if self._model_uid_to_worker[model_uid].address == address:\n                            dead_models.append(model_uid)\n                    logger.error(\n                        \"Worker timeout. address: %s, influenced models: %s\",\n                        address,\n                        dead_models,\n                    )\n                    self._worker_status.pop(address)\n                    self._worker_address_to_worker.pop(address)\n            await asyncio.sleep(5)\n\n    @log\n    async def terminate_model(self, model_uid: str):\n        assert model_uid in self._model_uid_to_worker\n\n        worker_ref = self._model_uid_to_worker[model_uid]\n        await worker_ref.terminate_model(model_uid=model_uid)\n        del self._model_uid_to_worker[model_uid]\n\n    @log\n    async def get_model(self, model_uid: str) -> xo.ActorRefType[\"ModelActor\"]:\n        worker_ref = self._model_uid_to_worker[model_uid]\n        return await worker_ref.get_model(model_uid=model_uid)\n\n    @log\n    async def describe_model(self, model_uid: str):\n        worker_ref = self._model_uid_to_worker[model_uid]\n        return await worker_ref.describe_model(model_uid=model_uid)\n\n    @log\n    async def list_models(self) -> List[Tuple[str, ModelSpec]]:\n        ret = []\n        for worker in self._worker_address_to_worker.values():\n            ret.extend(await worker.list_models())\n        return ret\n\n    @log\n    async def add_worker(self, worker_address: str):\n        assert worker_address not in self._worker_address_to_worker\n\n        worker_ref = await xo.actor_ref(address=worker_address, uid=WorkerActor.uid())\n        self._worker_address_to_worker[worker_address] = worker_ref\n\n    async def report_worker_status(\n        self, worker_address: str, status: Dict[str, ResourceStatus]\n    ):\n        self._worker_status[worker_address] = WorkerStatus(\n            update_time=time.time(), status=status\n        )\n\n\nclass WorkerActor(xo.Actor):\n    def __init__(self, supervisor_address: str, subpool_addresses: List[str]):\n        super().__init__()\n        self._supervisor_address = supervisor_address\n        self._supervisor_ref = None\n        self._model_uid_to_model: Dict[str, xo.ActorRefType[\"ModelActor\"]] = {}\n        self._model_uid_to_model_spec: Dict[str, ModelSpec] = {}\n        self._subpool_address_to_model_uids: Dict[str, Set[str]] = dict(\n            [(subpool_address, set()) for subpool_address in subpool_addresses]\n        )\n        logger.debug(f\"Worker actor initialized with subpools: {subpool_addresses}\")\n\n    @classmethod\n    def uid(cls) -> str:\n        return \"worker\"\n\n    async def __post_create__(self):\n        self._supervisor_ref: xo.ActorRefType[\"SupervisorActor\"] = await xo.actor_ref(\n            address=self._supervisor_address, uid=SupervisorActor.uid()\n        )\n        await self._supervisor_ref.add_worker(self.address)\n        self._upload_task = asyncio.create_task(self._periodical_report_status())\n\n    async def __pre_destroy__(self):\n        self._upload_task.cancel()\n\n    async def get_model_count(self) -> int:\n        return len(self._model_uid_to_model)\n\n    def _choose_subpool(self) -> str:\n        min_running_model_count = None\n        target_subpool_address = None\n        for subpool_address in self._subpool_address_to_model_uids:\n            running_model_count = len(\n                self._subpool_address_to_model_uids[subpool_address]\n            )\n            if (\n                min_running_model_count is None\n                or running_model_count < min_running_model_count\n            ):\n                min_running_model_count = running_model_count\n                target_subpool_address = subpool_address\n\n        if target_subpool_address:\n            logger.debug(\n                \"Subpool selected: %s, running model count: %d\",\n                target_subpool_address,\n                min_running_model_count,\n            )\n            return target_subpool_address\n\n        raise RuntimeError(\"No available slot found\")\n\n    @log\n    async def launch_builtin_model(\n        self,\n        model_uid: str,\n        model_name: str,\n        model_size_in_billions: Optional[int],\n        model_format: Optional[str],\n        quantization: Optional[str],\n        **kwargs,\n    ) -> xo.ActorRefType[\"ModelActor\"]:\n        assert model_uid not in self._model_uid_to_model\n\n        from ..model import MODEL_FAMILIES\n\n        for model_family in MODEL_FAMILIES:\n            model_spec = model_family.match(\n                model_name=model_name,\n                model_format=model_format,\n                model_size_in_billions=model_size_in_billions,\n                quantization=quantization,\n            )\n\n            if model_spec is None:\n                continue\n\n            cls = model_family.cls\n            save_path = model_family.cache(\n                model_spec.model_size_in_billions, model_spec.quantization\n            )\n            model = cls(model_uid, model_spec, save_path, kwargs)\n            subpool_address = self._choose_subpool()\n            model_ref = await xo.create_actor(\n                ModelActor, address=subpool_address, uid=model_uid, model=model\n            )\n            await model_ref.load()\n            self._model_uid_to_model[model_uid] = model_ref\n            self._model_uid_to_model_spec[model_uid] = model_spec\n            self._subpool_address_to_model_uids[subpool_address].add(model_uid)\n            return model_ref\n\n        raise ValueError(\n            f\"Model not found, name: {model_name}, format: {model_format},\"\n            f\" size: {model_size_in_billions}, quantization: {quantization}\"\n        )\n\n    @log\n    async def terminate_model(self, model_uid: str):\n        assert model_uid in self._model_uid_to_model\n\n        model_ref = self._model_uid_to_model[model_uid]\n        await xo.destroy_actor(model_ref)\n        del self._model_uid_to_model[model_uid]\n        del self._model_uid_to_model_spec[model_uid]\n        for subpool_address in self._subpool_address_to_model_uids:\n            if model_uid in self._subpool_address_to_model_uids[subpool_address]:\n                self._subpool_address_to_model_uids[subpool_address].remove(model_uid)\n\n    @log\n    async def list_models(self) -> List[Tuple[str, ModelSpec]]:\n        ret = []\n        for k, v in self._model_uid_to_model_spec.items():\n            ret.append((k, v))\n        return ret\n\n    @log\n    async def get_model(self, model_uid: str) -> xo.ActorRefType[\"ModelActor\"]:\n        return self._model_uid_to_model[model_uid]\n\n    @log\n    async def describe_model(self, model_uid: str) -> ModelSpec:\n        return self._model_uid_to_model_spec[model_uid]\n\n    async def report_status(self):\n        status = await asyncio.to_thread(gather_node_info)\n        await self._supervisor_ref.report_worker_status(self.address, status)\n\n    async def _periodical_report_status(self):\n        while True:\n            try:\n                await self.report_status()\n            except asyncio.CancelledError:  # pragma: no cover\n                break\n            except RuntimeError as ex:  # pragma: no cover\n                if \"cannot schedule new futures\" not in str(ex):\n                    # when atexit is triggered, the default pool might be shutdown\n                    # and to_thread will fail\n                    break\n            except (\n                Exception\n            ) as ex:  # pragma: no cover  # noqa: E722  # nosec  # pylint: disable=bare-except\n                logger.error(f\"Failed to upload node info: {ex}\")\n            try:\n                await asyncio.sleep(DEFAULT_NODE_CHECK_INTERVAL)\n            except asyncio.CancelledError:  # pragma: no cover\n                break\n", "path": "xinference/core/service.py"}], "after_files": [{"content": "# Copyright 2022-2023 XProbe Inc.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport time\nfrom dataclasses import dataclass\nfrom logging import getLogger\nfrom typing import Callable, Dict, List, Optional, Set, Tuple\n\nimport xoscar as xo\n\nfrom ..core import ModelActor\nfrom ..model import ModelSpec\nfrom .resource import ResourceStatus, gather_node_info\n\nlogger = getLogger(__name__)\n\n\nDEFAULT_NODE_DEAD_TIMEOUT = 30\nDEFAULT_NODE_CHECK_INTERVAL = 1\n\n\ndef log(func: Callable):\n    # TODO: support non-async function\n    import time\n    from functools import wraps\n\n    @wraps(func)\n    async def wrapped(*args, **kwargs):\n        logger.debug(f\"Enter {func.__name__}, args: {args}, kwargs: {kwargs}\")\n        start = time.time()\n        ret = await func(*args, **kwargs)\n        logger.debug(\n            f\"Leave {func.__name__}, elapsed time: {int(time.time() - start)} ms\"\n        )\n        return ret\n\n    return wrapped\n\n\n@dataclass\nclass WorkerStatus:\n    update_time: float\n    status: Dict[str, ResourceStatus]\n\n\nclass SupervisorActor(xo.Actor):\n    def __init__(self):\n        super().__init__()\n        self._worker_address_to_worker: Dict[str, xo.ActorRefType[WorkerActor]] = {}\n        self._model_uid_to_worker: Dict[str, xo.ActorRefType[WorkerActor]] = {}\n        self._worker_status: Dict[str, WorkerStatus] = {}\n\n    @classmethod\n    def uid(cls) -> str:\n        return \"supervisor\"\n\n    async def __post_create__(self):\n        self._check_dead_nodes_task = asyncio.create_task(self._check_dead_nodes())\n\n    async def __pre_destroy__(self):\n        self._check_dead_nodes_task.cancel()\n\n    async def _choose_worker(self) -> xo.ActorRefType[\"WorkerActor\"]:\n        # TODO: better allocation strategy.\n        min_running_model_count = None\n        target_worker = None\n        for worker in self._worker_address_to_worker.values():\n            running_model_count = await worker.get_model_count()\n            if (\n                min_running_model_count is None\n                or running_model_count < min_running_model_count\n            ):\n                min_running_model_count = running_model_count\n                target_worker = worker\n\n        if target_worker:\n            return target_worker\n\n        raise RuntimeError(\"No available worker found\")\n\n    async def launch_builtin_model(\n        self,\n        model_uid: str,\n        model_name: str,\n        model_size_in_billions: Optional[int],\n        model_format: Optional[str],\n        quantization: Optional[str],\n        **kwargs,\n    ) -> xo.ActorRefType[\"ModelActor\"]:\n        logger.debug(\n            (\n                f\"Enter launch_builtin_model, model_uid: %s, model_name: %s, model_size: %s, \"\n                f\"model_format: %s, quantization: %s\"\n            ),\n            model_uid,\n            model_name,\n            str(model_size_in_billions) if model_size_in_billions else \"\",\n            model_format,\n            quantization,\n        )\n\n        assert model_uid not in self._model_uid_to_worker\n\n        worker_ref = await self._choose_worker()\n        model_ref = yield worker_ref.launch_builtin_model(\n            model_uid=model_uid,\n            model_name=model_name,\n            model_size_in_billions=model_size_in_billions,\n            model_format=model_format,\n            quantization=quantization,\n            **kwargs,\n        )\n        # TODO: not protected.\n        self._model_uid_to_worker[model_uid] = worker_ref\n\n        raise xo.Return(model_ref)\n\n    async def _check_dead_nodes(self):\n        while True:\n            for address, status in self._worker_status.items():\n                if time.time() - status.update_time > DEFAULT_NODE_DEAD_TIMEOUT:\n                    dead_models = []\n                    for model_uid in self._model_uid_to_worker:\n                        if self._model_uid_to_worker[model_uid].address == address:\n                            dead_models.append(model_uid)\n                    logger.error(\n                        \"Worker timeout. address: %s, influenced models: %s\",\n                        address,\n                        dead_models,\n                    )\n                    self._worker_status.pop(address)\n                    self._worker_address_to_worker.pop(address)\n            await asyncio.sleep(5)\n\n    @log\n    async def terminate_model(self, model_uid: str):\n        assert model_uid in self._model_uid_to_worker\n\n        worker_ref = self._model_uid_to_worker[model_uid]\n        await worker_ref.terminate_model(model_uid=model_uid)\n        del self._model_uid_to_worker[model_uid]\n\n    @log\n    async def get_model(self, model_uid: str) -> xo.ActorRefType[\"ModelActor\"]:\n        worker_ref = self._model_uid_to_worker[model_uid]\n        return await worker_ref.get_model(model_uid=model_uid)\n\n    @log\n    async def describe_model(self, model_uid: str):\n        worker_ref = self._model_uid_to_worker[model_uid]\n        return await worker_ref.describe_model(model_uid=model_uid)\n\n    @log\n    async def list_models(self) -> List[Tuple[str, ModelSpec]]:\n        ret = []\n        for worker in self._worker_address_to_worker.values():\n            ret.extend(await worker.list_models())\n        return ret\n\n    @log\n    async def add_worker(self, worker_address: str):\n        assert worker_address not in self._worker_address_to_worker\n\n        worker_ref = await xo.actor_ref(address=worker_address, uid=WorkerActor.uid())\n        self._worker_address_to_worker[worker_address] = worker_ref\n\n    async def report_worker_status(\n        self, worker_address: str, status: Dict[str, ResourceStatus]\n    ):\n        self._worker_status[worker_address] = WorkerStatus(\n            update_time=time.time(), status=status\n        )\n\n\nclass WorkerActor(xo.Actor):\n    def __init__(self, supervisor_address: str, subpool_addresses: List[str]):\n        super().__init__()\n        self._supervisor_address = supervisor_address\n        self._supervisor_ref = None\n        self._model_uid_to_model: Dict[str, xo.ActorRefType[\"ModelActor\"]] = {}\n        self._model_uid_to_model_spec: Dict[str, ModelSpec] = {}\n        self._subpool_address_to_model_uids: Dict[str, Set[str]] = dict(\n            [(subpool_address, set()) for subpool_address in subpool_addresses]\n        )\n        logger.debug(f\"Worker actor initialized with subpools: {subpool_addresses}\")\n\n    @classmethod\n    def uid(cls) -> str:\n        return \"worker\"\n\n    async def __post_create__(self):\n        self._supervisor_ref: xo.ActorRefType[\"SupervisorActor\"] = await xo.actor_ref(\n            address=self._supervisor_address, uid=SupervisorActor.uid()\n        )\n        await self._supervisor_ref.add_worker(self.address)\n        self._upload_task = asyncio.create_task(self._periodical_report_status())\n\n    async def __pre_destroy__(self):\n        self._upload_task.cancel()\n\n    async def get_model_count(self) -> int:\n        return len(self._model_uid_to_model)\n\n    def _choose_subpool(self) -> str:\n        min_running_model_count = None\n        target_subpool_address = None\n        for subpool_address in self._subpool_address_to_model_uids:\n            running_model_count = len(\n                self._subpool_address_to_model_uids[subpool_address]\n            )\n            if (\n                min_running_model_count is None\n                or running_model_count < min_running_model_count\n            ):\n                min_running_model_count = running_model_count\n                target_subpool_address = subpool_address\n\n        if target_subpool_address:\n            logger.debug(\n                \"Subpool selected: %s, running model count: %d\",\n                target_subpool_address,\n                min_running_model_count,\n            )\n            return target_subpool_address\n\n        raise RuntimeError(\"No available slot found\")\n\n    @log\n    async def launch_builtin_model(\n        self,\n        model_uid: str,\n        model_name: str,\n        model_size_in_billions: Optional[int],\n        model_format: Optional[str],\n        quantization: Optional[str],\n        **kwargs,\n    ) -> xo.ActorRefType[\"ModelActor\"]:\n        assert model_uid not in self._model_uid_to_model\n\n        from ..model import MODEL_FAMILIES\n\n        for model_family in MODEL_FAMILIES:\n            model_spec = model_family.match(\n                model_name=model_name,\n                model_format=model_format,\n                model_size_in_billions=model_size_in_billions,\n                quantization=quantization,\n            )\n\n            if model_spec is None:\n                continue\n\n            save_path = await asyncio.to_thread(\n                model_family.cache,\n                model_spec.model_size_in_billions,\n                model_spec.quantization,\n            )\n\n            cls = model_family.cls\n            model = cls(model_uid, model_spec, save_path, kwargs)\n            subpool_address = self._choose_subpool()\n            model_ref = await xo.create_actor(\n                ModelActor, address=subpool_address, uid=model_uid, model=model\n            )\n            await model_ref.load()\n            self._model_uid_to_model[model_uid] = model_ref\n            self._model_uid_to_model_spec[model_uid] = model_spec\n            self._subpool_address_to_model_uids[subpool_address].add(model_uid)\n            return model_ref\n\n        raise ValueError(\n            f\"Model not found, name: {model_name}, format: {model_format},\"\n            f\" size: {model_size_in_billions}, quantization: {quantization}\"\n        )\n\n    @log\n    async def terminate_model(self, model_uid: str):\n        assert model_uid in self._model_uid_to_model\n\n        model_ref = self._model_uid_to_model[model_uid]\n        await xo.destroy_actor(model_ref)\n        del self._model_uid_to_model[model_uid]\n        del self._model_uid_to_model_spec[model_uid]\n        for subpool_address in self._subpool_address_to_model_uids:\n            if model_uid in self._subpool_address_to_model_uids[subpool_address]:\n                self._subpool_address_to_model_uids[subpool_address].remove(model_uid)\n\n    @log\n    async def list_models(self) -> List[Tuple[str, ModelSpec]]:\n        ret = []\n        for k, v in self._model_uid_to_model_spec.items():\n            ret.append((k, v))\n        return ret\n\n    @log\n    async def get_model(self, model_uid: str) -> xo.ActorRefType[\"ModelActor\"]:\n        return self._model_uid_to_model[model_uid]\n\n    @log\n    async def describe_model(self, model_uid: str) -> ModelSpec:\n        return self._model_uid_to_model_spec[model_uid]\n\n    async def report_status(self):\n        status = await asyncio.to_thread(gather_node_info)\n        await self._supervisor_ref.report_worker_status(self.address, status)\n\n    async def _periodical_report_status(self):\n        while True:\n            try:\n                await self.report_status()\n            except asyncio.CancelledError:  # pragma: no cover\n                break\n            except RuntimeError as ex:  # pragma: no cover\n                if \"cannot schedule new futures\" not in str(ex):\n                    # when atexit is triggered, the default pool might be shutdown\n                    # and to_thread will fail\n                    break\n            except (\n                Exception\n            ) as ex:  # pragma: no cover  # noqa: E722  # nosec  # pylint: disable=bare-except\n                logger.error(f\"Failed to upload node info: {ex}\")\n            try:\n                await asyncio.sleep(DEFAULT_NODE_CHECK_INTERVAL)\n            except asyncio.CancelledError:  # pragma: no cover\n                break\n", "path": "xinference/core/service.py"}]}