wehub-resource-sync
127 lines
4.1 KiB
Python
127 lines
4.1 KiB
Python
import json
|
|
from abc import abstractmethod
|
|
from typing import Any
|
|
|
|
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: str | None = None, digest: str | None = 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())
|
|
|
|
def _get_source_type(self) -> str:
|
|
"""Returns the type of the dataset's underlying source."""
|
|
|
|
return self.source._get_source_type()
|
|
|
|
@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) -> Any | None:
|
|
"""
|
|
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) -> Any | None:
|
|
"""
|
|
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"),
|
|
)
|