We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
mlserver >= 1.3 doesn't work with MLflow
After the release of `mlserver==1.3.1`, one of our tests in `https://github.com/mlflow/mlflow` started failing with the following error. We're investigating the cause.

https://github.com/mlflow/mlflow/actions/runs/4716614587/jobs/8364498235?pr=8241#step:7:8138

```
2023-04-17 03:00:01,885 [mlserver.parallel] DEBUG - Starting response processing loop...
2023-04-17 03:00:01,894 [mlserver.rest] INFO - HTTP server running on http://0.0.0.0:8080/
2023-04-17 03:00:01,943 [mlserver.metrics] INFO - Metrics server running on http://0.0.0.0:8082/
2023-04-17 03:00:01,944 [mlserver.metrics] INFO - Prometheus scraping endpoint can be accessed on http://0.0.0.0:8082/metrics
2023-04-17 03:00:01,963 [mlserver.grpc] INFO - gRPC server running on http://0.0.0.0:8081/
INFO:     172.17.0.1:59636 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:59644 - "GET /ping HTTP/1.1" 404 Not Found
2023-04-17 03:00:03,801 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
2023-04-17 03:00:03,802 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
2023-04-17 03:00:03,804 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
INFO:     172.17.0.1:59646 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:59660 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:59674 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:59690 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:59702 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:59708 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:59712 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:59718 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:42498 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:42512 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:42528 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:42530 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:42536 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:42544 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:42546 - "GET /ping HTTP/1.1" 404 Not Found
INFO:     172.17.0.1:42552 - "GET /ping HTTP/1.1" 404 Not Found
...
```

The error above indicates that the `/ping` endpoint is not registered.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mlserver/parallel/registry.py`
Content:
```
1 import asyncio
2 import os
3 import shutil
4 
5 from typing import Optional, Dict, List
6 
7 from ..settings import ModelSettings
8 from ..utils import to_absolute_path
9 from ..model import MLModel
10 from ..settings import Settings
11 from ..env import Environment, compute_hash
12 from ..registry import model_initialiser
13 
14 from .errors import EnvironmentNotFound
15 from .logging import logger
16 from .pool import InferencePool, InferencePoolHook
17 
18 ENV_HASH_ATTR = "__env_hash__"
19 
20 
21 def _set_environment_hash(model: MLModel, env_hash: Optional[str]):
22     setattr(model, ENV_HASH_ATTR, env_hash)
23 
24 
25 def _get_environment_hash(model: MLModel) -> Optional[str]:
26     return getattr(model, ENV_HASH_ATTR, None)
27 
28 
29 def _get_env_tarball(model: MLModel) -> Optional[str]:
30     model_settings = model.settings
31     if model_settings.parameters is None:
32         return None
33 
34     env_tarball = model_settings.parameters.environment_tarball
35     if env_tarball is None:
36         return None
37 
38     return to_absolute_path(model_settings, env_tarball)
39 
40 
41 class InferencePoolRegistry:
42     """
43     Keeps track of the different inference pools loaded in the server.
44     Each inference pool will generally be used to load a different environment.
45     """
46 
47     def __init__(
48         self, settings: Settings, on_worker_stop: List[InferencePoolHook] = []
49     ):
50         self._settings = settings
51         self._on_worker_stop = on_worker_stop
52         self._default_pool = InferencePool(
53             self._settings, on_worker_stop=on_worker_stop
54         )
55         self._pools: Dict[str, InferencePool] = {}
56 
57         os.makedirs(self._settings.environments_dir, exist_ok=True)
58 
59     async def _get_or_create(self, model: MLModel) -> InferencePool:
60         env_tarball = _get_env_tarball(model)
61         if not env_tarball:
62             return self._default_pool
63 
64         env_hash = await compute_hash(env_tarball)
65         if env_hash in self._pools:
66             return self._pools[env_hash]
67 
68         env = await self._extract_tarball(env_hash, env_tarball)
69         pool = InferencePool(
70             self._settings, env=env, on_worker_stop=self._on_worker_stop
71         )
72         self._pools[env_hash] = pool
73         return pool
74 
75     async def _extract_tarball(self, env_hash: str, env_tarball: str) -> Environment:
76         env_path = self._get_env_path(env_hash)
77         if os.path.isdir(env_path):
78             # If env has already been extracted, use that
79             return Environment(env_path, env_hash)
80 
81         os.makedirs(env_path)
82         return await Environment.from_tarball(env_tarball, env_path, env_hash)
83 
84     def _get_env_path(self, env_hash: str) -> str:
85         return os.path.join(self._settings.environments_dir, env_hash)
86 
87     async def _find(self, model: MLModel) -> InferencePool:
88         env_hash = _get_environment_hash(model)
89         if not env_hash:
90             return self._default_pool
91 
92         if env_hash not in self._pools:
93             raise EnvironmentNotFound(model, env_hash)
94 
95         return self._pools[env_hash]
96 
97     def _should_load_model(self, model_settings: ModelSettings):
98         if model_settings.parallel_workers is not None:
99             logger.warning(
100                 "DEPRECATED!! The `parallel_workers` setting at the model-level "
101                 "has now been deprecated and moved "
102                 "to the top-level server "
103                 "settings. "
104                 "This field will be removed in MLServer 1.2.0. "
105                 "To access the new field, you can either update the "
106                 "`settings.json` file, or update the `MLSERVER_PARALLEL_WORKERS` "
107                 "environment variable. "
108                 f"The current value of the server-level's `parallel_workers` field is "
109                 f"'{self._settings.parallel_workers}'."
110             )
111 
112             # NOTE: This is a remnant from the previous architecture for parallel
113             # workers, where each worker had its own pool.
114             # For backwards compatibility, we will respect when a model disables
115             # parallel inference.
116             if model_settings.parallel_workers <= 0:
117                 return False
118 
119         if not self._settings.parallel_workers:
120             return False
121 
122         return True
123 
124     def model_initialiser(self, model_settings: ModelSettings) -> MLModel:
125         """
126         Used to initialise a model object in the ModelRegistry.
127         """
128         if not self._should_load_model(model_settings):
129             # If parallel inference should not be used, instantiate the model
130             # as normal.
131             return model_initialiser(model_settings)
132 
133         # Otherwise, return a dummy model for now and wait for the load_model
134         # hook to create the actual thing.
135         # This avoids instantiating the model's actual class within the
136         # main process.
137         return MLModel(model_settings)
138 
139     async def load_model(self, model: MLModel) -> MLModel:
140         if not self._should_load_model(model.settings):
141             # Skip load if model has disabled parallel workers
142             return model
143 
144         # TODO: If load fails, should we remove pool if empty?
145         pool = await self._get_or_create(model)
146         loaded = await pool.load_model(model)
147         _set_environment_hash(loaded, pool.env_hash)
148         return loaded
149 
150     async def reload_model(self, old_model: MLModel, new_model: MLModel) -> MLModel:
151         if not self._should_load_model(new_model.settings):
152             # TODO: What would happen if old_model had parallel inference
153             # enabled and is disabled in new_model (and viceversa)?
154             # Skip reload if model has disabled parallel workers
155             return new_model
156 
157         old_hash = _get_environment_hash(old_model)
158         new_pool = await self._get_or_create(new_model)
159 
160         loaded = await new_pool.reload_model(old_model, new_model)
161         _set_environment_hash(loaded, new_pool.env_hash)
162         if old_hash != new_pool.env_hash:
163             # Environment has changed in the new version, so unload the old one
164             await self.unload_model(old_model)
165 
166         return loaded
167 
168     async def unload_model(self, model: MLModel) -> MLModel:
169         if not self._should_load_model(model.settings):
170             # Skip unload if model has disabled parallel workers
171             return model
172 
173         pool = await self._find(model)
174         unloaded = await pool.unload_model(model)
175 
176         if pool != self._default_pool and pool.empty():
177             logger.info(f"Inference pool with hash '{pool.env_hash}' is now empty")
178             await self._close_pool(pool.env_hash)
179 
180         return unloaded
181 
182     async def close(self):
183         await asyncio.gather(
184             self._close_pool(None),
185             *[self._close_pool(env_hash) for env_hash in self._pools],
186         )
187 
188     async def _close_pool(self, env_hash: Optional[str] = None):
189         pool = self._default_pool
190         pool_name = "default inference pool"
191         if env_hash:
192             pool = self._pools[env_hash]
193             pool_name = f"inference pool with hash '{env_hash}'"
194 
195         logger.info(f"Waiting for shutdown of {pool_name}...")
196         await pool.close()
197         logger.info(f"Shutdown of {pool_name} complete")
198 
199         if env_hash:
200             del self._pools[env_hash]
201             env_path = self._get_env_path(env_hash)
202             shutil.rmtree(env_path)
203
```
--- 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/mlserver/parallel/registry.py b/mlserver/parallel/registry.py
--- a/mlserver/parallel/registry.py
+++ b/mlserver/parallel/registry.py
@@ -130,6 +130,12 @@
             # as normal.
             return model_initialiser(model_settings)
 
+        parameters = model_settings.parameters
+        if not parameters or not parameters.environment_tarball:
+            # If model is not using a custom environment, instantiate the model
+            # as normal.
+            return model_initialiser(model_settings)
+
         # Otherwise, return a dummy model for now and wait for the load_model
         # hook to create the actual thing.
         # This avoids instantiating the model's actual class within the

{"golden_diff": "diff --git a/mlserver/parallel/registry.py b/mlserver/parallel/registry.py\n--- a/mlserver/parallel/registry.py\n+++ b/mlserver/parallel/registry.py\n@@ -130,6 +130,12 @@\n             # as normal.\n             return model_initialiser(model_settings)\n \n+        parameters = model_settings.parameters\n+        if not parameters or not parameters.environment_tarball:\n+            # If model is not using a custom environment, instantiate the model\n+            # as normal.\n+            return model_initialiser(model_settings)\n+\n         # Otherwise, return a dummy model for now and wait for the load_model\n         # hook to create the actual thing.\n         # This avoids instantiating the model's actual class within the\n", "issue": "mlserver >= 1.3 doesn't work with MLflow\nAfter the release of `mlserver==1.3.1`, one of our tests in `https://github.com/mlflow/mlflow` started failing with the following error. We're investigating the cause.\r\n\r\nhttps://github.com/mlflow/mlflow/actions/runs/4716614587/jobs/8364498235?pr=8241#step:7:8138\r\n\r\n```\r\n2023-04-17 03:00:01,885 [mlserver.parallel] DEBUG - Starting response processing loop...\r\n2023-04-17 03:00:01,894 [mlserver.rest] INFO - HTTP server running on http://0.0.0.0:8080/\r\n2023-04-17 03:00:01,943 [mlserver.metrics] INFO - Metrics server running on http://0.0.0.0:8082/\r\n2023-04-17 03:00:01,944 [mlserver.metrics] INFO - Prometheus scraping endpoint can be accessed on http://0.0.0.0:8082/metrics\r\n2023-04-17 03:00:01,963 [mlserver.grpc] INFO - gRPC server running on http://0.0.0.0:8081/\r\nINFO:     172.17.0.1:59636 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:59644 - \"GET /ping HTTP/1.1\" 404 Not Found\r\n2023-04-17 03:00:03,801 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.\r\n2023-04-17 03:00:03,802 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.\r\n2023-04-17 03:00:03,804 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.\r\nINFO:     172.17.0.1:59646 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:59660 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:59674 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:59690 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:59702 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:59708 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:59712 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:59718 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:42498 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:42512 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:42528 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:42530 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:42536 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:42544 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:42546 - \"GET /ping HTTP/1.1\" 404 Not Found\r\nINFO:     172.17.0.1:42552 - \"GET /ping HTTP/1.1\" 404 Not Found\r\n...\r\n```\r\n\r\nThe error above indicates that the `/ping` endpoint is not registered.\n", "before_files": [{"content": "import asyncio\nimport os\nimport shutil\n\nfrom typing import Optional, Dict, List\n\nfrom ..settings import ModelSettings\nfrom ..utils import to_absolute_path\nfrom ..model import MLModel\nfrom ..settings import Settings\nfrom ..env import Environment, compute_hash\nfrom ..registry import model_initialiser\n\nfrom .errors import EnvironmentNotFound\nfrom .logging import logger\nfrom .pool import InferencePool, InferencePoolHook\n\nENV_HASH_ATTR = \"__env_hash__\"\n\n\ndef _set_environment_hash(model: MLModel, env_hash: Optional[str]):\n    setattr(model, ENV_HASH_ATTR, env_hash)\n\n\ndef _get_environment_hash(model: MLModel) -> Optional[str]:\n    return getattr(model, ENV_HASH_ATTR, None)\n\n\ndef _get_env_tarball(model: MLModel) -> Optional[str]:\n    model_settings = model.settings\n    if model_settings.parameters is None:\n        return None\n\n    env_tarball = model_settings.parameters.environment_tarball\n    if env_tarball is None:\n        return None\n\n    return to_absolute_path(model_settings, env_tarball)\n\n\nclass InferencePoolRegistry:\n    \"\"\"\n    Keeps track of the different inference pools loaded in the server.\n    Each inference pool will generally be used to load a different environment.\n    \"\"\"\n\n    def __init__(\n        self, settings: Settings, on_worker_stop: List[InferencePoolHook] = []\n    ):\n        self._settings = settings\n        self._on_worker_stop = on_worker_stop\n        self._default_pool = InferencePool(\n            self._settings, on_worker_stop=on_worker_stop\n        )\n        self._pools: Dict[str, InferencePool] = {}\n\n        os.makedirs(self._settings.environments_dir, exist_ok=True)\n\n    async def _get_or_create(self, model: MLModel) -> InferencePool:\n        env_tarball = _get_env_tarball(model)\n        if not env_tarball:\n            return self._default_pool\n\n        env_hash = await compute_hash(env_tarball)\n        if env_hash in self._pools:\n            return self._pools[env_hash]\n\n        env = await self._extract_tarball(env_hash, env_tarball)\n        pool = InferencePool(\n            self._settings, env=env, on_worker_stop=self._on_worker_stop\n        )\n        self._pools[env_hash] = pool\n        return pool\n\n    async def _extract_tarball(self, env_hash: str, env_tarball: str) -> Environment:\n        env_path = self._get_env_path(env_hash)\n        if os.path.isdir(env_path):\n            # If env has already been extracted, use that\n            return Environment(env_path, env_hash)\n\n        os.makedirs(env_path)\n        return await Environment.from_tarball(env_tarball, env_path, env_hash)\n\n    def _get_env_path(self, env_hash: str) -> str:\n        return os.path.join(self._settings.environments_dir, env_hash)\n\n    async def _find(self, model: MLModel) -> InferencePool:\n        env_hash = _get_environment_hash(model)\n        if not env_hash:\n            return self._default_pool\n\n        if env_hash not in self._pools:\n            raise EnvironmentNotFound(model, env_hash)\n\n        return self._pools[env_hash]\n\n    def _should_load_model(self, model_settings: ModelSettings):\n        if model_settings.parallel_workers is not None:\n            logger.warning(\n                \"DEPRECATED!! The `parallel_workers` setting at the model-level \"\n                \"has now been deprecated and moved \"\n                \"to the top-level server \"\n                \"settings. \"\n                \"This field will be removed in MLServer 1.2.0. \"\n                \"To access the new field, you can either update the \"\n                \"`settings.json` file, or update the `MLSERVER_PARALLEL_WORKERS` \"\n                \"environment variable. \"\n                f\"The current value of the server-level's `parallel_workers` field is \"\n                f\"'{self._settings.parallel_workers}'.\"\n            )\n\n            # NOTE: This is a remnant from the previous architecture for parallel\n            # workers, where each worker had its own pool.\n            # For backwards compatibility, we will respect when a model disables\n            # parallel inference.\n            if model_settings.parallel_workers <= 0:\n                return False\n\n        if not self._settings.parallel_workers:\n            return False\n\n        return True\n\n    def model_initialiser(self, model_settings: ModelSettings) -> MLModel:\n        \"\"\"\n        Used to initialise a model object in the ModelRegistry.\n        \"\"\"\n        if not self._should_load_model(model_settings):\n            # If parallel inference should not be used, instantiate the model\n            # as normal.\n            return model_initialiser(model_settings)\n\n        # Otherwise, return a dummy model for now and wait for the load_model\n        # hook to create the actual thing.\n        # This avoids instantiating the model's actual class within the\n        # main process.\n        return MLModel(model_settings)\n\n    async def load_model(self, model: MLModel) -> MLModel:\n        if not self._should_load_model(model.settings):\n            # Skip load if model has disabled parallel workers\n            return model\n\n        # TODO: If load fails, should we remove pool if empty?\n        pool = await self._get_or_create(model)\n        loaded = await pool.load_model(model)\n        _set_environment_hash(loaded, pool.env_hash)\n        return loaded\n\n    async def reload_model(self, old_model: MLModel, new_model: MLModel) -> MLModel:\n        if not self._should_load_model(new_model.settings):\n            # TODO: What would happen if old_model had parallel inference\n            # enabled and is disabled in new_model (and viceversa)?\n            # Skip reload if model has disabled parallel workers\n            return new_model\n\n        old_hash = _get_environment_hash(old_model)\n        new_pool = await self._get_or_create(new_model)\n\n        loaded = await new_pool.reload_model(old_model, new_model)\n        _set_environment_hash(loaded, new_pool.env_hash)\n        if old_hash != new_pool.env_hash:\n            # Environment has changed in the new version, so unload the old one\n            await self.unload_model(old_model)\n\n        return loaded\n\n    async def unload_model(self, model: MLModel) -> MLModel:\n        if not self._should_load_model(model.settings):\n            # Skip unload if model has disabled parallel workers\n            return model\n\n        pool = await self._find(model)\n        unloaded = await pool.unload_model(model)\n\n        if pool != self._default_pool and pool.empty():\n            logger.info(f\"Inference pool with hash '{pool.env_hash}' is now empty\")\n            await self._close_pool(pool.env_hash)\n\n        return unloaded\n\n    async def close(self):\n        await asyncio.gather(\n            self._close_pool(None),\n            *[self._close_pool(env_hash) for env_hash in self._pools],\n        )\n\n    async def _close_pool(self, env_hash: Optional[str] = None):\n        pool = self._default_pool\n        pool_name = \"default inference pool\"\n        if env_hash:\n            pool = self._pools[env_hash]\n            pool_name = f\"inference pool with hash '{env_hash}'\"\n\n        logger.info(f\"Waiting for shutdown of {pool_name}...\")\n        await pool.close()\n        logger.info(f\"Shutdown of {pool_name} complete\")\n\n        if env_hash:\n            del self._pools[env_hash]\n            env_path = self._get_env_path(env_hash)\n            shutil.rmtree(env_path)\n", "path": "mlserver/parallel/registry.py"}], "after_files": [{"content": "import asyncio\nimport os\nimport shutil\n\nfrom typing import Optional, Dict, List\n\nfrom ..settings import ModelSettings\nfrom ..utils import to_absolute_path\nfrom ..model import MLModel\nfrom ..settings import Settings\nfrom ..env import Environment, compute_hash\nfrom ..registry import model_initialiser\n\nfrom .errors import EnvironmentNotFound\nfrom .logging import logger\nfrom .pool import InferencePool, InferencePoolHook\n\nENV_HASH_ATTR = \"__env_hash__\"\n\n\ndef _set_environment_hash(model: MLModel, env_hash: Optional[str]):\n    setattr(model, ENV_HASH_ATTR, env_hash)\n\n\ndef _get_environment_hash(model: MLModel) -> Optional[str]:\n    return getattr(model, ENV_HASH_ATTR, None)\n\n\ndef _get_env_tarball(model: MLModel) -> Optional[str]:\n    model_settings = model.settings\n    if model_settings.parameters is None:\n        return None\n\n    env_tarball = model_settings.parameters.environment_tarball\n    if env_tarball is None:\n        return None\n\n    return to_absolute_path(model_settings, env_tarball)\n\n\nclass InferencePoolRegistry:\n    \"\"\"\n    Keeps track of the different inference pools loaded in the server.\n    Each inference pool will generally be used to load a different environment.\n    \"\"\"\n\n    def __init__(\n        self, settings: Settings, on_worker_stop: List[InferencePoolHook] = []\n    ):\n        self._settings = settings\n        self._on_worker_stop = on_worker_stop\n        self._default_pool = InferencePool(\n            self._settings, on_worker_stop=on_worker_stop\n        )\n        self._pools: Dict[str, InferencePool] = {}\n\n        os.makedirs(self._settings.environments_dir, exist_ok=True)\n\n    async def _get_or_create(self, model: MLModel) -> InferencePool:\n        env_tarball = _get_env_tarball(model)\n        if not env_tarball:\n            return self._default_pool\n\n        env_hash = await compute_hash(env_tarball)\n        if env_hash in self._pools:\n            return self._pools[env_hash]\n\n        env = await self._extract_tarball(env_hash, env_tarball)\n        pool = InferencePool(\n            self._settings, env=env, on_worker_stop=self._on_worker_stop\n        )\n        self._pools[env_hash] = pool\n        return pool\n\n    async def _extract_tarball(self, env_hash: str, env_tarball: str) -> Environment:\n        env_path = self._get_env_path(env_hash)\n        if os.path.isdir(env_path):\n            # If env has already been extracted, use that\n            return Environment(env_path, env_hash)\n\n        os.makedirs(env_path)\n        return await Environment.from_tarball(env_tarball, env_path, env_hash)\n\n    def _get_env_path(self, env_hash: str) -> str:\n        return os.path.join(self._settings.environments_dir, env_hash)\n\n    async def _find(self, model: MLModel) -> InferencePool:\n        env_hash = _get_environment_hash(model)\n        if not env_hash:\n            return self._default_pool\n\n        if env_hash not in self._pools:\n            raise EnvironmentNotFound(model, env_hash)\n\n        return self._pools[env_hash]\n\n    def _should_load_model(self, model_settings: ModelSettings):\n        if model_settings.parallel_workers is not None:\n            logger.warning(\n                \"DEPRECATED!! The `parallel_workers` setting at the model-level \"\n                \"has now been deprecated and moved \"\n                \"to the top-level server \"\n                \"settings. \"\n                \"This field will be removed in MLServer 1.2.0. \"\n                \"To access the new field, you can either update the \"\n                \"`settings.json` file, or update the `MLSERVER_PARALLEL_WORKERS` \"\n                \"environment variable. \"\n                f\"The current value of the server-level's `parallel_workers` field is \"\n                f\"'{self._settings.parallel_workers}'.\"\n            )\n\n            # NOTE: This is a remnant from the previous architecture for parallel\n            # workers, where each worker had its own pool.\n            # For backwards compatibility, we will respect when a model disables\n            # parallel inference.\n            if model_settings.parallel_workers <= 0:\n                return False\n\n        if not self._settings.parallel_workers:\n            return False\n\n        return True\n\n    def model_initialiser(self, model_settings: ModelSettings) -> MLModel:\n        \"\"\"\n        Used to initialise a model object in the ModelRegistry.\n        \"\"\"\n        if not self._should_load_model(model_settings):\n            # If parallel inference should not be used, instantiate the model\n            # as normal.\n            return model_initialiser(model_settings)\n\n        parameters = model_settings.parameters\n        if not parameters or not parameters.environment_tarball:\n            # If model is not using a custom environment, instantiate the model\n            # as normal.\n            return model_initialiser(model_settings)\n\n        # Otherwise, return a dummy model for now and wait for the load_model\n        # hook to create the actual thing.\n        # This avoids instantiating the model's actual class within the\n        # main process.\n        return MLModel(model_settings)\n\n    async def load_model(self, model: MLModel) -> MLModel:\n        if not self._should_load_model(model.settings):\n            # Skip load if model has disabled parallel workers\n            return model\n\n        # TODO: If load fails, should we remove pool if empty?\n        pool = await self._get_or_create(model)\n        loaded = await pool.load_model(model)\n        _set_environment_hash(loaded, pool.env_hash)\n        return loaded\n\n    async def reload_model(self, old_model: MLModel, new_model: MLModel) -> MLModel:\n        if not self._should_load_model(new_model.settings):\n            # TODO: What would happen if old_model had parallel inference\n            # enabled and is disabled in new_model (and viceversa)?\n            # Skip reload if model has disabled parallel workers\n            return new_model\n\n        old_hash = _get_environment_hash(old_model)\n        new_pool = await self._get_or_create(new_model)\n\n        loaded = await new_pool.reload_model(old_model, new_model)\n        _set_environment_hash(loaded, new_pool.env_hash)\n        if old_hash != new_pool.env_hash:\n            # Environment has changed in the new version, so unload the old one\n            await self.unload_model(old_model)\n\n        return loaded\n\n    async def unload_model(self, model: MLModel) -> MLModel:\n        if not self._should_load_model(model.settings):\n            # Skip unload if model has disabled parallel workers\n            return model\n\n        pool = await self._find(model)\n        unloaded = await pool.unload_model(model)\n\n        if pool != self._default_pool and pool.empty():\n            logger.info(f\"Inference pool with hash '{pool.env_hash}' is now empty\")\n            await self._close_pool(pool.env_hash)\n\n        return unloaded\n\n    async def close(self):\n        await asyncio.gather(\n            self._close_pool(None),\n            *[self._close_pool(env_hash) for env_hash in self._pools],\n        )\n\n    async def _close_pool(self, env_hash: Optional[str] = None):\n        pool = self._default_pool\n        pool_name = \"default inference pool\"\n        if env_hash:\n            pool = self._pools[env_hash]\n            pool_name = f\"inference pool with hash '{env_hash}'\"\n\n        logger.info(f\"Waiting for shutdown of {pool_name}...\")\n        await pool.close()\n        logger.info(f\"Shutdown of {pool_name} complete\")\n\n        if env_hash:\n            del self._pools[env_hash]\n            env_path = self._get_env_path(env_hash)\n            shutil.rmtree(env_path)\n", "path": "mlserver/parallel/registry.py"}]}