wehub-resource-sync
4075 lines
149 KiB
Python
4075 lines
149 KiB
Python
import pytest
|
|
|
|
pytestmark = pytest.mark.skip(reason="FileStore is no longer supported")
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import posixpath
|
|
import random
|
|
import re
|
|
import shutil
|
|
import time
|
|
import uuid
|
|
from copy import deepcopy
|
|
from pathlib import Path
|
|
from typing import NamedTuple
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.entities import (
|
|
AssessmentSource,
|
|
AssessmentSourceType,
|
|
Dataset,
|
|
DatasetInput,
|
|
Expectation,
|
|
ExperimentTag,
|
|
Feedback,
|
|
InputTag,
|
|
LifecycleStage,
|
|
Metric,
|
|
Param,
|
|
RunData,
|
|
RunStatus,
|
|
RunTag,
|
|
TraceInfo,
|
|
TraceLocation,
|
|
TraceState,
|
|
ViewType,
|
|
_DatasetSummary,
|
|
)
|
|
from mlflow.entities.assessment import ExpectationValue, FeedbackValue
|
|
from mlflow.entities.trace_status import TraceStatus
|
|
from mlflow.exceptions import MissingConfigException, MlflowException
|
|
from mlflow.models import Model
|
|
from mlflow.protos.databricks_pb2 import (
|
|
INTERNAL_ERROR,
|
|
INVALID_PARAMETER_VALUE,
|
|
RESOURCE_DOES_NOT_EXIST,
|
|
ErrorCode,
|
|
)
|
|
from mlflow.store.entities.paged_list import PagedList
|
|
from mlflow.store.tracking import SEARCH_MAX_RESULTS_DEFAULT
|
|
from mlflow.store.tracking.file_store import FileStore, MissingConfigException
|
|
from mlflow.tracing.constant import (
|
|
MAX_CHARS_IN_TRACE_INFO_TAGS_VALUE,
|
|
TRACE_SCHEMA_VERSION_KEY,
|
|
TraceMetadataKey,
|
|
TraceTagKey,
|
|
)
|
|
from mlflow.tracking._tracking_service.utils import _use_tracking_uri
|
|
from mlflow.utils.file_utils import TempDir, path_to_local_file_uri
|
|
from mlflow.utils.mlflow_tags import (
|
|
MLFLOW_DATASET_CONTEXT,
|
|
MLFLOW_LOGGED_MODELS,
|
|
MLFLOW_RUN_NAME,
|
|
)
|
|
from mlflow.utils.name_utils import _EXPERIMENT_ID_FIXED_WIDTH, _GENERATOR_PREDICATES
|
|
from mlflow.utils.os import is_windows
|
|
from mlflow.utils.time import get_current_time_millis
|
|
from mlflow.utils.uri import append_to_uri_path
|
|
from mlflow.utils.validation import MAX_EXPERIMENT_NAME_LENGTH
|
|
from mlflow.utils.yaml_utils import read_yaml, safe_edit_yaml, write_yaml
|
|
|
|
from tests.helper_functions import random_int, random_str
|
|
|
|
FILESTORE_PACKAGE = "mlflow.store.tracking.file_store"
|
|
|
|
|
|
@pytest.fixture
|
|
def store(tmp_path):
|
|
return FileStore(str(tmp_path.joinpath("mlruns")))
|
|
|
|
|
|
@pytest.fixture
|
|
def store_and_trace_info(store):
|
|
exp_id = store.create_experiment("test")
|
|
timestamp_ms = get_current_time_millis()
|
|
return store, store.start_trace(
|
|
TraceInfo(
|
|
trace_id=f"tr-{uuid.uuid4()}",
|
|
trace_location=TraceLocation.from_experiment_id(exp_id),
|
|
request_time=timestamp_ms,
|
|
execution_duration=0,
|
|
state=TraceState.OK,
|
|
tags={},
|
|
trace_metadata={},
|
|
client_request_id=f"tr-{uuid.uuid4()}",
|
|
request_preview=None,
|
|
response_preview=None,
|
|
),
|
|
)
|
|
|
|
|
|
class TraceInfos(NamedTuple):
|
|
trace_infos: list[TraceInfo]
|
|
store: FileStore
|
|
exp_id: str
|
|
trace_ids: list[str]
|
|
timestamps: list[int]
|
|
|
|
|
|
@pytest.fixture
|
|
def generate_trace_infos(store):
|
|
exp_id = store.create_experiment("test")
|
|
timestamps = list(range(0, 100, 10))
|
|
trace_infos = []
|
|
trace_ids = []
|
|
for i, timestamp in enumerate(timestamps):
|
|
if i < 5:
|
|
state = TraceState.OK
|
|
execution_duration = 10
|
|
else:
|
|
state = TraceState.ERROR
|
|
execution_duration = 20
|
|
|
|
metadata = {TraceMetadataKey.SOURCE_RUN: f"run_{i}"} if i >= 5 else {}
|
|
|
|
trace_info = TraceInfo(
|
|
trace_id=f"tr-{uuid.uuid4()}",
|
|
trace_location=TraceLocation.from_experiment_id(exp_id),
|
|
request_time=timestamp,
|
|
execution_duration=execution_duration,
|
|
state=state,
|
|
tags={TraceTagKey.TRACE_NAME: f"trace_{i}", "test_tag": f"tag_{i}"},
|
|
trace_metadata=metadata,
|
|
)
|
|
trace_info = store.start_trace(trace_info)
|
|
trace_infos.append(trace_info)
|
|
trace_ids.append(trace_info.trace_id)
|
|
return TraceInfos(trace_infos, store, exp_id, trace_ids, timestamps)
|
|
|
|
|
|
def create_experiments(store, experiment_names):
|
|
ids = []
|
|
for name in experiment_names:
|
|
# ensure that the field `creation_time` is distinct for search ordering
|
|
time.sleep(0.001)
|
|
ids.append(store.create_experiment(name))
|
|
return ids
|
|
|
|
|
|
def test_file_store_deprecation_warning(tmp_path):
|
|
with pytest.warns(FutureWarning, match="filesystem tracking backend.*is deprecated"):
|
|
FileStore(str(tmp_path / "mlruns"))
|
|
|
|
|
|
def test_valid_root(store):
|
|
store._check_root_dir()
|
|
shutil.rmtree(store.root_directory)
|
|
with pytest.raises(Exception, match=r"does not exist"):
|
|
store._check_root_dir()
|
|
|
|
|
|
def test_attempting_to_remove_default_experiment(store):
|
|
def _is_default_in_experiments(view_type):
|
|
search_result = store.search_experiments(view_type=view_type)
|
|
ids = [experiment.experiment_id for experiment in search_result]
|
|
return FileStore.DEFAULT_EXPERIMENT_ID in ids
|
|
|
|
assert _is_default_in_experiments(ViewType.ACTIVE_ONLY)
|
|
|
|
# Ensure experiment deletion of default id raises
|
|
with pytest.raises(MlflowException, match="Cannot delete the default experiment"):
|
|
store.delete_experiment(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
|
|
|
|
def test_search_experiments_view_type(store):
|
|
experiment_names = ["a", "b"]
|
|
experiment_ids = create_experiments(store, experiment_names)
|
|
store.delete_experiment(experiment_ids[1])
|
|
|
|
experiments = store.search_experiments(view_type=ViewType.ACTIVE_ONLY)
|
|
assert [e.name for e in experiments] == ["a", "Default"]
|
|
experiments = store.search_experiments(view_type=ViewType.DELETED_ONLY)
|
|
assert [e.name for e in experiments] == ["b"]
|
|
experiments = store.search_experiments(view_type=ViewType.ALL)
|
|
assert [e.name for e in experiments] == ["b", "a", "Default"]
|
|
|
|
|
|
def test_search_experiments_filter_by_attribute(store):
|
|
experiment_names = ["a", "ab", "Abc"]
|
|
create_experiments(store, experiment_names)
|
|
|
|
experiments = store.search_experiments(filter_string="name = 'a'")
|
|
assert [e.name for e in experiments] == ["a"]
|
|
experiments = store.search_experiments(filter_string="attribute.name = 'a'")
|
|
assert [e.name for e in experiments] == ["a"]
|
|
experiments = store.search_experiments(filter_string="attribute.`name` = 'a'")
|
|
assert [e.name for e in experiments] == ["a"]
|
|
experiments = store.search_experiments(filter_string="attribute.`name` != 'a'")
|
|
assert [e.name for e in experiments] == ["Abc", "ab", "Default"]
|
|
experiments = store.search_experiments(filter_string="name LIKE 'a%'")
|
|
assert [e.name for e in experiments] == ["ab", "a"]
|
|
experiments = store.search_experiments(
|
|
filter_string="name ILIKE 'a%'", order_by=["last_update_time asc"]
|
|
)
|
|
assert [e.name for e in experiments] == ["a", "ab", "Abc"]
|
|
experiments = store.search_experiments(filter_string="name ILIKE 'a%'")
|
|
assert [e.name for e in experiments] == ["Abc", "ab", "a"]
|
|
experiments = store.search_experiments(filter_string="name ILIKE 'a%' AND name ILIKE '%b'")
|
|
assert [e.name for e in experiments] == ["ab"]
|
|
|
|
|
|
def test_search_experiments_filter_by_time_attribute(store):
|
|
# Sleep to ensure that the first experiment has a different creation_time than the default
|
|
# experiment and eliminate flakiness.
|
|
time.sleep(0.001)
|
|
time_before_create1 = get_current_time_millis()
|
|
exp_id1 = store.create_experiment("1")
|
|
exp1 = store.get_experiment(exp_id1)
|
|
time.sleep(0.001)
|
|
time_before_create2 = get_current_time_millis()
|
|
exp_id2 = store.create_experiment("2")
|
|
exp2 = store.get_experiment(exp_id2)
|
|
|
|
experiments = store.search_experiments(filter_string=f"creation_time = {exp1.creation_time}")
|
|
assert [e.experiment_id for e in experiments] == [exp_id1]
|
|
|
|
experiments = store.search_experiments(filter_string=f"creation_time != {exp1.creation_time}")
|
|
assert [e.experiment_id for e in experiments] == [
|
|
exp_id2,
|
|
store.DEFAULT_EXPERIMENT_ID,
|
|
]
|
|
|
|
experiments = store.search_experiments(filter_string=f"creation_time >= {time_before_create1}")
|
|
assert [e.experiment_id for e in experiments] == [exp_id2, exp_id1]
|
|
|
|
experiments = store.search_experiments(filter_string=f"creation_time < {time_before_create2}")
|
|
assert [e.experiment_id for e in experiments] == [
|
|
exp_id1,
|
|
store.DEFAULT_EXPERIMENT_ID,
|
|
]
|
|
|
|
now = get_current_time_millis()
|
|
experiments = store.search_experiments(filter_string=f"creation_time > {now}")
|
|
assert experiments == []
|
|
|
|
time.sleep(0.001)
|
|
time_before_rename = get_current_time_millis()
|
|
store.rename_experiment(exp_id1, "new_name")
|
|
experiments = store.search_experiments(
|
|
filter_string=f"last_update_time >= {time_before_rename}"
|
|
)
|
|
assert [e.experiment_id for e in experiments] == [exp_id1]
|
|
|
|
experiments = store.search_experiments(
|
|
filter_string=f"last_update_time <= {get_current_time_millis()}"
|
|
)
|
|
assert {e.experiment_id for e in experiments} == {
|
|
exp_id1,
|
|
exp_id2,
|
|
store.DEFAULT_EXPERIMENT_ID,
|
|
}
|
|
|
|
experiments = store.search_experiments(
|
|
filter_string=f"last_update_time = {exp2.last_update_time}"
|
|
)
|
|
assert [e.experiment_id for e in experiments] == [exp_id2]
|
|
|
|
|
|
def test_search_experiments_filter_by_attribute_and_tag(store):
|
|
store.create_experiment("exp1", tags=[ExperimentTag("a", "1"), ExperimentTag("b", "2")])
|
|
store.create_experiment("exp2", tags=[ExperimentTag("a", "3"), ExperimentTag("b", "4")])
|
|
experiments = store.search_experiments(filter_string="name ILIKE 'exp%' AND tag.a = '1'")
|
|
assert [e.name for e in experiments] == ["exp1"]
|
|
|
|
|
|
def test_search_experiments_filter_by_tag(store):
|
|
experiments = [
|
|
("exp1", [ExperimentTag("key", "value")]),
|
|
("exp2", [ExperimentTag("key", "vaLue")]),
|
|
("exp3", [ExperimentTag("k e y", "value")]),
|
|
]
|
|
for name, tags in experiments:
|
|
# sleep to enforce deterministic ordering based on last_update_time (creation_time due to
|
|
# no mutation of experiment state)
|
|
time.sleep(0.001)
|
|
store.create_experiment(name, tags=tags)
|
|
|
|
experiments = store.search_experiments(filter_string="tag.key = 'value'")
|
|
assert [e.name for e in experiments] == ["exp1"]
|
|
experiments = store.search_experiments(filter_string="tag.`k e y` = 'value'")
|
|
assert [e.name for e in experiments] == ["exp3"]
|
|
experiments = store.search_experiments(filter_string="tag.\"k e y\" = 'value'")
|
|
assert [e.name for e in experiments] == ["exp3"]
|
|
experiments = store.search_experiments(filter_string="tag.key != 'value'")
|
|
assert [e.name for e in experiments] == ["exp2"]
|
|
experiments = store.search_experiments(filter_string="tag.key LIKE 'val%'")
|
|
assert [e.name for e in experiments] == ["exp1"]
|
|
experiments = store.search_experiments(filter_string="tag.key LIKE '%Lue'")
|
|
assert [e.name for e in experiments] == ["exp2"]
|
|
experiments = store.search_experiments(filter_string="tag.key ILIKE '%alu%'")
|
|
assert [e.name for e in experiments] == ["exp2", "exp1"]
|
|
experiments = store.search_experiments(
|
|
filter_string="tag.key LIKE 'va%' AND tags.key LIKE '%Lue'"
|
|
)
|
|
assert [e.name for e in experiments] == ["exp2"]
|
|
|
|
|
|
def test_search_experiments_filter_by_tag_is_null(store):
|
|
experiments = [
|
|
("exp1", [ExperimentTag("key1", "value"), ExperimentTag("key2", "value")]),
|
|
("exp2", [ExperimentTag("key1", "value")]),
|
|
("exp3", []),
|
|
]
|
|
for name, tags in experiments:
|
|
time.sleep(0.001)
|
|
store.create_experiment(name, tags=tags)
|
|
|
|
# IS NOT NULL: experiments that have key1
|
|
results = store.search_experiments(filter_string="tag.key1 IS NOT NULL")
|
|
assert [e.name for e in results] == ["exp2", "exp1"]
|
|
|
|
# IS NULL: experiments that don't have key2 (includes Default)
|
|
results = store.search_experiments(filter_string="tag.key2 IS NULL")
|
|
assert [e.name for e in results] == ["exp3", "exp2", "Default"]
|
|
|
|
# Combined IS NOT NULL and IS NULL
|
|
results = store.search_experiments(filter_string="tag.key1 IS NOT NULL AND tag.key2 IS NULL")
|
|
assert [e.name for e in results] == ["exp2"]
|
|
|
|
# Combined with value filter
|
|
results = store.search_experiments(filter_string="tag.key1 = 'value' AND tag.key2 IS NULL")
|
|
assert [e.name for e in results] == ["exp2"]
|
|
|
|
# Error: IS NULL on attribute
|
|
with pytest.raises(MlflowException, match="IS NULL / IS NOT NULL is only supported for tags"):
|
|
store.search_experiments(filter_string="name IS NULL")
|
|
|
|
|
|
def test_search_experiments_order_by(store):
|
|
experiment_names = ["x", "y", "z"]
|
|
create_experiments(store, experiment_names)
|
|
|
|
# Test the case where an experiment does not have a creation time by simulating a time of
|
|
# `None`. This is applicable to experiments created in older versions of MLflow where the
|
|
# `creation_time` attribute did not exist
|
|
with mock.patch(
|
|
"mlflow.store.tracking.file_store.get_current_time_millis",
|
|
return_value=None,
|
|
):
|
|
store.create_experiment("n")
|
|
|
|
experiments = store.search_experiments(order_by=["name"])
|
|
assert [e.name for e in experiments] == ["Default", "n", "x", "y", "z"]
|
|
|
|
experiments = store.search_experiments(order_by=["name ASC"])
|
|
assert [e.name for e in experiments] == ["Default", "n", "x", "y", "z"]
|
|
|
|
experiments = store.search_experiments(order_by=["name DESC"])
|
|
assert [e.name for e in experiments] == ["z", "y", "x", "n", "Default"]
|
|
|
|
experiments = store.search_experiments(order_by=["creation_time DESC"])
|
|
assert [e.name for e in experiments] == ["z", "y", "x", "Default", "n"]
|
|
|
|
experiments = store.search_experiments(order_by=["creation_time ASC"])
|
|
assert [e.name for e in experiments] == ["Default", "x", "y", "z", "n"]
|
|
|
|
experiments = store.search_experiments(order_by=["name", "last_update_time asc"])
|
|
assert [e.name for e in experiments] == ["Default", "n", "x", "y", "z"]
|
|
|
|
|
|
def test_search_experiments_order_by_time_attribute(store):
|
|
# Sleep to ensure that the first experiment has a different creation_time than the default
|
|
# experiment and eliminate flakiness.
|
|
time.sleep(0.001)
|
|
exp_id1 = store.create_experiment("1")
|
|
time.sleep(0.001)
|
|
exp_id2 = store.create_experiment("2")
|
|
|
|
experiments = store.search_experiments(order_by=["creation_time"])
|
|
assert [e.experiment_id for e in experiments] == [
|
|
store.DEFAULT_EXPERIMENT_ID,
|
|
exp_id1,
|
|
exp_id2,
|
|
]
|
|
|
|
experiments = store.search_experiments(order_by=["creation_time DESC"])
|
|
assert [e.experiment_id for e in experiments] == [
|
|
exp_id2,
|
|
exp_id1,
|
|
store.DEFAULT_EXPERIMENT_ID,
|
|
]
|
|
|
|
experiments = store.search_experiments(order_by=["last_update_time"])
|
|
assert [e.experiment_id for e in experiments] == [
|
|
store.DEFAULT_EXPERIMENT_ID,
|
|
exp_id1,
|
|
exp_id2,
|
|
]
|
|
|
|
time.sleep(0.001)
|
|
store.rename_experiment(exp_id1, "new_name")
|
|
experiments = store.search_experiments(order_by=["last_update_time"])
|
|
assert [e.experiment_id for e in experiments] == [
|
|
store.DEFAULT_EXPERIMENT_ID,
|
|
exp_id2,
|
|
exp_id1,
|
|
]
|
|
|
|
|
|
def test_search_experiments_max_results(store):
|
|
experiment_names = list(map(str, range(9)))
|
|
create_experiments(store, experiment_names)
|
|
reversed_experiment_names = experiment_names[::-1]
|
|
|
|
experiments = store.search_experiments()
|
|
assert [e.name for e in experiments] == reversed_experiment_names + ["Default"]
|
|
experiments = store.search_experiments(max_results=3)
|
|
assert [e.name for e in experiments] == reversed_experiment_names[:3]
|
|
|
|
|
|
def test_search_experiments_max_results_validation(store):
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=r"Invalid value None for parameter 'max_results' supplied. "
|
|
r"It must be a positive integer",
|
|
):
|
|
store.search_experiments(max_results=None)
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=r"Invalid value 0 for parameter 'max_results' supplied. "
|
|
r"It must be a positive integer",
|
|
):
|
|
store.search_experiments(max_results=0)
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=r"Invalid value 1000000 for parameter 'max_results' supplied. "
|
|
r"It must be at most 50000",
|
|
):
|
|
store.search_experiments(max_results=1_000_000)
|
|
|
|
|
|
def test_search_experiments_pagination(store):
|
|
experiment_names = list(map(str, range(9)))
|
|
create_experiments(store, experiment_names)
|
|
reversed_experiment_names = experiment_names[::-1]
|
|
|
|
experiments = store.search_experiments(max_results=4)
|
|
assert [e.name for e in experiments] == reversed_experiment_names[:4]
|
|
assert experiments.token is not None
|
|
|
|
experiments = store.search_experiments(max_results=4, page_token=experiments.token)
|
|
assert [e.name for e in experiments] == reversed_experiment_names[4:8]
|
|
assert experiments.token is not None
|
|
|
|
experiments = store.search_experiments(max_results=4, page_token=experiments.token)
|
|
assert [e.name for e in experiments] == reversed_experiment_names[8:] + ["Default"]
|
|
assert experiments.token is None
|
|
|
|
|
|
def _verify_experiment(fs, exp_id, exp_data):
|
|
exp = fs.get_experiment(exp_id)
|
|
assert exp.experiment_id == exp_id
|
|
assert exp.name == exp_data[exp_id]["name"]
|
|
assert exp.artifact_location == exp_data[exp_id]["artifact_location"]
|
|
|
|
|
|
def _verify_logged(store, run_id, metrics, params, tags):
|
|
run = store.get_run(run_id)
|
|
all_metrics = sum((store.get_metric_history(run_id, key) for key in run.data.metrics), [])
|
|
assert len(all_metrics) == len(metrics)
|
|
logged_metrics = [(m.key, m.value, m.timestamp, m.step) for m in all_metrics]
|
|
assert set(logged_metrics) == {(m.key, m.value, m.timestamp, m.step) for m in metrics}
|
|
logged_tags = set(run.data.tags.items())
|
|
assert {(tag.key, tag.value) for tag in tags} <= logged_tags
|
|
assert len(run.data.params) == len(params)
|
|
assert set(run.data.params.items()) == {(param.key, param.value) for param in params}
|
|
|
|
|
|
def _create_root(store):
|
|
test_root = store.root_directory
|
|
experiments = [str(random_int(100, int(1e9))) for _ in range(3)]
|
|
exp_data = {}
|
|
run_data = {}
|
|
# Include default experiment
|
|
experiments.append(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
default_exp_folder = os.path.join(test_root, str(FileStore.DEFAULT_EXPERIMENT_ID))
|
|
if os.path.exists(default_exp_folder):
|
|
shutil.rmtree(default_exp_folder)
|
|
|
|
for exp in experiments:
|
|
# create experiment
|
|
exp_folder = os.path.join(test_root, str(exp))
|
|
os.makedirs(exp_folder)
|
|
current_time = get_current_time_millis()
|
|
d = {
|
|
"experiment_id": exp,
|
|
"name": random_str(),
|
|
"artifact_location": exp_folder,
|
|
"lifecycle_stage": LifecycleStage.ACTIVE,
|
|
"creation_time": current_time,
|
|
"last_update_time": current_time,
|
|
}
|
|
exp_data[exp] = d
|
|
write_yaml(exp_folder, FileStore.META_DATA_FILE_NAME, d)
|
|
# add runs
|
|
exp_data[exp]["runs"] = []
|
|
for _ in range(2):
|
|
run_id = uuid.uuid4().hex
|
|
exp_data[exp]["runs"].append(run_id)
|
|
run_folder = os.path.join(exp_folder, run_id)
|
|
os.makedirs(run_folder)
|
|
run_info = {
|
|
"run_uuid": run_id,
|
|
"run_id": run_id,
|
|
"run_name": "name",
|
|
"experiment_id": exp,
|
|
"user_id": random_str(random_int(10, 25)),
|
|
"status": random.choice(RunStatus.all_status()),
|
|
"start_time": random_int(1, 10),
|
|
"end_time": random_int(20, 30),
|
|
"deleted_time": random_int(20, 30),
|
|
"tags": [],
|
|
"artifact_uri": os.path.join(run_folder, FileStore.ARTIFACTS_FOLDER_NAME),
|
|
"lifecycle_stage": LifecycleStage.ACTIVE,
|
|
}
|
|
write_yaml(run_folder, FileStore.META_DATA_FILE_NAME, run_info)
|
|
run_data[run_id] = run_info
|
|
# tags
|
|
os.makedirs(os.path.join(run_folder, FileStore.TAGS_FOLDER_NAME))
|
|
# params
|
|
params_folder = os.path.join(run_folder, FileStore.PARAMS_FOLDER_NAME)
|
|
os.makedirs(params_folder)
|
|
params = {}
|
|
for _ in range(5):
|
|
param_name = random_str(random_int(10, 12))
|
|
param_value = random_str(random_int(10, 15))
|
|
param_file = os.path.join(params_folder, param_name)
|
|
with open(param_file, "w") as f:
|
|
f.write(param_value)
|
|
params[param_name] = param_value
|
|
run_data[run_id]["params"] = params
|
|
# metrics
|
|
metrics_folder = os.path.join(run_folder, FileStore.METRICS_FOLDER_NAME)
|
|
os.makedirs(metrics_folder)
|
|
metrics = {}
|
|
for _ in range(3):
|
|
metric_name = random_str(random_int(10, 12))
|
|
timestamp = get_current_time_millis()
|
|
metric_file = os.path.join(metrics_folder, metric_name)
|
|
values = []
|
|
for _ in range(10):
|
|
metric_value = random_int(100, 2000)
|
|
timestamp += random_int(10000, 2000000)
|
|
values.append((timestamp, metric_value))
|
|
with open(metric_file, "a") as f:
|
|
f.write(f"{timestamp} {metric_value}\n")
|
|
metrics[metric_name] = values
|
|
run_data[run_id]["metrics"] = metrics
|
|
# artifacts
|
|
os.makedirs(os.path.join(run_folder, FileStore.ARTIFACTS_FOLDER_NAME))
|
|
|
|
return experiments, exp_data, run_data
|
|
|
|
|
|
def create_test_run(store):
|
|
return store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
|
|
|
|
def test_record_logged_model(store):
|
|
run_id = create_test_run(store).info.run_id
|
|
m = Model(artifact_path="model/path", run_id=run_id, flavors={"tf": "flavor body"})
|
|
store.record_logged_model(run_id, m)
|
|
_verify_logged(
|
|
store,
|
|
run_id=run_id,
|
|
params=[],
|
|
metrics=[],
|
|
tags=[RunTag(MLFLOW_LOGGED_MODELS, json.dumps([m.get_tags_dict()]))],
|
|
)
|
|
m2 = Model(
|
|
artifact_path="some/other/path",
|
|
run_id=run_id,
|
|
flavors={"R": {"property": "value"}},
|
|
)
|
|
store.record_logged_model(run_id, m2)
|
|
_verify_logged(
|
|
store,
|
|
run_id,
|
|
params=[],
|
|
metrics=[],
|
|
tags=[
|
|
RunTag(
|
|
MLFLOW_LOGGED_MODELS,
|
|
json.dumps([m.get_tags_dict(), m2.get_tags_dict()]),
|
|
)
|
|
],
|
|
)
|
|
m3 = Model(
|
|
artifact_path="some/other/path2",
|
|
run_id=run_id,
|
|
flavors={"R2": {"property": "value"}},
|
|
)
|
|
store.record_logged_model(run_id, m3)
|
|
_verify_logged(
|
|
store,
|
|
run_id,
|
|
params=[],
|
|
metrics=[],
|
|
tags=[
|
|
RunTag(
|
|
MLFLOW_LOGGED_MODELS,
|
|
json.dumps([m.get_tags_dict(), m2.get_tags_dict(), m3.get_tags_dict()]),
|
|
)
|
|
],
|
|
)
|
|
m4 = Model(
|
|
artifact_path="some/other/path3",
|
|
run_id=run_id,
|
|
flavors={"python_function": {"config": {"a": 1}, "code": "code"}},
|
|
)
|
|
store.record_logged_model(run_id, m4)
|
|
assert all("config" not in v for v in m4.get_tags_dict().get("flavors", {}).values())
|
|
_verify_logged(
|
|
store,
|
|
run_id,
|
|
params=[],
|
|
metrics=[],
|
|
tags=[
|
|
RunTag(
|
|
MLFLOW_LOGGED_MODELS,
|
|
json.dumps([
|
|
m.get_tags_dict(),
|
|
m2.get_tags_dict(),
|
|
m3.get_tags_dict(),
|
|
m4.get_tags_dict(),
|
|
]),
|
|
)
|
|
],
|
|
)
|
|
with pytest.raises(
|
|
TypeError,
|
|
match="Argument 'mlflow_model' should be mlflow.models.Model, got '<class 'dict'>'",
|
|
):
|
|
store.record_logged_model(run_id, m.get_tags_dict())
|
|
|
|
|
|
def test_hard_delete_logged_model(store):
|
|
exp_id = store.create_experiment("exp")
|
|
model = store.create_logged_model(experiment_id=exp_id)
|
|
store.delete_logged_model(model.model_id)
|
|
model_dir = store._get_model_dir(exp_id, model.model_id)
|
|
assert os.path.exists(model_dir)
|
|
store._hard_delete_logged_model(model.model_id)
|
|
assert not os.path.exists(model_dir)
|
|
|
|
|
|
def test_get_deleted_logged_models(store):
|
|
exp_id = store.create_experiment("exp")
|
|
model = store.create_logged_model(experiment_id=exp_id)
|
|
assert store._get_deleted_logged_models() == []
|
|
store.delete_logged_model(model.model_id)
|
|
assert store._get_deleted_logged_models(older_than=1000000) == []
|
|
assert store._get_deleted_logged_models() == [model.model_id]
|
|
|
|
|
|
def test_get_experiment(store):
|
|
experiments, exp_data, _ = _create_root(store)
|
|
for exp_id in experiments:
|
|
_verify_experiment(store, exp_id, exp_data)
|
|
|
|
# test that fake experiments dont exist.
|
|
# look for random experiment ids between 8000, 15000 since created ones are (100, 2000)
|
|
for exp_id in {random_int(8000, 15000) for x in range(20)}:
|
|
with pytest.raises(Exception, match=f"Could not find experiment with ID {exp_id}"):
|
|
store.get_experiment(str(exp_id))
|
|
|
|
|
|
def test_get_experiment_int_experiment_id_backcompat(store):
|
|
_, exp_data, _ = _create_root(store)
|
|
exp_id = FileStore.DEFAULT_EXPERIMENT_ID
|
|
root_dir = os.path.join(store.root_directory, exp_id)
|
|
with safe_edit_yaml(root_dir, "meta.yaml", _experiment_id_edit_func):
|
|
_verify_experiment(store, exp_id, exp_data)
|
|
|
|
|
|
def test_get_experiment_retries_for_transient_empty_yaml_read(store):
|
|
exp_name = random_str()
|
|
exp_id = store.create_experiment(exp_name)
|
|
|
|
mock_empty_call_count = 0
|
|
|
|
def mock_read_yaml_impl(*args, **kwargs):
|
|
nonlocal mock_empty_call_count
|
|
if mock_empty_call_count < 2:
|
|
mock_empty_call_count += 1
|
|
return None
|
|
else:
|
|
return read_yaml(*args, **kwargs)
|
|
|
|
with mock.patch(
|
|
"mlflow.store.tracking.file_store.read_yaml", side_effect=mock_read_yaml_impl
|
|
) as mock_read_yaml:
|
|
fetched_experiment = store.get_experiment(exp_id)
|
|
assert fetched_experiment.experiment_id == exp_id
|
|
assert fetched_experiment.name == exp_name
|
|
assert mock_read_yaml.call_count == 3
|
|
|
|
|
|
def test_get_experiment_by_name(store):
|
|
experiments, exp_data, _ = _create_root(store)
|
|
for exp_id in experiments:
|
|
name = exp_data[exp_id]["name"]
|
|
exp = store.get_experiment_by_name(name)
|
|
assert exp.experiment_id == exp_id
|
|
assert exp.name == exp_data[exp_id]["name"]
|
|
assert exp.artifact_location == exp_data[exp_id]["artifact_location"]
|
|
|
|
# test that fake experiments dont exist.
|
|
# look up experiments with names of length 15 since created ones are of length 10
|
|
for exp_names in {random_str(15) for x in range(20)}:
|
|
exp = store.get_experiment_by_name(exp_names)
|
|
assert exp is None
|
|
|
|
exp_id = experiments[0]
|
|
store.delete_experiment(exp_id)
|
|
assert store.get_experiment_by_name(exp_data[exp_id]["name"]).experiment_id == exp_id
|
|
|
|
|
|
def test_create_additional_experiment_generates_random_fixed_length_id(store):
|
|
store._get_active_experiments = mock.Mock(return_value=[])
|
|
store._get_deleted_experiments = mock.Mock(return_value=[])
|
|
store._create_experiment_with_id = mock.Mock()
|
|
store.create_experiment(random_str())
|
|
store._create_experiment_with_id.assert_called_once()
|
|
experiment_id = store._create_experiment_with_id.call_args[0][1]
|
|
assert len(experiment_id) == _EXPERIMENT_ID_FIXED_WIDTH
|
|
|
|
|
|
def test_create_experiment(store):
|
|
# fs = FileStore(helper.test_root)
|
|
|
|
# Error cases
|
|
with pytest.raises(Exception, match="Invalid experiment name: 'None'"):
|
|
store.create_experiment(None)
|
|
with pytest.raises(Exception, match="Invalid experiment name: ''"):
|
|
store.create_experiment("")
|
|
with pytest.raises(MlflowException, match=r"'name' exceeds the maximum length"):
|
|
store.create_experiment(name="x" * (MAX_EXPERIMENT_NAME_LENGTH + 1))
|
|
name = random_str(25) # since existing experiments are 10 chars long
|
|
time_before_create = get_current_time_millis()
|
|
created_id = store.create_experiment(name)
|
|
# test that newly created experiment id is random but of a fixed length
|
|
assert len(created_id) == _EXPERIMENT_ID_FIXED_WIDTH
|
|
|
|
# get the new experiment (by id) and verify (by name)
|
|
exp1 = store.get_experiment(created_id)
|
|
assert exp1.name == name
|
|
assert exp1.artifact_location == path_to_local_file_uri(
|
|
posixpath.join(store.root_directory, created_id)
|
|
)
|
|
assert exp1.creation_time >= time_before_create
|
|
assert exp1.last_update_time == exp1.creation_time
|
|
|
|
# get the new experiment (by name) and verify (by id)
|
|
exp2 = store.get_experiment_by_name(name)
|
|
assert exp2.experiment_id == created_id
|
|
assert exp2.creation_time == exp1.creation_time
|
|
assert exp2.last_update_time == exp1.last_update_time
|
|
|
|
|
|
def test_create_experiment_with_tags_works_correctly(store):
|
|
created_id = store.create_experiment(
|
|
"heresAnExperiment",
|
|
"heresAnArtifact",
|
|
[ExperimentTag("key1", "val1"), ExperimentTag("key2", "val2")],
|
|
)
|
|
experiment = store.get_experiment(created_id)
|
|
assert len(experiment.tags) == 2
|
|
assert experiment.tags["key1"] == "val1"
|
|
assert experiment.tags["key2"] == "val2"
|
|
|
|
|
|
def test_create_duplicate_experiments(store):
|
|
experiments, exp_data, _ = _create_root(store)
|
|
for exp_id in experiments:
|
|
name = exp_data[exp_id]["name"]
|
|
with pytest.raises(Exception, match=f"Experiment '{name}' already exists"):
|
|
store.create_experiment(name)
|
|
|
|
|
|
def _extract_ids(experiments):
|
|
return [e.experiment_id for e in experiments]
|
|
|
|
|
|
def test_delete_restore_experiment(store):
|
|
experiments, _, _ = _create_root(store)
|
|
exp1_id = experiments[random_int(0, len(experiments) - 2)] # never select default experiment
|
|
exp1 = store.get_experiment(exp1_id)
|
|
|
|
# test deleting experiment
|
|
store.delete_experiment(exp1_id)
|
|
assert exp1_id not in _extract_ids(store.search_experiments(view_type=ViewType.ACTIVE_ONLY))
|
|
assert exp1_id in _extract_ids(store.search_experiments(view_type=ViewType.DELETED_ONLY))
|
|
assert exp1_id in _extract_ids(store.search_experiments(view_type=ViewType.ALL))
|
|
deleted_exp1 = store.get_experiment(exp1_id)
|
|
assert deleted_exp1.last_update_time > exp1.last_update_time
|
|
assert deleted_exp1.lifecycle_stage == LifecycleStage.DELETED
|
|
|
|
# test if setting lifecycle_stage is persisted correctly
|
|
deleted_exp1_dir = store._get_experiment_path(
|
|
experiment_id=exp1_id, view_type=ViewType.DELETED_ONLY
|
|
)
|
|
deleted_exp1_meta = FileStore._read_yaml(
|
|
root=deleted_exp1_dir, file_name=FileStore.META_DATA_FILE_NAME
|
|
)
|
|
assert deleted_exp1_meta["lifecycle_stage"] == LifecycleStage.DELETED
|
|
for run in store.search_runs(
|
|
experiment_ids=[exp1_id], filter_string="", run_view_type=ViewType.ALL
|
|
):
|
|
assert run.info.lifecycle_stage == LifecycleStage.DELETED
|
|
|
|
# test restoring experiment
|
|
store.restore_experiment(exp1_id)
|
|
assert exp1_id in _extract_ids(store.search_experiments(view_type=ViewType.ACTIVE_ONLY))
|
|
assert exp1_id not in _extract_ids(store.search_experiments(view_type=ViewType.DELETED_ONLY))
|
|
assert exp1_id in _extract_ids(store.search_experiments(view_type=ViewType.ALL))
|
|
restored1_exp1 = store.get_experiment(exp1_id)
|
|
assert restored1_exp1.experiment_id == exp1_id
|
|
assert restored1_exp1.name == exp1.name
|
|
assert restored1_exp1.last_update_time > exp1.last_update_time
|
|
assert restored1_exp1.lifecycle_stage == LifecycleStage.ACTIVE
|
|
restored2_exp1 = store.get_experiment_by_name(exp1.name)
|
|
assert restored2_exp1.experiment_id == exp1_id
|
|
assert restored2_exp1.name == exp1.name
|
|
|
|
# test if setting lifecycle_stage is persisted correctly
|
|
restored_exp1_dir = store._get_experiment_path(
|
|
experiment_id=exp1_id, view_type=ViewType.ACTIVE_ONLY
|
|
)
|
|
restored_exp1_meta = FileStore._read_yaml(
|
|
root=restored_exp1_dir, file_name=FileStore.META_DATA_FILE_NAME
|
|
)
|
|
assert restored_exp1_meta["lifecycle_stage"] == LifecycleStage.ACTIVE
|
|
for run in store.search_runs(
|
|
experiment_ids=[exp1_id], filter_string="", run_view_type=ViewType.ALL
|
|
):
|
|
assert run.info.lifecycle_stage == LifecycleStage.ACTIVE
|
|
|
|
|
|
def test_rename_experiment(store):
|
|
experiments, _, _ = _create_root(store)
|
|
exp_id = store.create_experiment("test_rename")
|
|
|
|
# Error cases
|
|
with pytest.raises(Exception, match="Invalid experiment name: 'None'"):
|
|
store.rename_experiment(exp_id, None)
|
|
# test that names of existing experiments are checked before renaming
|
|
other_exp_id = None
|
|
for exp in experiments:
|
|
if exp != exp_id:
|
|
other_exp_id = exp
|
|
break
|
|
name = store.get_experiment(other_exp_id).name
|
|
with pytest.raises(Exception, match=f"Experiment '{name}' already exists"):
|
|
store.rename_experiment(exp_id, name)
|
|
|
|
exp_name = store.get_experiment(exp_id).name
|
|
new_name = exp_name + "!!!"
|
|
assert exp_name != new_name
|
|
assert store.get_experiment(exp_id).name == exp_name
|
|
store.rename_experiment(exp_id, new_name)
|
|
assert store.get_experiment(exp_id).name == new_name
|
|
|
|
# Ensure that we cannot rename deleted experiments.
|
|
store.delete_experiment(exp_id)
|
|
with pytest.raises(
|
|
Exception, match="Cannot rename experiment in non-active lifecycle stage"
|
|
) as e:
|
|
store.rename_experiment(exp_id, exp_name)
|
|
assert "non-active lifecycle" in str(e.value)
|
|
assert store.get_experiment(exp_id).name == new_name
|
|
|
|
# Restore the experiment, and confirm that we can now rename it.
|
|
exp1 = store.get_experiment(exp_id)
|
|
time.sleep(0.01)
|
|
store.restore_experiment(exp_id)
|
|
restored_exp1 = store.get_experiment(exp_id)
|
|
assert restored_exp1.name == new_name
|
|
assert restored_exp1.last_update_time > exp1.last_update_time
|
|
|
|
exp1 = store.get_experiment(exp_id)
|
|
time.sleep(0.01)
|
|
store.rename_experiment(exp_id, exp_name)
|
|
renamed_exp1 = store.get_experiment(exp_id)
|
|
assert renamed_exp1.name == exp_name
|
|
assert renamed_exp1.last_update_time > exp1.last_update_time
|
|
|
|
|
|
def test_delete_restore_run(store):
|
|
experiments, exp_data, _ = _create_root(store)
|
|
exp_id = experiments[random_int(0, len(experiments) - 1)]
|
|
run_id = exp_data[exp_id]["runs"][0]
|
|
_, run_dir = store._find_run_root(run_id)
|
|
# Should not throw.
|
|
assert store.get_run(run_id).info.lifecycle_stage == "active"
|
|
# Verify that run deletion is idempotent by deleting twice
|
|
store.delete_run(run_id)
|
|
store.delete_run(run_id)
|
|
assert store.get_run(run_id).info.lifecycle_stage == "deleted"
|
|
meta = read_yaml(run_dir, FileStore.META_DATA_FILE_NAME)
|
|
assert "deleted_time" in meta
|
|
assert meta["deleted_time"] is not None
|
|
# Verify that run restoration is idempotent by restoring twice
|
|
store.restore_run(run_id)
|
|
store.restore_run(run_id)
|
|
assert store.get_run(run_id).info.lifecycle_stage == "active"
|
|
meta = read_yaml(run_dir, FileStore.META_DATA_FILE_NAME)
|
|
assert "deleted_time" not in meta
|
|
|
|
|
|
def test_hard_delete_run(store):
|
|
# fs = FileStore(helper.test_root)
|
|
experiments, exp_data, _ = _create_root(store)
|
|
exp_id = experiments[random_int(0, len(experiments) - 1)]
|
|
run_id = exp_data[exp_id]["runs"][0]
|
|
store._hard_delete_run(run_id)
|
|
with pytest.raises(MlflowException, match=f"Run '{run_id}' not found"):
|
|
store.get_run(run_id)
|
|
with pytest.raises(MlflowException, match=f"Run '{run_id}' not found"):
|
|
store.get_all_tags(run_id)
|
|
with pytest.raises(MlflowException, match=f"Run '{run_id}' not found"):
|
|
store.get_all_metrics(run_id)
|
|
with pytest.raises(MlflowException, match=f"Run '{run_id}' not found"):
|
|
store.get_all_params(run_id)
|
|
|
|
|
|
def test_get_deleted_runs(store):
|
|
experiments, exp_data, _ = _create_root(store)
|
|
exp_id = experiments[0]
|
|
run_id = exp_data[exp_id]["runs"][0]
|
|
store.delete_run(run_id)
|
|
deleted_runs = store._get_deleted_runs()
|
|
assert len(deleted_runs) == 1
|
|
assert deleted_runs[0] == run_id
|
|
|
|
|
|
def test_create_run_in_deleted_experiment(store):
|
|
exp_id = store.create_experiment("test")
|
|
store.delete_experiment(exp_id)
|
|
with pytest.raises(Exception, match="Could not create run under non-active experiment"):
|
|
store.create_run(exp_id, "user", 0, [], "name")
|
|
|
|
|
|
def test_create_run_returns_expected_run_data(store):
|
|
no_tags_run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
assert isinstance(no_tags_run.data, RunData)
|
|
assert len(no_tags_run.data.tags) == 1
|
|
|
|
run_name = no_tags_run.info.run_name
|
|
assert run_name.split("-")[0] in _GENERATOR_PREDICATES
|
|
|
|
run_name = no_tags_run.info.run_name
|
|
assert run_name.split("-")[0] in _GENERATOR_PREDICATES
|
|
|
|
tags_dict = {
|
|
"my_first_tag": "first",
|
|
"my-second-tag": "2nd",
|
|
}
|
|
tags_entities = [RunTag(key, value) for key, value in tags_dict.items()]
|
|
tags_run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=tags_entities,
|
|
run_name=None,
|
|
)
|
|
assert isinstance(tags_run.data, RunData)
|
|
assert tags_run.data.tags == {**tags_dict, MLFLOW_RUN_NAME: tags_run.info.run_name}
|
|
|
|
name_empty_str_run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=tags_entities,
|
|
run_name="",
|
|
)
|
|
run_name = name_empty_str_run.info.run_name
|
|
assert run_name.split("-")[0] in _GENERATOR_PREDICATES
|
|
|
|
|
|
def test_create_run_sets_name(store):
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="my name",
|
|
)
|
|
|
|
run = store.get_run(run.info.run_id)
|
|
assert run.info.run_name == "my name"
|
|
assert run.data.tags.get(MLFLOW_RUN_NAME) == "my name"
|
|
|
|
run_id = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
run_name=None,
|
|
tags=[RunTag(MLFLOW_RUN_NAME, "test")],
|
|
).info.run_id
|
|
run = store.get_run(run_id)
|
|
assert run.info.run_name == "test"
|
|
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=re.escape(
|
|
"Both 'run_name' argument and 'mlflow.runName' tag are specified, but with "
|
|
"different values (run_name='my name', mlflow.runName='test')."
|
|
),
|
|
):
|
|
store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
run_name="my name",
|
|
tags=[RunTag(MLFLOW_RUN_NAME, "test")],
|
|
)
|
|
|
|
|
|
def _experiment_id_edit_func(old_dict):
|
|
old_dict["experiment_id"] = int(old_dict["experiment_id"])
|
|
return old_dict
|
|
|
|
|
|
def _verify_run(store, run_id, run_data):
|
|
run = store.get_run(run_id)
|
|
run_info = run_data[run_id]
|
|
run_info.pop("metrics", None)
|
|
run_info.pop("params", None)
|
|
run_info.pop("tags", None)
|
|
run_info.pop("deleted_time", None)
|
|
run_info["lifecycle_stage"] = LifecycleStage.ACTIVE
|
|
run_info["status"] = RunStatus.to_string(run_info["status"])
|
|
# get a copy of run_info as we need to remove the `deleted_time`
|
|
# key without actually deleting it from self.run_data
|
|
_run_info = run_info.copy()
|
|
_run_info.pop("deleted_time", None)
|
|
_run_info.pop("run_uuid", None)
|
|
assert _run_info == dict(run.info)
|
|
|
|
|
|
def test_get_run(store):
|
|
experiments, exp_data, run_data = _create_root(store)
|
|
for exp_id in experiments:
|
|
runs = exp_data[exp_id]["runs"]
|
|
for run_id in runs:
|
|
_verify_run(store, run_id, run_data)
|
|
|
|
|
|
def test_get_run_returns_name_in_info(store):
|
|
run_id = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="my name",
|
|
).info.run_id
|
|
|
|
get_run = store.get_run(run_id)
|
|
assert get_run.info.run_name == "my name"
|
|
|
|
|
|
def test_get_run_retries_for_transient_empty_yaml_read(store):
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
|
|
mock_empty_call_count = 0
|
|
|
|
def mock_read_yaml_impl(*args, **kwargs):
|
|
nonlocal mock_empty_call_count
|
|
if mock_empty_call_count < 2:
|
|
mock_empty_call_count += 1
|
|
return None
|
|
else:
|
|
return read_yaml(*args, **kwargs)
|
|
|
|
with mock.patch(
|
|
"mlflow.store.tracking.file_store.read_yaml", side_effect=mock_read_yaml_impl
|
|
) as mock_read_yaml:
|
|
fetched_run = store.get_run(run.info.run_id)
|
|
assert fetched_run.info.run_id == run.info.run_id
|
|
assert fetched_run.info.artifact_uri == run.info.artifact_uri
|
|
assert mock_read_yaml.call_count == 3
|
|
|
|
|
|
def test_get_run_int_experiment_id_backcompat(store):
|
|
_, exp_data, run_data = _create_root(store)
|
|
exp_id = FileStore.DEFAULT_EXPERIMENT_ID
|
|
run_id = exp_data[exp_id]["runs"][0]
|
|
root_dir = os.path.join(store.root_directory, exp_id, run_id)
|
|
with safe_edit_yaml(root_dir, "meta.yaml", _experiment_id_edit_func):
|
|
_verify_run(store, run_id, run_data)
|
|
|
|
|
|
def test_update_run_renames_run(store):
|
|
run_id = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="first name",
|
|
).info.run_id
|
|
store.update_run_info(run_id, RunStatus.FINISHED, 1000, "new name")
|
|
get_run = store.get_run(run_id)
|
|
assert get_run.info.run_name == "new name"
|
|
|
|
|
|
def test_update_run_does_not_rename_run_with_none_name(store):
|
|
run_id = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="first name",
|
|
).info.run_id
|
|
store.update_run_info(run_id, RunStatus.FINISHED, 1000, None)
|
|
get_run = store.get_run(run_id)
|
|
assert get_run.info.run_name == "first name"
|
|
|
|
|
|
def test_log_metric_allows_multiple_values_at_same_step_and_run_data_uses_max_step_value(
|
|
store,
|
|
):
|
|
run_id = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="first name",
|
|
).info.run_id
|
|
|
|
metric_name = "test-metric-1"
|
|
# Check that we get the max of (step, timestamp, value) in that order
|
|
tuples_to_log = [
|
|
(0, 100, 1000),
|
|
(3, 40, 100), # larger step wins even though it has smaller value
|
|
(3, 50, 10), # larger timestamp wins even though it has smaller value
|
|
(3, 50, 20), # tiebreak by max value
|
|
(3, 50, 20), # duplicate metrics with same (step, timestamp, value) are ok
|
|
# verify that we can log steps out of order / negative steps
|
|
(-3, 900, 900),
|
|
(-1, 800, 800),
|
|
]
|
|
for step, timestamp, value in reversed(tuples_to_log):
|
|
store.log_metric(run_id, Metric(metric_name, value, timestamp, step))
|
|
|
|
metric_history = store.get_metric_history(run_id, metric_name)
|
|
logged_tuples = [(m.step, m.timestamp, m.value) for m in metric_history]
|
|
assert set(logged_tuples) == set(tuples_to_log)
|
|
|
|
run_data = store.get_run(run_id).data
|
|
run_metrics = run_data.metrics
|
|
assert len(run_metrics) == 1
|
|
assert run_metrics[metric_name] == 20
|
|
metric_obj = run_data._metric_objs[0]
|
|
assert metric_obj.key == metric_name
|
|
assert metric_obj.step == 3
|
|
assert metric_obj.timestamp == 50
|
|
assert metric_obj.value == 20
|
|
|
|
|
|
def test_log_metric_with_non_numeric_value_raises_exception(store):
|
|
run_id = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="first name",
|
|
).info.run_id
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=r"Invalid value \"string\" for parameter \'value\' supplied",
|
|
):
|
|
store.log_metric(run_id, Metric("test", "string", 0, 0))
|
|
|
|
|
|
def test_get_all_metrics(store):
|
|
experiments, exp_data, run_data = _create_root(store)
|
|
for exp_id in experiments:
|
|
runs = exp_data[exp_id]["runs"]
|
|
for run_id in runs:
|
|
run_info = run_data[run_id]
|
|
metrics = store.get_all_metrics(run_id)
|
|
metrics_dict = run_info.pop("metrics")
|
|
for metric in metrics:
|
|
expected_timestamp, expected_value = max(metrics_dict[metric.key])
|
|
assert metric.timestamp == expected_timestamp
|
|
assert metric.value == expected_value
|
|
|
|
|
|
def test_get_metric_history(store):
|
|
experiments, exp_data, run_data = _create_root(store)
|
|
for exp_id in experiments:
|
|
runs = exp_data[exp_id]["runs"]
|
|
for run_id in runs:
|
|
run_info = run_data[run_id]
|
|
metrics = run_info.pop("metrics")
|
|
for metric_name, values in metrics.items():
|
|
metric_history = store.get_metric_history(run_id, metric_name)
|
|
sorted_values = sorted(values, reverse=True)
|
|
for metric in metric_history:
|
|
timestamp, metric_value = sorted_values.pop()
|
|
assert metric.timestamp == timestamp
|
|
assert metric.key == metric_name
|
|
assert metric.value == metric_value
|
|
|
|
|
|
def test_get_metric_history_with_max_results(store):
|
|
exp_id = store.create_experiment("test_max_results")
|
|
run = store.create_run(exp_id, user_id="user", start_time=0, tags=[], run_name="test")
|
|
run_id = run.info.run_id
|
|
|
|
metric_key = "test_metric"
|
|
for i in range(5):
|
|
metric = Metric(key=metric_key, value=float(i), timestamp=1000 + i, step=i)
|
|
store.log_metric(run_id, metric)
|
|
|
|
# Test without max_results - should return all 5 metrics
|
|
all_metrics = store.get_metric_history(run_id, metric_key)
|
|
assert len(all_metrics) == 5
|
|
|
|
# Test with max_results=3 - should return only first 3 metrics
|
|
limited_metrics = store.get_metric_history(run_id, metric_key, max_results=3)
|
|
assert len(limited_metrics) == 3
|
|
|
|
all_values = [m.value for m in all_metrics]
|
|
limited_values = [m.value for m in limited_metrics]
|
|
assert limited_values == all_values[:3]
|
|
|
|
# Test with max_results=0 - should return no metrics
|
|
no_metrics = store.get_metric_history(run_id, metric_key, max_results=0)
|
|
assert len(no_metrics) == 0
|
|
|
|
# Test with max_results larger than available metrics - should return all metrics
|
|
more_metrics = store.get_metric_history(run_id, metric_key, max_results=10)
|
|
assert len(more_metrics) == 5
|
|
|
|
|
|
def test_get_metric_history_with_page_token(store):
|
|
exp_id = store.create_experiment("test_page_token")
|
|
run = store.create_run(exp_id, user_id="user", start_time=0, tags=[], run_name="test")
|
|
run_id = run.info.run_id
|
|
|
|
metric_key = "test_metric"
|
|
for i in range(10):
|
|
metric = Metric(key=metric_key, value=float(i), timestamp=1000 + i, step=i)
|
|
store.log_metric(run_id, metric)
|
|
|
|
page_size = 4
|
|
|
|
first_page = store.get_metric_history(
|
|
run_id, metric_key, max_results=page_size, page_token=None
|
|
)
|
|
assert isinstance(first_page, PagedList)
|
|
assert first_page.token is not None
|
|
assert len(first_page) == 4
|
|
|
|
second_page = store.get_metric_history(
|
|
run_id, metric_key, max_results=page_size, page_token=first_page.token
|
|
)
|
|
assert isinstance(first_page, PagedList)
|
|
assert second_page.token is not None
|
|
assert len(second_page) == 4
|
|
|
|
third_page = store.get_metric_history(
|
|
run_id, metric_key, max_results=page_size, page_token=second_page.token
|
|
)
|
|
assert isinstance(first_page, PagedList)
|
|
assert third_page.token is None
|
|
assert len(third_page) == 2
|
|
|
|
all_paginated_metrics = first_page + second_page + third_page
|
|
assert len(all_paginated_metrics) == 10
|
|
|
|
for i, metric in enumerate(all_paginated_metrics):
|
|
assert metric.value == float(i)
|
|
assert metric.step == i
|
|
assert metric.timestamp == 1000 + i
|
|
|
|
# Test with invalid page_token
|
|
with pytest.raises(MlflowException, match="Invalid page token"):
|
|
store.get_metric_history(run_id, metric_key, page_token="invalid_token")
|
|
|
|
# Test pagination without max_results (should return all in one page)
|
|
result = store.get_metric_history(run_id, metric_key, page_token=None)
|
|
assert len(result) == 10
|
|
assert result.token is None
|
|
|
|
|
|
def _search(
|
|
fs,
|
|
experiment_id,
|
|
filter_str=None,
|
|
run_view_type=ViewType.ALL,
|
|
max_results=SEARCH_MAX_RESULTS_DEFAULT,
|
|
):
|
|
return [
|
|
r.info.run_id
|
|
for r in fs.search_runs([experiment_id], filter_str, run_view_type, max_results)
|
|
]
|
|
|
|
|
|
def test_search_runs(store):
|
|
# replace with test with code is implemented
|
|
experiments, _, _ = _create_root(store)
|
|
# Expect 2 runs for each experiment
|
|
assert len(_search(store, experiments[0], run_view_type=ViewType.ACTIVE_ONLY)) == 2
|
|
assert len(_search(store, experiments[0])) == 2
|
|
assert len(_search(store, experiments[0], run_view_type=ViewType.DELETED_ONLY)) == 0
|
|
|
|
|
|
def test_search_tags(store):
|
|
experiments, _, _ = _create_root(store)
|
|
experiment_id = experiments[0]
|
|
r1 = store.create_run(experiment_id, "user", 0, [], "name").info.run_id
|
|
r2 = store.create_run(experiment_id, "user", 0, [], "name").info.run_id
|
|
|
|
store.set_tag(r1, RunTag("generic_tag", "p_val"))
|
|
store.set_tag(r2, RunTag("generic_tag", "p_val"))
|
|
|
|
store.set_tag(r1, RunTag("generic_2", "some value"))
|
|
store.set_tag(r2, RunTag("generic_2", "another value"))
|
|
|
|
store.set_tag(r1, RunTag("p_a", "abc"))
|
|
store.set_tag(r2, RunTag("p_b", "ABC"))
|
|
|
|
# test search returns both runs
|
|
assert sorted(
|
|
[r1, r2],
|
|
) == sorted(_search(store, experiment_id, filter_str="tags.generic_tag = 'p_val'"))
|
|
# test search returns appropriate run (same key different values per run)
|
|
assert _search(store, experiment_id, filter_str="tags.generic_2 = 'some value'") == [r1]
|
|
assert _search(store, experiment_id, filter_str="tags.generic_2='another value'") == [r2]
|
|
assert _search(store, experiment_id, filter_str="tags.generic_tag = 'wrong_val'") == []
|
|
assert _search(store, experiment_id, filter_str="tags.generic_tag != 'p_val'") == []
|
|
assert sorted([r1, r2]) == sorted(
|
|
_search(store, experiment_id, filter_str="tags.generic_tag != 'wrong_val'"),
|
|
)
|
|
assert sorted([r1, r2]) == sorted(
|
|
_search(store, experiment_id, filter_str="tags.generic_2 != 'wrong_val'"),
|
|
)
|
|
assert _search(store, experiment_id, filter_str="tags.p_a = 'abc'") == [r1]
|
|
assert _search(store, experiment_id, filter_str="tags.p_b = 'ABC'") == [r2]
|
|
|
|
assert _search(store, experiment_id, filter_str="tags.generic_2 LIKE '%other%'") == [r2]
|
|
assert _search(store, experiment_id, filter_str="tags.generic_2 LIKE 'other%'") == []
|
|
assert _search(store, experiment_id, filter_str="tags.generic_2 LIKE '%other'") == []
|
|
assert _search(store, experiment_id, filter_str="tags.generic_2 ILIKE '%OTHER%'") == [r2]
|
|
|
|
|
|
def test_search_with_max_results(store):
|
|
exp = store.create_experiment("search_with_max_results")
|
|
|
|
runs = [store.create_run(exp, "user", r, [], "name").info.run_id for r in range(10)]
|
|
runs.reverse()
|
|
|
|
assert runs[:10] == _search(store, exp)
|
|
for n in [0, 1, 2, 4, 8, 10, 20, 50, 100, 500, 1000, 1200, 2000]:
|
|
assert runs[: min(1200, n)] == _search(store, exp, max_results=n)
|
|
|
|
with pytest.raises(
|
|
MlflowException, match="Invalid value for request parameter max_results. It "
|
|
):
|
|
_search(store, exp, None, max_results=int(1e10))
|
|
|
|
|
|
def test_search_with_deterministic_max_results(store):
|
|
exp = store.create_experiment("test_search_with_deterministic_max_results")
|
|
|
|
# Create 10 runs with the same start_time.
|
|
# Sort based on run_id
|
|
runs = sorted([store.create_run(exp, "user", 1000, [], "name").info.run_id for r in range(10)])
|
|
for n in [0, 1, 2, 4, 8, 10, 20]:
|
|
assert runs[: min(10, n)] == _search(store, exp, max_results=n)
|
|
|
|
|
|
def test_search_runs_pagination(store):
|
|
exp = store.create_experiment("test_search_runs_pagination")
|
|
# test returned token behavior
|
|
runs = sorted([store.create_run(exp, "user", 1000, [], "name").info.run_id for r in range(10)])
|
|
result = store.search_runs([exp], None, ViewType.ALL, max_results=4)
|
|
assert [r.info.run_id for r in result] == runs[0:4]
|
|
assert result.token is not None
|
|
result = store.search_runs([exp], None, ViewType.ALL, max_results=4, page_token=result.token)
|
|
assert [r.info.run_id for r in result] == runs[4:8]
|
|
assert result.token is not None
|
|
result = store.search_runs([exp], None, ViewType.ALL, max_results=4, page_token=result.token)
|
|
assert [r.info.run_id for r in result] == runs[8:]
|
|
assert result.token is None
|
|
|
|
|
|
def test_search_runs_run_name(store):
|
|
exp_id = store.create_experiment("test_search_runs_pagination")
|
|
run1 = store.create_run(exp_id, user_id="user", start_time=1000, tags=[], run_name="run_name1")
|
|
run2 = store.create_run(exp_id, user_id="user", start_time=1000, tags=[], run_name="run_name2")
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="attributes.run_name = 'run_name1'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert [r.info.run_id for r in result] == [run1.info.run_id]
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="tags.`mlflow.runName` = 'run_name2'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert [r.info.run_id for r in result] == [run2.info.run_id]
|
|
|
|
store.update_run_info(
|
|
run1.info.run_id,
|
|
RunStatus.FINISHED,
|
|
end_time=run1.info.end_time,
|
|
run_name="new_run_name1",
|
|
)
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="attributes.run_name = 'new_run_name1'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert [r.info.run_id for r in result] == [run1.info.run_id]
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="attributes.`run name` = 'new_run_name1'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert [r.info.run_id for r in result] == [run1.info.run_id]
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="attributes.`Run name` = 'new_run_name1'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert [r.info.run_id for r in result] == [run1.info.run_id]
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="attributes.`Run Name` = 'new_run_name1'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert [r.info.run_id for r in result] == [run1.info.run_id]
|
|
|
|
# TODO: Test attribute-based search after set_tag
|
|
|
|
# Test run name filter works for runs logged in MLflow <= 1.29.0
|
|
run_meta_path = Path(store.root_directory, exp_id, run1.info.run_id, "meta.yaml")
|
|
without_run_name = run_meta_path.read_text().replace("run_name: new_run_name1\n", "")
|
|
run_meta_path.write_text(without_run_name)
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="attributes.run_name = 'new_run_name1'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert [r.info.run_id for r in result] == [run1.info.run_id]
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="tags.`mlflow.runName` = 'new_run_name1'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert [r.info.run_id for r in result] == [run1.info.run_id]
|
|
|
|
|
|
def test_search_runs_run_id(store):
|
|
exp_id = store.create_experiment("test_search_runs_run_id")
|
|
# Set start_time to ensure the search result is deterministic
|
|
run1 = store.create_run(exp_id, user_id="user", start_time=1, tags=[], run_name="1")
|
|
run2 = store.create_run(exp_id, user_id="user", start_time=2, tags=[], run_name="2")
|
|
run_id1 = run1.info.run_id
|
|
run_id2 = run2.info.run_id
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string=f"attributes.run_id = '{run_id1}'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert [r.info.run_id for r in result] == [run_id1]
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string=f"attributes.run_id != '{run_id1}'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert [r.info.run_id for r in result] == [run_id2]
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string=f"attributes.run_id IN ('{run_id1}')",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert [r.info.run_id for r in result] == [run_id1]
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string=f"attributes.run_id NOT IN ('{run_id1}')",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert [r.info.run_id for r in result] == [run_id2]
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string=f"run_name = '{run1.info.run_name}' AND run_id IN ('{run_id1}')",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert [r.info.run_id for r in result] == [run_id1]
|
|
|
|
for filter_string in [
|
|
f"attributes.run_id IN ('{run_id1}','{run_id2}')",
|
|
f"attributes.run_id IN ('{run_id1}', '{run_id2}')",
|
|
f"attributes.run_id IN ('{run_id1}', '{run_id2}')",
|
|
]:
|
|
result = store.search_runs(
|
|
[exp_id], filter_string=filter_string, run_view_type=ViewType.ACTIVE_ONLY
|
|
)
|
|
assert [r.info.run_id for r in result] == [run_id2, run_id1]
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string=f"attributes.run_id NOT IN ('{run_id1}', '{run_id2}')",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert result == []
|
|
|
|
|
|
def test_search_runs_start_time_alias(store):
|
|
exp_id = store.create_experiment("test_search_runs_start_time_alias")
|
|
# Set start_time to ensure the search result is deterministic
|
|
run1 = store.create_run(exp_id, user_id="user", start_time=1, tags=[], run_name="name")
|
|
run2 = store.create_run(exp_id, user_id="user", start_time=2, tags=[], run_name="name")
|
|
run_id1 = run1.info.run_id
|
|
run_id2 = run2.info.run_id
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="attributes.run_name = 'name'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
order_by=["attributes.start_time DESC"],
|
|
)
|
|
assert [r.info.run_id for r in result] == [run_id2, run_id1]
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="attributes.run_name = 'name'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
order_by=["attributes.created ASC"],
|
|
)
|
|
assert [r.info.run_id for r in result] == [run_id1, run_id2]
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="attributes.run_name = 'name'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
order_by=["attributes.Created DESC"],
|
|
)
|
|
assert [r.info.run_id for r in result] == [run_id2, run_id1]
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="attributes.start_time > 0",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert {r.info.run_id for r in result} == {run_id1, run_id2}
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="attributes.created > 1",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert [r.info.run_id for r in result] == [run_id2]
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="attributes.Created > 2",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert result == []
|
|
|
|
|
|
def test_search_runs_datasets(store):
|
|
exp_id = store.create_experiment("12345dataset")
|
|
|
|
run1 = store.create_run(
|
|
experiment_id=exp_id,
|
|
user_id="user1",
|
|
start_time=1,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
run2 = store.create_run(
|
|
experiment_id=exp_id,
|
|
user_id="user2",
|
|
start_time=3,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
run3 = store.create_run(
|
|
experiment_id=exp_id,
|
|
user_id="user3",
|
|
start_time=2,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
|
|
dataset1 = Dataset(
|
|
name="name1",
|
|
digest="digest1",
|
|
source_type="st1",
|
|
source="source1",
|
|
schema="schema1",
|
|
profile="profile1",
|
|
)
|
|
dataset2 = Dataset(
|
|
name="name2",
|
|
digest="digest2",
|
|
source_type="st2",
|
|
source="source2",
|
|
schema="schema2",
|
|
profile="profile2",
|
|
)
|
|
dataset3 = Dataset(
|
|
name="name3",
|
|
digest="digest3",
|
|
source_type="st3",
|
|
source="source3",
|
|
schema="schema3",
|
|
profile="profile3",
|
|
)
|
|
|
|
test_tag = [InputTag(key=MLFLOW_DATASET_CONTEXT, value="test")]
|
|
train_tag = [InputTag(key=MLFLOW_DATASET_CONTEXT, value="train")]
|
|
eval_tag = [InputTag(key=MLFLOW_DATASET_CONTEXT, value="eval")]
|
|
|
|
inputs_run1 = [DatasetInput(dataset1, train_tag), DatasetInput(dataset2, eval_tag)]
|
|
inputs_run2 = [DatasetInput(dataset1, train_tag), DatasetInput(dataset3, eval_tag)]
|
|
inputs_run3 = [DatasetInput(dataset2, test_tag)]
|
|
|
|
store.log_inputs(run1.info.run_id, inputs_run1)
|
|
store.log_inputs(run2.info.run_id, inputs_run2)
|
|
store.log_inputs(run3.info.run_id, inputs_run3)
|
|
run_id1 = run1.info.run_id
|
|
run_id2 = run2.info.run_id
|
|
run_id3 = run3.info.run_id
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="dataset.name = 'name1'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert {r.info.run_id for r in result} == {run_id2, run_id1}
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="dataset.digest = 'digest2'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert {r.info.run_id for r in result} == {run_id3, run_id1}
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="dataset.name = 'name4'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert {r.info.run_id for r in result} == set()
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="dataset.context = 'train'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert {r.info.run_id for r in result} == {run_id2, run_id1}
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="dataset.context = 'test'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert {r.info.run_id for r in result} == {run_id3}
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="dataset.context = 'test' and dataset.name = 'name2'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert {r.info.run_id for r in result} == {run_id3}
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="dataset.name = 'name2' and dataset.context = 'test'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert {r.info.run_id for r in result} == {run_id3}
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="datasets.name IN ('name1', 'name2')",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert {r.info.run_id for r in result} == {run_id3, run_id1, run_id2}
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="datasets.digest IN ('digest1', 'digest2')",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert {r.info.run_id for r in result} == {run_id3, run_id1, run_id2}
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="datasets.name LIKE 'Name%'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert {r.info.run_id for r in result} == set()
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="datasets.name ILIKE 'Name%'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert {r.info.run_id for r in result} == {run_id3, run_id1, run_id2}
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="datasets.context ILIKE 'test%'",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert {r.info.run_id for r in result} == {run_id3}
|
|
|
|
result = store.search_runs(
|
|
[exp_id],
|
|
filter_string="datasets.context IN ('test', 'train')",
|
|
run_view_type=ViewType.ACTIVE_ONLY,
|
|
)
|
|
assert {r.info.run_id for r in result} == {run_id3, run_id1, run_id2}
|
|
|
|
|
|
def test_weird_param_names(store):
|
|
WEIRD_PARAM_NAME = "this is/a weird/but valid param"
|
|
_, exp_data, _ = _create_root(store)
|
|
run_id = exp_data[FileStore.DEFAULT_EXPERIMENT_ID]["runs"][0]
|
|
store.log_param(run_id, Param(WEIRD_PARAM_NAME, "Value"))
|
|
run = store.get_run(run_id)
|
|
assert run.data.params[WEIRD_PARAM_NAME] == "Value"
|
|
|
|
|
|
def test_log_param_empty_str(store):
|
|
PARAM_NAME = "new param"
|
|
_, exp_data, _ = _create_root(store)
|
|
run_id = exp_data[FileStore.DEFAULT_EXPERIMENT_ID]["runs"][0]
|
|
store.log_param(run_id, Param(PARAM_NAME, ""))
|
|
run = store.get_run(run_id)
|
|
assert run.data.params[PARAM_NAME] == ""
|
|
|
|
|
|
def test_log_param_with_newline(store):
|
|
param_name = "new param"
|
|
param_value = "a string\nwith multiple\nlines"
|
|
_, exp_data, _ = _create_root(store)
|
|
run_id = exp_data[FileStore.DEFAULT_EXPERIMENT_ID]["runs"][0]
|
|
store.log_param(run_id, Param(param_name, param_value))
|
|
run = store.get_run(run_id)
|
|
assert run.data.params[param_name] == param_value
|
|
|
|
|
|
def test_log_param_enforces_value_immutability(store):
|
|
param_name = "new param"
|
|
_, exp_data, _ = _create_root(store)
|
|
run_id = exp_data[FileStore.DEFAULT_EXPERIMENT_ID]["runs"][0]
|
|
store.log_param(run_id, Param(param_name, "value1"))
|
|
# Duplicate calls to `log_param` with the same key and value should succeed
|
|
store.log_param(run_id, Param(param_name, "value1"))
|
|
with pytest.raises(
|
|
MlflowException, match="Changing param values is not allowed. Param with key="
|
|
) as e:
|
|
store.log_param(run_id, Param(param_name, "value2"))
|
|
assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
|
|
run = store.get_run(run_id)
|
|
assert run.data.params[param_name] == "value1"
|
|
|
|
|
|
def test_log_param_max_length_value(store, monkeypatch):
|
|
param_name = "new param"
|
|
param_value = "x" * 6000
|
|
_, exp_data, _ = _create_root(store)
|
|
run_id = exp_data[FileStore.DEFAULT_EXPERIMENT_ID]["runs"][0]
|
|
store.log_param(run_id, Param(param_name, param_value))
|
|
run = store.get_run(run_id)
|
|
assert run.data.params[param_name] == param_value
|
|
monkeypatch.setenv("MLFLOW_TRUNCATE_LONG_VALUES", "false")
|
|
with pytest.raises(MlflowException, match="exceeds the maximum length"):
|
|
store.log_param(run_id, Param(param_name, "x" * 6001))
|
|
|
|
monkeypatch.setenv("MLFLOW_TRUNCATE_LONG_VALUES", "true")
|
|
store.log_param(run_id, Param(param_name, "x" * 6001))
|
|
|
|
|
|
def test_weird_metric_names(store):
|
|
WEIRD_METRIC_NAME = "this is/a weird/but valid metric"
|
|
_, exp_data, _ = _create_root(store)
|
|
run_id = exp_data[FileStore.DEFAULT_EXPERIMENT_ID]["runs"][0]
|
|
store.log_metric(run_id, Metric(WEIRD_METRIC_NAME, 10, 1234, 0))
|
|
run = store.get_run(run_id)
|
|
assert run.data.metrics[WEIRD_METRIC_NAME] == 10
|
|
history = store.get_metric_history(run_id, WEIRD_METRIC_NAME)
|
|
assert len(history) == 1
|
|
metric = history[0]
|
|
assert metric.key == WEIRD_METRIC_NAME
|
|
assert metric.value == 10
|
|
assert metric.timestamp == 1234
|
|
|
|
|
|
def test_weird_tag_names(store):
|
|
WEIRD_TAG_NAME = "this is/a weird/but valid tag"
|
|
_, exp_data, _ = _create_root(store)
|
|
run_id = exp_data[FileStore.DEFAULT_EXPERIMENT_ID]["runs"][0]
|
|
store.set_tag(run_id, RunTag(WEIRD_TAG_NAME, "Muhahaha!"))
|
|
run = store.get_run(run_id)
|
|
assert run.data.tags[WEIRD_TAG_NAME] == "Muhahaha!"
|
|
|
|
|
|
def test_set_experiment_tags(store):
|
|
experiments, _, _ = _create_root(store)
|
|
store.set_experiment_tag(FileStore.DEFAULT_EXPERIMENT_ID, ExperimentTag("tag0", "value0"))
|
|
store.set_experiment_tag(FileStore.DEFAULT_EXPERIMENT_ID, ExperimentTag("tag1", "value1"))
|
|
experiment = store.get_experiment(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
assert len(experiment.tags) == 2
|
|
assert experiment.tags["tag0"] == "value0"
|
|
assert experiment.tags["tag1"] == "value1"
|
|
# test that updating a tag works
|
|
store.set_experiment_tag(FileStore.DEFAULT_EXPERIMENT_ID, ExperimentTag("tag0", "value00000"))
|
|
experiment = store.get_experiment(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
assert experiment.tags["tag0"] == "value00000"
|
|
assert experiment.tags["tag1"] == "value1"
|
|
# test that setting a tag on 1 experiment does not impact another experiment.
|
|
exp_id = None
|
|
for exp in experiments:
|
|
if exp != FileStore.DEFAULT_EXPERIMENT_ID:
|
|
exp_id = exp
|
|
break
|
|
experiment = store.get_experiment(exp_id)
|
|
assert len(experiment.tags) == 0
|
|
# setting a tag on different experiments maintains different values across experiments
|
|
store.set_experiment_tag(exp_id, ExperimentTag("tag1", "value11111"))
|
|
experiment = store.get_experiment(exp_id)
|
|
assert len(experiment.tags) == 1
|
|
assert experiment.tags["tag1"] == "value11111"
|
|
experiment = store.get_experiment(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
assert experiment.tags["tag0"] == "value00000"
|
|
assert experiment.tags["tag1"] == "value1"
|
|
# test can set multi-line tags
|
|
store.set_experiment_tag(exp_id, ExperimentTag("multiline_tag", "value2\nvalue2\nvalue2"))
|
|
experiment = store.get_experiment(exp_id)
|
|
assert experiment.tags["multiline_tag"] == "value2\nvalue2\nvalue2"
|
|
# test cannot set tags on deleted experiments
|
|
store.delete_experiment(exp_id)
|
|
with pytest.raises(MlflowException, match="must be in the 'active' lifecycle_stage"):
|
|
store.set_experiment_tag(exp_id, ExperimentTag("should", "notset"))
|
|
|
|
|
|
def test_delete_experiment_tags(store):
|
|
experiments, _, _ = _create_root(store)
|
|
store.set_experiment_tag(FileStore.DEFAULT_EXPERIMENT_ID, ExperimentTag("tag0", "value0"))
|
|
store.set_experiment_tag(FileStore.DEFAULT_EXPERIMENT_ID, ExperimentTag("tag1", "value1"))
|
|
experiment = store.get_experiment(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
assert len(experiment.tags) == 2
|
|
assert experiment.tags["tag0"] == "value0"
|
|
assert experiment.tags["tag1"] == "value1"
|
|
# test that deleting a tag works
|
|
store.delete_experiment_tag(FileStore.DEFAULT_EXPERIMENT_ID, "tag0")
|
|
experiment = store.get_experiment(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
assert "tag0" not in experiment.tags.keys()
|
|
assert len(experiment.tags) == 1
|
|
|
|
|
|
def test_set_tags(store):
|
|
_, exp_data, _ = _create_root(store)
|
|
run_id = exp_data[FileStore.DEFAULT_EXPERIMENT_ID]["runs"][0]
|
|
store.set_tag(run_id, RunTag("tag0", "value0"))
|
|
store.set_tag(run_id, RunTag("tag1", "value1"))
|
|
tags = store.get_run(run_id).data.tags
|
|
assert tags["tag0"] == "value0"
|
|
assert tags["tag1"] == "value1"
|
|
|
|
# Can overwrite tags.
|
|
store.set_tag(run_id, RunTag("tag0", "value2"))
|
|
tags = store.get_run(run_id).data.tags
|
|
assert tags["tag0"] == "value2"
|
|
assert tags["tag1"] == "value1"
|
|
|
|
# Can set multiline tags.
|
|
store.set_tag(run_id, RunTag("multiline_tag", "value2\nvalue2\nvalue2"))
|
|
tags = store.get_run(run_id).data.tags
|
|
assert tags["multiline_tag"] == "value2\nvalue2\nvalue2"
|
|
|
|
|
|
def test_delete_tags(store):
|
|
experiments, exp_data, _ = _create_root(store)
|
|
exp_id = experiments[random_int(0, len(experiments) - 1)]
|
|
run_id = exp_data[exp_id]["runs"][0]
|
|
store.set_tag(run_id, RunTag("tag0", "value0"))
|
|
store.set_tag(run_id, RunTag("tag1", "value1"))
|
|
tags = store.get_run(run_id).data.tags
|
|
assert tags["tag0"] == "value0"
|
|
assert tags["tag1"] == "value1"
|
|
store.delete_tag(run_id, "tag0")
|
|
new_tags = store.get_run(run_id).data.tags
|
|
assert "tag0" not in new_tags.keys()
|
|
# test that you cannot delete tags that don't exist.
|
|
with pytest.raises(MlflowException, match="No tag with name"):
|
|
store.delete_tag(run_id, "fakeTag")
|
|
# test that you cannot delete tags for nonexistent runs
|
|
with pytest.raises(MlflowException, match=r"Run .+ not found"):
|
|
store.delete_tag("random_id", "tag0")
|
|
store.delete_run(run_id)
|
|
# test that you cannot delete tags for deleted runs.
|
|
assert store.get_run(run_id).info.lifecycle_stage == LifecycleStage.DELETED
|
|
with pytest.raises(MlflowException, match="must be in 'active' lifecycle_stage"):
|
|
store.delete_tag(run_id, "tag0")
|
|
|
|
|
|
def test_unicode_tag(store):
|
|
_, exp_data, _ = _create_root(store)
|
|
run_id = exp_data[FileStore.DEFAULT_EXPERIMENT_ID]["runs"][0]
|
|
value = "𝐼 𝓈𝑜𝓁𝑒𝓂𝓃𝓁𝓎 𝓈𝓌𝑒𝒶𝓇 𝓉𝒽𝒶𝓉 𝐼 𝒶𝓂 𝓊𝓅 𝓉𝑜 𝓃𝑜 𝑔𝑜𝑜𝒹"
|
|
store.set_tag(run_id, RunTag("message", value))
|
|
tags = store.get_run(run_id).data.tags
|
|
assert tags["message"] == value
|
|
|
|
|
|
def test_get_deleted_run(store):
|
|
"""
|
|
Getting metrics/tags/params/run info should be allowed on deleted runs.
|
|
"""
|
|
experiments, exp_data, _ = _create_root(store)
|
|
exp_id = experiments[random_int(0, len(experiments) - 1)]
|
|
run_id = exp_data[exp_id]["runs"][0]
|
|
store.delete_run(run_id)
|
|
assert store.get_run(run_id)
|
|
|
|
|
|
def test_set_deleted_run(store):
|
|
"""
|
|
Setting metrics/tags/params/updating run info should not be allowed on deleted runs.
|
|
"""
|
|
experiments, exp_data, _ = _create_root(store)
|
|
exp_id = experiments[random_int(0, len(experiments) - 1)]
|
|
run_id = exp_data[exp_id]["runs"][0]
|
|
store.delete_run(run_id)
|
|
|
|
assert store.get_run(run_id).info.lifecycle_stage == LifecycleStage.DELETED
|
|
match = "must be in 'active' lifecycle_stage"
|
|
with pytest.raises(MlflowException, match=match):
|
|
store.set_tag(run_id, RunTag("a", "b"))
|
|
with pytest.raises(MlflowException, match=match):
|
|
store.log_metric(run_id, Metric("a", 0.0, timestamp=0, step=0))
|
|
with pytest.raises(MlflowException, match=match):
|
|
store.log_param(run_id, Param("a", "b"))
|
|
|
|
|
|
def test_default_experiment_attempted_deletion(store):
|
|
_create_root(store)
|
|
with pytest.raises(MlflowException, match="Cannot delete the default experiment"):
|
|
store.delete_experiment(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
experiment = store.get_experiment(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
assert experiment.lifecycle_stage == LifecycleStage.ACTIVE
|
|
test_id = store.create_experiment("test")
|
|
store.delete_experiment(test_id)
|
|
test_experiment = store.get_experiment(test_id)
|
|
assert test_experiment.lifecycle_stage == LifecycleStage.DELETED
|
|
|
|
|
|
def test_malformed_experiment(store):
|
|
exp_0 = store.get_experiment(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
assert exp_0.experiment_id == FileStore.DEFAULT_EXPERIMENT_ID
|
|
|
|
experiments = len(store.search_experiments(view_type=ViewType.ALL))
|
|
|
|
# delete metadata file.
|
|
path = os.path.join(store.root_directory, str(exp_0.experiment_id), "meta.yaml")
|
|
os.remove(path)
|
|
with pytest.raises(MissingConfigException, match="does not exist"):
|
|
store.get_experiment(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
|
|
assert len(store.search_experiments(view_type=ViewType.ALL)) == experiments - 1
|
|
|
|
|
|
def test_malformed_run(store):
|
|
_, exp_data, _ = _create_root(store)
|
|
exp_0 = store.get_experiment(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
all_runs = _search(store, exp_0.experiment_id)
|
|
|
|
all_run_ids = exp_data[exp_0.experiment_id]["runs"]
|
|
assert len(all_runs) == len(all_run_ids)
|
|
|
|
# delete metadata file.
|
|
bad_run_id = exp_data[exp_0.experiment_id]["runs"][0]
|
|
path = os.path.join(
|
|
store.root_directory, str(exp_0.experiment_id), str(bad_run_id), "meta.yaml"
|
|
)
|
|
os.remove(path)
|
|
with pytest.raises(MissingConfigException, match="does not exist"):
|
|
store.get_run(bad_run_id)
|
|
|
|
valid_runs = _search(store, exp_0.experiment_id)
|
|
assert len(valid_runs) == len(all_runs) - 1
|
|
|
|
for rid in all_run_ids:
|
|
if rid != bad_run_id:
|
|
store.get_run(rid)
|
|
|
|
|
|
def test_malformed_metric(store):
|
|
exp_id = FileStore.DEFAULT_EXPERIMENT_ID
|
|
run_id = store.create_run(
|
|
experiment_id=exp_id,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="first name",
|
|
).info.run_id
|
|
store.log_metric(run_id, Metric("test", 1, 0, 0))
|
|
with (
|
|
mock.patch("mlflow.store.tracking.file_store.read_file_lines", return_value=["0 1 0 2\n"]),
|
|
pytest.raises(
|
|
MlflowException,
|
|
match=f"Metric 'test' is malformed; persisted metric data contained "
|
|
f"4 fields. Expected 2, 3, or 5 fields. "
|
|
f"Experiment id: {exp_id}",
|
|
),
|
|
):
|
|
store.get_metric_history(run_id, "test")
|
|
|
|
|
|
def test_mismatching_experiment_id(store):
|
|
exp_0 = store.get_experiment(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
assert exp_0.experiment_id == FileStore.DEFAULT_EXPERIMENT_ID
|
|
|
|
experiments = len(store.search_experiments(view_type=ViewType.ALL))
|
|
|
|
# mv experiment folder
|
|
target = "1"
|
|
path_orig = os.path.join(store.root_directory, str(exp_0.experiment_id))
|
|
path_new = os.path.join(store.root_directory, str(target))
|
|
os.rename(path_orig, path_new)
|
|
|
|
with pytest.raises(MlflowException, match="Could not find experiment with ID"):
|
|
store.get_experiment(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
|
|
with pytest.raises(MlflowException, match="does not exist"):
|
|
store.get_experiment(target)
|
|
assert len(store.search_experiments(view_type=ViewType.ALL)) == experiments - 1
|
|
|
|
|
|
def test_bad_experiment_id_recorded_for_run(store):
|
|
_, exp_data, _ = _create_root(store)
|
|
exp_0 = store.get_experiment(FileStore.DEFAULT_EXPERIMENT_ID)
|
|
all_runs = _search(store, exp_0.experiment_id)
|
|
|
|
all_run_ids = exp_data[exp_0.experiment_id]["runs"]
|
|
assert len(all_runs) == len(all_run_ids)
|
|
|
|
# change experiment pointer in run
|
|
bad_run_id = str(exp_data[exp_0.experiment_id]["runs"][0])
|
|
path = os.path.join(store.root_directory, str(exp_0.experiment_id), bad_run_id)
|
|
experiment_data = read_yaml(path, "meta.yaml")
|
|
experiment_data["experiment_id"] = 1
|
|
write_yaml(path, "meta.yaml", experiment_data, True)
|
|
|
|
with pytest.raises(MlflowException, match="metadata is in invalid state"):
|
|
store.get_run(bad_run_id)
|
|
|
|
valid_runs = _search(store, exp_0.experiment_id)
|
|
assert len(valid_runs) == len(all_runs) - 1
|
|
|
|
for rid in all_run_ids:
|
|
if rid != bad_run_id:
|
|
store.get_run(rid)
|
|
|
|
|
|
def test_log_batch(store):
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
run_id = run.info.run_id
|
|
metric_entities = [Metric("m1", 0.87, 12345, 0), Metric("m2", 0.49, 12345, 0)]
|
|
param_entities = [Param("p1", "p1val"), Param("p2", "p2val")]
|
|
tag_entities = [RunTag("t1", "t1val"), RunTag("t2", "t2val")]
|
|
store.log_batch(
|
|
run_id=run_id, metrics=metric_entities, params=param_entities, tags=tag_entities
|
|
)
|
|
_verify_logged(store, run_id, metric_entities, param_entities, tag_entities)
|
|
|
|
|
|
def test_log_batch_max_length_value(store, monkeypatch):
|
|
param_entities = [Param("long param", "x" * 6000), Param("short param", "xyz")]
|
|
expected_param_entities = [
|
|
Param("long param", "x" * 6000),
|
|
Param("short param", "xyz"),
|
|
]
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
store.log_batch(run.info.run_id, (), param_entities, ())
|
|
_verify_logged(store, run.info.run_id, (), expected_param_entities, ())
|
|
|
|
monkeypatch.setenv("MLFLOW_TRUNCATE_LONG_VALUES", "false")
|
|
param_entities = [Param("long param", "x" * 6001), Param("short param", "xyz")]
|
|
with pytest.raises(MlflowException, match="exceeds the maximum length"):
|
|
store.log_batch(run.info.run_id, (), param_entities, ())
|
|
|
|
monkeypatch.setenv("MLFLOW_TRUNCATE_LONG_VALUES", "true")
|
|
store.log_batch(run.info.run_id, (), param_entities, ())
|
|
|
|
|
|
def test_log_batch_internal_error(store):
|
|
# Verify that internal errors during log_batch result in MlflowExceptions
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
|
|
def _raise_exception_fn(*args, **kwargs):
|
|
raise Exception("Some internal error")
|
|
|
|
with (
|
|
mock.patch(FILESTORE_PACKAGE + ".FileStore._log_run_metric") as log_metric_mock,
|
|
mock.patch(FILESTORE_PACKAGE + ".FileStore._log_run_param") as log_param_mock,
|
|
mock.patch(FILESTORE_PACKAGE + ".FileStore._set_run_tag") as set_tag_mock,
|
|
):
|
|
log_metric_mock.side_effect = _raise_exception_fn
|
|
log_param_mock.side_effect = _raise_exception_fn
|
|
set_tag_mock.side_effect = _raise_exception_fn
|
|
for kwargs in [
|
|
{"metrics": [Metric("a", 3, 1, 0)]},
|
|
{"params": [Param("b", "c")]},
|
|
{"tags": [RunTag("c", "d")]},
|
|
]:
|
|
log_batch_kwargs = {"metrics": [], "params": [], "tags": []}
|
|
log_batch_kwargs.update(kwargs)
|
|
with pytest.raises(MlflowException, match="Some internal error") as e:
|
|
store.log_batch(run.info.run_id, **log_batch_kwargs)
|
|
assert e.value.error_code == ErrorCode.Name(INTERNAL_ERROR)
|
|
|
|
|
|
def test_log_batch_nonexistent_run(store):
|
|
nonexistent_uuid = uuid.uuid4().hex
|
|
with pytest.raises(MlflowException, match=f"Run '{nonexistent_uuid}' not found") as e:
|
|
store.log_batch(nonexistent_uuid, [], [], [])
|
|
assert e.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
|
|
|
|
|
|
def test_log_batch_params_idempotency(store):
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
params = [Param("p-key", "p-val")]
|
|
store.log_batch(run.info.run_id, metrics=[], params=params, tags=[])
|
|
store.log_batch(run.info.run_id, metrics=[], params=params, tags=[])
|
|
_verify_logged(store, run.info.run_id, metrics=[], params=params, tags=[])
|
|
|
|
|
|
def test_log_batch_tags_idempotency(store):
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
store.log_batch(run.info.run_id, metrics=[], params=[], tags=[RunTag("t-key", "t-val")])
|
|
store.log_batch(run.info.run_id, metrics=[], params=[], tags=[RunTag("t-key", "t-val")])
|
|
_verify_logged(store, run.info.run_id, metrics=[], params=[], tags=[RunTag("t-key", "t-val")])
|
|
|
|
|
|
def test_log_batch_allows_tag_overwrite(store):
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
store.log_batch(run.info.run_id, metrics=[], params=[], tags=[RunTag("t-key", "val")])
|
|
store.log_batch(run.info.run_id, metrics=[], params=[], tags=[RunTag("t-key", "newval")])
|
|
_verify_logged(store, run.info.run_id, metrics=[], params=[], tags=[RunTag("t-key", "newval")])
|
|
|
|
|
|
def test_log_batch_same_metric_repeated_single_req(store):
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
metric0 = Metric(key="metric-key", value=1, timestamp=2, step=0)
|
|
metric1 = Metric(key="metric-key", value=2, timestamp=3, step=0)
|
|
store.log_batch(run.info.run_id, params=[], metrics=[metric0, metric1], tags=[])
|
|
_verify_logged(store, run.info.run_id, params=[], metrics=[metric0, metric1], tags=[])
|
|
|
|
|
|
def test_log_batch_same_metric_repeated_multiple_reqs(store):
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
metric0 = Metric(key="metric-key", value=1, timestamp=2, step=0)
|
|
metric1 = Metric(key="metric-key", value=2, timestamp=3, step=0)
|
|
store.log_batch(run.info.run_id, params=[], metrics=[metric0], tags=[])
|
|
_verify_logged(store, run.info.run_id, params=[], metrics=[metric0], tags=[])
|
|
store.log_batch(run.info.run_id, params=[], metrics=[metric1], tags=[])
|
|
_verify_logged(store, run.info.run_id, params=[], metrics=[metric0, metric1], tags=[])
|
|
|
|
|
|
def test_log_batch_allows_tag_overwrite_single_req(store):
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
tags = [RunTag("t-key", "val"), RunTag("t-key", "newval")]
|
|
store.log_batch(run.info.run_id, metrics=[], params=[], tags=tags)
|
|
_verify_logged(store, run.info.run_id, metrics=[], params=[], tags=[tags[-1]])
|
|
|
|
|
|
def test_log_batch_accepts_empty_payload(store):
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
store.log_batch(run.info.run_id, metrics=[], params=[], tags=[])
|
|
_verify_logged(store, run.info.run_id, metrics=[], params=[], tags=[])
|
|
|
|
|
|
def test_log_batch_with_duplicate_params_errors_no_partial_write(store):
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
with pytest.raises(MlflowException, match="Duplicate parameter keys have been submitted") as e:
|
|
store.log_batch(
|
|
run.info.run_id,
|
|
metrics=[],
|
|
params=[Param("a", "1"), Param("a", "2")],
|
|
tags=[],
|
|
)
|
|
assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
|
|
_verify_logged(store, run.info.run_id, metrics=[], params=[], tags=[])
|
|
|
|
|
|
def test_update_run_name(store):
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
run_id = run.info.run_id
|
|
|
|
assert run.info.run_name == "name"
|
|
assert run.data.tags.get(MLFLOW_RUN_NAME) == "name"
|
|
|
|
store.update_run_info(run_id, RunStatus.FINISHED, 100, "new name")
|
|
run = store.get_run(run_id)
|
|
assert run.info.run_name == "new name"
|
|
assert run.data.tags.get(MLFLOW_RUN_NAME) == "new name"
|
|
|
|
store.update_run_info(run_id, RunStatus.FINISHED, 100, None)
|
|
run = store.get_run(run_id)
|
|
assert run.info.run_name == "new name"
|
|
assert run.data.tags.get(MLFLOW_RUN_NAME) == "new name"
|
|
|
|
store.delete_tag(run_id, MLFLOW_RUN_NAME)
|
|
run = store.get_run(run_id)
|
|
assert run.info.run_name == "new name"
|
|
assert run.data.tags.get(MLFLOW_RUN_NAME) is None
|
|
|
|
store.update_run_info(run_id, RunStatus.FINISHED, 100, "another name")
|
|
run = store.get_run(run_id)
|
|
assert run.data.tags.get(MLFLOW_RUN_NAME) == "another name"
|
|
assert run.info.run_name == "another name"
|
|
|
|
store.set_tag(run_id, RunTag(MLFLOW_RUN_NAME, "yet another name"))
|
|
run = store.get_run(run_id)
|
|
assert run.info.run_name == "yet another name"
|
|
assert run.data.tags.get(MLFLOW_RUN_NAME) == "yet another name"
|
|
|
|
store.log_batch(run_id, metrics=[], params=[], tags=[RunTag(MLFLOW_RUN_NAME, "batch name")])
|
|
run = store.get_run(run_id)
|
|
assert run.info.run_name == "batch name"
|
|
assert run.data.tags.get(MLFLOW_RUN_NAME) == "batch name"
|
|
|
|
|
|
def test_get_metric_history_on_non_existent_metric_key(store):
|
|
run = store.create_run(
|
|
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name="name",
|
|
)
|
|
run_id = run.info.run_id
|
|
test_metrics = store.get_metric_history(run_id, "test_metric")
|
|
assert isinstance(test_metrics, PagedList)
|
|
assert test_metrics == []
|
|
|
|
|
|
def test_experiment_with_default_root_artifact_uri(tmp_path):
|
|
file_store_root_uri = path_to_local_file_uri(tmp_path)
|
|
file_store = FileStore(file_store_root_uri)
|
|
experiment_id = file_store.create_experiment(name="test", artifact_location="test")
|
|
experiment_info = file_store.get_experiment(experiment_id)
|
|
if is_windows():
|
|
assert experiment_info.artifact_location == Path.cwd().joinpath("test").as_uri()
|
|
else:
|
|
assert experiment_info.artifact_location == str(Path.cwd().joinpath("test"))
|
|
|
|
|
|
def test_experiment_with_relative_artifact_uri(tmp_path):
|
|
file_store_root_uri = append_to_uri_path(path_to_local_file_uri(tmp_path), "experiments")
|
|
artifacts_root_uri = append_to_uri_path(path_to_local_file_uri(tmp_path), "artifacts")
|
|
file_store = FileStore(file_store_root_uri, artifacts_root_uri)
|
|
experiment_id = file_store.create_experiment(name="test")
|
|
experiment_info = file_store.get_experiment(experiment_id)
|
|
assert experiment_info.artifact_location == append_to_uri_path(
|
|
artifacts_root_uri, experiment_id
|
|
)
|
|
|
|
|
|
def _assert_create_run_appends_to_artifact_uri_path_correctly(
|
|
artifact_root_uri, expected_artifact_uri_format
|
|
):
|
|
with TempDir() as tmp:
|
|
fs = FileStore(tmp.path(), artifact_root_uri)
|
|
exp_id = fs.create_experiment("exp")
|
|
run = fs.create_run(
|
|
experiment_id=exp_id, user_id="user", start_time=0, tags=[], run_name="name"
|
|
)
|
|
cwd = Path.cwd().as_posix()
|
|
drive = Path.cwd().drive
|
|
if is_windows() and expected_artifact_uri_format.startswith("file:"):
|
|
cwd = f"/{cwd}"
|
|
drive = f"{drive}/"
|
|
assert run.info.artifact_uri == expected_artifact_uri_format.format(
|
|
e=exp_id, r=run.info.run_id, cwd=cwd, drive=drive
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows")
|
|
@pytest.mark.parametrize(
|
|
("input_uri", "expected_uri"),
|
|
[
|
|
(
|
|
"\\my_server/my_path/my_sub_path",
|
|
"file:///{drive}my_server/my_path/my_sub_path/{e}/{r}/artifacts",
|
|
),
|
|
("path/to/local/folder", "file://{cwd}/path/to/local/folder/{e}/{r}/artifacts"),
|
|
(
|
|
"/path/to/local/folder",
|
|
"file:///{drive}path/to/local/folder/{e}/{r}/artifacts",
|
|
),
|
|
(
|
|
"#path/to/local/folder?",
|
|
"file://{cwd}/{e}/{r}/artifacts#path/to/local/folder?",
|
|
),
|
|
(
|
|
"file:///path/to/local/folder",
|
|
"file:///{drive}path/to/local/folder/{e}/{r}/artifacts",
|
|
),
|
|
(
|
|
"file:///path/to/local/folder?param=value#fragment",
|
|
"file:///{drive}path/to/local/folder/{e}/{r}/artifacts?param=value#fragment",
|
|
),
|
|
(
|
|
"file:path/to/local/folder",
|
|
"file://{cwd}/path/to/local/folder/{e}/{r}/artifacts",
|
|
),
|
|
(
|
|
"file:path/to/local/folder?param=value",
|
|
"file://{cwd}/path/to/local/folder/{e}/{r}/artifacts?param=value",
|
|
),
|
|
],
|
|
)
|
|
def test_create_run_appends_to_artifact_local_path_file_uri_correctly_on_windows(
|
|
input_uri, expected_uri
|
|
):
|
|
_assert_create_run_appends_to_artifact_uri_path_correctly(input_uri, expected_uri)
|
|
|
|
|
|
@pytest.mark.skipif(is_windows(), reason="This test fails on Windows")
|
|
@pytest.mark.parametrize(
|
|
("input_uri", "expected_uri"),
|
|
[
|
|
("path/to/local/folder", "{cwd}/path/to/local/folder/{e}/{r}/artifacts"),
|
|
("/path/to/local/folder", "/path/to/local/folder/{e}/{r}/artifacts"),
|
|
("#path/to/local/folder?", "{cwd}/#path/to/local/folder?/{e}/{r}/artifacts"),
|
|
(
|
|
"file:///path/to/local/folder",
|
|
"file:///path/to/local/folder/{e}/{r}/artifacts",
|
|
),
|
|
(
|
|
"file:///path/to/local/folder?param=value#fragment",
|
|
"file:///path/to/local/folder/{e}/{r}/artifacts?param=value#fragment",
|
|
),
|
|
(
|
|
"file:path/to/local/folder",
|
|
"file://{cwd}/path/to/local/folder/{e}/{r}/artifacts",
|
|
),
|
|
(
|
|
"file:path/to/local/folder?param=value",
|
|
"file://{cwd}/path/to/local/folder/{e}/{r}/artifacts?param=value",
|
|
),
|
|
],
|
|
)
|
|
def test_create_run_appends_to_artifact_local_path_file_uri_correctly(input_uri, expected_uri):
|
|
_assert_create_run_appends_to_artifact_uri_path_correctly(input_uri, expected_uri)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("input_uri", "expected_uri"),
|
|
[
|
|
("s3://bucket/path/to/root", "s3://bucket/path/to/root/{e}/{r}/artifacts"),
|
|
(
|
|
"s3://bucket/path/to/root?creds=mycreds",
|
|
"s3://bucket/path/to/root/{e}/{r}/artifacts?creds=mycreds",
|
|
),
|
|
(
|
|
"dbscheme+driver://root@host/dbname?creds=mycreds#myfragment",
|
|
"dbscheme+driver://root@host/dbname/{e}/{r}/artifacts?creds=mycreds#myfragment",
|
|
),
|
|
(
|
|
"dbscheme+driver://root:password@hostname.com?creds=mycreds#myfragment",
|
|
"dbscheme+driver://root:password@hostname.com/{e}/{r}/artifacts"
|
|
"?creds=mycreds#myfragment",
|
|
),
|
|
(
|
|
"dbscheme+driver://root:password@hostname.com/mydb?creds=mycreds#myfragment",
|
|
"dbscheme+driver://root:password@hostname.com/mydb/{e}/{r}/artifacts"
|
|
"?creds=mycreds#myfragment",
|
|
),
|
|
],
|
|
)
|
|
def test_create_run_appends_to_artifact_uri_path_correctly(input_uri, expected_uri):
|
|
_assert_create_run_appends_to_artifact_uri_path_correctly(input_uri, expected_uri)
|
|
|
|
|
|
def _assert_create_experiment_appends_to_artifact_uri_path_correctly(
|
|
artifact_root_uri, expected_artifact_uri_format
|
|
):
|
|
with TempDir() as tmp:
|
|
fs = FileStore(tmp.path(), artifact_root_uri)
|
|
exp_id = fs.create_experiment("exp")
|
|
exp = fs.get_experiment(exp_id)
|
|
cwd = Path.cwd().as_posix()
|
|
drive = Path.cwd().drive
|
|
if is_windows() and expected_artifact_uri_format.startswith("file:"):
|
|
cwd = f"/{cwd}"
|
|
drive = f"{drive}/"
|
|
|
|
assert exp.artifact_location == expected_artifact_uri_format.format(
|
|
e=exp_id, cwd=cwd, drive=drive
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows")
|
|
@pytest.mark.parametrize(
|
|
("input_uri", "expected_uri"),
|
|
[
|
|
(
|
|
"\\my_server/my_path/my_sub_path",
|
|
"file:///{drive}my_server/my_path/my_sub_path/{e}",
|
|
),
|
|
("path/to/local/folder", "file://{cwd}/path/to/local/folder/{e}"),
|
|
("/path/to/local/folder", "file:///{drive}path/to/local/folder/{e}"),
|
|
("#path/to/local/folder?", "file://{cwd}/{e}#path/to/local/folder?"),
|
|
("file:path/to/local/folder", "file://{cwd}/path/to/local/folder/{e}"),
|
|
("file:///path/to/local/folder", "file:///{drive}path/to/local/folder/{e}"),
|
|
(
|
|
"file:path/to/local/folder?param=value",
|
|
"file://{cwd}/path/to/local/folder/{e}?param=value",
|
|
),
|
|
(
|
|
"file:///path/to/local/folder?param=value#fragment",
|
|
"file:///{drive}path/to/local/folder/{e}?param=value#fragment",
|
|
),
|
|
],
|
|
)
|
|
def test_create_experiment_appends_to_artifact_local_path_file_uri_correctly_on_windows(
|
|
input_uri, expected_uri
|
|
):
|
|
_assert_create_experiment_appends_to_artifact_uri_path_correctly(input_uri, expected_uri)
|
|
|
|
|
|
@pytest.mark.skipif(is_windows(), reason="This test fails on Windows")
|
|
@pytest.mark.parametrize(
|
|
("input_uri", "expected_uri"),
|
|
[
|
|
("path/to/local/folder", "{cwd}/path/to/local/folder/{e}"),
|
|
("/path/to/local/folder", "/path/to/local/folder/{e}"),
|
|
("#path/to/local/folder?", "{cwd}/#path/to/local/folder?/{e}"),
|
|
("file:path/to/local/folder", "file://{cwd}/path/to/local/folder/{e}"),
|
|
("file:///path/to/local/folder", "file:///path/to/local/folder/{e}"),
|
|
(
|
|
"file:path/to/local/folder?param=value",
|
|
"file://{cwd}/path/to/local/folder/{e}?param=value",
|
|
),
|
|
(
|
|
"file:///path/to/local/folder?param=value#fragment",
|
|
"file:///path/to/local/folder/{e}?param=value#fragment",
|
|
),
|
|
],
|
|
)
|
|
def test_create_experiment_appends_to_artifact_local_path_file_uri_correctly(
|
|
input_uri, expected_uri
|
|
):
|
|
_assert_create_experiment_appends_to_artifact_uri_path_correctly(input_uri, expected_uri)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("input_uri", "expected_uri"),
|
|
[
|
|
("s3://bucket/path/to/root", "s3://bucket/path/to/root/{e}"),
|
|
(
|
|
"s3://bucket/path/to/root?creds=mycreds",
|
|
"s3://bucket/path/to/root/{e}?creds=mycreds",
|
|
),
|
|
(
|
|
"dbscheme+driver://root@host/dbname?creds=mycreds#myfragment",
|
|
"dbscheme+driver://root@host/dbname/{e}?creds=mycreds#myfragment",
|
|
),
|
|
(
|
|
"dbscheme+driver://root:password@hostname.com?creds=mycreds#myfragment",
|
|
"dbscheme+driver://root:password@hostname.com/{e}?creds=mycreds#myfragment",
|
|
),
|
|
(
|
|
"dbscheme+driver://root:password@hostname.com/mydb?creds=mycreds#myfragment",
|
|
"dbscheme+driver://root:password@hostname.com/mydb/{e}?creds=mycreds#myfragment",
|
|
),
|
|
],
|
|
)
|
|
def test_create_experiment_appends_to_artifact_uri_path_correctly(input_uri, expected_uri):
|
|
_assert_create_experiment_appends_to_artifact_uri_path_correctly(input_uri, expected_uri)
|
|
|
|
|
|
def assert_dataset_inputs_equal(inputs1: list[DatasetInput], inputs2: list[DatasetInput]):
|
|
inputs1 = sorted(inputs1, key=lambda inp: (inp.dataset.name, inp.dataset.digest))
|
|
inputs2 = sorted(inputs2, key=lambda inp: (inp.dataset.name, inp.dataset.digest))
|
|
assert len(inputs1) == len(inputs2)
|
|
for idx, inp1 in enumerate(inputs1):
|
|
inp2 = inputs2[idx]
|
|
assert dict(inp1.dataset) == dict(inp2.dataset)
|
|
tags1 = sorted(inp1.tags, key=lambda tag: tag.key)
|
|
tags2 = sorted(inp2.tags, key=lambda tag: tag.key)
|
|
for idx, tag1 in enumerate(tags1):
|
|
tag2 = tags2[idx]
|
|
assert tag1.key == tag2.key
|
|
assert tag1.value == tag2.value
|
|
|
|
|
|
def test_log_inputs_and_retrieve_runs_behaves_as_expected(store):
|
|
exp_id = store.create_experiment("12345dataset")
|
|
|
|
run1 = store.create_run(
|
|
experiment_id=exp_id,
|
|
user_id="user1",
|
|
start_time=1,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
run2 = store.create_run(
|
|
experiment_id=exp_id,
|
|
user_id="user2",
|
|
start_time=3,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
run3 = store.create_run(
|
|
experiment_id=exp_id,
|
|
user_id="user3",
|
|
start_time=2,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
|
|
dataset1 = Dataset(
|
|
name="name1",
|
|
digest="digest1",
|
|
source_type="st1",
|
|
source="source1",
|
|
schema="schema1",
|
|
profile="profile1",
|
|
)
|
|
dataset2 = Dataset(
|
|
name="name2",
|
|
digest="digest2",
|
|
source_type="st2",
|
|
source="source2",
|
|
schema="schema2",
|
|
profile="profile2",
|
|
)
|
|
dataset3 = Dataset(
|
|
name="name3",
|
|
digest="digest3",
|
|
source_type="st3",
|
|
source="source3",
|
|
schema="schema3",
|
|
profile="profile3",
|
|
)
|
|
|
|
tags1 = [InputTag(key="key1", value="value1"), InputTag(key="key2", value="value2")]
|
|
tags2 = [InputTag(key="key3", value="value3"), InputTag(key="key4", value="value4")]
|
|
tags3 = [InputTag(key="key5", value="value5"), InputTag(key="key6", value="value6")]
|
|
|
|
inputs_run1 = [DatasetInput(dataset1, tags1), DatasetInput(dataset2, tags1)]
|
|
inputs_run2 = [DatasetInput(dataset1, tags2), DatasetInput(dataset3, tags3)]
|
|
inputs_run3 = [DatasetInput(dataset2, tags3)]
|
|
|
|
store.log_inputs(run1.info.run_id, inputs_run1)
|
|
store.log_inputs(run2.info.run_id, inputs_run2)
|
|
store.log_inputs(run3.info.run_id, inputs_run3)
|
|
|
|
run1 = store.get_run(run1.info.run_id)
|
|
assert_dataset_inputs_equal(run1.inputs.dataset_inputs, inputs_run1)
|
|
run2 = store.get_run(run2.info.run_id)
|
|
assert_dataset_inputs_equal(run2.inputs.dataset_inputs, inputs_run2)
|
|
run3 = store.get_run(run3.info.run_id)
|
|
assert_dataset_inputs_equal(run3.inputs.dataset_inputs, inputs_run3)
|
|
|
|
search_results_1 = store.search_runs(
|
|
[exp_id], None, ViewType.ALL, max_results=4, order_by=["start_time ASC"]
|
|
)
|
|
run1 = search_results_1[0]
|
|
assert_dataset_inputs_equal(run1.inputs.dataset_inputs, inputs_run1)
|
|
run2 = search_results_1[2]
|
|
assert_dataset_inputs_equal(run2.inputs.dataset_inputs, inputs_run2)
|
|
run3 = search_results_1[1]
|
|
assert_dataset_inputs_equal(run3.inputs.dataset_inputs, inputs_run3)
|
|
|
|
search_results_2 = store.search_runs(
|
|
[exp_id], None, ViewType.ALL, max_results=4, order_by=["start_time DESC"]
|
|
)
|
|
run1 = search_results_2[2]
|
|
assert_dataset_inputs_equal(run1.inputs.dataset_inputs, inputs_run1)
|
|
run2 = search_results_2[0]
|
|
assert_dataset_inputs_equal(run2.inputs.dataset_inputs, inputs_run2)
|
|
run3 = search_results_2[1]
|
|
assert_dataset_inputs_equal(run3.inputs.dataset_inputs, inputs_run3)
|
|
|
|
|
|
def test_log_input_multiple_times_does_not_overwrite_tags_or_dataset(store):
|
|
exp_id = store.create_experiment("dataset_no_overwrite")
|
|
|
|
run = store.create_run(
|
|
experiment_id=exp_id,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
dataset = Dataset(
|
|
name="name",
|
|
digest="digest",
|
|
source_type="st",
|
|
source="source",
|
|
schema="schema",
|
|
profile="profile",
|
|
)
|
|
tags = [InputTag(key="key1", value="value1"), InputTag(key="key2", value="value2")]
|
|
store.log_inputs(run.info.run_id, [DatasetInput(dataset, tags)])
|
|
|
|
for i in range(3):
|
|
# Since the dataset name and digest are the same as the previously logged dataset,
|
|
# no changes should be made
|
|
overwrite_dataset = Dataset(
|
|
name="name",
|
|
digest="digest",
|
|
source_type=f"st{i}",
|
|
source=f"source{i}",
|
|
schema=f"schema{i}",
|
|
profile=f"profile{i}",
|
|
)
|
|
# Since the dataset has already been logged as an input to the run, no changes should be
|
|
# made to the input tags
|
|
overwrite_tags = [
|
|
InputTag(key=f"key{i}", value=f"value{i}"),
|
|
InputTag(key=f"key{i + 1}", value=f"value{i + 1}"),
|
|
]
|
|
store.log_inputs(run.info.run_id, [DatasetInput(overwrite_dataset, overwrite_tags)])
|
|
|
|
run = store.get_run(run.info.run_id)
|
|
assert_dataset_inputs_equal(run.inputs.dataset_inputs, [DatasetInput(dataset, tags)])
|
|
|
|
# Logging a dataset with a different name or digest to the original run should result
|
|
# in the addition of another dataset input
|
|
other_name_dataset = Dataset(
|
|
name="other_name",
|
|
digest="digest",
|
|
source_type="st",
|
|
source="source",
|
|
schema="schema",
|
|
profile="profile",
|
|
)
|
|
other_name_input_tags = [InputTag(key="k1", value="v1")]
|
|
store.log_inputs(run.info.run_id, [DatasetInput(other_name_dataset, other_name_input_tags)])
|
|
|
|
other_digest_dataset = Dataset(
|
|
name="name",
|
|
digest="other_digest",
|
|
source_type="st",
|
|
source="source",
|
|
schema="schema",
|
|
profile="profile",
|
|
)
|
|
other_digest_input_tags = [InputTag(key="k2", value="v2")]
|
|
store.log_inputs(run.info.run_id, [DatasetInput(other_digest_dataset, other_digest_input_tags)])
|
|
|
|
run = store.get_run(run.info.run_id)
|
|
assert_dataset_inputs_equal(
|
|
run.inputs.dataset_inputs,
|
|
[
|
|
DatasetInput(dataset, tags),
|
|
DatasetInput(other_name_dataset, other_name_input_tags),
|
|
DatasetInput(other_digest_dataset, other_digest_input_tags),
|
|
],
|
|
)
|
|
|
|
# Logging the same dataset with different tags to new runs should result in each run
|
|
# having its own new input tags and the same dataset input
|
|
for i in range(3):
|
|
new_run = store.create_run(
|
|
experiment_id=exp_id,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
new_tags = [
|
|
InputTag(key=f"key{i}", value=f"value{i}"),
|
|
InputTag(key=f"key{i + 1}", value=f"value{i + 1}"),
|
|
]
|
|
store.log_inputs(new_run.info.run_id, [DatasetInput(dataset, new_tags)])
|
|
new_run = store.get_run(new_run.info.run_id)
|
|
assert_dataset_inputs_equal(
|
|
new_run.inputs.dataset_inputs, [DatasetInput(dataset, new_tags)]
|
|
)
|
|
|
|
|
|
def test_log_inputs_uses_expected_input_and_dataset_ids_for_storage(store):
|
|
"""
|
|
This test verifies that the FileStore uses expected IDs as folder names to represent datasets
|
|
and run inputs. This is very important because the IDs are used to deduplicate inputs and
|
|
datasets if the same dataset is logged to multiple runs or the same dataset is logged
|
|
multiple times as an input to the same run with different tags.
|
|
|
|
**If this test fails, be very careful before removing or changing asserts. Unintended changes
|
|
could result in user-visible duplication of datasets and run inputs.**
|
|
"""
|
|
exp_id = store.create_experiment("dataset_expected_ids")
|
|
|
|
run1 = store.create_run(
|
|
experiment_id=exp_id,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
run2 = store.create_run(
|
|
experiment_id=exp_id,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
|
|
experiment_dir = store._get_experiment_path(exp_id, assert_exists=True)
|
|
datasets_dir = os.path.join(experiment_dir, FileStore.DATASETS_FOLDER_NAME)
|
|
|
|
def assert_expected_dataset_storage_ids_present(storage_ids):
|
|
assert set(os.listdir(datasets_dir)) == set(storage_ids)
|
|
|
|
def assert_expected_input_storage_ids_present(run, dataset_storage_ids):
|
|
run_dir = store._get_run_dir(run.info.experiment_id, run.info.run_id)
|
|
inputs_dir = os.path.join(run_dir, FileStore.INPUTS_FOLDER_NAME)
|
|
expected_input_storage_ids = []
|
|
for dataset_storage_id in dataset_storage_ids:
|
|
md5 = hashlib.md5(dataset_storage_id.encode("utf-8"), usedforsecurity=False)
|
|
md5.update(run.info.run_id.encode("utf-8"))
|
|
expected_input_storage_ids.append(md5.hexdigest())
|
|
assert set(os.listdir(inputs_dir)) == set(expected_input_storage_ids)
|
|
|
|
tags = [InputTag(key="key", value="value")]
|
|
|
|
dataset1 = Dataset(
|
|
name="name",
|
|
digest="digest",
|
|
source_type="st",
|
|
source="source",
|
|
schema="schema",
|
|
profile="profile",
|
|
)
|
|
store.log_inputs(run1.info.run_id, [DatasetInput(dataset1, tags)])
|
|
expected_dataset1_storage_id = "efa4363cd8179759e8c7f113aebdd340"
|
|
assert_expected_dataset_storage_ids_present([expected_dataset1_storage_id])
|
|
assert_expected_input_storage_ids_present(run1, [expected_dataset1_storage_id])
|
|
|
|
dataset2 = Dataset(
|
|
name="name",
|
|
digest="digest_other",
|
|
source_type="st2",
|
|
source="source2",
|
|
schema="schema2",
|
|
profile="profile2",
|
|
)
|
|
expected_dataset2_storage_id = "419804e8e153199481c3e509de1fef8f"
|
|
store.log_inputs(run2.info.run_id, [DatasetInput(dataset2)])
|
|
assert_expected_dataset_storage_ids_present([
|
|
expected_dataset1_storage_id,
|
|
expected_dataset2_storage_id,
|
|
])
|
|
assert_expected_input_storage_ids_present(run2, [expected_dataset2_storage_id])
|
|
|
|
dataset3 = Dataset(
|
|
name="name_other",
|
|
digest="digest",
|
|
source_type="st",
|
|
source="source",
|
|
schema="schema",
|
|
profile="profile",
|
|
)
|
|
expected_dataset3_storage_id = "bc5dd0841d8898512d988fe3f984313c"
|
|
store.log_inputs(
|
|
run2.info.run_id,
|
|
[DatasetInput(dataset1), DatasetInput(dataset2), DatasetInput(dataset3, tags)],
|
|
)
|
|
assert_expected_dataset_storage_ids_present([
|
|
expected_dataset1_storage_id,
|
|
expected_dataset2_storage_id,
|
|
expected_dataset3_storage_id,
|
|
])
|
|
assert_expected_input_storage_ids_present(
|
|
run2,
|
|
[
|
|
expected_dataset1_storage_id,
|
|
expected_dataset2_storage_id,
|
|
expected_dataset3_storage_id,
|
|
],
|
|
)
|
|
|
|
|
|
def test_log_inputs_handles_case_when_no_datasets_are_specified(store):
|
|
exp_id = store.create_experiment("log_input_no_datasets")
|
|
run = store.create_run(
|
|
experiment_id=exp_id,
|
|
user_id="user",
|
|
start_time=0,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
store.log_inputs(run.info.run_id)
|
|
store.log_inputs(run.info.run_id, datasets=None)
|
|
|
|
|
|
def test_search_datasets(store):
|
|
exp_id1 = store.create_experiment("test_search_datasets_1")
|
|
# Create an additional experiment to ensure we filter on specified experiment
|
|
# and search works on multiple experiments.
|
|
exp_id2 = store.create_experiment("test_search_datasets_2")
|
|
|
|
run1 = store.create_run(
|
|
experiment_id=exp_id1,
|
|
user_id="user",
|
|
start_time=1,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
run2 = store.create_run(
|
|
experiment_id=exp_id1,
|
|
user_id="user",
|
|
start_time=2,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
run3 = store.create_run(
|
|
experiment_id=exp_id2,
|
|
user_id="user",
|
|
start_time=3,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
|
|
dataset1 = Dataset(
|
|
name="name1",
|
|
digest="digest1",
|
|
source_type="st1",
|
|
source="source1",
|
|
schema="schema1",
|
|
profile="profile1",
|
|
)
|
|
dataset2 = Dataset(
|
|
name="name2",
|
|
digest="digest2",
|
|
source_type="st2",
|
|
source="source2",
|
|
schema="schema2",
|
|
profile="profile2",
|
|
)
|
|
dataset3 = Dataset(
|
|
name="name3",
|
|
digest="digest3",
|
|
source_type="st3",
|
|
source="source3",
|
|
schema="schema3",
|
|
profile="profile3",
|
|
)
|
|
dataset4 = Dataset(
|
|
name="name4",
|
|
digest="digest4",
|
|
source_type="st4",
|
|
source="source4",
|
|
schema="schema4",
|
|
profile="profile4",
|
|
)
|
|
|
|
test_tag = [InputTag(key=MLFLOW_DATASET_CONTEXT, value="test")]
|
|
train_tag = [InputTag(key=MLFLOW_DATASET_CONTEXT, value="train")]
|
|
eval_tag = [InputTag(key=MLFLOW_DATASET_CONTEXT, value="eval")]
|
|
no_context_tag = [InputTag(key="not_context", value="test")]
|
|
|
|
inputs_run1 = [
|
|
DatasetInput(dataset1, train_tag),
|
|
DatasetInput(dataset2, eval_tag),
|
|
DatasetInput(dataset4, no_context_tag),
|
|
]
|
|
inputs_run2 = [
|
|
DatasetInput(dataset1, train_tag),
|
|
DatasetInput(dataset2, test_tag),
|
|
]
|
|
inputs_run3 = [DatasetInput(dataset3, train_tag)]
|
|
|
|
store.log_inputs(run1.info.run_id, inputs_run1)
|
|
store.log_inputs(run2.info.run_id, inputs_run2)
|
|
store.log_inputs(run3.info.run_id, inputs_run3)
|
|
|
|
# Verify actual and expected results are same size and that all elements are equal.
|
|
def assert_has_same_elements(actual_list, expected_list):
|
|
assert len(actual_list) == len(expected_list)
|
|
for actual in actual_list:
|
|
# Verify the expected results list contains same element.
|
|
isEqual = False
|
|
for expected in expected_list:
|
|
isEqual = actual == expected
|
|
if isEqual:
|
|
break
|
|
assert isEqual
|
|
|
|
# Verify no results from exp_id2 are returned.
|
|
results = store._search_datasets([exp_id1])
|
|
expected_results = [
|
|
_DatasetSummary(exp_id1, dataset1.name, dataset1.digest, "train"),
|
|
_DatasetSummary(exp_id1, dataset2.name, dataset2.digest, "eval"),
|
|
_DatasetSummary(exp_id1, dataset2.name, dataset2.digest, "test"),
|
|
_DatasetSummary(exp_id1, dataset4.name, dataset4.digest, None),
|
|
]
|
|
assert_has_same_elements(results, expected_results)
|
|
|
|
# Verify results from both experiment are returned.
|
|
results = store._search_datasets([exp_id1, exp_id2])
|
|
expected_results.append(_DatasetSummary(exp_id2, dataset3.name, dataset3.digest, "train"))
|
|
assert_has_same_elements(results, expected_results)
|
|
|
|
|
|
def test_search_datasets_returns_no_more_than_max_results(store):
|
|
exp_id = store.create_experiment("test_search_datasets")
|
|
run = store.create_run(
|
|
experiment_id=exp_id,
|
|
user_id="user",
|
|
start_time=1,
|
|
tags=[],
|
|
run_name=None,
|
|
)
|
|
inputs = []
|
|
# We intentionally add more than 1000 datasets here to test we only return 1000.
|
|
for i in range(1010):
|
|
dataset = Dataset(
|
|
name="name" + str(i),
|
|
digest="digest" + str(i),
|
|
source_type="st" + str(i),
|
|
source="source" + str(i),
|
|
schema="schema" + str(i),
|
|
profile="profile" + str(i),
|
|
)
|
|
input_tag = [InputTag(key=MLFLOW_DATASET_CONTEXT, value=str(i))]
|
|
inputs.append(DatasetInput(dataset, input_tag))
|
|
|
|
store.log_inputs(run.info.run_id, inputs)
|
|
|
|
results = store._search_datasets([exp_id])
|
|
assert len(results) == 1000
|
|
|
|
|
|
def test_legacy_start_trace_v2(store):
|
|
exp_id = store.create_experiment("test")
|
|
timestamp_ms = get_current_time_millis()
|
|
tags = {"some_key": "test"}
|
|
trace_info = store.deprecated_start_trace_v2(exp_id, timestamp_ms, {}, tags)
|
|
assert trace_info.request_id is not None
|
|
assert trace_info.experiment_id == exp_id
|
|
assert trace_info.timestamp_ms == timestamp_ms
|
|
assert trace_info.execution_time_ms is None
|
|
assert trace_info.status == TraceStatus.IN_PROGRESS
|
|
assert trace_info.tags == tags
|
|
|
|
with pytest.raises(MlflowException, match=r"Experiment fake_exp_id does not exist."):
|
|
store.deprecated_start_trace_v2("fake_exp_id", timestamp_ms, {}, {})
|
|
|
|
|
|
def test_legacy_end_trace(store_and_trace_info):
|
|
store, trace = store_and_trace_info
|
|
timestamp_ms = get_current_time_millis()
|
|
request_metadata = {
|
|
TraceMetadataKey.INPUTS: {"query": "test"},
|
|
TraceMetadataKey.OUTPUTS: "test",
|
|
TRACE_SCHEMA_VERSION_KEY: "2",
|
|
}
|
|
tags = {TraceTagKey.TRACE_NAME: "mlflow_trace"}
|
|
trace_info = store.deprecated_end_trace_v2(
|
|
trace.request_id, timestamp_ms, TraceStatus.OK, request_metadata, tags
|
|
)
|
|
assert trace_info.request_id == trace.request_id
|
|
assert trace_info.timestamp_ms == trace.timestamp_ms
|
|
assert trace_info.execution_time_ms == timestamp_ms - trace.timestamp_ms
|
|
assert trace_info.status == TraceStatus.OK
|
|
assert trace_info.request_metadata == trace.request_metadata | request_metadata
|
|
assert trace_info.tags == trace.tags | tags
|
|
|
|
with pytest.raises(MlflowException, match=r"Trace with ID 'fake_request_id' not found"):
|
|
store.deprecated_end_trace_v2(
|
|
"fake_request_id", timestamp_ms, TraceStatus.OK, request_metadata, tags
|
|
)
|
|
|
|
|
|
def test_start_trace(store):
|
|
exp_id = store.create_experiment("test_start_trace")
|
|
timestamp_ms = get_current_time_millis()
|
|
trace_info = TraceInfo(
|
|
trace_id=f"tr-{uuid.uuid4()}",
|
|
trace_location=TraceLocation.from_experiment_id(exp_id),
|
|
request_time=timestamp_ms,
|
|
execution_duration=100,
|
|
state=TraceState.OK,
|
|
tags={},
|
|
trace_metadata={},
|
|
client_request_id=f"tr-{uuid.uuid4()}",
|
|
request_preview=None,
|
|
response_preview=None,
|
|
)
|
|
new_trace_info = store.start_trace(trace_info)
|
|
|
|
assert new_trace_info.trace_id == trace_info.trace_id
|
|
assert new_trace_info.experiment_id == exp_id
|
|
assert new_trace_info.timestamp_ms == timestamp_ms
|
|
assert new_trace_info.execution_time_ms == 100
|
|
assert new_trace_info.state == TraceState.OK
|
|
assert new_trace_info.tags["mlflow.artifactLocation"] is not None
|
|
assert new_trace_info.client_request_id == trace_info.client_request_id
|
|
|
|
|
|
def test_get_trace_info(store_and_trace_info):
|
|
store, trace = store_and_trace_info
|
|
trace_info = store.get_trace_info(trace.request_id)
|
|
assert trace_info == trace
|
|
|
|
with pytest.raises(MlflowException, match=r"Trace with ID 'fake_request_id' not found"):
|
|
store.get_trace_info("fake_request_id")
|
|
|
|
mock_trace_info = deepcopy(trace_info)
|
|
mock_trace_info.trace_id = "invalid_request_id"
|
|
with (
|
|
mock.patch(
|
|
"mlflow.store.tracking.file_store.FileStore._get_trace_info_from_dir",
|
|
return_value=mock_trace_info,
|
|
),
|
|
pytest.raises(
|
|
MlflowException,
|
|
match=rf"Trace with ID '{trace.request_id}' metadata is in invalid state.",
|
|
),
|
|
):
|
|
store.get_trace_info(trace.trace_id)
|
|
|
|
|
|
def test_set_trace_tag(store_and_trace_info):
|
|
store, trace = store_and_trace_info
|
|
store.set_trace_tag(trace.trace_id, "some_key", "a")
|
|
trace_info = store.get_trace_info(trace.trace_id)
|
|
assert trace_info.tags["some_key"] == "a"
|
|
|
|
# test overwrite
|
|
store.set_trace_tag(trace.trace_id, "some_key", "test")
|
|
trace_info = store.get_trace_info(trace.trace_id)
|
|
assert trace_info.tags["some_key"] == "test"
|
|
|
|
# test value written as string
|
|
store.set_trace_tag(trace.trace_id, "int_key", 1234)
|
|
trace_info = store.get_trace_info(trace.trace_id)
|
|
assert trace_info.tags["int_key"] == "1234"
|
|
|
|
# test value length
|
|
store.set_trace_tag(trace.trace_id, "key", "v" * MAX_CHARS_IN_TRACE_INFO_TAGS_VALUE)
|
|
trace_info = store.get_trace_info(trace.trace_id)
|
|
assert trace_info.tags["key"] == "v" * MAX_CHARS_IN_TRACE_INFO_TAGS_VALUE
|
|
|
|
with pytest.raises(MlflowException, match=r"Missing value for required parameter \'key\'"):
|
|
store.set_trace_tag(trace.trace_id, None, "test")
|
|
|
|
|
|
def test_delete_trace_tag(store_and_trace_info):
|
|
store, trace = store_and_trace_info
|
|
store.set_trace_tag(trace.trace_id, "some_key", "a")
|
|
store.delete_trace_tag(trace.trace_id, "some_key")
|
|
trace_info = store.get_trace_info(trace.trace_id)
|
|
assert "some_key" not in trace_info.tags
|
|
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=rf"No tag with name: invalid_key in trace with ID {trace.trace_id}.",
|
|
):
|
|
store.delete_trace_tag(trace.trace_id, "invalid_key")
|
|
|
|
|
|
def test_delete_traces(store):
|
|
exp_id = store.create_experiment("test")
|
|
trace_ids = []
|
|
timestamps = list(range(90, -1, -10))
|
|
for i in range(10):
|
|
trace_info = TraceInfo(
|
|
trace_id=f"tr-{uuid.uuid4()}",
|
|
trace_location=TraceLocation.from_experiment_id(exp_id),
|
|
request_time=timestamps[i],
|
|
state=TraceState.OK,
|
|
)
|
|
trace_info = store.start_trace(trace_info)
|
|
trace_ids.append(trace_info.trace_id)
|
|
|
|
assert store.delete_traces(exp_id, max_timestamp_millis=0) == 1
|
|
assert len(store.search_traces([exp_id])[0]) == 9
|
|
|
|
# delete with max_timestamp_millis
|
|
# if max_traces < number of traces with timestamp < max_timestamp_millis,
|
|
# delete older traces first
|
|
assert store.delete_traces(exp_id, max_timestamp_millis=50, max_traces=2) == 2
|
|
assert len(store.search_traces([exp_id])[0]) == 7
|
|
assert store.delete_traces(exp_id, max_timestamp_millis=50) == 4
|
|
assert len(store.search_traces([exp_id])[0]) == 3
|
|
|
|
# delete with trace_ids
|
|
assert store.delete_traces(exp_id, trace_ids=[trace_ids[3]]) == 1
|
|
assert len(store.search_traces([exp_id])[0]) == 2
|
|
assert store.delete_traces(exp_id, trace_ids=["non_existing_trace_id"]) == 0
|
|
assert len(store.search_traces([exp_id])[0]) == 2
|
|
assert store.delete_traces(exp_id, trace_ids=trace_ids) == 2
|
|
assert len(store.search_traces([exp_id])[0]) == 0
|
|
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=r"Either `max_timestamp_millis` or `trace_ids` must be specified.",
|
|
):
|
|
store.delete_traces(exp_id)
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=r"Only one of `max_timestamp_millis` and `trace_ids` can be specified.",
|
|
):
|
|
store.delete_traces(exp_id, max_timestamp_millis=100, trace_ids=trace_ids)
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=r"`max_traces` can't be specified if `trace_ids` is specified.",
|
|
):
|
|
store.delete_traces(exp_id, max_traces=2, trace_ids=trace_ids)
|
|
with pytest.raises(
|
|
MlflowException, match=r"`max_traces` must be a positive integer, received 0"
|
|
):
|
|
store.delete_traces(exp_id, 100, max_traces=0)
|
|
with pytest.raises(MlflowException, match=r"Experiment non_existing_exp does not exist."):
|
|
store.delete_traces("non_existing_exp", 100, 2)
|
|
|
|
|
|
def _validate_search_traces(store, exp_ids, filter_string, expected_traces, order_by=None):
|
|
traces, _ = store.search_traces(exp_ids, filter_string, order_by=order_by)
|
|
assert traces == expected_traces
|
|
|
|
|
|
def test_search_traces(store):
|
|
traces, _token = store.search_traces(["0"])
|
|
assert traces == []
|
|
|
|
|
|
def test_search_traces_filter(generate_trace_infos):
|
|
trace_infos = generate_trace_infos.trace_infos
|
|
store = generate_trace_infos.store
|
|
exp_id = generate_trace_infos.exp_id
|
|
trace_ids = generate_trace_infos.trace_ids
|
|
|
|
# by default sort by timestamp_ms DESC, request_id ASC
|
|
_validate_search_traces(store, [exp_id], None, trace_infos[::-1])
|
|
_validate_search_traces(store, [exp_id], "", trace_infos[::-1])
|
|
|
|
# filter by name
|
|
_validate_search_traces(store, [exp_id], "name = 'trace_0'", trace_infos[:1])
|
|
_validate_search_traces(store, [exp_id], "name != 'trace_0'", trace_infos[1:][::-1])
|
|
|
|
# filter by status
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id],
|
|
"status IN ('IN_PROGRESS', 'OK')",
|
|
(trace_infos[:5])[::-1],
|
|
)
|
|
_validate_search_traces(
|
|
store, [exp_id], "status NOT IN ('IN_PROGRESS', 'OK')", trace_infos[5:][::-1]
|
|
)
|
|
# filter by status w/ attributes. or trace. prefix
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id],
|
|
"trace.status = 'ERROR'",
|
|
trace_infos[5:][::-1],
|
|
)
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id],
|
|
"attributes.status IN ('IN_PROGRESS', 'OK')",
|
|
(trace_infos[:5])[::-1],
|
|
)
|
|
|
|
# filter by timestamp
|
|
for timestamp_key in ["timestamp", "timestamp_ms"]:
|
|
_validate_search_traces(store, [exp_id], f"{timestamp_key} < 10", trace_infos[:1])
|
|
_validate_search_traces(store, [exp_id], f"{timestamp_key} <= 0", trace_infos[:1])
|
|
_validate_search_traces(store, [exp_id], f"{timestamp_key} > 0", trace_infos[1:][::-1])
|
|
_validate_search_traces(store, [exp_id], f"{timestamp_key} >= 10", trace_infos[1:][::-1])
|
|
_validate_search_traces(store, [exp_id], f"{timestamp_key} = 100", [])
|
|
_validate_search_traces(store, [exp_id], f"{timestamp_key} != 100", trace_infos[::-1])
|
|
|
|
# filter by request_id
|
|
_validate_search_traces(store, [exp_id], f"request_id = '{trace_ids[0]}'", [trace_infos[0]])
|
|
_validate_search_traces(
|
|
store, [exp_id], f"request_id != '{trace_ids[0]}'", trace_infos[1:][::-1]
|
|
)
|
|
_validate_search_traces(store, [exp_id], f"request_id IN ('{trace_ids[0]}')", [trace_infos[0]])
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id],
|
|
f"request_id NOT IN ('{trace_ids[0]}')",
|
|
trace_infos[1:][::-1],
|
|
)
|
|
|
|
# filter by execution_time
|
|
for execution_time_key in ["execution_time", "execution_time_ms"]:
|
|
_validate_search_traces(
|
|
store, [exp_id], f"{execution_time_key} = 10", trace_infos[:5][::-1]
|
|
)
|
|
# value None is always seen as not-match
|
|
_validate_search_traces(
|
|
store, [exp_id], f"{execution_time_key} != 10", trace_infos[5:][::-1]
|
|
)
|
|
_validate_search_traces(
|
|
store, [exp_id], f"{execution_time_key} > 10", trace_infos[5:][::-1]
|
|
)
|
|
_validate_search_traces(store, [exp_id], f"{execution_time_key} < 10", [])
|
|
_validate_search_traces(store, [exp_id], f"{execution_time_key} >= 10", trace_infos[::-1])
|
|
_validate_search_traces(
|
|
store, [exp_id], f"{execution_time_key} <= 10", trace_infos[:5][::-1]
|
|
)
|
|
|
|
# filter by run_id
|
|
_validate_search_traces(store, [exp_id], "run_id = 'run_5'", [trace_infos[5]])
|
|
_validate_search_traces(store, [exp_id], "run_id != 'run_5'", trace_infos[6:][::-1])
|
|
|
|
# filter by tag
|
|
for tag_identifier in ["tag", "tags"]:
|
|
_validate_search_traces(
|
|
store, [exp_id], f"{tag_identifier}.test_tag = 'tag_0'", [trace_infos[0]]
|
|
)
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id],
|
|
f"{tag_identifier}.test_tag != 'tag_0'",
|
|
trace_infos[1:][::-1],
|
|
)
|
|
_validate_search_traces(store, [exp_id], f"{tag_identifier}.test_tag = '123'", [])
|
|
|
|
# multiple filter conditions
|
|
_validate_search_traces(
|
|
store, [exp_id], "status = 'OK' AND timestamp <= 10", trace_infos[:2][::-1]
|
|
)
|
|
|
|
|
|
def test_search_traces_filter_trace_metadata(store):
|
|
exp_id = store.create_experiment("test")
|
|
timestamp_ms_1 = get_current_time_millis()
|
|
trace_info_1 = store.start_trace(
|
|
TraceInfo(
|
|
trace_id=f"tr-{uuid.uuid4()}",
|
|
trace_location=TraceLocation.from_experiment_id(exp_id),
|
|
request_time=timestamp_ms_1,
|
|
state=TraceState.OK,
|
|
trace_metadata={
|
|
TraceMetadataKey.INPUTS: "inputs1",
|
|
TraceMetadataKey.OUTPUTS: "outputs1",
|
|
},
|
|
),
|
|
)
|
|
time.sleep(0.001) # ensure unique timestamps
|
|
timestamp_ms_2 = get_current_time_millis()
|
|
trace_info_2 = store.start_trace(
|
|
TraceInfo(
|
|
trace_id=f"tr-{uuid.uuid4()}",
|
|
trace_location=TraceLocation.from_experiment_id(exp_id),
|
|
request_time=timestamp_ms_2,
|
|
state=TraceState.OK,
|
|
trace_metadata={
|
|
TraceMetadataKey.INPUTS: "inputs2",
|
|
TraceMetadataKey.OUTPUTS: "outputs2",
|
|
},
|
|
),
|
|
)
|
|
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id],
|
|
f"request_metadata.{TraceMetadataKey.INPUTS} = 'inputs1'",
|
|
[trace_info_1],
|
|
)
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id],
|
|
f"request_metadata.{TraceMetadataKey.OUTPUTS} = 'outputs1'",
|
|
[trace_info_1],
|
|
)
|
|
# not equal
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id],
|
|
f"request_metadata.{TraceMetadataKey.INPUTS} != 'inputs1'",
|
|
[trace_info_2],
|
|
)
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id],
|
|
f"request_metadata.{TraceMetadataKey.INPUTS} != 'test'",
|
|
[trace_info_2, trace_info_1],
|
|
)
|
|
|
|
# backtick
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id],
|
|
f"request_metadata.`{TraceMetadataKey.INPUTS}` = 'inputs1'",
|
|
[trace_info_1],
|
|
)
|
|
|
|
# alias
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id],
|
|
f"metadata.{TraceMetadataKey.INPUTS} = 'inputs1'",
|
|
[trace_info_1],
|
|
)
|
|
|
|
|
|
def test_search_traces_with_like_ilike_filters(generate_trace_infos):
|
|
trace_infos = generate_trace_infos.trace_infos
|
|
store = generate_trace_infos.store
|
|
exp_id = generate_trace_infos.exp_id
|
|
|
|
# Test LIKE operator for trace name (case-sensitive)
|
|
_validate_search_traces(store, [exp_id], "name LIKE 'trace_%'", trace_infos[::-1])
|
|
_validate_search_traces(store, [exp_id], "name LIKE 'trace_0'", [trace_infos[0]])
|
|
_validate_search_traces(store, [exp_id], "name LIKE 'trace_1%'", [trace_infos[1]])
|
|
_validate_search_traces(store, [exp_id], "name LIKE 'TRACE_%'", []) # case-sensitive
|
|
|
|
# Test ILIKE operator for trace name (case-insensitive)
|
|
_validate_search_traces(store, [exp_id], "name ILIKE 'TRACE_%'", trace_infos[::-1])
|
|
_validate_search_traces(store, [exp_id], "name ILIKE 'TRACE_0'", [trace_infos[0]])
|
|
_validate_search_traces(store, [exp_id], "name ILIKE 'TrAcE_1'", [trace_infos[1]])
|
|
|
|
# Test LIKE operator for tags
|
|
_validate_search_traces(store, [exp_id], "tag.test_tag LIKE 'tag_%'", trace_infos[::-1])
|
|
_validate_search_traces(store, [exp_id], "tag.test_tag LIKE 'tag_0'", [trace_infos[0]])
|
|
_validate_search_traces(store, [exp_id], "tag.test_tag LIKE 'TAG_%'", []) # case-sensitive
|
|
|
|
# Test ILIKE operator for tags using both 'tag' and 'tags' prefix
|
|
_validate_search_traces(store, [exp_id], "tag.test_tag ILIKE 'TAG_%'", trace_infos[::-1])
|
|
_validate_search_traces(store, [exp_id], "tags.test_tag ILIKE 'TAG_0'", [trace_infos[0]])
|
|
|
|
# Test LIKE/ILIKE for run_id
|
|
_validate_search_traces(store, [exp_id], "run_id LIKE 'run_%'", trace_infos[5:][::-1])
|
|
_validate_search_traces(store, [exp_id], "run_id LIKE 'run_5'", [trace_infos[5]])
|
|
_validate_search_traces(store, [exp_id], "run_id ILIKE 'RUN_5'", [trace_infos[5]])
|
|
_validate_search_traces(store, [exp_id], "run_id ILIKE 'RUN_%'", trace_infos[5:][::-1])
|
|
|
|
# Test combined filters with LIKE/ILIKE
|
|
_validate_search_traces(
|
|
store, [exp_id], "name LIKE 'trace_%' AND status = 'OK'", trace_infos[:5][::-1]
|
|
)
|
|
_validate_search_traces(
|
|
store, [exp_id], "tag.test_tag ILIKE 'TAG_%' AND timestamp < 20", trace_infos[:2][::-1]
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("filter_string", "error"),
|
|
[
|
|
("invalid", r"Invalid clause\(s\) in filter string"),
|
|
("name = 'foo' AND invalid", r"Invalid clause\(s\) in filter string"),
|
|
("foo.bar = 'baz'", r"Invalid entity type 'foo'"),
|
|
("invalid = 'foo'", r"Invalid attribute key 'invalid'"),
|
|
("trace.tags.foo = 'bar'", r"Invalid attribute key 'tags\.foo'"),
|
|
("trace.status < 'OK'", r"Invalid comparator '<'"),
|
|
("name IN ('foo', 'bar')", r"Invalid comparator 'IN'"),
|
|
("feedback.correctness = 'true'", r"Assessment filtering requires database support"),
|
|
],
|
|
)
|
|
def test_search_traces_invalid_filter(generate_trace_infos, filter_string, error):
|
|
store = generate_trace_infos.store
|
|
exp_id = generate_trace_infos.exp_id
|
|
|
|
# Invalid filter key
|
|
with pytest.raises(MlflowException, match=error):
|
|
store.search_traces([exp_id], filter_string)
|
|
|
|
|
|
def test_search_traces_order(generate_trace_infos):
|
|
trace_infos = generate_trace_infos.trace_infos
|
|
store = generate_trace_infos.store
|
|
exp_id = generate_trace_infos.exp_id
|
|
timestamps = generate_trace_infos.timestamps
|
|
# order by timestamp
|
|
for timestamp_key in ["timestamp", "timestamp_ms"]:
|
|
_validate_search_traces(store, [exp_id], "", trace_infos, order_by=[f"{timestamp_key} ASC"])
|
|
_validate_search_traces(
|
|
store, [exp_id], "", trace_infos[::-1], order_by=[f"{timestamp_key} DESC"]
|
|
)
|
|
|
|
# order by execution time
|
|
for execution_time_key in ["execution_time", "execution_time_ms"]:
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id],
|
|
"",
|
|
trace_infos[::-1],
|
|
order_by=[f"{execution_time_key} DESC"],
|
|
)
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id],
|
|
"",
|
|
trace_infos[:5][::-1] + trace_infos[5:][::-1],
|
|
order_by=[f"{execution_time_key} ASC"],
|
|
)
|
|
|
|
# order by status
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id],
|
|
"",
|
|
trace_infos[:5][::-1] + trace_infos[5:][::-1],
|
|
order_by=["status DESC"],
|
|
)
|
|
_validate_search_traces(store, [exp_id], "", trace_infos[::-1], order_by=["status ASC"])
|
|
|
|
# order by request_id
|
|
expected_trace_infos = sorted(trace_infos, key=lambda x: x.request_id)
|
|
_validate_search_traces(store, [exp_id], "", expected_trace_infos, order_by=["request_id ASC"])
|
|
expected_trace_infos = sorted(trace_infos, key=lambda x: x.request_id, reverse=True)
|
|
_validate_search_traces(store, [exp_id], "", expected_trace_infos, order_by=["request_id DESC"])
|
|
|
|
# order by experiment_id
|
|
exp_id2 = store.create_experiment("test2")
|
|
trace_info = store.start_trace(
|
|
TraceInfo(
|
|
trace_id=f"tr-{uuid.uuid4()}",
|
|
trace_location=TraceLocation.from_experiment_id(exp_id2),
|
|
request_time=timestamps[-1],
|
|
state=TraceState.OK,
|
|
),
|
|
)
|
|
trace_infos.append(trace_info)
|
|
|
|
order = exp_id2 > exp_id
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id, exp_id2],
|
|
"",
|
|
trace_infos[::-1] if order else trace_infos[:10][::-1] + [trace_infos[-1]],
|
|
order_by=["experiment_id DESC"],
|
|
)
|
|
_validate_search_traces(
|
|
store,
|
|
[exp_id, exp_id2],
|
|
"",
|
|
trace_infos[:10][::-1] + [trace_infos[-1]] if order else trace_infos[::-1],
|
|
order_by=["experiment_id ASC"],
|
|
)
|
|
|
|
|
|
def test_search_traces_raise_errors(generate_trace_infos):
|
|
store = generate_trace_infos.store
|
|
exp_id = generate_trace_infos.exp_id
|
|
|
|
# unsupported order_by keys
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=r"Invalid order_by entity `tag` with key `mlflow.traceName`",
|
|
):
|
|
store.search_traces([exp_id], "", order_by=["name DESC"])
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=r"Invalid order_by entity `request_metadata` "
|
|
rf"with key `{TraceMetadataKey.SOURCE_RUN}`",
|
|
):
|
|
store.search_traces([exp_id], "", order_by=["run_id ASC"])
|
|
|
|
|
|
def test_search_traces_pagination(generate_trace_infos):
|
|
trace_infos = generate_trace_infos.trace_infos
|
|
store = generate_trace_infos.store
|
|
exp_id = generate_trace_infos.exp_id
|
|
|
|
# test returned token behavior
|
|
traces, token = store.search_traces([exp_id], None, max_results=5)
|
|
assert traces == trace_infos[::-1][:5]
|
|
assert token is not None
|
|
traces, token = store.search_traces([exp_id], None, max_results=5, page_token=token)
|
|
assert traces == trace_infos[::-1][5:]
|
|
assert token is None
|
|
|
|
|
|
def test_traces_not_listed_as_runs(tmp_path):
|
|
with _use_tracking_uri(tmp_path.joinpath("mlruns").as_uri()):
|
|
client = mlflow.MlflowClient()
|
|
with mlflow.start_run() as run:
|
|
client.start_trace("test")
|
|
|
|
with mock.patch("mlflow.store.tracking.file_store.logging.debug") as mock_debug:
|
|
client.search_runs([run.info.experiment_id], "", ViewType.ALL, max_results=1)
|
|
mock_debug.assert_not_called()
|
|
|
|
|
|
def test_create_and_get_assessment(store_and_trace_info):
|
|
store, trace_info = store_and_trace_info
|
|
|
|
feedback = Feedback(
|
|
trace_id=trace_info.request_id,
|
|
name="correctness",
|
|
value=True,
|
|
rationale="The response is correct and well-formatted",
|
|
source=AssessmentSource(
|
|
source_type=AssessmentSourceType.HUMAN, source_id="evaluator@company.com"
|
|
),
|
|
metadata={"project": "test-project", "version": "1.0"},
|
|
span_id="span-123",
|
|
)
|
|
|
|
created_feedback = store.create_assessment(feedback)
|
|
assert created_feedback.assessment_id is not None
|
|
assert created_feedback.assessment_id.startswith("a-")
|
|
assert created_feedback.trace_id == trace_info.request_id
|
|
assert created_feedback.create_time_ms is not None
|
|
assert created_feedback.name == "correctness"
|
|
assert created_feedback.value is True
|
|
assert created_feedback.rationale == "The response is correct and well-formatted"
|
|
assert created_feedback.metadata == {"project": "test-project", "version": "1.0"}
|
|
assert created_feedback.span_id == "span-123"
|
|
assert created_feedback.valid
|
|
|
|
expectation = Expectation(
|
|
trace_id=trace_info.request_id,
|
|
name="expected_response",
|
|
value="The capital of France is Paris.",
|
|
source=AssessmentSource(
|
|
source_type=AssessmentSourceType.HUMAN, source_id="annotator@company.com"
|
|
),
|
|
metadata={"context": "geography-qa", "difficulty": "easy"},
|
|
span_id="span-456",
|
|
)
|
|
|
|
created_expectation = store.create_assessment(expectation)
|
|
assert created_expectation.assessment_id != created_feedback.assessment_id
|
|
assert created_expectation.trace_id == trace_info.request_id
|
|
assert created_expectation.value == "The capital of France is Paris."
|
|
assert created_expectation.metadata == {"context": "geography-qa", "difficulty": "easy"}
|
|
assert created_expectation.span_id == "span-456"
|
|
assert created_expectation.valid
|
|
|
|
retrieved_feedback = store.get_assessment(trace_info.request_id, created_feedback.assessment_id)
|
|
assert retrieved_feedback.name == "correctness"
|
|
assert retrieved_feedback.value is True
|
|
assert retrieved_feedback.rationale == "The response is correct and well-formatted"
|
|
assert retrieved_feedback.metadata == {"project": "test-project", "version": "1.0"}
|
|
assert retrieved_feedback.span_id == "span-123"
|
|
assert retrieved_feedback.trace_id == trace_info.request_id
|
|
assert retrieved_feedback.valid
|
|
|
|
retrieved_expectation = store.get_assessment(
|
|
trace_info.request_id, created_expectation.assessment_id
|
|
)
|
|
assert retrieved_expectation.value == "The capital of France is Paris."
|
|
assert retrieved_expectation.metadata == {"context": "geography-qa", "difficulty": "easy"}
|
|
assert retrieved_expectation.span_id == "span-456"
|
|
assert retrieved_expectation.trace_id == trace_info.request_id
|
|
assert retrieved_expectation.valid is None
|
|
|
|
|
|
def test_get_assessment_errors(store_and_trace_info):
|
|
store, trace_info = store_and_trace_info
|
|
|
|
with pytest.raises(MlflowException, match=r"Trace with ID 'fake_trace' not found"):
|
|
store.get_assessment("fake_trace", "fake_assessment")
|
|
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=r"Assessment with ID 'fake_assessment' not found for trace",
|
|
):
|
|
store.get_assessment(trace_info.request_id, "fake_assessment")
|
|
|
|
|
|
def test_update_assessment_feedback(store_and_trace_info):
|
|
store, trace_info = store_and_trace_info
|
|
|
|
original_feedback = Feedback(
|
|
trace_id=trace_info.request_id,
|
|
name="correctness",
|
|
value=True,
|
|
rationale="Original rationale",
|
|
source=AssessmentSource(
|
|
source_type=AssessmentSourceType.HUMAN, source_id="evaluator@company.com"
|
|
),
|
|
metadata={"project": "test-project", "version": "1.0"},
|
|
span_id="span-123",
|
|
)
|
|
|
|
created_feedback = store.create_assessment(original_feedback)
|
|
original_id = created_feedback.assessment_id
|
|
|
|
updated_feedback = store.update_assessment(
|
|
trace_id=trace_info.request_id,
|
|
assessment_id=original_id,
|
|
name="correctness_updated",
|
|
feedback=FeedbackValue(value=False),
|
|
rationale="Updated rationale",
|
|
metadata={"project": "test-project", "version": "2.0", "new_field": "added"},
|
|
)
|
|
|
|
assert updated_feedback.assessment_id == original_id
|
|
assert updated_feedback.name == "correctness_updated"
|
|
assert updated_feedback.value is False
|
|
assert updated_feedback.rationale == "Updated rationale"
|
|
assert updated_feedback.metadata == {
|
|
"project": "test-project",
|
|
"version": "2.0",
|
|
"new_field": "added",
|
|
}
|
|
assert updated_feedback.span_id == "span-123"
|
|
assert updated_feedback.source.source_id == "evaluator@company.com"
|
|
assert updated_feedback.valid is True
|
|
|
|
retrieved = store.get_assessment(trace_info.request_id, original_id)
|
|
assert retrieved.value is False
|
|
assert retrieved.name == "correctness_updated"
|
|
assert retrieved.rationale == "Updated rationale"
|
|
|
|
|
|
def test_update_assessment_expectation(store_and_trace_info):
|
|
store, trace_info = store_and_trace_info
|
|
|
|
original_expectation = Expectation(
|
|
trace_id=trace_info.request_id,
|
|
name="expected_response",
|
|
value="The capital of France is Paris.",
|
|
source=AssessmentSource(
|
|
source_type=AssessmentSourceType.HUMAN, source_id="annotator@company.com"
|
|
),
|
|
metadata={"context": "geography-qa"},
|
|
span_id="span-456",
|
|
)
|
|
|
|
created_expectation = store.create_assessment(original_expectation)
|
|
original_id = created_expectation.assessment_id
|
|
|
|
updated_expectation = store.update_assessment(
|
|
trace_id=trace_info.request_id,
|
|
assessment_id=original_id,
|
|
expectation=ExpectationValue(value="The capital and largest city of France is Paris."),
|
|
metadata={"context": "geography-qa", "updated": "true"},
|
|
)
|
|
|
|
assert updated_expectation.assessment_id == original_id
|
|
assert updated_expectation.name == "expected_response"
|
|
assert updated_expectation.value == "The capital and largest city of France is Paris."
|
|
assert updated_expectation.metadata == {"context": "geography-qa", "updated": "true"}
|
|
assert updated_expectation.span_id == "span-456"
|
|
assert updated_expectation.source.source_id == "annotator@company.com"
|
|
|
|
|
|
def test_update_assessment_partial_fields(store_and_trace_info):
|
|
store, trace_info = store_and_trace_info
|
|
|
|
original_feedback = Feedback(
|
|
trace_id=trace_info.request_id,
|
|
name="quality",
|
|
value=5,
|
|
rationale="Original rationale",
|
|
source=AssessmentSource(source_type=AssessmentSourceType.CODE),
|
|
metadata={"scorer": "automated"},
|
|
)
|
|
|
|
created_feedback = store.create_assessment(original_feedback)
|
|
original_id = created_feedback.assessment_id
|
|
|
|
updated_feedback = store.update_assessment(
|
|
trace_id=trace_info.request_id,
|
|
assessment_id=original_id,
|
|
rationale="Updated rationale only",
|
|
)
|
|
|
|
assert updated_feedback.assessment_id == original_id
|
|
assert updated_feedback.name == "quality"
|
|
assert updated_feedback.value == 5
|
|
assert updated_feedback.rationale == "Updated rationale only"
|
|
assert updated_feedback.metadata == {"scorer": "automated"}
|
|
|
|
|
|
def test_update_assessment_type_validation(store_and_trace_info):
|
|
store, trace_info = store_and_trace_info
|
|
|
|
feedback = Feedback(
|
|
trace_id=trace_info.request_id,
|
|
name="test_feedback",
|
|
value="original",
|
|
source=AssessmentSource(source_type=AssessmentSourceType.CODE),
|
|
)
|
|
created_feedback = store.create_assessment(feedback)
|
|
|
|
with pytest.raises(
|
|
MlflowException, match=r"Cannot update expectation value on a Feedback assessment"
|
|
):
|
|
store.update_assessment(
|
|
trace_id=trace_info.request_id,
|
|
assessment_id=created_feedback.assessment_id,
|
|
expectation=ExpectationValue(value="This should fail"),
|
|
)
|
|
|
|
expectation = Expectation(
|
|
trace_id=trace_info.request_id,
|
|
name="test_expectation",
|
|
value="original_expected",
|
|
source=AssessmentSource(source_type=AssessmentSourceType.HUMAN),
|
|
)
|
|
created_expectation = store.create_assessment(expectation)
|
|
|
|
with pytest.raises(
|
|
MlflowException, match=r"Cannot update feedback value on an Expectation assessment"
|
|
):
|
|
store.update_assessment(
|
|
trace_id=trace_info.request_id,
|
|
assessment_id=created_expectation.assessment_id,
|
|
feedback=FeedbackValue(value="This should fail"),
|
|
)
|
|
|
|
|
|
def test_update_assessment_errors(store_and_trace_info):
|
|
store, trace_info = store_and_trace_info
|
|
|
|
with pytest.raises(MlflowException, match=r"Trace with ID 'fake_trace' not found"):
|
|
store.update_assessment(
|
|
trace_id="fake_trace", assessment_id="fake_assessment", rationale="This should fail"
|
|
)
|
|
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=r"Assessment with ID 'fake_assessment' not found for trace",
|
|
):
|
|
store.update_assessment(
|
|
trace_id=trace_info.request_id,
|
|
assessment_id="fake_assessment",
|
|
rationale="This should fail",
|
|
)
|
|
|
|
|
|
def test_update_assessment_metadata_merging(store_and_trace_info):
|
|
store, trace_info = store_and_trace_info
|
|
|
|
original = Feedback(
|
|
trace_id=trace_info.request_id,
|
|
name="test",
|
|
value="original",
|
|
source=AssessmentSource(source_type=AssessmentSourceType.CODE),
|
|
metadata={"keep": "this", "override": "old_value", "remove_me": "will_stay"},
|
|
)
|
|
|
|
created = store.create_assessment(original)
|
|
|
|
updated = store.update_assessment(
|
|
trace_id=trace_info.request_id,
|
|
assessment_id=created.assessment_id,
|
|
metadata={"override": "new_value", "new_key": "new_value"},
|
|
)
|
|
|
|
expected_metadata = {
|
|
"keep": "this",
|
|
"override": "new_value",
|
|
"remove_me": "will_stay",
|
|
"new_key": "new_value",
|
|
}
|
|
assert updated.metadata == expected_metadata
|
|
|
|
|
|
def test_update_assessment_timestamps(store_and_trace_info):
|
|
store, trace_info = store_and_trace_info
|
|
|
|
original = Feedback(
|
|
trace_id=trace_info.request_id,
|
|
name="test",
|
|
value="original",
|
|
source=AssessmentSource(source_type=AssessmentSourceType.CODE),
|
|
)
|
|
|
|
created = store.create_assessment(original)
|
|
original_create_time = created.create_time_ms
|
|
original_update_time = created.last_update_time_ms
|
|
|
|
time.sleep(0.001)
|
|
|
|
updated = store.update_assessment(
|
|
trace_id=trace_info.request_id,
|
|
assessment_id=created.assessment_id,
|
|
name="updated_name",
|
|
)
|
|
|
|
assert updated.create_time_ms == original_create_time
|
|
assert updated.last_update_time_ms > original_update_time
|
|
|
|
|
|
def test_create_assessment_with_overrides(store_and_trace_info):
|
|
store, trace_info = store_and_trace_info
|
|
|
|
original_feedback = Feedback(
|
|
trace_id=trace_info.request_id,
|
|
name="quality",
|
|
value="poor",
|
|
source=AssessmentSource(source_type=AssessmentSourceType.LLM_JUDGE),
|
|
)
|
|
|
|
created_original = store.create_assessment(original_feedback)
|
|
|
|
override_feedback = Feedback(
|
|
trace_id=trace_info.request_id,
|
|
name="quality",
|
|
value="excellent",
|
|
source=AssessmentSource(source_type=AssessmentSourceType.HUMAN),
|
|
overrides=created_original.assessment_id,
|
|
)
|
|
|
|
created_override = store.create_assessment(override_feedback)
|
|
|
|
assert created_override.overrides == created_original.assessment_id
|
|
assert created_override.value == "excellent"
|
|
assert created_override.valid is True
|
|
|
|
retrieved_original = store.get_assessment(trace_info.request_id, created_original.assessment_id)
|
|
assert retrieved_original.valid is False
|
|
assert retrieved_original.value == "poor"
|
|
|
|
|
|
def test_create_assessment_override_nonexistent(store_and_trace_info):
|
|
store, trace_info = store_and_trace_info
|
|
|
|
override_feedback = Feedback(
|
|
trace_id=trace_info.request_id,
|
|
name="quality",
|
|
value="excellent",
|
|
source=AssessmentSource(source_type=AssessmentSourceType.HUMAN),
|
|
overrides="nonexistent-assessment-id",
|
|
)
|
|
|
|
with pytest.raises(
|
|
MlflowException, match=r"Assessment with ID 'nonexistent-assessment-id' not found"
|
|
):
|
|
store.create_assessment(override_feedback)
|
|
|
|
|
|
def test_delete_assessment_idempotent(store_and_trace_info):
|
|
store, trace_info = store_and_trace_info
|
|
|
|
feedback = Feedback(
|
|
trace_id=trace_info.request_id,
|
|
name="test",
|
|
value="test_value",
|
|
source=AssessmentSource(source_type=AssessmentSourceType.CODE),
|
|
)
|
|
|
|
created_feedback = store.create_assessment(feedback)
|
|
|
|
retrieved = store.get_assessment(trace_info.request_id, created_feedback.assessment_id)
|
|
assert retrieved.assessment_id == created_feedback.assessment_id
|
|
|
|
store.delete_assessment(trace_info.request_id, created_feedback.assessment_id)
|
|
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=rf"Assessment with ID '{created_feedback.assessment_id}' not found for trace",
|
|
):
|
|
store.get_assessment(trace_info.request_id, created_feedback.assessment_id)
|
|
|
|
store.delete_assessment(trace_info.request_id, created_feedback.assessment_id)
|
|
store.delete_assessment(trace_info.request_id, "fake_assessment_id")
|
|
|
|
|
|
def test_delete_assessment_override_behavior(store_and_trace_info):
|
|
store, trace_info = store_and_trace_info
|
|
|
|
original = store.create_assessment(
|
|
Feedback(
|
|
trace_id=trace_info.request_id,
|
|
name="original",
|
|
value="original_value",
|
|
source=AssessmentSource(source_type=AssessmentSourceType.CODE),
|
|
),
|
|
)
|
|
|
|
override = store.create_assessment(
|
|
Feedback(
|
|
trace_id=trace_info.request_id,
|
|
name="override",
|
|
value="override_value",
|
|
source=AssessmentSource(source_type=AssessmentSourceType.HUMAN),
|
|
overrides=original.assessment_id,
|
|
),
|
|
)
|
|
|
|
assert store.get_assessment(trace_info.request_id, original.assessment_id).valid is False
|
|
assert store.get_assessment(trace_info.request_id, override.assessment_id).valid is True
|
|
|
|
store.delete_assessment(trace_info.request_id, override.assessment_id)
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
store.get_assessment(trace_info.request_id, override.assessment_id)
|
|
assert store.get_assessment(trace_info.request_id, original.assessment_id).valid is True
|
|
|
|
|
|
def test_get_experiment_missing_and_empty_metadata_file(tmp_path):
|
|
fs = FileStore(str(tmp_path))
|
|
|
|
exp_id = "Demo_Experiment"
|
|
exp_dir = tmp_path / exp_id
|
|
exp_dir.mkdir()
|
|
|
|
# Missing meta.yaml — should raise MissingConfigException about missing file
|
|
with pytest.raises(
|
|
MissingConfigException, match=rf"Yaml file '.*{exp_id}[\\/]+meta.yaml' does not exist."
|
|
):
|
|
fs._get_experiment(exp_id)
|
|
|
|
# Create an empty meta.yaml
|
|
(exp_dir / FileStore.META_DATA_FILE_NAME).write_text("")
|
|
# Should raise MissingConfigException about invalid metadata
|
|
with pytest.raises(MissingConfigException, match=rf"Experiment {exp_id} is invalid with empty"):
|
|
fs._get_experiment(exp_id)
|
|
|
|
|
|
def test_malicious_meta_yaml_in_artifact_folder_path_traversal(tmp_path):
|
|
"""
|
|
Regression test for ZDI-CAN-26649: Directory traversal via malicious meta.yaml.
|
|
|
|
Attack flow that should be blocked:
|
|
1. Create experiment with artifact_location pointing to FileStore root
|
|
2. Create a run - artifacts go to {root}/{run_id}/artifacts/
|
|
3. Plant malicious meta.yaml in artifacts folder with arbitrary artifact_uri
|
|
4. Try to use "artifacts" as run_uuid to access files via the malicious artifact_uri
|
|
|
|
The fix validates that run directories have required subdirectories (metrics/, params/,
|
|
artifacts/), which artifact folders do not have.
|
|
"""
|
|
root_dir = tmp_path / "mlruns"
|
|
root_dir.mkdir()
|
|
fs = FileStore(str(root_dir))
|
|
|
|
exp_id = fs.create_experiment("malicious_exp", artifact_location=str(root_dir))
|
|
run = fs.create_run(
|
|
experiment_id=exp_id, user_id="attacker", start_time=0, tags=[], run_name=""
|
|
)
|
|
run_id = run.info.run_id
|
|
|
|
assert Path(run.info.artifact_uri) == root_dir / run_id / "artifacts"
|
|
|
|
artifacts_dir = root_dir / run_id / "artifacts"
|
|
artifacts_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
target_dir = tmp_path / "sensitive_data"
|
|
target_dir.mkdir()
|
|
|
|
malicious_meta = {
|
|
"run_id": "artifacts",
|
|
"run_uuid": "artifacts",
|
|
"experiment_id": run_id,
|
|
"user_id": "attacker",
|
|
"status": 1,
|
|
"start_time": 0,
|
|
"end_time": None,
|
|
"lifecycle_stage": "active",
|
|
"artifact_uri": str(target_dir),
|
|
"tags": [],
|
|
}
|
|
write_yaml(str(artifacts_dir), "meta.yaml", malicious_meta)
|
|
|
|
# The fix should prevent the artifact folder from being treated as a run directory
|
|
with pytest.raises(MlflowException, match="Run 'artifacts' not found"):
|
|
fs.get_run("artifacts")
|