Repository: mlflow/mlflow

Issue:
[BUG] input_example not working properly for pyfunc.PythonModel
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Where did you encounter this bug?

Databricks

### Willingness to contribute

No. I cannot contribute a bug fix at this time.

### MLflow version

mlflow==2.8.0

### System information

- Databricks runtime 14.2

### Describe the problem

The use of `input_example` to infer schemas for custom python models, i.e. `pyfunc.PythonModel`, is not working. For example, the example in the [documentation](https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html) does not work. The error message below indicates that the position is not aligned correctly. Context is required to call `predict`. Instead, only `model_input` is given. 



### Tracking information

_No response_

### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
from typing import List, Dict
import mlflow

class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input: List[str]) -> List[str]:
        return [i.upper() for i in model_input]


mlflow.pyfunc.save_model("model", python_model=MyModel(), input_example=["a"])
```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
File <command-3509881627683303>, line 10
      6     def predict(self, context, model_input: List[str]) -> List[str]:
      7         return [i.upper() for i in model_input]
---> 10 mlflow.pyfunc.save_model("model", python_model=MyModel(), input_example=["a"])

File /databricks/python/lib/python3.10/site-packages/mlflow/pyfunc/__init__.py:1902, in save_model(path, loader_module, data_path, code_path, conda_env, mlflow_model, python_model, artifacts, signature, input_example, pip_requirements, extra_pip_requirements, metadata, model_config, **kwargs)
   1900 elif isinstance(python_model, PythonModel):
   1901     input_arg_index = 1  # second argument
-> 1902     if signature := _infer_signature_from_type_hints(
   1903         python_model.predict,
   1904         input_arg_index=input_arg_index,
   1905         input_example=input_example,
   1906     ):
   1907         mlflow_model.signature = signature
   1908     elif input_example is not None:

File /databricks/python/lib/python3.10/site-packages/mlflow/models/signature.py:319, in _infer_signature_from_type_hints(func, input_arg_index, input_example)
    317     output_example = func(input_example, params=params) if input_example else None
    318 else:
--> 319     output_example = func(input_example) if input_example else None
    320 output_schema = (
    321     _infer_schema_from_type_hint(hints.output, output_example) if hints.output else None
    322 )
    323 if not any([input_schema, output_schema, params_schema]):

TypeError: MyModel.predict() missing 1 required positional argument: 'model_input'
```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
```


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/deployments`: MLflow Deployments client APIs, server, and third-party Deployments integrations
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations

Discussion:
harupy: let's cherry-pick this into 2.9.2

Patch:
diff --git a/mlflow/models/signature.py b/mlflow/models/signature.py
index b7c585cc3..a9ca4712a 100644
--- a/mlflow/models/signature.py
+++ b/mlflow/models/signature.py
@@ -313,10 +313,17 @@ def _infer_signature_from_type_hints(func, input_arg_index, input_example=None):
 
     input_schema = _infer_schema_from_type_hint(hints.input, input_example) if hints.input else None
     params_schema = _infer_param_schema(params) if params else None
-    if params and params_key in inspect.signature(func).parameters:
-        output_example = func(input_example, params=params) if input_example else None
+    input_arg_name = _get_arg_names(func)[input_arg_index]
+    if input_example:
+        inputs = {input_arg_name: input_example}
+        if params and params_key in inspect.signature(func).parameters:
+            inputs[params_key] = params
+        # This is for PythonModel's predict function
+        if input_arg_index == 1:
+            inputs["context"] = None
+        output_example = func(**inputs)
     else:
-        output_example = func(input_example) if input_example else None
+        output_example = None
     output_schema = (
         _infer_schema_from_type_hint(hints.output, output_example) if hints.output else None
     )
diff --git a/mlflow/utils/uri.py b/mlflow/utils/uri.py
index 331648b8e..ecaa228c7 100644
--- a/mlflow/utils/uri.py
+++ b/mlflow/utils/uri.py
@@ -430,7 +430,10 @@ def validate_path_is_safe(path):
     """
     from mlflow.utils.file_utils import local_file_uri_to_path
 
-    exc = MlflowException(f"Invalid path: {path}", error_code=INVALID_PARAMETER_VALUE)
+    # We must decode URL before validating it
+    path = urllib.parse.unquote(path)
+
+    exc = MlflowException("Invalid path", error_code=INVALID_PARAMETER_VALUE)
     if any((s in path) for s in ("#", "%23")):
         raise exc
 


Test patch:
diff --git a/tests/pyfunc/test_model_export_with_class_and_artifacts.py b/tests/pyfunc/test_model_export_with_class_and_artifacts.py
index d0ad9bf2e..4a54e9a81 100644
--- a/tests/pyfunc/test_model_export_with_class_and_artifacts.py
+++ b/tests/pyfunc/test_model_export_with_class_and_artifacts.py
@@ -1464,3 +1464,19 @@ def test_load_model_fails_for_feature_store_models(tmp_path):
         MlflowException, match="mlflow.pyfunc.load_model is not supported for Feature Store models"
     ):
         mlflow.pyfunc.load_model(f"runs:/{run.info.run_id}/model")
+
+
+def test_pyfunc_model_infer_signature_from_type_hints(model_path):
+    class TestModel(mlflow.pyfunc.PythonModel):
+        def predict(self, context, model_input: List[str], params=None) -> List[str]:
+            return model_input
+
+    with mlflow.start_run():
+        model_info = mlflow.pyfunc.log_model(
+            python_model=TestModel(),
+            artifact_path="test_model",
+            input_example=["a"],
+        )
+    pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)
+    assert pyfunc_model.metadata.get_input_schema().to_dict() == [{"type": "string"}]
+    pd.testing.assert_frame_equal(pyfunc_model.predict(["a", "b"]), pd.DataFrame(["a", "b"]))
diff --git a/tests/server/test_handlers.py b/tests/server/test_handlers.py
index 1bc5bd047..a6bbbace8 100644
--- a/tests/server/test_handlers.py
+++ b/tests/server/test_handlers.py
@@ -38,11 +38,12 @@ from mlflow.protos.model_registry_pb2 import (
     UpdateRegisteredModel,
 )
 from mlflow.protos.service_pb2 import CreateExperiment, SearchRuns
-from mlflow.server import BACKEND_STORE_URI_ENV_VAR, app
+from mlflow.server import BACKEND_STORE_URI_ENV_VAR, SERVE_ARTIFACTS_ENV_VAR, app
 from mlflow.server.handlers import (
     _create_experiment,
     _create_model_version,
     _create_registered_model,
+    _delete_artifact_mlflow_artifacts,
     _delete_model_version,
     _delete_model_version_tag,
     _delete_registered_model,
@@ -105,6 +106,11 @@ def mock_model_registry_store():
         yield mock_store
 
 
+@pytest.fixture
+def enable_serve_artifacts(monkeypatch):
+    monkeypatch.setenv(SERVE_ARTIFACTS_ENV_VAR, "true")
+
+
 def test_health():
     with app.test_client() as c:
         response = c.get("/health")
@@ -777,3 +783,22 @@ def test_get_model_version_by_alias(mock_get_request_message, mock_model_registr
     _, args = mock_model_registry_store.get_model_version_by_alias.call_args
     assert args == {"name": name, "alias": alias}
     assert json.loads(resp.get_data()) == {"model_version": jsonify(mvd)}
+
+
+@pytest.mark.parametrize(
+    "path",
+    [
+        "/path",
+        "path/../to/file",
+        "/etc/passwd",
+        "/etc/passwd%00.jpg",
+        "/file://etc/passwd",
+        "%2E%2E%2F%2E%2E%2Fpath",
+    ],
+)
+def test_delete_artifact_mlflow_artifacts_throws_for_malicious_path(enable_serve_artifacts, path):
+    response = _delete_artifact_mlflow_artifacts(path)
+    assert response.status_code == 400
+    json_response = json.loads(response.get_data())
+    assert json_response["error_code"] == ErrorCode.Name(INVALID_PARAMETER_VALUE)
+    assert json_response["message"] == "Invalid path"
diff --git a/tests/store/artifact/test_http_artifact_repo.py b/tests/store/artifact/test_http_artifact_repo.py
index 8249f083d..f6b9a9d72 100644
--- a/tests/store/artifact/test_http_artifact_repo.py
+++ b/tests/store/artifact/test_http_artifact_repo.py
@@ -246,7 +246,7 @@ def test_list_artifacts(http_artifact_repo):
             http_artifact_repo.list_artifacts()
 
 
-@pytest.mark.parametrize("path", ["/tmp/path", "../../path"])
+@pytest.mark.parametrize("path", ["/tmp/path", "../../path", "%2E%2E%2Fpath"])
 def test_list_artifacts_malicious_path(http_artifact_repo, path):
     with mock.patch(
         "mlflow.store.artifact.http_artifact_repo.http_request",
@@ -259,7 +259,7 @@ def test_list_artifacts_malicious_path(http_artifact_repo, path):
             200,
         ),
     ):
-        with pytest.raises(MlflowException, match=f"Invalid path: {path}"):
+        with pytest.raises(MlflowException, match="Invalid path"):
             http_artifact_repo.list_artifacts()
 
 
diff --git a/tests/tracking/test_rest_tracking.py b/tests/tracking/test_rest_tracking.py
index d56eecd77..5f952dda4 100644
--- a/tests/tracking/test_rest_tracking.py
+++ b/tests/tracking/test_rest_tracking.py
@@ -35,7 +35,6 @@ from mlflow.entities import (
 )
 from mlflow.exceptions import MlflowException
 from mlflow.models import Model
-from mlflow.server.handlers import validate_path_is_safe
 from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
 from mlflow.utils import mlflow_tags
 from mlflow.utils.file_utils import TempDir, path_to_local_file_uri
@@ -514,152 +513,6 @@ def test_set_tag_validation(mlflow_client):
     assert response.status_code == 200
 
 
-@pytest.mark.parametrize(
-    "path",
-    [
-        "path",
-        "path/",
-        "path/to/file",
-    ],
-)
-def test_validate_path_is_safe_good(path):
-    validate_path_is_safe(path)
-
-
-@pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows")
-@pytest.mark.parametrize(
-    "path",
-    [
-        # relative path from current directory of C: drive
-        ".../...//",
-    ],
-)
-def test_validate_path_is_safe_windows_good(path):
-    validate_path_is_safe(path)
-
-
-@pytest.mark.skipif(is_windows(), reason="This test does not pass on Windows")
-@pytest.mark.parametrize(
-    "path",
-    [
-        "/path",
-        "../path",
-        "../../path",
-        "./../path",
-        "path/../to/file",
-        "path/../../to/file",
-        "file://a#/..//tmp",
-        "file://a%23/..//tmp/",
-        "/etc/passwd",
-        "/etc/passwd%00.jpg",
-        "/etc/passwd%00.html",
-        "/etc/passwd%00.txt",
-        "/etc/passwd%00.php",
-        "/etc/passwd%00.asp",
-        "/file://etc/passwd",
-    ],
-)
-def test_validate_path_is_safe_bad(path):
-    with pytest.raises(MlflowException, match="Invalid path"):
-        validate_path_is_safe(path)
-
-
-@pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows")
-@pytest.mark.parametrize(
-    "path",
-    [
-        r"../path",
-        r"../../path",
-        r"./../path",
-        r"path/../to/file",
-        r"path/../../to/file",
-        r"..\path",
-        r"..\..\path",
-        r".\..\path",
-        r"path\..\to\file",
-        r"path\..\..\to\file",
-        # Drive-relative paths
-        r"C:path",
-        r"C:path/",
-        r"C:path/to/file",
-        r"C:../path/to/file",
-        r"C:\path",
-        r"C:/path",
-        r"C:\path\to\file",
-        r"C:\path/to/file",
-        r"C:\path\..\to\file",
-        r"C:/path/../to/file",
-        # UNC(Universal Naming Convention) paths
-        r"\\path\to\file",
-        r"\\path/to/file",
-        r"\\.\\C:\path\to\file",
-        r"\\?\C:\path\to\file",
-        r"\\?\UNC/path/to/file",
-        # Other potential attackable paths
-        r"/etc/password",
-        r"/path",
-        r"/etc/passwd%00.jpg",
-        r"/etc/passwd%00.html",
-        r"/etc/passwd%00.txt",
-        r"/etc/passwd%00.php",
-        r"/etc/passwd%00.asp",
-        r"/Windows/no/such/path",
-        r"/file://etc/passwd",
-        r"/file:c:/passwd",
-        r"/file://d:/windows/win.ini",
-        r"/file://./windows/win.ini",
-        r"file://c:/boot.ini",
-        r"file://C:path",
-        r"file://C:path/",
-        r"file://C:path/to/file",
-        r"file:///C:/Windows/System32/",
-        r"file:///etc/passwd",
-        r"file:///d:/windows/repair/sam",
-        r"file:///proc/version",
-        r"file:///inetpub/wwwroot/global.asa",
-        r"/file://../windows/win.ini",
-        r"../etc/passwd",
-        r"..\Windows\System32\\",
-        r"C:\Windows\System32\\",
-        r"/etc/passwd",
-        r"::Windows\System32",
-        r"..\..\..\..\Windows\System32\\",
-        r"../Windows/System32",
-        r"....\\",
-        r"\\?\C:\Windows\System32\\",
-        r"\\.\C:\Windows\System32\\",
-        r"\\UNC\Server\Share\\",
-        r"\\Server\Share\folder\\",
-        r"\\127.0.0.1\c$\Windows\\",
-        r"\\localhost\c$\Windows\\",
-        r"\\smbserver\share\path\\",
-        r"..\\?\C:\Windows\System32\\",
-        r"C:/Windows/../Windows/System32/",
-        r"C:\Windows\..\Windows\System32\\",
-        r"../../../../../../../../../../../../Windows/System32",
-        r"../../../../../../../../../../../../etc/passwd",
-        r"../../../../../../../../../../../../var/www/html/index.html",
-        r"../../../../../../../../../../../../usr/local/etc/openvpn/server.conf",
-        r"../../../../../../../../../../../../Program Files (x86)",
-        r"/../../../../../../../../../../../../Windows/System32",
-        r"/Windows\../etc/passwd",
-        r"/Windows\..\Windows\System32\\",
-        r"/Windows\..\Windows\System32\cmd.exe",
-        r"/Windows\..\Windows\System32\msconfig.exe",
-        r"/Windows\..\Windows\System32\regedit.exe",
-        r"/Windows\..\Windows\System32\taskmgr.exe",
-        r"/Windows\..\Windows\System32\control.exe",
-        r"/Windows\..\Windows\System32\services.msc",
-        r"/Windows\..\Windows\System32\diskmgmt.msc",
-        r"/Windows\..\Windows\System32\eventvwr.msc",
-        r"/Windows/System32/drivers/etc/hosts",
-    ],
-)
-def test_validate_path_is_safe_windows_bad(path):
-    with pytest.raises(MlflowException, match="Invalid path"):
-        validate_path_is_safe(path)
-
-
 def test_path_validation(mlflow_client):
     experiment_id = mlflow_client.create_experiment("tags validation")
     created_run = mlflow_client.create_run(experiment_id)
@@ -670,7 +523,7 @@ def test_path_validation(mlflow_client):
         assert resp.status_code == 400
         assert response.json() == {
             "error_code": "INVALID_PARAMETER_VALUE",
-            "message": f"Invalid path: {invalid_path}",
+            "message": "Invalid path",
         }
 
     response = requests.get(
diff --git a/tests/utils/test_uri.py b/tests/utils/test_uri.py
index b776795c6..26aecc514 100644
--- a/tests/utils/test_uri.py
+++ b/tests/utils/test_uri.py
@@ -24,6 +24,7 @@ from mlflow.utils.uri import (
     is_valid_dbfs_uri,
     remove_databricks_profile_info_from_artifact_uri,
     resolve_uri_if_local,
+    validate_path_is_safe,
 )
 
 
@@ -724,3 +725,152 @@ def test_correctly_detect_fuse_and_uc_uris(uri):
 )
 def test_negative_detection(uri):
     assert not is_fuse_or_uc_volumes_uri(uri)
+
+
+@pytest.mark.parametrize(
+    "path",
+    [
+        "path",
+        "path/",
+        "path/to/file",
+    ],
+)
+def test_validate_path_is_safe_good(path):
+    validate_path_is_safe(path)
+
+
+@pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows")
+@pytest.mark.parametrize(
+    "path",
+    [
+        # relative path from current directory of C: drive
+        ".../...//",
+    ],
+)
+def test_validate_path_is_safe_windows_good(path):
+    validate_path_is_safe(path)
+
+
+@pytest.mark.skipif(is_windows(), reason="This test does not pass on Windows")
+@pytest.mark.parametrize(
+    "path",
+    [
+        "/path",
+        "../path",
+        "../../path",
+        "./../path",
+        "path/../to/file",
+        "path/../../to/file",
+        "file://a#/..//tmp",
+        "file://a%23/..//tmp/",
+        "/etc/passwd",
+        "/etc/passwd%00.jpg",
+        "/etc/passwd%00.html",
+        "/etc/passwd%00.txt",
+        "/etc/passwd%00.php",
+        "/etc/passwd%00.asp",
+        "/file://etc/passwd",
+        # Encoded paths with '..'
+        "%2E%2E%2Fpath",
+        "%2E%2E%2F%2E%2E%2Fpath",
+    ],
+)
+def test_validate_path_is_safe_bad(path):
+    with pytest.raises(MlflowException, match="Invalid path"):
+        validate_path_is_safe(path)
+
+
+@pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows")
+@pytest.mark.parametrize(
+    "path",
+    [
+        r"../path",
+        r"../../path",
+        r"./../path",
+        r"path/../to/file",
+        r"path/../../to/file",
+        r"..\path",
+        r"..\..\path",
+        r".\..\path",
+        r"path\..\to\file",
+        r"path\..\..\to\file",
+        # Drive-relative paths
+        r"C:path",
+        r"C:path/",
+        r"C:path/to/file",
+        r"C:../path/to/file",
+        r"C:\path",
+        r"C:/path",
+        r"C:\path\to\file",
+        r"C:\path/to/file",
+        r"C:\path\..\to\file",
+        r"C:/path/../to/file",
+        # UNC(Universal Naming Convention) paths
+        r"\\path\to\file",
+        r"\\path/to/file",
+        r"\\.\\C:\path\to\file",
+        r"\\?\C:\path\to\file",
+        r"\\?\UNC/path/to/file",
+        # Other potential attackable paths
+        r"/etc/password",
+        r"/path",
+        r"/etc/passwd%00.jpg",
+        r"/etc/passwd%00.html",
+        r"/etc/passwd%00.txt",
+        r"/etc/passwd%00.php",
+        r"/etc/passwd%00.asp",
+        r"/Windows/no/such/path",
+        r"/file://etc/passwd",
+        r"/file:c:/passwd",
+        r"/file://d:/windows/win.ini",
+        r"/file://./windows/win.ini",
+        r"file://c:/boot.ini",
+        r"file://C:path",
+        r"file://C:path/",
+        r"file://C:path/to/file",
+        r"file:///C:/Windows/System32/",
+        r"file:///etc/passwd",
+        r"file:///d:/windows/repair/sam",
+        r"file:///proc/version",
+        r"file:///inetpub/wwwroot/global.asa",
+        r"/file://../windows/win.ini",
+        r"../etc/passwd",
+        r"..\Windows\System32\\",
+        r"C:\Windows\System32\\",
+        r"/etc/passwd",
+        r"::Windows\System32",
+        r"..\..\..\..\Windows\System32\\",
+        r"../Windows/System32",
+        r"....\\",
+        r"\\?\C:\Windows\System32\\",
+        r"\\.\C:\Windows\System32\\",
+        r"\\UNC\Server\Share\\",
+        r"\\Server\Share\folder\\",
+        r"\\127.0.0.1\c$\Windows\\",
+        r"\\localhost\c$\Windows\\",
+        r"\\smbserver\share\path\\",
+        r"..\\?\C:\Windows\System32\\",
+        r"C:/Windows/../Windows/System32/",
+        r"C:\Windows\..\Windows\System32\\",
+        r"../../../../../../../../../../../../Windows/System32",
+        r"../../../../../../../../../../../../etc/passwd",
+        r"../../../../../../../../../../../../var/www/html/index.html",
+        r"../../../../../../../../../../../../usr/local/etc/openvpn/server.conf",
+        r"../../../../../../../../../../../../Program Files (x86)",
+        r"/../../../../../../../../../../../../Windows/System32",
+        r"/Windows\../etc/passwd",
+        r"/Windows\..\Windows\System32\\",
+        r"/Windows\..\Windows\System32\cmd.exe",
+        r"/Windows\..\Windows\System32\msconfig.exe",
+        r"/Windows\..\Windows\System32\regedit.exe",
+        r"/Windows\..\Windows\System32\taskmgr.exe",
+        r"/Windows\..\Windows\System32\control.exe",
+        r"/Windows\..\Windows\System32\services.msc",
+        r"/Windows\..\Windows\System32\diskmgmt.msc",
+        r"/Windows\..\Windows\System32\eventvwr.msc",
+        r"/Windows/System32/drivers/etc/hosts",
+    ],
+)
+def test_validate_path_is_safe_windows_bad(path):
+    with pytest.raises(MlflowException, match="Invalid path"):
+        validate_path_is_safe(path)