Skip to content

Commit

Permalink
remove step feature for local workspaces
Browse files Browse the repository at this point in the history
  • Loading branch information
Pranjali Basmatkar authored and Pranjali Basmatkar committed Jul 12, 2023
1 parent 01077eb commit ab72783
Showing 1 changed file with 25 additions and 0 deletions.
25 changes: 25 additions & 0 deletions tango/workspace.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
from urllib.parse import ParseResult, urlparse

import pytz
from sqlitedict import SqliteDict
import shutil

from .common import Registrable
from .common.from_params import FromParams
Expand Down Expand Up @@ -419,6 +421,29 @@ def step_result(self, step_name: str) -> Any:
return self.step_cache[run.steps[step_name]]
raise KeyError(f"No step named '{step_name}' found in previous runs")


def remove_step(self, step_name: str) -> Any:
"""
Get Step name (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
"""
# get path to dir dynamically
sqlite_path = self.dir / "stepinfo.sqlite"
with SqliteDict(sqlite_path) as d:
try:
step_location = d[step_name].result_location
# remove step info from the sqlite dict
del d[step_name]
d.commit()
# remove cache directory
try:
shutil.rmtree(step_location)
except OSError:
raise OSError('Step Cache folder not found')
return('Step deleted')
except KeyError:
raise KeyError(f"No step named '{step_name}' found")

def capture_logs_for_run(self, name: str) -> ContextManager[None]:
"""
Should return a context manager that can be used to capture the logs for a run.
Expand Down

0 comments on commit ab72783

Please sign in to comment.