genesis-flow 1.0.0__py3-none-any.whl

mlflow/utils/autologging_utils/versioning.py

@@ -0,0 +1,95 @@
+import importlib
+import importlib.metadata
+import re
+from typing import Literal
+
+from packaging.version import InvalidVersion, Version
+
+from mlflow.ml_package_versions import _ML_PACKAGE_VERSIONS, FLAVOR_TO_MODULE_NAME
+from mlflow.utils.databricks_utils import is_in_databricks_runtime
+
+
+def _check_version_in_range(ver, min_ver, max_ver):
+    return Version(min_ver) <= Version(ver) <= Version(max_ver)
+
+
+def _check_spark_version_in_range(ver, min_ver, max_ver):
+    """
+    Utility function for allowing late addition release changes to PySpark minor version increments
+    to be accepted, provided that the previous minor version has been previously validated.
+    For example, if version 3.2.1 has been validated as functional with MLflow, an upgrade of
+    PySpark's minor version to 3.2.2 will still provide a valid version check.
+    """
+    parsed_ver = Version(ver)
+    if parsed_ver > Version(min_ver):
+        ver = f"{parsed_ver.major}.{parsed_ver.minor}"
+    return _check_version_in_range(ver, min_ver, max_ver)
+
+
+def _violates_pep_440(ver):
+    try:
+        Version(ver)
+        return False
+    except InvalidVersion:
+        return True
+
+
+def _is_pre_or_dev_release(ver):
+    v = Version(ver)
+    return v.is_devrelease or v.is_prerelease
+
+
+def _strip_dev_version_suffix(version):
+    return re.sub(r"(\.?)dev.*", "", version)
+
+
+def get_min_max_version_and_pip_release(
+    flavor_name: str, category: Literal["autologging", "models"] = "autologging"
+):
+    if flavor_name == "pyspark.ml":
+        # pyspark.ml is a special case of spark flavor
+        flavor_name = "spark"
+
+    min_version = _ML_PACKAGE_VERSIONS[flavor_name][category]["minimum"]
+    max_version = _ML_PACKAGE_VERSIONS[flavor_name][category]["maximum"]
+    pip_release = _ML_PACKAGE_VERSIONS[flavor_name]["package_info"]["pip_release"]
+    return min_version, max_version, pip_release
+
+
+def is_flavor_supported_for_associated_package_versions(flavor_name):
+    """
+    Returns:
+        True if the specified flavor is supported for the currently-installed versions of its
+        associated packages.
+    """
+    module_name = FLAVOR_TO_MODULE_NAME[flavor_name]
+
+    try:
+        actual_version = importlib.import_module(module_name).__version__
+    except AttributeError:
+        try:
+            # NB: Module name is not necessarily the same as the package name. However,
+            # we assume they are the same here for simplicity. If they are not the same,
+            # this will fail and fallback to 'True', which is not a disaster.
+            actual_version = importlib.metadata.version(module_name)
+        except importlib.metadata.PackageNotFoundError:
+            # Some package (e.g. dspy) do not publish version info in a standard format.
+            # For this case, we assume the package version is supported by MLflow.
+            return True
+
+    # In Databricks, treat 'pyspark 3.x.y.dev0' as 'pyspark 3.x.y'
+    if module_name == "pyspark" and is_in_databricks_runtime():
+        actual_version = _strip_dev_version_suffix(actual_version)
+
+    if _violates_pep_440(actual_version) or _is_pre_or_dev_release(actual_version):
+        return False
+    min_version, max_version, _ = get_min_max_version_and_pip_release(flavor_name)
+
+    if module_name == "pyspark" and is_in_databricks_runtime():
+        # MLflow 1.25.0 is known to be compatible with PySpark 3.3.0 on Databricks, despite the
+        # fact that PySpark 3.3.0 was not available in PyPI at the time of the MLflow 1.25.0 release
+        if Version(max_version) < Version("3.3.0"):
+            max_version = "3.3.0"
+        return _check_spark_version_in_range(actual_version, min_version, max_version)
+    else:
+        return _check_version_in_range(actual_version, min_version, max_version)

mlflow/utils/checkpoint_utils.py

@@ -0,0 +1,206 @@
+import logging
+import os
+import posixpath
+
+import mlflow
+from mlflow.exceptions import MlflowException
+from mlflow.utils.autologging_utils import (
+    ExceptionSafeAbstractClass,
+)
+from mlflow.utils.file_utils import TempDir
+from mlflow.utils.mlflow_tags import LATEST_CHECKPOINT_ARTIFACT_TAG_KEY
+
+_logger = logging.getLogger(__name__)
+
+
+_CHECKPOINT_DIR = "checkpoints"
+_CHECKPOINT_METRIC_FILENAME = "checkpoint_metrics.json"
+_CHECKPOINT_MODEL_FILENAME = "checkpoint"
+_LATEST_CHECKPOINT_PREFIX = "latest_"
+_CHECKPOINT_EPOCH_PREFIX = "epoch_"
+_CHECKPOINT_GLOBAL_STEP_PREFIX = "global_step_"
+_WEIGHT_ONLY_CHECKPOINT_SUFFIX = ".weights"
+
+
+class MlflowModelCheckpointCallbackBase(metaclass=ExceptionSafeAbstractClass):
+    """Callback base class for automatic model checkpointing to MLflow.
+
+    You must implement "save_checkpoint" method to save the model as the checkpoint file.
+    and you must call `check_and_save_checkpoint_if_needed` method in relevant
+    callback events to trigger automatic checkpointing.
+
+    Args:
+        checkpoint_file_suffix: checkpoint file suffix.
+        monitor: In automatic model checkpointing, the metric name to monitor if
+            you set `model_checkpoint_save_best_only` to True.
+        save_best_only: If True, automatic model checkpointing only saves when
+            the model is considered the "best" model according to the quantity
+            monitored and previous checkpoint model is overwritten.
+        mode: one of {"min", "max"}. In automatic model checkpointing,
+            if save_best_only=True, the decision to overwrite the current save file is made
+            based on either the maximization or the minimization of the monitored quantity.
+        save_weights_only: In automatic model checkpointing, if True, then
+            only the model's weights will be saved. Otherwise, the optimizer states,
+            lr-scheduler states, etc are added in the checkpoint too.
+        save_freq: `"epoch"` or integer. When using `"epoch"`, the callback
+            saves the model after each epoch. When using integer, the callback
+            saves the model at end of this many batches. Note that if the saving isn't
+            aligned to epochs, the monitored metric may potentially be less reliable (it
+            could reflect as little as 1 batch, since the metrics get reset
+            every epoch). Defaults to `"epoch"`.
+    """
+
+    def __init__(
+        self,
+        checkpoint_file_suffix,
+        monitor,
+        mode,
+        save_best_only,
+        save_weights_only,
+        save_freq,
+    ):
+        self.checkpoint_file_suffix = checkpoint_file_suffix
+        self.monitor = monitor
+        self.mode = mode
+        self.save_best_only = save_best_only
+        self.save_weights_only = save_weights_only
+        self.save_freq = save_freq
+        self.last_monitor_value = None
+
+        self.mlflow_tracking_uri = mlflow.get_tracking_uri()
+
+        if self.save_best_only:
+            if self.monitor is None:
+                raise MlflowException(
+                    "If checkpoint 'save_best_only' config is set to True, you need to set "
+                    "'monitor' config as well."
+                )
+            if self.mode not in ["min", "max"]:
+                raise MlflowException(
+                    "If checkpoint 'save_best_only' config is set to True, you need to set "
+                    "'mode' config and available modes includes 'min' and 'max', but you set "
+                    f"'mode' to '{self.mode}'."
+                )
+
+    def _is_new_checkpoint_better(self, new_monitor_value):
+        if self.last_monitor_value is None:
+            return True
+
+        if self.mode == "min":
+            return new_monitor_value < self.last_monitor_value
+
+        return new_monitor_value > self.last_monitor_value
+
+    def save_checkpoint(self, filepath: str):
+        raise NotImplementedError()
+
+    def check_and_save_checkpoint_if_needed(self, current_epoch, global_step, metric_dict):
+        # For distributed model training, trainer workers need to use the driver process
+        # mlflow_tracking_uri.
+        # Note that `self.mlflow_tracking_uri` value is assigned in the driver process
+        # then it is pickled to trainer workers.
+        mlflow.set_tracking_uri(self.mlflow_tracking_uri)
+
+        if self.save_best_only:
+            if self.monitor not in metric_dict:
+                _logger.warning(
+                    "Checkpoint logging is skipped, because checkpoint 'save_best_only' config is "
+                    "True, it requires to compare the monitored metric value, but the provided "
+                    "monitored metric value is not available."
+                )
+                return
+
+            new_monitor_value = metric_dict[self.monitor]
+            if not self._is_new_checkpoint_better(new_monitor_value):
+                # Current checkpoint is worse than last saved checkpoint,
+                # so skip checkpointing.
+                return
+
+            self.last_monitor_value = new_monitor_value
+
+        suffix = self.checkpoint_file_suffix
+
+        if self.save_best_only:
+            if self.save_weights_only:
+                checkpoint_model_filename = (
+                    f"{_LATEST_CHECKPOINT_PREFIX}{_CHECKPOINT_MODEL_FILENAME}"
+                    f"{_WEIGHT_ONLY_CHECKPOINT_SUFFIX}{suffix}"
+                )
+            else:
+                checkpoint_model_filename = (
+                    f"{_LATEST_CHECKPOINT_PREFIX}{_CHECKPOINT_MODEL_FILENAME}{suffix}"
+                )
+            checkpoint_metrics_filename = (
+                f"{_LATEST_CHECKPOINT_PREFIX}{_CHECKPOINT_METRIC_FILENAME}"
+            )
+            checkpoint_artifact_dir = _CHECKPOINT_DIR
+        else:
+            if self.save_freq == "epoch":
+                sub_dir_name = f"{_CHECKPOINT_EPOCH_PREFIX}{current_epoch}"
+            else:
+                sub_dir_name = f"{_CHECKPOINT_GLOBAL_STEP_PREFIX}{global_step}"
+
+            if self.save_weights_only:
+                checkpoint_model_filename = (
+                    f"{_CHECKPOINT_MODEL_FILENAME}{_WEIGHT_ONLY_CHECKPOINT_SUFFIX}{suffix}"
+                )
+            else:
+                checkpoint_model_filename = f"{_CHECKPOINT_MODEL_FILENAME}{suffix}"
+            checkpoint_metrics_filename = _CHECKPOINT_METRIC_FILENAME
+            checkpoint_artifact_dir = f"{_CHECKPOINT_DIR}/{sub_dir_name}"
+
+        mlflow.set_tag(
+            LATEST_CHECKPOINT_ARTIFACT_TAG_KEY,
+            f"{checkpoint_artifact_dir}/{checkpoint_model_filename}",
+        )
+
+        mlflow.log_dict(
+            {**metric_dict, "epoch": current_epoch, "global_step": global_step},
+            f"{checkpoint_artifact_dir}/{checkpoint_metrics_filename}",
+        )
+
+        with TempDir() as tmp_dir:
+            tmp_model_save_path = os.path.join(tmp_dir.path(), checkpoint_model_filename)
+            self.save_checkpoint(tmp_model_save_path)
+            mlflow.log_artifact(tmp_model_save_path, checkpoint_artifact_dir)
+
+
+def download_checkpoint_artifact(run_id=None, epoch=None, global_step=None, dst_path=None):
+    from mlflow.client import MlflowClient
+    from mlflow.utils.mlflow_tags import LATEST_CHECKPOINT_ARTIFACT_TAG_KEY
+
+    client = MlflowClient()
+
+    if run_id is None:
+        run = mlflow.active_run()
+        if run is None:
+            raise MlflowException(
+                "There is no active run, please provide the 'run_id' argument for "
+                "'load_checkpoint' invocation."
+            )
+        run_id = run.info.run_id
+    else:
+        run = client.get_run(run_id)
+
+    latest_checkpoint_artifact_path = run.data.tags.get(LATEST_CHECKPOINT_ARTIFACT_TAG_KEY)
+    if latest_checkpoint_artifact_path is None:
+        raise MlflowException("There is no logged checkpoint artifact in the current run.")
+
+    checkpoint_filename = posixpath.basename(latest_checkpoint_artifact_path)
+
+    if epoch is not None and global_step is not None:
+        raise MlflowException(
+            "Only one of 'epoch' and 'global_step' can be set for 'load_checkpoint'."
+        )
+    elif global_step is not None:
+        checkpoint_artifact_path = (
+            f"{_CHECKPOINT_DIR}/{_CHECKPOINT_GLOBAL_STEP_PREFIX}{global_step}/{checkpoint_filename}"
+        )
+    elif epoch is not None:
+        checkpoint_artifact_path = (
+            f"{_CHECKPOINT_DIR}/{_CHECKPOINT_EPOCH_PREFIX}{epoch}/{checkpoint_filename}"
+        )
+    else:
+        checkpoint_artifact_path = latest_checkpoint_artifact_path
+
+    return client.download_artifacts(run_id, checkpoint_artifact_path, dst_path=dst_path)

mlflow/utils/class_utils.py

@@ -0,0 +1,6 @@
+import importlib
+
+
+def _get_class_from_string(fully_qualified_class_name):
+    module, class_name = fully_qualified_class_name.rsplit(".", maxsplit=1)
+    return getattr(importlib.import_module(module), class_name)

mlflow/utils/cli_args.py

@@ -0,0 +1,257 @@
+"""
+Definitions of click options shared by several CLI commands.
+"""
+
+import warnings
+
+import click
+
+from mlflow.environment_variables import MLFLOW_DISABLE_ENV_MANAGER_CONDA_WARNING
+from mlflow.utils import env_manager as _EnvManager
+
+MODEL_PATH = click.option(
+    "--model-path",
+    "-m",
+    default=None,
+    metavar="PATH",
+    required=True,
+    help="Path to the model. The path is relative to the run with the given "
+    "run-id or local filesystem path without run-id.",
+)
+
+_model_uri_help_string = (
+    "URI to the model. A local path, a 'runs:/' URI, or a"
+    " remote storage URI (e.g., an 's3://' URI). For more information"
+    " about supported remote URIs for model artifacts, see"
+    " https://mlflow.org/docs/latest/tracking.html#artifact-stores"
+)
+
+MODEL_URI_BUILD_DOCKER = click.option(
+    "--model-uri",
+    "-m",
+    metavar="URI",
+    default=None,
+    required=False,
+    help="[Optional] " + _model_uri_help_string,
+)
+
+MODEL_URI = click.option(
+    "--model-uri",
+    "-m",
+    metavar="URI",
+    required=True,
+    help=_model_uri_help_string,
+)
+
+MLFLOW_HOME = click.option(
+    "--mlflow-home",
+    default=None,
+    metavar="PATH",
+    help="Path to local clone of MLflow project. Use for development only.",
+)
+
+RUN_ID = click.option(
+    "--run-id",
+    "-r",
+    default=None,
+    required=False,
+    metavar="ID",
+    help="ID of the MLflow run that generated the referenced content.",
+)
+
+
+def _resolve_env_manager(_, __, env_manager):
+    if env_manager is not None:
+        _EnvManager.validate(env_manager)
+        if env_manager == _EnvManager.CONDA and not MLFLOW_DISABLE_ENV_MANAGER_CONDA_WARNING.get():
+            warnings.warn(
+                (
+                    "Use of conda is discouraged. If you use it, please ensure that your use of "
+                    "conda complies with Anaconda's terms of service "
+                    "(https://legal.anaconda.com/policies/en/?name=terms-of-service). "
+                    "virtualenv is the recommended tool for environment reproducibility. "
+                    f"To suppress this warning, set the {MLFLOW_DISABLE_ENV_MANAGER_CONDA_WARNING} "
+                    "environment variable to 'TRUE'."
+                ),
+                UserWarning,
+                stacklevel=2,
+            )
+        return env_manager
+
+    return None
+
+
+def _create_env_manager_option(help_string, default=None):
+    return click.option(
+        "--env-manager",
+        default=default,
+        type=click.UNPROCESSED,
+        callback=_resolve_env_manager,
+        help=help_string,
+    )
+
+
+ENV_MANAGER = _create_env_manager_option(
+    default=_EnvManager.VIRTUALENV,
+    # '\b' prevents rewrapping text:
+    # https://click.palletsprojects.com/en/8.1.x/documentation/#preventing-rewrapping
+    help_string="""
+If specified, create an environment for MLmodel using the specified
+environment manager. The following values are supported:
+
+\b
+- local: use the local environment
+- virtualenv: use virtualenv (and pyenv for Python version management)
+- conda: use conda
+
+If unspecified, default to virtualenv.
+""",
+)
+
+ENV_MANAGER_PROJECTS = _create_env_manager_option(
+    help_string="""
+If specified, create an environment for MLproject using the specified
+environment manager. The following values are supported:
+
+\b
+- local: use the local environment
+- virtualenv: use virtualenv (and pyenv for Python version management)
+- conda: use conda
+
+If unspecified, the appropriate environment manager is automatically selected based on
+the project configuration. For example, if `MLproject.yaml` contains a `python_env` key,
+virtualenv is used.
+""",
+)
+
+ENV_MANAGER_DOCKERFILE = _create_env_manager_option(
+    default=None,
+    # '\b' prevents rewrapping text:
+    # https://click.palletsprojects.com/en/8.1.x/documentation/#preventing-rewrapping
+    help_string="""
+If specified, create an environment for MLmodel using the specified
+environment manager. The following values are supported:
+
+\b
+- local: use the local environment
+- virtualenv: use virtualenv (and pyenv for Python version management)
+- conda: use conda
+
+If unspecified, default to None, then MLflow will automatically pick the env manager
+based on the model's flavor configuration.
+If model-uri is specified: if python version is specified in the flavor configuration
+and no java installation is required, then we use local environment. Otherwise we use virtualenv.
+If no model-uri is provided, we use virtualenv.
+""",
+)
+
+
+INSTALL_MLFLOW = click.option(
+    "--install-mlflow",
+    is_flag=True,
+    default=False,
+    help="If specified and there is a conda or virtualenv environment to be activated "
+    "mlflow will be installed into the environment after it has been "
+    "activated. The version of installed mlflow will be the same as "
+    "the one used to invoke this command.",
+)
+
+HOST = click.option(
+    "--host",
+    "-h",
+    envvar="MLFLOW_HOST",
+    metavar="HOST",
+    default="127.0.0.1",
+    help="The network address to listen on (default: 127.0.0.1). "
+    "Use 0.0.0.0 to bind to all addresses if you want to access the tracking "
+    "server from other machines.",
+)
+
+PORT = click.option(
+    "--port",
+    "-p",
+    envvar="MLFLOW_PORT",
+    default=5000,
+    help="The port to listen on (default: 5000).",
+)
+
+TIMEOUT = click.option(
+    "--timeout",
+    "-t",
+    envvar="MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT",
+    default=60,
+    help="Timeout in seconds to serve a request (default: 60).",
+)
+
+# We use None to disambiguate manually selecting "4"
+WORKERS = click.option(
+    "--workers",
+    "-w",
+    envvar="MLFLOW_WORKERS",
+    default=None,
+    help="Number of gunicorn worker processes to handle requests (default: 4).",
+)
+
+MODELS_WORKERS = click.option(
+    "--workers",
+    "-w",
+    envvar="MLFLOW_MODELS_WORKERS",
+    default=None,
+    help="Number of uvicorn workers to handle requests when serving mlflow models (default: 1).",
+)
+
+ENABLE_MLSERVER = click.option(
+    "--enable-mlserver",
+    is_flag=True,
+    default=False,
+    help=(
+        "Enable serving with MLServer through the v2 inference protocol. "
+        "You can use environment variables to configure MLServer. "
+        "(See https://mlserver.readthedocs.io/en/latest/reference/settings.html)"
+    ),
+)
+
+ARTIFACTS_DESTINATION = click.option(
+    "--artifacts-destination",
+    envvar="MLFLOW_ARTIFACTS_DESTINATION",
+    metavar="URI",
+    default="./mlartifacts",
+    help=(
+        "The base artifact location from which to resolve artifact upload/download/list requests "
+        "(e.g. 's3://my-bucket'). Defaults to a local './mlartifacts' directory. This option only "
+        "applies when the tracking server is configured to stream artifacts and the experiment's "
+        "artifact root location is http or mlflow-artifacts URI."
+    ),
+)
+
+SERVE_ARTIFACTS = click.option(
+    "--serve-artifacts/--no-serve-artifacts",
+    envvar="MLFLOW_SERVE_ARTIFACTS",
+    is_flag=True,
+    default=True,
+    help="Enables serving of artifact uploads, downloads, and list requests "
+    "by routing these requests to the storage location that is specified by "
+    "'--artifacts-destination' directly through a proxy. The default location that "
+    "these requests are served from is a local './mlartifacts' directory which can be "
+    "overridden via the '--artifacts-destination' argument. To disable artifact serving, "
+    "specify `--no-serve-artifacts`. Default: True",
+)
+
+NO_CONDA = click.option(
+    "--no-conda",
+    is_flag=True,
+    help="If specified, use local environment.",
+)
+
+INSTALL_JAVA = click.option(
+    "--install-java",
+    is_flag=False,
+    flag_value=True,
+    default=None,
+    type=bool,
+    help="Installs Java in the image if needed. Default is None, "
+    "allowing MLflow to determine installation. Flavors requiring "
+    "Java, such as Spark, enable this automatically. "
+    "Note: This option only works with the UBUNTU base image; "
+    "Python base images do not support Java installation.",
+)