1  import logging
  2  import os
  3  import shutil
  4  from io import StringIO
  5  from typing import ForwardRef, get_args, get_origin
  6  
  7  from mlflow.exceptions import MlflowException
  8  from mlflow.models.flavor_backend_registry import get_flavor_backend
  9  from mlflow.utils import env_manager as _EnvManager
 10  from mlflow.utils.databricks_utils import is_databricks_connect
 11  from mlflow.utils.file_utils import TempDir
 12  
 13  _logger = logging.getLogger(__name__)
 14  UV_INSTALLATION_INSTRUCTIONS = (
 15      "Run `pip install uv` to install uv. See "
 16      "https://docs.astral.sh/uv/getting-started/installation for other installation methods."
 17  )
 18  
 19  
 20  def build_docker(
 21      model_uri=None,
 22      name="mlflow-pyfunc",
 23      env_manager=_EnvManager.VIRTUALENV,
 24      mlflow_home=None,
 25      install_java=False,
 26      install_mlflow=False,
 27      enable_mlserver=False,
 28      base_image=None,
 29  ):
 30      """
 31      Builds a Docker image whose default entrypoint serves an MLflow model at port 8080, using the
 32      python_function flavor. The container serves the model referenced by ``model_uri``, if
 33      specified. If ``model_uri`` is not specified, an MLflow Model directory must be mounted as a
 34      volume into the /opt/ml/model directory in the container.
 35  
 36      .. important::
 37  
 38          Since MLflow 2.10.1, the Docker image built with ``--model-uri`` does **not install Java**
 39          for improved performance, unless the model flavor is one of ``["johnsnowlabs", "h2o"
 40          "spark"]``. If you need to install Java for other flavors, e.g. custom Python model
 41          that uses SparkML, please specify ``install-java=True`` to enforce Java installation.
 42          For earlier versions, Java is always installed to the image.
 43  
 44  
 45      .. warning::
 46  
 47          If ``model_uri`` is unspecified, the resulting image doesn't support serving models with
 48          the RFunc server.
 49  
 50      NB: by default, the container will start nginx and uvicorn processes. If you don't need the
 51      nginx process to be started (for instance if you deploy your container to Google Cloud Run),
 52      you can disable it via the DISABLE_NGINX environment variable:
 53  
 54      .. code:: bash
 55  
 56          docker run -p 5001:8080 -e DISABLE_NGINX=true "my-image-name"
 57  
 58      See https://www.mlflow.org/docs/latest/python_api/mlflow.pyfunc.html for more information on the
 59      'python_function' flavor.
 60  
 61      Args:
 62          model_uri: URI to the model. A local path, a 'runs:/' URI, or a remote storage URI (e.g.,
 63              an 's3://' URI). For more information about supported remote URIs for model artifacts,
 64              see https://mlflow.org/docs/latest/tracking.html#artifact-stores
 65          name: Name of the Docker image to build. Defaults to 'mlflow-pyfunc'.
 66          env_manager: If specified, create an environment for MLmodel using the specified environment
 67              manager. The following values are supported: (1) virtualenv (default): use virtualenv
 68              and pyenv for Python version management (2) conda: use conda (3) local: use the local
 69              environment without creating a new one.
 70          mlflow_home: Path to local clone of MLflow project. Use for development only.
 71          install_java: If specified, install Java in the image. Default is False in order to
 72              reduce both the image size and the build time. Model flavors requiring Java will enable
 73              this setting automatically, such as the Spark flavor. (This argument is only available
 74              in MLflow 2.10.1 and later. In earlier versions, Java is always installed to the image.)
 75          install_mlflow: If specified and there is a conda or virtualenv environment to be activated
 76              mlflow will be installed into the environment after it has been activated.
 77              The version of installed mlflow will be the same as the one used to invoke this command.
 78          enable_mlserver: If specified, the image will be built with the Seldon MLserver as backend.
 79          base_image: Base image for the Docker image. If not specified, the default image is either
 80              UBUNTU_BASE_IMAGE = "ubuntu:22.04" or PYTHON_SLIM_BASE_IMAGE = "python:{version}-slim"
 81              Note: If custom image is used, there are no guarantees that the image will work. You
 82              may find greater compatibility by building your image on top of the ubuntu images. In
 83              addition, you must install Java and virtualenv to have the image work properly.
 84      """
 85      get_flavor_backend(model_uri, docker_build=True, env_manager=env_manager).build_image(
 86          model_uri,
 87          name,
 88          mlflow_home=mlflow_home,
 89          install_java=install_java,
 90          install_mlflow=install_mlflow,
 91          enable_mlserver=enable_mlserver,
 92          base_image=base_image,
 93      )
 94  
 95  
 96  _CONTENT_TYPE_CSV = "csv"
 97  _CONTENT_TYPE_JSON = "json"
 98  
 99  
100  def predict(
101      model_uri,
102      input_data=None,
103      input_path=None,
104      content_type=_CONTENT_TYPE_JSON,
105      output_path=None,
106      env_manager=_EnvManager.VIRTUALENV,
107      install_mlflow=False,
108      pip_requirements_override=None,
109      extra_envs=None,
110      # TODO: add an option to force recreating the env
111  ):
112      """
113      Generate predictions in json format using a saved MLflow model. For information about the input
114      data formats accepted by this function, see the following documentation:
115      https://www.mlflow.org/docs/latest/models.html#built-in-deployment-tools.
116  
117      .. note::
118  
119          To increase verbosity for debugging purposes (in order to inspect the full dependency
120          resolver operations when processing transient dependencies), consider setting the following
121          environment variables:
122  
123          .. code-block:: bash
124  
125              # For virtualenv
126              export PIP_VERBOSE=1
127  
128              # For uv
129              export RUST_LOG=uv=debug
130  
131          See also:
132  
133          - https://pip.pypa.io/en/stable/topics/configuration/#environment-variables
134          - https://docs.astral.sh/uv/configuration/environment
135  
136      Args:
137          model_uri: URI to the model. A local path, a local or remote URI e.g. runs:/, s3://.
138          input_data: Input data for prediction. Must be valid input for the PyFunc model. Refer
139              to the :py:func:`mlflow.pyfunc.PyFuncModel.predict()` for the supported input types.
140  
141              .. note::
142                  If this API fails due to errors in input_data, use
143                  `mlflow.models.convert_input_example_to_serving_input` to manually validate
144                  your input data.
145          input_path: Path to a file containing input data. If provided, 'input_data' must be None.
146          content_type: Content type of the input data. Can be one of {'json', 'csv'}.
147          output_path: File to output results to as json. If not provided, output to stdout.
148          env_manager: Specify a way to create an environment for MLmodel inference:
149  
150              - "virtualenv" (default): use virtualenv (and pyenv for Python version management)
151              - "uv": use uv
152              - "local": use the local environment
153              - "conda": use conda
154  
155          install_mlflow: If specified and there is a conda or virtualenv environment to be activated
156              mlflow will be installed into the environment after it has been activated. The version
157              of installed mlflow will be the same as the one used to invoke this command.
158          pip_requirements_override: If specified, install the specified python dependencies to the
159              model inference environment. This is particularly useful when you want to add extra
160              dependencies or try different versions of the dependencies defined in the logged model.
161  
162              .. tip::
163                  After validating the pip requirements override works as expected, you can update
164                  the logged model's dependency using `mlflow.models.update_model_requirements` API
165                  without re-logging it. Note that a registered model is immutable, so you need to
166                  register a new model version with the updated model.
167          extra_envs: If specified, a dictionary of extra environment variables will be passed to the
168              model inference environment. This is useful for testing what environment variables are
169              needed for the model to run correctly. By default, environment variables existing in the
170              current os.environ are passed, and this parameter can be used to override them.
171  
172              .. note::
173                  If your model dependencies include pre-release versions such as `mlflow==3.2.0rc0`
174                  and you are using `uv` as the environment manager, set `UV_PRERELEASE` environment
175                  variable to "allow" in `extra_envs` to allow installing pre-release versions.
176                  e.g. `extra_envs={"UV_PRERELEASE": "allow"}`.
177  
178              .. note::
179                  This parameter is only supported when `env_manager` is set to "virtualenv",
180                  "conda" or "uv".
181  
182      Code example:
183  
184      .. code-block:: python
185  
186          import mlflow
187  
188          run_id = "..."
189  
190          mlflow.models.predict(
191              model_uri=f"runs:/{run_id}/model",
192              input_data={"x": 1, "y": 2},
193              content_type="json",
194          )
195  
196          # Run prediction with "uv" as the environment manager
197          mlflow.models.predict(
198              model_uri=f"runs:/{run_id}/model",
199              input_data={"x": 1, "y": 2},
200              env_manager="uv",
201          )
202  
203          # Run prediction with additional pip dependencies and extra environment variables
204          mlflow.models.predict(
205              model_uri=f"runs:/{run_id}/model",
206              input_data={"x": 1, "y": 2},
207              content_type="json",
208              pip_requirements_override=["scikit-learn==0.23.2"],
209              extra_envs={"OPENAI_API_KEY": "some_value"},
210          )
211  
212          # Run prediction with output_path
213          mlflow.models.predict(
214              model_uri=f"runs:/{run_id}/model",
215              input_data={"x": 1, "y": 2},
216              env_manager="uv",
217              output_path="output.json",
218          )
219  
220          # Run prediction with pre-release versions
221          mlflow.models.predict(
222              model_uri=f"runs:/{run_id}/model",
223              input_data={"x": 1, "y": 2},
224              env_manager="uv",
225              extra_envs={"UV_PRERELEASE": "allow"},
226          )
227  
228      """
229      # to avoid circular imports
230      from mlflow.pyfunc import _PREBUILD_ENV_ROOT_LOCATION
231  
232      if content_type not in [_CONTENT_TYPE_JSON, _CONTENT_TYPE_CSV]:
233          raise MlflowException.invalid_parameter_value(
234              f"Content type must be one of {_CONTENT_TYPE_JSON} or {_CONTENT_TYPE_CSV}."
235          )
236      if extra_envs and env_manager not in (
237          _EnvManager.VIRTUALENV,
238          _EnvManager.CONDA,
239          _EnvManager.UV,
240      ):
241          raise MlflowException.invalid_parameter_value(
242              "Extra environment variables are only supported when env_manager is "
243              f"set to '{_EnvManager.VIRTUALENV}', '{_EnvManager.CONDA}' or '{_EnvManager.UV}'."
244          )
245      if env_manager == _EnvManager.UV:
246          if not shutil.which("uv"):
247              raise MlflowException(
248                  f"Found '{env_manager}' as env_manager, but the 'uv' command is not found in the "
249                  f"PATH. {UV_INSTALLATION_INSTRUCTIONS} Alternatively, you can use 'virtualenv' or "
250                  "'conda' as the environment manager, but note their performances are not "
251                  "as good as 'uv'."
252              )
253      else:
254          _logger.info(
255              f"It is highly recommended to use `{_EnvManager.UV}` as the environment manager for "
256              "predicting with MLflow models as its performance is significantly better than other "
257              f"environment managers. {UV_INSTALLATION_INSTRUCTIONS}"
258          )
259  
260      is_dbconnect_mode = is_databricks_connect()
261      if is_dbconnect_mode:
262          if env_manager not in (_EnvManager.VIRTUALENV, _EnvManager.UV):
263              raise MlflowException(
264                  f"Databricks Connect only supports '{_EnvManager.VIRTUALENV}' or '{_EnvManager.UV}'"
265                  f" as the environment manager. Got {env_manager}."
266              )
267          pyfunc_backend_env_root_config = {
268              "create_env_root_dir": False,
269              "env_root_dir": _PREBUILD_ENV_ROOT_LOCATION,
270          }
271      else:
272          pyfunc_backend_env_root_config = {"create_env_root_dir": True}
273  
274      def _predict(_input_path: str):
275          return get_flavor_backend(
276              model_uri,
277              env_manager=env_manager,
278              install_mlflow=install_mlflow,
279              **pyfunc_backend_env_root_config,
280          ).predict(
281              model_uri=model_uri,
282              input_path=_input_path,
283              output_path=output_path,
284              content_type=content_type,
285              pip_requirements_override=pip_requirements_override,
286              extra_envs=extra_envs,
287          )
288  
289      if input_data is not None and input_path is not None:
290          raise MlflowException.invalid_parameter_value(
291              "Both input_data and input_path are provided. Only one of them should be specified."
292          )
293      elif input_data is not None:
294          input_data = _serialize_input_data(input_data, content_type)
295  
296          # Write input data to a temporary file
297          with TempDir() as tmp:
298              input_path = os.path.join(tmp.path(), f"input.{content_type}")
299              with open(input_path, "w") as f:
300                  f.write(input_data)
301  
302              _predict(input_path)
303      else:
304          _predict(input_path)
305  
306  
307  def _get_pyfunc_supported_input_types():
308      # Importing here as the util module depends on optional packages not available in mlflow-skinny
309      import mlflow.models.utils as base_module
310  
311      supported_input_types = []
312      for input_type in get_args(base_module.PyFuncInput):
313          if isinstance(input_type, type):
314              supported_input_types.append(input_type)
315          elif isinstance(input_type, ForwardRef):
316              name = input_type.__forward_arg__
317              if hasattr(base_module, name):
318                  cls = getattr(base_module, name)
319                  supported_input_types.append(cls)
320          else:
321              # typing instances like List, Dict, Tuple, etc.
322              supported_input_types.append(get_origin(input_type))
323      return tuple(supported_input_types)
324  
325  
326  def _serialize_input_data(input_data, content_type):
327      # build-docker command is available in mlflow-skinny (which doesn't contain pandas)
328      # so we shouldn't import pandas at the top level
329      import pandas as pd
330  
331      # this introduces numpy as dependency, we shouldn't import it at the top level
332      # as it is not available in mlflow-skinny
333      from mlflow.models.utils import convert_input_example_to_serving_input
334  
335      valid_input_types = {
336          _CONTENT_TYPE_CSV: (str, list, dict, pd.DataFrame),
337          _CONTENT_TYPE_JSON: _get_pyfunc_supported_input_types(),
338      }.get(content_type)
339  
340      if not isinstance(input_data, valid_input_types):
341          raise MlflowException.invalid_parameter_value(
342              f"Input data must be one of {valid_input_types} when content type is '{content_type}', "
343              f"but got {type(input_data)}."
344          )
345  
346      if content_type == _CONTENT_TYPE_CSV:
347          if isinstance(input_data, str):
348              _validate_csv_string(input_data)
349              return input_data
350          else:
351              try:
352                  return pd.DataFrame(input_data).to_csv(index=False)
353              except Exception as e:
354                  raise MlflowException.invalid_parameter_value(
355                      "Failed to serialize input data to CSV format."
356                  ) from e
357  
358      try:
359          # rely on convert_input_example_to_serving_input to validate
360          # the input_data is valid type for the loaded pyfunc model
361          return convert_input_example_to_serving_input(input_data)
362      except Exception as e:
363          raise MlflowException.invalid_parameter_value(
364              "Invalid input data, please make sure the data is acceptable by the "
365              "loaded pyfunc model. Use `mlflow.models.convert_input_example_to_serving_input` "
366              "to manually validate your input data."
367          ) from e
368  
369  
370  def _validate_csv_string(input_data: str):
371      """
372      Validate the string must be the path to a CSV file.
373      """
374      try:
375          import pandas as pd
376  
377          pd.read_csv(StringIO(input_data))
378      except Exception as e:
379          raise MlflowException.invalid_parameter_value(
380              message="Failed to deserialize input string data to Pandas DataFrame."
381          ) from e