Files
mlflow--mlflow/tests/data/test_tensorflow_dataset.py
2026-07-13 13:22:34 +08:00

396 lines
15 KiB
Python

import json
import numpy as np
import pytest
import tensorflow as tf
import mlflow.data
from mlflow.data.code_dataset_source import CodeDatasetSource
from mlflow.data.evaluation_dataset import EvaluationDataset
from mlflow.data.pyfunc_dataset_mixin import PyFuncInputsOutputs
from mlflow.data.schema import TensorDatasetSchema
from mlflow.data.tensorflow_dataset import TensorFlowDataset
from mlflow.exceptions import MlflowException
from mlflow.types.utils import _infer_schema
from tests.resources.data.dataset_source import SampleDatasetSource
def test_dataset_construction_validates_features_and_targets():
x = np.random.sample((100, 2))
tf_dataset = tf.data.Dataset.from_tensors(x)
tf_tensor = tf.convert_to_tensor(x)
with pytest.raises(
MlflowException,
match="features' must be an instance of tf.data.Dataset or a TensorFlow Tensor.*NoneType",
):
mlflow.data.from_tensorflow(features=None)
with pytest.raises(
MlflowException,
match="features' must be an instance of tf.data.Dataset or a TensorFlow Tensor.*str",
):
mlflow.data.from_tensorflow(features="foo")
with pytest.raises(
MlflowException,
match="features' must be an instance of tf.data.Dataset or a TensorFlow Tensor.*str",
):
mlflow.data.from_tensorflow(features="foo", targets=tf_tensor)
mlflow.data.from_tensorflow(features=tf_tensor, targets=tf_tensor)
mlflow.data.from_tensorflow(features=tf_tensor, targets=None)
with pytest.raises(
MlflowException,
match=(
"If 'features' is a TensorFlow Tensor, then 'targets' must also be a TensorFlow"
" Tensor.*str"
),
):
mlflow.data.from_tensorflow(features=tf_tensor, targets="foo")
with pytest.raises(
MlflowException,
match=(
"If 'features' is a TensorFlow Tensor, then 'targets' must also be a TensorFlow"
" Tensor.*Dataset"
),
):
mlflow.data.from_tensorflow(features=tf_tensor, targets=tf_dataset)
mlflow.data.from_tensorflow(features=tf_dataset, targets=tf_dataset)
mlflow.data.from_tensorflow(features=tf_dataset, targets=None)
with pytest.raises(
MlflowException,
match=(
"If 'features' is an instance of tf.data.Dataset, then 'targets' must also be an"
" instance of tf.data.Dataset.*str"
),
):
mlflow.data.from_tensorflow(features=tf_dataset, targets="foo")
with pytest.raises(
MlflowException,
match=(
"If 'features' is an instance of tf.data.Dataset, then 'targets' must also be an"
" instance of tf.data.Dataset.*Tensor"
),
):
mlflow.data.from_tensorflow(features=tf_dataset, targets=tf_tensor)
def test_conversion_to_json():
source_uri = "test:/my/test/uri"
x = np.random.sample((100, 2))
tf_dataset = tf.data.Dataset.from_tensors(x)
source = SampleDatasetSource._resolve(source_uri)
dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
dataset_json = dataset.to_json()
parsed_json = json.loads(dataset_json)
assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
assert parsed_json["name"] == dataset.name
assert parsed_json["digest"] == dataset.digest
assert parsed_json["source"] == dataset.source.to_json()
assert parsed_json["source_type"] == dataset.source._get_source_type()
assert parsed_json["profile"] == json.dumps(dataset.profile)
parsed_schema = json.loads(parsed_json["schema"])
assert TensorDatasetSchema.from_dict(parsed_schema) == dataset.schema
@pytest.mark.parametrize(
("features", "targets"),
[
(
tf.data.Dataset.from_tensors({
"a": np.random.sample((100, 2)),
"b": np.random.sample((100, 4)),
}),
tf.data.Dataset.from_tensors({
"c": np.random.sample((100, 1)),
"d": np.random.sample((100,)),
}),
),
(
tf.data.Dataset.from_tensors((
np.random.sample((100, 2)),
np.random.sample((100, 4)),
)),
tf.data.Dataset.from_tensors((
np.random.sample((100, 1)),
np.random.sample((100,)),
)),
),
(
tf.data.Dataset.from_tensors((
np.random.sample((100, 2)),
np.random.sample((100, 4)),
)),
tf.data.Dataset.from_tensors({
"c": np.random.sample((100, 1)),
"d": np.random.sample((100,)),
}),
),
(
tf.data.Dataset.from_tensors((
np.random.sample((100, 2)),
np.random.sample((100, 4)),
)),
None,
),
],
)
def test_conversion_to_json_with_multi_tensor_datasets(features, targets):
source_uri = "test:/my/test/uri"
source = SampleDatasetSource._resolve(source_uri)
dataset = TensorFlowDataset(features=features, targets=targets, source=source, name="testname")
dataset_json = dataset.to_json()
parsed_json = json.loads(dataset_json)
assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
assert parsed_json["name"] == dataset.name
assert parsed_json["digest"] == dataset.digest
assert parsed_json["source"] == dataset.source.to_json()
assert parsed_json["source_type"] == dataset.source._get_source_type()
assert parsed_json["profile"] == json.dumps(dataset.profile)
parsed_schema = json.loads(parsed_json["schema"])
assert TensorDatasetSchema.from_dict(parsed_schema) == dataset.schema
def test_schema_and_profile_with_multi_tensor_tuple_datasets():
features_dataset = tf.data.Dataset.from_tensors((
np.random.sample((100, 2)),
np.random.sample((100, 4)),
))
targets_dataset = tf.data.Dataset.from_tensors((
np.random.sample((100, 1)),
np.random.sample((100,)),
))
source_uri = "test:/my/test/uri"
source = SampleDatasetSource._resolve(source_uri)
dataset = TensorFlowDataset(
features=features_dataset, targets=targets_dataset, source=source, name="testname"
)
assert dataset.schema.features == _infer_schema({
"0": np.random.sample((100, 2)),
"1": np.random.sample((100, 4)),
})
assert dataset.schema.targets == _infer_schema({
"0": np.random.sample((100, 1)),
"1": np.random.sample((100,)),
})
assert dataset.profile == {
"features_cardinality": 1,
"targets_cardinality": 1,
}
assert dataset.profile == {
"features_cardinality": features_dataset.cardinality().numpy(),
"targets_cardinality": targets_dataset.cardinality().numpy(),
}
def test_schema_and_profile_with_multi_tensor_dict_datasets():
features_dataset = tf.data.Dataset.from_tensors({
"a": np.random.sample((100, 2)),
"b": np.random.sample((100, 4)),
})
targets_dataset = tf.data.Dataset.from_tensors({
"c": np.random.sample((100, 1)),
"d": np.random.sample((100,)),
})
source_uri = "test:/my/test/uri"
source = SampleDatasetSource._resolve(source_uri)
dataset = TensorFlowDataset(
features=features_dataset, targets=targets_dataset, source=source, name="testname"
)
assert dataset.schema.features == _infer_schema({
"a": np.random.sample((100, 2)),
"b": np.random.sample((100, 4)),
})
assert dataset.schema.targets == _infer_schema({
"c": np.random.sample((100, 1)),
"d": np.random.sample((100,)),
})
assert dataset.profile == {
"features_cardinality": 1,
"targets_cardinality": 1,
}
assert dataset.profile == {
"features_cardinality": features_dataset.cardinality().numpy(),
"targets_cardinality": targets_dataset.cardinality().numpy(),
}
def test_digest_property_has_expected_value():
source_uri = "test:/my/test/uri"
x = [[1, 2, 3], [4, 5, 6]]
tf_dataset = tf.data.Dataset.from_tensors(x)
source = SampleDatasetSource._resolve(source_uri)
dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
assert dataset.digest == dataset._compute_digest()
assert dataset.digest == "666a9820"
def test_data_property_has_expected_value():
source_uri = "test:/my/test/uri"
x = [[1, 2, 3], [4, 5, 6]]
tf_dataset = tf.data.Dataset.from_tensors(x)
source = SampleDatasetSource._resolve(source_uri)
dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
assert dataset.data == tf_dataset
def test_source_property_has_expected_value():
source_uri = "test:/my/test/uri"
x = [[1, 2, 3], [4, 5, 6]]
tf_dataset = tf.data.Dataset.from_tensors(x)
source = SampleDatasetSource._resolve(source_uri)
dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
assert dataset.source == source
def test_profile_property_has_expected_value_dataset():
source_uri = "test:/my/test/uri"
x = [[1, 2, 3], [4, 5, 6]]
tf_dataset = tf.data.Dataset.from_tensors(x)
source = SampleDatasetSource._resolve(source_uri)
dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
assert dataset.profile == {
"features_cardinality": tf_dataset.cardinality().numpy(),
}
def test_profile_property_has_expected_value_tensors():
source_uri = "test:/my/test/uri"
x = [[1, 2, 3], [4, 5, 6]]
tf_tensor = tf.convert_to_tensor(x)
source = SampleDatasetSource._resolve(source_uri)
dataset = TensorFlowDataset(features=tf_tensor, source=source, name="testname")
assert dataset.profile == {
"features_cardinality": tf.size(tf_tensor).numpy(),
}
def test_to_pyfunc():
source_uri = "test:/my/test/uri"
x = np.random.sample((100, 2))
tf_dataset = tf.data.Dataset.from_tensors(x)
source = SampleDatasetSource._resolve(source_uri)
dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
assert isinstance(dataset.to_pyfunc(), PyFuncInputsOutputs)
def test_to_evaluation_dataset():
source_uri = "test:/my/test/uri"
x = np.random.sample((2, 2))
y = np.random.sample((2, 1))
x_tensors = tf.convert_to_tensor(x)
y_tensors = tf.convert_to_tensor(y)
source = SampleDatasetSource._resolve(source_uri)
dataset = TensorFlowDataset(
features=x_tensors, source=source, targets=y_tensors, name="testname"
)
evaluation_dataset = dataset.to_evaluation_dataset()
assert isinstance(evaluation_dataset, EvaluationDataset)
assert np.array_equal(evaluation_dataset.features_data, dataset.data.numpy())
assert np.array_equal(evaluation_dataset.labels_data, dataset.targets.numpy())
def test_to_evaluation_dataset_with_tensorflow_dataset_data():
source_uri = "test:/my/test/uri"
x = np.random.sample((2, 2))
y = np.random.sample((2, 1))
x_tf_data = tf.data.Dataset.from_tensors(x)
y_tf_data = tf.data.Dataset.from_tensors(y)
source = SampleDatasetSource._resolve(source_uri)
dataset = TensorFlowDataset(
features=x_tf_data, source=source, targets=y_tf_data, name="testname"
)
with pytest.raises(
MlflowException, match="Data must be a Tensor to convert to an EvaluationDataset"
):
dataset.to_evaluation_dataset()
def test_from_tensorflow_dataset_constructs_expected_dataset():
x = np.random.sample((100, 2))
tf_dataset = tf.data.Dataset.from_tensors(x)
mlflow_ds = mlflow.data.from_tensorflow(tf_dataset, source="my_source")
assert isinstance(mlflow_ds, TensorFlowDataset)
assert mlflow_ds.data == tf_dataset
assert mlflow_ds.schema == TensorDatasetSchema(
features=_infer_schema(next(tf_dataset.as_numpy_iterator()))
)
assert mlflow_ds.profile == {
"features_cardinality": tf_dataset.cardinality().numpy(),
}
def test_from_tensorflow_dataset_with_targets_constructs_expected_dataset():
x = np.random.sample((100, 2))
y = np.random.sample((100, 1))
tf_dataset_x = tf.data.Dataset.from_tensors(x)
tf_dataset_y = tf.data.Dataset.from_tensors(y)
mlflow_ds = mlflow.data.from_tensorflow(tf_dataset_x, source="my_source", targets=tf_dataset_y)
assert isinstance(mlflow_ds, TensorFlowDataset)
assert mlflow_ds.data == tf_dataset_x
assert mlflow_ds.targets == tf_dataset_y
assert mlflow_ds.schema == TensorDatasetSchema(
features=_infer_schema(next(tf_dataset_x.as_numpy_iterator())),
targets=_infer_schema(next(tf_dataset_y.as_numpy_iterator())),
)
assert mlflow_ds.profile == {
"features_cardinality": tf_dataset_x.cardinality().numpy(),
"targets_cardinality": tf_dataset_y.cardinality().numpy(),
}
def test_from_tensorflow_tensor_constructs_expected_dataset():
x = np.random.sample((100, 2))
tf_tensor = tf.convert_to_tensor(x)
mlflow_ds = mlflow.data.from_tensorflow(tf_tensor, source="my_source")
assert isinstance(mlflow_ds, TensorFlowDataset)
# compare if two tensors are equal using tensorflow utils
assert tf.reduce_all(tf.math.equal(mlflow_ds.data, tf_tensor))
assert mlflow_ds.schema == TensorDatasetSchema(features=_infer_schema(tf_tensor.numpy()))
assert mlflow_ds.profile == {
"features_cardinality": tf.size(tf_tensor).numpy(),
}
def test_from_tensorflow_tensor_with_targets_constructs_expected_dataset():
x = np.random.sample((100, 2))
y = np.random.sample((100, 1))
tf_tensor_x = tf.convert_to_tensor(x)
tf_tensor_y = tf.convert_to_tensor(y)
mlflow_ds = mlflow.data.from_tensorflow(tf_tensor_x, source="my_source", targets=tf_tensor_y)
assert isinstance(mlflow_ds, TensorFlowDataset)
assert tf.reduce_all(tf.math.equal(mlflow_ds.data, tf_tensor_x))
assert tf.reduce_all(tf.math.equal(mlflow_ds.targets, tf_tensor_y))
assert mlflow_ds.schema == TensorDatasetSchema(
features=_infer_schema(tf_tensor_x.numpy()),
targets=_infer_schema(tf_tensor_y.numpy()),
)
assert mlflow_ds.profile == {
"features_cardinality": tf.size(tf_tensor_x).numpy(),
"targets_cardinality": tf.size(tf_tensor_y).numpy(),
}
def test_from_tensorflow_no_source_specified():
x = np.random.sample((100, 2))
tf_dataset = tf.data.Dataset.from_tensors(x)
mlflow_ds = mlflow.data.from_tensorflow(tf_dataset)
assert isinstance(mlflow_ds, TensorFlowDataset)
assert isinstance(mlflow_ds.source, CodeDatasetSource)
assert "mlflow.source.name" in mlflow_ds.source.to_json()
def test_digest_computation_succeeds_with_none_element_in_numpy_iterator():
x = np.array([[0, 1], [1, 2]])
tf_dataset = tf.data.Dataset.from_tensors(x)
tf_dataset.as_numpy_iterator = lambda: [None, x]
mlflow_ds = mlflow.data.from_tensorflow(tf_dataset)
assert mlflow_ds.digest == "bc8ef018"