wehub-resource-sync
74 lines
2.4 KiB
Python
74 lines
2.4 KiB
Python
import json
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
import mlflow.data
|
|
from mlflow.exceptions import MlflowException
|
|
|
|
from tests.resources.data.dataset_source import SampleDatasetSource
|
|
|
|
|
|
def test_load(tmp_path):
|
|
assert SampleDatasetSource("test:" + str(tmp_path)).load() == str(tmp_path)
|
|
|
|
|
|
def test_conversion_to_json_and_back():
|
|
uri = "test:/my/test/uri"
|
|
source = SampleDatasetSource._resolve(uri)
|
|
source_json = source.to_json()
|
|
assert json.loads(source_json)["uri"] == uri
|
|
reloaded_source = SampleDatasetSource.from_json(source_json)
|
|
assert reloaded_source.uri == source.uri
|
|
|
|
|
|
def test_get_source_obtains_expected_file_source(tmp_path):
|
|
df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
|
|
path = tmp_path / "temp.csv"
|
|
df.to_csv(path)
|
|
pandas_ds = mlflow.data.from_pandas(df, source=path)
|
|
|
|
source1 = mlflow.data.get_source(pandas_ds)
|
|
assert json.loads(source1.to_json()) == json.loads(pandas_ds.source.to_json())
|
|
|
|
with mlflow.start_run() as r:
|
|
mlflow.log_input(pandas_ds)
|
|
|
|
run = mlflow.get_run(r.info.run_id)
|
|
|
|
ds_input = run.inputs.dataset_inputs[0]
|
|
source2 = mlflow.data.get_source(ds_input)
|
|
assert json.loads(source2.to_json()) == json.loads(pandas_ds.source.to_json())
|
|
|
|
ds_entity = run.inputs.dataset_inputs[0].dataset
|
|
source3 = mlflow.data.get_source(ds_entity)
|
|
assert json.loads(source3.to_json()) == json.loads(pandas_ds.source.to_json())
|
|
|
|
assert source1.load() == source2.load() == source3.load() == str(path)
|
|
|
|
|
|
def test_get_source_obtains_expected_code_source():
|
|
df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
|
|
pandas_ds = mlflow.data.from_pandas(df)
|
|
|
|
source1 = mlflow.data.get_source(pandas_ds)
|
|
assert json.loads(source1.to_json()) == json.loads(pandas_ds.source.to_json())
|
|
|
|
with mlflow.start_run() as r:
|
|
mlflow.log_input(pandas_ds)
|
|
|
|
run = mlflow.get_run(r.info.run_id)
|
|
|
|
ds_input = run.inputs.dataset_inputs[0]
|
|
source2 = mlflow.data.get_source(ds_input)
|
|
assert json.loads(source2.to_json()) == json.loads(pandas_ds.source.to_json())
|
|
|
|
ds_entity = run.inputs.dataset_inputs[0].dataset
|
|
source3 = mlflow.data.get_source(ds_entity)
|
|
assert json.loads(source3.to_json()) == json.loads(pandas_ds.source.to_json())
|
|
|
|
|
|
def test_get_source_throws_for_invalid_input(tmp_path):
|
|
with pytest.raises(MlflowException, match="Unrecognized dataset type.*str"):
|
|
mlflow.data.get_source(str(tmp_path))
|