star

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 [], {}