1  import logging
  2  
  3  import click
  4  
  5  from mlflow.mcp.decorator import mlflow_mcp
  6  from mlflow.models import python_api
  7  from mlflow.models.flavor_backend_registry import get_flavor_backend
  8  from mlflow.models.model import update_model_requirements
  9  from mlflow.utils import cli_args
 10  from mlflow.utils import env_manager as _EnvManager
 11  
 12  _logger = logging.getLogger(__name__)
 13  
 14  
 15  @click.group("models")
 16  def commands():
 17      """
 18      Deploy MLflow models locally.
 19  
 20      To deploy a model associated with a run on a tracking server, set the MLFLOW_TRACKING_URI
 21      environment variable to the URL of the desired server.
 22      """
 23  
 24  
 25  @commands.command("serve")
 26  @mlflow_mcp(tool_name="serve_model")
 27  @cli_args.MODEL_URI
 28  @cli_args.PORT
 29  @cli_args.HOST
 30  @cli_args.TIMEOUT
 31  @cli_args.MODELS_WORKERS
 32  @cli_args.ENV_MANAGER
 33  @cli_args.NO_CONDA
 34  @cli_args.INSTALL_MLFLOW
 35  @cli_args.ENABLE_MLSERVER
 36  def serve(
 37      model_uri,
 38      port,
 39      host,
 40      timeout,
 41      workers,
 42      env_manager=None,
 43      no_conda=False,
 44      install_mlflow=False,
 45      enable_mlserver=False,
 46  ):
 47      """
 48      Serve a model saved with MLflow by launching a webserver on the specified host and port.
 49      The command supports models with the ``python_function`` or ``crate`` (R Function) flavor.
 50      For information about the input data formats accepted by the webserver, see the following
 51      documentation: https://www.mlflow.org/docs/latest/models.html#built-in-deployment-tools.
 52  
 53      .. warning::
 54  
 55          Models built using MLflow 1.x will require adjustments to the endpoint request payload
 56          if executed in an environment that has MLflow 2.x installed. In 1.x, a request payload
 57          was in the format: ``{'columns': [str], 'data': [[...]]}``. 2.x models require
 58          payloads that are defined by the structural-defining keys of either ``dataframe_split``,
 59          ``instances``, ``inputs`` or ``dataframe_records``. See the examples below for
 60          demonstrations of the changes to the invocation API endpoint in 2.0.
 61  
 62      .. note::
 63  
 64          Requests made in pandas DataFrame structures can be made in either `split` or `records`
 65          oriented formats.
 66          See https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.to_json.html for
 67          detailed information on orientation formats for converting a pandas DataFrame to json.
 68  
 69      Example:
 70  
 71      .. code-block:: bash
 72  
 73          $ mlflow models serve -m runs:/my-run-id/model-path &
 74  
 75          # records orientation input format for serializing a pandas DataFrame
 76          $ curl http://127.0.0.1:5000/invocations -H 'Content-Type: application/json' -d '{
 77              "dataframe_records": [{"a":1, "b":2}, {"a":3, "b":4}, {"a":5, "b":6}]
 78          }'
 79  
 80          # split orientation input format for serializing a pandas DataFrame
 81          $ curl http://127.0.0.1:5000/invocations -H 'Content-Type: application/json' -d '{
 82              "dataframe_split": {"columns": ["a", "b"],
 83                                  "index": [0, 1, 2],
 84                                  "data": [[1, 2], [3, 4], [5, 6]]}
 85          }'
 86  
 87          # inputs format for List submission of array, tensor, or DataFrame data
 88          $ curl http://127.0.0.1:5000/invocations -H 'Content-Type: application/json' -d '{
 89              "inputs": [[1, 2], [3, 4], [5, 6]]
 90          }'
 91  
 92          # instances format for submission of Tensor data
 93          curl http://127.0.0.1:5000/invocations -H 'Content-Type: application/json' -d '{
 94              "instances": [
 95                  {"a": "t1", "b": [1, 2, 3]},
 96                  {"a": "t2", "b": [4, 5, 6]},
 97                  {"a": "t3", "b": [7, 8, 9]}
 98              ]
 99          }'
100  
101      """
102      env_manager = _EnvManager.LOCAL if no_conda else env_manager
103  
104      return get_flavor_backend(
105          model_uri, env_manager=env_manager, workers=workers, install_mlflow=install_mlflow
106      ).serve(
107          model_uri=model_uri, port=port, host=host, timeout=timeout, enable_mlserver=enable_mlserver
108      )
109  
110  
111  class KeyValueType(click.ParamType):
112      name = "key=value"
113  
114      def convert(self, value, param, ctx):
115          if "=" not in value:
116              self.fail(f"{value!r} is not a valid key value pair, expecting `key=value`", param, ctx)
117          return value.split("=", 1)
118  
119  
120  @commands.command("predict")
121  @mlflow_mcp(tool_name="predict_with_model")
122  @cli_args.MODEL_URI
123  @click.option(
124      "--input-path", "-i", default=None, help="CSV containing pandas DataFrame to predict against."
125  )
126  @click.option(
127      "--output-path",
128      "-o",
129      default=None,
130      help="File to output results to as json file. If not provided, output to stdout.",
131  )
132  @click.option(
133      "--content-type",
134      "-t",
135      default="json",
136      help="Content type of the input file. Can be one of {'json', 'csv'}.",
137  )
138  @cli_args.ENV_MANAGER
139  @cli_args.INSTALL_MLFLOW
140  @click.option(
141      "--pip-requirements-override",
142      "-r",
143      default=None,
144      help="Specify packages and versions to override the dependencies defined "
145      "in the model. Must be a comma-separated string like x==y,z==a.",
146  )
147  @click.option(
148      "--env",
149      default=None,
150      type=KeyValueType(),
151      multiple=True,
152      help="Extra environment variables to set when running the model. Must be "
153      "key value pairs, e.g. `--env key=value`.",
154  )
155  def predict(
156      model_uri,
157      input_data=None,
158      input_path=None,
159      content_type=python_api._CONTENT_TYPE_JSON,
160      output_path=None,
161      env_manager=_EnvManager.VIRTUALENV,
162      install_mlflow=False,
163      pip_requirements_override=None,
164      env=None,
165  ):
166      """
167      Generate predictions in json format using a saved MLflow model. For information about the input
168      data formats accepted by this function, see the following documentation:
169      https://www.mlflow.org/docs/latest/models.html#built-in-deployment-tools.
170      """
171      return python_api.predict(
172          model_uri=model_uri,
173          input_data=input_data,
174          input_path=input_path,
175          content_type=content_type,
176          output_path=output_path,
177          env_manager=env_manager,
178          install_mlflow=install_mlflow,
179          pip_requirements_override=pip_requirements_override,
180          extra_envs=dict(env),
181      )
182  
183  
184  @commands.command("prepare-env")
185  @mlflow_mcp(tool_name="prepare_model_env")
186  @cli_args.MODEL_URI
187  @cli_args.ENV_MANAGER
188  @cli_args.INSTALL_MLFLOW
189  def prepare_env(
190      model_uri,
191      env_manager,
192      install_mlflow,
193  ):
194      """
195      Performs any preparation necessary to predict or serve the model, for example
196      downloading dependencies or initializing a conda environment. After preparation,
197      calling predict or serve should be fast.
198      """
199      return get_flavor_backend(
200          model_uri, env_manager=env_manager, install_mlflow=install_mlflow
201      ).prepare_env(model_uri=model_uri)
202  
203  
204  @commands.command("generate-dockerfile")
205  @mlflow_mcp(tool_name="generate_model_dockerfile")
206  @cli_args.MODEL_URI_BUILD_DOCKER
207  @click.option(
208      "--output-directory",
209      "-d",
210      default="mlflow-dockerfile",
211      help="Output directory where the generated Dockerfile is stored.",
212  )
213  @cli_args.ENV_MANAGER_DOCKERFILE
214  @cli_args.MLFLOW_HOME
215  @cli_args.INSTALL_JAVA
216  @cli_args.INSTALL_MLFLOW
217  @cli_args.ENABLE_MLSERVER
218  def generate_dockerfile(
219      model_uri,
220      output_directory,
221      env_manager,
222      mlflow_home,
223      install_java,
224      install_mlflow,
225      enable_mlserver,
226  ):
227      """
228      Generates a directory with Dockerfile whose default entrypoint serves an MLflow model at port
229      8080 using the python_function flavor. The generated Dockerfile is written to the specified
230      output directory, along with the model (if specified). This Dockerfile defines an image that
231      is equivalent to the one produced by ``mlflow models build-docker``.
232      """
233      if model_uri:
234          _logger.info("Generating Dockerfile for model %s", model_uri)
235      else:
236          _logger.info("Generating Dockerfile")
237      backend = get_flavor_backend(model_uri, docker_build=True, env_manager=env_manager)
238      if backend.can_build_image():
239          backend.generate_dockerfile(
240              model_uri,
241              output_directory,
242              mlflow_home=mlflow_home,
243              install_java=install_java,
244              install_mlflow=install_mlflow,
245              enable_mlserver=enable_mlserver,
246          )
247          _logger.info("Generated Dockerfile in directory %s", output_directory)
248      else:
249          _logger.error(
250              "Cannot build docker image for selected backend",
251              extra={"backend": backend.__class__.__name__},
252          )
253          raise NotImplementedError("Cannot build docker image for selected backend")
254  
255  
256  @commands.command("build-docker")
257  @mlflow_mcp(tool_name="build_model_docker")
258  @cli_args.MODEL_URI_BUILD_DOCKER
259  @click.option("--name", "-n", default="mlflow-pyfunc-servable", help="Name to use for built image")
260  @cli_args.ENV_MANAGER
261  @cli_args.MLFLOW_HOME
262  @cli_args.INSTALL_JAVA
263  @cli_args.INSTALL_MLFLOW
264  @cli_args.ENABLE_MLSERVER
265  def build_docker(**kwargs):
266      """
267      Builds a Docker image whose default entrypoint serves an MLflow model at port 8080, using the
268      python_function flavor. The container serves the model referenced by ``--model-uri``, if
269      specified when ``build-docker`` is called. If ``--model-uri`` is not specified when build_docker
270      is called, an MLflow Model directory must be mounted as a volume into the /opt/ml/model
271      directory in the container.
272  
273      Building a Docker image with ``--model-uri``:
274  
275      .. code:: bash
276  
277          # Build a Docker image named 'my-image-name' that serves the model from run 'some-run-uuid'
278          # at run-relative artifact path 'my-model'
279          mlflow models build-docker --model-uri "runs:/some-run-uuid/my-model" --name "my-image-name"
280          # Serve the model
281          docker run -p 5001:8080 "my-image-name"
282  
283      Building a Docker image without ``--model-uri``:
284  
285      .. code:: bash
286  
287          # Build a generic Docker image named 'my-image-name'
288          mlflow models build-docker --name "my-image-name"
289          # Mount the model stored in '/local/path/to/artifacts/model' and serve it
290          docker run --rm -p 5001:8080 -v /local/path/to/artifacts/model:/opt/ml/model "my-image-name"
291  
292      .. important::
293  
294          Since MLflow 2.10.1, the Docker image built with ``--model-uri`` does **not install Java**
295          for improved performance, unless the model flavor is one of ``["johnsnowlabs", "h2o",
296          "spark"]``. If you need to install Java for other flavors, e.g. custom Python model
297          that uses SparkML, please specify the ``--install-java`` flag to enforce Java installation.
298  
299      NB: by default, the container will start nginx and uvicorn processes. If you don't need the
300      nginx process to be started (for instance if you deploy your container to Google Cloud Run),
301      you can disable it via the DISABLE_NGINX environment variable:
302  
303      .. code:: bash
304  
305          docker run -p 5001:8080 -e DISABLE_NGINX=true "my-image-name"
306  
307      By default, the number of uvicorn workers is set to CPU count. If you want to set a custom
308      number of workers, you can set the MLFLOW_MODELS_WORKERS environment variable:
309  
310      .. code:: bash
311  
312          docker run -p 5001:8080 -e MLFLOW_MODELS_WORKERS=4 "my-image-name"
313  
314      See https://www.mlflow.org/docs/latest/python_api/mlflow.pyfunc.html for more information on the
315      'python_function' flavor.
316      """
317      python_api.build_docker(**kwargs)
318  
319  
320  @commands.command("update-pip-requirements")
321  @mlflow_mcp(tool_name="update_model_pip_requirements")
322  @cli_args.MODEL_URI
323  @click.argument("operation", type=click.Choice(["add", "remove"]))
324  @click.argument("requirement_strings", type=str, nargs=-1)
325  def update_pip_requirements(model_uri, operation, requirement_strings):
326      """
327      Add or remove requirements from a model's conda.yaml and requirements.txt files.
328      If using a remote tracking server, please make sure to set the MLFLOW_TRACKING_URI
329      environment variable to the URL of the desired server.
330  
331      REQUIREMENT_STRINGS is a list of pip requirements specifiers.
332      See below for examples.
333  
334      Sample usage:
335  
336      .. code::
337  
338          # Add requirements using the model's "runs:/" URI
339  
340          mlflow models update-pip-requirements -m runs:/<run_id>/<model_path> \\
341              add "pandas==1.0.0" "scikit-learn" "mlflow >= 2.8, != 2.9.0"
342  
343          # Remove requirements from a local model
344  
345          mlflow models update-pip-requirements -m /path/to/local/model \\
346              remove "torchvision" "pydantic"
347  
348      Note that model registry URIs (i.e. URIs in the form ``models:/``) are not
349      supported, as artifacts in the model registry are intended to be read-only.
350      Editing requirements is read-only artifact repositories is also not supported.
351  
352      If adding requirements, the function will overwrite any existing requirements
353      that overlap, or else append the new requirements to the existing list.
354  
355      If removing requirements, the function will ignore any version specifiers,
356      and remove all the specified package names. Any requirements that are not
357      found in the existing files will be ignored.
358      """
359      update_model_requirements(model_uri, operation, requirement_strings)
360  
361      _logger.info(f"Successfully updated the requirements for the model at {model_uri}!")