def test_create_run_backwards_compatability(
mocker,
run_name_tag,
expected_run_name,
parent_run_id_tag,
expected_parent_run_id,
):
mocker.patch("mlflow_faculty.tracking.mlflow_timestamp_to_datetime")
mocker.patch("mlflow_faculty.tracking.mlflow_tag_to_faculty_tag")
mock_client = mocker.Mock()
mocker.patch("faculty.client", return_value=mock_client)
mocker.patch("mlflow_faculty.tracking.faculty_run_to_mlflow_run")
tags = []
if run_name_tag is not None:
tags.append(RunTag(key=MLFLOW_RUN_NAME, value=run_name_tag))
if parent_run_id_tag is not None:
tags.append(RunTag(key=MLFLOW_PARENT_RUN_ID, value=parent_run_id_tag))
store = FacultyRestStore(STORE_URI)
store.create_run(
EXPERIMENT_ID,
"unused-mlflow-user-id",
RUN_STARTED_AT_MILLISECONDS,
tags,
)
args, _ = mock_client.create_run.call_args
assert args[2] == expected_run_name
assert args[4] == expected_parent_run_id
def test_set_experiment_tag(mocker):
mocker.patch("faculty.client")
store = FacultyRestStore(STORE_URI)
with pytest.raises(NotImplementedError, match="not supported"):
store.set_experiment_tag(EXPERIMENT_ID, mocker.Mock())
def test_restore_experiment_invalid_experiment_id(mocker):
mocker.patch("faculty.client")
store = FacultyRestStore(STORE_URI)
with pytest.raises(ValueError):
store.restore_experiment("invalid-experiment-id")
def test_get_metric_history(mocker):
metric_key = "metric_key"
first_faculty_metric = mocker.Mock()
second_faculty_metric = mocker.Mock()
mock_client = mocker.Mock(get_metric_history=mocker.Mock(
return_value=[first_faculty_metric, second_faculty_metric]))
mocker.patch("faculty.client", return_value=mock_client)
first_mlflow_metric = mocker.Mock()
second_mlflow_metric = mocker.Mock()
metric_converter = mocker.patch(
"mlflow_faculty.tracking.faculty_metric_to_mlflow_metric",
side_effect=[first_mlflow_metric, second_mlflow_metric],
)
store = FacultyRestStore(STORE_URI)
returned_metric_history = store.get_metric_history(RUN_UUID_HEX_STR,
metric_key)
mock_client.get_metric_history.assert_called_once_with(
PROJECT_ID, RUN_UUID, metric_key)
metric_converter.assert_has_calls([
mocker.call(first_faculty_metric),
mocker.call(second_faculty_metric)
])
assert returned_metric_history == [
first_mlflow_metric,
second_mlflow_metric,
]
def test_create_run_experiment_deleted(mocker):
mocker.patch(
"mlflow_faculty.tracking.mlflow_timestamp_to_datetime",
return_value=mocker.Mock(),
)
mocker.patch(
"mlflow_faculty.tracking.mlflow_tag_to_faculty_tag",
return_value=mocker.Mock(),
)
mock_client = mocker.Mock()
exception = faculty.clients.experiment.ExperimentDeleted(
message="message", experiment_id="test-id")
mock_client.create_run.side_effect = exception
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(MlflowException, match="experiment"):
store.create_run(
EXPERIMENT_ID,
"unused-mlflow-user-id",
RUN_STARTED_AT_MILLISECONDS,
tags=[],
)
def test_update_run_info(mocker, run_status, faculty_run_status):
faculty_run = mocker.Mock()
mock_client = mocker.Mock()
mock_client.update_run_info.return_value = faculty_run
mocker.patch("faculty.client", return_value=mock_client)
mlflow_run_info = mocker.Mock()
mlflow_run = mocker.Mock()
mlflow_run.info = mlflow_run_info
run_converter_mock = mocker.patch(
"mlflow_faculty.tracking.faculty_run_to_mlflow_run",
return_value=mlflow_run,
)
store = FacultyRestStore(STORE_URI)
returned_run_info = store.update_run_info(
RUN_UUID_HEX_STR, run_status, RUN_ENDED_AT_MILLISECONDS
)
mock_client.update_run_info.assert_called_once_with(
PROJECT_ID, RUN_UUID, faculty_run_status, RUN_ENDED_AT
)
run_converter_mock.assert_called_once_with(faculty_run)
assert returned_run_info == mlflow_run_info
def test_restore_experiment(mocker):
mock_client = mocker.Mock()
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
store.restore_experiment(EXPERIMENT_ID)
mock_client.restore.assert_called_once_with(PROJECT_ID, EXPERIMENT_ID)
def test_restore_run_failures(mocker, response, message):
mock_client = mocker.Mock()
mock_client.restore_runs.return_value = response
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(MlflowException, match=message):
store.restore_run(RUN_UUID_HEX_STR)
def test_get_metric_history_invalid_run_id(mocker):
mock_client = mocker.Mock()
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(ValueError):
store.get_metric_history("invalid-run-id", "metric-key")
def test_rename_experiment_invalid_experiment_id(mocker):
mock_client = mocker.Mock()
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(ValueError):
store.rename_experiment("invalid-experiment-id", "new name")
def test_log_batch_invalid_run_id(mocker):
mock_client = mocker.Mock()
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(ValueError):
store.log_batch("invalid-run-id")
def test_restore_run_client_error(mocker):
mock_client = mocker.Mock()
mock_client.restore_runs.side_effect = HttpError(mocker.Mock(), "An error")
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(MlflowException, match="An error"):
store.restore_run(RUN_UUID_HEX_STR)
def test_rename_experiment_client_error(mocker):
mock_client = mocker.Mock()
mock_client.update.side_effect = HttpError(mocker.Mock(), "Error")
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(MlflowException, match="Error"):
store.rename_experiment(EXPERIMENT_ID, "new name")
def test_update_run_info_invalid_run_id(mocker):
mock_client = mocker.Mock()
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(ValueError):
store.update_run_info("invalid-run-id", RunStatus.RUNNING,
RUN_ENDED_AT_MILLISECONDS)
def test_create_experiment_client_error(mocker):
mock_client = mocker.Mock()
mock_client.create.side_effect = HttpError(mocker.Mock(), "Error")
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(MlflowException, match="Error"):
store.create_experiment(NAME, ARTIFACT_LOCATION)
def test_list_experiments_client_error(mocker):
mock_client = mocker.Mock()
mock_client.list.side_effect = HttpError(mocker.Mock(), "Error")
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(MlflowException, match="Error"):
store.list_experiments()
def test_log_batch_empty(mocker):
mock_client = mocker.Mock()
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
store.log_batch(RUN_UUID_HEX_STR)
mock_client.log_run_data.assert_called_once_with(
PROJECT_ID, RUN_UUID, metrics=[], params=[], tags=[]
)
def test_restore_run(mocker):
mock_client = mocker.Mock()
mock_client.restore_runs.return_value = RestoreExperimentRunsResponse(
restored_run_ids=[RUN_UUID], conflicted_run_ids=[])
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
store.restore_run(RUN_UUID_HEX_STR)
mock_client.restore_runs.assert_called_once_with(PROJECT_ID, [RUN_UUID])
def test_get_metric_history_client_error(mocker):
mock_client = mocker.Mock()
mock_client.get_metric_history.side_effect = HttpError(
mocker.Mock(), "Dummy client error.")
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(MlflowException, match="Dummy client error."):
store.get_metric_history(RUN_UUID_HEX_STR, "metric-key")
def test_create_experiment_name_conflict(mocker):
exception = ExperimentNameConflict("bad name")
mock_client = mocker.Mock()
mock_client.create.side_effect = exception
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(MlflowException, match=str(exception)):
store.create_experiment(NAME, ARTIFACT_LOCATION)
def test_rename_experiment_name_conflict(mocker):
exception = ExperimentNameConflict("bad name")
mock_client = mocker.Mock()
mock_client.update.side_effect = exception
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(MlflowException, match=str(exception)):
store.rename_experiment(EXPERIMENT_ID, "bad name")
def test_search_runs(mocker, max_results):
mock_faculty_runs = [mocker.Mock() for _ in range(4)]
run_page_1 = ListExperimentRunsResponse(
runs=mock_faculty_runs[:2],
pagination=mocker.Mock(next=Page(start=2, limit=2)),
)
run_page_2 = ListExperimentRunsResponse(
runs=mock_faculty_runs[2:], pagination=mocker.Mock(next=None)
)
mock_client = mocker.Mock()
mock_client.query_runs.side_effect = [run_page_1, run_page_2]
mocker.patch("faculty.client", return_value=mock_client)
experiment_ids = mocker.Mock()
filter_string = mocker.Mock()
run_view_type = mocker.Mock()
mock_filter = mocker.Mock()
build_filter_mock = mocker.patch(
"mlflow_faculty.tracking.build_search_runs_filter",
return_value=mock_filter,
)
mock_mlflow_runs = [mocker.Mock() for _ in range(4)]
run_converter_mock = mocker.patch(
"mlflow_faculty.tracking.faculty_run_to_mlflow_run",
side_effect=mock_mlflow_runs,
)
store = FacultyRestStore(STORE_URI)
runs, page_token = store._search_runs(
experiment_ids,
filter_string,
run_view_type,
max_results,
order_by=None,
page_token=None,
)
assert runs == mock_mlflow_runs[:max_results]
assert page_token is None
build_filter_mock.assert_called_once_with(
experiment_ids, filter_string, run_view_type
)
mock_client.query_runs.assert_has_calls(
[
mocker.call(PROJECT_ID, mock_filter),
mocker.call(PROJECT_ID, mock_filter, start=2, limit=2),
]
)
run_converter_mock.assert_has_calls(
[mocker.call(run) for run in mock_faculty_runs[:max_results]]
)
def test_get_experiment_client_error(mocker):
mock_client = mocker.Mock()
mock_client.get.side_effect = HttpError(
mocker.Mock(), "Experiment with ID _ not found in project _")
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(MlflowException,
match="Experiment with ID _ not found in project _"):
store.get_experiment(EXPERIMENT_ID)
def test_create_experiment(mocker):
mock_client = mocker.Mock()
mock_client.create.return_value = FACULTY_EXPERIMENT
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
experiment_id = store.create_experiment(NAME, ARTIFACT_LOCATION)
assert experiment_id == str(FACULTY_EXPERIMENT.id)
mock_client.create.assert_called_once_with(
PROJECT_ID, NAME, artifact_location=ARTIFACT_LOCATION)
def test_create_run_invalid_experiment_id(mocker):
mock_client = mocker.Mock()
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(ValueError):
store.create_run(
"invalid-experiment-id",
"unused-mlflow-user-id",
RUN_STARTED_AT_MILLISECONDS,
tags=[],
)
def test_create_run_invalid_parent_run_id(mocker):
mock_client = mocker.Mock()
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(ValueError):
store.create_run(
EXPERIMENT_ID,
"unused-mlflow-user-id",
RUN_STARTED_AT_MILLISECONDS,
[RunTag(key=MLFLOW_PARENT_RUN_ID, value="invalid-uuid")],
)
def test_get_run_client_error(mocker):
mock_client = mocker.Mock()
mock_client.get_run.side_effect = HttpError(
mocker.Mock(), "Experiment run with ID _ not found in project _")
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(
MlflowException,
match="Experiment run with ID _ not found in project _",
):
store.get_run(RUN_UUID_HEX_STR)
def test_update_run_info_client_error(mocker):
mock_client = mocker.Mock()
mock_client.update_run_info.side_effect = HttpError(
mocker.Mock(), "Experiment run with ID _ not found in project _")
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(
MlflowException,
match="Experiment run with ID _ not found in project _",
):
store.update_run_info(RUN_UUID_HEX_STR, "RUNNING",
RUN_ENDED_AT_MILLISECONDS)
def test_init_invalid_uri_in_netloc():
store_uri = "faculty://{}".format(PROJECT_ID)
expected_error_message = ("Invalid URI {}. Netloc is reserved. "
"Did you mean 'faculty:{}".format(
store_uri, PROJECT_ID))
with pytest.raises(ValueError, match=expected_error_message):
FacultyRestStore(store_uri)
def test_init(mocker, store_uri):
mocker.patch("faculty.client")
store = FacultyRestStore(store_uri)
assert store._project_id == PROJECT_ID
faculty.client.assert_called_once_with("experiment")
