def test_creation_and_hydration(self):
(ri1, run_id, experiment_id, user_id, status, start_time, end_time,
lifecycle_stage, artifact_uri) = self._create()
self._check(ri1, run_id, experiment_id, user_id, status, start_time,
end_time, lifecycle_stage, artifact_uri)
as_dict = {
"run_uuid": run_id,
"run_id": run_id,
"experiment_id": experiment_id,
"user_id": user_id,
"status": status,
"start_time": start_time,
"end_time": end_time,
"lifecycle_stage": lifecycle_stage,
"artifact_uri": artifact_uri
}
self.assertEqual(dict(ri1), as_dict)
proto = ri1.to_proto()
ri2 = RunInfo.from_proto(proto)
self._check(ri2, run_id, experiment_id, user_id, status, start_time,
end_time, lifecycle_stage, artifact_uri)
ri3 = RunInfo.from_dictionary(as_dict)
self._check(ri3, run_id, experiment_id, user_id, status, start_time,
end_time, lifecycle_stage, artifact_uri)
# Test that we can add a field to RunInfo and still deserialize it from a dictionary
dict_copy_0 = as_dict.copy()
dict_copy_0["my_new_field"] = "new field value"
ri4 = RunInfo.from_dictionary(dict_copy_0)
self._check(ri4, run_id, experiment_id, user_id, status, start_time,
end_time, lifecycle_stage, artifact_uri)
def update_run_info(self, run_id, run_status, end_time):
""" Updates the metadata of the specified run. """
req_body = message_to_json(
UpdateRun(run_uuid=run_id,
run_id=run_id,
status=run_status,
end_time=end_time))
response_proto = self._call_endpoint(UpdateRun, req_body)
return RunInfo.from_proto(response_proto.run_info)
def _read_persisted_run_info_dict(run_info_dict):
dict_copy = run_info_dict.copy()
if 'lifecycle_stage' not in dict_copy:
dict_copy['lifecycle_stage'] = LifecycleStage.ACTIVE
# 'status' is stored as an integer enum in meta file, but RunInfo.status field is a string.
# converting to string before hydrating RunInfo.
# If 'status' value not recorded in files, mark it as 'RUNNING' (default)
dict_copy['status'] = RunStatus.to_string(
run_info_dict.get('status', RunStatus.RUNNING))
# 'experiment_id' was changed from int to string, so we must cast to string
# when reading legacy run_infos
if isinstance(dict_copy["experiment_id"], int):
dict_copy["experiment_id"] = str(dict_copy["experiment_id"])
return RunInfo.from_dictionary(dict_copy)
def _create():
run_id = str(uuid.uuid4())
experiment_id = str(random_int(10, 2000))
user_id = random_str(random_int(10, 25))
status = RunStatus.to_string(random.choice(RunStatus.all_status()))
start_time = random_int(1, 10)
end_time = start_time + random_int(1, 10)
lifecycle_stage = LifecycleStage.ACTIVE
artifact_uri = random_str(random_int(10, 40))
ri = RunInfo(run_uuid=run_id,
run_id=run_id,
experiment_id=experiment_id,
user_id=user_id,
status=status,
start_time=start_time,
end_time=end_time,
lifecycle_stage=lifecycle_stage,
artifact_uri=artifact_uri)
return (ri, run_id, experiment_id, user_id, status, start_time,
end_time, lifecycle_stage, artifact_uri)
def create_run(run_id="",
exp_id="",
uid="",
start=0,
end=0,
metrics=None,
params=None,
tags=None,
status=RunStatus.FINISHED,
a_uri=None):
return Run(
RunInfo(run_uuid=run_id,
run_id=run_id,
experiment_id=exp_id,
user_id=uid,
status=status,
start_time=start,
end_time=end,
lifecycle_stage=LifecycleStage.ACTIVE,
artifact_uri=a_uri),
RunData(metrics=metrics, params=params, tags=tags))
def to_mlflow_entity(self):
"""
Convert DB model to corresponding MLflow entity.
:return: :py:class:`mlflow.entities.Run`.
"""
run_info = RunInfo(run_uuid=self.run_uuid,
run_id=self.run_uuid,
experiment_id=str(self.experiment_id),
user_id=self.user_id,
status=self.status,
start_time=self.start_time,
end_time=self.end_time,
lifecycle_stage=self.lifecycle_stage,
artifact_uri=self.artifact_uri)
run_data = RunData(
metrics=[m.to_mlflow_entity() for m in self.latest_metrics],
params=[p.to_mlflow_entity() for p in self.params],
tags=[t.to_mlflow_entity() for t in self.tags])
return Run(run_info=run_info, run_data=run_data)
def test_string_repr(self):
run_info = RunInfo(run_uuid="hi",
run_id="hi",
experiment_id=0,
user_id="user-id",
status=RunStatus.FAILED,
start_time=0,
end_time=1,
lifecycle_stage=LifecycleStage.ACTIVE)
metrics = [
Metric(key="key-%s" % i, value=i, timestamp=0, step=i)
for i in range(3)
]
run_data = RunData(metrics=metrics, params=[], tags=[])
run1 = Run(run_info, run_data)
expected = (
"<Run: data=<RunData: metrics={'key-0': 0, 'key-1': 1, 'key-2': 2}, "
"params={}, tags={}>, info=<RunInfo: artifact_uri=None, end_time=1, "
"experiment_id=0, "
"lifecycle_stage='active', run_id='hi', run_uuid='hi', "
"start_time=0, status=4, user_id='user-id'>>")
assert str(run1) == expected
def create_run(self, experiment_id, user_id, start_time, tags):
"""
Creates a run with the specified attributes.
"""
experiment_id = FileStore.DEFAULT_EXPERIMENT_ID if experiment_id is None else experiment_id
experiment = self.get_experiment(experiment_id)
if experiment is None:
raise MlflowException(
"Could not create run under experiment with ID %s - no such experiment "
"exists." % experiment_id,
databricks_pb2.RESOURCE_DOES_NOT_EXIST)
if experiment.lifecycle_stage != LifecycleStage.ACTIVE:
raise MlflowException(
"Could not create run under non-active experiment with ID "
"%s." % experiment_id, databricks_pb2.INVALID_STATE)
run_uuid = uuid.uuid4().hex
artifact_uri = self._get_artifact_dir(experiment_id, run_uuid)
run_info = RunInfo(run_uuid=run_uuid,
run_id=run_uuid,
experiment_id=experiment_id,
artifact_uri=artifact_uri,
user_id=user_id,
status=RunStatus.to_string(RunStatus.RUNNING),
start_time=start_time,
end_time=None,
lifecycle_stage=LifecycleStage.ACTIVE)
# Persist run metadata and create directories for logging metrics, parameters, artifacts
run_dir = self._get_run_dir(run_info.experiment_id, run_info.run_id)
mkdir(run_dir)
run_info_dict = _make_persisted_run_info_dict(run_info)
write_yaml(run_dir, FileStore.META_DATA_FILE_NAME, run_info_dict)
mkdir(run_dir, FileStore.METRICS_FOLDER_NAME)
mkdir(run_dir, FileStore.PARAMS_FOLDER_NAME)
mkdir(run_dir, FileStore.ARTIFACTS_FOLDER_NAME)
for tag in tags:
self.set_tag(run_uuid, tag)
return self.get_run(run_id=run_uuid)
def test_searchable_attributes(self):
self.assertSequenceEqual(set(["status", "artifact_uri"]),
set(RunInfo.get_searchable_attributes()))