def test_databricks_rest_store_get_experiment_by_name(self):
creds = MlflowHostCreds("https://hello")
store = DatabricksRestStore(lambda: creds)
with mock.patch("mlflow.utils.rest_utils.http_request") as mock_http:
# Verify that Databricks REST client won't fall back to ListExperiments for 500-level
# errors that are not ENDPOINT_NOT_FOUND
def rate_limit_response_fn(*args, **kwargs):
# pylint: disable=unused-argument
raise MlflowException("Some internal error!", INTERNAL_ERROR)
mock_http.side_effect = rate_limit_response_fn
with pytest.raises(MlflowException) as exc_info:
store.get_experiment_by_name("abc")
assert exc_info.value.error_code == ErrorCode.Name(INTERNAL_ERROR)
assert exc_info.value.message == "Some internal error!"
expected_message0 = GetExperimentByName(experiment_name="abc")
self._verify_requests(
mock_http,
creds,
"experiments/get-by-name",
"GET",
message_to_json(expected_message0),
)
assert mock_http.call_count == 1
def _get_experiment_by_name():
request_message = _get_request_message(GetExperimentByName())
response_message = GetExperimentByName.Response()
store_exp = _get_tracking_store().get_experiment_by_name(
request_message.experiment_name)
if store_exp is None:
raise MlflowException(
"Could not find experiment with name '%s'" %
request_message.experiment_name,
error_code=RESOURCE_DOES_NOT_EXIST,
)
experiment = store_exp.to_proto()
response_message.experiment.MergeFrom(experiment)
response = Response(mimetype="application/json")
response.set_data(message_to_json(response_message))
return response
def get_experiment_by_name(self, experiment_name):
try:
req_body = message_to_json(GetExperimentByName(experiment_name=experiment_name))
response_proto = self._call_endpoint(GetExperimentByName, req_body)
return Experiment.from_proto(response_proto.experiment)
except MlflowException as e:
if e.error_code == databricks_pb2.ErrorCode.Name(
databricks_pb2.RESOURCE_DOES_NOT_EXIST
):
return None
elif e.error_code == databricks_pb2.ErrorCode.Name(databricks_pb2.ENDPOINT_NOT_FOUND):
# Fall back to using ListExperiments-based implementation.
for experiment in self.list_experiments(ViewType.ALL):
if experiment.name == experiment_name:
return experiment
return None
raise e
def test_get_experiment_by_name(self, store_class):
creds = MlflowHostCreds('https://hello')
store = store_class(lambda: creds)
with mock.patch('mlflow.store.rest_store.http_request') as mock_http:
response = mock.MagicMock
response.status_code = 200
experiment = Experiment(
experiment_id="123", name="abc", artifact_location="/abc",
lifecycle_stage=LifecycleStage.ACTIVE)
response.text = json.dumps({
"experiment": json.loads(message_to_json(experiment.to_proto()))})
mock_http.return_value = response
result = store.get_experiment_by_name("abc")
expected_message0 = GetExperimentByName(experiment_name="abc")
self._verify_requests(mock_http, creds,
"experiments/get-by-name", "GET",
message_to_json(expected_message0))
assert result.experiment_id == experiment.experiment_id
assert result.name == experiment.name
assert result.artifact_location == experiment.artifact_location
assert result.lifecycle_stage == experiment.lifecycle_stage
# Test GetExperimentByName against nonexistent experiment
mock_http.reset_mock()
nonexistent_exp_response = mock.MagicMock
nonexistent_exp_response.status_code = 404
nonexistent_exp_response.text =\
MlflowException("Exp doesn't exist!", RESOURCE_DOES_NOT_EXIST).serialize_as_json()
mock_http.return_value = nonexistent_exp_response
assert store.get_experiment_by_name("nonexistent-experiment") is None
expected_message1 = GetExperimentByName(experiment_name="nonexistent-experiment")
self._verify_requests(mock_http, creds,
"experiments/get-by-name", "GET",
message_to_json(expected_message1))
assert mock_http.call_count == 1
# Test REST client behavior against a mocked old server, which has handler for
# ListExperiments but not GetExperimentByName
mock_http.reset_mock()
list_exp_response = mock.MagicMock
list_exp_response.text = json.dumps({
"experiments": [json.loads(message_to_json(experiment.to_proto()))]})
list_exp_response.status_code = 200
def response_fn(*args, **kwargs):
# pylint: disable=unused-argument
if kwargs.get('endpoint') == "/api/2.0/mlflow/experiments/get-by-name":
raise MlflowException("GetExperimentByName is not implemented",
ENDPOINT_NOT_FOUND)
else:
return list_exp_response
mock_http.side_effect = response_fn
result = store.get_experiment_by_name("abc")
expected_message2 = ListExperiments(view_type=ViewType.ALL)
self._verify_requests(mock_http, creds,
"experiments/get-by-name", "GET",
message_to_json(expected_message0))
self._verify_requests(mock_http, creds,
"experiments/list", "GET",
message_to_json(expected_message2))
assert result.experiment_id == experiment.experiment_id
assert result.name == experiment.name
assert result.artifact_location == experiment.artifact_location
assert result.lifecycle_stage == experiment.lifecycle_stage
# Verify that REST client won't fall back to ListExperiments for 429 errors (hitting
# rate limits)
mock_http.reset_mock()
def rate_limit_response_fn(*args, **kwargs):
# pylint: disable=unused-argument
raise MlflowException("Hit rate limit on GetExperimentByName",
REQUEST_LIMIT_EXCEEDED)
mock_http.side_effect = rate_limit_response_fn
with pytest.raises(MlflowException) as exc_info:
store.get_experiment_by_name("imspamming")
assert exc_info.value.error_code == ErrorCode.Name(REQUEST_LIMIT_EXCEEDED)
assert mock_http.call_count == 1
