Allow MLflow URI with scheme (#839)

Author: regen100

mlserver/utils.py

@@ -1,6 +1,7 @@
 import os
 import uuid
 import asyncio
+import urllib.parse
 
 from asyncio import Task
 from typing import Callable, Dict, Optional, List, Type
@@ -21,22 +22,26 @@ async def get_model_uri(
     if not model_uri:
         raise InvalidModelURI(settings.name)
 
-    full_model_uri = _to_absolute_path(settings._source, model_uri)
-    if os.path.isfile(full_model_uri):
-        return full_model_uri
+    model_uri_components = urllib.parse.urlparse(model_uri, scheme="file")
+    if model_uri_components.scheme != "file":
+        return model_uri
+
+    full_model_path = _to_absolute_path(settings._source, model_uri_components.path)
+    if os.path.isfile(full_model_path):
+        return full_model_path
 
-    if os.path.isdir(full_model_uri):
-        # If full_model_uri is a folder, search for a well-known model filename
+    if os.path.isdir(full_model_path):
+        # If full_model_path is a folder, search for a well-known model filename
         for fname in wellknown_filenames:
-            model_path = os.path.join(full_model_uri, fname)
+            model_path = os.path.join(full_model_path, fname)
             if os.path.isfile(model_path):
                 return model_path
 
         # If none, return the folder
-        return full_model_uri
+        return full_model_path
 
     # Otherwise, the uri is neither a file nor a folder
-    raise InvalidModelURI(settings.name, full_model_uri)
+    raise InvalidModelURI(settings.name, full_model_path)
 
 
 def _to_absolute_path(source: Optional[str], model_uri: str) -> str:

runtimes/mlflow/tests/conftest.py

@@ -87,8 +87,10 @@ def pytorch_model_uri() -> str:
     return model_path
 
 
-@pytest.fixture
-def model_settings(model_uri: str) -> ModelSettings:
+@pytest.fixture(params=["", "file:"])
+def model_settings(model_uri: str, request: pytest.FixtureRequest) -> ModelSettings:
+    scheme = request.param
+    model_uri = scheme + model_uri
     return ModelSettings(
         name="mlflow-model",
         implementation=MLflowRuntime,

tests/test_utils.py

@@ -16,10 +16,12 @@
 from mlserver.types import InferenceRequest, InferenceResponse, Parameters
 from mlserver.settings import ModelSettings, ModelParameters
 
-
-@pytest.mark.parametrize(
-    "uri, source, expected",
-    [
+test_get_model_uri_paramaters = [
+    ("s3://bucket/key", None, "s3://bucket/key"),
+    ("s3://bucket/key", "/mnt/models/model-settings.json", "s3://bucket/key"),
+]
+for scheme in ["", "file:"]:
+    for uri, source, expected in [
         ("my-model.bin", None, "my-model.bin"),
         (
             "my-model.bin",
@@ -36,7 +38,13 @@
             "/mnt/models/model-settings.json",
             "/an/absolute/path/my-model.bin",
         ),
-    ],
+    ]:
+        test_get_model_uri_paramaters.append((scheme + uri, source, expected))
+
+
+@pytest.mark.parametrize(
+    "uri, source, expected",
+    test_get_model_uri_paramaters,
 )
 async def test_get_model_uri(uri: str, source: Optional[str], expected: str):
     model_settings = ModelSettings(