wehub-resource-sync
405 lines
16 KiB
Python
405 lines
16 KiB
Python
import json
|
|
import logging
|
|
from functools import cached_property
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from packaging.version import Version
|
|
|
|
from mlflow.data.dataset import Dataset
|
|
from mlflow.data.dataset_source import DatasetSource
|
|
from mlflow.data.delta_dataset_source import DeltaDatasetSource
|
|
from mlflow.data.digest_utils import get_normalized_md5_digest
|
|
from mlflow.data.evaluation_dataset import EvaluationDataset
|
|
from mlflow.data.pyfunc_dataset_mixin import PyFuncConvertibleDatasetMixin, PyFuncInputsOutputs
|
|
from mlflow.data.spark_dataset_source import SparkDatasetSource
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos.databricks_pb2 import INTERNAL_ERROR, INVALID_PARAMETER_VALUE
|
|
from mlflow.types import Schema
|
|
from mlflow.types.utils import _infer_schema
|
|
|
|
if TYPE_CHECKING:
|
|
import pyspark
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SparkDataset(Dataset, PyFuncConvertibleDatasetMixin):
|
|
"""
|
|
Represents a Spark dataset (e.g. data derived from a Spark Table / file directory or Delta
|
|
Table) for use with MLflow Tracking.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
df: "pyspark.sql.DataFrame",
|
|
source: DatasetSource,
|
|
targets: str | None = None,
|
|
name: str | None = None,
|
|
digest: str | None = None,
|
|
predictions: str | None = None,
|
|
):
|
|
if targets is not None and targets not in df.columns:
|
|
raise MlflowException(
|
|
f"The specified Spark dataset does not contain the specified targets column"
|
|
f" '{targets}'.",
|
|
INVALID_PARAMETER_VALUE,
|
|
)
|
|
if predictions is not None and predictions not in df.columns:
|
|
raise MlflowException(
|
|
f"The specified Spark dataset does not contain the specified predictions column"
|
|
f" '{predictions}'.",
|
|
INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
self._df = df
|
|
self._targets = targets
|
|
self._predictions = predictions
|
|
super().__init__(source=source, name=name, digest=digest)
|
|
|
|
def _compute_digest(self) -> str:
|
|
"""
|
|
Computes a digest for the dataset. Called if the user doesn't supply
|
|
a digest when constructing the dataset.
|
|
"""
|
|
# Retrieve a semantic hash of the DataFrame's logical plan, which is much more efficient
|
|
# and deterministic than hashing DataFrame records
|
|
import numpy as np
|
|
import pyspark
|
|
|
|
# Spark 3.1.0+ has a semanticHash() method on DataFrame
|
|
if Version(pyspark.__version__) >= Version("3.1.0"):
|
|
semantic_hash = self._df.semanticHash()
|
|
else:
|
|
semantic_hash = self._df._jdf.queryExecution().analyzed().semanticHash()
|
|
return get_normalized_md5_digest([np.int64(semantic_hash)])
|
|
|
|
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({"mlflow_colspec": 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 df(self):
|
|
"""The Spark DataFrame instance.
|
|
|
|
Returns:
|
|
The Spark DataFrame instance.
|
|
|
|
"""
|
|
return self._df
|
|
|
|
@property
|
|
def targets(self) -> str | None:
|
|
"""The name of the Spark DataFrame column containing targets (labels) for supervised
|
|
learning.
|
|
|
|
Returns:
|
|
The string name of the Spark DataFrame column containing targets.
|
|
"""
|
|
return self._targets
|
|
|
|
@property
|
|
def predictions(self) -> str | None:
|
|
"""
|
|
The name of the predictions column. May be ``None`` if no predictions column
|
|
was specified when the dataset was created.
|
|
"""
|
|
return self._predictions
|
|
|
|
@property
|
|
def source(self) -> SparkDatasetSource | DeltaDatasetSource:
|
|
"""
|
|
Spark dataset source information.
|
|
|
|
Returns:
|
|
An instance of
|
|
:py:class:`SparkDatasetSource <mlflow.data.spark_dataset_source.SparkDatasetSource>` or
|
|
:py:class:`DeltaDatasetSource <mlflow.data.delta_dataset_source.DeltaDatasetSource>`.
|
|
"""
|
|
return self._source
|
|
|
|
@property
|
|
def profile(self) -> Any | None:
|
|
"""
|
|
A profile of the dataset. May be None if no profile is available.
|
|
"""
|
|
try:
|
|
from pyspark.rdd import BoundedFloat
|
|
|
|
# Use Spark RDD countApprox to get approximate count since count() may be expensive.
|
|
# Note that we call the Scala RDD API because the PySpark API does not respect the
|
|
# specified timeout. Reference code:
|
|
# https://spark.apache.org/docs/3.4.0/api/python/_modules/pyspark/rdd.html
|
|
# #RDD.countApprox. This is confirmed to work in all Spark 3.x versions
|
|
py_rdd = self.df.rdd
|
|
drdd = py_rdd.mapPartitions(lambda it: [float(sum(1 for i in it))])
|
|
jrdd = drdd.mapPartitions(lambda it: [float(sum(it))])._to_java_object_rdd()
|
|
jdrdd = drdd.ctx._jvm.JavaDoubleRDD.fromRDD(jrdd.rdd())
|
|
timeout_millis = 5000
|
|
confidence = 0.9
|
|
approx_count_operation = jdrdd.sumApprox(timeout_millis, confidence)
|
|
approx_count_result = approx_count_operation.initialValue()
|
|
approx_count_float = BoundedFloat(
|
|
mean=approx_count_result.mean(),
|
|
confidence=approx_count_result.confidence(),
|
|
low=approx_count_result.low(),
|
|
high=approx_count_result.high(),
|
|
)
|
|
approx_count = int(approx_count_float)
|
|
if approx_count <= 0:
|
|
# An approximate count of zero likely indicates that the count timed
|
|
# out before an estimate could be made. In this case, we use the value
|
|
# "unknown" so that users don't think the dataset is empty
|
|
approx_count = "unknown"
|
|
|
|
return {
|
|
"approx_count": approx_count,
|
|
}
|
|
except Exception as e:
|
|
_logger.warning(
|
|
"Encountered an unexpected exception while computing Spark dataset profile."
|
|
" Exception: %s",
|
|
e,
|
|
)
|
|
|
|
@cached_property
|
|
def schema(self) -> Schema | None:
|
|
"""
|
|
The MLflow ColSpec schema of the Spark dataset.
|
|
"""
|
|
try:
|
|
return _infer_schema(self._df)
|
|
except Exception as e:
|
|
_logger.warning("Failed to infer schema for Spark dataset. Exception: %s", e)
|
|
return None
|
|
|
|
def to_pyfunc(self) -> PyFuncInputsOutputs:
|
|
"""
|
|
Converts the Spark DataFrame to pandas and splits the resulting
|
|
:py:class:`pandas.DataFrame` into: 1. a :py:class:`pandas.DataFrame` of features and
|
|
2. a :py:class:`pandas.Series` of targets.
|
|
|
|
To avoid overuse of driver memory, only the first 10,000 DataFrame rows are selected.
|
|
"""
|
|
df = self._df.limit(10000).toPandas()
|
|
if self._targets is not None:
|
|
if self._targets not in df.columns:
|
|
raise MlflowException(
|
|
f"Failed to convert Spark dataset to pyfunc inputs and outputs because"
|
|
f" the pandas representation of the Spark dataset does not contain the"
|
|
f" specified targets column '{self._targets}'.",
|
|
# This is an internal error because we should have validated the presence of
|
|
# the target column in the Hugging Face dataset at construction time
|
|
INTERNAL_ERROR,
|
|
)
|
|
inputs = df.drop(columns=self._targets)
|
|
outputs = df[self._targets]
|
|
return PyFuncInputsOutputs(inputs=inputs, outputs=outputs)
|
|
else:
|
|
return PyFuncInputsOutputs(inputs=df, outputs=None)
|
|
|
|
def to_evaluation_dataset(self, path=None, feature_names=None) -> EvaluationDataset:
|
|
"""
|
|
Converts the dataset to an EvaluationDataset for model evaluation. Required
|
|
for use with mlflow.evaluate().
|
|
"""
|
|
return EvaluationDataset(
|
|
data=self._df.limit(10000).toPandas(),
|
|
targets=self._targets,
|
|
path=path,
|
|
feature_names=feature_names,
|
|
predictions=self._predictions,
|
|
name=self.name,
|
|
digest=self.digest,
|
|
)
|
|
|
|
|
|
def load_delta(
|
|
path: str | None = None,
|
|
table_name: str | None = None,
|
|
version: str | None = None,
|
|
targets: str | None = None,
|
|
name: str | None = None,
|
|
digest: str | None = None,
|
|
) -> SparkDataset:
|
|
"""
|
|
Loads a :py:class:`SparkDataset <mlflow.data.spark_dataset.SparkDataset>` from a Delta table
|
|
for use with MLflow Tracking.
|
|
|
|
Args:
|
|
path: The path to the Delta table. Either ``path`` or ``table_name`` must be specified.
|
|
table_name: The name of the Delta table. Either ``path`` or ``table_name`` must be
|
|
specified.
|
|
version: The Delta table version. If not specified, the version will be inferred.
|
|
targets: Optional. The name of the Delta table column containing targets (labels) for
|
|
supervised learning.
|
|
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.
|
|
|
|
Returns:
|
|
An instance of :py:class:`SparkDataset <mlflow.data.spark_dataset.SparkDataset>`.
|
|
"""
|
|
from mlflow.data.spark_delta_utils import (
|
|
_try_get_delta_table_latest_version_from_path,
|
|
_try_get_delta_table_latest_version_from_table_name,
|
|
)
|
|
|
|
if (path, table_name).count(None) != 1:
|
|
raise MlflowException(
|
|
"Must specify exactly one of `table_name` or `path`.",
|
|
INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
if version is None:
|
|
if path is not None:
|
|
version = _try_get_delta_table_latest_version_from_path(path)
|
|
else:
|
|
version = _try_get_delta_table_latest_version_from_table_name(table_name)
|
|
|
|
if name is None and table_name is not None:
|
|
name = table_name + (f"@v{version}" if version is not None else "")
|
|
|
|
source = DeltaDatasetSource(path=path, delta_table_name=table_name, delta_table_version=version)
|
|
df = source.load()
|
|
|
|
return SparkDataset(
|
|
df=df,
|
|
source=source,
|
|
targets=targets,
|
|
name=name,
|
|
digest=digest,
|
|
)
|
|
|
|
|
|
def from_spark(
|
|
df: "pyspark.sql.DataFrame",
|
|
path: str | None = None,
|
|
table_name: str | None = None,
|
|
version: str | None = None,
|
|
sql: str | None = None,
|
|
targets: str | None = None,
|
|
name: str | None = None,
|
|
digest: str | None = None,
|
|
predictions: str | None = None,
|
|
) -> SparkDataset:
|
|
"""
|
|
Given a Spark DataFrame, constructs a
|
|
:py:class:`SparkDataset <mlflow.data.spark_dataset.SparkDataset>` object for use with
|
|
MLflow Tracking.
|
|
|
|
Args:
|
|
df: The Spark DataFrame from which to construct a SparkDataset.
|
|
path: The path of the Spark or Delta source that the DataFrame originally came from. Note
|
|
that the path does not have to match the DataFrame exactly, since the DataFrame may have
|
|
been modified by Spark operations. This is used to reload the dataset upon request via
|
|
:py:func:`SparkDataset.source.load()
|
|
<mlflow.data.spark_dataset_source.SparkDatasetSource.load>`. If none of ``path``,
|
|
``table_name``, or ``sql`` are specified, a CodeDatasetSource is used, which will source
|
|
information from the run context.
|
|
table_name: The name of the Spark or Delta table that the DataFrame originally came from.
|
|
Note that the table does not have to match the DataFrame exactly, since the DataFrame
|
|
may have been modified by Spark operations. This is used to reload the dataset upon
|
|
request via :py:func:`SparkDataset.source.load()
|
|
<mlflow.data.spark_dataset_source.SparkDatasetSource.load>`. If none of ``path``,
|
|
``table_name``, or ``sql`` are specified, a CodeDatasetSource is used, which will source
|
|
information from the run context.
|
|
version: If the DataFrame originally came from a Delta table, specifies the version of the
|
|
Delta table. This is used to reload the dataset upon request via
|
|
:py:func:`SparkDataset.source.load()
|
|
<mlflow.data.spark_dataset_source.SparkDatasetSource.load>`. ``version`` cannot be
|
|
specified if ``sql`` is specified.
|
|
sql: The Spark SQL statement that was originally used to construct the DataFrame. Note that
|
|
the Spark SQL statement does not have to match the DataFrame exactly, since the
|
|
DataFrame may have been modified by Spark operations. This is used to reload the dataset
|
|
upon request via :py:func:`SparkDataset.source.load()
|
|
<mlflow.data.spark_dataset_source.SparkDatasetSource.load>`. If none of ``path``,
|
|
``table_name``, or ``sql`` are specified, a CodeDatasetSource is used, which will source
|
|
information from the run context.
|
|
targets: Optional. The name of the Data Frame column containing targets (labels) for
|
|
supervised learning.
|
|
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.
|
|
predictions: Optional. The name of the column containing model predictions,
|
|
if the dataset contains model predictions. If specified, this column
|
|
must be present in the dataframe (``df``).
|
|
|
|
Returns:
|
|
An instance of :py:class:`SparkDataset <mlflow.data.spark_dataset.SparkDataset>`.
|
|
"""
|
|
from mlflow.data.code_dataset_source import CodeDatasetSource
|
|
from mlflow.data.spark_delta_utils import (
|
|
_is_delta_table,
|
|
_is_delta_table_path,
|
|
_try_get_delta_table_latest_version_from_path,
|
|
_try_get_delta_table_latest_version_from_table_name,
|
|
)
|
|
from mlflow.tracking.context import registry
|
|
|
|
if (path, table_name, sql).count(None) < 2:
|
|
raise MlflowException(
|
|
"Must specify at most one of `path`, `table_name`, or `sql`.",
|
|
INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
if (sql, version).count(None) == 0:
|
|
raise MlflowException(
|
|
"`version` may not be specified when `sql` is specified. `version` may only be"
|
|
" specified when `table_name` or `path` is specified.",
|
|
INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
if sql is not None:
|
|
source = SparkDatasetSource(sql=sql)
|
|
elif path is not None:
|
|
if _is_delta_table_path(path):
|
|
version = version or _try_get_delta_table_latest_version_from_path(path)
|
|
source = DeltaDatasetSource(path=path, delta_table_version=version)
|
|
elif version is None:
|
|
source = SparkDatasetSource(path=path)
|
|
else:
|
|
raise MlflowException(
|
|
f"Version '{version}' was specified, but the path '{path}' does not refer"
|
|
f" to a Delta table.",
|
|
INVALID_PARAMETER_VALUE,
|
|
)
|
|
elif table_name is not None:
|
|
if _is_delta_table(table_name):
|
|
version = version or _try_get_delta_table_latest_version_from_table_name(table_name)
|
|
source = DeltaDatasetSource(
|
|
delta_table_name=table_name,
|
|
delta_table_version=version,
|
|
)
|
|
elif version is None:
|
|
source = SparkDatasetSource(table_name=table_name)
|
|
else:
|
|
raise MlflowException(
|
|
f"Version '{version}' was specified, but could not find a Delta table with name"
|
|
f" '{table_name}'.",
|
|
INVALID_PARAMETER_VALUE,
|
|
)
|
|
else:
|
|
context_tags = registry.resolve_tags()
|
|
source = CodeDatasetSource(tags=context_tags)
|
|
|
|
return SparkDataset(
|
|
df=df,
|
|
source=source,
|
|
targets=targets,
|
|
name=name,
|
|
digest=digest,
|
|
predictions=predictions,
|
|
)
|