diff --git a/tango/step_caches/local_step_cache.py b/tango/step_caches/local_step_cache.py index aaece269b..c65e9dcb6 100644 --- a/tango/step_caches/local_step_cache.py +++ b/tango/step_caches/local_step_cache.py @@ -160,11 +160,11 @@ def __setitem__(self, step: Step, value: Any) -> None: pass raise - def __delitem__(self, step_unique_id) -> None: - location = str(self.dir) + "/" + str(step_unique_id) + def __delitem__(self, step: Union[Step, StepInfo]) -> None: + location = str(self.dir) + "/" + str(step.unique_id) try: shutil.rmtree(location) - assert not os.path.exists(location) + self._remove_from_cache(step.unique_id) except OSError: raise OSError("Step Cache folder not found") diff --git a/tango/step_caches/memory_step_cache.py b/tango/step_caches/memory_step_cache.py index 75e39b2f0..184b7a4ce 100644 --- a/tango/step_caches/memory_step_cache.py +++ b/tango/step_caches/memory_step_cache.py @@ -35,12 +35,11 @@ def __setitem__(self, step: Step, value: Any) -> None: UserWarning, ) - def __delitem__(self, step_unique_id) -> None: - if step_unique_id in self.cache: - del self.cache[step_unique_id] - assert step_unique_id not in self.cache + def __delitem__(self, step: Union[Step, StepInfo]) -> None: + if step.unique_id in self.cache: + del self.cache[step.unique_id] else: - raise KeyError(f"{step_unique_id} not present in the memory cache. Can't be deleted") + raise KeyError(f"{step.unique_id} not present in the memory cache. Cannot be deleted.") def __contains__(self, step: object) -> bool: if isinstance(step, (Step, StepInfo)): diff --git a/tango/step_caches/remote_step_cache.py b/tango/step_caches/remote_step_cache.py index 2f6fb5e86..74950fcc0 100644 --- a/tango/step_caches/remote_step_cache.py +++ b/tango/step_caches/remote_step_cache.py @@ -156,10 +156,10 @@ def __setitem__(self, step: Step, value: Any) -> None: # Finally, add to in-memory caches. self._add_to_cache(step.unique_id, value) - def __delitem__(self, step_unique_id) -> None: + def __delitem__(self, step: Union[Step, StepInfo]) -> None: # check and delete local cache dir - if self.step_dir(step_unique_id).is_dir(): - shutil.rmtree(self.step_dir(step_unique_id)) + if self.step_dir(step.unique_id).is_dir(): + shutil.rmtree(self.step_dir(step.unique_id)) # remove from memory cache - self._remove_from_cache(key=step_unique_id) + self._remove_from_cache(key=step.unique_id) diff --git a/tango/workspace.py b/tango/workspace.py index 050c7e587..1261b267b 100644 --- a/tango/workspace.py +++ b/tango/workspace.py @@ -420,7 +420,7 @@ def step_result(self, step_name: str) -> Any: raise KeyError(f"No step named '{step_name}' found in previous runs") @abstractmethod - def step_cache_remove(self, step_unique_id: str): + def remove_step(self, step_unique_id: str): """ Removes cached step using the given unique step id :raises KeyError: If there is no step with the given name. diff --git a/tango/workspaces/local_workspace.py b/tango/workspaces/local_workspace.py index d6469f0ee..61b4113e4 100644 --- a/tango/workspaces/local_workspace.py +++ b/tango/workspaces/local_workspace.py @@ -322,18 +322,17 @@ def step_failed(self, step: Step, e: BaseException) -> None: lock.release() del self.locks[step] - def step_cache_remove(self, step_unique_id: str) -> None: + def remove_step(self, step_unique_id: str) -> None: """ Get Step unique id from the user and remove the step information from cache :raises KeyError: If no step with the unique name found in the cache dir """ with SqliteDict(self.step_info_file) as d: try: - assert step_unique_id in d + step_info = self.step_info(step_unique_id) del d[step_unique_id] d.commit() - assert step_unique_id not in d - self.cache.__delitem__(step_unique_id) + del self.cache[step_info] except KeyError: raise KeyError(f"No step named '{step_unique_id}' found") diff --git a/tango/workspaces/memory_workspace.py b/tango/workspaces/memory_workspace.py index 687b8f8bd..47b29c077 100644 --- a/tango/workspaces/memory_workspace.py +++ b/tango/workspaces/memory_workspace.py @@ -98,15 +98,15 @@ def step_failed(self, step: Step, e: BaseException) -> None: existing_step_info.end_time = utc_now_datetime() existing_step_info.error = exception_to_string(e) - def step_cache_remove(self, step_unique_id: str) -> None: + def remove_step(self, step_unique_id: str) -> None: """ Get Step unique id from the user and remove the step information from memory cache :raises KeyError: If no step with the unique name found in the cache dir """ try: + step_info = self.step_info(step_unique_id) del self.unique_id_to_info[step_unique_id] - assert step_unique_id not in self.unique_id_to_info - del self.step_cache[step_unique_id] + del self.step_cache[step_info] except KeyError: raise KeyError(f"{step_unique_id} step info not found, step cache cannot be deleted") diff --git a/tango/workspaces/remote_workspace.py b/tango/workspaces/remote_workspace.py index 75e1bda69..49d351a9f 100644 --- a/tango/workspaces/remote_workspace.py +++ b/tango/workspaces/remote_workspace.py @@ -174,7 +174,7 @@ def step_failed(self, step: Step, e: BaseException) -> None: finally: self.locks.pop(step).release() - def step_cache_remove(self, step_unique_id: str) -> None: + def remove_step(self, step_unique_id: str) -> None: """ Get Step unique id from the user and remove the step information from cache :raises KeyError: If no step with the unique name found in the cache dir @@ -185,7 +185,7 @@ def step_cache_remove(self, step_unique_id: str) -> None: self._remove_step_info(step_info) # remove cache info - del self.cache[step_unique_id] + del self.cache[step_info] except KeyError: raise KeyError(f"No step named '{step_unique_id}' found.") return None diff --git a/tests/integrations/beaker/workspace_test.py b/tests/integrations/beaker/workspace_test.py index 4c68bcef0..f6f0d71d2 100644 --- a/tests/integrations/beaker/workspace_test.py +++ b/tests/integrations/beaker/workspace_test.py @@ -4,6 +4,7 @@ from tango.workspace import Workspace + def test_from_url(beaker_workspace: str): workspace = Workspace.from_url(f"beaker://{beaker_workspace}") assert isinstance(workspace, BeakerWorkspace) diff --git a/tests/integrations/gs/workspace_test.py b/tests/integrations/gs/workspace_test.py index 28318db8e..290776a04 100644 --- a/tests/integrations/gs/workspace_test.py +++ b/tests/integrations/gs/workspace_test.py @@ -40,3 +40,31 @@ def test_direct_usage(self): workspace.step_finished(step, 1.0) assert workspace.step_info(step).state == StepState.COMPLETED assert workspace.step_result_for_run(run.name, "float") == 1.0 + + def test_remove_step(self): + workspace = GSWorkspace(GS_BUCKET_NAME) + step = FloatStep(step_name="float", result=1.0) + step_info = workspace.step_info(step) + + workspace.step_starting(step) + workspace.step_finished(step, 1.0) + bucket_artifact = workspace.Constants.step_artifact_name(step_info) + ds_entity = workspace._ds.get(key=workspace._ds.key("stepinfo", step_info.unique_id)) + cache = workspace.step_cache + + assert workspace.client.artifacts(prefix=bucket_artifact) is not None + assert ds_entity is not None + assert step in cache + + workspace.remove_step(step.unique_id) + cache = workspace.step_cache + + ds_entity = workspace._ds.get(key=workspace._ds.key("stepinfo", step_info.unique_id)) + + try: + workspace.client.artifacts(prefix=bucket_artifact) + except KeyError: + pass + #to assert that the artifact is no longer present in the bucket + assert ds_entity is None + assert step not in cache diff --git a/tests/workspaces/local_workspace_cache_remove_test.py b/tests/workspaces/local_workspace_cache_remove_test.py deleted file mode 100644 index 9d16b6886..000000000 --- a/tests/workspaces/local_workspace_cache_remove_test.py +++ /dev/null @@ -1,24 +0,0 @@ -import pytest - -from tango import Step -from tango.common.testing import TangoTestCase -from tango.step_info import StepState -from tango.workspaces import LocalWorkspace - - -class AdditionStep(Step): - def run(self, a: int, b: int) -> int: # type: ignore - return a + b - - -class Test_Cache_Remove_Workspace(TangoTestCase): - def test_step_cache_remove(self): - workspace = LocalWorkspace(self.TEST_DIR) - step = AdditionStep(a=1, b=2) - step_info = workspace.step_info(step) - assert step_info.state == StepState.INCOMPLETE - result = step.result(workspace) - step_info = workspace.step_info(step) - assert step_info.state == StepState.COMPLETED - step_unique_id = step.unique_id - workspace.step_cache_remove(step_unique_id) diff --git a/tests/workspaces/local_workspace_test.py b/tests/workspaces/local_workspace_test.py index b987a87c3..b48777649 100644 --- a/tests/workspaces/local_workspace_test.py +++ b/tests/workspaces/local_workspace_test.py @@ -1,4 +1,5 @@ from shutil import copytree +from sqlitedict import SqliteDict import pytest @@ -8,6 +9,7 @@ from tango.workspaces import LocalWorkspace + class AdditionStep(Step): def run(self, a: int, b: int) -> int: # type: ignore return a + b @@ -73,3 +75,23 @@ def test_local_workspace_upgrade_v1_to_v2(self): while len(dependencies) > 0: step_info = workspace.step_info(dependencies.pop()) dependencies.extend(step_info.dependencies) + + def test_remove_step(self): + workspace = LocalWorkspace(self.TEST_DIR) + step = AdditionStep(a=1, b=2) + workspace.step_starting(step) + workspace.step_finished(step, 1.0) + + with SqliteDict(workspace.step_info_file) as d: + assert step.unique_id in d + + cache = workspace.step_cache + assert step in cache + + workspace.remove_step(step.unique_id) + + with SqliteDict(workspace.step_info_file) as d: + assert step.unique_id not in d + + cache = workspace.step_cache + assert step not in cache \ No newline at end of file diff --git a/tests/workspaces/memory_workspace_cache_remove_test.py b/tests/workspaces/memory_workspace_cache_remove_test.py deleted file mode 100644 index 91dfe4b4e..000000000 --- a/tests/workspaces/memory_workspace_cache_remove_test.py +++ /dev/null @@ -1,22 +0,0 @@ -# import pytest - -from tango import Step -from tango.step_info import StepState -from tango.workspaces import MemoryWorkspace - - -class AdditionStep(Step): - def run(self, a: int, b: int) -> int: # type: ignore - return a + b - - -def test_step_cache_remove(): - workspace = MemoryWorkspace() - step1 = AdditionStep(a=1, b=2) - step_info = workspace.step_info(step1) - assert step_info.state == StepState.INCOMPLETE - result1 = step1.result(workspace) - step_info = workspace.step_info(step1) - assert step_info.state == StepState.COMPLETED - step1_unique_id = step1.unique_id - workspace.step_cache_remove(step1_unique_id) diff --git a/tests/workspaces/memory_workspace_test.py b/tests/workspaces/memory_workspace_test.py new file mode 100644 index 000000000..9a84d8af8 --- /dev/null +++ b/tests/workspaces/memory_workspace_test.py @@ -0,0 +1,20 @@ +from tango.common.testing.steps import FloatStep +from tango.workspaces import MemoryWorkspace + + +def test_remove_step(): + workspace = MemoryWorkspace() + step = FloatStep(step_name="float", result=1.0) + + workspace.step_starting(step) + workspace.step_finished(step, 1.0) + cache = workspace.step_cache + + assert step.unique_id in workspace.unique_id_to_info + assert step in cache + + workspace.remove_step(step.unique_id) + cache = workspace.step_cache + + assert step.unique_id not in workspace.unique_id_to_info + assert step not in cache \ No newline at end of file