Skip to content

Commit

Permalink
Fix delete_experiment and restore_experiment in FileStore impleme…
Browse files Browse the repository at this point in the history 
…ntation (#8178)

* Fix delete_experiment and restore_experiment in file_store
implementation (#8177)

Signed-off-by: Marius Schlegel <marius.schlegel@tu-ilmenau.de>

* Sort imports with isort

Signed-off-by: Marius Schlegel <marius.schlegel@tu-ilmenau.de>

* Replace exception with logger warning

Signed-off-by: Marius Schlegel <marius.schlegel@tu-ilmenau.de>

* Fix implementation for passing tests

Signed-off-by: Marius Schlegel <marius.schlegel@tu-ilmenau.de>

* Sort imports with isort

Signed-off-by: Marius Schlegel <marius.schlegel@tu-ilmenau.de>

* Fixes and cleanup

Signed-off-by: Marius Schlegel <marius.schlegel@tu-ilmenau.de>

* Add new tests and fix existing tests

Signed-off-by: Marius Schlegel <marius.schlegel@tu-ilmenau.de>

* Revert "Sort imports with isort"

This reverts commit 537a6d2.

Signed-off-by: Marius Schlegel <marius.schlegel@tu-ilmenau.de>

* Revert "Sort imports with isort"

This reverts commit fc161a0.

Signed-off-by: Marius Schlegel <marius.schlegel@tu-ilmenau.de>

---------

Signed-off-by: Marius Schlegel <marius.schlegel@tu-ilmenau.de>
  • Loading branch information
@mariusschlegel
mariusschlegel committed Apr 13, 2023
1 parent 3762014 commit 24e7c07
Show file tree
Hide file tree
Showing 2 changed files with 69 additions and 35 deletions.
23 changes: 18 additions & 5 deletions mlflow/store/tracking/file_store.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -381,10 +381,6 @@ def _get_experiment(self, experiment_id, view_type=ViewType.ALL):
databricks_pb2.RESOURCE_DOES_NOT_EXIST,
)
meta = FileStore._read_yaml(experiment_dir, FileStore.META_DATA_FILE_NAME)
if experiment_dir.startswith(self.trash_folder):
meta["lifecycle_stage"] = LifecycleStage.DELETED
else:
meta["lifecycle_stage"] = LifecycleStage.ACTIVE
meta["tags"] = self.get_all_experiment_tags(experiment_id)
experiment = _read_persisted_experiment_dict(meta)
if experiment_id != experiment.experiment_id:
Expand Down Expand Up @@ -429,7 +425,16 @@ def delete_experiment(self, experiment_id):
databricks_pb2.RESOURCE_DOES_NOT_EXIST,
)
experiment = self._get_experiment(experiment_id)
experiment._set_last_update_time(get_current_time_millis())
experiment._lifecycle_stage = LifecycleStage.DELETED
deletion_time = get_current_time_millis()
experiment._set_last_update_time(deletion_time)
runs = self._list_run_infos(experiment_id, view_type=ViewType.ACTIVE_ONLY)
for run_info in runs:
if run_info is not None:
new_info = run_info._copy_with_overrides(lifecycle_stage=LifecycleStage.DELETED)
self._overwrite_run_info(new_info, deleted_time=deletion_time)
else:
logging.warning("Run metadata is in invalid state.")
meta_dir = os.path.join(self.root_directory, experiment_id)
overwrite_yaml(
root=meta_dir,
Expand Down Expand Up @@ -463,7 +468,15 @@ def restore_experiment(self, experiment_id):
mv(experiment_dir, self.root_directory)
experiment = self._get_experiment(experiment_id)
meta_dir = os.path.join(self.root_directory, experiment_id)
experiment._lifecycle_stage = LifecycleStage.ACTIVE
experiment._set_last_update_time(get_current_time_millis())
runs = self._list_run_infos(experiment_id, view_type=ViewType.DELETED_ONLY)
for run_info in runs:
if run_info is not None:
new_info = run_info._copy_with_overrides(lifecycle_stage=LifecycleStage.ACTIVE)
self._overwrite_run_info(new_info, deleted_time=None)
else:
logging.warning("Run metadata is in invalid state.")
overwrite_yaml(
root=meta_dir,
file_name=FileStore.META_DATA_FILE_NAME,
Expand Down
81 changes: 51 additions & 30 deletions tests/store/tracking/test_file_store.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -354,6 +354,7 @@ def _create_root(store):
"experiment_id": exp,
"name": random_str(),
"artifact_location": exp_folder,
"lifecycle_stage": LifecycleStage.ACTIVE,
"creation_time": current_time,
"last_update_time": current_time,
}
Expand All @@ -378,6 +379,7 @@ def _create_root(store):
"deleted_time": random_int(20, 30),
"tags": [],
"artifact_uri": os.path.join(run_folder, FileStore.ARTIFACTS_FOLDER_NAME),
"lifecycle_stage": LifecycleStage.ACTIVE,
}
write_yaml(run_folder, FileStore.META_DATA_FILE_NAME, run_info)
run_data[run_id] = run_info
Expand Down Expand Up @@ -590,39 +592,58 @@ def _extract_ids(experiments):


def test_delete_restore_experiment(store):
exp_id = store.create_experiment("test_delete")
exp_name = store.get_experiment(exp_id).name

exp1 = store.get_experiment(exp_id)
time.sleep(0.001)

# delete it
store.delete_experiment(exp_id)
assert exp_id not in _extract_ids(store.search_experiments(view_type=ViewType.ACTIVE_ONLY))
assert exp_id in _extract_ids(store.search_experiments(view_type=ViewType.DELETED_ONLY))
assert exp_id in _extract_ids(store.search_experiments(view_type=ViewType.ALL))
assert store.get_experiment(exp_id).lifecycle_stage == LifecycleStage.DELETED

deleted_exp1 = store.get_experiment(exp_id)
experiments, _, _ = _create_root(store)
exp1_id = experiments[random_int(0, len(experiments) - 2)] # never select default experiment
exp1 = store.get_experiment(exp1_id)

# test deleting experiment
store.delete_experiment(exp1_id)
assert exp1_id not in _extract_ids(store.search_experiments(view_type=ViewType.ACTIVE_ONLY))
assert exp1_id in _extract_ids(store.search_experiments(view_type=ViewType.DELETED_ONLY))
assert exp1_id in _extract_ids(store.search_experiments(view_type=ViewType.ALL))
deleted_exp1 = store.get_experiment(exp1_id)
assert deleted_exp1.last_update_time > exp1.last_update_time
assert deleted_exp1.lifecycle_stage == LifecycleStage.DELETED

# restore it
exp1 = store.get_experiment(exp_id)
time.sleep(0.01)
store.restore_experiment(exp_id)
restored_1 = store.get_experiment(exp_id)
assert restored_1.experiment_id == exp_id
assert restored_1.name == exp_name
assert restored_1.last_update_time > exp1.last_update_time

restored_2 = store.get_experiment_by_name(exp_name)
assert restored_2.experiment_id == exp_id
assert restored_2.name == exp_name
assert exp_id in _extract_ids(store.search_experiments(view_type=ViewType.ACTIVE_ONLY))
assert exp_id not in _extract_ids(store.search_experiments(view_type=ViewType.DELETED_ONLY))
assert exp_id in _extract_ids(store.search_experiments(view_type=ViewType.ALL))
assert store.get_experiment(exp_id).lifecycle_stage == LifecycleStage.ACTIVE
# test if setting lifecycle_stage is persisted correctly
deleted_exp1_dir = store._get_experiment_path(
experiment_id=exp1_id, view_type=ViewType.DELETED_ONLY
)
deleted_exp1_meta = FileStore._read_yaml(
root=deleted_exp1_dir, file_name=FileStore.META_DATA_FILE_NAME
)
assert deleted_exp1_meta["lifecycle_stage"] == LifecycleStage.DELETED
for run in store.search_runs(
experiment_ids=[exp1_id], filter_string="", run_view_type=ViewType.ALL
):
assert run.info.lifecycle_stage == LifecycleStage.DELETED

# test restoring experiment
store.restore_experiment(exp1_id)
assert exp1_id in _extract_ids(store.search_experiments(view_type=ViewType.ACTIVE_ONLY))
assert exp1_id not in _extract_ids(store.search_experiments(view_type=ViewType.DELETED_ONLY))
assert exp1_id in _extract_ids(store.search_experiments(view_type=ViewType.ALL))
restored1_exp1 = store.get_experiment(exp1_id)
assert restored1_exp1.experiment_id == exp1_id
assert restored1_exp1.name == exp1.name
assert restored1_exp1.last_update_time > exp1.last_update_time
assert restored1_exp1.lifecycle_stage == LifecycleStage.ACTIVE
restored2_exp1 = store.get_experiment_by_name(exp1.name)
assert restored2_exp1.experiment_id == exp1_id
assert restored2_exp1.name == exp1.name

# test if setting lifecycle_stage is persisted correctly
restored_exp1_dir = store._get_experiment_path(
experiment_id=exp1_id, view_type=ViewType.ACTIVE_ONLY
)
restored_exp1_meta = FileStore._read_yaml(
root=restored_exp1_dir, file_name=FileStore.META_DATA_FILE_NAME
)
assert restored_exp1_meta["lifecycle_stage"] == LifecycleStage.ACTIVE
for run in store.search_runs(
experiment_ids=[exp1_id], filter_string="", run_view_type=ViewType.ALL
):
assert run.info.lifecycle_stage == LifecycleStage.ACTIVE


def test_rename_experiment(store):
Expand Down

0 comments on commit 24e7c07

Please sign in to comment.