genesis-flow 1.0.0__py3-none-any.whl

mlflow/utils/conda.py

@@ -0,0 +1,354 @@
+import hashlib
+import json
+import logging
+import os
+
+import yaml
+
+from mlflow.environment_variables import MLFLOW_CONDA_CREATE_ENV_CMD, MLFLOW_CONDA_HOME
+from mlflow.exceptions import ExecutionException
+from mlflow.utils import process
+from mlflow.utils.environment import Environment
+from mlflow.utils.os import is_windows
+
+_logger = logging.getLogger(__name__)
+
+CONDA_EXE = "CONDA_EXE"
+
+
+def get_conda_command(conda_env_name):
+    #  Checking for newer conda versions
+    if not is_windows() and (CONDA_EXE in os.environ or MLFLOW_CONDA_HOME.defined):
+        conda_path = get_conda_bin_executable("conda")
+        activate_conda_env = [f"source {os.path.dirname(conda_path)}/../etc/profile.d/conda.sh"]
+        activate_conda_env += [f"conda activate {conda_env_name} 1>&2"]
+    else:
+        activate_path = get_conda_bin_executable("activate")
+        # in case os name is not 'nt', we are not running on windows. It introduces
+        # bash command otherwise.
+        if not is_windows():
+            return [f"source {activate_path} {conda_env_name} 1>&2"]
+        else:
+            return [f"conda activate {conda_env_name}"]
+    return activate_conda_env
+
+
+def get_conda_bin_executable(executable_name):
+    """
+    Return path to the specified executable, assumed to be discoverable within the 'bin'
+    subdirectory of a conda installation.
+
+    The conda home directory (expected to contain a 'bin' subdirectory) is configurable via the
+    ``mlflow.projects.MLFLOW_CONDA_HOME`` environment variable. If
+    ``mlflow.projects.MLFLOW_CONDA_HOME`` is unspecified, this method simply returns the passed-in
+    executable name.
+    """
+    if conda_home := MLFLOW_CONDA_HOME.get():
+        return os.path.join(conda_home, f"bin/{executable_name}")
+    # Use CONDA_EXE as per https://github.com/conda/conda/issues/7126
+    if conda_exe := os.getenv(CONDA_EXE):
+        conda_bin_dir = os.path.dirname(conda_exe)
+        return os.path.join(conda_bin_dir, executable_name)
+    return executable_name
+
+
+def _get_conda_env_name(conda_env_path, env_id=None, env_root_dir=None):
+    if conda_env_path:
+        with open(conda_env_path) as f:
+            conda_env_contents = f.read()
+    else:
+        conda_env_contents = ""
+
+    if env_id:
+        conda_env_contents += env_id
+
+    env_name = "mlflow-{}".format(
+        hashlib.sha1(conda_env_contents.encode("utf-8"), usedforsecurity=False).hexdigest()
+    )
+    if env_root_dir:
+        env_root_dir = os.path.normpath(env_root_dir)
+        # Generate env name with format "mlflow-{conda_env_contents_hash}-{env_root_dir_hash}"
+        # hashing `conda_env_contents` and `env_root_dir` separately helps debugging
+        env_name += "-{}".format(
+            hashlib.sha1(env_root_dir.encode("utf-8"), usedforsecurity=False).hexdigest()
+        )
+
+    return env_name
+
+
+def _get_conda_executable_for_create_env():
+    """
+    Returns the executable that should be used to create environments. This is "conda"
+    by default, but it can be set to something else by setting the environment variable
+
+    """
+
+    return get_conda_bin_executable(MLFLOW_CONDA_CREATE_ENV_CMD.get())
+
+
+def _list_conda_environments(extra_env=None):
+    """Return a list of names of conda environments.
+
+    Args:
+        extra_env: Extra environment variables for running "conda env list" command.
+
+    """
+    prc = process._exec_cmd(
+        [get_conda_bin_executable("conda"), "env", "list", "--json"], extra_env=extra_env
+    )
+    return list(map(os.path.basename, json.loads(prc.stdout).get("envs", [])))
+
+
+_CONDA_ENVS_DIR = "conda_envs"
+_CONDA_CACHE_PKGS_DIR = "conda_cache_pkgs"
+_PIP_CACHE_DIR = "pip_cache_pkgs"
+
+
+def _create_conda_env(
+    conda_env_path,
+    conda_env_create_path,
+    project_env_name,
+    conda_extra_env_vars,
+    capture_output,
+):
+    if conda_env_path:
+        process._exec_cmd(
+            [
+                conda_env_create_path,
+                "env",
+                "create",
+                "-n",
+                project_env_name,
+                "--file",
+                conda_env_path,
+            ],
+            extra_env=conda_extra_env_vars,
+            capture_output=capture_output,
+        )
+    else:
+        process._exec_cmd(
+            [
+                conda_env_create_path,
+                "create",
+                "--channel",
+                "conda-forge",
+                "--yes",
+                "--override-channels",
+                "-n",
+                project_env_name,
+                "python",
+            ],
+            extra_env=conda_extra_env_vars,
+            capture_output=capture_output,
+        )
+
+    return Environment(get_conda_command(project_env_name), conda_extra_env_vars)
+
+
+def _create_conda_env_retry(
+    conda_env_path, conda_env_create_path, project_env_name, conda_extra_env_vars, _capture_output
+):
+    """
+    `conda env create` command can fail due to network issues such as `ConnectionResetError`
+    while collecting package metadata. This function retries the command up to 3 times.
+    """
+    num_attempts = 3
+    for attempt in range(num_attempts):
+        try:
+            return _create_conda_env(
+                conda_env_path,
+                conda_env_create_path,
+                project_env_name,
+                conda_extra_env_vars,
+                capture_output=True,
+            )
+        except process.ShellCommandException as e:
+            error_str = str(e)
+            if (num_attempts - attempt - 1) > 0 and (
+                "ConnectionResetError" in error_str or "ChunkedEncodingError" in error_str
+            ):
+                _logger.warning("Conda env creation failed due to network issue. Retrying...")
+                continue
+            raise
+
+
+def _get_conda_extra_env_vars(env_root_dir=None):
+    """
+    Given the `env_root_dir` (See doc of PyFuncBackend constructor argument `env_root_dir`),
+    return a dict of environment variables which are used to config conda to generate envs
+    under the expected `env_root_dir`.
+    """
+    if env_root_dir is None:
+        return None
+
+    # Create isolated conda package cache dir "conda_pkgs" under the env_root_dir
+    # for each python process.
+    # Note: shared conda package cache dir causes race condition issues:
+    # See https://github.com/conda/conda/issues/8870
+    # See https://docs.conda.io/projects/conda/en/latest/user-guide/configuration/use-condarc.html#specify-environment-directories-envs-dirs
+    # and https://docs.conda.io/projects/conda/en/latest/user-guide/configuration/use-condarc.html#specify-package-directories-pkgs-dirs
+
+    conda_envs_path = os.path.join(env_root_dir, _CONDA_ENVS_DIR)
+    conda_pkgs_path = os.path.join(env_root_dir, _CONDA_CACHE_PKGS_DIR)
+    pip_cache_dir = os.path.join(env_root_dir, _PIP_CACHE_DIR)
+
+    os.makedirs(conda_envs_path, exist_ok=True)
+    os.makedirs(conda_pkgs_path, exist_ok=True)
+    os.makedirs(pip_cache_dir, exist_ok=True)
+
+    return {
+        "CONDA_ENVS_PATH": conda_envs_path,
+        "CONDA_PKGS_DIRS": conda_pkgs_path,
+        "PIP_CACHE_DIR": pip_cache_dir,
+        # PIP_NO_INPUT=1 makes pip run in non-interactive mode,
+        # otherwise pip might prompt "yes or no" and ask stdin input
+        "PIP_NO_INPUT": "1",
+    }
+
+
+def get_or_create_conda_env(
+    conda_env_path,
+    env_id=None,
+    capture_output=False,
+    env_root_dir=None,
+    pip_requirements_override=None,
+    extra_envs=None,
+):
+    """Given a `Project`, creates a conda environment containing the project's dependencies if such
+    a conda environment doesn't already exist. Returns the name of the conda environment.
+
+    Args:
+        conda_env_path: Path to a conda yaml file.
+        env_id: Optional string that is added to the contents of the yaml file before
+            calculating the hash. It can be used to distinguish environments that have the
+            same conda dependencies but are supposed to be different based on the context.
+            For example, when serving the model we may install additional dependencies to the
+            environment after the environment has been activated.
+        capture_output: Specify the capture_output argument while executing the
+            "conda env create" command.
+        env_root_dir: See doc of PyFuncBackend constructor argument `env_root_dir`.
+        pip_requirements_override: If specified, install the specified python dependencies to
+            the environment (upgrade if already installed).
+        extra_envs: If specified, a dictionary of extra environment variables will be passed to the
+            model inference environment.
+
+    Returns:
+        The name of the conda environment.
+
+    """
+
+    conda_path = get_conda_bin_executable("conda")
+    conda_env_create_path = _get_conda_executable_for_create_env()
+
+    try:
+        # Checks if Conda executable exists
+        process._exec_cmd([conda_path, "--help"], throw_on_error=False, extra_env=extra_envs)
+    except OSError:
+        raise ExecutionException(
+            f"Could not find Conda executable at {conda_path}. "
+            "Ensure Conda is installed as per the instructions at "
+            "https://conda.io/projects/conda/en/latest/"
+            "user-guide/install/index.html. "
+            "You can also configure MLflow to look for a specific "
+            f"Conda executable by setting the {MLFLOW_CONDA_HOME} environment variable "
+            "to the path of the Conda executable"
+        )
+
+    try:
+        # Checks if executable for environment creation exists
+        process._exec_cmd(
+            [conda_env_create_path, "--help"], throw_on_error=False, extra_env=extra_envs
+        )
+    except OSError:
+        raise ExecutionException(
+            f"You have set the env variable {MLFLOW_CONDA_CREATE_ENV_CMD}, but "
+            f"{conda_env_create_path} does not exist or it is not working properly. "
+            f"Note that {conda_env_create_path} and the conda executable need to be "
+            "in the same conda environment. You can change the search path by"
+            f"modifying the env variable {MLFLOW_CONDA_HOME}"
+        )
+
+    conda_extra_env_vars = _get_conda_extra_env_vars(env_root_dir)
+    if extra_envs:
+        conda_extra_env_vars.update(extra_envs)
+
+    # Include the env_root_dir hash in the project_env_name,
+    # this is for avoid conda env name conflicts between different CONDA_ENVS_PATH.
+    project_env_name = _get_conda_env_name(conda_env_path, env_id=env_id, env_root_dir=env_root_dir)
+    if env_root_dir is not None:
+        project_env_path = os.path.join(env_root_dir, _CONDA_ENVS_DIR, project_env_name)
+    else:
+        project_env_path = project_env_name
+
+    if project_env_name in _list_conda_environments(conda_extra_env_vars):
+        _logger.info("Conda environment %s already exists.", project_env_path)
+        return Environment(get_conda_command(project_env_name), conda_extra_env_vars)
+
+    _logger.info("=== Creating conda environment %s ===", project_env_path)
+    try:
+        _create_conda_env_func = (
+            # Retry conda env creation in a pytest session to avoid flaky test failures
+            _create_conda_env_retry if "PYTEST_CURRENT_TEST" in os.environ else _create_conda_env
+        )
+        conda_env = _create_conda_env_func(
+            conda_env_path,
+            conda_env_create_path,
+            project_env_name,
+            conda_extra_env_vars,
+            capture_output,
+        )
+
+        if pip_requirements_override:
+            _logger.info(
+                "Installing additional dependencies specified"
+                f"by pip_requirements_override: {pip_requirements_override}"
+            )
+            cmd = [
+                conda_path,
+                "install",
+                "-n",
+                project_env_name,
+                "--yes",
+                *pip_requirements_override,
+            ]
+            process._exec_cmd(cmd, extra_env=conda_extra_env_vars, capture_output=capture_output)
+
+        return conda_env
+
+    except Exception:
+        try:
+            if project_env_name in _list_conda_environments(conda_extra_env_vars):
+                _logger.warning(
+                    "Encountered unexpected error while creating conda environment. Removing %s.",
+                    project_env_path,
+                )
+                process._exec_cmd(
+                    [
+                        conda_path,
+                        "remove",
+                        "--yes",
+                        "--name",
+                        project_env_name,
+                        "--all",
+                    ],
+                    extra_env=conda_extra_env_vars,
+                    capture_output=False,
+                )
+        except Exception as e:
+            _logger.warning(
+                "Removing conda environment %s failed (error: %s)",
+                project_env_path,
+                repr(e),
+            )
+        raise
+
+
+def _get_conda_dependencies(conda_yaml_path):
+    """Extracts conda dependencies from a conda yaml file.
+
+    Args:
+        conda_yaml_path: Conda yaml file path.
+    """
+    with open(conda_yaml_path) as f:
+        conda_yaml = yaml.safe_load(f)
+        return [d for d in conda_yaml.get("dependencies", []) if isinstance(d, str)]

mlflow/utils/credentials.py

@@ -0,0 +1,231 @@
+import configparser
+import getpass
+import logging
+import os
+from typing import NamedTuple, Optional
+
+from mlflow.environment_variables import (
+    MLFLOW_TRACKING_AUTH,
+    MLFLOW_TRACKING_AWS_SIGV4,
+    MLFLOW_TRACKING_CLIENT_CERT_PATH,
+    MLFLOW_TRACKING_INSECURE_TLS,
+    MLFLOW_TRACKING_PASSWORD,
+    MLFLOW_TRACKING_SERVER_CERT_PATH,
+    MLFLOW_TRACKING_TOKEN,
+    MLFLOW_TRACKING_USERNAME,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.utils.rest_utils import MlflowHostCreds
+
+_logger = logging.getLogger(__name__)
+
+
+class MlflowCreds(NamedTuple):
+    username: Optional[str]
+    password: Optional[str]
+
+
+def _get_credentials_path() -> str:
+    return os.path.expanduser("~/.mlflow/credentials")
+
+
+def _read_mlflow_creds_from_file() -> tuple[Optional[str], Optional[str]]:
+    path = _get_credentials_path()
+    if not os.path.exists(path):
+        return None, None
+
+    config = configparser.ConfigParser()
+    config.read(path)
+    if "mlflow" not in config:
+        return None, None
+
+    mlflow_cfg = config["mlflow"]
+    username_key = MLFLOW_TRACKING_USERNAME.name.lower()
+    password_key = MLFLOW_TRACKING_PASSWORD.name.lower()
+    return mlflow_cfg.get(username_key), mlflow_cfg.get(password_key)
+
+
+def _read_mlflow_creds_from_env() -> tuple[Optional[str], Optional[str]]:
+    return MLFLOW_TRACKING_USERNAME.get(), MLFLOW_TRACKING_PASSWORD.get()
+
+
+def read_mlflow_creds() -> MlflowCreds:
+    username_file, password_file = _read_mlflow_creds_from_file()
+    username_env, password_env = _read_mlflow_creds_from_env()
+    return MlflowCreds(
+        username=username_env or username_file,
+        password=password_env or password_file,
+    )
+
+
+def get_default_host_creds(store_uri):
+    creds = read_mlflow_creds()
+    return MlflowHostCreds(
+        host=store_uri,
+        username=creds.username,
+        password=creds.password,
+        token=MLFLOW_TRACKING_TOKEN.get(),
+        aws_sigv4=MLFLOW_TRACKING_AWS_SIGV4.get(),
+        auth=MLFLOW_TRACKING_AUTH.get(),
+        ignore_tls_verification=MLFLOW_TRACKING_INSECURE_TLS.get(),
+        client_cert_path=MLFLOW_TRACKING_CLIENT_CERT_PATH.get(),
+        server_cert_path=MLFLOW_TRACKING_SERVER_CERT_PATH.get(),
+    )
+
+
+def login(backend: str = "databricks", interactive: bool = True) -> None:
+    """Configure MLflow server authentication and connect MLflow to tracking server.
+
+    This method provides a simple way to connect MLflow to its tracking server. Currently only
+    Databricks tracking server is supported. Users will be prompted to enter the credentials if no
+    existing Databricks profile is found, and the credentials will be saved to `~/.databrickscfg`.
+
+    Args:
+        backend: string, the backend of the tracking server. Currently only "databricks" is
+            supported.
+
+        interactive: bool, controls request for user input on missing credentials. If true, user
+            input will be requested if no credentials are found, otherwise an exception will be
+            raised if no credentials are found.
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+
+        mlflow.login()
+        with mlflow.start_run():
+            mlflow.log_param("p", 0)
+    """
+    from mlflow.tracking import set_tracking_uri
+
+    if backend == "databricks":
+        _databricks_login(interactive)
+        set_tracking_uri("databricks")
+    else:
+        raise MlflowException(
+            f"Currently only 'databricks' backend is supported, received `backend={backend}`."
+        )
+
+
+def _validate_databricks_auth():
+    # Check if databricks credentials are valid.
+    try:
+        from databricks.sdk import WorkspaceClient
+    except ImportError:
+        raise ImportError(
+            "Databricks SDK is not installed. To use `mlflow.login()`, please install "
+            "databricks-sdk by `pip install databricks-sdk`."
+        )
+
+    try:
+        w = WorkspaceClient()
+        if "community" in w.config.host:
+            # Databricks community edition cannot use `w.current_user.me()` for auth validation.
+            w.clusters.list_zones()
+        else:
+            # If credentials are invalid, `w.current_user.me()` will throw an error.
+            w.current_user.me()
+        _logger.info(
+            f"Successfully connected to MLflow hosted tracking server! Host: {w.config.host}."
+        )
+    except Exception as e:
+        raise MlflowException(f"Failed to validate databricks credentials: {e}")
+
+
+def _overwrite_or_create_databricks_profile(
+    file_name,
+    profile,
+    profile_name="DEFAULT",
+):
+    """Overwrite or create a profile in the databricks config file.
+
+    Args:
+        file_name: string, the file name of the databricks config file, usually `~/.databrickscfg`.
+        profile: dict, contains the authentiacation profile information.
+        profile_name: string, the name of the profile to be overwritten or created.
+    """
+    profile_name = f"[{profile_name}]"
+    lines = []
+    # Read `file_name` if the file exists, otherwise `lines=[]`.
+    if os.path.exists(file_name):
+        with open(file_name) as file:
+            lines = file.readlines()
+    start_index = -1
+    end_index = -1
+    # Find the start and end indices of the profile to overwrite.
+    for i in range(len(lines)):
+        if lines[i].strip() == profile_name:
+            start_index = i
+            break
+
+    if start_index != -1:
+        for i in range(start_index + 1, len(lines)):
+            # Reach an empty line or a new profile.
+            if lines[i].strip() == "" or lines[i].startswith("["):
+                end_index = i
+                break
+        end_index = end_index if end_index != -1 else len(lines)
+        del lines[start_index : end_index + 1]
+
+    # Write the new profile to the top of the file.
+    new_profile = []
+    new_profile.append(profile_name + "\n")
+    new_profile.append(f"host = {profile['host']}\n")
+    if "token" in profile:
+        new_profile.append(f"token = {profile['token']}\n")
+    else:
+        new_profile.append(f"username = {profile['username']}\n")
+        new_profile.append(f"password = {profile['password']}\n")
+    new_profile.append("\n")
+    lines = new_profile + lines
+
+    # Write back the modified lines to the file.
+    with open(file_name, "w") as file:
+        file.writelines(lines)
+
+
+def _databricks_login(interactive):
+    """Set up databricks authentication."""
+    try:
+        # Failed validation will throw an error.
+        _validate_databricks_auth()
+        return
+    except Exception:
+        if interactive:
+            _logger.info("No valid Databricks credentials found, please enter your credentials...")
+        else:
+            raise MlflowException(
+                "No valid Databricks credentials found while running in non-interactive mode."
+            )
+
+    while True:
+        host = input("Databricks Host (should begin with https://): ")
+        if not host.startswith("https://"):
+            _logger.error("Invalid host: {host}, host must begin with https://, please retry.")
+        break
+
+    profile = {"host": host}
+    if "community" in host:
+        # Databricks community edition requires username and password for authentication.
+        username = input("Username: ")
+        password = getpass.getpass("Password: ")
+        profile["username"] = username
+        profile["password"] = password
+    else:
+        # Production or staging Databricks requires personal token for authentication.
+        token = getpass.getpass("Token: ")
+        profile["token"] = token
+
+    file_name = os.environ.get(
+        "DATABRICKS_CONFIG_FILE", f"{os.path.expanduser('~')}/.databrickscfg"
+    )
+    profile_name = os.environ.get("DATABRICKS_CONFIG_PROFILE", "DEFAULT")
+    _overwrite_or_create_databricks_profile(file_name, profile, profile_name)
+
+    try:
+        # Failed validation will throw an error.
+        _validate_databricks_auth()
+    except Exception as e:
+        # If user entered invalid auth, we will raise an error and ask users to retry.
+        raise MlflowException(f"`mlflow.login()` failed with error: {e}")

mlflow/utils/data_utils.py

@@ -0,0 +1,17 @@
+import urllib.parse
+
+
+def parse_s3_uri(uri):
+    """Parse an S3 URI, returning (bucket, path)"""
+    parsed = urllib.parse.urlparse(uri)
+    if parsed.scheme != "s3":
+        raise Exception(f"Not an S3 URI: {uri}")
+    path = parsed.path
+    if path.startswith("/"):
+        path = path[1:]
+    return parsed.netloc, path
+
+
+def is_uri(string):
+    parsed_uri = urllib.parse.urlparse(string)
+    return len(parsed_uri.scheme) > 0