wehub-resource-sync
444 lines
16 KiB
Python
444 lines
16 KiB
Python
import cgi
|
|
import os
|
|
import pathlib
|
|
import subprocess
|
|
import tempfile
|
|
from contextlib import contextmanager
|
|
from io import BytesIO
|
|
from typing import NamedTuple
|
|
|
|
import pytest
|
|
import requests
|
|
|
|
import mlflow
|
|
from mlflow import MlflowClient
|
|
from mlflow.artifacts import download_artifacts
|
|
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
|
from mlflow.utils.os import is_windows
|
|
|
|
from tests.helper_functions import LOCALHOST, get_safe_port, kill_process_tree
|
|
from tests.tracking.integration_test_utils import _await_server_up_or_die
|
|
|
|
|
|
@contextmanager
|
|
def _launch_server(host, port, backend_store_uri, default_artifact_root, artifacts_destination):
|
|
extra_cmd = [] if is_windows() else ["--gunicorn-opts", "--log-level debug"]
|
|
cmd = [
|
|
"mlflow",
|
|
"server",
|
|
"--host",
|
|
host,
|
|
"--port",
|
|
str(port),
|
|
"--backend-store-uri",
|
|
backend_store_uri,
|
|
"--default-artifact-root",
|
|
default_artifact_root,
|
|
"--artifacts-destination",
|
|
artifacts_destination,
|
|
*extra_cmd,
|
|
]
|
|
with subprocess.Popen(cmd) as process:
|
|
try:
|
|
_await_server_up_or_die(port)
|
|
yield process
|
|
finally:
|
|
kill_process_tree(process.pid)
|
|
|
|
|
|
class ArtifactsServer(NamedTuple):
|
|
backend_store_uri: str
|
|
default_artifact_root: str
|
|
artifacts_destination: str
|
|
url: str
|
|
process: subprocess.Popen
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def artifacts_server():
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
port = get_safe_port()
|
|
backend_store_uri = f"sqlite:///{os.path.join(tmpdir, 'mlruns.db')}"
|
|
artifacts_destination = os.path.join(tmpdir, "mlartifacts")
|
|
url = f"http://{LOCALHOST}:{port}"
|
|
default_artifact_root = f"{url}/api/2.0/mlflow-artifacts/artifacts"
|
|
# Initialize the database before launching the server process
|
|
s = SqlAlchemyStore(backend_store_uri, default_artifact_root)
|
|
s.engine.dispose()
|
|
with _launch_server(
|
|
LOCALHOST,
|
|
port,
|
|
backend_store_uri,
|
|
default_artifact_root,
|
|
("file:///" + artifacts_destination if is_windows() else artifacts_destination),
|
|
) as process:
|
|
yield ArtifactsServer(
|
|
backend_store_uri, default_artifact_root, artifacts_destination, url, process
|
|
)
|
|
|
|
|
|
def read_file(path):
|
|
with open(path) as f:
|
|
return f.read()
|
|
|
|
|
|
def upload_file(path, url, headers=None):
|
|
with open(path, "rb") as f:
|
|
requests.put(url, data=f, headers=headers).raise_for_status()
|
|
|
|
|
|
def download_file(url, local_path, headers=None):
|
|
with requests.get(url, stream=True, headers=headers) as r:
|
|
r.raise_for_status()
|
|
assert r.headers["X-Content-Type-Options"] == "nosniff"
|
|
assert "Content-Type" in r.headers
|
|
assert "Content-Disposition" in r.headers
|
|
with open(local_path, "wb") as f:
|
|
for chunk in r.iter_content(chunk_size=8192):
|
|
f.write(chunk)
|
|
return r
|
|
|
|
|
|
def test_mlflow_artifacts_rest_apis(artifacts_server, tmp_path):
|
|
default_artifact_root = artifacts_server.default_artifact_root
|
|
artifacts_destination = artifacts_server.artifacts_destination
|
|
|
|
# Upload artifacts
|
|
file_a = tmp_path.joinpath("a.txt")
|
|
file_a.write_text("0")
|
|
upload_file(file_a, f"{default_artifact_root}/a.txt")
|
|
assert os.path.exists(os.path.join(artifacts_destination, "a.txt"))
|
|
assert read_file(os.path.join(artifacts_destination, "a.txt")) == "0"
|
|
|
|
file_b = tmp_path.joinpath("b.txt")
|
|
file_b.write_text("1")
|
|
upload_file(file_b, f"{default_artifact_root}/dir/b.txt")
|
|
assert os.path.join(artifacts_destination, "dir", "b.txt")
|
|
assert read_file(os.path.join(artifacts_destination, "dir", "b.txt")) == "1"
|
|
|
|
# Download artifacts
|
|
local_dir = tmp_path.joinpath("folder")
|
|
local_dir.mkdir()
|
|
local_path_a = local_dir.joinpath("a.txt")
|
|
download_file(f"{default_artifact_root}/a.txt", local_path_a)
|
|
assert read_file(local_path_a) == "0"
|
|
|
|
local_path_b = local_dir.joinpath("b.txt")
|
|
download_file(f"{default_artifact_root}/dir/b.txt", local_path_b)
|
|
assert read_file(local_path_b) == "1"
|
|
|
|
# List artifacts
|
|
resp = requests.get(default_artifact_root)
|
|
assert resp.json() == {
|
|
"files": [
|
|
{"path": "a.txt", "is_dir": False, "file_size": 1},
|
|
{"path": "dir", "is_dir": True},
|
|
]
|
|
}
|
|
resp = requests.get(default_artifact_root, params={"path": "dir"})
|
|
assert resp.json() == {"files": [{"path": "b.txt", "is_dir": False, "file_size": 1}]}
|
|
|
|
|
|
def test_log_artifact(artifacts_server, tmp_path):
|
|
url = artifacts_server.url
|
|
artifacts_destination = artifacts_server.artifacts_destination
|
|
mlflow.set_tracking_uri(url)
|
|
|
|
tmp_path = tmp_path.joinpath("a.txt")
|
|
tmp_path.write_text("0")
|
|
|
|
with mlflow.start_run() as run:
|
|
mlflow.log_artifact(tmp_path)
|
|
|
|
experiment_id = "0"
|
|
run_artifact_root = os.path.join(
|
|
artifacts_destination, experiment_id, run.info.run_id, "artifacts"
|
|
)
|
|
dest_path = os.path.join(run_artifact_root, tmp_path.name)
|
|
assert os.path.exists(dest_path)
|
|
assert read_file(dest_path) == "0"
|
|
|
|
with mlflow.start_run() as run:
|
|
mlflow.log_artifact(tmp_path, artifact_path="artifact_path")
|
|
|
|
run_artifact_root = os.path.join(
|
|
artifacts_destination, experiment_id, run.info.run_id, "artifacts"
|
|
)
|
|
dest_path = os.path.join(run_artifact_root, "artifact_path", tmp_path.name)
|
|
assert os.path.exists(dest_path)
|
|
assert read_file(dest_path) == "0"
|
|
|
|
|
|
def test_log_artifacts(artifacts_server, tmp_path):
|
|
url = artifacts_server.url
|
|
mlflow.set_tracking_uri(url)
|
|
|
|
tmp_path.joinpath("a.txt").write_text("0")
|
|
d = tmp_path.joinpath("dir")
|
|
d.mkdir()
|
|
d.joinpath("b.txt").write_text("1")
|
|
|
|
with mlflow.start_run() as run:
|
|
mlflow.log_artifacts(tmp_path)
|
|
|
|
client = MlflowClient()
|
|
artifacts = [a.path for a in client.list_artifacts(run.info.run_id)]
|
|
assert sorted(artifacts) == ["a.txt", "dir"]
|
|
artifacts = [a.path for a in client.list_artifacts(run.info.run_id, "dir")]
|
|
assert artifacts == ["dir/b.txt"]
|
|
|
|
# With `artifact_path`
|
|
with mlflow.start_run() as run:
|
|
mlflow.log_artifacts(tmp_path, artifact_path="artifact_path")
|
|
|
|
artifacts = [a.path for a in client.list_artifacts(run.info.run_id)]
|
|
assert artifacts == ["artifact_path"]
|
|
artifacts = [a.path for a in client.list_artifacts(run.info.run_id, "artifact_path")]
|
|
assert sorted(artifacts) == ["artifact_path/a.txt", "artifact_path/dir"]
|
|
artifacts = [a.path for a in client.list_artifacts(run.info.run_id, "artifact_path/dir")]
|
|
assert artifacts == ["artifact_path/dir/b.txt"]
|
|
|
|
|
|
def test_list_artifacts(artifacts_server, tmp_path):
|
|
url = artifacts_server.url
|
|
mlflow.set_tracking_uri(url)
|
|
|
|
tmp_path_a = tmp_path.joinpath("a.txt")
|
|
tmp_path_a.write_text("0")
|
|
tmp_path_b = tmp_path.joinpath("b.txt")
|
|
tmp_path_b.write_text("1")
|
|
client = MlflowClient()
|
|
with mlflow.start_run() as run:
|
|
assert client.list_artifacts(run.info.run_id) == []
|
|
mlflow.log_artifact(tmp_path_a)
|
|
mlflow.log_artifact(tmp_path_b, "dir")
|
|
|
|
artifacts = [a.path for a in client.list_artifacts(run.info.run_id)]
|
|
assert sorted(artifacts) == ["a.txt", "dir"]
|
|
artifacts = [a.path for a in client.list_artifacts(run.info.run_id, "dir")]
|
|
assert artifacts == ["dir/b.txt"]
|
|
|
|
|
|
def test_download_artifacts(artifacts_server, tmp_path):
|
|
url = artifacts_server.url
|
|
mlflow.set_tracking_uri(url)
|
|
|
|
tmp_path_a = tmp_path.joinpath("a.txt")
|
|
tmp_path_a.write_text("0")
|
|
tmp_path_b = tmp_path.joinpath("b.txt")
|
|
tmp_path_b.write_text("1")
|
|
with mlflow.start_run() as run:
|
|
mlflow.log_artifact(tmp_path_a)
|
|
mlflow.log_artifact(tmp_path_b, "dir")
|
|
|
|
dest_path = download_artifacts(run_id=run.info.run_id, artifact_path="")
|
|
assert sorted(os.listdir(dest_path)) == ["a.txt", "dir"]
|
|
assert read_file(os.path.join(dest_path, "a.txt")) == "0"
|
|
dest_path = download_artifacts(run_id=run.info.run_id, artifact_path="dir")
|
|
assert os.listdir(dest_path) == ["b.txt"]
|
|
assert read_file(os.path.join(dest_path, "b.txt")) == "1"
|
|
|
|
|
|
def is_github_actions():
|
|
return "GITHUB_ACTIONS" in os.environ
|
|
|
|
|
|
@pytest.mark.skipif(is_windows(), reason="This example doesn't work on Windows")
|
|
def test_mlflow_artifacts_example(tmp_path):
|
|
root = pathlib.Path(mlflow.__file__).parents[1]
|
|
# On GitHub Actions, remove generated images to save disk space
|
|
rmi_option = "--rmi all" if is_github_actions() else ""
|
|
cmd = f"""
|
|
err=0
|
|
trap 'err=1' ERR
|
|
./build.sh
|
|
docker compose run -v ${{PWD}}/example.py:/app/example.py client python example.py
|
|
docker compose logs
|
|
docker compose down {rmi_option} --volumes --remove-orphans
|
|
test $err = 0
|
|
"""
|
|
script_path = tmp_path.joinpath("test.sh")
|
|
script_path.write_text(cmd)
|
|
subprocess.run(
|
|
["bash", script_path],
|
|
check=True,
|
|
cwd=os.path.join(root, "examples", "mlflow_artifacts"),
|
|
)
|
|
|
|
|
|
def test_rest_tracking_api_list_artifacts_with_proxied_artifacts(artifacts_server, tmp_path):
|
|
def list_artifacts_via_rest_api(url, run_id, path=None):
|
|
if path:
|
|
resp = requests.get(url, params={"run_id": run_id, "path": path})
|
|
else:
|
|
resp = requests.get(url, params={"run_id": run_id})
|
|
resp.raise_for_status()
|
|
return resp.json()
|
|
|
|
url = artifacts_server.url
|
|
mlflow.set_tracking_uri(url)
|
|
api = f"{url}/api/2.0/mlflow/artifacts/list"
|
|
|
|
tmp_path_a = tmp_path.joinpath("a.txt")
|
|
tmp_path_a.write_text("0")
|
|
tmp_path_b = tmp_path.joinpath("b.txt")
|
|
tmp_path_b.write_text("1")
|
|
mlflow.set_experiment("rest_list_api_test")
|
|
with mlflow.start_run() as run:
|
|
mlflow.log_artifact(tmp_path_a)
|
|
mlflow.log_artifact(tmp_path_b, "dir")
|
|
|
|
list_artifacts_response = list_artifacts_via_rest_api(url=api, run_id=run.info.run_id)
|
|
assert list_artifacts_response.get("files") == [
|
|
{"path": "a.txt", "is_dir": False, "file_size": 1},
|
|
{"path": "dir", "is_dir": True},
|
|
]
|
|
assert list_artifacts_response.get("root_uri") == run.info.artifact_uri
|
|
|
|
nested_list_artifacts_response = list_artifacts_via_rest_api(
|
|
url=api, run_id=run.info.run_id, path="dir"
|
|
)
|
|
assert nested_list_artifacts_response.get("files") == [
|
|
{"path": "dir/b.txt", "is_dir": False, "file_size": 1},
|
|
]
|
|
assert list_artifacts_response.get("root_uri") == run.info.artifact_uri
|
|
|
|
|
|
def test_rest_get_artifact_api_proxied_with_artifacts(artifacts_server, tmp_path):
|
|
url = artifacts_server.url
|
|
mlflow.set_tracking_uri(url)
|
|
tmp_path_a = tmp_path.joinpath("a.txt")
|
|
tmp_path_a.write_text("abcdefg")
|
|
|
|
mlflow.set_experiment("rest_get_artifact_api_test")
|
|
with mlflow.start_run() as run:
|
|
mlflow.log_artifact(tmp_path_a)
|
|
|
|
get_artifact_response = requests.get(
|
|
url=f"{url}/get-artifact", params={"run_id": run.info.run_id, "path": "a.txt"}
|
|
)
|
|
get_artifact_response.raise_for_status()
|
|
assert get_artifact_response.text == "abcdefg"
|
|
|
|
|
|
def test_rest_get_model_version_artifact_api_proxied_artifact_root(artifacts_server):
|
|
url = artifacts_server.url
|
|
artifact_file = pathlib.Path(artifacts_server.artifacts_destination, "a.txt")
|
|
artifact_file.parent.mkdir(exist_ok=True, parents=True)
|
|
artifact_file.write_text("abcdefg")
|
|
|
|
name = "GetModelVersionTest"
|
|
mlflow_client = MlflowClient(artifacts_server.backend_store_uri)
|
|
mlflow_client.create_registered_model(name)
|
|
# An artifact root with scheme http, https, or mlflow-artifacts is a proxied artifact root
|
|
mlflow_client.create_model_version(name, "mlflow-artifacts:", 1)
|
|
|
|
get_model_version_artifact_response = requests.get(
|
|
url=f"{url}/model-versions/get-artifact",
|
|
params={"name": name, "version": "1", "path": "a.txt"},
|
|
)
|
|
get_model_version_artifact_response.raise_for_status()
|
|
assert get_model_version_artifact_response.text == "abcdefg"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("filename", "expected_mime_type"),
|
|
[
|
|
("a.txt", "text/plain"),
|
|
("b.pkl", "application/octet-stream"),
|
|
("c.png", "image/png"),
|
|
("d.pdf", "application/pdf"),
|
|
("MLmodel", "text/plain"),
|
|
("mlproject", "text/plain"),
|
|
],
|
|
)
|
|
def test_mime_type_for_download_artifacts_api(
|
|
artifacts_server, tmp_path, filename, expected_mime_type
|
|
):
|
|
default_artifact_root = artifacts_server.default_artifact_root
|
|
url = artifacts_server.url
|
|
test_file = tmp_path.joinpath(filename)
|
|
test_file.touch()
|
|
upload_file(test_file, f"{default_artifact_root}/dir/{filename}")
|
|
download_response = download_file(f"{default_artifact_root}/dir/{filename}", test_file)
|
|
|
|
_, params = cgi.parse_header(download_response.headers["Content-Disposition"])
|
|
assert params["filename"] == filename
|
|
assert download_response.headers["Content-Type"] == expected_mime_type
|
|
|
|
mlflow.set_tracking_uri(url)
|
|
with mlflow.start_run() as run:
|
|
mlflow.log_artifact(test_file)
|
|
artifact_response = requests.get(
|
|
url=f"{url}/get-artifact", params={"run_id": run.info.run_id, "path": filename}
|
|
)
|
|
artifact_response.raise_for_status()
|
|
_, params = cgi.parse_header(artifact_response.headers["Content-Disposition"])
|
|
assert params["filename"] == filename
|
|
assert artifact_response.headers["Content-Type"] == expected_mime_type
|
|
assert artifact_response.headers["X-Content-Type-Options"] == "nosniff"
|
|
|
|
|
|
def test_rest_get_artifact_api_log_image(artifacts_server):
|
|
url = artifacts_server.url
|
|
mlflow.set_tracking_uri(url)
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
image = np.random.randint(0, 256, size=(100, 100, 3), dtype=np.uint8)
|
|
|
|
with mlflow.start_run() as run:
|
|
mlflow.log_image(image, key="dog", step=20, timestamp=100, synchronous=True)
|
|
|
|
artifact_list_response = requests.get(
|
|
url=f"{url}/ajax-api/2.0/mlflow/artifacts/list",
|
|
params={"path": "images", "run_id": run.info.run_id},
|
|
)
|
|
artifact_list_response.raise_for_status()
|
|
|
|
for file in artifact_list_response.json()["files"]:
|
|
path = file["path"]
|
|
get_artifact_response = requests.get(
|
|
url=f"{url}/get-artifact", params={"run_id": run.info.run_id, "path": path}
|
|
)
|
|
get_artifact_response.raise_for_status()
|
|
assert (
|
|
"attachment; filename=dog+step+20+timestamp+100"
|
|
in get_artifact_response.headers["Content-Disposition"]
|
|
)
|
|
if path.endswith("png"):
|
|
loaded_image = np.asarray(
|
|
Image.open(BytesIO(get_artifact_response.content)), dtype=np.uint8
|
|
)
|
|
np.testing.assert_array_equal(loaded_image, image)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("filename", "requested_mime_type", "responded_mime_type"),
|
|
[
|
|
("b.pkl", "text/html", "application/octet-stream"),
|
|
("c.png", "text/html", "image/png"),
|
|
("d.pdf", "text/html", "application/pdf"),
|
|
],
|
|
)
|
|
def test_server_overrides_requested_mime_type(
|
|
artifacts_server, tmp_path, filename, requested_mime_type, responded_mime_type
|
|
):
|
|
default_artifact_root = artifacts_server.default_artifact_root
|
|
test_file = tmp_path.joinpath(filename)
|
|
test_file.touch()
|
|
upload_file(
|
|
test_file,
|
|
f"{default_artifact_root}/dir/{filename}",
|
|
)
|
|
download_response = download_file(
|
|
f"{default_artifact_root}/dir/{filename}",
|
|
test_file,
|
|
headers={"Accept": requested_mime_type},
|
|
)
|
|
|
|
_, params = cgi.parse_header(download_response.headers["Content-Disposition"])
|
|
assert params["filename"] == filename
|
|
assert download_response.headers["Content-Type"] == responded_mime_type
|