/ mlflow / utils / cli_args.py
cli_args.py
  1  """
  2  Definitions of click options shared by several CLI commands.
  3  """
  4  
  5  import warnings
  6  
  7  import click
  8  
  9  from mlflow.environment_variables import MLFLOW_DISABLE_ENV_MANAGER_CONDA_WARNING
 10  from mlflow.utils import env_manager as _EnvManager
 11  
 12  MODEL_PATH = click.option(
 13      "--model-path",
 14      "-m",
 15      default=None,
 16      metavar="PATH",
 17      required=True,
 18      help="Path to the model. The path is relative to the run with the given "
 19      "run-id or local filesystem path without run-id.",
 20  )
 21  
 22  _model_uri_help_string = (
 23      "URI to the model. A local path, a 'runs:/' URI, or a"
 24      " remote storage URI (e.g., an 's3://' URI). For more information"
 25      " about supported remote URIs for model artifacts, see"
 26      " https://mlflow.org/docs/latest/tracking.html#artifact-stores"
 27  )
 28  
 29  MODEL_URI_BUILD_DOCKER = click.option(
 30      "--model-uri",
 31      "-m",
 32      metavar="URI",
 33      default=None,
 34      required=False,
 35      help="[Optional] " + _model_uri_help_string,
 36  )
 37  
 38  MODEL_URI = click.option(
 39      "--model-uri",
 40      "-m",
 41      metavar="URI",
 42      required=True,
 43      help=_model_uri_help_string,
 44  )
 45  
 46  MLFLOW_HOME = click.option(
 47      "--mlflow-home",
 48      default=None,
 49      metavar="PATH",
 50      help="Path to local clone of MLflow project. Use for development only.",
 51  )
 52  
 53  RUN_ID = click.option(
 54      "--run-id",
 55      "-r",
 56      default=None,
 57      required=False,
 58      metavar="ID",
 59      help="ID of the MLflow run that generated the referenced content.",
 60  )
 61  
 62  
 63  def _resolve_env_manager(_, __, env_manager):
 64      if env_manager is not None:
 65          _EnvManager.validate(env_manager)
 66          if env_manager == _EnvManager.CONDA and not MLFLOW_DISABLE_ENV_MANAGER_CONDA_WARNING.get():
 67              warnings.warn(
 68                  (
 69                      "Use of conda is discouraged. If you use it, please ensure that your use of "
 70                      "conda complies with Anaconda's terms of service "
 71                      "(https://legal.anaconda.com/policies/en/?name=terms-of-service). "
 72                      "virtualenv is the recommended tool for environment reproducibility. "
 73                      f"To suppress this warning, set the {MLFLOW_DISABLE_ENV_MANAGER_CONDA_WARNING} "
 74                      "environment variable to 'TRUE'."
 75                  ),
 76                  UserWarning,
 77                  stacklevel=2,
 78              )
 79          return env_manager
 80  
 81      return None
 82  
 83  
 84  def _create_env_manager_option(help_string, default=None):
 85      return click.option(
 86          "--env-manager",
 87          default=default,
 88          type=click.UNPROCESSED,
 89          callback=_resolve_env_manager,
 90          help=help_string,
 91      )
 92  
 93  
 94  ENV_MANAGER = _create_env_manager_option(
 95      default=_EnvManager.VIRTUALENV,
 96      # '\b' prevents rewrapping text:
 97      # https://click.palletsprojects.com/en/8.1.x/documentation/#preventing-rewrapping
 98      help_string="""
 99  If specified, create an environment for MLmodel using the specified
100  environment manager. The following values are supported:
101  
102  \b
103  - local: use the local environment
104  - virtualenv: use venv (and pyenv for Python version management)
105  - uv: use uv
106  - conda: use conda
107  
108  If unspecified, default to virtualenv.
109  """,
110  )
111  
112  ENV_MANAGER_PROJECTS = _create_env_manager_option(
113      help_string="""
114  If specified, create an environment for MLproject using the specified
115  environment manager. The following values are supported:
116  
117  \b
118  - local: use the local environment
119  - virtualenv: use venv (and pyenv for Python version management)
120  - uv: use uv
121  - conda: use conda
122  
123  If unspecified, the appropriate environment manager is automatically selected based on
124  the project configuration. For example, if `MLproject.yaml` contains a `python_env` key,
125  virtualenv is used.
126  """,
127  )
128  
129  ENV_MANAGER_DOCKERFILE = _create_env_manager_option(
130      default=None,
131      # '\b' prevents rewrapping text:
132      # https://click.palletsprojects.com/en/8.1.x/documentation/#preventing-rewrapping
133      help_string="""
134  If specified, create an environment for MLmodel using the specified
135  environment manager. The following values are supported:
136  
137  \b
138  - local: use the local environment
139  - virtualenv: use venv (and pyenv for Python version management)
140  - uv: use uv
141  - conda: use conda
142  
143  If unspecified, default to None, then MLflow will automatically pick the env manager
144  based on the model's flavor configuration.
145  If model-uri is specified: if python version is specified in the flavor configuration
146  and no java installation is required, then we use local environment. Otherwise we use virtualenv.
147  If no model-uri is provided, we use virtualenv.
148  """,
149  )
150  
151  
152  INSTALL_MLFLOW = click.option(
153      "--install-mlflow",
154      is_flag=True,
155      default=False,
156      help="If specified and there is a conda, virtualenv, or uv environment to be activated "
157      "mlflow will be installed into the environment after it has been "
158      "activated. The version of installed mlflow will be the same as "
159      "the one used to invoke this command.",
160  )
161  
162  HOST = click.option(
163      "--host",
164      "-h",
165      envvar="MLFLOW_HOST",
166      metavar="HOST",
167      default="127.0.0.1",
168      help="The network interface to bind the server to (default: 127.0.0.1). "
169      "This controls which network interfaces accept connections. "
170      "Use '127.0.0.1' for local-only access, or '0.0.0.0' to allow connections from any network. "
171      "NOTE: This is NOT a security setting - it only controls network binding. "
172      "To restrict which clients can connect, use --allowed-hosts.",
173  )
174  
175  PORT = click.option(
176      "--port",
177      "-p",
178      envvar="MLFLOW_PORT",
179      default=5000,
180      help="The port to listen on (default: 5000).",
181  )
182  
183  TIMEOUT = click.option(
184      "--timeout",
185      "-t",
186      envvar="MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT",
187      default=60,
188      help="Timeout in seconds to serve a request (default: 60).",
189  )
190  
191  # We use None to disambiguate manually selecting "4"
192  WORKERS = click.option(
193      "--workers",
194      "-w",
195      envvar="MLFLOW_WORKERS",
196      default=None,
197      help="Number of worker processes to handle requests (default: 4).",
198  )
199  
200  MODELS_WORKERS = click.option(
201      "--workers",
202      "-w",
203      envvar="MLFLOW_MODELS_WORKERS",
204      default=None,
205      help="Number of uvicorn workers to handle requests when serving mlflow models (default: 1).",
206  )
207  
208  ENABLE_MLSERVER = click.option(
209      "--enable-mlserver",
210      is_flag=True,
211      default=False,
212      help=(
213          "Enable serving with MLServer through the v2 inference protocol. "
214          "You can use environment variables to configure MLServer. "
215          "(See https://mlserver.readthedocs.io/en/latest/reference/settings.html)"
216      ),
217  )
218  
219  ARTIFACTS_DESTINATION = click.option(
220      "--artifacts-destination",
221      envvar="MLFLOW_ARTIFACTS_DESTINATION",
222      metavar="URI",
223      default="./mlartifacts",
224      help=(
225          "The base artifact location from which to resolve artifact upload/download/list requests "
226          "(e.g. 's3://my-bucket'). Defaults to a local './mlartifacts' directory. This option only "
227          "applies when the tracking server is configured to stream artifacts and the experiment's "
228          "artifact root location is http or mlflow-artifacts URI."
229      ),
230  )
231  
232  SERVE_ARTIFACTS = click.option(
233      "--serve-artifacts/--no-serve-artifacts",
234      envvar="MLFLOW_SERVE_ARTIFACTS",
235      is_flag=True,
236      default=True,
237      help="Enables serving of artifact uploads, downloads, and list requests "
238      "by routing these requests to the storage location that is specified by "
239      "'--artifacts-destination' directly through a proxy. The default location that "
240      "these requests are served from is a local './mlartifacts' directory which can be "
241      "overridden via the '--artifacts-destination' argument. To disable artifact serving, "
242      "specify `--no-serve-artifacts`. Default: True",
243  )
244  
245  NO_CONDA = click.option(
246      "--no-conda",
247      is_flag=True,
248      help="If specified, use local environment.",
249  )
250  
251  INSTALL_JAVA = click.option(
252      "--install-java",
253      is_flag=False,
254      flag_value=True,
255      default=None,
256      type=bool,
257      help="Installs Java in the image if needed. Default is None, "
258      "allowing MLflow to determine installation. Flavors requiring "
259      "Java, such as Spark, enable this automatically. "
260      "Note: This option only works with the UBUNTU base image; "
261      "Python base images do not support Java installation.",
262  )
263  
264  # Security-related options for MLflow server
265  ALLOWED_HOSTS = click.option(
266      "--allowed-hosts",
267      envvar="MLFLOW_SERVER_ALLOWED_HOSTS",
268      default=None,
269      help="Comma-separated list of allowed Host headers to prevent DNS rebinding attacks "
270      "(default: localhost + private IPs). "
271      "DNS rebinding allows attackers to trick your browser into accessing internal services. "
272      "Examples: 'mlflow.company.com,10.0.0.100:5000'. "
273      "Supports wildcards: 'mlflow.company.com,192.168.*,app-*.internal.com'. "
274      "Use '*' to allow ALL hosts (not recommended for production). "
275      "Default allows: localhost (all ports), private IPs (10.*, 192.168.*, 172.16-31.*). "
276      "Set this when exposing MLflow beyond localhost to prevent host header attacks.",
277  )
278  
279  CORS_ALLOWED_ORIGINS = click.option(
280      "--cors-allowed-origins",
281      envvar="MLFLOW_SERVER_CORS_ALLOWED_ORIGINS",
282      default=None,
283      help="Comma-separated list of allowed CORS origins to prevent cross-site request attacks "
284      "(default: localhost origins on any port). "
285      "CORS attacks allow malicious websites to make requests to your MLflow server using your "
286      "credentials. Examples: 'https://app.company.com,https://notebook.company.com'. "
287      "Default allows: http://localhost:* (any port), http://127.0.0.1:*, http://[::1]:*. "
288      "Set this when you have web applications on different domains that need to access MLflow. "
289      "Use '*' to allow ALL origins (DANGEROUS - only for development!).",
290  )
291  
292  DISABLE_SECURITY_MIDDLEWARE = click.option(
293      "--disable-security-middleware",
294      envvar="MLFLOW_SERVER_DISABLE_SECURITY_MIDDLEWARE",
295      is_flag=True,
296      default=False,
297      help="DANGEROUS: Disable all security middleware including CORS protection and host "
298      "validation. This completely removes security protections and should only be used for "
299      "testing. When disabled, your MLflow server is vulnerable to CORS attacks, DNS rebinding, "
300      "and clickjacking. Instead, prefer configuring specific security settings with "
301      "--cors-allowed-origins and --allowed-hosts.",
302  )
303  
304  X_FRAME_OPTIONS = click.option(
305      "--x-frame-options",
306      envvar="MLFLOW_SERVER_X_FRAME_OPTIONS",
307      default="SAMEORIGIN",
308      help="X-Frame-Options header value for clickjacking protection. "
309      "Options: 'SAMEORIGIN' (default - allows embedding only from same origin), "
310      "'DENY' (prevents all embedding), 'NONE' (disables header - allows embedding from anywhere). "
311      "Set to 'NONE' if you need to embed MLflow UI in iframes from different origins.",
312  )