From 15656581f5acd61f8e72976f9640fef1c0b471ca Mon Sep 17 00:00:00 2001 From: Pranjali Basmatkar Date: Fri, 14 Jul 2023 11:17:52 -0700 Subject: [PATCH] remove cached step for memory workspaces --- tango/step_caches/memory_step_cache.py | 6 ++++++ tango/workspaces/memory_workspace.py | 11 +++++++++++ 2 files changed, 17 insertions(+) diff --git a/tango/step_caches/memory_step_cache.py b/tango/step_caches/memory_step_cache.py index e57751c70..7844ea13e 100644 --- a/tango/step_caches/memory_step_cache.py +++ b/tango/step_caches/memory_step_cache.py @@ -35,6 +35,12 @@ 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] + else: + raise KeyError(f"{step_unique_id} not present in the memory cache. Can't be deleted") + def __contains__(self, step: object) -> bool: if isinstance(step, (Step, StepInfo)): return step.unique_id in self.cache diff --git a/tango/workspaces/memory_workspace.py b/tango/workspaces/memory_workspace.py index bcbb89498..7c039773a 100644 --- a/tango/workspaces/memory_workspace.py +++ b/tango/workspaces/memory_workspace.py @@ -98,6 +98,17 @@ 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: + """ + 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: + del self.unique_id_to_info[step_unique_id] + self.step_cache.__delitem__(step_unique_id) + except KeyError: + raise KeyError(f"{step_unique_id} step info not found, step cache cannot be deleted") + def register_run(self, targets: Iterable[Step], name: Optional[str] = None) -> Run: if name is None: name = petname.generate()