Repository: SeldonIO/MLServer

Issue:
Expose `unload` api  in `MLModel` 
Currently there is no explicit user-defined unload utility that the user can define. This is required in cases where the user wants to do explicit clean up when the mlserver model is unloaded. e.g. closing connections / releasing resources etc. 


Discussion:
Based on https://github.com/SeldonIO/MLServer/issues/986#issuecomment-1587422911, as part of this issue we could also make sure that GPU mem is freed on the HF runtime.

Resolving patch:
diff --git a/mlserver/model.py b/mlserver/model.py
--- a/mlserver/model.py
+++ b/mlserver/model.py
@@ -79,6 +79,20 @@ async def predict(self, payload: InferenceRequest) -> InferenceResponse:
         """
         raise NotImplementedError("predict() method not implemented")
 
+    async def unload(self) -> bool:
+        """
+        Method responsible for unloading the model, freeing any resources (e.g.
+        CPU memory, GPU memory, etc.).
+        This method will be called on each of the parallel workers (when
+        :doc:`parallel inference </user-guide/parallel-inference>`) is
+        enabled).
+        A return value of ``True`` will mean the model is now unloaded.
+
+        **This method should be overriden to implement your custom unload
+        logic.**
+        """
+        return True
+
     @property
     def name(self) -> str:
         """
diff --git a/mlserver/registry.py b/mlserver/registry.py
--- a/mlserver/registry.py
+++ b/mlserver/registry.py
@@ -186,6 +186,7 @@ async def _reload_model(self, old_model: MLModel, new_model: MLModel):
         if old_model == self.default:
             self._clear_default()
 
+        old_model.ready = not await old_model.unload()
         logger.info(f"Reloaded model '{new_model.name}' succesfully.")
 
     async def unload(self):
@@ -221,6 +222,8 @@ async def _unload_model(self, model: MLModel):
         if model == self.default:
             self._clear_default()
 
+        model.ready = not await model.unload()
+
     def _find_model(self, version: Optional[str] = None) -> Optional[MLModel]:
         if version:
             if version not in self._versions:
diff --git a/runtimes/huggingface/mlserver_huggingface/runtime.py b/runtimes/huggingface/mlserver_huggingface/runtime.py
--- a/runtimes/huggingface/mlserver_huggingface/runtime.py
+++ b/runtimes/huggingface/mlserver_huggingface/runtime.py
@@ -1,4 +1,5 @@
 import asyncio
+import torch
 
 from mlserver.model import MLModel
 from mlserver.settings import ModelSettings
@@ -50,6 +51,21 @@ async def predict(self, payload: InferenceRequest) -> InferenceResponse:
             payload=prediction, default_codec=HuggingfaceRequestCodec
         )
 
+    async def unload(self) -> bool:
+        # TODO: Free up Tensorflow's GPU memory
+        is_torch = self._model.framework == "pt"
+        if not is_torch:
+            return True
+
+        uses_gpu = torch.cuda.is_available() and self._model.device != -1
+        if not uses_gpu:
+            # Nothing to free
+            return True
+
+        # Free up Torch's GPU memory
+        torch.cuda.empty_cache()
+        return True
+
     def _merge_metadata(self) -> None:
         meta = METADATA.get(self.hf_settings.task)
         if meta:


Test patch:
diff --git a/runtimes/huggingface/tests/test_runtime.py b/runtimes/huggingface/tests/test_runtime.py
--- a/runtimes/huggingface/tests/test_runtime.py
+++ b/runtimes/huggingface/tests/test_runtime.py
@@ -27,6 +27,14 @@ async def test_load(future_runtime: Awaitable[HuggingFaceRuntime]):
     assert isinstance(runtime._model, QuestionAnsweringPipeline)
 
 
+async def test_unload(future_runtime: Awaitable[HuggingFaceRuntime]):
+    runtime = await future_runtime
+    assert runtime.ready
+
+    unloaded = await runtime.unload()
+    assert unloaded
+
+
 async def test_infer(
     future_runtime: Awaitable[HuggingFaceRuntime], inference_request: InferenceRequest
 ):
diff --git a/tests/test_registry.py b/tests/test_registry.py
--- a/tests/test_registry.py
+++ b/tests/test_registry.py
@@ -97,6 +97,7 @@ async def test_reload_model(
     assert new_model != existing_model
     assert new_model == reloaded_model
     assert reloaded_model.ready
+    assert not existing_model.ready
 
     for callback in model_registry._on_model_load:
         callback.assert_not_called()
@@ -138,6 +139,14 @@ async def test_load_multi_version(
         callback.assert_not_called_with(existing_model)
 
 
+async def test_unload(model_registry: MultiModelRegistry, sum_model: MLModel, mocker):
+    spy = mocker.spy(sum_model, "unload")
+    await model_registry.unload(sum_model.name)
+
+    assert not sum_model.ready
+    spy.assert_called_once()
+
+
 @pytest.mark.parametrize(
     "versions_to_unload",
     [