[WIP] Use StorageManager in execute_DAG

docs/guide/storage.rst

@@ -1,5 +1,106 @@
 The GUFE Storage System
 =======================
 
+Storage lifetimes
+-----------------
 
-Storage docs.
+The storage system in GUFE is heavily tied to the GUFE protocol system. The
+key concept here is that the different levels of the GUFE protocol system;
+campaign, DAG, and unit; inherently imply different lifetimes for the data
+that is created. Those different lifetimes define the stages of the GUFE
+storage system.
+
+In an abstract sense, as used by protocol developers, these three levels
+correspond to three lifetimes of data:
+
+* ``scratch``: This is temporary data that is only needed for the lifetime
+  of a :class:`.ProtocolUnit`. This data is not guaranteed to be available
+  beyond the single :class:`.ProtocolUnit` where it is created, but may be
+  reused within that :class:`.ProtocolUnit`.
+* ``shared``: This is data that is shared between different units in a
+  :class:`.ProtocolDAG`. For example, a single equilibration stage might be
+  shared between multiple production runs. The output snapshot of the
+  equilibration would be suitable for as something to put in ``shared``
+  data. This data is guaranteed to be present from when it is created until
+  the end of the :class:`.ProtocolDAG`, but is not guaranteed to exist after
+  the :class:`.ProtocolDAG` terminates.
+* ``permanent``: This is the data that will be needed beyond the scope of a
+  single rough estimate of the calculation. This could include anything that
+  an extension of the simulation would require, or things that require
+  network-scale analysis. Anything stored here will be usable after the
+  calculation has completed.
+
+The ``scratch`` area is always a local directory. However, ``shared`` and
+``permanent`` can be external (remote) resources, using the
+:class:`.ExternalResource` API.
+
+As a practical matter, the GUFE storage system can be handled with a
+:class:`.StorageManager`. This automates some aspects of the transfer
+between stages of the GUFE storage system, and simplifies the API for
+protocol authors.  In detail, this provides protocol authors with
+``PathLike`` objects for ``scratch``, ``shared``, and ``permanent``. All
+three of these objects actually point to special subdirectories of the
+local scratch space for a specific unit, but are managed by context
+managers at the executor level, which handle the process of moving objects
+from local staging directories to the actual ``shared`` and ``permanent``
+locations, which can be external resources.
+
+
+External resource utilities
+---------------------------
+
+For flexible data storage, GUFE defines the :class:`.ExternalResource` API,
+which allows data be stored/loaded in a way that is agnostic to the
+underlying data store, as long as the store can be represented as a
+key-value store. Withing GUFE, we provide several examples, including
+:class:`.FileStorage` and :class:`.MemoryStorage` (primarily useful for
+testing.) The specific ``shared`` and ``permanent`` resources, as provided
+to the executor, can be instances of an :class:`.ExternalResource`.
+
+.. note::
+
+   The ``shared`` space must be a resource where an uploaded object is
+   instantaneously available, otherwise later protocol units may fail if the
+   shared result is unavailable. This means that files or approaches based
+   on ``scp`` or ``sftp`` are fine, but things like cloud storage, where the
+   existence of a new document may take time to propagate through the
+   network, are not recommended for ``shared``.
+
+
+Details: Manangement of the Storage Lifetime
+--------------------------------------------
+
+The concepts of the storage lifetimes are important for protocol authors,
+but details of implementation are left to the specific executor. In order to
+facilitate correct treatment of the storage lifecycle, GUFE provides a few
+helpers. The information in this section is mostly of interest to authors of
+executors. The helpers are:
+
+* :class:`.StorageManager`: This is the overall façade interface for
+  interacting with the rest of the storage lifecycle tools. It provides two
+  methods to generate context managers; one for the :class:`.ProtocolDAG`
+  level of the lifecycle, and one for the :class:`.ProtocoUnit` level of the
+  lifecycle. This class is designed for the use case that the entire DAG is
+  run in serial within a single process. Subclasses of this can be created
+  for other execution architectures, where the main logic changes would be
+  in the methods that return those context managers.
+* :class:`.StagingRegistry`: This handles the logic around staging paths
+  within a :class:`.ProtocolUnit`. Think of this as an abstract
+  representation of a local directory. Paths within it register with it, and
+  it handles deletion of the temporary local files when not needed, as well
+  as the download of remote files when necessary for reading. There are two
+  important subclasses of this: :class:`.SharedStaging` for a ``shared``
+  resource, and :class:`.PermanentStaging` for a ``permanent`` resource.
+* :class:`.StagingPath`: This represents a file within the
+  :class:`.StagingRegistry`. It contains both the key (label) used in the
+  key-value store, as well as the actual local path to the file. When its
+  ``__fspath__`` method is called, it registers itself with its
+  :class:`.StagingRegistry`, which handles managing it over its lifecycle.
+
+In practice, the executor uses the :class:`.StorageManager` to create a
+:class:`.DAGContextManager` at the level of a DAG, and then uses the
+:class:`.DAGContextManager` to create a context to run a unit. That context
+creates a :class:`.SharedStaging` and a :class:`.PermanentStaging`
+associated with the specific unit. Those staging directories, with the
+scratch directory, are provided to the :class:`.ProtocolUnit`, so that
+these are the only objects protocol authors need to interact with.

gufe/protocols/protocoldag.py

@@ -17,6 +17,13 @@
     ProtocolUnit, ProtocolUnitResult, ProtocolUnitFailure, Context
 )
 
+from ..storage.storagemanager import StorageManager
+from ..storage.externalresource.filestorage import FileStorage
+from ..storage.externalresource.base import ExternalStorage
+
+import logging
+_logger = logging.getLogger(__name__)
+
 
 class DAGMixin:
     _protocol_units: list[ProtocolUnit]
@@ -345,9 +352,67 @@ def _from_dict(cls, dct: dict):
         return cls(**dct)
 
 
+class ReproduceOldBehaviorStorageManager(StorageManager):
+    # Default behavior has scratch at {dag_label}/scratch/{unit_label} and
+    # shared at {dag_label}/{unit_label}. This little class makes changes
+    # that get us back to the original behavior of this class: scratch at
+    # {dag_label}/scratch_{unit_label} and shared at
+    # {dag_label}/shared_{unit_label}.
+    def _scratch_loc(self, dag_label, unit_label, attempt):
+        return (
+            self.scratch_root
+            / f"{dag_label}/scratch_{unit_label}_attempt_{attempt}"
+        )
+
+    def make_label(self, dag_label, unit_label, attempt):
+        return f"{dag_label}/shared_{unit_label}_attempt_{attempt}"
+
+    @classmethod
+    def from_old_args(
+        cls,
+        shared_basedir: PathLike,
+        scratch_basedir: PathLike, *,
+        keep_shared: bool = False,
+        keep_scratch: bool = False,
+    ):
+        """
+        Create an new storage manager based on the old execute_DAG args.
+
+        Parameters
+        ----------
+        shared_basedir : Path
+            Filesystem path to use for shared space that persists across whole DAG
+            execution. Used by a `ProtocolUnit` to pass file contents to dependent
+            class:``ProtocolUnit`` instances.
+        scratch_basedir : Path
+            Filesystem path to use for `ProtocolUnit` `scratch` space.
+        keep_shared : bool
+            If True, don't remove shared directories for `ProtocolUnit`s after
+            the `ProtocolDAG` is executed.
+        keep_scratch : bool
+            If True, don't remove scratch directories for a `ProtocolUnit` after
+            it is executed.
+        """
+        # doing this here makes it easier to test than putting in
+        # execute_DAG
+        shared_basedir = Path(shared_basedir)
+        shared = FileStorage(shared_basedir.parent)
+        storage_manager = cls(
+            scratch_root=scratch_basedir,
+            shared_root=shared,
+            permanent_root=shared,
+            keep_scratch=keep_scratch,
+            keep_shared=keep_shared,
+            keep_staging=True,
+            keep_empty_dirs=True,
+            staging=Path(""),  # use the actual directories as the staging
+        )
+        return storage_manager
+
+
 def execute_DAG(protocoldag: ProtocolDAG, *,
-                shared_basedir: Path,
-                scratch_basedir: Path,
+                shared_basedir: PathLike,
+                scratch_basedir: PathLike,
                 keep_shared: bool = False,
                 keep_scratch: bool = False,
                 raise_error: bool = True,
@@ -385,52 +450,92 @@ def execute_DAG(protocoldag: ProtocolDAG, *,
         The result of executing the `ProtocolDAG`.
 
     """
+    # the directory given as shared_root is actually the directory for this
+    # DAG; the "shared_root" for the storage manager is the parent. We'll
+    # force permanent to be the same.
+    storage_manager = ReproduceOldBehaviorStorageManager.from_old_args(
+        shared_basedir=shared_basedir,
+        scratch_basedir=scratch_basedir,
+        keep_shared=keep_shared,
+        keep_scratch=keep_scratch
+    )
+    dag_label = shared_basedir.name
+    return new_execute_DAG(protocoldag, dag_label, storage_manager,
+                           raise_error, n_retries)
+
+
+def new_execute_DAG(  # TODO: this is a terrible name
+    protocoldag: ProtocolDAG,
+    dag_label: str,
+    storage_manager: StorageManager,
+    raise_error: bool = False,
+    n_retries: int = 0
+) -> ProtocolDAGResult:
+    """
+    Locally execute a full :class:`ProtocolDAG` in serial and in-process.
+
+    Alternate input signature to generalize execute_DAG
+
+    Parameters
+    ----------
+    protocoldag : ProtocolDAG
+        The :class:``ProtocolDAG`` to execute.
+    dag_label : str
+        Label to use for the DAG
+    storage_manager : StorageManager
+        The :class:`.StorageManager` to handle storing files.
+    raise_error : bool
+        If True, raise an exception if a ProtocolUnit fails, default True
+        if False, any exceptions will be stored as `ProtocolUnitFailure`
+        objects inside the returned `ProtocolDAGResult`
+    n_retries : int
+        the number of times to attempt, default 0, i.e. try once and only once
+
+    Returns
+    -------
+    ProtocolDAGResult
+        The result of executing the `ProtocolDAG`.
+    """
+    # this simplifies setup of execute_DAG by allowing you to directly
+    # provide the storage_manager; the extra options in the old one just
+    # configure the storage_manager
     if n_retries < 0:
         raise ValueError("Must give positive number of retries")
 
     # iterate in DAG order
     results: dict[GufeKey, ProtocolUnitResult] = {}
     all_results = []  # successes AND failures
-    shared_paths = []
-    for unit in protocoldag.protocol_units:
-        # translate each `ProtocolUnit` in input into corresponding
-        # `ProtocolUnitResult`
-        inputs = _pu_to_pur(unit.inputs, results)
-
-        attempt = 0
-        while attempt <= n_retries:
-            shared = shared_basedir / f'shared_{str(unit.key)}_attempt_{attempt}'
-            shared_paths.append(shared)
-            shared.mkdir()
-
-            scratch = scratch_basedir / f'scratch_{str(unit.key)}_attempt_{attempt}'
-            scratch.mkdir()
-
-            context = Context(shared=shared,
-                              scratch=scratch)
-
-            # execute
-            result = unit.execute(
-                    context=context,
-                    raise_error=raise_error,
-                    **inputs)
-            all_results.append(result)
-
-            if not keep_scratch:
-                shutil.rmtree(scratch)
-
-            if result.ok():
-                # attach result to this `ProtocolUnit`
-                results[unit.key] = result
-                break
-            attempt += 1
 
-        if not result.ok():
-            break
-
-    if not keep_shared:
-        for shared_path in shared_paths:
-            shutil.rmtree(shared_path)
+    with storage_manager.running_dag(dag_label) as dag_ctx:
+        for unit in protocoldag.protocol_units:
+            # import pdb; pdb.set_trace()
+            attempt = 0
+            while attempt <= n_retries:
+                # translate each `ProtocolUnit` in input into corresponding
+                # `ProtocolUnitResult`
+                inputs = _pu_to_pur(unit.inputs, results)
+
+                label = storage_manager.make_label(dag_label, unit.key,
+                                                   attempt=attempt)
+                with dag_ctx.running_unit(
+                    dag_label, unit.key, attempt=attempt
+                ) as context:
+                    _logger.info(f"Starting unit {label}")
+                    _logger.info(context)
+                    result = unit.execute(
+                            context=context,
+                            raise_error=raise_error,
+                            **inputs)
+                    all_results.append(result)
+
+                if result.ok():
+                    # attach result to this `ProtocolUnit`
+                    results[unit.key] = result
+                    break
+                attempt += 1
+
+            if not result.ok():
+                break
 
     return ProtocolDAGResult(
             name=protocoldag.name, 

gufe/protocols/protocolunit.py

@@ -23,6 +23,8 @@
     GufeTokenizable, GufeKey, TOKENIZABLE_REGISTRY
 )
 
+from ..storage.stagingregistry import StagingPath
+
 
 @dataclass
 class Context:
@@ -31,7 +33,8 @@ class Context:
 
     """
     scratch: PathLike
-    shared: PathLike
+    shared: StagingPath
+    permanent: StagingPath
 
 
 def _list_dependencies(inputs, cls):