Repository: caikit/caikit
Pull request #582 (2023-11-14 16:23:47)

Add HTTP Request Aborter
<!--  Thanks for sending a pull request!  Here are some tips for you:
1. Make sure to read the Contributing Guide before submitting your PR: https://github.com/caikit/caikit/blob/main/CONTRIBUTING.md
2. If this PR closes another issue, add 'closes #<issue number>' somewhere in the PR summary. GitHub will automatically close that issue when this PR gets merged. Alternatively, adding 'refs #<issue number>' will not close the issue, but help provide the reviewer more context.-->

**What this PR does / why we need it**:
This PR supports #575 by adding a new `RequestAborter` class to abort in-progress HTTP requests. This allows the HTTP server to stop work when a client is no longer listening for the response. This is especially important for long running inference requests in-which the client might timeout. 

**Special notes for your reviewer**:
I put the `RequestAborter` implementation in the `work_management` module but it might make more sense keeping it in `http_server`.

**If applicable**:
- [ ] this PR contains documentation
- [ ] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/runtime/http_server/http_server.py b/caikit/runtime/http_server/http_server.py
index 259f045d..91a3b0d1 100644
--- a/caikit/runtime/http_server/http_server.py
+++ b/caikit/runtime/http_server/http_server.py
@@ -17,7 +17,7 @@ The server is responsible for binding caikit workloads to a consistent REST/SSE
 API based on the task definitions available at boot.
 """
 # Standard
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from enum import Enum
 from functools import partial
@@ -56,6 +56,7 @@ from .pydantic_wrapper import (
     pydantic_from_request,
     pydantic_to_dataobject,
 )
+from .request_aborter import HttpRequestAborter
 from .utils import convert_json_schema_to_multipart, flatten_json_schema
 from caikit.config import get_config
 from caikit.core.data_model import DataBase
@@ -460,18 +461,26 @@ class RuntimeHTTPServer(RuntimeServerBase):
                     model_id
                 )
 
-                # TODO: use `async_wrap_*`?
-                call = partial(
-                    self.global_predict_servicer.predict_model,
-                    model_id=model_id,
-                    request_name=rpc.request.name,
-                    inference_func_name=model.get_inference_signature(
-                        output_streaming=False, input_streaming=False, task=rpc.task
-                    ).method_name,
-                    **request_params,
+                aborter_context = (
+                    HttpRequestAborter(context)
+                    if get_config().runtime.use_abortable_threads
+                    else nullcontext()
                 )
-                result = await loop.run_in_executor(None, call)
-                log.debug4("Response from model %s is %s", model_id, result)
+
+                with aborter_context as aborter:
+                    # TODO: use `async_wrap_*`?
+                    call = partial(
+                        self.global_predict_servicer.predict_model,
+                        model_id=model_id,
+                        request_name=rpc.request.name,
+                        inference_func_name=model.get_inference_signature(
+                            output_streaming=False, input_streaming=False, task=rpc.task
+                        ).method_name,
+                        aborter=aborter,
+                        **request_params,
+                    )
+                    result = await loop.run_in_executor(None, call)
+                    log.debug4("Response from model %s is %s", model_id, result)
 
                 if response_data_object.supports_file_operations:
                     return self._format_file_response(result)
@@ -526,20 +535,30 @@ class RuntimeHTTPServer(RuntimeServerBase):
                         model_id
                     )
 
-                    log.debug("In stream generator for %s", rpc.name)
-                    async for result in async_wrap_iter(
-                        self.global_predict_servicer.predict_model(
-                            model_id=model_id,
-                            request_name=rpc.request.name,
-                            inference_func_name=model.get_inference_signature(
-                                output_streaming=True, input_streaming=False
-                            ).method_name,
-                            **request_params,
-                        )
-                    ):
-                        yield ServerSentEvent(
-                            data=result.to_json(), event=StreamEventTypes.MESSAGE.value
-                        )
+                    aborter_context = (
+                        HttpRequestAborter(context)
+                        if get_config().runtime.use_abortable_threads
+                        else nullcontext()
+                    )
+
+                    with aborter_context as aborter:
+                        log.debug("In stream generator for %s", rpc.name)
+                        async for result in async_wrap_iter(
+                            self.global_predict_servicer.predict_model(
+                                model_id=model_id,
+                                request_name=rpc.request.name,
+                                inference_func_name=model.get_inference_signature(
+                                    output_streaming=True, input_streaming=False
+                                ).method_name,
+                                aborter=aborter,
+                                **request_params,
+                            )
+                        ):
+                            yield ServerSentEvent(
+                                data=result.to_json(),
+                                event=StreamEventTypes.MESSAGE.value,
+                            )
+
                     return
                 except HTTPException as err:
                     raise err
@@ -710,8 +729,8 @@ class RuntimeHTTPServer(RuntimeServerBase):
         return {
             "requestBody": {
                 "content": {
-                    "multipart/form-data": {"schema": multipart_schema},
                     "application/json": {"schema": parsed_schema},
+                    "multipart/form-data": {"schema": multipart_schema},
                 },
                 "required": True,
             }
diff --git a/caikit/runtime/http_server/request_aborter.py b/caikit/runtime/http_server/request_aborter.py
new file mode 100644
index 00000000..49030158
--- /dev/null
+++ b/caikit/runtime/http_server/request_aborter.py
@@ -0,0 +1,116 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This module helps us know when a HTTP client call is cancelled, and we need to stop or undo work
+"""
+# Standard
+from typing import Optional
+import asyncio
+import threading
+
+# Third Party
+from fastapi import Request
+
+# First Party
+import alog
+
+# Local
+from caikit.runtime.work_management.abortable_action import ActionAborter
+
+log = alog.use_channel("REQUEST-ABORTER")
+
+
+class HttpRequestAborter(ActionAborter):
+    """
+    In order to actually interrupt threads doing the work, events can be registered with an
+    instance of this class in order ton receive notification on request disconnection. This allows
+    work to be terminated when a client time's out or stops listening.
+
+    IFF the client request has been terminated, `must_abort` will return True.
+    """
+
+    def __init__(
+        self,
+        context: Request,
+        loop: Optional[asyncio.AbstractEventLoop] = None,
+        poll_time: Optional[float] = 0.25,
+    ):
+        """Initialize a Aborter and start the watch loop
+
+        Args:
+            context: starlette.Request
+               The HTTP Request to be watched
+            loop: Optional[asyncio.AbstractEventLoop]
+               The asyncio loop to run tasks on. If not provided use the running loop
+           poll_time: Optional[int]
+               The time between disconnect checks
+        """
+
+        self.context = context
+        self.event_loop = loop or asyncio.get_running_loop()
+        self.poll_time = poll_time
+
+        # Set initial
+        self.is_terminated = threading.Event()
+        # Add an empty list for event variables that will be notified on termination
+        self.events: list[threading.Event] = []
+
+        # Start request aborter task. Hold onto a reference of the task to ensure
+        # it isn't garbage collected
+        log.debug("<RUN81824293>", "Watching for request disconnect")
+        self.task = self.event_loop.create_task(self.watch_for_disconnect())
+
+    def __enter__(self):
+        """Helper function to enable context manager support"""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Automatically abort aborter when exiting contextmanager"""
+        self.abort()
+
+    async def watch_for_disconnect(self):
+        """Wait for a context to be disconnected"""
+
+        while True:
+            # Short circuit incase thread terminated externally
+            if self.is_terminated.is_set():
+                log.debug3(
+                    "<RUN81824293>", "RequestAborter has already been terminated"
+                )
+                return
+
+            is_disconnected = await self.context.is_disconnected()
+            if is_disconnected:
+                log.debug("<RUN81824293>", "Client disconnected, terminating action")
+                self.is_terminated.set()
+                for event in self.events:
+                    event.set()
+                return
+            log.debug4("<RUN81824293>", "Client still connected, sleeping aborter")
+            await asyncio.sleep(self.poll_time)
+
+    def abort(self):
+        """Helper function to stop aborter before the request has terminated"""
+        self.is_terminated.set()
+
+    def must_abort(self):
+        return self.is_terminated.is_set()
+
+    def add_event(self, event: threading.Event):
+        self.events.append(event)
+
+        # Sanity check: If we have already terminated, notify anything waiting on this condition
+        if self.must_abort():
+            event.set()
diff --git a/caikit/runtime/model_management/model_manager.py b/caikit/runtime/model_management/model_manager.py
index ba94cda9..c3d0c46b 100644
--- a/caikit/runtime/model_management/model_manager.py
+++ b/caikit/runtime/model_management/model_manager.py
@@ -158,7 +158,7 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
         wait: bool = True,
         aborter: Optional[ActionAborter] = None,
         retries: Optional[int] = None,
-    ) -> LoadedModel:
+    ) -> int:
         """Load a model using model_path (in Cloud Object Storage) & give it a model ID
         Args:
             model_id (str):  Model ID string for the model to load.
@@ -166,11 +166,9 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
             model_type (str): Type of the model to load.
             wait (bool): Wait for the model to finish loading
             aborter (Optional[ActionAborter]): The aborter to use for this load
-            retries (Optional[int]): Number of times to retry on load failure
-            return_loaded_model (bool): Return the LoadedModel instance instead
-                of the size
+            retries: Optional[int]: Number of times to retry on load failure
         Returns:
-            model (LoadedModel): The LoadedModel instance
+            Model_size (int) : Size of the loaded model in bytes
         """
         with LOAD_MODEL_DURATION_SUMMARY.labels(model_type=model_type).time():
 
@@ -181,7 +179,7 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
             model = self.loaded_models.get(model_id)
             if model is not None:
                 log.debug("Model '%s' is already loaded", model_id)
-                return model
+                return model.size()
 
             # Grab the mutation lock and load the model if needed
             with self._loaded_models_lock:
@@ -226,8 +224,8 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
             if wait:
                 model.wait()
 
-            # Return the loaded model handle
-            return model
+            # Return the model's size
+            return model.size()
 
     def sync_local_models(self, wait: bool = False):
         """Sync in-memory models with models in the configured local_model_dir
@@ -378,8 +376,8 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
             )
 
         # Now retrieve the model and fall back to lazy loading
-        loaded_model = self.loaded_models.get(model_id)
-        if not loaded_model and self._lazy_load_local_models:
+        model_loaded = model_id in self.loaded_models
+        if not model_loaded and self._lazy_load_local_models:
             local_model_path = os.path.join(self._local_models_dir, model_id)
             log.debug2(
                 "Lazy loading local model %s from %s", model_id, local_model_path
@@ -391,16 +389,17 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
             if not os.path.exists(local_model_path):
                 log.debug2("Attempting to load ephemeral model %s", model_id)
                 local_model_path = model_id
-            loaded_model = self.load_model(
+            self.load_model(
                 model_id=model_id,
                 local_model_path=local_model_path,
                 model_type=self._LOCAL_MODEL_TYPE,
                 wait=True,
                 retries=get_config().runtime.lazy_load_retries,
             )
+            model_loaded = True
 
         # If still not loaded, there's nothing to find, so raise NOT_FOUND
-        if not loaded_model:
+        if not model_loaded:
             msg = f"Model '{model_id}' not loaded"
             log.debug(
                 {"log_code": "<RUN61105243D>", "message": msg, "model_id": model_id}
@@ -411,7 +410,7 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
 
         # NOTE: If the model is partially loaded, this call will wait on the
         #   model future in the LoadedModel
-        return loaded_model.model()
+        return self.loaded_models[model_id].model()
 
     ## Implementation Details ##
 
@@ -444,8 +443,10 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
         unload_models = [
             model_id
             for model_id, loaded_model in self.loaded_models.items()
-            if loaded_model.path().startswith(self._local_models_dir)
-            and not os.path.exists(loaded_model.path())
+            if model_id not in disk_models
+            and loaded_model.path().startswith(
+                self._local_models_dir,
+            )
         ]
         log.debug("Unloaded local models: %s", unload_models)
 
@@ -482,8 +483,8 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
                 try:
                     loaded_model.wait()
                 except CaikitRuntimeException as err:
-                    log.debug(
-                        "<RUN56627485D>",
+                    log.warning(
+                        "<RUN56627485W>",
                         "Failed to load model %s: %s",
                         model_id,
                         repr(err),
diff --git a/caikit/runtime/servicers/info_servicer.py b/caikit/runtime/servicers/info_servicer.py
index b4974dec..6d261edb 100644
--- a/caikit/runtime/servicers/info_servicer.py
+++ b/caikit/runtime/servicers/info_servicer.py
@@ -46,14 +46,8 @@ class InfoServicer:
 
     def get_runtime_info_impl(self) -> RuntimeInfoResponse:
         """Get information on versions of libraries and server from config"""
+        python_packages = {}
         config_version_info = get_config().runtime.version_info or {}
-        python_packages = {
-            package: version
-            for package, version in config_version_info.get(
-                "python_packages", {}
-            ).items()
-            if package != "all"
-        }
         all_packages = (config_version_info.get("python_packages") or {}).get("all")
 
         for lib, dist_names in importlib_metadata.packages_distributions().items():
diff --git a/caikit/runtime/servicers/model_runtime_servicer.py b/caikit/runtime/servicers/model_runtime_servicer.py
index 340c53a5..cc62b5eb 100644
--- a/caikit/runtime/servicers/model_runtime_servicer.py
+++ b/caikit/runtime/servicers/model_runtime_servicer.py
@@ -66,12 +66,11 @@ class ModelRuntimeServicerImpl(model_runtime_pb2_grpc.ModelRuntimeServicer):
                     request.modelType,
                     aborter=aborter,
                 )
-                loaded_model = work.do()
+                model_size = work.do()
             else:
-                loaded_model = self.model_manager.load_model(
+                model_size = self.model_manager.load_model(
                     request.modelId, request.modelPath, request.modelType
                 )
-            model_size = loaded_model.size()
 
             log.info(
                 {


Test patch:
diff --git a/tests/runtime/http_server/test_http_server.py b/tests/runtime/http_server/test_http_server.py
index f4dbd066..b99fb0c9 100644
--- a/tests/runtime/http_server/test_http_server.py
+++ b/tests/runtime/http_server/test_http_server.py
@@ -484,6 +484,20 @@ def test_inference_other_task(other_task_model_id, client):
     assert json_response["farewell"] == "goodbye: world 42 times"
 
 
+def test_inference_no_aborter(other_task_model_id, client):
+    """Simple check that we can ping a model"""
+    with temp_config(
+        {"runtime": {"use_abortable_threads": False}}, merge_strategy="merge"
+    ):
+        json_input = {"model_id": other_task_model_id, "inputs": {"name": "world"}}
+        response = client.post(
+            f"/api/v1/task/other",
+            json=json_input,
+        )
+        json_response = json.loads(response.content.decode(response.default_encoding))
+        assert response.status_code == 200, json_response
+
+
 def test_output_file_task(file_task_model_id, client):
     """Simple check that we can get a file output"""
     # cGRmZGF0Yf//AA== is b"pdfdata\xff\xff\x00" base64 encoded
@@ -789,27 +803,6 @@ def test_runtime_info_ok_response_all_packages(runtime_http_server):
             assert "py_to_proto" in json_response["python_packages"]
 
 
-def test_runtime_info_ok_custom_python_packages(runtime_http_server):
-    """Make sure the runtime info returns version data"""
-    with temp_config(
-        {"runtime": {"version_info": {"python_packages": {"custom_package": "0.1.0"}}}},
-        merge_strategy="merge",
-    ):
-        with TestClient(runtime_http_server.app) as client:
-            response = client.get(http_server.RUNTIME_INFO_ENDPOINT)
-            assert response.status_code == 200
-
-            json_response = json.loads(
-                response.content.decode(response.default_encoding)
-            )
-            # runtime_version not added if not set
-            assert json_response["runtime_version"] == ""
-            # custom library is set while other random packages are not
-            assert "caikit" in json_response["python_packages"]
-            assert json_response["python_packages"]["custom_package"] == "0.1.0"
-            assert "py_to_proto" not in json_response["python_packages"]
-
-
 def test_http_server_shutdown_with_model_poll(open_port):
     """Test that a SIGINT successfully shuts down the running server"""
     with tempfile.TemporaryDirectory() as workdir:
diff --git a/tests/runtime/http_server/test_request_aborter.py b/tests/runtime/http_server/test_request_aborter.py
new file mode 100644
index 00000000..29eb458a
--- /dev/null
+++ b/tests/runtime/http_server/test_request_aborter.py
@@ -0,0 +1,108 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Standard
+import asyncio
+import datetime
+import threading
+import time
+
+# Third Party
+from fastapi import FastAPI, Request
+from requests.exceptions import ReadTimeout
+import pytest
+import requests
+import uvicorn
+
+# Local
+from caikit.runtime.http_server.request_aborter import HttpRequestAborter
+
+
+def get_time_remaining(start_time: datetime.datetime, timeout: int = 10) -> int:
+    now = datetime.datetime.now()
+    return ((start_time + datetime.timedelta(seconds=timeout)) - now).seconds
+
+
+def test_request_aborter(open_port):
+    # Get start time to ensure test doesn't hang
+    start_time = datetime.datetime.now()
+
+    # Create FastAPI app for testing
+    app = FastAPI()
+
+    # Initialize synchronization variables for tracking request process
+    abort_event = threading.Event()
+    request_finished = threading.Event()
+
+    # Define an endpoint that sleeps until the client disconnects.
+    @app.get("/test_abort")
+    async def test_aborter(context: Request):
+        # Create aborter and add parent event
+        TEST_ABORTER = HttpRequestAborter(context, poll_time=0.001)
+        TEST_ABORTER.add_event(abort_event)
+
+        # Wait for client to disconnect
+        while not abort_event.is_set():
+            await asyncio.sleep(0.001)
+
+        # Assign TEST_ABORTER to the parent function. This allows the test to have
+        # access to this object without using globals
+        test_request_aborter.TEST_ABORTER = TEST_ABORTER
+        request_finished.set()
+
+    # Start up a local uvicorn server in a thread
+    config = uvicorn.Config(
+        app,
+        host="0.0.0.0",
+        port=open_port,
+        log_level="trace",
+        log_config=None,
+        timeout_graceful_shutdown=None,
+    )
+    server = uvicorn.Server(config=config)
+    server_thread = threading.Thread(target=server.run)
+    server_thread.start()
+
+    server_exception = None
+    try:
+        # Wait for uvicorn to start
+        while not server.started:
+            if get_time_remaining(start_time) < 0:
+                raise TimeoutError("Server did not start in time")
+
+            time.sleep(0.01)
+
+        # Try the endpoint but timeout after 1 second
+        with pytest.raises(ReadTimeout):
+            requests.get(
+                f"http://localhost:{open_port}/test_abort",
+                timeout=0.1,
+            )
+
+        # Wait for the request to finish/abort
+        request_finished.wait(get_time_remaining(start_time))
+
+        # Assert the request aborter actually aborted
+        assert test_request_aborter.TEST_ABORTER.must_abort()
+        assert abort_event.is_set()
+        assert request_finished.is_set()
+
+    except Exception as exc:
+        server_exception = exc
+    finally:
+        # Clean up the server
+        server.should_exit = True
+        server_thread.join()
+
+    if server_exception:
+        raise server_exception
diff --git a/tests/runtime/model_management/test_model_manager.py b/tests/runtime/model_management/test_model_manager.py
index 265b76a8..43096725 100644
--- a/tests/runtime/model_management/test_model_manager.py
+++ b/tests/runtime/model_management/test_model_manager.py
@@ -110,7 +110,7 @@ def test_load_model_ok_response():
         model_id=model_id,
         local_model_path=Fixtures.get_good_model_path(),
         model_type=Fixtures.get_good_model_type(),
-    ).size()
+    )
     assert model_size > 0
 
 
@@ -127,7 +127,7 @@ def test_load_model_no_size_update():
         model_id=model_id,
         local_model_path=Fixtures.get_good_model_path(),
         model_type=Fixtures.get_good_model_type(),
-    ).size()
+    )
     assert model_size > 0
     loaded_model = MODEL_MANAGER.loaded_models[model_id]
     assert loaded_model.size() == model_size
@@ -485,85 +485,6 @@ def test_model_manager_disk_caching_periodic_sync(good_model_path):
             assert mgr_one_unloaded and mgr_two_unloaded
 
 
-def test_nested_local_model_load_unload(good_model_path):
-    """Test that a model can be loaded in a subdirectory of the local_models_dir
-    and that the periodic sync does not unload the model.
-    """
-    with TemporaryDirectory() as cache_dir:
-        with non_singleton_model_managers(
-            1,
-            {
-                "runtime": {
-                    "local_models_dir": cache_dir,
-                    "lazy_load_local_models": True,
-                    "lazy_load_poll_period_seconds": 0,
-                },
-            },
-            "merge",
-        ) as managers:
-            manager = managers[0]
-
-            # Copy the model into a nested model directory
-            model_name = os.path.join("parent", os.path.basename(good_model_path))
-            model_cache_path = os.path.join(cache_dir, model_name)
-            assert not os.path.exists(model_cache_path)
-            shutil.copytree(good_model_path, model_cache_path)
-
-            # Trigger the periodic sync and make sure the model is NOT loaded
-            assert model_name not in manager.loaded_models
-            manager.sync_local_models(wait=True)
-            assert model_name not in manager.loaded_models
-
-            # Explicitly ask to load the nested model name to trigger the lazy
-            # load
-            model = manager.retrieve_model(model_name)
-            assert model
-            assert model_name in manager.loaded_models
-
-            # Re-trigger the sync and make sure the model does not get unloaded
-            manager.sync_local_models(wait=True)
-            assert model_name in manager.loaded_models
-
-
-def test_model_unload_race(good_model_path):
-    """Test that if a model gets unloaded _while_ it's actively being loaded
-    (before retrieve_model completes, but after load_model completes), no
-    exception is raised.
-    """
-    with TemporaryDirectory() as cache_dir:
-        with non_singleton_model_managers(
-            1,
-            {
-                "runtime": {
-                    "local_models_dir": cache_dir,
-                    "lazy_load_local_models": True,
-                    "lazy_load_poll_period_seconds": 0,
-                },
-            },
-            "merge",
-        ) as managers:
-            manager = managers[0]
-
-            # Copy the model to the local_models_dir
-            model_id = random_test_id()
-            model_cache_path = os.path.join(cache_dir, model_id)
-            shutil.copytree(good_model_path, model_cache_path)
-
-            # Patch the manager's load_model to immediately unload the model
-            orig_load_model = manager.load_model
-
-            def load_and_unload_model(self, model_id: str, *args, **kwargs):
-                res = orig_load_model(model_id, *args, **kwargs)
-                manager.unload_model(model_id)
-                return res
-
-            with patch.object(manager.__class__, "load_model", load_and_unload_model):
-
-                # Retrieve the model and make sure there's no error
-                assert manager.retrieve_model(model_id)
-                assert model_id not in manager.loaded_models
-
-
 def test_load_local_model_deleted_dir():
     """Make sure losing the local_models_dir out from under a running manager
     doesn't kill the whole thing
@@ -669,7 +590,7 @@ def test_load_model():
 
             model_size = MODEL_MANAGER.load_model(
                 model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE
-            ).size()
+            )
             assert expected_model_size == model_size
             mock_loader.load_model.assert_called_once()
             call_args = mock_loader.load_model.call_args
@@ -833,12 +754,12 @@ def test_reload_partially_loaded():
             mock_loader.load_model.return_value = loaded_model
             model_size = MODEL_MANAGER.load_model(
                 model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE, wait=False
-            ).size()
+            )
             assert model_size == special_model_size
             assert (
                 MODEL_MANAGER.load_model(
                     model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE, wait=False
-                ).size()
+                )
                 == special_model_size
             )
 
diff --git a/tests/runtime/servicers/test_model_runtime_servicer_impl.py b/tests/runtime/servicers/test_model_runtime_servicer_impl.py
index 2e63a1b8..d7f7b8d5 100644
--- a/tests/runtime/servicers/test_model_runtime_servicer_impl.py
+++ b/tests/runtime/servicers/test_model_runtime_servicer_impl.py
@@ -57,7 +57,7 @@ class TestModelRuntimeServicerImpl(unittest.TestCase):
             ]
         )
         mock_manager = MagicMock()
-        mock_manager.load_model.size.return_value = 1
+        mock_manager.load_model.return_value = 1
 
         with patch.object(self.servicer, "model_manager", mock_manager):
             response = self.servicer.loadModel(request, context)
@@ -75,7 +75,7 @@ class TestModelRuntimeServicerImpl(unittest.TestCase):
             get_config().inference_plugin.model_mesh.max_model_concurrency
         )
         mock_manager = MagicMock()
-        mock_manager.load_model.size.return_value = 1
+        mock_manager.load_model.return_value = 1
 
         with patch.object(self.servicer, "model_manager", mock_manager):
             response = self.servicer.loadModel(request, context)