genesis-flow 1.0.0__py3-none-any.whl

mlflow/data/code_dataset_source.py

@@ -0,0 +1,40 @@
+from typing import Any
+
+from typing_extensions import Self
+
+from mlflow.data.dataset_source import DatasetSource
+
+
+class CodeDatasetSource(DatasetSource):
+    def __init__(
+        self,
+        tags: dict[Any, Any],
+    ):
+        self._tags = tags
+
+    @staticmethod
+    def _get_source_type() -> str:
+        return "code"
+
+    def load(self, **kwargs):
+        """
+        Load is not implemented for Code Dataset Source.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def _can_resolve(raw_source: Any):
+        return False
+
+    @classmethod
+    def _resolve(cls, raw_source: str) -> Self:
+        raise NotImplementedError
+
+    def to_dict(self) -> dict[Any, Any]:
+        return {"tags": self._tags}
+
+    @classmethod
+    def from_dict(cls, source_dict: dict[Any, Any]) -> Self:
+        return cls(
+            tags=source_dict.get("tags"),
+        )

mlflow/data/dataset.py

@@ -0,0 +1,123 @@
+import json
+from abc import abstractmethod
+from typing import Any, Optional
+
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.entities import Dataset as DatasetEntity
+
+
+class Dataset:
+    """
+    Represents a dataset for use with MLflow Tracking, including the name, digest (hash),
+    schema, and profile of the dataset as well as source information (e.g. the S3 bucket or
+    managed Delta table from which the dataset was derived). Most datasets expose features
+    and targets for training and evaluation as well.
+    """
+
+    def __init__(
+        self, source: DatasetSource, name: Optional[str] = None, digest: Optional[str] = None
+    ):
+        """
+        Base constructor for a dataset. All subclasses must call this constructor.
+        """
+        self._name = name
+        self._source = source
+        # Note: Subclasses should call super() once they've initialized all of
+        # the class attributes necessary for digest computation
+        self._digest = digest or self._compute_digest()
+
+    @abstractmethod
+    def _compute_digest(self) -> str:
+        """Computes a digest for the dataset. Called if the user doesn't supply
+        a digest when constructing the dataset.
+
+        Returns:
+            A string digest for the dataset. We recommend a maximum digest length
+            of 10 characters with an ideal length of 8 characters.
+
+        """
+
+    def to_dict(self) -> dict[str, str]:
+        """Create config dictionary for the dataset.
+
+        Subclasses should override this method to provide additional fields in the config dict,
+        e.g., schema, profile, etc.
+
+        Returns a string dictionary containing the following fields: name, digest, source, source
+        type.
+        """
+        return {
+            "name": self.name,
+            "digest": self.digest,
+            "source": self.source.to_json(),
+            "source_type": self.source._get_source_type(),
+        }
+
+    def to_json(self) -> str:
+        """
+        Obtains a JSON string representation of the :py:class:`Dataset
+        <mlflow.data.dataset.Dataset>`.
+
+        Returns:
+            A JSON string representation of the :py:class:`Dataset <mlflow.data.dataset.Dataset>`.
+        """
+
+        return json.dumps(self.to_dict())
+
+    @property
+    def name(self) -> str:
+        """
+        The name of the dataset, e.g. ``"iris_data"``, ``"myschema.mycatalog.mytable@v1"``, etc.
+        """
+        if self._name is not None:
+            return self._name
+        else:
+            return "dataset"
+
+    @property
+    def digest(self) -> str:
+        """
+        A unique hash or fingerprint of the dataset, e.g. ``"498c7496"``.
+        """
+        return self._digest
+
+    @property
+    def source(self) -> DatasetSource:
+        """
+        Information about the dataset's source, represented as an instance of
+        :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`. For example, this
+        may be the S3 location or the name of the managed Delta Table from which the dataset
+        was derived.
+        """
+        return self._source
+
+    @property
+    @abstractmethod
+    def profile(self) -> Optional[Any]:
+        """
+        Optional summary statistics for the dataset, such as the number of rows in a table, the
+        mean / median / std of each table column, etc.
+        """
+
+    @property
+    @abstractmethod
+    def schema(self) -> Optional[Any]:
+        """
+        Optional dataset schema, such as an instance of :py:class:`mlflow.types.Schema` representing
+        the features and targets of the dataset.
+        """
+
+    def _to_mlflow_entity(self) -> DatasetEntity:
+        """
+        Returns:
+            A `mlflow.entities.Dataset` instance representing the dataset.
+        """
+        dataset_dict = self.to_dict()
+        return DatasetEntity(
+            name=dataset_dict["name"],
+            digest=dataset_dict["digest"],
+            source_type=dataset_dict["source_type"],
+            source=dataset_dict["source"],
+            schema=dataset_dict.get("schema"),
+            profile=dataset_dict.get("profile"),
+        )

mlflow/data/dataset_registry.py

@@ -0,0 +1,168 @@
+import inspect
+import warnings
+from contextlib import suppress
+from typing import Callable, Optional
+
+import mlflow.data
+from mlflow.data.dataset import Dataset
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.utils.plugins import get_entry_points
+
+
+class DatasetRegistry:
+    def __init__(self):
+        self.constructors = {}
+
+    def register_constructor(
+        self,
+        constructor_fn: Callable[[Optional[str], Optional[str]], Dataset],
+        constructor_name: Optional[str] = None,
+    ) -> str:
+        """Registers a dataset constructor.
+
+        Args:
+            constructor_fn: A function that accepts at least the following
+                inputs and returns an instance of a subclass of
+                :py:class:`mlflow.data.dataset.Dataset`:
+
+                - name: Optional. A string dataset name
+                - digest: Optional. A string dataset digest.
+
+            constructor_name: The name of the constructor, e.g.
+                "from_spark". The name must begin with the
+                string "from_" or "load_". If unspecified, the `__name__`
+                attribute of the `constructor_fn` is used instead and must
+                begin with the string "from_" or "load_".
+
+        Returns:
+            The name of the registered constructor, e.g. "from_pandas" or "load_delta".
+        """
+        if constructor_name is None:
+            constructor_name = constructor_fn.__name__
+        DatasetRegistry._validate_constructor(constructor_fn, constructor_name)
+        self.constructors[constructor_name] = constructor_fn
+        return constructor_name
+
+    def register_entrypoints(self):
+        """
+        Registers dataset sources defined as Python entrypoints. For reference, see
+        https://mlflow.org/docs/latest/plugins.html#defining-a-plugin.
+        """
+        for entrypoint in get_entry_points("mlflow.dataset_constructor"):
+            try:
+                self.register_constructor(
+                    constructor_fn=entrypoint.load(), constructor_name=entrypoint.name
+                )
+            except Exception as exc:
+                warnings.warn(
+                    f"Failure attempting to register dataset constructor"
+                    f' "{entrypoint.name}": {exc}.',
+                    stacklevel=2,
+                )
+
+    @staticmethod
+    def _validate_constructor(
+        constructor_fn: Callable[[Optional[str], Optional[str]], Dataset],
+        constructor_name: str,
+    ):
+        if not constructor_name.startswith("load_") and not constructor_name.startswith("from_"):
+            raise MlflowException(
+                f"Invalid dataset constructor name: {constructor_name}."
+                f" Constructor name must start with 'load_' or 'from_'.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        signature = inspect.signature(constructor_fn)
+        parameters = signature.parameters
+        for expected_kwarg in ["name", "digest"]:
+            if expected_kwarg not in parameters or parameters[expected_kwarg].kind not in [
+                inspect.Parameter.KEYWORD_ONLY,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            ]:
+                raise MlflowException(
+                    f"Invalid dataset constructor function: {constructor_fn.__name__}. Function"
+                    f" must define an optional parameter named '{expected_kwarg}'.",
+                    INVALID_PARAMETER_VALUE,
+                )
+
+        if not issubclass(signature.return_annotation, Dataset):
+            raise MlflowException(
+                f"Invalid dataset constructor function: {constructor_fn.__name__}. Function must"
+                f" have a return type annotation that is a subclass of"
+                f" :py:class:`mlflow.data.dataset.Dataset`.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+
+def register_constructor(
+    constructor_fn: Callable[[Optional[str], Optional[str]], Dataset],
+    constructor_name: Optional[str] = None,
+) -> str:
+    """Registers a dataset constructor.
+
+    Args:
+        constructor_fn: A function that accepts at least the following
+            inputs and returns an instance of a subclass of
+            :py:class:`mlflow.data.dataset.Dataset`:
+
+            - name: Optional. A string dataset name
+            - digest: Optional. A string dataset digest.
+
+        constructor_name: The name of the constructor, e.g.
+            "from_spark". The name must begin with the
+            string "from_" or "load_". If unspecified, the `__name__`
+            attribute of the `constructor_fn` is used instead and must
+            begin with the string "from_" or "load_".
+
+    Returns:
+        The name of the registered constructor, e.g. "from_pandas" or "load_delta".
+
+    """
+    registered_constructor_name = _dataset_registry.register_constructor(
+        constructor_fn=constructor_fn, constructor_name=constructor_name
+    )
+    setattr(mlflow.data, registered_constructor_name, constructor_fn)
+    mlflow.data.__all__.append(registered_constructor_name)
+    return registered_constructor_name
+
+
+def get_registered_constructors() -> dict[str, Callable[[Optional[str], Optional[str]], Dataset]]:
+    """Obtains the registered dataset constructors.
+
+    Returns:
+        A dictionary mapping constructor names to constructor functions.
+
+    """
+    return _dataset_registry.constructors
+
+
+_dataset_registry = DatasetRegistry()
+_dataset_registry.register_entrypoints()
+
+# use contextlib suppress to ignore import errors
+with suppress(ImportError):
+    from mlflow.data.pandas_dataset import from_pandas
+
+    _dataset_registry.register_constructor(from_pandas)
+with suppress(ImportError):
+    from mlflow.data.numpy_dataset import from_numpy
+
+    _dataset_registry.register_constructor(from_numpy)
+with suppress(ImportError):
+    from mlflow.data.huggingface_dataset import from_huggingface
+
+    _dataset_registry.register_constructor(from_huggingface)
+with suppress(ImportError):
+    from mlflow.data.tensorflow_dataset import from_tensorflow
+
+    _dataset_registry.register_constructor(from_tensorflow)
+with suppress(ImportError):
+    from mlflow.data.spark_dataset import from_spark, load_delta
+
+    _dataset_registry.register_constructor(load_delta)
+    _dataset_registry.register_constructor(from_spark)
+with suppress(ImportError):
+    from mlflow.data.polars_dataset import from_polars
+
+    _dataset_registry.register_constructor(from_polars)

mlflow/data/dataset_source.py

@@ -0,0 +1,110 @@
+import json
+from abc import abstractmethod
+from typing import Any
+
+
+class DatasetSource:
+    """
+    Represents the source of a dataset used in MLflow Tracking, providing information such as
+    cloud storage location, delta table name / version, etc.
+    """
+
+    @staticmethod
+    @abstractmethod
+    def _get_source_type() -> str:
+        """Obtains a string representing the source type of the dataset.
+
+        Returns:
+            A string representing the source type of the dataset, e.g. "s3", "delta_table", ...
+
+        """
+
+    @abstractmethod
+    def load(self) -> Any:
+        """
+        Loads files / objects referred to by the DatasetSource. For example, depending on the type
+        of :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`, this may download
+        source CSV files from S3 to the local filesystem, load a source Delta Table as a Spark
+        DataFrame, etc.
+
+        Returns:
+            The downloaded source, e.g. a local filesystem path, a Spark DataFrame, etc.
+
+        """
+
+    @staticmethod
+    @abstractmethod
+    def _can_resolve(raw_source: Any) -> bool:
+        """Determines whether this type of DatasetSource can be resolved from a specified raw source
+        object. For example, an S3DatasetSource can be resolved from an S3 URI like
+        "s3://mybucket/path/to/iris/data" but not from an Azure Blob Storage URI like
+        "wasbs:/account@host.blob.core.windows.net".
+
+        Args:
+            raw_source: The raw source, e.g. a string like "s3://mybucket/path/to/iris/data".
+
+        Returns:
+            True if this DatasetSource can resolve the raw source, False otherwise.
+
+        """
+
+    @classmethod
+    @abstractmethod
+    def _resolve(cls, raw_source: Any) -> "DatasetSource":
+        """Constructs an instance of the DatasetSource from a raw source object, such as a
+        string URI like "s3://mybucket/path/to/iris/data" or a delta table identifier
+        like "my.delta.table@2".
+
+        Args:
+            raw_source: The raw source, e.g. a string like "s3://mybucket/path/to/iris/data".
+
+        Returns:
+            A DatasetSource instance derived from the raw_source.
+
+        """
+
+    @abstractmethod
+    def to_dict(self) -> dict[str, Any]:
+        """Obtains a JSON-compatible dictionary representation of the DatasetSource.
+
+        Returns:
+            A JSON-compatible dictionary representation of the DatasetSource.
+
+        """
+
+    def to_json(self) -> str:
+        """
+        Obtains a JSON string representation of the
+        :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`.
+
+        Returns:
+            A JSON string representation of the
+            :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`.
+        """
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    @abstractmethod
+    def from_dict(cls, source_dict: dict[Any, Any]) -> "DatasetSource":
+        """Constructs an instance of the DatasetSource from a dictionary representation.
+
+        Args:
+            source_dict: A dictionary representation of the DatasetSource.
+
+        Returns:
+            A DatasetSource instance.
+
+        """
+
+    @classmethod
+    def from_json(cls, source_json: str) -> "DatasetSource":
+        """Constructs an instance of the DatasetSource from a JSON string representation.
+
+        Args:
+            source_json: A JSON string representation of the DatasetSource.
+
+        Returns:
+            A DatasetSource instance.
+
+        """
+        return cls.from_dict(json.loads(source_json))

mlflow/data/dataset_source_registry.py

@@ -0,0 +1,219 @@
+import warnings
+from typing import Any, Optional
+
+from mlflow.data.artifact_dataset_sources import register_artifact_dataset_sources
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.data.http_dataset_source import HTTPDatasetSource
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST
+from mlflow.utils.plugins import get_entry_points
+
+
+class DatasetSourceRegistry:
+    def __init__(self):
+        self.sources = []
+
+    def register(self, source: DatasetSource):
+        """Registers a DatasetSource for use with MLflow Tracking.
+
+        Args:
+            source: The DatasetSource to register.
+        """
+        self.sources.append(source)
+
+    def register_entrypoints(self):
+        """
+        Registers dataset sources defined as Python entrypoints. For reference, see
+        https://mlflow.org/docs/latest/plugins.html#defining-a-plugin.
+        """
+        for entrypoint in get_entry_points("mlflow.dataset_source"):
+            try:
+                self.register(entrypoint.load())
+            except (AttributeError, ImportError) as exc:
+                warnings.warn(
+                    "Failure attempting to register dataset constructor"
+                    + f' "{entrypoint}": {exc}',
+                    stacklevel=2,
+                )
+
+    def resolve(
+        self, raw_source: Any, candidate_sources: Optional[list[DatasetSource]] = None
+    ) -> DatasetSource:
+        """Resolves a raw source object, such as a string URI, to a DatasetSource for use with
+        MLflow Tracking.
+
+        Args:
+            raw_source: The raw source, e.g. a string like "s3://mybucket/path/to/iris/data" or a
+                HuggingFace :py:class:`datasets.Dataset` object.
+            candidate_sources: A list of DatasetSource classes to consider as potential sources
+                when resolving the raw source. Subclasses of the specified candidate sources are
+                also considered. If unspecified, all registered sources are considered.
+
+        Raises:
+            MlflowException: If no DatasetSource class can resolve the raw source.
+
+        Returns:
+            The resolved DatasetSource.
+        """
+        matching_sources = []
+        for source in self.sources:
+            if candidate_sources and not any(
+                issubclass(source, candidate_src) for candidate_src in candidate_sources
+            ):
+                continue
+            try:
+                if source._can_resolve(raw_source):
+                    matching_sources.append(source)
+            except Exception as e:
+                warnings.warn(
+                    f"Failed to determine whether {source.__name__} can resolve source"
+                    f" information for '{raw_source}'. Exception: {e}",
+                    stacklevel=2,
+                )
+                continue
+
+        if len(matching_sources) > 1:
+            source_class_names_str = ", ".join([source.__name__ for source in matching_sources])
+            warnings.warn(
+                f"The specified dataset source can be interpreted in multiple ways:"
+                f" {source_class_names_str}. MLflow will assume that this is a"
+                f" {matching_sources[-1].__name__} source.",
+                stacklevel=2,
+            )
+
+        for matching_source in reversed(matching_sources):
+            try:
+                return matching_source._resolve(raw_source)
+            except Exception as e:
+                warnings.warn(
+                    f"Encountered an unexpected error while using {matching_source.__name__} to"
+                    f" resolve source information for '{raw_source}'. Exception: {e}",
+                    stacklevel=2,
+                )
+                continue
+
+        raise MlflowException(
+            f"Could not find a source information resolver for the specified"
+            f" dataset source: {raw_source}.",
+            RESOURCE_DOES_NOT_EXIST,
+        )
+
+    def get_source_from_json(self, source_json: str, source_type: str) -> DatasetSource:
+        """Parses and returns a DatasetSource object from its JSON representation.
+
+        Args:
+            source_json: The JSON representation of the DatasetSource.
+            source_type: The string type of the DatasetSource, which indicates how to parse the
+                source JSON.
+        """
+        for source in reversed(self.sources):
+            if source._get_source_type() == source_type:
+                return source.from_json(source_json)
+
+        raise MlflowException(
+            f"Could not parse dataset source from JSON due to unrecognized"
+            f" source type: {source_type}.",
+            RESOURCE_DOES_NOT_EXIST,
+        )
+
+
+def register_dataset_source(source: DatasetSource):
+    """Registers a DatasetSource for use with MLflow Tracking.
+
+    Args:
+        source: The DatasetSource to register.
+    """
+    _dataset_source_registry.register(source)
+
+
+def resolve_dataset_source(
+    raw_source: Any, candidate_sources: Optional[list[DatasetSource]] = None
+) -> DatasetSource:
+    """Resolves a raw source object, such as a string URI, to a DatasetSource for use with
+    MLflow Tracking.
+
+    Args:
+        raw_source: The raw source, e.g. a string like "s3://mybucket/path/to/iris/data" or a
+            HuggingFace :py:class:`datasets.Dataset` object.
+        candidate_sources: A list of DatasetSource classes to consider as potential sources
+            when resolving the raw source. Subclasses of the specified candidate
+            sources are also considered. If unspecified, all registered sources
+            are considered.
+
+    Raises:
+        MlflowException: If no DatasetSource class can resolve the raw source.
+
+    Returns:
+        The resolved DatasetSource.
+    """
+    return _dataset_source_registry.resolve(
+        raw_source=raw_source, candidate_sources=candidate_sources
+    )
+
+
+def get_dataset_source_from_json(source_json: str, source_type: str) -> DatasetSource:
+    """Parses and returns a DatasetSource object from its JSON representation.
+
+    Args:
+        source_json: The JSON representation of the DatasetSource.
+        source_type: The string type of the DatasetSource, which indicates how to parse the
+            source JSON.
+    """
+    return _dataset_source_registry.get_source_from_json(
+        source_json=source_json, source_type=source_type
+    )
+
+
+def get_registered_sources() -> list[DatasetSource]:
+    """Obtains the registered dataset sources.
+
+    Returns:
+        A list of registered dataset sources.
+
+    """
+    return _dataset_source_registry.sources
+
+
+# NB: The ordering here is important. The last dataset source to be registered takes precedence
+# when resolving dataset information for a raw source (e.g. a string like "s3://mybucket/my/path").
+# Dataset sources derived from artifact repositories are the most generic / provide the most
+# general information about dataset source locations, so they are registered first. More specific
+# source information is provided by specialized dataset platform sources like
+# HuggingFaceDatasetSource, so these sources are registered next. Finally, externally-defined
+# dataset sources are registered last because externally-defined behavior should take precedence
+# over any internally-defined generic behavior
+_dataset_source_registry = DatasetSourceRegistry()
+register_artifact_dataset_sources()
+_dataset_source_registry.register(HTTPDatasetSource)
+_dataset_source_registry.register_entrypoints()
+
+try:
+    from mlflow.data.huggingface_dataset_source import HuggingFaceDatasetSource
+
+    _dataset_source_registry.register(HuggingFaceDatasetSource)
+except ImportError:
+    pass
+try:
+    from mlflow.data.spark_dataset_source import SparkDatasetSource
+
+    _dataset_source_registry.register(SparkDatasetSource)
+except ImportError:
+    pass
+try:
+    from mlflow.data.delta_dataset_source import DeltaDatasetSource
+
+    _dataset_source_registry.register(DeltaDatasetSource)
+except ImportError:
+    pass
+try:
+    from mlflow.data.code_dataset_source import CodeDatasetSource
+
+    _dataset_source_registry.register(CodeDatasetSource)
+except ImportError:
+    pass
+try:
+    from mlflow.data.uc_volume_dataset_source import UCVolumeDatasetSource
+
+    _dataset_source_registry.register(UCVolumeDatasetSource)
+except ImportError:
+    pass