def test_failed_http_request(self, request):
response = mock.MagicMock
response.status_code = 404
response.text = '{"error_code": "RESOURCE_DOES_NOT_EXIST", "message": "No experiment"}'
request.return_value = response
store = RestStore(lambda: MlflowHostCreds('https://hello'))
with pytest.raises(MlflowException) as cm:
store.list_experiments()
assert "RESOURCE_DOES_NOT_EXIST: No experiment" in str(cm.value)
def test_failed_http_request(self, request):
response = mock.MagicMock
response.status_code = 404
response.text = '{"error_code": "RESOURCE_DOES_NOT_EXIST", "message": "No experiment"}'
request.return_value = response
store = RestStore({'hostname': 'https://hello'})
with self.assertRaises(RestException) as cm:
store.list_experiments()
self.assertIn("RESOURCE_DOES_NOT_EXIST: No experiment",
str(cm.exception))
def _get_databricks_rest_store(store_uri):
parsed_uri = urllib.parse.urlparse(store_uri)
profile = None
if parsed_uri.scheme == 'databricks':
profile = parsed_uri.hostname
http_request_kwargs = rest_utils.get_databricks_http_request_kwargs_or_fail(profile)
return RestStore(http_request_kwargs)
def test_successful_http_request(self, request):
def mock_request(**kwargs):
# Filter out None arguments
kwargs = dict((k, v) for k, v in six.iteritems(kwargs) if v is not None)
assert kwargs == {
'method': 'GET',
'url': 'https://hello/api/2.0/preview/mlflow/experiments/list',
}
response = mock.MagicMock
response.status_code = 200
response.text = '{"experiments": [{"name": "Exp!"}]}'
return response
request.side_effect = mock_request
store = RestStore({'hostname': 'https://hello'})
experiments = store.list_experiments()
assert experiments[0].name == "Exp!"
def test_response_with_unknown_fields(self, request):
experiment_json = {
"experiment_id": 1,
"name": "My experiment",
"artifact_location": "foo",
"OMG_WHAT_IS_THIS_FIELD": "Hooly cow",
}
response = mock.MagicMock
response.status_code = 200
experiments = {"experiments": [experiment_json]}
response.text = json.dumps(experiments)
request.return_value = response
store = RestStore(lambda: MlflowHostCreds('https://hello'))
experiments = store.list_experiments()
assert len(experiments) == 1
assert experiments[0].name == 'My experiment'
def _get_rest_store(store_uri):
def get_default_host_creds():
return rest_utils.MlflowHostCreds(
host=store_uri,
username=os.environ.get(_TRACKING_USERNAME_ENV_VAR),
password=os.environ.get(_TRACKING_PASSWORD_ENV_VAR),
token=os.environ.get(_TRACKING_TOKEN_ENV_VAR),
ignore_tls_verification=os.environ.get(_TRACKING_INSECURE_TLS_ENV_VAR) == 'true',
)
return RestStore(get_default_host_creds)
def test_successful_http_request(self, request):
def mock_request(**kwargs):
# Filter out None arguments
kwargs = dict((k, v) for k, v in six.iteritems(kwargs) if v is not None)
assert kwargs == {
'method': 'GET',
'params': {'view_type': 'ACTIVE_ONLY'},
'url': 'https://hello/api/2.0/mlflow/experiments/list',
'headers': _DEFAULT_HEADERS,
'verify': True,
}
response = mock.MagicMock
response.status_code = 200
response.text = '{"experiments": [{"name": "Exp!", "lifecycle_stage": "active"}]}'
return response
request.side_effect = mock_request
store = RestStore(lambda: MlflowHostCreds('https://hello'))
experiments = store.list_experiments()
assert experiments[0].name == "Exp!"
def test_dbfs_instantiation():
artifact_repository_registry = ArtifactRepositoryRegistry()
mock_dbfs_constructor = mock.Mock()
artifact_repository_registry.register("dbfs", mock_dbfs_constructor)
mock_get_host_creds = mock.Mock()
rest_store = RestStore(mock_get_host_creds)
mock_dbfs_repo = artifact_repository_registry.get_artifact_repository(
artifact_uri="dbfs://test-path", store=rest_store
)
assert mock_dbfs_repo == mock_dbfs_constructor.return_value
mock_dbfs_constructor.assert_called_once_with("dbfs://test-path", mock_get_host_creds)
from mlflow import log_metric, log_param, log_artifacts, get_artifact_uri, active_run,\
get_tracking_uri, log_artifact
if __name__ == "__main__":
print("Running {} with tracking URI {}".format(sys.argv[0],
get_tracking_uri()))
log_param("param1", 5)
log_metric("foo", 5)
log_metric("foo", 6)
log_metric("foo", 7)
log_metric("random_int", random.randint(0, 100))
run = active_run()
print("In run with UUID: %s" % run.info.run_uuid)
tracking_uri = get_tracking_uri()
if tracking_uri.startswith("http://"):
store = RestStore(get_tracking_uri())
metric_obj = store.get_metric(run.info.run_uuid, "foo")
metric_history = store.get_metric_history(run.info.run_uuid, "foo")
param_obj = store.get_param(run.info.run_uuid, "param1")
print("Got metric %s, %s" % (metric_obj.key, metric_obj.value))
print("Got param %s, %s" % (param_obj.key, param_obj.value))
print("Got metric history %s" % metric_history)
local_dir = tempfile.mkdtemp()
message = "test artifact written during run %s within artifact URI %s\n" \
% (active_run().info.run_uuid, get_artifact_uri())
try:
file_path = os.path.join(local_dir, "some_output_file.txt")
with open(file_path, "w") as handle:
handle.write(message)
log_artifacts(local_dir, "some_subdir")
log_artifact(file_path, "another_dir")
def test_requestor(self, request):
response = mock.MagicMock
response.status_code = 200
response.text = '{}'
request.return_value = response
creds = MlflowHostCreds('https://hello')
store = RestStore(lambda: creds)
user_name = "mock user"
source_name = "rest test"
source_name_patch = mock.patch(
"mlflow.tracking.context.default_context._get_source_name",
return_value=source_name)
source_type_patch = mock.patch(
"mlflow.tracking.context.default_context._get_source_type",
return_value=SourceType.LOCAL)
with mock.patch('mlflow.store.rest_store.http_request') as mock_http, \
mock.patch('mlflow.tracking.utils._get_store', return_value=store), \
mock.patch('mlflow.tracking.context.default_context._get_user',
return_value=user_name), \
mock.patch('time.time', return_value=13579), \
source_name_patch, source_type_patch:
with mlflow.start_run(experiment_id="43"):
cr_body = message_to_json(
CreateRun(experiment_id="43",
user_id=user_name,
start_time=13579000,
tags=[
ProtoRunTag(key='mlflow.source.name',
value=source_name),
ProtoRunTag(key='mlflow.source.type',
value='LOCAL'),
ProtoRunTag(key='mlflow.user',
value=user_name)
]))
expected_kwargs = self._args(creds, "runs/create", "POST",
cr_body)
assert mock_http.call_count == 1
actual_kwargs = mock_http.call_args[1]
# Test the passed tag values separately from the rest of the request
# Tag order is inconsistent on Python 2 and 3, but the order does not matter
expected_tags = expected_kwargs['json'].pop('tags')
actual_tags = actual_kwargs['json'].pop('tags')
assert (sorted(expected_tags,
key=lambda t: t['key']) == sorted(
actual_tags, key=lambda t: t['key']))
assert expected_kwargs == actual_kwargs
with mock.patch('mlflow.store.rest_store.http_request') as mock_http:
store.log_param("some_uuid", Param("k1", "v1"))
body = message_to_json(
LogParam(run_uuid="some_uuid",
run_id="some_uuid",
key="k1",
value="v1"))
self._verify_requests(mock_http, creds, "runs/log-parameter",
"POST", body)
with mock.patch('mlflow.store.rest_store.http_request') as mock_http:
store.set_tag("some_uuid", RunTag("t1", "abcd" * 1000))
body = message_to_json(
SetTag(run_uuid="some_uuid",
run_id="some_uuid",
key="t1",
value="abcd" * 1000))
self._verify_requests(mock_http, creds, "runs/set-tag", "POST",
body)
with mock.patch('mlflow.store.rest_store.http_request') as mock_http:
store.delete_tag("some_uuid", "t1")
body = message_to_json(DeleteTag(run_id="some_uuid", key="t1"))
self._verify_requests(mock_http, creds, "runs/delete-tag", "POST",
body)
with mock.patch('mlflow.store.rest_store.http_request') as mock_http:
store.log_metric("u2", Metric("m1", 0.87, 12345, 3))
body = message_to_json(
LogMetric(run_uuid="u2",
run_id="u2",
key="m1",
value=0.87,
timestamp=12345,
step=3))
self._verify_requests(mock_http, creds, "runs/log-metric", "POST",
body)
with mock.patch('mlflow.store.rest_store.http_request') as mock_http:
metrics = [
Metric("m1", 0.87, 12345, 0),
Metric("m2", 0.49, 12345, -1),
Metric("m3", 0.58, 12345, 2)
]
params = [Param("p1", "p1val"), Param("p2", "p2val")]
tags = [RunTag("t1", "t1val"), RunTag("t2", "t2val")]
store.log_batch(run_id="u2",
metrics=metrics,
params=params,
tags=tags)
metric_protos = [metric.to_proto() for metric in metrics]
param_protos = [param.to_proto() for param in params]
tag_protos = [tag.to_proto() for tag in tags]
body = message_to_json(
LogBatch(run_id="u2",
metrics=metric_protos,
params=param_protos,
tags=tag_protos))
self._verify_requests(mock_http, creds, "runs/log-batch", "POST",
body)
with mock.patch('mlflow.store.rest_store.http_request') as mock_http:
store.delete_run("u25")
self._verify_requests(mock_http, creds, "runs/delete", "POST",
message_to_json(DeleteRun(run_id="u25")))
with mock.patch('mlflow.store.rest_store.http_request') as mock_http:
store.restore_run("u76")
self._verify_requests(mock_http, creds, "runs/restore", "POST",
message_to_json(RestoreRun(run_id="u76")))
with mock.patch('mlflow.store.rest_store.http_request') as mock_http:
store.delete_experiment("0")
self._verify_requests(
mock_http, creds, "experiments/delete", "POST",
message_to_json(DeleteExperiment(experiment_id="0")))
with mock.patch('mlflow.store.rest_store.http_request') as mock_http:
store.restore_experiment("0")
self._verify_requests(
mock_http, creds, "experiments/restore", "POST",
message_to_json(RestoreExperiment(experiment_id="0")))
with mock.patch('mlflow.store.rest_store.http_request') as mock_http:
response = mock.MagicMock
response.text = '{"runs": ["1a", "2b", "3c"], "next_page_token": "67890fghij"}'
mock_http.return_value = response
result = store.search_runs(["0", "1"],
"params.p1 = 'a'",
ViewType.ACTIVE_ONLY,
max_results=10,
order_by=["a"],
page_token="12345abcde")
expected_message = SearchRuns(experiment_ids=["0", "1"],
filter="params.p1 = 'a'",
run_view_type=ViewType.to_proto(
ViewType.ACTIVE_ONLY),
max_results=10,
order_by=["a"],
page_token="12345abcde")
self._verify_requests(mock_http, creds, "runs/search", "POST",
message_to_json(expected_message))
assert result.token == "67890fghij"
def _get_rest_store(store_uri):
return RestStore({'hostname': store_uri})
def _get_rest_store(store_uri):
return RestStore(store_uri)
def _get_databricks_rest_store(store_uri):
profile = get_db_profile_from_uri(store_uri)
return RestStore(lambda: get_databricks_host_creds(profile))
def test_requestor(self, request):
response = mock.MagicMock
response.status_code = 200
response.text = '{}'
request.return_value = response
creds = MlflowHostCreds('https://hello')
store = RestStore(lambda: creds)
user_name = "mock user"
run_name = "rest run"
source_name = "rest test"
with mock.patch('mlflow.store.rest_store.http_request_safe') as mock_http, \
mock.patch('mlflow.tracking.utils._get_store', return_value=store), \
mock.patch('mlflow.tracking.client._get_user_id', return_value=user_name), \
mock.patch('time.time', return_value=13579):
with mlflow.start_run(experiment_id="43",
run_name=run_name,
source_name=source_name):
cr_body = message_to_json(
CreateRun(experiment_id="43",
run_name='',
user_id=user_name,
source_type=SourceType.LOCAL,
source_name=source_name,
start_time=13579000,
tags=[
ProtoRunTag(key='mlflow.source.name',
value=source_name),
ProtoRunTag(key='mlflow.source.type',
value='LOCAL')
]))
st_body = message_to_json(
SetTag(run_uuid='', key='mlflow.runName', value=run_name))
assert mock_http.call_count == 2
exp_calls = [("runs/create", "POST", cr_body),
("runs/set-tag", "POST", st_body)]
self._verify_request_has_calls(mock_http, creds, exp_calls)
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
store.log_param("some_uuid", Param("k1", "v1"))
body = message_to_json(
LogParam(run_uuid="some_uuid", key="k1", value="v1"))
self._verify_requests(mock_http, creds, "runs/log-parameter",
"POST", body)
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
store.set_tag("some_uuid", RunTag("t1", "abcd" * 1000))
body = message_to_json(
SetTag(run_uuid="some_uuid", key="t1", value="abcd" * 1000))
self._verify_requests(mock_http, creds, "runs/set-tag", "POST",
body)
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
store.log_metric("u2", Metric("m1", 0.87, 12345, 3))
body = message_to_json(
LogMetric(run_uuid="u2",
key="m1",
value=0.87,
timestamp=12345,
step=3))
self._verify_requests(mock_http, creds, "runs/log-metric", "POST",
body)
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
metrics = [
Metric("m1", 0.87, 12345, 0),
Metric("m2", 0.49, 12345, -1),
Metric("m3", 0.58, 12345, 2)
]
params = [Param("p1", "p1val"), Param("p2", "p2val")]
tags = [RunTag("t1", "t1val"), RunTag("t2", "t2val")]
store.log_batch(run_id="u2",
metrics=metrics,
params=params,
tags=tags)
metric_protos = [metric.to_proto() for metric in metrics]
param_protos = [param.to_proto() for param in params]
tag_protos = [tag.to_proto() for tag in tags]
body = message_to_json(
LogBatch(run_id="u2",
metrics=metric_protos,
params=param_protos,
tags=tag_protos))
self._verify_requests(mock_http, creds, "runs/log-batch", "POST",
body)
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
store.delete_run("u25")
self._verify_requests(mock_http, creds, "runs/delete", "POST",
message_to_json(DeleteRun(run_id="u25")))
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
store.restore_run("u76")
self._verify_requests(mock_http, creds, "runs/restore", "POST",
message_to_json(RestoreRun(run_id="u76")))
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
store.delete_experiment("0")
self._verify_requests(
mock_http, creds, "experiments/delete", "POST",
message_to_json(DeleteExperiment(experiment_id="0")))
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
store.restore_experiment("0")
self._verify_requests(
mock_http, creds, "experiments/restore", "POST",
message_to_json(RestoreExperiment(experiment_id="0")))
from mlflow import log_metric, log_param, log_artifacts, get_artifact_uri, active_run,\
get_tracking_uri,set_tracking_uri, log_artifact
from mlflow.sklearn import log_model
if __name__ == "__main__":
print("Running {} with tracking URI {}".format(sys.argv[0], sys.argv[1]))
set_tracking_uri(sys.argv[1])
log_param("my_param", 5)
log_metric("my_metric_name", 5)
log_metric("my_metric_name", 1)
log_metric("my_metric_name", 22)
run = active_run()
print("In run with UUID: %s" % run.info.run_uuid)
tracking_uri = get_tracking_uri()
if tracking_uri.startswith("http://"):
store = RestStore({'hostname': tracking_uri})
metric_obj = store.get_metric(run.info.run_uuid, "my_metric_name")
metric_history = store.get_metric_history(run.info.run_uuid,
"my_metric_name")
param_obj = store.get_param(run.info.run_uuid, "my_param")
print("Got metric %s, %s" % (metric_obj.key, metric_obj.value))
print("Got param %s, %s" % (param_obj.key, param_obj.value))
print("Got metric history %s" % metric_history)
local_dir = tempfile.mkdtemp()
message = "test artifact written during run %s within artifact URI %s\n" \
% (active_run().info.run_uuid, get_artifact_uri())
try:
file_path = os.path.join(local_dir, "some_output_file.txt")
with open(file_path, "w") as handle:
handle.write(message)
log_artifacts(local_dir, "some_subdir")
def test_get_host_creds_from_default_store_rest_store():
with mock.patch('mlflow.tracking.utils._get_store') as get_store_mock:
get_store_mock.return_value = RestStore(
lambda: MlflowHostCreds('http://host'))
assert isinstance(_get_host_creds_from_default_store()(),
MlflowHostCreds)
def test_requestor(self, request):
response = mock.MagicMock
response.status_code = 200
response.text = '{}'
request.return_value = response
creds = MlflowHostCreds('https://hello')
store = RestStore(lambda: creds)
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
store.log_param("some_uuid", Param("k1", "v1"))
body = message_to_json(
LogParam(run_uuid="some_uuid", key="k1", value="v1"))
self._verify_requests(mock_http, creds, "runs/log-parameter",
"POST", body)
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
store.set_tag("some_uuid", RunTag("t1", "abcd" * 1000))
body = message_to_json(
SetTag(run_uuid="some_uuid", key="t1", value="abcd" * 1000))
self._verify_requests(mock_http, creds, "runs/set-tag", "POST",
body)
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
store.log_metric("u2", Metric("m1", 0.87, 12345))
body = message_to_json(
LogMetric(run_uuid="u2", key="m1", value=0.87,
timestamp=12345))
self._verify_requests(mock_http, creds, "runs/log-metric", "POST",
body)
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
store.delete_run("u25")
self._verify_requests(mock_http, creds, "runs/delete", "POST",
message_to_json(DeleteRun(run_id="u25")))
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
store.restore_run("u76")
self._verify_requests(mock_http, creds, "runs/restore", "POST",
message_to_json(RestoreRun(run_id="u76")))
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
store.delete_experiment(0)
self._verify_requests(
mock_http, creds, "experiments/delete", "POST",
message_to_json(DeleteExperiment(experiment_id=0)))
with mock.patch(
'mlflow.store.rest_store.http_request_safe') as mock_http:
store.restore_experiment(0)
self._verify_requests(
mock_http, creds, "experiments/restore", "POST",
message_to_json(RestoreExperiment(experiment_id=0)))
