wehub-resource-sync
138 lines
5.6 KiB
Python
138 lines
5.6 KiB
Python
import json
|
|
import os
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from mlflow.data.dataset_source_registry import get_dataset_source_from_json, resolve_dataset_source
|
|
from mlflow.data.filesystem_dataset_source import FileSystemDatasetSource
|
|
from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("source_uri", "source_type", "source_class_name"),
|
|
[
|
|
("/tmp/path/to/my/local/file.txt", "local", "LocalArtifactDatasetSource"),
|
|
("file:///tmp/path/to/my/local/directory", "local", "LocalArtifactDatasetSource"),
|
|
("s3://mybucket/path/to/my/file.txt", "s3", "S3ArtifactDatasetSource"),
|
|
("gs://mybucket/path/to/my/dir", "gs", "GCSArtifactDatasetSource"),
|
|
("wasbs://user@host.blob.core.windows.net/dir", "wasbs", "AzureBlobArtifactDatasetSource"),
|
|
("ftp://mysite.com/path/to/my/file.txt", "ftp", "FTPArtifactDatasetSource"),
|
|
("sftp://mysite.com/path/to/my/dir", "sftp", "SFTPArtifactDatasetSource"),
|
|
("hdfs://host_name:8020/hdfs/path/to/my/file.txt", "hdfs", "HdfsArtifactDatasetSource"),
|
|
("viewfs://host_name:8020/path/to/my/dir", "viewfs", "HdfsArtifactDatasetSource"),
|
|
],
|
|
)
|
|
def test_expected_artifact_dataset_sources_are_registered_and_resolvable(
|
|
source_uri, source_type, source_class_name
|
|
):
|
|
dataset_source = resolve_dataset_source(source_uri)
|
|
assert isinstance(dataset_source, FileSystemDatasetSource)
|
|
assert dataset_source._get_source_type() == source_type
|
|
assert type(dataset_source).__name__ == source_class_name
|
|
assert type(dataset_source).__qualname__ == source_class_name
|
|
assert dataset_source.uri == source_uri
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("source_uri", "source_type"),
|
|
[
|
|
("/tmp/path/to/my/local/file.txt", "local"),
|
|
("file:///tmp/path/to/my/local/directory", "local"),
|
|
("s3://mybucket/path/to/my/file.txt", "s3"),
|
|
("gs://mybucket/path/to/my/dir", "gs"),
|
|
("wasbs://user@host.blob.core.windows.net/dir", "wasbs"),
|
|
("ftp://mysite.com/path/to/my/file.txt", "ftp"),
|
|
("sftp://mysite.com/path/to/my/dir", "sftp"),
|
|
("hdfs://host_name:8020/hdfs/path/to/my/file.txt", "hdfs"),
|
|
("viewfs://host_name:8020/path/to/my/dir", "viewfs"),
|
|
],
|
|
)
|
|
def test_to_and_from_json(source_uri, source_type):
|
|
dataset_source = resolve_dataset_source(source_uri)
|
|
assert dataset_source._get_source_type() == source_type
|
|
source_json = dataset_source.to_json()
|
|
|
|
parsed_source_json = json.loads(source_json)
|
|
assert parsed_source_json["uri"] == source_uri
|
|
|
|
reloaded_source = get_dataset_source_from_json(
|
|
source_json, source_type=dataset_source._get_source_type()
|
|
)
|
|
assert isinstance(reloaded_source, FileSystemDatasetSource)
|
|
assert type(dataset_source) == type(reloaded_source)
|
|
assert reloaded_source.uri == dataset_source.uri
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("source_uri", "source_type"),
|
|
[
|
|
("/tmp/path/to/my/local/file.txt", "local"),
|
|
("file:///tmp/path/to/my/local/directory", "local"),
|
|
("s3://mybucket/path/to/my/file.txt", "s3"),
|
|
("gs://mybucket/path/to/my/dir", "gs"),
|
|
("wasbs://user@host.blob.core.windows.net/dir", "wasbs"),
|
|
("ftp://mysite.com/path/to/my/file.txt", "ftp"),
|
|
("sftp://mysite.com/path/to/my/dir", "sftp"),
|
|
("hdfs://host_name:8020/hdfs/path/to/my/file.txt", "hdfs"),
|
|
("viewfs://host_name:8020/path/to/my/dir", "viewfs"),
|
|
],
|
|
)
|
|
def test_load_makes_expected_mlflow_artifacts_download_call(source_uri, source_type, tmp_path):
|
|
dataset_source = resolve_dataset_source(source_uri)
|
|
assert dataset_source._get_source_type() == source_type
|
|
|
|
with mock.patch("mlflow.artifacts.download_artifacts") as download_imp_mock:
|
|
dataset_source.load()
|
|
download_imp_mock.assert_called_once_with(artifact_uri=source_uri, dst_path=None)
|
|
|
|
with mock.patch("mlflow.artifacts.download_artifacts") as download_imp_mock:
|
|
dataset_source.load(dst_path=str(tmp_path))
|
|
download_imp_mock.assert_called_once_with(artifact_uri=source_uri, dst_path=str(tmp_path))
|
|
|
|
|
|
@pytest.mark.parametrize("dst_path", [None, "dst"])
|
|
def test_local_load(dst_path, tmp_path):
|
|
if dst_path is not None:
|
|
dst_path = str(tmp_path / dst_path)
|
|
|
|
# Test string file paths
|
|
file_path = str(tmp_path / "myfile.txt")
|
|
with open(file_path, "w") as f:
|
|
f.write("text")
|
|
|
|
file_dataset_source = resolve_dataset_source(file_path)
|
|
assert file_dataset_source._get_source_type() == "local"
|
|
assert file_dataset_source.load(dst_path=dst_path) == dst_path or file_path
|
|
with open(file_path) as f:
|
|
assert f.read() == "text"
|
|
|
|
# Test directory paths with pathlib.Path
|
|
dir_path = tmp_path / "mydir"
|
|
os.makedirs(dir_path)
|
|
|
|
dir_dataset_source = resolve_dataset_source(dir_path)
|
|
assert file_dataset_source._get_source_type() == "local"
|
|
assert dir_dataset_source.load() == dst_path or str(dir_path)
|
|
|
|
|
|
@pytest.mark.parametrize("dst_path", [None, "dst"])
|
|
def test_s3_load(mock_s3_bucket, dst_path, tmp_path):
|
|
if dst_path is not None:
|
|
dst_path = str(tmp_path / dst_path)
|
|
|
|
file_path = str(tmp_path / "myfile.txt")
|
|
with open(file_path, "w") as f:
|
|
f.write("text")
|
|
|
|
S3ArtifactRepository(f"s3://{mock_s3_bucket}").log_artifact(file_path)
|
|
|
|
s3_source_uri = f"s3://{mock_s3_bucket}/myfile.txt"
|
|
s3_dataset_source = resolve_dataset_source(s3_source_uri)
|
|
assert s3_dataset_source._get_source_type() == "s3"
|
|
downloaded_source = s3_dataset_source.load(dst_path=dst_path)
|
|
if dst_path is not None:
|
|
assert downloaded_source == os.path.join(dst_path, "myfile.txt")
|
|
with open(downloaded_source) as f:
|
|
assert f.read() == "text"
|