def test_logging_failures_are_handled_as_expected():
experiment_name = "test_run_creation_termination"
MlflowClient().create_experiment(experiment_name)
experiment_id = MlflowClient().get_experiment_by_name(
experiment_name).experiment_id
with mock.patch(
"mlflow.utils.autologging_utils.client.MlflowClient.log_batch"
) as log_batch_mock:
log_batch_mock.side_effect = Exception("Batch logging failed!")
client = MlflowAutologgingQueueingClient()
pending_run_id = client.create_run(experiment_id=experiment_id)
client.log_metrics(run_id=pending_run_id, metrics={"a": 1})
client.set_terminated(run_id=pending_run_id, status="KILLED")
with pytest.raises(MlflowException) as exc:
client.flush()
runs = mlflow.search_runs(experiment_ids=[experiment_id],
output_format="list")
assert len(runs) == 1
run = runs[0]
# Verify that metrics are absent due to the failure of batch logging
assert not run.data.metrics
# Verify that the run termination operation was still performed successfully
assert run.info.status == "KILLED"
assert "Failed to perform one or more operations on the run with ID {run_id}".format(
run_id=run.info.run_id) in str(exc.value)
assert "Batch logging failed!" in str(exc.value)
def test_client_run_creation_and_termination_are_successful():
experiment_name = "test_run_creation_termination"
MlflowClient().create_experiment(experiment_name)
experiment_id = MlflowClient().get_experiment_by_name(experiment_name).experiment_id
client = MlflowAutologgingQueueingClient()
pending_run_id = client.create_run(experiment_id=experiment_id, start_time=5, tags={"a": "b"})
client.set_terminated(run_id=pending_run_id, status="FINISHED", end_time=6)
client.flush()
runs = mlflow.search_runs(experiment_ids=[experiment_id], output_format="list")
assert len(runs) == 1
run = runs[0]
assert run.info.start_time == 5
assert run.info.end_time == 6
assert run.info.status == "FINISHED"
assert {"a": "b"}.items() <= run.data.tags.items()
def test_flush_clears_pending_operations():
with mock.patch("mlflow.utils.autologging_utils.client.MlflowClient",
autospec=True) as mlflow_client_mock:
client = MlflowAutologgingQueueingClient()
pending_run_id = client.create_run(experiment_id=5)
client.log_params(run_id=pending_run_id, params={"a": "b"})
client.log_metrics(run_id=pending_run_id, metrics={"c": 1})
client.set_terminated(run_id=pending_run_id, status="FINISHED")
client.flush()
logging_call_count_1 = len(mlflow_client_mock.method_calls)
# Verify that at least 3 calls have been made to MLflow logging APIs as a result
# of the flush (i.e. log_batch, create_run, and set_terminated)
assert logging_call_count_1 >= 3
client.flush()
logging_call_count_2 = len(mlflow_client_mock.method_calls)
# Verify that performing a second flush did not result in any additional logging API calls,
# since no new run content was added prior to the flush
assert logging_call_count_2 == logging_call_count_1