From ff7cd9bbb27e5058a5d275e3446cf65d12d3add0 Mon Sep 17 00:00:00 2001 From: Louis Moulin Date: Wed, 7 Jan 2026 17:16:47 +0100 Subject: [PATCH 1/5] feat: Add automatic dependency installation for MLflow models Signed-off-by: Louis Moulin --- runtimes/mlflow/mlserver_mlflow/runtime.py | 50 ++++++++++++++++------ 1 file changed, 36 insertions(+), 14 deletions(-) diff --git a/runtimes/mlflow/mlserver_mlflow/runtime.py b/runtimes/mlflow/mlserver_mlflow/runtime.py index ffb322fe0..9c0bbf1a8 100644 --- a/runtimes/mlflow/mlserver_mlflow/runtime.py +++ b/runtimes/mlflow/mlserver_mlflow/runtime.py @@ -1,34 +1,33 @@ -import mlflow - +import asyncio from io import StringIO -from fastapi import Depends, Header, Request, Response -from mlflow.version import VERSION +import mlflow +from fastapi import Depends, Header, Request, Response from mlflow.exceptions import MlflowException from mlflow.pyfunc.scoring_server import ( - CONTENT_TYPES, CONTENT_TYPE_CSV, CONTENT_TYPE_JSON, - parse_csv_input, + CONTENT_TYPES, _split_data_and_params, infer_and_parse_data, + parse_csv_input, predictions_to_json, ) - -from mlserver.types import InferenceRequest, InferenceResponse -from mlserver.model import MLModel -from mlserver.utils import get_model_uri -from mlserver.handlers import custom_handler +from mlflow.version import VERSION from mlserver.errors import InferenceError -from mlserver.settings import ModelParameters +from mlserver.handlers import custom_handler from mlserver.logging import logger +from mlserver.model import MLModel +from mlserver.settings import ModelParameters +from mlserver.types import InferenceRequest, InferenceResponse +from mlserver.utils import get_model_uri from .codecs import TensorDictCodec from .metadata import ( - to_metadata_tensors, - to_model_content_type, DefaultInputPrefix, DefaultOutputPrefix, + to_metadata_tensors, + to_model_content_type, ) @@ -155,6 +154,10 @@ async def invocations( async def load(self) -> bool: # TODO: Log info message model_uri = await get_model_uri(self._settings) + + # INFO: Install dependencies before loading the model + await self._install_dependencies(model_uri) + self._model = mlflow.pyfunc.load_model(model_uri) self._input_schema = self._model.metadata.get_input_schema() @@ -163,6 +166,25 @@ async def load(self) -> bool: return True + async def _install_dependencies(self, model_uri: str): + model_dependencies = mlflow.pyfunc.mlflow.artifacts.download_artifacts( + f"{model_uri}/requirements.txt" + ) + + logger.info(f"Dependencies file: {model_dependencies}") + await self._run_uv_add(model_dependencies) + cmd = ["uv", "pip", "install", "-r", model_dependencies] + logger.debug(f"Running command: {' '.join(cmd)}") + + # Execute command asynchronously + process = await asyncio.create_subprocess_exec( + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE + ) + stdout, stderr = await process.communicate() + if process.returncode != 0: + error_msg = stderr.decode() if stderr else "Unknown error" + raise RuntimeError(f"Failed to install dependencies: {error_msg}") + def _sync_metadata(self) -> None: # Update metadata from model signature (if present) if self._signature is None: From 1cfaba6e06fd48dd30ec32b1cb23339e2eadbf19 Mon Sep 17 00:00:00 2001 From: Louis Moulin Date: Wed, 7 Jan 2026 17:20:29 +0100 Subject: [PATCH 2/5] revert some format stuff Signed-off-by: Louis Moulin --- runtimes/mlflow/mlserver_mlflow/runtime.py | 26 ++++++++++++---------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/runtimes/mlflow/mlserver_mlflow/runtime.py b/runtimes/mlflow/mlserver_mlflow/runtime.py index 9c0bbf1a8..c9ade0af1 100644 --- a/runtimes/mlflow/mlserver_mlflow/runtime.py +++ b/runtimes/mlflow/mlserver_mlflow/runtime.py @@ -1,33 +1,35 @@ +import mlflow + import asyncio from io import StringIO - -import mlflow from fastapi import Depends, Header, Request, Response + +from mlflow.version import VERSION from mlflow.exceptions import MlflowException from mlflow.pyfunc.scoring_server import ( + CONTENT_TYPES, CONTENT_TYPE_CSV, CONTENT_TYPE_JSON, - CONTENT_TYPES, + parse_csv_input, _split_data_and_params, infer_and_parse_data, - parse_csv_input, predictions_to_json, ) -from mlflow.version import VERSION -from mlserver.errors import InferenceError -from mlserver.handlers import custom_handler -from mlserver.logging import logger -from mlserver.model import MLModel -from mlserver.settings import ModelParameters + from mlserver.types import InferenceRequest, InferenceResponse +from mlserver.model import MLModel from mlserver.utils import get_model_uri +from mlserver.handlers import custom_handler +from mlserver.errors import InferenceError +from mlserver.settings import ModelParameters +from mlserver.logging import logger from .codecs import TensorDictCodec from .metadata import ( - DefaultInputPrefix, - DefaultOutputPrefix, to_metadata_tensors, to_model_content_type, + DefaultInputPrefix, + DefaultOutputPrefix, ) From 713866139e3dd52dbebccbeb74884a65b4b1e943 Mon Sep 17 00:00:00 2001 From: Louis Moulin Date: Wed, 7 Jan 2026 17:26:05 +0100 Subject: [PATCH 3/5] switch to pip Signed-off-by: Louis Moulin --- runtimes/mlflow/mlserver_mlflow/runtime.py | 3 +-- runtimes/mlflow/uv.lock | 3 +++ 2 files changed, 4 insertions(+), 2 deletions(-) create mode 100644 runtimes/mlflow/uv.lock diff --git a/runtimes/mlflow/mlserver_mlflow/runtime.py b/runtimes/mlflow/mlserver_mlflow/runtime.py index c9ade0af1..5eca76940 100644 --- a/runtimes/mlflow/mlserver_mlflow/runtime.py +++ b/runtimes/mlflow/mlserver_mlflow/runtime.py @@ -174,8 +174,7 @@ async def _install_dependencies(self, model_uri: str): ) logger.info(f"Dependencies file: {model_dependencies}") - await self._run_uv_add(model_dependencies) - cmd = ["uv", "pip", "install", "-r", model_dependencies] + cmd = ["pip", "install", "-r", model_dependencies] logger.debug(f"Running command: {' '.join(cmd)}") # Execute command asynchronously diff --git a/runtimes/mlflow/uv.lock b/runtimes/mlflow/uv.lock new file mode 100644 index 000000000..7518fc90b --- /dev/null +++ b/runtimes/mlflow/uv.lock @@ -0,0 +1,3 @@ +version = 1 +revision = 3 +requires-python = ">=3.12" From a3b1b7492e878fa016d3da2808c17f4755a8d334 Mon Sep 17 00:00:00 2001 From: Louis Moulin Date: Wed, 7 Jan 2026 17:34:00 +0100 Subject: [PATCH 4/5] add feature gate and defaut to False for auto install Signed-off-by: Louis Moulin --- runtimes/mlflow/mlserver_mlflow/runtime.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/runtimes/mlflow/mlserver_mlflow/runtime.py b/runtimes/mlflow/mlserver_mlflow/runtime.py index 5eca76940..439af7455 100644 --- a/runtimes/mlflow/mlserver_mlflow/runtime.py +++ b/runtimes/mlflow/mlserver_mlflow/runtime.py @@ -158,7 +158,8 @@ async def load(self) -> bool: model_uri = await get_model_uri(self._settings) # INFO: Install dependencies before loading the model - await self._install_dependencies(model_uri) + if getattr(self._settings.parameters.extra, "auto_install_dependencies", False): + await self._install_dependencies(model_uri) self._model = mlflow.pyfunc.load_model(model_uri) @@ -173,9 +174,7 @@ async def _install_dependencies(self, model_uri: str): f"{model_uri}/requirements.txt" ) - logger.info(f"Dependencies file: {model_dependencies}") cmd = ["pip", "install", "-r", model_dependencies] - logger.debug(f"Running command: {' '.join(cmd)}") # Execute command asynchronously process = await asyncio.create_subprocess_exec( From b2210b589f461fe0732a3ad16dd805620c8a96e3 Mon Sep 17 00:00:00 2001 From: Louis Moulin Date: Wed, 7 Jan 2026 17:38:34 +0100 Subject: [PATCH 5/5] little format Signed-off-by: Louis Moulin --- runtimes/mlflow/mlserver_mlflow/runtime.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runtimes/mlflow/mlserver_mlflow/runtime.py b/runtimes/mlflow/mlserver_mlflow/runtime.py index 439af7455..373e4e641 100644 --- a/runtimes/mlflow/mlserver_mlflow/runtime.py +++ b/runtimes/mlflow/mlserver_mlflow/runtime.py @@ -157,8 +157,8 @@ async def load(self) -> bool: # TODO: Log info message model_uri = await get_model_uri(self._settings) - # INFO: Install dependencies before loading the model if getattr(self._settings.parameters.extra, "auto_install_dependencies", False): + # INFO: Install dependencies before loading the model await self._install_dependencies(model_uri) self._model = mlflow.pyfunc.load_model(model_uri)