genesis-flow 1.0.0__py3-none-any.whl

mlflow/data/tensorflow_dataset.py

@@ -0,0 +1,350 @@
+import json
+import logging
+from functools import cached_property
+from typing import Any, Optional, Union
+
+import numpy as np
+
+from mlflow.data.dataset import Dataset
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.data.digest_utils import (
+    MAX_ROWS,
+    compute_numpy_digest,
+    get_normalized_md5_digest,
+)
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.data.pyfunc_dataset_mixin import PyFuncConvertibleDatasetMixin, PyFuncInputsOutputs
+from mlflow.data.schema import TensorDatasetSchema
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR, INVALID_PARAMETER_VALUE
+from mlflow.types.schema import Schema
+from mlflow.types.utils import _infer_schema
+
+_logger = logging.getLogger(__name__)
+
+
+class TensorFlowDataset(Dataset, PyFuncConvertibleDatasetMixin):
+    """
+    Represents a TensorFlow dataset for use with MLflow Tracking.
+    """
+
+    def __init__(
+        self,
+        features,
+        source: DatasetSource,
+        targets=None,
+        name: Optional[str] = None,
+        digest: Optional[str] = None,
+    ):
+        """
+        Args:
+            features: A TensorFlow dataset or tensor of features.
+            source: The source of the TensorFlow dataset.
+            targets: A TensorFlow dataset or tensor of targets. Optional.
+            name: The name of the dataset. E.g. "wiki_train". If unspecified, a name is
+                automatically generated.
+            digest: The digest (hash, fingerprint) of the dataset. If unspecified, a digest
+                is automatically computed.
+        """
+        import tensorflow as tf
+
+        if not isinstance(features, tf.data.Dataset) and not tf.is_tensor(features):
+            raise MlflowException(
+                f"'features' must be an instance of tf.data.Dataset or a TensorFlow Tensor."
+                f" Found: {type(features)}.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        if tf.is_tensor(features) and targets is not None and not tf.is_tensor(targets):
+            raise MlflowException(
+                f"If 'features' is a TensorFlow Tensor, then 'targets' must also be a TensorFlow"
+                f" Tensor. Found: {type(targets)}.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        if (
+            isinstance(features, tf.data.Dataset)
+            and targets is not None
+            and not isinstance(targets, tf.data.Dataset)
+        ):
+            raise MlflowException(
+                "If 'features' is an instance of tf.data.Dataset, then 'targets' must also be an"
+                f" instance of tf.data.Dataset. Found: {type(targets)}.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        self._features = features
+        self._targets = targets
+        super().__init__(source=source, name=name, digest=digest)
+
+    def _compute_tensorflow_dataset_digest(  # noqa: D417
+        self,
+        dataset,
+        targets=None,
+    ) -> str:
+        """Computes a digest for the given Tensorflow dataset.
+
+        Args:
+            dataset: A Tensorflow dataset.
+
+        Returns:
+            A string digest.
+        """
+        import pandas as pd
+        import tensorflow as tf
+
+        hashable_elements = []
+
+        def hash_tf_dataset_iterator_element(element):
+            if element is None:
+                return
+            flat_element = tf.nest.flatten(element)
+            flattened_array = np.concatenate([x.flatten() for x in flat_element])
+            trimmed_array = flattened_array[0:MAX_ROWS]
+            try:
+                hashable_elements.append(pd.util.hash_array(trimmed_array))
+            except TypeError:
+                hashable_elements.append(np.int64(trimmed_array.size))
+
+        for element in dataset.as_numpy_iterator():
+            hash_tf_dataset_iterator_element(element)
+        if targets is not None:
+            for element in targets.as_numpy_iterator():
+                hash_tf_dataset_iterator_element(element)
+
+        return get_normalized_md5_digest(hashable_elements)
+
+    def _compute_tensor_digest(
+        self,
+        tensor_data,
+        tensor_targets,
+    ) -> str:
+        """Computes a digest for the given Tensorflow tensor.
+
+        Args:
+            tensor_data: A Tensorflow tensor, representing the features.
+            tensor_targets: A Tensorflow tensor, representing the targets. Optional.
+
+        Returns:
+            A string digest.
+        """
+        if tensor_targets is None:
+            return compute_numpy_digest(tensor_data.numpy())
+        else:
+            return compute_numpy_digest(tensor_data.numpy(), tensor_targets.numpy())
+
+    def _compute_digest(self) -> str:
+        """
+        Computes a digest for the dataset. Called if the user doesn't supply
+        a digest when constructing the dataset.
+        """
+        import tensorflow as tf
+
+        if isinstance(self._features, tf.data.Dataset):
+            return self._compute_tensorflow_dataset_digest(self._features, self._targets)
+        return self._compute_tensor_digest(self._features, self._targets)
+
+    def to_dict(self) -> dict[str, str]:
+        """Create config dictionary for the dataset.
+
+        Returns a string dictionary containing the following fields: name, digest, source, source
+        type, schema, and profile.
+        """
+        schema = json.dumps(self.schema.to_dict()) if self.schema else None
+        config = super().to_dict()
+        config.update(
+            {
+                "schema": schema,
+                "profile": json.dumps(self.profile),
+            }
+        )
+        return config
+
+    @property
+    def data(self):
+        """
+        The underlying TensorFlow data.
+        """
+        return self._features
+
+    @property
+    def source(self) -> DatasetSource:
+        """
+        The source of the dataset.
+        """
+        return self._source
+
+    @property
+    def targets(self):
+        """
+        The targets of the dataset.
+        """
+        return self._targets
+
+    @property
+    def profile(self) -> Optional[Any]:
+        """
+        A profile of the dataset. May be None if no profile is available.
+        """
+        import tensorflow as tf
+
+        profile = {
+            "features_cardinality": int(self._features.cardinality().numpy())
+            if isinstance(self._features, tf.data.Dataset)
+            else int(tf.size(self._features).numpy()),
+        }
+        if self._targets is not None:
+            profile.update(
+                {
+                    "targets_cardinality": int(self._targets.cardinality().numpy())
+                    if isinstance(self._targets, tf.data.Dataset)
+                    else int(tf.size(self._targets).numpy()),
+                }
+            )
+        return profile
+
+    @cached_property
+    def schema(self) -> Optional[TensorDatasetSchema]:
+        """
+        An MLflow TensorSpec schema representing the tensor dataset
+        """
+        try:
+            features_schema = TensorFlowDataset._get_tf_object_schema(self._features)
+            targets_schema = None
+            if self._targets is not None:
+                targets_schema = TensorFlowDataset._get_tf_object_schema(self._targets)
+            return TensorDatasetSchema(features=features_schema, targets=targets_schema)
+        except Exception as e:
+            _logger.warning("Failed to infer schema for TensorFlow dataset. Exception: %s", e)
+            return None
+
+    @staticmethod
+    def _get_tf_object_schema(tf_object) -> Schema:
+        import tensorflow as tf
+
+        if isinstance(tf_object, tf.data.Dataset):
+            numpy_data = next(tf_object.as_numpy_iterator())
+            if isinstance(numpy_data, np.ndarray):
+                return _infer_schema(numpy_data)
+            elif isinstance(numpy_data, dict):
+                return TensorFlowDataset._get_schema_from_tf_dataset_dict_numpy_data(numpy_data)
+            elif isinstance(numpy_data, tuple):
+                return TensorFlowDataset._get_schema_from_tf_dataset_tuple_numpy_data(numpy_data)
+            else:
+                raise MlflowException(
+                    f"Failed to infer schema for tf.data.Dataset due to unrecognized numpy iterator"
+                    f" data type. Numpy iterator data types 'np.ndarray', 'dict', and 'tuple' are"
+                    f" supported. Found: {type(numpy_data)}.",
+                    INVALID_PARAMETER_VALUE,
+                )
+        elif tf.is_tensor(tf_object):
+            return _infer_schema(tf_object.numpy())
+        else:
+            raise MlflowException(
+                f"Cannot infer schema of an object that is not an instance of tf.data.Dataset or"
+                f" a TensorFlow Tensor. Found: {type(tf_object)}",
+                INTERNAL_ERROR,
+            )
+
+    @staticmethod
+    def _get_schema_from_tf_dataset_dict_numpy_data(numpy_data: dict[Any, Any]) -> Schema:
+        if not all(isinstance(data_element, np.ndarray) for data_element in numpy_data.values()):
+            raise MlflowException(
+                "Failed to infer schema for tf.data.Dataset. Schemas can only be inferred"
+                " if the dataset consists of tensors. Ragged tensors, tensor arrays, and"
+                " other types are not supported. Additionally, datasets with nested tensors"
+                " are not supported.",
+                INVALID_PARAMETER_VALUE,
+            )
+        return _infer_schema(numpy_data)
+
+    @staticmethod
+    def _get_schema_from_tf_dataset_tuple_numpy_data(numpy_data: tuple[Any]) -> Schema:
+        if not all(isinstance(data_element, np.ndarray) for data_element in numpy_data):
+            raise MlflowException(
+                "Failed to infer schema for tf.data.Dataset. Schemas can only be inferred"
+                " if the dataset consists of tensors. Ragged tensors, tensor arrays, and"
+                " other types are not supported. Additionally, datasets with nested tensors"
+                " are not supported.",
+                INVALID_PARAMETER_VALUE,
+            )
+        return _infer_schema(
+            {
+                # MLflow Schemas currently require each tensor to have a name, if more than
+                # one tensor is defined. Accordingly, use the index as the name
+                str(i): data_element
+                for i, data_element in enumerate(numpy_data)
+            }
+        )
+
+    def to_pyfunc(self) -> PyFuncInputsOutputs:
+        """
+        Converts the dataset to a collection of pyfunc inputs and outputs for model
+        evaluation. Required for use with mlflow.evaluate().
+        """
+        return PyFuncInputsOutputs(self._features, self._targets)
+
+    def to_evaluation_dataset(self, path=None, feature_names=None) -> EvaluationDataset:
+        """
+        Converts the dataset to an EvaluationDataset for model evaluation. Only supported if the
+        dataset is a Tensor. Required for use with mlflow.evaluate().
+        """
+        import tensorflow as tf
+
+        # check that data and targets are Tensors
+        if not tf.is_tensor(self._features):
+            raise MlflowException("Data must be a Tensor to convert to an EvaluationDataset.")
+        if self._targets is not None and not tf.is_tensor(self._targets):
+            raise MlflowException("Targets must be a Tensor to convert to an EvaluationDataset.")
+        return EvaluationDataset(
+            data=self._features.numpy(),
+            targets=self._targets.numpy() if self._targets is not None else None,
+            path=path,
+            feature_names=feature_names,
+            name=self.name,
+            digest=self.digest,
+        )
+
+
+def from_tensorflow(
+    features,
+    source: Optional[Union[str, DatasetSource]] = None,
+    targets=None,
+    name: Optional[str] = None,
+    digest: Optional[str] = None,
+) -> TensorFlowDataset:
+    """Constructs a TensorFlowDataset object from TensorFlow data, optional targets, and source.
+
+    If the source is path like, then this will construct a DatasetSource object from the source
+    path. Otherwise, the source is assumed to be a DatasetSource object.
+
+    Args:
+        features: A TensorFlow dataset or tensor of features.
+        source: The source from which the data was derived, e.g. a filesystem
+            path, an S3 URI, an HTTPS URL, a delta table name with version, or
+            spark table etc. If source is not a path like string,
+            pass in a DatasetSource object directly. If no source is specified,
+            a CodeDatasetSource is used, which will source information from the run
+            context.
+        targets: A TensorFlow dataset or tensor of targets. Optional.
+        name: The name of the dataset. If unspecified, a name is generated.
+        digest: A dataset digest (hash). If unspecified, a digest is computed
+            automatically.
+    """
+    from mlflow.data.code_dataset_source import CodeDatasetSource
+    from mlflow.data.dataset_source_registry import resolve_dataset_source
+    from mlflow.tracking.context import registry
+
+    if source is not None:
+        if isinstance(source, DatasetSource):
+            resolved_source = source
+        else:
+            resolved_source = resolve_dataset_source(
+                source,
+            )
+    else:
+        context_tags = registry.resolve_tags()
+        resolved_source = CodeDatasetSource(tags=context_tags)
+    return TensorFlowDataset(
+        features=features, source=resolved_source, targets=targets, name=name, digest=digest
+    )

mlflow/data/uc_volume_dataset_source.py

@@ -0,0 +1,81 @@
+import logging
+from typing import Any
+
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.exceptions import MlflowException
+
+_logger = logging.getLogger(__name__)
+
+
+class UCVolumeDatasetSource(DatasetSource):
+    """Represents the source of a dataset stored in Databricks Unified Catalog Volume.
+
+    If you are using a delta table, please use `mlflow.data.delta_dataset_source.DeltaDatasetSource`
+    instead. This `UCVolumeDatasetSource` does not provide loading function, and is mostly useful
+    when you are logging a `mlflow.data.meta_dataset.MetaDataset` to MLflow, i.e., you want
+    to log the source of dataset to MLflow without loading the dataset.
+
+    Args:
+        path: the UC path of your data. It should be a valid UC path following the pattern
+            "/Volumes/{catalog}/{schema}/{volume}/{file_path}". For example,
+            "/Volumes/MyCatalog/MySchema/MyVolume/MyFile.json".
+    """
+
+    def __init__(self, path: str):
+        self.path = path
+        self._verify_uc_path_is_valid()
+
+    def _verify_uc_path_is_valid(self):
+        """Verify if the path exists in Databricks Unified Catalog."""
+        try:
+            from databricks.sdk import WorkspaceClient
+
+            w = WorkspaceClient()
+        except ImportError:
+            _logger.warning(
+                "Cannot verify the path of `UCVolumeDatasetSource` because of missing"
+                "`databricks-sdk`. Please install `databricks-sdk` via "
+                "`pip install -U databricks-sdk`. This does not block creating "
+                "`UCVolumeDatasetSource`, but your `UCVolumeDatasetSource` might be invalid."
+            )
+            return
+        except Exception:
+            _logger.warning(
+                "Cannot verify the path of `UCVolumeDatasetSource` due to a connection failure "
+                "with Databricks workspace. Please run `mlflow.login()` to log in to Databricks. "
+                "This does not block creating `UCVolumeDatasetSource`, but your "
+                "`UCVolumeDatasetSource` might be invalid."
+            )
+            return
+
+        try:
+            # Check if `self.path` points to a valid UC file.
+            w.files.get_metadata(self.path)
+        except Exception:
+            try:
+                # Check if `self.path` points to a valid UC directory.
+                w.files.get_directory_metadata(self.path)
+                # Append a slash to `self.path` to indicate it's a directory.
+                self.path += "/" if not self.path.endswith("/") else ""
+            except Exception:
+                # Neither file nor directory exists, we throw an exception.
+                raise MlflowException(f"{self.path} does not exist in Databricks Unified Catalog.")
+
+    @staticmethod
+    def _get_source_type() -> str:
+        return "uc_volume"
+
+    @staticmethod
+    def _can_resolve(raw_source: Any):
+        raise NotImplementedError
+
+    @classmethod
+    def _resolve(cls, raw_source: str):
+        raise NotImplementedError
+
+    def to_dict(self) -> dict[Any, Any]:
+        return {"path": self.path}
+
+    @classmethod
+    def from_dict(cls, source_dict: dict[Any, Any]) -> "UCVolumeDatasetSource":
+        return cls(**source_dict)

mlflow/db.py

@@ -0,0 +1,27 @@
+import click
+
+
+@click.group("db")
+def commands():
+    """
+    Commands for managing an MLflow tracking database.
+    """
+
+
+@commands.command()
+@click.argument("url")
+def upgrade(url):
+    """
+    Upgrade the schema of an MLflow tracking database to the latest supported version.
+
+    **IMPORTANT**: Schema migrations can be slow and are not guaranteed to be transactional -
+    **always take a backup of your database before running migrations**. The migrations README,
+    which is located at
+    https://github.com/mlflow/mlflow/blob/master/mlflow/store/db_migrations/README.md, describes
+    large migrations and includes information about how to estimate their performance and
+    recover from failures.
+    """
+    import mlflow.store.db.utils
+
+    engine = mlflow.store.db.utils.create_sqlalchemy_engine_with_retry(url)
+    mlflow.store.db.utils._upgrade_db(engine)

mlflow/dspy/__init__.py

@@ -0,0 +1,17 @@
+from mlflow.dspy.autolog import autolog
+from mlflow.version import IS_TRACING_SDK_ONLY
+
+__all__ = ["autolog"]
+
+# Import model logging APIs only if mlflow skinny or full package is installed,
+# i.e., skip if only mlflow-tracing package is installed.
+if not IS_TRACING_SDK_ONLY:
+    from mlflow.dspy.load import _load_pyfunc, load_model
+    from mlflow.dspy.save import log_model, save_model
+
+    __all__ += [
+        "save_model",
+        "log_model",
+        "load_model",
+        "_load_pyfunc",
+    ]

mlflow/dspy/autolog.py

@@ -0,0 +1,197 @@
+import importlib
+
+from packaging.version import Version
+
+import mlflow
+from mlflow.dspy.constant import FLAVOR_NAME
+from mlflow.tracing.provider import trace_disabled
+from mlflow.tracing.utils import construct_full_inputs
+from mlflow.utils.annotations import experimental
+from mlflow.utils.autologging_utils import (
+    autologging_integration,
+    get_autologging_config,
+    safe_patch,
+)
+
+
+@experimental(version="2.18.0")
+def autolog(
+    log_traces: bool = True,
+    log_traces_from_compile: bool = False,
+    log_traces_from_eval: bool = True,
+    log_compiles: bool = False,
+    log_evals: bool = False,
+    disable: bool = False,
+    silent: bool = False,
+):
+    """
+    Enables (or disables) and configures autologging from DSPy to MLflow. Currently, the
+    MLflow DSPy flavor only supports autologging for tracing.
+
+    Args:
+        log_traces: If ``True``, traces are logged for DSPy models by using. If ``False``,
+            no traces are collected during inference. Default to ``True``.
+        log_traces_from_compile: If ``True``, traces are logged when compiling (optimizing)
+            DSPy programs. If ``False``, traces are only logged from normal model inference and
+            disabled when compiling. Default to ``False``.
+        log_traces_from_eval: If ``True``, traces are logged for DSPy models when running DSPy's
+            `built-in evaluator <https://dspy.ai/learn/evaluation/metrics/#evaluation>`_.
+            If ``False``, traces are only logged from normal model inference and disabled when
+            running the evaluator. Default to ``True``.
+        log_compiles: If ``True``, information about the optimization process is logged when
+            `Teleprompter.compile()` is called.
+        log_evals: If ``True``, information about the evaluation call is logged when
+            `Evaluate.__call__()` is called.
+        disable: If ``True``, disables the DSPy autologging integration. If ``False``,
+            enables the DSPy autologging integration.
+        silent: If ``True``, suppress all event logs and warnings from MLflow during DSPy
+            autologging. If ``False``, show all events and warnings.
+    """
+    # NB: The @autologging_integration annotation is used for adding shared logic. However, one
+    # caveat is that the wrapped function is NOT executed when disable=True is passed. This prevents
+    # us from running cleaning up logging when autologging is turned off. To workaround this, we
+    # annotate _autolog() instead of this entrypoint, and define the cleanup logic outside it.
+    # This needs to be called before doing any safe-patching (otherwise safe-patch will be no-op).
+    # TODO: since this implementation is inconsistent, explore a universal way to solve the issue.
+    _autolog(
+        log_traces=log_traces,
+        log_traces_from_compile=log_traces_from_compile,
+        log_traces_from_eval=log_traces_from_eval,
+        log_compiles=log_compiles,
+        log_evals=log_evals,
+        disable=disable,
+        silent=silent,
+    )
+
+    import dspy
+
+    from mlflow.dspy.callback import MlflowCallback
+    from mlflow.dspy.util import log_dspy_dataset, save_dspy_module_state
+
+    # Enable tracing by setting the MlflowCallback
+    if not disable:
+        if not any(isinstance(c, MlflowCallback) for c in dspy.settings.callbacks):
+            dspy.settings.configure(callbacks=[*dspy.settings.callbacks, MlflowCallback()])
+    else:
+        dspy.settings.configure(
+            callbacks=[c for c in dspy.settings.callbacks if not isinstance(c, MlflowCallback)]
+        )
+
+    # Patch teleprompter/evaluator not to generate traces by default
+    def patch_fn(original, self, *args, **kwargs):
+        # NB: Since calling mlflow.dspy.autolog() again does not unpatch a function, we need to
+        # check this flag at runtime to determine if we should generate traces.
+        # method to disable tracing for compile and evaluate by default
+        @trace_disabled
+        def _trace_disabled_fn(self, *args, **kwargs):
+            return original(self, *args, **kwargs)
+
+        def _compile_fn(self, *args, **kwargs):
+            if callback := _active_callback():
+                callback.optimizer_stack_level += 1
+            try:
+                if get_autologging_config(FLAVOR_NAME, "log_traces_from_compile"):
+                    result = original(self, *args, **kwargs)
+                else:
+                    result = _trace_disabled_fn(self, *args, **kwargs)
+                return result
+            finally:
+                if callback:
+                    callback.optimizer_stack_level -= 1
+                    if callback.optimizer_stack_level == 0:
+                        # Reset the callback state after the completion of root compile
+                        callback.reset()
+
+        if isinstance(self, Teleprompter):
+            if not get_autologging_config(FLAVOR_NAME, "log_compiles"):
+                return _compile_fn(self, *args, **kwargs)
+
+            program = _compile_fn(self, *args, **kwargs)
+            # Save the state of the best model in json format
+            # so that users can see the demonstrations and instructions.
+            save_dspy_module_state(program, "best_model.json")
+
+            # Teleprompter.get_params is introduced in dspy 2.6.15
+            params = (
+                self.get_params()
+                if Version(importlib.metadata.version("dspy")) >= Version("2.6.15")
+                else {}
+            )
+            # Construct the dict of arguments passed to the compile call
+            inputs = construct_full_inputs(original, self, *args, **kwargs)
+            # Update params with the arguments passed to the compile call
+            params.update(inputs)
+            mlflow.log_params(
+                {k: v for k, v in inputs.items() if isinstance(v, (int, float, str, bool))}
+            )
+
+            if trainset := inputs.get("trainset"):
+                log_dspy_dataset(trainset, "trainset.json")
+            if valset := inputs.get("valset"):
+                log_dspy_dataset(valset, "valset.json")
+            return program
+
+        if isinstance(self, Teleprompter) and get_autologging_config(
+            FLAVOR_NAME, "log_traces_from_compile"
+        ):
+            return original(self, *args, **kwargs)
+
+        if isinstance(self, Evaluate) and get_autologging_config(
+            FLAVOR_NAME, "log_traces_from_eval"
+        ):
+            return original(self, *args, **kwargs)
+
+        return _trace_disabled_fn(self, *args, **kwargs)
+
+    from dspy.evaluate import Evaluate
+    from dspy.teleprompt import Teleprompter
+
+    compile_patch = "compile"
+    for cls in Teleprompter.__subclasses__():
+        # NB: This is to avoid the abstraction inheritance of superclasses that are defined
+        # only for the purposes of abstraction. The recursion behavior of the
+        # __subclasses__ dunder method will target the appropriate subclasses we need to patch.
+        if hasattr(cls, compile_patch):
+            safe_patch(
+                FLAVOR_NAME,
+                cls,
+                compile_patch,
+                patch_fn,
+                manage_run=get_autologging_config(FLAVOR_NAME, "log_compiles"),
+            )
+
+    call_patch = "__call__"
+    if hasattr(Evaluate, call_patch):
+        safe_patch(
+            FLAVOR_NAME,
+            Evaluate,
+            call_patch,
+            patch_fn,
+        )
+
+
+# This is required by mlflow.autolog()
+autolog.integration_name = FLAVOR_NAME
+
+
+@autologging_integration(FLAVOR_NAME)
+def _autolog(
+    log_traces: bool,
+    log_traces_from_compile: bool,
+    log_traces_from_eval: bool,
+    log_compiles: bool,
+    log_evals: bool,
+    disable: bool = False,
+    silent: bool = False,
+):
+    pass
+
+
+def _active_callback():
+    import dspy
+
+    from mlflow.dspy.callback import MlflowCallback
+
+    for callback in dspy.settings.callbacks:
+        if isinstance(callback, MlflowCallback):
+            return callback