Skip to content

Commit

Permalink
change to unit taking in the label
Browse files Browse the repository at this point in the history
dwhswenson committed May 31, 2023
1 parent b805aac commit ed5e83c
Showing 2 changed files with 5 additions and 6 deletions.
9 changes: 4 additions & 5 deletions gufe/storage/storagemanager.py
Original file line number Diff line number Diff line change
@@ -27,14 +27,14 @@ def _storage_path_conflict(external, path):
class _AbstractDAGContextManager:
@classmethod
@contextmanager
def running_dag(cls, storage_manager, dag_label):
def running_dag(cls, storage_manager, dag_label: str):
raise NotImplementedError()

@contextmanager
def running_unit(cls, unit):
def running_unit(cls, unit_label: str):
raise NotImplementedError()

DAGContextManager = Type[_DAGStorageManager]
DAGContextManager = Type[_AbstractDAGContextManager]


class _DAGStorageManager(_AbstractDAGContextManager):
@@ -74,15 +74,14 @@ def running_dag(cls, storage_manager, dag_label):
d.cleanup()

@contextmanager
def running_unit(self, unit):
def running_unit(self, unit_label: str):
"""Unit level of the storage lifecycle.
This provides the holding directories used for scratch, shared, and
permanent. At the end of the unit, it transfers anything from shared
to the real shared external storage, cleans up the scratch
directory and the shared holding directory.
"""
unit_label = unit.key
scratch = self.manager.get_scratch(self.dag_label, unit_label)
shared = self.manager.get_shared(self.dag_label, unit_label)
permanent = self.manager.get_permanent(self.dag_label, unit_label)
2 changes: 1 addition & 1 deletion gufe/tests/storage/test_storagemanager.py
Original file line number Diff line number Diff line change
@@ -74,7 +74,7 @@ def test_lifecycle(request, manager, dag_units):
all_files = {barfile, bazfile, foofile, foo2file}
with storage_manager.running_dag("dag_label") as dag_ctx:
for unit in dag_units:
with dag_ctx.running_unit(unit) as (scratch, shared, permanent):
with dag_ctx.running_unit(unit.key) as (scratch, shared, permanent):
results.append(unit.run(scratch, shared, permanent))

# check that the expected files are found in staging

0 comments on commit ed5e83c

Please sign in to comment.