suggested changes + tests

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")
 

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)):

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)

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.

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")
 

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")
 

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

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)

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

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

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