Repository: caikit/caikit
Pull request #554 (2023-10-31 19:33:03)

🐛 return 404 instead of 500 for incorrect model id with multi model finder
<!--  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**:

In case of `local_model_finder`, if we send in a non-existent model, it throws a `FileNotFoundError` if the model_path is not found - https://github.com/caikit/caikit/blob/main/caikit/core/model_management/local_model_finder.py#L71

`multi_model_finder` on the other hand, returns a `None` if the model is not found - 
https://github.com/caikit/caikit/blob/main/caikit/core/model_management/multi_model_finder.py#L137

This was causing a scenario to occur in which `multi_model_finder` was throwing a 500 if we sent it a non-existent model:

> {"channel": "MODEL-LOADER", "exception": "Traceback (most recent call last):\n  File \"/home/runner/work/caikit/caikit/caikit/runtime/model_management/model_loader.py\", line 118, in _load_module\n    model = MODEL_MANAGER.load(model_path)\n  File \"/home/runner/work/caikit/caikit/caikit/core/model_manager.py\", line 273, in load\n    return self._do_load(\n  File \"/home/runner/work/caikit/caikit/caikit/core/model_manager.py\", line 459, in _do_load\n    error.value_check(\n  File \"/home/runner/work/caikit/caikit/caikit/core/exceptions/error_handler.py\", line 384, in value_check\n    self(\n  File \"/home/runner/work/caikit/caikit/caikit/core/exceptions/error_handler.py\", line 130, in log_raise\n    raise exception\nValueError: value check failed: Unable to find a ModuleConfig for not_an_id", "level": "error", "log_code": "<RUN62912924E>", "message": "load failed when processing path: not_an_id with error: ValueError('value check failed: Unable to find a ModuleConfig for not_an_id')", "model_id": "not_an_id", "num_indent": 0, "thread_id": 140647025604160, "timestamp": "2023-11-01T19:17:57.688955"}


This PR adds a `ValueError` exception block while loading models, that can catch the `ValueError` thrown by `multi_model_finder` in case it doesn't find a model.

**Special notes for your reviewer**:

**If applicable**:
- [ ] this PR contains documentation
- [x] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/runtime/model_management/model_loader.py b/caikit/runtime/model_management/model_loader.py
index 2a3c3996..0a7d4006 100644
--- a/caikit/runtime/model_management/model_loader.py
+++ b/caikit/runtime/model_management/model_loader.py
@@ -135,6 +135,18 @@ class ModelLoader:
                 StatusCode.NOT_FOUND,
                 f"Model {model_id} not found. Nested error: {fnfe}",
             ) from fnfe
+        except ValueError as ve:
+            log_dict = {
+                "log_code": "<RUN38617724E>",
+                "message": "load failed to find model: %s with error: %s"
+                % (model_path, repr(ve)),
+                "model_id": model_id,
+            }
+            log.error(log_dict)
+            raise CaikitRuntimeException(
+                StatusCode.NOT_FOUND,
+                f"Model {model_id} not found. Nested error: {ve}",
+            ) from ve
         except Exception as ex:
             log_dict = {
                 "log_code": "<RUN62912924E>",


Test patch:
diff --git a/tests/core/model_management/test_multi_model_finder.py b/tests/core/model_management/test_multi_model_finder.py
index 0d3ddf31..2d9cfeff 100644
--- a/tests/core/model_management/test_multi_model_finder.py
+++ b/tests/core/model_management/test_multi_model_finder.py
@@ -28,6 +28,7 @@ from caikit.core import ModelManager
 from caikit.core.model_management.multi_model_finder import MultiModelFinder
 from tests.conftest import temp_config
 from tests.core.helpers import TestFinder
+import caikit
 
 ## Helpers #####################################################################
 
@@ -113,11 +114,17 @@ def test_multi_model_finder_first_not_found(test_finder_config, good_model_path)
         assert finder.find_model(good_model_path)
 
 
-def test_multi_model_finder_not_found():
+def test_multi_model_finder_not_found(reset_globals):
     """Make sure that a simple proxy to local works"""
     with temp_config_finder() as finder:
         assert isinstance(finder, MultiModelFinder)
         assert finder.find_model("not/a/valid/path") is None
+        with pytest.raises(ValueError) as e:
+            caikit.core.load("bad/path/to/model")
+        assert (
+            e.value.args[0]
+            == "value check failed: Unable to find a ModuleConfig for bad/path/to/model"
+        )
 
 
 @pytest.mark.parametrize(
diff --git a/tests/runtime/http_server/test_http_server.py b/tests/runtime/http_server/test_http_server.py
index c8ee703d..f6fe625c 100644
--- a/tests/runtime/http_server/test_http_server.py
+++ b/tests/runtime/http_server/test_http_server.py
@@ -16,11 +16,9 @@ Tests for the caikit HTTP server
 """
 # Standard
 from contextlib import contextmanager
-from dataclasses import dataclass
 from io import BytesIO
 from pathlib import Path
 from typing import Dict
-from unittest.mock import patch
 import json
 import os
 import signal
@@ -29,13 +27,13 @@ import zipfile
 
 # Third Party
 from fastapi.testclient import TestClient
-import numpy as np
 import pytest
 import requests
 import tls_test_tools
 
 # Local
 from caikit.core import MODEL_MANAGER, DataObjectBase, dataobject
+from caikit.core.model_management.multi_model_finder import MultiModelFinder
 from caikit.runtime import http_server
 from caikit.runtime.http_server.http_server import StreamEventTypes
 from tests.conftest import temp_config
@@ -44,6 +42,9 @@ from tests.runtime.conftest import (
     register_trained_model,
     runtime_http_test_server,
 )
+from tests.runtime.model_management.test_model_manager import (
+    non_singleton_model_managers,
+)
 
 ## Fixtures #####################################################################
 
@@ -576,7 +577,7 @@ def test_inference_streaming_sample_module_actual_server_throws(
 
 
 def test_no_model_id(client):
-    """Simple check that we can ping a model"""
+    """Simple check to make sure we return a 400 if no model_id in payload"""
     response = client.post(
         f"/api/v1/task/sample",
         json={"inputs": {"name": "world"}},
@@ -604,7 +605,8 @@ def test_inference_multi_task_module(multi_task_model_id, client):
 
 
 def test_model_not_found(client):
-    """Simple check that we can ping a model"""
+    """Simple error check to make sure we return a 404 in case of
+    incorrect model_id"""
     response = client.post(
         f"/api/v1/task/sample",
         json={"model_id": "not_an_id", "inputs": {"name": "world"}},
@@ -612,6 +614,45 @@ def test_model_not_found(client):
     assert response.status_code == 404
 
 
+def test_model_not_found_with_lazy_load_multi_model_finder(open_port):
+    """An error check to make sure we return a 404 in case of
+    incorrect model_id while using multi model finder with lazy load enabled"""
+    with tempfile.TemporaryDirectory() as workdir:
+        # NOTE: This test requires that the ModelManager class not be a singleton.
+        #   To accomplish this, the singleton instance is temporarily removed.
+        with non_singleton_model_managers(
+            1,
+            {
+                "runtime": {
+                    "local_models_dir": workdir,
+                    "lazy_load_local_models": True,
+                },
+                "model_management": {
+                    "finders": {
+                        "default": {
+                            "type": MultiModelFinder.name,
+                            "config": {
+                                "finder_priority": ["local"],
+                            },
+                        },
+                        "local": {"type": "LOCAL"},
+                    }
+                },
+            },
+            "merge",
+        ):
+            with runtime_http_test_server(open_port) as server:
+                # double checking that our local model_management change took affect
+                assert (
+                    server.global_predict_servicer._model_manager._lazy_load_local_models
+                )
+                response = requests.post(
+                    f"http://localhost:{server.port}/api/v1/task/sample",
+                    json={"model_id": "not_an_id", "inputs": {"name": "world"}},
+                )
+                assert response.status_code == 404
+
+
 def test_inference_sample_task_incorrect_input(sample_task_model_id, client):
     """Test that with an incorrect input, the test doesn't throw but
     instead returns None"""
diff --git a/tests/runtime/model_management/test_model_loader.py b/tests/runtime/model_management/test_model_loader.py
index d17da8af..592d6b2f 100644
--- a/tests/runtime/model_management/test_model_loader.py
+++ b/tests/runtime/model_management/test_model_loader.py
@@ -114,7 +114,7 @@ def test_load_invalid_model_error_response(model_loader):
             local_model_path=Fixtures.get_bad_model_archive_path(),
             model_type="not_real",
         ).wait()
-    assert context.value.status_code == grpc.StatusCode.INTERNAL
+    assert context.value.status_code == grpc.StatusCode.NOT_FOUND
     assert model_id in context.value.message
 
 
@@ -308,7 +308,7 @@ def test_load_model_without_waiting_deferred_error(model_loader):
     )
     with pytest.raises(CaikitRuntimeException) as context:
         loaded_model.model()
-    assert context.value.status_code == grpc.StatusCode.INTERNAL
+    assert context.value.status_code == grpc.StatusCode.NOT_FOUND
     assert model_id in context.value.message
 
 
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index 5df682c8..6dfead16 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -835,7 +835,7 @@ def test_load_model_badmodel_error_response(runtime_grpc_server):
             modelKey="baz",
         )
         stub.loadModel(load_model_request)
-    assert context.value.code() == grpc.StatusCode.INTERNAL
+    assert context.value.code() == grpc.StatusCode.NOT_FOUND
 
 
 def test_unload_model_ok_response(sample_task_model_id, runtime_grpc_server):