diff --git a/runtimes/mlflow/mlserver_mlflow/runtime.py b/runtimes/mlflow/mlserver_mlflow/runtime.py index ffb322fe0..373e4e641 100644 --- a/runtimes/mlflow/mlserver_mlflow/runtime.py +++ b/runtimes/mlflow/mlserver_mlflow/runtime.py @@ -1,5 +1,6 @@ import mlflow +import asyncio from io import StringIO from fastapi import Depends, Header, Request, Response @@ -155,6 +156,11 @@ async def invocations( async def load(self) -> bool: # TODO: Log info message model_uri = await get_model_uri(self._settings) + + 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) self._input_schema = self._model.metadata.get_input_schema() @@ -163,6 +169,22 @@ 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" + ) + + cmd = ["pip", "install", "-r", model_dependencies] + + # 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: 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"