star

venv/lib/python3.9/site-packages/mlflow/projects/__init__.py

@@ -0,0 +1,449 @@
+"""
+The ``mlflow.projects`` module provides an API for running MLflow projects locally or remotely.
+"""
+
+import json
+import logging
+import os
+
+import yaml
+
+import mlflow.projects.databricks
+import mlflow.utils.uri
+from mlflow import tracking
+from mlflow.entities import RunStatus
+from mlflow.exceptions import ExecutionException, MlflowException
+from mlflow.projects.backend import loader
+from mlflow.projects.submitted_run import SubmittedRun
+from mlflow.projects.utils import (
+    MLFLOW_LOCAL_BACKEND_RUN_ID_CONFIG,
+    PROJECT_BUILD_IMAGE,
+    PROJECT_DOCKER_ARGS,
+    PROJECT_DOCKER_AUTH,
+    PROJECT_ENV_MANAGER,
+    PROJECT_STORAGE_DIR,
+    PROJECT_SYNCHRONOUS,
+    fetch_and_validate_project,
+    get_entry_point_command,
+    get_or_create_run,
+    get_run_env_vars,
+    load_project,
+)
+from mlflow.tracking.fluent import _get_experiment_id
+from mlflow.utils import env_manager as _EnvManager
+from mlflow.utils.mlflow_tags import (
+    MLFLOW_DOCKER_IMAGE_ID,
+    MLFLOW_PROJECT_BACKEND,
+    MLFLOW_PROJECT_ENV,
+    MLFLOW_RUN_NAME,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+def _resolve_experiment_id(experiment_name=None, experiment_id=None):
+    """
+    Resolve experiment.
+
+    Verifies either one or other is specified - cannot be both selected.
+
+    If ``experiment_name`` is provided and does not exist, an experiment
+    of that name is created and its id is returned.
+
+    Args:
+        experiment_name: Name of experiment under which to launch the run.
+        experiment_id: ID of experiment under which to launch the run.
+
+    Returns:
+        str
+    """
+
+    if experiment_name and experiment_id:
+        raise MlflowException("Specify only one of 'experiment_name' or 'experiment_id'.")
+
+    if experiment_id:
+        return str(experiment_id)
+
+    if experiment_name:
+        client = tracking.MlflowClient()
+        exp = client.get_experiment_by_name(experiment_name)
+        if exp:
+            return exp.experiment_id
+        else:
+            _logger.info("'%s' does not exist. Creating a new experiment", experiment_name)
+            return client.create_experiment(experiment_name)
+
+    return _get_experiment_id()
+
+
+def _run(
+    uri,
+    experiment_id,
+    entry_point,
+    version,
+    parameters,
+    docker_args,
+    backend_name,
+    backend_config,
+    storage_dir,
+    env_manager,
+    synchronous,
+    run_name,
+    build_image,
+    docker_auth,
+):
+    """
+    Helper that delegates to the project-running method corresponding to the passed-in backend.
+    Returns a ``SubmittedRun`` corresponding to the project run.
+    """
+    tracking_store_uri = tracking.get_tracking_uri()
+    backend_config[PROJECT_ENV_MANAGER] = env_manager
+    backend_config[PROJECT_SYNCHRONOUS] = synchronous
+    backend_config[PROJECT_DOCKER_ARGS] = docker_args
+    backend_config[PROJECT_STORAGE_DIR] = storage_dir
+    backend_config[PROJECT_BUILD_IMAGE] = build_image
+    backend_config[PROJECT_DOCKER_AUTH] = docker_auth
+    # TODO: remove this check once kubernetes execution has been refactored
+    if backend_name not in {"databricks", "kubernetes"}:
+        backend = loader.load_backend(backend_name)
+        if backend:
+            submitted_run = backend.run(
+                uri,
+                entry_point,
+                parameters,
+                version,
+                backend_config,
+                tracking_store_uri,
+                experiment_id,
+            )
+            tracking.MlflowClient().set_tag(
+                submitted_run.run_id, MLFLOW_PROJECT_BACKEND, backend_name
+            )
+            if run_name is not None:
+                tracking.MlflowClient().set_tag(submitted_run.run_id, MLFLOW_RUN_NAME, run_name)
+            return submitted_run
+
+    work_dir = fetch_and_validate_project(uri, version, entry_point, parameters)
+    project = load_project(work_dir)
+    _validate_execution_environment(project, backend_name)
+
+    active_run = get_or_create_run(
+        None, uri, experiment_id, work_dir, version, entry_point, parameters
+    )
+
+    if run_name is not None:
+        tracking.MlflowClient().set_tag(active_run.info.run_id, MLFLOW_RUN_NAME, run_name)
+
+    if backend_name == "databricks":
+        tracking.MlflowClient().set_tag(
+            active_run.info.run_id, MLFLOW_PROJECT_BACKEND, "databricks"
+        )
+        from mlflow.projects.databricks import run_databricks, run_databricks_spark_job
+
+        if project.databricks_spark_job_spec is not None:
+            return run_databricks_spark_job(
+                remote_run=active_run,
+                uri=uri,
+                work_dir=work_dir,
+                experiment_id=experiment_id,
+                cluster_spec=backend_config,
+                project_spec=project,
+                entry_point=entry_point,
+                parameters=parameters,
+            )
+
+        return run_databricks(
+            remote_run=active_run,
+            uri=uri,
+            entry_point=entry_point,
+            work_dir=work_dir,
+            parameters=parameters,
+            experiment_id=experiment_id,
+            cluster_spec=backend_config,
+            env_manager=env_manager,
+        )
+
+    elif backend_name == "kubernetes":
+        from mlflow.projects import kubernetes as kb
+        from mlflow.projects.docker import (
+            build_docker_image,
+            validate_docker_env,
+            validate_docker_installation,
+        )
+
+        tracking.MlflowClient().set_tag(active_run.info.run_id, MLFLOW_PROJECT_ENV, "docker")
+        tracking.MlflowClient().set_tag(
+            active_run.info.run_id, MLFLOW_PROJECT_BACKEND, "kubernetes"
+        )
+        validate_docker_env(project)
+        validate_docker_installation()
+        kube_config = _parse_kubernetes_config(backend_config)
+        image = build_docker_image(
+            work_dir=work_dir,
+            repository_uri=kube_config["repository-uri"],
+            base_image=project.docker_env.get("image"),
+            run_id=active_run.info.run_id,
+            build_image=build_image,
+            docker_auth=docker_auth,
+        )
+        image_digest = kb.push_image_to_registry(image.tags[0])
+        tracking.MlflowClient().set_tag(
+            active_run.info.run_id, MLFLOW_DOCKER_IMAGE_ID, image_digest
+        )
+        return kb.run_kubernetes_job(
+            project.name,
+            active_run,
+            image.tags[0],
+            image_digest,
+            get_entry_point_command(project, entry_point, parameters, storage_dir),
+            get_run_env_vars(
+                run_id=active_run.info.run_uuid, experiment_id=active_run.info.experiment_id
+            ),
+            kube_config.get("kube-context", None),
+            kube_config["kube-job-template"],
+        )
+
+    supported_backends = ["databricks", "kubernetes"] + list(loader.MLFLOW_BACKENDS.keys())
+    raise ExecutionException(
+        f"Got unsupported execution mode {backend_name}. Supported values: {supported_backends}"
+    )
+
+
+def run(
+    uri,
+    entry_point="main",
+    version=None,
+    parameters=None,
+    docker_args=None,
+    experiment_name=None,
+    experiment_id=None,
+    backend="local",
+    backend_config=None,
+    storage_dir=None,
+    synchronous=True,
+    run_id=None,
+    run_name=None,
+    env_manager=None,
+    build_image=False,
+    docker_auth=None,
+):
+    """
+    Run an MLflow project. The project can be local or stored at a Git URI.
+
+    MLflow provides built-in support for running projects locally or remotely on a Databricks or
+    Kubernetes cluster. You can also run projects against other targets by installing an appropriate
+    third-party plugin. See `Community Plugins <../plugins.html#community-plugins>`_ for more
+    information.
+
+    For information on using this method in chained workflows, see `Building Multistep Workflows
+    <../projects.html#building-multistep-workflows>`_.
+
+    Raises:
+        :py:class:`mlflow.exceptions.ExecutionException` If a run launched in blocking mode
+            is unsuccessful.
+
+    Args:
+        uri: URI of project to run. A local filesystem path
+            or a Git repository URI (e.g. https://github.com/mlflow/mlflow-example)
+            pointing to a project directory containing an MLproject file.
+        entry_point: Entry point to run within the project. If no entry point with the specified
+            name is found, runs the project file ``entry_point`` as a script,
+            using "python" to run ``.py`` files and the default shell (specified by
+            environment variable ``$SHELL``) to run ``.sh`` files.
+        version: For Git-based projects, either a commit hash or a branch name.
+        parameters: Parameters (dictionary) for the entry point command.
+        docker_args: Arguments (dictionary) for the docker command.
+        experiment_name: Name of experiment under which to launch the run.
+        experiment_id: ID of experiment under which to launch the run.
+        backend: Execution backend for the run: MLflow provides built-in support for "local",
+            "databricks", and "kubernetes" (experimental) backends. If running against
+            Databricks, will run against a Databricks workspace determined as follows:
+            if a Databricks tracking URI of the form ``databricks://profile`` has been set
+            (e.g. by setting the MLFLOW_TRACKING_URI environment variable), will run
+            against the workspace specified by <profile>. Otherwise, runs against the
+            workspace specified by the default Databricks CLI profile.
+        backend_config: A dictionary, or a path to a JSON file (must end in '.json'), which will
+            be passed as config to the backend. The exact content which should be
+            provided is different for each execution backend and is documented
+            at https://www.mlflow.org/docs/latest/projects.html.
+        storage_dir: Used only if ``backend`` is "local". MLflow downloads artifacts from
+            distributed URIs passed to parameters of type ``path`` to subdirectories of
+            ``storage_dir``.
+        synchronous: Whether to block while waiting for a run to complete. Defaults to True.
+            Note that if ``synchronous`` is False and ``backend`` is "local", this
+            method will return, but the current process will block when exiting until
+            the local run completes. If the current process is interrupted, any
+            asynchronous runs launched via this method will be terminated. If
+            ``synchronous`` is True and the run fails, the current process will
+            error out as well.
+        run_id: Note: this argument is used internally by the MLflow project APIs and should
+            not be specified. If specified, the run ID will be used instead of
+            creating a new run.
+        run_name: The name to give the MLflow Run associated with the project execution.
+            If ``None``, the MLflow Run name is left unset.
+        env_manager: Specify an environment manager to create a new environment for the run and
+            install project dependencies within that environment. The following values
+            are supported:
+
+            - local: use the local environment
+            - virtualenv: use virtualenv (and pyenv for Python version management)
+            - conda: use conda
+
+            If unspecified, MLflow automatically determines the environment manager to
+            use by inspecting files in the project directory. For example, if
+            ``python_env.yaml`` is present, virtualenv will be used.
+        build_image: Whether to build a new docker image of the project or to reuse an existing
+            image. Default: False (reuse an existing image)
+        docker_auth: A dictionary representing information to authenticate with a Docker
+            registry. See `docker.client.DockerClient.login
+            <https://docker-py.readthedocs.io/en/stable/client.html#docker.client.DockerClient.login>`_
+            for available options.
+
+    Returns:
+        :py:class:`mlflow.projects.SubmittedRun` exposing information (e.g. run ID)
+        about the launched run.
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+
+        project_uri = "https://github.com/mlflow/mlflow-example"
+        params = {"alpha": 0.5, "l1_ratio": 0.01}
+
+        # Run MLflow project and create a reproducible conda environment
+        # on a local host
+        mlflow.run(project_uri, parameters=params)
+
+    .. code-block:: text
+        :caption: Output
+
+        ...
+        ...
+        Elasticnet model (alpha=0.500000, l1_ratio=0.010000):
+        RMSE: 0.788347345611717
+        MAE: 0.6155576449938276
+        R2: 0.19729662005412607
+        ... mlflow.projects: === Run (ID '6a5109febe5e4a549461e149590d0a7c') succeeded ===
+    """
+    backend_config_dict = backend_config if backend_config is not None else {}
+    if (
+        backend_config
+        and type(backend_config) != dict
+        and os.path.splitext(backend_config)[-1] == ".json"
+    ):
+        with open(backend_config) as handle:
+            try:
+                backend_config_dict = json.load(handle)
+            except ValueError:
+                _logger.error(
+                    "Error when attempting to load and parse JSON cluster spec from file %s",
+                    backend_config,
+                )
+                raise
+
+    if env_manager is not None:
+        _EnvManager.validate(env_manager)
+
+    if backend == "databricks":
+        mlflow.projects.databricks.before_run_validations(mlflow.get_tracking_uri(), backend_config)
+    elif backend == "local" and run_id is not None:
+        backend_config_dict[MLFLOW_LOCAL_BACKEND_RUN_ID_CONFIG] = run_id
+
+    experiment_id = _resolve_experiment_id(
+        experiment_name=experiment_name, experiment_id=experiment_id
+    )
+
+    submitted_run_obj = _run(
+        uri=uri,
+        experiment_id=experiment_id,
+        entry_point=entry_point,
+        version=version,
+        parameters=parameters,
+        docker_args=docker_args,
+        backend_name=backend,
+        backend_config=backend_config_dict,
+        env_manager=env_manager,
+        storage_dir=storage_dir,
+        synchronous=synchronous,
+        run_name=run_name,
+        build_image=build_image,
+        docker_auth=docker_auth,
+    )
+    if synchronous:
+        _wait_for(submitted_run_obj)
+    return submitted_run_obj
+
+
+def _wait_for(submitted_run_obj):
+    """Wait on the passed-in submitted run, reporting its status to the tracking server."""
+    run_id = submitted_run_obj.run_id
+    active_run = None
+    # Note: there's a small chance we fail to report the run's status to the tracking server if
+    # we're interrupted before we reach the try block below
+    try:
+        active_run = tracking.MlflowClient().get_run(run_id) if run_id is not None else None
+        if submitted_run_obj.wait():
+            _logger.info("=== Run (ID '%s') succeeded ===", run_id)
+            _maybe_set_run_terminated(active_run, "FINISHED")
+        else:
+            _maybe_set_run_terminated(active_run, "FAILED")
+            raise ExecutionException(f"Run (ID '{run_id}') failed")
+    except KeyboardInterrupt:
+        _logger.error("=== Run (ID '%s') interrupted, cancelling run ===", run_id)
+        submitted_run_obj.cancel()
+        _maybe_set_run_terminated(active_run, "FAILED")
+        raise
+
+
+def _maybe_set_run_terminated(active_run, status):
+    """
+    If the passed-in active run is defined and still running (i.e. hasn't already been terminated
+    within user code), mark it as terminated with the passed-in status.
+    """
+    if active_run is None:
+        return
+    run_id = active_run.info.run_id
+    cur_status = tracking.MlflowClient().get_run(run_id).info.status
+    if RunStatus.is_terminated(cur_status):
+        return
+    tracking.MlflowClient().set_terminated(run_id, status)
+
+
+def _validate_execution_environment(project, backend):
+    if project.docker_env and backend == "databricks":
+        raise ExecutionException(
+            "Running docker-based projects on Databricks is not yet supported."
+        )
+
+
+def _parse_kubernetes_config(backend_config):
+    """
+    Creates build context tarfile containing Dockerfile and project code, returning path to tarfile
+    """
+    if not backend_config:
+        raise ExecutionException("Backend_config file not found.")
+    kube_config = backend_config.copy()
+    if "kube-job-template-path" not in backend_config.keys():
+        raise ExecutionException(
+            "'kube-job-template-path' attribute must be specified in backend_config."
+        )
+    kube_job_template = backend_config["kube-job-template-path"]
+    if os.path.exists(kube_job_template):
+        with open(kube_job_template) as job_template:
+            yaml_obj = yaml.safe_load(job_template.read())
+        kube_job_template = yaml_obj
+        kube_config["kube-job-template"] = kube_job_template
+    else:
+        raise ExecutionException(f"Could not find 'kube-job-template-path': {kube_job_template}")
+    if "kube-context" not in backend_config.keys():
+        _logger.debug(
+            "Could not find kube-context in backend_config."
+            " Using current context or in-cluster config."
+        )
+    if "repository-uri" not in backend_config.keys():
+        raise ExecutionException("Could not find 'repository-uri' in backend_config.")
+    return kube_config
+
+
+__all__ = ["run", "SubmittedRun"]

venv/lib/python3.9/site-packages/mlflow/projects/_project_spec.py

@@ -0,0 +1,363 @@
+"""Internal utilities for parsing MLproject YAML files."""
+
+import os
+
+import yaml
+
+from mlflow.exceptions import ExecutionException, MlflowException
+from mlflow.projects import env_type
+from mlflow.tracking import artifact_utils
+from mlflow.utils import data_utils
+from mlflow.utils.environment import _PYTHON_ENV_FILE_NAME
+from mlflow.utils.file_utils import get_local_path_or_none
+from mlflow.utils.string_utils import is_string_type, quote
+
+MLPROJECT_FILE_NAME = "mlproject"
+DEFAULT_CONDA_FILE_NAME = "conda.yaml"
+
+
+def _find_mlproject(directory):
+    filenames = os.listdir(directory)
+    for filename in filenames:
+        if filename.lower() == MLPROJECT_FILE_NAME:
+            return os.path.join(directory, filename)
+    return None
+
+
+def load_project(directory):
+    mlproject_path = _find_mlproject(directory)
+
+    # TODO: Validate structure of YAML loaded from the file
+    yaml_obj = {}
+    if mlproject_path is not None:
+        with open(mlproject_path) as mlproject_file:
+            yaml_obj = yaml.safe_load(mlproject_file)
+
+    # Validate the project config does't contain multiple environment fields
+    env_fields = set(yaml_obj.keys()).intersection(env_type.ALL)
+    if len(env_fields) > 1:
+        raise ExecutionException(
+            f"Project cannot contain multiple environment fields: {env_fields}"
+        )
+
+    project_name = yaml_obj.get("name")
+
+    # Parse entry points
+    entry_points = {}
+    for name, entry_point_yaml in yaml_obj.get("entry_points", {}).items():
+        parameters = entry_point_yaml.get("parameters", {})
+        command = entry_point_yaml.get("command")
+        entry_points[name] = EntryPoint(name, parameters, command)
+
+    databricks_spark_job_yaml = yaml_obj.get("databricks_spark_job")
+    if databricks_spark_job_yaml is not None:
+        python_file = databricks_spark_job_yaml.get("python_file")
+
+        if python_file is None and not entry_points:
+            raise MlflowException(
+                "Databricks Spark job requires either 'databricks_spark_job.python_file' "
+                "setting or 'entry_points' setting."
+            )
+        if python_file is not None and entry_points:
+            raise MlflowException(
+                "Databricks Spark job does not allow setting both "
+                "'databricks_spark_job.python_file' and 'entry_points'."
+            )
+
+        for entry_point in entry_points.values():
+            for param in entry_point.parameters.values():
+                if param.type == "path":
+                    raise MlflowException(
+                        "Databricks Spark job does not support entry point parameter of 'path' "
+                        f"type. '{param.name}' value type is invalid."
+                    )
+
+        if env_type.DOCKER in yaml_obj:
+            raise MlflowException(
+                "Databricks Spark job does not support setting docker environment."
+            )
+
+        if env_type.PYTHON in yaml_obj:
+            raise MlflowException(
+                "Databricks Spark job does not support setting python environment."
+            )
+
+        if env_type.CONDA in yaml_obj:
+            raise MlflowException(
+                "Databricks Spark job does not support setting conda environment."
+            )
+
+        databricks_spark_job_spec = DatabricksSparkJobSpec(
+            python_file=databricks_spark_job_yaml.get("python_file"),
+            parameters=databricks_spark_job_yaml.get("parameters", []),
+            python_libraries=databricks_spark_job_yaml.get("python_libraries", []),
+        )
+        return Project(
+            databricks_spark_job_spec=databricks_spark_job_spec,
+            name=project_name,
+            entry_points=entry_points,
+        )
+
+    # Validate config if docker_env parameter is present
+    docker_env = yaml_obj.get(env_type.DOCKER)
+    if docker_env:
+        if not docker_env.get("image"):
+            raise ExecutionException(
+                "Project configuration (MLproject file) was invalid: Docker "
+                "environment specified but no image attribute found."
+            )
+        if docker_env.get("volumes"):
+            if not (
+                isinstance(docker_env["volumes"], list)
+                and all(isinstance(i, str) for i in docker_env["volumes"])
+            ):
+                raise ExecutionException(
+                    "Project configuration (MLproject file) was invalid: "
+                    "Docker volumes must be a list of strings, "
+                    """e.g.: '["/path1/:/path1", "/path2/:/path2"])"""
+                )
+        if docker_env.get("environment"):
+            if not (
+                isinstance(docker_env["environment"], list)
+                and all(isinstance(i, (list, str)) for i in docker_env["environment"])
+            ):
+                raise ExecutionException(
+                    "Project configuration (MLproject file) was invalid: "
+                    "environment must be a list containing either strings (to copy environment "
+                    "variables from host system) or lists of string pairs (to define new "
+                    "environment variables)."
+                    """E.g.: '[["NEW_VAR", "new_value"], "VAR_TO_COPY_FROM_HOST"])"""
+                )
+        return Project(
+            env_type=env_type.DOCKER,
+            env_config_path=None,
+            entry_points=entry_points,
+            docker_env=docker_env,
+            name=project_name,
+        )
+
+    python_env = yaml_obj.get(env_type.PYTHON)
+    if python_env:
+        python_env_path = os.path.join(directory, python_env)
+        if not os.path.exists(python_env_path):
+            raise ExecutionException(
+                f"Project specified python_env file {python_env_path}, but no such file was found."
+            )
+        return Project(
+            env_type=env_type.PYTHON,
+            env_config_path=python_env_path,
+            entry_points=entry_points,
+            docker_env=None,
+            name=project_name,
+        )
+
+    conda_path = yaml_obj.get(env_type.CONDA)
+    if conda_path:
+        conda_env_path = os.path.join(directory, conda_path)
+        if not os.path.exists(conda_env_path):
+            raise ExecutionException(
+                f"Project specified conda environment file {conda_env_path}, but no such "
+                "file was found."
+            )
+        return Project(
+            env_type=env_type.CONDA,
+            env_config_path=conda_env_path,
+            entry_points=entry_points,
+            docker_env=None,
+            name=project_name,
+        )
+
+    default_python_env_path = os.path.join(directory, _PYTHON_ENV_FILE_NAME)
+    if os.path.exists(default_python_env_path):
+        return Project(
+            env_type=env_type.PYTHON,
+            env_config_path=default_python_env_path,
+            entry_points=entry_points,
+            docker_env=None,
+            name=project_name,
+        )
+
+    default_conda_path = os.path.join(directory, DEFAULT_CONDA_FILE_NAME)
+    if os.path.exists(default_conda_path):
+        return Project(
+            env_type=env_type.CONDA,
+            env_config_path=default_conda_path,
+            entry_points=entry_points,
+            docker_env=None,
+            name=project_name,
+        )
+
+    return Project(
+        env_type=env_type.PYTHON,
+        env_config_path=None,
+        entry_points=entry_points,
+        docker_env=None,
+        name=project_name,
+    )
+
+
+class Project:
+    """A project specification loaded from an MLproject file in the passed-in directory."""
+
+    def __init__(
+        self,
+        name,
+        env_type=None,
+        env_config_path=None,
+        entry_points=None,
+        docker_env=None,
+        databricks_spark_job_spec=None,
+    ):
+        self.env_type = env_type
+        self.env_config_path = env_config_path
+        self._entry_points = entry_points
+        self.docker_env = docker_env
+        self.name = name
+        self.databricks_spark_job_spec = databricks_spark_job_spec
+
+    def get_entry_point(self, entry_point):
+        if self.databricks_spark_job_spec:
+            if self.databricks_spark_job_spec.python_file is not None:
+                # If Databricks Spark job is configured with python_file field,
+                # it does not need to configure entry_point section
+                # and the 'entry_point' param in 'mlflow run' command is ignored
+                return None
+
+            if self._entry_points is None or entry_point not in self._entry_points:
+                raise MlflowException(
+                    f"The entry point '{entry_point}' is not defined in the Databricks spark job "
+                    f"MLproject file."
+                )
+
+        if entry_point in self._entry_points:
+            return self._entry_points[entry_point]
+        _, file_extension = os.path.splitext(entry_point)
+        ext_to_cmd = {".py": "python", ".sh": os.environ.get("SHELL", "bash")}
+        if file_extension in ext_to_cmd:
+            command = f"{ext_to_cmd[file_extension]} {quote(entry_point)}"
+            if not is_string_type(command):
+                command = command.encode("utf-8")
+            return EntryPoint(name=entry_point, parameters={}, command=command)
+        elif file_extension == ".R":
+            command = f"Rscript -e \"mlflow::mlflow_source('{quote(entry_point)}')\" --args"
+            return EntryPoint(name=entry_point, parameters={}, command=command)
+        raise ExecutionException(
+            "Could not find {0} among entry points {1} or interpret {0} as a "
+            "runnable script. Supported script file extensions: "
+            "{2}".format(entry_point, list(self._entry_points.keys()), list(ext_to_cmd.keys()))
+        )
+
+
+class EntryPoint:
+    """An entry point in an MLproject specification."""
+
+    def __init__(self, name, parameters, command):
+        self.name = name
+        self.parameters = {k: Parameter(k, v) for (k, v) in parameters.items()}
+        self.command = command
+
+    def _validate_parameters(self, user_parameters):
+        missing_params = []
+        for name in self.parameters:
+            if name not in user_parameters and self.parameters[name].default is None:
+                missing_params.append(name)
+        if missing_params:
+            raise ExecutionException(
+                "No value given for missing parameters: {}".format(
+                    ", ".join([f"'{name}'" for name in missing_params])
+                )
+            )
+
+    def compute_parameters(self, user_parameters, storage_dir):
+        """
+        Given a dict mapping user-specified param names to values, computes parameters to
+        substitute into the command for this entry point. Returns a tuple (params, extra_params)
+        where `params` contains key-value pairs for parameters specified in the entry point
+        definition, and `extra_params` contains key-value pairs for additional parameters passed
+        by the user.
+
+        Note that resolving parameter values can be a heavy operation, e.g. if a remote URI is
+        passed for a parameter of type `path`, we download the URI to a local path within
+        `storage_dir` and substitute in the local path as the parameter value.
+
+        If `storage_dir` is `None`, report path will be return as parameter.
+        """
+        if user_parameters is None:
+            user_parameters = {}
+        # Validate params before attempting to resolve parameter values
+        self._validate_parameters(user_parameters)
+        final_params = {}
+        extra_params = {}
+
+        parameter_keys = list(self.parameters.keys())
+        for key in parameter_keys:
+            param_obj = self.parameters[key]
+            key_position = parameter_keys.index(key)
+            value = user_parameters[key] if key in user_parameters else self.parameters[key].default
+            final_params[key] = param_obj.compute_value(value, storage_dir, key_position)
+        for key in user_parameters:
+            if key not in final_params:
+                extra_params[key] = user_parameters[key]
+        return self._sanitize_param_dict(final_params), self._sanitize_param_dict(extra_params)
+
+    def compute_command(self, user_parameters, storage_dir):
+        params, extra_params = self.compute_parameters(user_parameters, storage_dir)
+        command_with_params = self.command.format(**params)
+        command_arr = [command_with_params]
+        command_arr.extend([f"--{key} {value}" for key, value in extra_params.items()])
+        return " ".join(command_arr)
+
+    @staticmethod
+    def _sanitize_param_dict(param_dict):
+        return {str(key): quote(str(value)) for key, value in param_dict.items()}
+
+
+class Parameter:
+    """A parameter in an MLproject entry point."""
+
+    def __init__(self, name, yaml_obj):
+        self.name = name
+        if is_string_type(yaml_obj):
+            self.type = yaml_obj
+            self.default = None
+        else:
+            self.type = yaml_obj.get("type", "string")
+            self.default = yaml_obj.get("default")
+
+    def _compute_uri_value(self, user_param_value):
+        if not data_utils.is_uri(user_param_value):
+            raise ExecutionException(
+                f"Expected URI for parameter {self.name} but got {user_param_value}"
+            )
+        return user_param_value
+
+    def _compute_path_value(self, user_param_value, storage_dir, key_position):
+        local_path = get_local_path_or_none(user_param_value)
+        if local_path:
+            if not os.path.exists(local_path):
+                raise ExecutionException(
+                    f"Got value {user_param_value} for parameter {self.name}, but no such file or "
+                    "directory was found."
+                )
+            return os.path.abspath(local_path)
+        target_sub_dir = f"param_{key_position}"
+        download_dir = os.path.join(storage_dir, target_sub_dir)
+        os.mkdir(download_dir)
+        return artifact_utils._download_artifact_from_uri(
+            artifact_uri=user_param_value, output_path=download_dir
+        )
+
+    def compute_value(self, param_value, storage_dir, key_position):
+        if storage_dir and self.type == "path":
+            return self._compute_path_value(param_value, storage_dir, key_position)
+        elif self.type == "uri":
+            return self._compute_uri_value(param_value)
+        else:
+            return param_value
+
+
+class DatabricksSparkJobSpec:
+    def __init__(self, python_file, parameters, python_libraries):
+        self.python_file = python_file
+        self.parameters = parameters
+        self.python_libraries = python_libraries

venv/lib/python3.9/site-packages/mlflow/projects/backend/__init__.py

@@ -0,0 +1,8 @@
+"""
+This module defines developer APIs for defining pluggable execution backends
+for MLflow projects. See `MLflow Plugins <../../plugins.html>`_ for more information.
+"""
+
+from mlflow.projects.backend.abstract_backend import AbstractBackend
+
+__all__ = ["AbstractBackend"]

venv/lib/python3.9/site-packages/mlflow/projects/backend/abstract_backend.py

@@ -0,0 +1,50 @@
+from abc import ABCMeta, abstractmethod
+
+from mlflow.utils.annotations import developer_stable
+
+
+@developer_stable
+class AbstractBackend:
+    """
+    Abstract plugin class defining the interface needed to execute MLflow projects. You can define
+    subclasses of ``AbstractBackend`` and expose them as third-party plugins to enable running
+    MLflow projects against custom execution backends (e.g. to run projects against your team's
+    in-house cluster or job scheduler). See `MLflow Plugins <../../plugins.html>`_ for more
+    information.
+    """
+
+    __metaclass__ = ABCMeta
+
+    @abstractmethod
+    def run(
+        self,
+        project_uri,
+        entry_point,
+        params,
+        version,
+        backend_config,
+        tracking_uri,
+        experiment_id,
+    ):
+        """
+        Submit an entrypoint. It must return a SubmittedRun object to track the execution
+
+        Args:
+            project_uri: URI of the project to execute, e.g. a local filesystem path
+                or a Git repository URI like https://github.com/mlflow/mlflow-example
+            entry_point: Entry point to run within the project.
+            params: Dict of parameters to pass to the entry point
+            version: For git-based projects, either a commit hash or a branch name.
+            backend_config: A dictionary, or a path to a JSON file (must end in '.json'), which
+                will be passed as config to the backend. The exact content which
+                should be provided is different for each execution backend and is
+                documented at https://www.mlflow.org/docs/latest/projects.html.
+            tracking_uri: URI of tracking server against which to log run information related
+                to project execution.
+            experiment_id: ID of experiment under which to launch the run.
+
+        Returns:
+            A :py:class:`mlflow.projects.SubmittedRun`. This function is expected to run
+            the project asynchronously, i.e. it should trigger project execution and then
+            immediately return a `SubmittedRun` to track execution status.
+        """

venv/lib/python3.9/site-packages/mlflow/projects/backend/loader.py

@@ -0,0 +1,35 @@
+import logging
+
+from mlflow.projects.backend.local import LocalBackend
+from mlflow.utils.plugins import get_entry_points
+
+ENTRYPOINT_GROUP_NAME = "mlflow.project_backend"
+
+_logger = logging.getLogger(__name__)
+
+
+# Statically register backend defined in mlflow
+MLFLOW_BACKENDS = {
+    "local": LocalBackend,
+}
+
+
+def load_backend(backend_name):
+    # Static backends
+    if backend_name in MLFLOW_BACKENDS:
+        return MLFLOW_BACKENDS[backend_name]()
+
+    # backends from plugin
+    entrypoints = get_entry_points(ENTRYPOINT_GROUP_NAME)
+    if entrypoint := next((e for e in entrypoints if e.name == backend_name), None):
+        builder = entrypoint.load()
+        return builder()
+
+    # TODO Should be a error when all backends are migrated here
+    _logger.warning(
+        "Backend '%s' is not available. Available plugins are %s",
+        backend_name,
+        [*entrypoints, *MLFLOW_BACKENDS.keys()],
+    )
+
+    return None

venv/lib/python3.9/site-packages/mlflow/projects/backend/local.py

@@ -0,0 +1,427 @@
+import logging
+import os
+import platform
+import posixpath
+import subprocess
+import sys
+from pathlib import Path
+
+import mlflow
+from mlflow import tracking
+from mlflow.environment_variables import (
+    MLFLOW_KERBEROS_TICKET_CACHE,
+    MLFLOW_KERBEROS_USER,
+    MLFLOW_PYARROW_EXTRA_CONF,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.projects import env_type
+from mlflow.projects.backend.abstract_backend import AbstractBackend
+from mlflow.projects.submitted_run import LocalSubmittedRun
+from mlflow.projects.utils import (
+    MLFLOW_DOCKER_WORKDIR_PATH,
+    MLFLOW_LOCAL_BACKEND_RUN_ID_CONFIG,
+    PROJECT_BUILD_IMAGE,
+    PROJECT_DOCKER_ARGS,
+    PROJECT_DOCKER_AUTH,
+    PROJECT_ENV_MANAGER,
+    PROJECT_STORAGE_DIR,
+    PROJECT_SYNCHRONOUS,
+    fetch_and_validate_project,
+    get_entry_point_command,
+    get_or_create_run,
+    get_run_env_vars,
+    load_project,
+)
+from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
+from mlflow.store.artifact.azure_blob_artifact_repo import AzureBlobArtifactRepository
+from mlflow.store.artifact.gcs_artifact_repo import GCSArtifactRepository
+from mlflow.store.artifact.hdfs_artifact_repo import HdfsArtifactRepository
+from mlflow.store.artifact.local_artifact_repo import LocalArtifactRepository
+from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
+from mlflow.utils import env_manager as _EnvManager
+from mlflow.utils.conda import get_or_create_conda_env
+from mlflow.utils.databricks_utils import get_databricks_env_vars, is_in_databricks_runtime
+from mlflow.utils.environment import _PythonEnv
+from mlflow.utils.file_utils import get_or_create_nfs_tmp_dir
+from mlflow.utils.mlflow_tags import MLFLOW_PROJECT_ENV
+from mlflow.utils.os import is_windows
+from mlflow.utils.virtualenv import (
+    _PYENV_ROOT_DIR,
+    _VIRTUALENV_ENVS_DIR,
+    _create_virtualenv,
+    _get_mlflow_virtualenv_root,
+    _get_virtualenv_extra_env_vars,
+    _get_virtualenv_name,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+def _env_type_to_env_manager(env_typ):
+    if env_typ == env_type.CONDA:
+        return _EnvManager.CONDA
+    elif env_typ == env_type.PYTHON:
+        return _EnvManager.VIRTUALENV
+    elif env_typ == env_type.DOCKER:
+        return _EnvManager.LOCAL
+
+
+class LocalBackend(AbstractBackend):
+    def run(
+        self,
+        project_uri,
+        entry_point,
+        params,
+        version,
+        backend_config,
+        tracking_uri,
+        experiment_id,
+    ):
+        work_dir = fetch_and_validate_project(project_uri, version, entry_point, params)
+        project = load_project(work_dir)
+        if MLFLOW_LOCAL_BACKEND_RUN_ID_CONFIG in backend_config:
+            run_id = backend_config[MLFLOW_LOCAL_BACKEND_RUN_ID_CONFIG]
+        else:
+            run_id = None
+        active_run = get_or_create_run(
+            run_id, project_uri, experiment_id, work_dir, version, entry_point, params
+        )
+        command_args = []
+        command_separator = " "
+        env_manager = backend_config[PROJECT_ENV_MANAGER]
+        synchronous = backend_config[PROJECT_SYNCHRONOUS]
+        docker_args = backend_config[PROJECT_DOCKER_ARGS]
+        storage_dir = backend_config[PROJECT_STORAGE_DIR]
+        build_image = backend_config[PROJECT_BUILD_IMAGE]
+        docker_auth = backend_config[PROJECT_DOCKER_AUTH]
+        # Select an appropriate env manager for the project env type
+        if env_manager is None:
+            env_manager = _env_type_to_env_manager(project.env_type)
+        else:
+            if project.env_type == env_type.PYTHON and env_manager == _EnvManager.CONDA:
+                raise MlflowException.invalid_parameter_value(
+                    "python_env project cannot be executed using conda. Set `--env-manager` to "
+                    "'virtualenv' or 'local' to execute this project."
+                )
+
+        # If a docker_env attribute is defined in MLproject then it takes precedence over conda yaml
+        # environments, so the project will be executed inside a docker container.
+        if project.docker_env:
+            from mlflow.projects.docker import (
+                build_docker_image,
+                validate_docker_env,
+                validate_docker_installation,
+            )
+
+            tracking.MlflowClient().set_tag(active_run.info.run_id, MLFLOW_PROJECT_ENV, "docker")
+            validate_docker_env(project)
+            validate_docker_installation()
+            image = build_docker_image(
+                work_dir=work_dir,
+                repository_uri=project.name,
+                base_image=project.docker_env.get("image"),
+                run_id=active_run.info.run_id,
+                build_image=build_image,
+                docker_auth=docker_auth,
+            )
+            command_args += _get_docker_command(
+                image=image,
+                active_run=active_run,
+                docker_args=docker_args,
+                volumes=project.docker_env.get("volumes"),
+                user_env_vars=project.docker_env.get("environment"),
+            )
+        # Synchronously create a conda environment (even though this may take some time)
+        # to avoid failures due to multiple concurrent attempts to create the same conda env.
+        elif env_manager == _EnvManager.VIRTUALENV:
+            tracking.MlflowClient().set_tag(
+                active_run.info.run_id, MLFLOW_PROJECT_ENV, "virtualenv"
+            )
+            command_separator = " && "
+            if project.env_type == env_type.CONDA:
+                python_env = _PythonEnv.from_conda_yaml(project.env_config_path)
+            else:
+                python_env = (
+                    _PythonEnv.from_yaml(project.env_config_path)
+                    if project.env_config_path
+                    else _PythonEnv()
+                )
+
+            if is_in_databricks_runtime():
+                nfs_tmp_dir = get_or_create_nfs_tmp_dir()
+                env_root = Path(nfs_tmp_dir) / "envs"
+                pyenv_root_dir = str(env_root / _PYENV_ROOT_DIR)
+                virtualenv_root = env_root / _VIRTUALENV_ENVS_DIR
+                env_vars = _get_virtualenv_extra_env_vars(str(env_root))
+            else:
+                pyenv_root_dir = None
+                virtualenv_root = Path(_get_mlflow_virtualenv_root())
+                env_vars = None
+            work_dir_path = Path(work_dir)
+            env_name = _get_virtualenv_name(python_env, work_dir_path)
+            env_dir = virtualenv_root / env_name
+            activate_cmd = _create_virtualenv(
+                local_model_path=work_dir_path,
+                python_env=python_env,
+                env_dir=env_dir,
+                pyenv_root_dir=pyenv_root_dir,
+                env_manager=env_manager,
+                extra_env=env_vars,
+            )
+            command_args += [activate_cmd]
+        elif env_manager == _EnvManager.CONDA:
+            tracking.MlflowClient().set_tag(active_run.info.run_id, MLFLOW_PROJECT_ENV, "conda")
+            command_separator = " && "
+            conda_env = get_or_create_conda_env(project.env_config_path)
+            command_args += conda_env.get_activate_command()
+
+        # In synchronous mode, run the entry point command in a blocking fashion, sending status
+        # updates to the tracking server when finished. Note that the run state may not be
+        # persisted to the tracking server if interrupted
+        if synchronous:
+            command_args += get_entry_point_command(project, entry_point, params, storage_dir)
+            command_str = command_separator.join(command_args)
+            return _run_entry_point(
+                command_str, work_dir, experiment_id, run_id=active_run.info.run_id
+            )
+        # Otherwise, invoke `mlflow run` in a subprocess
+        return _invoke_mlflow_run_subprocess(
+            work_dir=work_dir,
+            entry_point=entry_point,
+            parameters=params,
+            experiment_id=experiment_id,
+            env_manager=env_manager,
+            docker_args=docker_args,
+            storage_dir=storage_dir,
+            run_id=active_run.info.run_id,
+        )
+
+
+def _invoke_mlflow_run_subprocess(
+    work_dir, entry_point, parameters, experiment_id, env_manager, docker_args, storage_dir, run_id
+):
+    """
+    Run an MLflow project asynchronously by invoking ``mlflow run`` in a subprocess, returning
+    a SubmittedRun that can be used to query run status.
+    """
+    _logger.info("=== Asynchronously launching MLflow run with ID %s ===", run_id)
+    mlflow_run_arr = _build_mlflow_run_cmd(
+        uri=work_dir,
+        entry_point=entry_point,
+        docker_args=docker_args,
+        storage_dir=storage_dir,
+        env_manager=env_manager,
+        run_id=run_id,
+        parameters=parameters,
+    )
+    env_vars = get_run_env_vars(run_id, experiment_id)
+    env_vars.update(get_databricks_env_vars(mlflow.get_tracking_uri()))
+    mlflow_run_subprocess = _run_mlflow_run_cmd(mlflow_run_arr, env_vars)
+    return LocalSubmittedRun(run_id, mlflow_run_subprocess)
+
+
+def _build_mlflow_run_cmd(
+    uri, entry_point, docker_args, storage_dir, env_manager, run_id, parameters
+):
+    """
+    Build and return an array containing an ``mlflow run`` command that can be invoked to locally
+    run the project at the specified URI.
+    """
+    mlflow_run_arr = ["mlflow", "run", uri, "-e", entry_point, "--run-id", run_id]
+    if docker_args is not None:
+        for key, value in docker_args.items():
+            args = key if isinstance(value, bool) else f"{key}={value}"
+            mlflow_run_arr.extend(["--docker-args", args])
+    if storage_dir is not None:
+        mlflow_run_arr.extend(["--storage-dir", storage_dir])
+    mlflow_run_arr.extend(["--env-manager", env_manager])
+    for key, value in parameters.items():
+        mlflow_run_arr.extend(["-P", f"{key}={value}"])
+    return mlflow_run_arr
+
+
+def _run_mlflow_run_cmd(mlflow_run_arr, env_map):
+    """
+    Invoke ``mlflow run`` in a subprocess, which in turn runs the entry point in a child process.
+    Returns a handle to the subprocess. Popen launched to invoke ``mlflow run``.
+    """
+    final_env = os.environ.copy()
+    final_env.update(env_map)
+    # Launch `mlflow run` command as the leader of its own process group so that we can do a
+    # best-effort cleanup of all its descendant processes if needed
+    if sys.platform == "win32":
+        return subprocess.Popen(
+            mlflow_run_arr,
+            env=final_env,
+            text=True,
+            creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
+        )
+    else:
+        return subprocess.Popen(mlflow_run_arr, env=final_env, text=True, preexec_fn=os.setsid)
+
+
+def _run_entry_point(command, work_dir, experiment_id, run_id):  # noqa: D417
+    """
+    Run an entry point command in a subprocess, returning a SubmittedRun that can be used to
+    query the run's status.
+
+    Args:
+        command: Entry point command to run
+        work_dir: Working directory in which to run the command
+        run_id: MLflow run ID associated with the entry point execution.
+    """
+    env = os.environ.copy()
+    env.update(get_run_env_vars(run_id, experiment_id))
+    env.update(get_databricks_env_vars(tracking_uri=mlflow.get_tracking_uri()))
+    _logger.info("=== Running command '%s' in run with ID '%s' === ", command, run_id)
+    # in case os name is not 'nt', we are not running on windows. It introduces
+    # bash command otherwise.
+    if not is_windows():
+        process = subprocess.Popen(["bash", "-c", command], close_fds=True, cwd=work_dir, env=env)
+    else:
+        # process = subprocess.Popen(command, close_fds=True, cwd=work_dir, env=env)
+        process = subprocess.Popen(["cmd", "/c", command], close_fds=True, cwd=work_dir, env=env)
+    return LocalSubmittedRun(run_id, process)
+
+
+def _get_docker_command(image, active_run, docker_args=None, volumes=None, user_env_vars=None):
+    from mlflow.projects.docker import get_docker_tracking_cmd_and_envs
+
+    docker_path = "docker"
+    cmd = [docker_path, "run", "--rm"]
+
+    if docker_args:
+        for name, value in docker_args.items():
+            # Passed just the name as boolean flag
+            if isinstance(value, bool) and value:
+                if len(name) == 1:
+                    cmd += ["-" + name]
+                else:
+                    cmd += ["--" + name]
+            else:
+                # Passed name=value
+                if len(name) == 1:
+                    cmd += ["-" + name, value]
+                else:
+                    cmd += ["--" + name, value]
+
+    env_vars = get_run_env_vars(
+        run_id=active_run.info.run_id, experiment_id=active_run.info.experiment_id
+    )
+    tracking_uri = tracking.get_tracking_uri()
+    tracking_cmds, tracking_envs = get_docker_tracking_cmd_and_envs(tracking_uri)
+    artifact_cmds, artifact_envs = _get_docker_artifact_storage_cmd_and_envs(
+        active_run.info.artifact_uri
+    )
+
+    cmd += tracking_cmds + artifact_cmds
+    env_vars.update(tracking_envs)
+    env_vars.update(artifact_envs)
+    if user_env_vars is not None:
+        for user_entry in user_env_vars:
+            if isinstance(user_entry, list):
+                # User has defined a new environment variable for the docker environment
+                env_vars[user_entry[0]] = user_entry[1]
+            else:
+                # User wants to copy an environment variable from system environment
+                system_var = os.environ.get(user_entry)
+                if system_var is None:
+                    raise MlflowException(
+                        "This project expects the {} environment variables to "
+                        "be set on the machine running the project, but {} was "
+                        "not set. Please ensure all expected environment variables "
+                        "are set".format(", ".join(user_env_vars), user_entry)
+                    )
+                env_vars[user_entry] = system_var
+
+    if volumes is not None:
+        for v in volumes:
+            cmd += ["-v", v]
+
+    for key, value in env_vars.items():
+        cmd += ["-e", f"{key}={value}"]
+    cmd += [image.tags[0]]
+    return cmd
+
+
+def _get_local_artifact_cmd_and_envs(artifact_repo):
+    artifact_dir = artifact_repo.artifact_dir
+    container_path = artifact_dir
+    if not os.path.isabs(container_path):
+        container_path = os.path.join(MLFLOW_DOCKER_WORKDIR_PATH, container_path)
+        container_path = os.path.normpath(container_path)
+    abs_artifact_dir = os.path.abspath(artifact_dir)
+    return ["-v", f"{abs_artifact_dir}:{container_path}"], {}
+
+
+def _get_s3_artifact_cmd_and_envs(artifact_repo):
+    if platform.system() == "Windows":
+        win_user_dir = os.environ["USERPROFILE"]
+        aws_path = os.path.join(win_user_dir, ".aws")
+    else:
+        aws_path = posixpath.expanduser("~/.aws")
+
+    volumes = []
+    if posixpath.exists(aws_path):
+        volumes = ["-v", "{}:{}".format(str(aws_path), "/.aws")]
+    envs = {
+        "AWS_SECRET_ACCESS_KEY": os.environ.get("AWS_SECRET_ACCESS_KEY"),
+        "AWS_ACCESS_KEY_ID": os.environ.get("AWS_ACCESS_KEY_ID"),
+        "MLFLOW_S3_ENDPOINT_URL": os.environ.get("MLFLOW_S3_ENDPOINT_URL"),
+        "MLFLOW_S3_IGNORE_TLS": os.environ.get("MLFLOW_S3_IGNORE_TLS"),
+    }
+    envs = {k: v for k, v in envs.items() if v is not None}
+    return volumes, envs
+
+
+def _get_azure_blob_artifact_cmd_and_envs(artifact_repo):
+    envs = {
+        "AZURE_STORAGE_CONNECTION_STRING": os.environ.get("AZURE_STORAGE_CONNECTION_STRING"),
+        "AZURE_STORAGE_ACCESS_KEY": os.environ.get("AZURE_STORAGE_ACCESS_KEY"),
+    }
+    envs = {k: v for k, v in envs.items() if v is not None}
+    return [], envs
+
+
+def _get_gcs_artifact_cmd_and_envs(artifact_repo):
+    cmds = []
+    envs = {}
+
+    if "GOOGLE_APPLICATION_CREDENTIALS" in os.environ:
+        credentials_path = os.environ["GOOGLE_APPLICATION_CREDENTIALS"]
+        cmds = ["-v", f"{credentials_path}:/.gcs"]
+        envs["GOOGLE_APPLICATION_CREDENTIALS"] = "/.gcs"
+    return cmds, envs
+
+
+def _get_hdfs_artifact_cmd_and_envs(artifact_repo):
+    cmds = []
+    envs = {
+        "MLFLOW_KERBEROS_TICKET_CACHE": MLFLOW_KERBEROS_TICKET_CACHE.get(),
+        "MLFLOW_KERBEROS_USER": MLFLOW_KERBEROS_USER.get(),
+        "MLFLOW_PYARROW_EXTRA_CONF": MLFLOW_PYARROW_EXTRA_CONF.get(),
+    }
+    envs = {k: v for k, v in envs.items() if v is not None}
+
+    if "MLFLOW_KERBEROS_TICKET_CACHE" in envs:
+        ticket_cache = envs["MLFLOW_KERBEROS_TICKET_CACHE"]
+        cmds = ["-v", f"{ticket_cache}:{ticket_cache}"]
+    return cmds, envs
+
+
+_artifact_storages = {
+    LocalArtifactRepository: _get_local_artifact_cmd_and_envs,
+    S3ArtifactRepository: _get_s3_artifact_cmd_and_envs,
+    AzureBlobArtifactRepository: _get_azure_blob_artifact_cmd_and_envs,
+    HdfsArtifactRepository: _get_hdfs_artifact_cmd_and_envs,
+    GCSArtifactRepository: _get_gcs_artifact_cmd_and_envs,
+}
+
+
+def _get_docker_artifact_storage_cmd_and_envs(artifact_uri):
+    artifact_repo = get_artifact_repository(artifact_uri)
+    _get_cmd_and_envs = _artifact_storages.get(type(artifact_repo))
+    if _get_cmd_and_envs is not None:
+        return _get_cmd_and_envs(artifact_repo)
+    else:
+        return [], {}

venv/lib/python3.9/site-packages/mlflow/projects/databricks.py

@@ -0,0 +1,611 @@
+import hashlib
+import json
+import logging
+import os
+import posixpath
+import re
+import tempfile
+import textwrap
+import time
+import uuid
+from pathlib import Path
+from shlex import quote
+
+from mlflow import tracking
+from mlflow.entities import RunStatus
+from mlflow.environment_variables import MLFLOW_EXPERIMENT_ID, MLFLOW_RUN_ID, MLFLOW_TRACKING_URI
+from mlflow.exceptions import ExecutionException, MlflowException
+from mlflow.projects.submitted_run import SubmittedRun
+from mlflow.projects.utils import MLFLOW_LOCAL_BACKEND_RUN_ID_CONFIG
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.utils import databricks_utils, file_utils, rest_utils
+from mlflow.utils.mlflow_tags import (
+    MLFLOW_DATABRICKS_RUN_URL,
+    MLFLOW_DATABRICKS_SHELL_JOB_ID,
+    MLFLOW_DATABRICKS_SHELL_JOB_RUN_ID,
+    MLFLOW_DATABRICKS_WEBAPP_URL,
+)
+from mlflow.utils.uri import is_databricks_uri, is_http_uri
+from mlflow.version import VERSION, is_release_version
+
+# Base directory within driver container for storing files related to MLflow
+DB_CONTAINER_BASE = "/databricks/mlflow"
+# Base directory within driver container for storing project archives
+DB_TARFILE_BASE = posixpath.join(DB_CONTAINER_BASE, "project-tars")
+# Base directory directory within driver container for storing extracted project directories
+DB_PROJECTS_BASE = posixpath.join(DB_CONTAINER_BASE, "projects")
+# Name to use for project directory when archiving it for upload to DBFS; the TAR will contain
+# a single directory with this name
+DB_TARFILE_ARCHIVE_NAME = "mlflow-project"
+# Base directory within DBFS for storing code for project runs for experiments
+DBFS_EXPERIMENT_DIR_BASE = "mlflow-experiments"
+
+
+_logger = logging.getLogger(__name__)
+
+_MLFLOW_GIT_URI_REGEX = re.compile(r"^git\+https://github.com/[\w-]+/mlflow")
+
+
+def _is_mlflow_git_uri(s):
+    return bool(_MLFLOW_GIT_URI_REGEX.match(s))
+
+
+def _contains_mlflow_git_uri(libraries):
+    for lib in libraries:
+        package = lib.get("pypi", {}).get("package")
+        if package and _is_mlflow_git_uri(package):
+            return True
+    return False
+
+
+def before_run_validations(tracking_uri, backend_config):
+    """Validations to perform before running a project on Databricks."""
+    if backend_config is None:
+        raise ExecutionException(
+            "Backend spec must be provided when launching MLflow project runs on Databricks."
+        )
+    elif "existing_cluster_id" in backend_config:
+        raise MlflowException(
+            message=(
+                "MLflow Project runs on Databricks must provide a *new cluster* specification."
+                " Project execution against existing clusters is not currently supported. For more"
+                " information, see https://mlflow.org/docs/latest/projects.html"
+                "#run-an-mlflow-project-on-databricks"
+            ),
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+    if not is_databricks_uri(tracking_uri) and not is_http_uri(tracking_uri):
+        raise ExecutionException(
+            "When running on Databricks, the MLflow tracking URI must be of the form "
+            "'databricks' or 'databricks://profile', or a remote HTTP URI accessible to both the "
+            "current client and code running on Databricks. Got local tracking URI "
+            f"{tracking_uri}. Please specify a valid tracking URI via mlflow.set_tracking_uri or "
+            "by setting the MLFLOW_TRACKING_URI environment variable."
+        )
+
+
+class DatabricksJobRunner:
+    """
+    Helper class for running an MLflow project as a Databricks Job.
+
+    Args:
+        databricks_profile: Optional Databricks CLI profile to use to fetch hostname &
+           authentication information when making Databricks API requests.
+    """
+
+    def __init__(self, databricks_profile_uri):
+        self.databricks_profile_uri = databricks_profile_uri
+
+    def _databricks_api_request(self, endpoint, method, **kwargs):
+        host_creds = databricks_utils.get_databricks_host_creds(self.databricks_profile_uri)
+        return rest_utils.http_request_safe(
+            host_creds=host_creds, endpoint=endpoint, method=method, **kwargs
+        )
+
+    def _jobs_runs_submit(self, req_body):
+        response = self._databricks_api_request(
+            endpoint="/api/2.0/jobs/runs/submit", method="POST", json=req_body
+        )
+        return json.loads(response.text)
+
+    def _upload_to_dbfs(self, src_path, dbfs_fuse_uri):
+        """
+        Upload the file at `src_path` to the specified DBFS URI within the Databricks workspace
+        corresponding to the default Databricks CLI profile.
+        """
+        _logger.info("=== Uploading project to DBFS path %s ===", dbfs_fuse_uri)
+        http_endpoint = dbfs_fuse_uri
+        with open(src_path, "rb") as f:
+            try:
+                self._databricks_api_request(endpoint=http_endpoint, method="POST", data=f)
+            except MlflowException as e:
+                if "Error 409" in e.message and "File already exists" in e.message:
+                    _logger.info("=== Did not overwrite existing DBFS path %s ===", dbfs_fuse_uri)
+                else:
+                    raise e
+
+    def _dbfs_path_exists(self, dbfs_path):
+        """
+        Return True if the passed-in path exists in DBFS for the workspace corresponding to the
+        default Databricks CLI profile. The path is expected to be a relative path to the DBFS root
+        directory, e.g. 'path/to/file'.
+        """
+        host_creds = databricks_utils.get_databricks_host_creds(self.databricks_profile_uri)
+        response = rest_utils.http_request(
+            host_creds=host_creds,
+            endpoint="/api/2.0/dbfs/get-status",
+            method="GET",
+            json={"path": f"/{dbfs_path}"},
+        )
+        try:
+            json_response_obj = json.loads(response.text)
+        except Exception:
+            raise MlflowException(
+                f"API request to check existence of file at DBFS path {dbfs_path} failed with "
+                f"status code {response.status_code}. Response body: {response.text}"
+            )
+        # If request fails with a RESOURCE_DOES_NOT_EXIST error, the file does not exist on DBFS
+        error_code_field = "error_code"
+        if error_code_field in json_response_obj:
+            if json_response_obj[error_code_field] == "RESOURCE_DOES_NOT_EXIST":
+                return False
+            raise ExecutionException(
+                f"Got unexpected error response when checking whether file {dbfs_path} "
+                f"exists in DBFS: {json_response_obj}"
+            )
+        return True
+
+    def _upload_project_to_dbfs(self, project_dir, experiment_id):  # noqa: D417
+        """
+        Tars a project directory into an archive in a temp dir and uploads it to DBFS, returning
+        the HDFS-style URI of the tarball in DBFS (e.g. dbfs:/path/to/tar).
+
+        Args:
+            project_dir: Path to a directory containing an MLflow project to upload to DBFS (e.g.
+                a directory containing an MLproject file).
+        """
+        with tempfile.TemporaryDirectory() as temp_tarfile_dir:
+            temp_tar_filename = os.path.join(temp_tarfile_dir, "project.tar.gz")
+
+            def custom_filter(x):
+                return None if os.path.basename(x.name) == "mlruns" else x
+
+            directory_size = file_utils._get_local_project_dir_size(project_dir)
+            _logger.info(
+                f"=== Creating tarball from {project_dir} in temp directory {temp_tarfile_dir} ==="
+            )
+            _logger.info(f"=== Total file size to compress: {directory_size} KB ===")
+            file_utils.make_tarfile(
+                temp_tar_filename, project_dir, DB_TARFILE_ARCHIVE_NAME, custom_filter=custom_filter
+            )
+            with open(temp_tar_filename, "rb") as tarred_project:
+                tarfile_hash = hashlib.sha256(tarred_project.read()).hexdigest()
+            # TODO: Get subdirectory for experiment from the tracking server
+            dbfs_path = posixpath.join(
+                DBFS_EXPERIMENT_DIR_BASE,
+                str(experiment_id),
+                "projects-code",
+                f"{tarfile_hash}.tar.gz",
+            )
+            tar_size = file_utils._get_local_file_size(temp_tar_filename)
+            dbfs_fuse_uri = posixpath.join("/dbfs", dbfs_path)
+            if not self._dbfs_path_exists(dbfs_path):
+                _logger.info(
+                    f"=== Uploading project tarball (size: {tar_size} KB) to {dbfs_fuse_uri} ==="
+                )
+                self._upload_to_dbfs(temp_tar_filename, dbfs_fuse_uri)
+                _logger.info("=== Finished uploading project to %s ===", dbfs_fuse_uri)
+            else:
+                _logger.info("=== Project already exists in DBFS ===")
+        return dbfs_fuse_uri
+
+    def _run_shell_command_job(self, project_uri, command, env_vars, cluster_spec):
+        """
+        Run the specified shell command on a Databricks cluster.
+
+        Args:
+            project_uri: URI of the project from which the shell command originates.
+            command: Shell command to run.
+            env_vars: Environment variables to set in the process running ``command``.
+            cluster_spec: Dictionary containing a `Databricks cluster specification
+                <https://docs.databricks.com/dev-tools/api/latest/jobs.html#clusterspec>`_
+                or a `Databricks new cluster specification
+                <https://docs.databricks.com/dev-tools/api/latest/jobs.html#jobsclusterspecnewcluster>`_
+                to use when launching a run. If you specify libraries, this function
+                will add MLflow to the library list. This function does not support
+                installation of conda environment libraries on the workers.
+
+        Returns:
+            ID of the Databricks job run. Can be used to query the run's status via the
+            Databricks `Runs Get <https://docs.databricks.com/api/latest/jobs.html#runs-get>`_ API.
+        """
+        if is_release_version():
+            mlflow_lib = {"pypi": {"package": f"mlflow=={VERSION}"}}
+        else:
+            # When running a non-release version as the client the same version will not be
+            # available within Databricks.
+            _logger.warning(
+                "Your client is running a non-release version of MLflow. "
+                "This version is not available on the databricks runtime. "
+                "MLflow will fallback the MLflow version provided by the runtime. "
+                "This might lead to unforeseen issues. "
+            )
+            mlflow_lib = {"pypi": {"package": f"'mlflow<={VERSION}'"}}
+
+        # Check syntax of JSON - if it contains libraries and new_cluster, pull those out
+        if "new_cluster" in cluster_spec:
+            # Libraries are optional, so we don't require that this be specified
+            cluster_spec_libraries = cluster_spec.get("libraries", [])
+            libraries = (
+                # This is for development purposes only. If the cluster spec already includes
+                # an MLflow Git URI, then we don't append `mlflow_lib` to avoid having
+                # two different pip requirements for mlflow.
+                cluster_spec_libraries
+                if _contains_mlflow_git_uri(cluster_spec_libraries)
+                else cluster_spec_libraries + [mlflow_lib]
+            )
+            cluster_spec = cluster_spec["new_cluster"]
+        else:
+            libraries = [mlflow_lib]
+
+        # Make jobs API request to launch run.
+        req_body_json = {
+            "run_name": f"MLflow Run for {project_uri}",
+            "new_cluster": cluster_spec,
+            "shell_command_task": {"command": command, "env_vars": env_vars},
+            "libraries": libraries,
+        }
+        _logger.info("=== Submitting a run to execute the MLflow project... ===")
+        run_submit_res = self._jobs_runs_submit(req_body_json)
+        return run_submit_res["run_id"]
+
+    def run_databricks_spark_job(
+        self,
+        project_uri,
+        work_dir,
+        experiment_id,
+        cluster_spec,
+        run_id,
+        project_spec,
+        entry_point,
+        parameters,
+    ):
+        from mlflow.utils.file_utils import get_or_create_tmp_dir
+
+        dbfs_fuse_uri = self._upload_project_to_dbfs(work_dir, experiment_id)
+
+        env_vars = {
+            MLFLOW_TRACKING_URI.name: "databricks",
+            MLFLOW_EXPERIMENT_ID.name: experiment_id,
+            MLFLOW_RUN_ID.name: run_id,
+        }
+        _logger.info(
+            "=== Running databricks spark job of project %s on Databricks ===", project_uri
+        )
+
+        if project_spec.databricks_spark_job_spec.python_file is not None:
+            if entry_point != "main" or parameters:
+                _logger.warning(
+                    "You configured Databricks spark job python_file and parameters within the "
+                    "MLProject file's databricks_spark_job section. '--entry-point' "
+                    "and '--param-list' arguments specified in the 'mlflow run' command are "
+                    "ignored."
+                )
+            job_code_file = project_spec.databricks_spark_job_spec.python_file
+            job_parameters = project_spec.databricks_spark_job_spec.parameters
+        else:
+            command = project_spec.get_entry_point(entry_point).compute_command(parameters, None)
+            command_splits = command.split(" ")
+            if command_splits[0] != "python":
+                raise MlflowException(
+                    "Databricks spark job only supports 'python' command in the entry point "
+                    "configuration."
+                )
+            job_code_file = command_splits[1]
+            job_parameters = command_splits[2:]
+
+        tmp_dir = Path(get_or_create_tmp_dir())
+        origin_job_code = (Path(work_dir) / job_code_file).read_text()
+        job_code_filename = f"{uuid.uuid4().hex}.py"
+        new_job_code_file = tmp_dir / job_code_filename
+
+        project_dir, extracting_tar_command = _get_project_dir_and_extracting_tar_command(
+            dbfs_fuse_uri
+        )
+
+        env_vars_str = json.dumps(env_vars)
+        new_job_code_file.write_text(
+            f"""
+import os
+import subprocess
+os.environ.update({env_vars_str})
+
+extracting_tar_command = \"\"\"
+{extracting_tar_command}
+\"\"\"
+subprocess.check_call(extracting_tar_command, shell=True)
+
+os.chdir('{project_dir}')
+{origin_job_code}
+"""
+        )
+
+        dbfs_job_code_file_path = posixpath.join(
+            DBFS_EXPERIMENT_DIR_BASE,
+            str(experiment_id),
+            "projects-code",
+            job_code_filename,
+        )
+        job_code_file_dbfs_fuse_uri = posixpath.join("/dbfs", dbfs_job_code_file_path)
+        if not self._dbfs_path_exists(dbfs_job_code_file_path):
+            self._upload_to_dbfs(str(new_job_code_file), job_code_file_dbfs_fuse_uri)
+
+        libraries_config = [
+            {"pypi": {"package": python_lib}}
+            for python_lib in project_spec.databricks_spark_job_spec.python_libraries
+        ]
+        # Make Databricks Spark jobs API request to launch run.
+        req_body_json = {
+            "run_name": f"MLflow Run for {project_uri}",
+            "new_cluster": cluster_spec,
+            "libraries": libraries_config,
+            "spark_python_task": {
+                "python_file": f"dbfs:/{dbfs_job_code_file_path}",
+                "parameters": job_parameters,
+            },
+        }
+
+        _logger.info("=== Submitting a run to execute the MLflow project... ===")
+        run_submit_res = self._jobs_runs_submit(req_body_json)
+        return run_submit_res["run_id"]
+
+    def run_databricks(
+        self,
+        uri,
+        entry_point,
+        work_dir,
+        parameters,
+        experiment_id,
+        cluster_spec,
+        run_id,
+        env_manager,
+    ):
+        tracking_uri = _get_tracking_uri_for_run()
+        dbfs_fuse_uri = self._upload_project_to_dbfs(work_dir, experiment_id)
+
+        env_vars = {
+            MLFLOW_TRACKING_URI.name: tracking_uri,
+            MLFLOW_EXPERIMENT_ID.name: experiment_id,
+        }
+        _logger.info("=== Running entry point %s of project %s on Databricks ===", entry_point, uri)
+        # Launch run on Databricks
+        command = _get_databricks_run_cmd(
+            dbfs_fuse_uri, run_id, entry_point, parameters, env_manager
+        )
+        return self._run_shell_command_job(uri, command, env_vars, cluster_spec)
+
+    def _get_status(self, databricks_run_id):
+        run_state = self.get_run_result_state(databricks_run_id)
+        if run_state is None:
+            return RunStatus.RUNNING
+        if run_state == "SUCCESS":
+            return RunStatus.FINISHED
+        return RunStatus.FAILED
+
+    def get_status(self, databricks_run_id):
+        return RunStatus.to_string(self._get_status(databricks_run_id))
+
+    def get_run_result_state(self, databricks_run_id):
+        """
+        Get the run result state (string) of a Databricks job run.
+
+        Args:
+            databricks_run_id: Integer Databricks job run ID.
+
+        Returns:
+            `RunResultState <https://docs.databricks.com/api/latest/jobs.html#runresultstate>`_ or
+            None if the run is still active.
+        """
+        res = self.jobs_runs_get(databricks_run_id)
+        return res["state"].get("result_state", None)
+
+    def jobs_runs_cancel(self, databricks_run_id):
+        response = self._databricks_api_request(
+            endpoint="/api/2.0/jobs/runs/cancel", method="POST", json={"run_id": databricks_run_id}
+        )
+        return json.loads(response.text)
+
+    def jobs_runs_get(self, databricks_run_id):
+        response = self._databricks_api_request(
+            endpoint="/api/2.0/jobs/runs/get", method="GET", params={"run_id": databricks_run_id}
+        )
+        return json.loads(response.text)
+
+
+def _get_tracking_uri_for_run():
+    uri = tracking.get_tracking_uri()
+    if uri.startswith("databricks"):
+        return "databricks"
+    return uri
+
+
+def _get_cluster_mlflow_run_cmd(project_dir, run_id, entry_point, parameters, env_manager):
+    cmd = [
+        "mlflow",
+        "run",
+        project_dir,
+        "--entry-point",
+        entry_point,
+    ]
+    if env_manager:
+        cmd += ["--env-manager", env_manager]
+    mlflow_run_arr = list(map(quote, cmd))
+    if run_id:
+        mlflow_run_arr.extend(["-c", json.dumps({MLFLOW_LOCAL_BACKEND_RUN_ID_CONFIG: run_id})])
+    if parameters:
+        for key, value in parameters.items():
+            mlflow_run_arr.extend(["-P", f"{key}={value}"])
+    return mlflow_run_arr
+
+
+def _get_project_dir_and_extracting_tar_command(dbfs_fuse_tar_uri):
+    # Strip ".gz" and ".tar" file extensions from base filename of the tarfile
+    tar_hash = posixpath.splitext(posixpath.splitext(posixpath.basename(dbfs_fuse_tar_uri))[0])[0]
+    container_tar_path = posixpath.abspath(
+        posixpath.join(DB_TARFILE_BASE, posixpath.basename(dbfs_fuse_tar_uri))
+    )
+    project_dir = posixpath.join(DB_PROJECTS_BASE, tar_hash)
+    command = textwrap.dedent(
+        f"""
+    # Make local directories in the container into which to copy/extract the tarred project
+    mkdir -p {DB_TARFILE_BASE} {DB_PROJECTS_BASE} &&
+    # Rsync from DBFS FUSE to avoid copying archive into local filesystem if it already exists
+    rsync -a -v --ignore-existing {dbfs_fuse_tar_uri} {DB_TARFILE_BASE} &&
+    # Extract project into a temporary directory. We don't extract directly into the desired
+    # directory as tar extraction isn't guaranteed to be atomic
+    cd $(mktemp -d) &&
+    tar --no-same-owner -xzvf {container_tar_path} &&
+    # Atomically move the extracted project into the desired directory
+    mv -T {DB_TARFILE_ARCHIVE_NAME} {project_dir}"""
+    )
+    return project_dir, command
+
+
+def _get_databricks_run_cmd(dbfs_fuse_tar_uri, run_id, entry_point, parameters, env_manager):
+    """
+    Generate MLflow CLI command to run on Databricks cluster in order to launch a run on Databricks.
+    """
+    project_dir, extracting_tar_command = _get_project_dir_and_extracting_tar_command(
+        dbfs_fuse_tar_uri
+    )
+    mlflow_run_arr = _get_cluster_mlflow_run_cmd(
+        project_dir,
+        run_id,
+        entry_point,
+        parameters,
+        env_manager,
+    )
+    mlflow_run_cmd = " ".join([quote(elem) for elem in mlflow_run_arr])
+    shell_command = textwrap.dedent(
+        f"""
+    export PATH=$PATH:$DB_HOME/python/bin &&
+    mlflow --version &&
+    {extracting_tar_command} &&
+    {mlflow_run_cmd}
+    """
+    )
+    return ["bash", "-c", shell_command]
+
+
+def run_databricks(
+    remote_run, uri, entry_point, work_dir, parameters, experiment_id, cluster_spec, env_manager
+):
+    """
+    Run the project at the specified URI on Databricks, returning a ``SubmittedRun`` that can be
+    used to query the run's status or wait for the resulting Databricks Job run to terminate.
+    """
+    run_id = remote_run.info.run_id
+    db_job_runner = DatabricksJobRunner(databricks_profile_uri=tracking.get_tracking_uri())
+    db_run_id = db_job_runner.run_databricks(
+        uri, entry_point, work_dir, parameters, experiment_id, cluster_spec, run_id, env_manager
+    )
+    submitted_run = DatabricksSubmittedRun(db_run_id, run_id, db_job_runner)
+    submitted_run._print_description_and_log_tags()
+    return submitted_run
+
+
+def run_databricks_spark_job(
+    remote_run,
+    uri,
+    work_dir,
+    experiment_id,
+    cluster_spec,
+    project_spec,
+    entry_point,
+    parameters,
+):
+    run_id = remote_run.info.run_id
+    db_job_runner = DatabricksJobRunner(databricks_profile_uri=tracking.get_tracking_uri())
+    db_run_id = db_job_runner.run_databricks_spark_job(
+        uri,
+        work_dir,
+        experiment_id,
+        cluster_spec,
+        run_id,
+        project_spec,
+        entry_point,
+        parameters,
+    )
+    submitted_run = DatabricksSubmittedRun(db_run_id, run_id, db_job_runner)
+    submitted_run._print_description_and_log_tags()
+    return submitted_run
+
+
+class DatabricksSubmittedRun(SubmittedRun):
+    """
+    Instance of SubmittedRun corresponding to a Databricks Job run launched to run an MLflow
+    project. Note that run_id may be None, e.g. if we did not launch the run against a tracking
+    server accessible to the local client.
+
+    Args:
+        databricks_run_id: Run ID of the launched Databricks Job.
+        mlflow_run_id: ID of the MLflow project run.
+        databricks_job_runner: Instance of ``DatabricksJobRunner`` used to make Databricks API
+            requests.
+    """
+
+    # How often to poll run status when waiting on a run
+    POLL_STATUS_INTERVAL = 30
+
+    def __init__(self, databricks_run_id, mlflow_run_id, databricks_job_runner):
+        super().__init__()
+        self._databricks_run_id = databricks_run_id
+        self._mlflow_run_id = mlflow_run_id
+        self._job_runner = databricks_job_runner
+
+    def _print_description_and_log_tags(self):
+        _logger.info(
+            "=== Launched MLflow run as Databricks job run with ID %s."
+            " Getting run status page URL... ===",
+            self._databricks_run_id,
+        )
+        run_info = self._job_runner.jobs_runs_get(self._databricks_run_id)
+        jobs_page_url = run_info["run_page_url"]
+        _logger.info("=== Check the run's status at %s ===", jobs_page_url)
+        host_creds = databricks_utils.get_databricks_host_creds(
+            self._job_runner.databricks_profile_uri
+        )
+        tracking.MlflowClient().set_tag(
+            self._mlflow_run_id, MLFLOW_DATABRICKS_RUN_URL, jobs_page_url
+        )
+        tracking.MlflowClient().set_tag(
+            self._mlflow_run_id, MLFLOW_DATABRICKS_SHELL_JOB_RUN_ID, self._databricks_run_id
+        )
+        tracking.MlflowClient().set_tag(
+            self._mlflow_run_id, MLFLOW_DATABRICKS_WEBAPP_URL, host_creds.host
+        )
+        job_id = run_info.get("job_id")
+        # In some releases of Databricks we do not return the job ID. We start including it in DB
+        # releases 2.80 and above.
+        if job_id is not None:
+            tracking.MlflowClient().set_tag(
+                self._mlflow_run_id, MLFLOW_DATABRICKS_SHELL_JOB_ID, job_id
+            )
+
+    @property
+    def run_id(self):
+        return self._mlflow_run_id
+
+    def wait(self):
+        result_state = self._job_runner.get_run_result_state(self._databricks_run_id)
+        while result_state is None:
+            time.sleep(self.POLL_STATUS_INTERVAL)
+            result_state = self._job_runner.get_run_result_state(self._databricks_run_id)
+        return result_state == "SUCCESS"
+
+    def cancel(self):
+        self._job_runner.jobs_runs_cancel(self._databricks_run_id)
+        self.wait()
+
+    def get_status(self):
+        return self._job_runner.get_status(self._databricks_run_id)

venv/lib/python3.9/site-packages/mlflow/projects/docker.py

@@ -0,0 +1,166 @@
+import logging
+import os
+import posixpath
+import shutil
+import subprocess
+import tempfile
+import urllib.parse
+import urllib.request
+
+import docker
+from mlflow import tracking
+from mlflow.environment_variables import MLFLOW_TRACKING_URI
+from mlflow.exceptions import ExecutionException
+from mlflow.projects.utils import MLFLOW_DOCKER_WORKDIR_PATH
+from mlflow.utils import file_utils, process
+from mlflow.utils.databricks_utils import get_databricks_env_vars
+from mlflow.utils.file_utils import _handle_readonly_on_windows
+from mlflow.utils.git_utils import get_git_commit
+from mlflow.utils.mlflow_tags import MLFLOW_DOCKER_IMAGE_ID, MLFLOW_DOCKER_IMAGE_URI
+
+_logger = logging.getLogger(__name__)
+
+_GENERATED_DOCKERFILE_NAME = "Dockerfile.mlflow-autogenerated"
+_MLFLOW_DOCKER_TRACKING_DIR_PATH = "/mlflow/tmp/mlruns"
+_PROJECT_TAR_ARCHIVE_NAME = "mlflow-project-docker-build-context"
+
+
+def validate_docker_installation():
+    """
+    Verify if Docker is installed and running on host machine.
+    """
+    if shutil.which("docker") is None:
+        raise ExecutionException(
+            "Could not find Docker executable. "
+            "Ensure Docker is installed as per the instructions "
+            "at https://docs.docker.com/install/overview/."
+        )
+
+    cmd = ["docker", "info"]
+    prc = process._exec_cmd(
+        cmd,
+        throw_on_error=False,
+        capture_output=False,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+    )
+    if prc.returncode != 0:
+        joined_cmd = " ".join(cmd)
+        raise ExecutionException(
+            f"Ran `{joined_cmd}` to ensure docker daemon is running but it failed "
+            f"with the following output:\n{prc.stdout}"
+        )
+
+
+def validate_docker_env(project):
+    if not project.name:
+        raise ExecutionException(
+            "Project name in MLProject must be specified when using docker for image tagging."
+        )
+    if not project.docker_env.get("image"):
+        raise ExecutionException(
+            "Project with docker environment must specify the docker image "
+            "to use via an 'image' field under the 'docker_env' field."
+        )
+
+
+def build_docker_image(work_dir, repository_uri, base_image, run_id, build_image, docker_auth):
+    """
+    Build a docker image containing the project in `work_dir`, using the base image.
+    """
+    image_uri = _get_docker_image_uri(repository_uri=repository_uri, work_dir=work_dir)
+    client = docker.from_env()
+    if docker_auth is not None:
+        client.login(**docker_auth)
+
+    if not build_image:
+        if not client.images.list(name=base_image):
+            _logger.info(f"Pulling {base_image}")
+            image = client.images.pull(base_image)
+        else:
+            _logger.info(f"{base_image} already exists")
+            image = client.images.get(base_image)
+        image_uri = base_image
+    else:
+        dockerfile = (
+            f"FROM {base_image}\n COPY {_PROJECT_TAR_ARCHIVE_NAME}/ {MLFLOW_DOCKER_WORKDIR_PATH}\n"
+            f" WORKDIR {MLFLOW_DOCKER_WORKDIR_PATH}\n"
+        )
+        build_ctx_path = _create_docker_build_ctx(work_dir, dockerfile)
+        with open(build_ctx_path, "rb") as docker_build_ctx:
+            _logger.info("=== Building docker image %s ===", image_uri)
+            image, _ = client.images.build(
+                tag=image_uri,
+                forcerm=True,
+                dockerfile=posixpath.join(_PROJECT_TAR_ARCHIVE_NAME, _GENERATED_DOCKERFILE_NAME),
+                fileobj=docker_build_ctx,
+                custom_context=True,
+                encoding="gzip",
+            )
+        try:
+            os.remove(build_ctx_path)
+        except Exception:
+            _logger.info("Temporary docker context file %s was not deleted.", build_ctx_path)
+    tracking.MlflowClient().set_tag(run_id, MLFLOW_DOCKER_IMAGE_URI, image_uri)
+    tracking.MlflowClient().set_tag(run_id, MLFLOW_DOCKER_IMAGE_ID, image.id)
+    return image
+
+
+def _get_docker_image_uri(repository_uri, work_dir):
+    """
+    Args:
+        repository_uri: The URI of the Docker repository with which to tag the image. The
+            repository URI is used as the prefix of the image URI.
+        work_dir: Path to the working directory in which to search for a git commit hash
+    """
+    repository_uri = repository_uri if repository_uri else "docker-project"
+    # Optionally include first 7 digits of git SHA in tag name, if available.
+    git_commit = get_git_commit(work_dir)
+    version_string = ":" + git_commit[:7] if git_commit else ""
+    return repository_uri + version_string
+
+
+def _create_docker_build_ctx(work_dir, dockerfile_contents):
+    """
+    Creates build context tarfile containing Dockerfile and project code, returning path to tarfile
+    """
+    directory = tempfile.mkdtemp()
+    try:
+        dst_path = os.path.join(directory, "mlflow-project-contents")
+        shutil.copytree(src=work_dir, dst=dst_path)
+        with open(os.path.join(dst_path, _GENERATED_DOCKERFILE_NAME), "w") as handle:
+            handle.write(dockerfile_contents)
+        _, result_path = tempfile.mkstemp()
+        file_utils.make_tarfile(
+            output_filename=result_path, source_dir=dst_path, archive_name=_PROJECT_TAR_ARCHIVE_NAME
+        )
+    finally:
+        shutil.rmtree(directory, onerror=_handle_readonly_on_windows)
+    return result_path
+
+
+def get_docker_tracking_cmd_and_envs(tracking_uri):
+    cmds = []
+    env_vars = {}
+
+    local_path, container_tracking_uri = _get_local_uri_or_none(tracking_uri)
+    if local_path is not None:
+        cmds = ["-v", f"{local_path}:{_MLFLOW_DOCKER_TRACKING_DIR_PATH}"]
+        env_vars[MLFLOW_TRACKING_URI.name] = container_tracking_uri
+    env_vars.update(get_databricks_env_vars(tracking_uri))
+    return cmds, env_vars
+
+
+def _get_local_uri_or_none(uri):
+    if uri == "databricks":
+        return None, None
+    parsed_uri = urllib.parse.urlparse(uri)
+    if not parsed_uri.netloc and parsed_uri.scheme in ("", "file", "sqlite"):
+        path = urllib.request.url2pathname(parsed_uri.path)
+        if parsed_uri.scheme == "sqlite":
+            uri = file_utils.path_to_local_sqlite_uri(_MLFLOW_DOCKER_TRACKING_DIR_PATH)
+        else:
+            uri = file_utils.path_to_local_file_uri(_MLFLOW_DOCKER_TRACKING_DIR_PATH)
+        return path, uri
+    else:
+        return None, None

venv/lib/python3.9/site-packages/mlflow/projects/env_type.py

@@ -0,0 +1,4 @@
+DOCKER = "docker_env"
+PYTHON = "python_env"
+CONDA = "conda_env"
+ALL = [DOCKER, PYTHON, CONDA]

venv/lib/python3.9/site-packages/mlflow/projects/kubernetes.py

@@ -0,0 +1,165 @@
+import logging
+import os
+import time
+from datetime import datetime
+from shlex import quote, split
+from threading import RLock
+
+import kubernetes
+from kubernetes.config.config_exception import ConfigException
+
+import docker
+from mlflow.entities import RunStatus
+from mlflow.exceptions import ExecutionException
+from mlflow.projects.submitted_run import SubmittedRun
+
+_logger = logging.getLogger(__name__)
+
+_DOCKER_API_TIMEOUT = 300
+
+
+def push_image_to_registry(image_tag):
+    client = docker.from_env(timeout=_DOCKER_API_TIMEOUT)
+    _logger.info("=== Pushing docker image %s ===", image_tag)
+    for line in client.images.push(repository=image_tag, stream=True, decode=True):
+        if "error" in line and line["error"]:
+            raise ExecutionException(
+                "Error while pushing to docker registry: {error}".format(error=line["error"])
+            )
+    return client.images.get_registry_data(image_tag).id
+
+
+def _get_kubernetes_job_definition(
+    project_name, image_tag, image_digest, command, env_vars, job_template
+):
+    container_image = image_tag + "@" + image_digest
+    timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S-%f")
+    job_name = f"{project_name}-{timestamp}"
+    _logger.info("=== Creating Job %s ===", job_name)
+    if os.environ.get("KUBE_MLFLOW_TRACKING_URI") is not None:
+        env_vars["MLFLOW_TRACKING_URI"] = os.environ["KUBE_MLFLOW_TRACKING_URI"]
+    environment_variables = [{"name": k, "value": v} for k, v in env_vars.items()]
+    job_template["metadata"]["name"] = job_name
+    job_template["spec"]["template"]["spec"]["containers"][0]["name"] = project_name
+    job_template["spec"]["template"]["spec"]["containers"][0]["image"] = container_image
+    job_template["spec"]["template"]["spec"]["containers"][0]["command"] = command
+    if "env" not in job_template["spec"]["template"]["spec"]["containers"][0].keys():
+        job_template["spec"]["template"]["spec"]["containers"][0]["env"] = []
+    job_template["spec"]["template"]["spec"]["containers"][0]["env"] += environment_variables
+    return job_template
+
+
+def _get_run_command(entrypoint_command):
+    formatted_command = []
+    for cmd in entrypoint_command:
+        formatted_command.extend([quote(s) for s in split(cmd)])
+    return formatted_command
+
+
+def _load_kube_context(context=None):
+    try:
+        # trying to load either the context passed as arg or, if None,
+        # the one provided as env var `KUBECONFIG` or in `~/.kube/config`
+        kubernetes.config.load_kube_config(context=context)
+    except (OSError, ConfigException) as e:
+        _logger.debug('Error loading kube context "%s": %s', context, e)
+        _logger.info("No valid kube config found, using in-cluster configuration")
+        kubernetes.config.load_incluster_config()
+
+
+def run_kubernetes_job(
+    project_name,
+    active_run,
+    image_tag,
+    image_digest,
+    command,
+    env_vars,
+    kube_context=None,
+    job_template=None,
+):
+    job_template = _get_kubernetes_job_definition(
+        project_name, image_tag, image_digest, _get_run_command(command), env_vars, job_template
+    )
+    job_name = job_template["metadata"]["name"]
+    job_namespace = job_template["metadata"]["namespace"]
+    _load_kube_context(context=kube_context)
+    api_instance = kubernetes.client.BatchV1Api()
+    api_instance.create_namespaced_job(namespace=job_namespace, body=job_template, pretty=True)
+    return KubernetesSubmittedRun(active_run.info.run_id, job_name, job_namespace)
+
+
+class KubernetesSubmittedRun(SubmittedRun):
+    """
+    Instance of SubmittedRun corresponding to a Kubernetes Job run launched to run an MLflow
+    project.
+
+    Args:
+        mlflow_run_id: ID of the MLflow project run.
+        job_name: Kubernetes job name.
+        job_namespace: Kubernetes job namespace.
+    """
+
+    # How often to poll run status when waiting on a run
+    POLL_STATUS_INTERVAL = 5
+
+    def __init__(self, mlflow_run_id, job_name, job_namespace):
+        super().__init__()
+        self._mlflow_run_id = mlflow_run_id
+        self._job_name = job_name
+        self._job_namespace = job_namespace
+        self._status = RunStatus.SCHEDULED
+        self._status_lock = RLock()
+        self._kube_api = kubernetes.client.BatchV1Api()
+
+    @property
+    def run_id(self):
+        return self._mlflow_run_id
+
+    def wait(self):
+        while not RunStatus.is_terminated(self._update_status()):
+            time.sleep(self.POLL_STATUS_INTERVAL)
+
+        return self._status == RunStatus.FINISHED
+
+    def _update_status(self):
+        api_response = self._kube_api.read_namespaced_job_status(
+            name=self._job_name, namespace=self._job_namespace, pretty=True
+        )
+        status = api_response.status
+        with self._status_lock:
+            if RunStatus.is_terminated(self._status):
+                return self._status
+            if self._status == RunStatus.SCHEDULED:
+                if api_response.status.start_time is None:
+                    _logger.info("Waiting for Job to start")
+                else:
+                    _logger.info("Job started.")
+                    self._status = RunStatus.RUNNING
+            if status.conditions is not None:
+                for condition in status.conditions:
+                    if condition.status == "True":
+                        _logger.info(condition.message)
+                        if condition.type == "Failed":
+                            self._status = RunStatus.FAILED
+                        elif condition.type == "Complete":
+                            self._status = RunStatus.FINISHED
+        return self._status
+
+    def get_status(self):
+        status = self._status
+        return status if RunStatus.is_terminated(status) else self._update_status()
+
+    def cancel(self):
+        with self._status_lock:
+            if not RunStatus.is_terminated(self._status):
+                _logger.info("Cancelling job.")
+                self._kube_api.delete_namespaced_job(
+                    name=self._job_name,
+                    namespace=self._job_namespace,
+                    body=kubernetes.client.V1DeleteOptions(),
+                    pretty=True,
+                )
+                self._status = RunStatus.KILLED
+                _logger.info("Job cancelled.")
+            else:
+                _logger.info("Attempting to cancel a job that is already terminated.")

venv/lib/python3.9/site-packages/mlflow/projects/submitted_run.py

@@ -0,0 +1,106 @@
+import logging
+import os
+import signal
+from abc import abstractmethod
+
+from mlflow.entities import RunStatus
+from mlflow.utils.annotations import developer_stable
+
+_logger = logging.getLogger(__name__)
+
+
+@developer_stable
+class SubmittedRun:
+    """
+    Wrapper around an MLflow project run (e.g. a subprocess running an entry point
+    command or a Databricks job run) and exposing methods for waiting on and cancelling the run.
+    This class defines the interface that the MLflow project runner uses to manage the lifecycle
+    of runs launched in different environments (e.g. runs launched locally or on Databricks).
+
+    ``SubmittedRun`` is not thread-safe. That is, concurrent calls to wait() / cancel()
+    from multiple threads may inadvertently kill resources (e.g. local processes) unrelated to the
+    run.
+
+    NOTE:
+
+        Subclasses of ``SubmittedRun`` must expose a ``run_id`` member containing the
+        run's MLflow run ID.
+    """
+
+    @abstractmethod
+    def wait(self):
+        """
+        Wait for the run to finish, returning True if the run succeeded and false otherwise. Note
+        that in some cases (e.g. remote execution on Databricks), we may wait until the remote job
+        completes rather than until the MLflow run completes.
+        """
+
+    @abstractmethod
+    def get_status(self):
+        """
+        Get status of the run.
+        """
+
+    @abstractmethod
+    def cancel(self):
+        """
+        Cancel the run (interrupts the command subprocess, cancels the Databricks run, etc) and
+        waits for it to terminate. The MLflow run status may not be set correctly
+        upon run cancellation.
+        """
+
+    @property
+    @abstractmethod
+    def run_id(self):
+        pass
+
+
+class LocalSubmittedRun(SubmittedRun):
+    """
+    Instance of ``SubmittedRun`` corresponding to a subprocess launched to run an entry point
+    command locally.
+    """
+
+    def __init__(self, run_id, command_proc):
+        super().__init__()
+        self._run_id = run_id
+        self.command_proc = command_proc
+
+    @property
+    def run_id(self):
+        return self._run_id
+
+    def wait(self):
+        return self.command_proc.wait() == 0
+
+    def cancel(self):
+        # Interrupt child process if it hasn't already exited
+        if self.command_proc.poll() is None:
+            # Kill the the process tree rooted at the child if it's the leader of its own process
+            # group, otherwise just kill the child
+            try:
+                if self.command_proc.pid == os.getpgid(self.command_proc.pid):
+                    os.killpg(self.command_proc.pid, signal.SIGTERM)
+                else:
+                    self.command_proc.terminate()
+            except OSError:
+                # The child process may have exited before we attempted to terminate it, so we
+                # ignore OSErrors raised during child process termination
+                _logger.info(
+                    "Failed to terminate child process (PID %s) corresponding to MLflow "
+                    "run with ID %s. The process may have already exited.",
+                    self.command_proc.pid,
+                    self._run_id,
+                )
+            self.command_proc.wait()
+
+    def _get_status(self):
+        exit_code = self.command_proc.poll()
+        if exit_code is None:
+            return RunStatus.RUNNING
+        if exit_code == 0:
+            return RunStatus.FINISHED
+        return RunStatus.FAILED
+
+    def get_status(self):
+        return RunStatus.to_string(self._get_status())

venv/lib/python3.9/site-packages/mlflow/projects/utils.py

@@ -0,0 +1,346 @@
+import logging
+import os
+import pathlib
+import re
+import shutil
+import tempfile
+import urllib.parse
+import zipfile
+from io import BytesIO
+
+from mlflow import tracking
+from mlflow.entities import Param, SourceType
+from mlflow.environment_variables import MLFLOW_EXPERIMENT_ID, MLFLOW_RUN_ID, MLFLOW_TRACKING_URI
+from mlflow.exceptions import ExecutionException
+from mlflow.projects import _project_spec
+from mlflow.tracking import fluent
+from mlflow.tracking.context.default_context import _get_user
+from mlflow.utils.git_utils import get_git_commit, get_git_repo_url
+from mlflow.utils.mlflow_tags import (
+    LEGACY_MLFLOW_GIT_BRANCH_NAME,
+    LEGACY_MLFLOW_GIT_REPO_URL,
+    MLFLOW_GIT_BRANCH,
+    MLFLOW_GIT_COMMIT,
+    MLFLOW_GIT_REPO_URL,
+    MLFLOW_PARENT_RUN_ID,
+    MLFLOW_PROJECT_ENTRY_POINT,
+    MLFLOW_SOURCE_NAME,
+    MLFLOW_SOURCE_TYPE,
+    MLFLOW_USER,
+)
+from mlflow.utils.rest_utils import augmented_raise_for_status
+
+_FILE_URI_REGEX = re.compile(r"^file://.+")
+_ZIP_URI_REGEX = re.compile(r".+\.zip$")
+MLFLOW_LOCAL_BACKEND_RUN_ID_CONFIG = "_mlflow_local_backend_run_id"
+MLFLOW_DOCKER_WORKDIR_PATH = "/mlflow/projects/code/"
+
+PROJECT_ENV_MANAGER = "ENV_MANAGER"
+PROJECT_SYNCHRONOUS = "SYNCHRONOUS"
+PROJECT_DOCKER_ARGS = "DOCKER_ARGS"
+PROJECT_STORAGE_DIR = "STORAGE_DIR"
+PROJECT_BUILD_IMAGE = "build_image"
+PROJECT_DOCKER_AUTH = "docker_auth"
+GIT_FETCH_DEPTH = 1
+
+
+_logger = logging.getLogger(__name__)
+
+
+def _parse_subdirectory(uri):
+    # Parses a uri and returns the uri and subdirectory as separate values.
+    # Uses '#' as a delimiter.
+    unquoted_uri = _strip_quotes(uri)
+    subdirectory = ""
+    parsed_uri = unquoted_uri
+    if "#" in unquoted_uri:
+        subdirectory = unquoted_uri[unquoted_uri.find("#") + 1 :]
+        parsed_uri = unquoted_uri[: unquoted_uri.find("#")]
+    if subdirectory and "." in subdirectory:
+        raise ExecutionException("'.' is not allowed in project subdirectory paths.")
+    return parsed_uri, subdirectory
+
+
+def _strip_quotes(uri):
+    return uri.strip("'\"")
+
+
+def _get_storage_dir(storage_dir):
+    if storage_dir is not None and not os.path.exists(storage_dir):
+        os.makedirs(storage_dir)
+    return tempfile.mkdtemp(dir=storage_dir)
+
+
+def _expand_uri(uri):
+    if _is_local_uri(uri):
+        return os.path.abspath(uri)
+    return uri
+
+
+def _is_file_uri(uri):
+    """Returns True if the passed-in URI is a file:// URI."""
+    return _FILE_URI_REGEX.match(uri)
+
+
+def _is_git_repo(path) -> bool:
+    """Returns True if passed-in path is a valid git repository"""
+    import git
+
+    try:
+        repo = git.Repo(path)
+        if len(repo.branches) > 0:
+            return True
+    except git.exc.InvalidGitRepositoryError:
+        pass
+    return False
+
+
+def _parse_file_uri(uri: str) -> str:
+    """Converts file URIs to filesystem paths"""
+    if _is_file_uri(uri):
+        parsed_file_uri = urllib.parse.urlparse(uri)
+        return str(
+            pathlib.Path(parsed_file_uri.netloc, parsed_file_uri.path, parsed_file_uri.fragment)
+        )
+    return uri
+
+
+def _is_local_uri(uri: str) -> bool:
+    """Returns True if passed-in URI should be interpreted as a folder on the local filesystem."""
+    resolved_uri = pathlib.Path(_parse_file_uri(uri)).resolve()
+    return resolved_uri.exists()
+
+
+def _is_zip_uri(uri):
+    """Returns True if the passed-in URI points to a ZIP file."""
+    return _ZIP_URI_REGEX.match(uri)
+
+
+def _is_valid_branch_name(work_dir, version):
+    """
+    Returns True if the ``version`` is the name of a branch in a Git project.
+    ``work_dir`` must be the working directory in a git repo.
+    """
+    if version is not None:
+        from git import Repo
+        from git.exc import GitCommandError
+
+        repo = Repo(work_dir, search_parent_directories=True)
+        try:
+            return repo.git.rev_parse("--verify", f"refs/heads/{version}") != ""
+        except GitCommandError:
+            return False
+    return False
+
+
+def fetch_and_validate_project(uri, version, entry_point, parameters):
+    parameters = parameters or {}
+    work_dir = _fetch_project(uri=uri, version=version)
+    project = _project_spec.load_project(work_dir)
+    if entry_point_obj := project.get_entry_point(entry_point):
+        entry_point_obj._validate_parameters(parameters)
+    return work_dir
+
+
+def load_project(work_dir):
+    return _project_spec.load_project(work_dir)
+
+
+def _fetch_project(uri, version=None):
+    """
+    Fetch a project into a local directory, returning the path to the local project directory.
+    """
+    parsed_uri, subdirectory = _parse_subdirectory(uri)
+    use_temp_dst_dir = _is_zip_uri(parsed_uri) or not _is_local_uri(parsed_uri)
+    dst_dir = tempfile.mkdtemp() if use_temp_dst_dir else _parse_file_uri(parsed_uri)
+
+    if use_temp_dst_dir:
+        _logger.info("=== Fetching project from %s into %s ===", uri, dst_dir)
+    if _is_zip_uri(parsed_uri):
+        parsed_uri = _parse_file_uri(parsed_uri)
+        _unzip_repo(
+            zip_file=(parsed_uri if _is_local_uri(parsed_uri) else _fetch_zip_repo(parsed_uri)),
+            dst_dir=dst_dir,
+        )
+    elif _is_local_uri(parsed_uri):
+        if use_temp_dst_dir:
+            shutil.copytree(parsed_uri, dst_dir, dirs_exist_ok=True)
+        if version is not None:
+            if not _is_git_repo(_parse_file_uri(parsed_uri)):
+                raise ExecutionException("Setting a version is only supported for Git project URIs")
+            _fetch_git_repo(parsed_uri, version, dst_dir)
+    else:
+        _fetch_git_repo(parsed_uri, version, dst_dir)
+    res = os.path.abspath(os.path.join(dst_dir, subdirectory))
+    if not os.path.exists(res):
+        raise ExecutionException(f"Could not find subdirectory {subdirectory} of {dst_dir}")
+    return res
+
+
+def _unzip_repo(zip_file, dst_dir):
+    with zipfile.ZipFile(zip_file) as zip_in:
+        zip_in.extractall(dst_dir)
+
+
+_HEAD_BRANCH_REGEX = re.compile(r"^\s*HEAD branch:\s+(?P<branch>\S+)")
+
+
+def _get_head_branch(remote_show_output):
+    for line in remote_show_output.splitlines():
+        match = _HEAD_BRANCH_REGEX.match(line)
+        if match:
+            return match.group("branch")
+
+
+def _fetch_git_repo(uri, version, dst_dir):
+    """
+    Clone the git repo at ``uri`` into ``dst_dir``, checking out commit ``version`` (or defaulting
+    to the head commit of the repository's master branch if version is unspecified).
+    Assumes authentication parameters are specified by the environment, e.g. by a Git credential
+    helper.
+    """
+    # We defer importing git until the last moment, because the import requires that the git
+    # executable is available on the PATH, so we only want to fail if we actually need it.
+    import git
+
+    repo = git.Repo.init(dst_dir)
+    origin = next((remote for remote in repo.remotes), None)
+    if origin is None:
+        origin = repo.create_remote("origin", uri)
+    if version is not None:
+        try:
+            origin.fetch(refspec=version, depth=GIT_FETCH_DEPTH, tags=True)
+            repo.git.checkout(version)
+        except git.exc.GitCommandError as e:
+            raise ExecutionException(
+                f"Unable to checkout version '{version}' of git repo {uri}"
+                "- please ensure that the version exists in the repo. "
+                f"Error: {e}"
+            )
+    else:
+        g = git.cmd.Git(dst_dir)
+        cmd = ["git", "remote", "show", "origin"]
+        output = g.execute(cmd)
+        head_branch = _get_head_branch(output)
+        if head_branch is None:
+            raise ExecutionException(
+                "Failed to find HEAD branch. Output of `{cmd}`:\n{output}".format(
+                    cmd=" ".join(cmd), output=output
+                )
+            )
+        origin.fetch(head_branch, depth=GIT_FETCH_DEPTH)
+        ref = origin.refs[0]
+        _logger.info("Fetched '%s' branch", head_branch)
+        repo.create_head(head_branch, ref)
+        repo.heads[head_branch].checkout()
+    repo.git.execute(command=["git", "submodule", "update", "--init", "--recursive"])
+
+
+def _fetch_zip_repo(uri):
+    import requests
+
+    # TODO (dbczumar): Replace HTTP resolution via ``requests.get`` with an invocation of
+    # ```mlflow.data.download_uri()`` when the API supports the same set of available stores as
+    # the artifact repository (Azure, FTP, etc). See the following issue:
+    # https://github.com/mlflow/mlflow/issues/763.
+    response = requests.get(uri)
+    try:
+        augmented_raise_for_status(response)
+    except requests.HTTPError as error:
+        raise ExecutionException(f"Unable to retrieve ZIP file. Reason: {error!s}")
+    return BytesIO(response.content)
+
+
+def get_or_create_run(run_id, uri, experiment_id, work_dir, version, entry_point, parameters):
+    if run_id:
+        return tracking.MlflowClient().get_run(run_id)
+    else:
+        return _create_run(uri, experiment_id, work_dir, version, entry_point, parameters)
+
+
+def _create_run(uri, experiment_id, work_dir, version, entry_point, parameters):
+    """
+    Create a ``Run`` against the current MLflow tracking server, logging metadata (e.g. the URI,
+    entry point, and parameters of the project) about the run. Return an ``ActiveRun`` that can be
+    used to report additional data about the run (metrics/params) to the tracking server.
+    """
+    if _is_local_uri(uri):
+        source_name = tracking._tracking_service.utils._get_git_url_if_present(_expand_uri(uri))
+    else:
+        source_name = _expand_uri(uri)
+    source_version = get_git_commit(work_dir)
+    existing_run = fluent.active_run()
+    parent_run_id = existing_run.info.run_id if existing_run else None
+
+    tags = {
+        MLFLOW_USER: _get_user(),
+        MLFLOW_SOURCE_NAME: source_name,
+        MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.PROJECT),
+        MLFLOW_PROJECT_ENTRY_POINT: entry_point,
+    }
+    if source_version is not None:
+        tags[MLFLOW_GIT_COMMIT] = source_version
+    if parent_run_id is not None:
+        tags[MLFLOW_PARENT_RUN_ID] = parent_run_id
+
+    repo_url = get_git_repo_url(work_dir)
+    if repo_url is not None:
+        tags[MLFLOW_GIT_REPO_URL] = repo_url
+        tags[LEGACY_MLFLOW_GIT_REPO_URL] = repo_url
+
+    # Add branch name tag if a branch is specified through -version
+    if _is_valid_branch_name(work_dir, version):
+        tags[MLFLOW_GIT_BRANCH] = version
+        tags[LEGACY_MLFLOW_GIT_BRANCH_NAME] = version
+    active_run = tracking.MlflowClient().create_run(experiment_id=experiment_id, tags=tags)
+
+    project = _project_spec.load_project(work_dir)
+    # Consolidate parameters for logging.
+    # `storage_dir` is `None` since we want to log actual path not downloaded local path
+    entry_point_obj = project.get_entry_point(entry_point)
+    if entry_point_obj:
+        final_params, extra_params = entry_point_obj.compute_parameters(
+            parameters, storage_dir=None
+        )
+        params_list = [
+            Param(key, value)
+            for key, value in list(final_params.items()) + list(extra_params.items())
+        ]
+        tracking.MlflowClient().log_batch(active_run.info.run_id, params=params_list)
+    return active_run
+
+
+def get_entry_point_command(project, entry_point, parameters, storage_dir):
+    """
+    Returns the shell command to execute in order to run the specified entry point.
+
+    Args:
+        project: Project containing the target entry point.
+        entry_point: Entry point to run.
+        parameters: Parameters (dictionary) for the entry point command.
+        storage_dir: Base local directory to use for downloading remote artifacts passed to
+            arguments of type 'path'. If None, a temporary base directory is used.
+    """
+    storage_dir_for_run = _get_storage_dir(storage_dir)
+    _logger.info(
+        "=== Created directory %s for downloading remote URIs passed to arguments of"
+        " type 'path' ===",
+        storage_dir_for_run,
+    )
+    commands = []
+    commands.append(
+        project.get_entry_point(entry_point).compute_command(parameters, storage_dir_for_run)
+    )
+    return commands
+
+
+def get_run_env_vars(run_id, experiment_id):
+    """
+    Returns a dictionary of environment variable key-value pairs to set in subprocess launched
+    to run MLflow projects.
+    """
+    return {
+        MLFLOW_RUN_ID.name: run_id,
+        MLFLOW_TRACKING_URI.name: tracking.get_tracking_uri(),
+        MLFLOW_EXPERIMENT_ID.name: str(experiment_id),
+    }