From 0800b0382261cb493e4c1309f5aa5787467a45ec Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Wed, 21 Jan 2026 11:56:36 -0700 Subject: [PATCH 01/15] docs: add information on storage Signed-off-by: Ethan Holz --- docs/concepts/index.rst | 1 + docs/concepts/storage.rst | 204 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 205 insertions(+) create mode 100644 docs/concepts/storage.rst diff --git a/docs/concepts/index.rst b/docs/concepts/index.rst index 1b9b879c..6e8007e0 100644 --- a/docs/concepts/index.rst +++ b/docs/concepts/index.rst @@ -12,3 +12,4 @@ utilize **gufe** APIs. tokenizables included_models serialization + storage diff --git a/docs/concepts/storage.rst b/docs/concepts/storage.rst new file mode 100644 index 00000000..a9454eae --- /dev/null +++ b/docs/concepts/storage.rst @@ -0,0 +1,204 @@ +.. _concepts-storage: + +How storage is handled in **gufe** +================================== + +**gufe** abstracts storage into a reusable storage interface using the :class:`.ExternalStorage` abstract base class. +This abstraction enables the storage of any file or byte stream using various storage backends without changing application code. + +Overview +-------- + +The storage system is designed to handle (files or byte data) that need to be stored in some location. +Instead of embedding the data, objects can store a reference (a *location* string) to where the data is stored externally. This approach provides several benefits: + +* **Efficiency**: Large objects don't need to be serialized multiple times +* **Flexibility**: Different storage backends (local filesystem, cloud storage, in-memory) can be used interchangeably +* **Deduplication**: The same data can be referenced by multiple objects +* **Lazy Loading**: Data is only loaded when needed + +The Storage Architecture +------------------------- + +The storage system consists of several key components: + +``ExternalStorage`` Base Class +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The :class:`.ExternalStorage` abstract base class defines the interface that all storage implementations must provide. This class provides: + +* **Store operations**: ``store_bytes()`` and ``store_path()`` to store data +* **Load operations**: ``load_stream()`` to retrieve data as a stream +* **Management**: ``exists()``, ``delete()``, and ``iter_contents()`` for managing stored data + +All storage operations use a *location* string as an identifier for the stored data. + +Storage Implementations +----------------------- + +**gufe** provides several built-in implementations of :class:`.ExternalStorage`: + +``FileStorage`` +~~~~~~~~~~~~~~~ + +The :class:`.FileStorage` implementation stores data on the local filesystem. It requires a root directory path and organizes stored files using the location string as a relative path: + +.. code-block:: python + + from pathlib import Path + from gufe.storage.externalresource import FileStorage + + # Create a file storage backend + storage = FileStorage(root_dir=Path("/path/to/storage")) + + # Store some data + data = b"Hello, World!" + storage.store_bytes("datasets/sample1.txt", data) + + # Check if data exists + if storage.exists("datasets/sample1.txt"): + # Load the data + with storage.load_stream("datasets/sample1.txt") as stream: + loaded_data = stream.read() + assert loaded_data == data + + # Delete the data + storage.delete("datasets/sample1.txt") + +``FileStorage`` automatically creates any necessary parent directories when storing files. + +``MemoryStorage`` +~~~~~~~~~~~~~~~~~ + +The :class:`.MemoryStorage` implementation stores data in a Python dictionary. This is primarily useful for testing and prototyping: + +.. code-block:: python + + from gufe.storage.externalresource import MemoryStorage + + # Create an in-memory storage backend + storage = MemoryStorage() + + # Store some data + data = b"Hello, World!" + storage.store_bytes("datasets/sample1.txt", data) + + # Load the data back + with storage.load_stream("datasets/sample1.txt") as stream: + loaded_data = stream.read() + +.. warning:: + ``MemoryStorage`` is not intended for production use and all data is lost when the Python process exits. + +StorageManager +-------------- + +The :class:`.StorageManager` class provides a higher-level interface for managing storage operations within a computational workflow. +It handles the transfer of files between a scratch directory and external storage: + +.. code-block:: python + + from pathlib import Path + from gufe.storage import StorageManager + from gufe.storage.externalresource import FileStorage + + # Set up storage + storage = FileStorage(root_dir=Path("/path/to/storage")) + scratch_dir = Path("/path/to/scratch") + + # Create a storage manager for a specific DAG and unit + manager = StorageManager( + scratch_dir=scratch_dir, + storage=storage, + dag_label="my_experiment", + unit_label="transformation_1" + ) + + # Register files for later transfer + manager.register("trajectory.dcd") + manager.register("results.json") + + # Transfer all registered files to external storage + manager._transfer() + + # Load files from external storage + trajectory_data = manager.load("my_experiment/transformation_1/trajectory.dcd") + +The ``StorageManager`` uses a namespace combining the ``dag_label`` and ``unit_label`` to organize files in the external storage backend. + +Implementing Custom Storage Backends +------------------------------------- + +To create a custom storage backend, subclass :class:`.ExternalStorage` and implement all the abstract methods: + +.. code-block:: python + + from gufe.storage.externalresource.base import ExternalStorage + from typing import ContextManager + + class MyCustomStorage(ExternalStorage): + """A custom storage implementation.""" + + def _store_bytes(self, location: str, byte_data: bytes): + """Store bytes at the given location.""" + # Implement storage logic + pass + + def _store_path(self, location: str, path): + """Store a file at the given path.""" + # Implement storage logic + pass + + def _load_stream(self, location: str) -> ContextManager: + """Return a context manager that yields a bytes-like object.""" + # Implement loading logic + pass + + def _exists(self, location: str) -> bool: + """Check if data exists at the location.""" + # Implement existence check + pass + + def _delete(self, location: str): + """Delete data at the location.""" + # Implement deletion logic + pass + + def _get_filename(self, location: str) -> str: + """Return a filename for the location.""" + # Implement filename generation + pass + + def _iter_contents(self, prefix: str = ""): + """Iterate over stored locations matching the prefix.""" + # Implement iteration logic + pass + + def _get_hexdigest(self, location: str) -> str: + """Return MD5 hexdigest of the data (optional override).""" + # Can override for performance improvements + pass + +.. note:: + All storage methods should be blocking operations, even if the underlying storage backend supports asynchronous operations. + +Error Handling +-------------- + +The storage system defines several exceptions in :mod:`gufe.storage.errors`: + +* :class:`.ExternalResourceError`: Base class for storage-related errors +* :class:`.MissingExternalResourceError`: Raised when attempting to access non-existent data +* :class:`.ChangedExternalResourceError`: Raised when metadata verification fails + +These exceptions can be caught and handled appropriately in application code: + +.. code-block:: python + + from gufe.storage.errors import MissingExternalResourceError + + try: + with storage.load_stream("nonexistent_file.txt") as stream: + data = stream.read() + except MissingExternalResourceError: + print("File not found in storage") From 5c428ebead55c5bf8a725ca1878f3814e28f479a Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Wed, 21 Jan 2026 13:34:55 -0700 Subject: [PATCH 02/15] docs: initial draft on context --- docs/concepts/context.rst | 202 ++++++++++++++++++++++++++++++++++++++ docs/concepts/index.rst | 1 + 2 files changed, 203 insertions(+) create mode 100644 docs/concepts/context.rst diff --git a/docs/concepts/context.rst b/docs/concepts/context.rst new file mode 100644 index 00000000..34b41cfd --- /dev/null +++ b/docs/concepts/context.rst @@ -0,0 +1,202 @@ +Context +======= + +``Context`` instances carry the execution environment for individual +``ProtocolUnit`` executions. They are created by the execution engine just +before a unit's ``_execute`` method is called and then discarded once the unit +returns. The class lives in ``gufe.protocols.protocolunit`` and acts as a thin +wrapper around two :class:`~gufe.storage.storagemanager.StorageManager` +objects and a scratch directory. + + +Why Context exists +------------------ + +``ProtocolUnit`` code frequently needs a few shared facilities: + +``scratch`` + A temporary directory that the unit can freely write to while it runs. + Files written here are considered ephemeral; the engine may delete them as + soon as the unit finishes. + +``shared`` + A :class:`~gufe.storage.storagemanager.StorageManager` backed by + short–lived :class:`~gufe.storage.externalresource.ExternalStorage`. Use + this to hand large files to downstream units without serializing the + payloads through Python return values. + +``permanent`` + Another ``StorageManager`` targeting long–term storage. Results saved here + survive beyond the life of the ``ProtocolDAG`` run (for example for + inspection or for reuse in future extensions). + +``stdout`` / ``stderr`` + Optional directories where the engine captures subprocess output triggered + by the unit. The directories are removed automatically when the context + closes. + +Keeping these handles bundled together and managed by a context manager lets +``ProtocolUnit`` implementers focus on domain logic while the engine ensures +storage gets flushed and temporary directories are cleaned. + + +Lifecycle +--------- + +``Context`` implements ``__enter__``/``__exit__`` so the execution engine can +rely on ``with Context(...) as ctx`` semantics internally. When ``__exit__`` +triggers the ``shared`` and ``permanent`` storage managers flush pending +transfers back to their underlying ``ExternalStorage`` instances. Any +``stdout`` or ``stderr`` capture directories are also removed. + +This means a ``ProtocolUnit`` should treat ``ctx.shared`` and +``ctx.permanent`` as handles that stay valid only for the duration of +``_execute``; once the method returns, the engine is responsible for moving the +files into their durable homes. + + +Using Context inside ProtocolUnits +---------------------------------- + +Every ``ProtocolUnit._execute`` definition must accept ``ctx: Context`` as its +first argument. Typical usage looks like the example below. + +.. code-block:: python + + from gufe import ProtocolUnit, Context + + class SimulationUnit(ProtocolUnit): + + @staticmethod + def _execute(ctx: Context, *, setup_result, lambda_window, settings): + scratch_path = ctx.scratch / f"lambda_{lambda_window}" + scratch_path.mkdir(exist_ok=True) + + # Read upstream artifacts from ctx.shared + system_file = setup_result.outputs["system_file"] + topology_file = setup_result.outputs["topology_file"] + + # Produce large payloads into ctx.shared so downstream units can + # access them without serialization + result_path = ctx.shared / f"window_{lambda_window}.npz" + save_simulation_outputs(result_path) + + # Return only lightweight metadata + return { + "lambda_window": lambda_window, + "result_path": str(result_path), + } + + +Choosing between shared and permanent storage +--------------------------------------------- + +Both ``ctx.shared`` and ``ctx.permanent`` expose the same ``StorageManager`` +API but they serve different audiences: + +``ctx.shared`` + Optimized for communication between units in the same DAG execution. The + execution backend is free to prune these assets once no downstream unit + references them. + +``ctx.permanent`` + Intended for outputs that should survive beyond the immediate DAG, such as + user-facing reports or artifacts that will seed future ``extends`` runs. + +As a rule of thumb, prefer ``ctx.shared`` unless you have a clear requirement +to keep the data after the ``Protocol`` run concludes. Small scalar values or +lightweight metadata should still be returned directly from ``_execute`` so +they become part of the ``ProtocolUnitResult`` record. + + +Interaction with Protocols +-------------------------- + +``Protocol`` instances do not instantiate ``Context`` directly; they declare +``ProtocolUnit`` objects via ``Protocol._create``. When an execution backend +(for example, ``gufe-client`` or a workflow manager) walks the resulting +``ProtocolDAG`` it constructs a ``Context`` for each unit using the DAG label, +unit label, scratch directory, and the configured +``ExternalStorage`` implementations. The backend might provide different +``ExternalStorage`` implementations (e.g., local filesystem, object store, +cluster scratch) depending on where the work runs, but the ``Context`` API seen +by ``ProtocolUnit`` authors stays consistent. + +Because the execution backend is in charge of creating the contexts, protocol +authors can rely on ``ctx`` always being populated with valid storage managers +and paths that are safe to write to from distributed workers. + + +Practical tips +-------------- + +* Keep return values small. Record numeric summaries, filenames, and status in + the ``dict`` returned by ``_execute``; stream heavy data through + ``ctx.shared`` or ``ctx.permanent``. +* Clean up unit-specific scratch files as you go if they consume lots of disk; + the execution engine only guarantees cleanup once ``_execute`` finishes. +* If you launch subprocesses, direct their output into ``ctx.stdout`` and + ``ctx.stderr`` (when provided) to make debugging easier. +* ``StorageManager`` exposes convenience helpers such as + :meth:`~gufe.storage.storagemanager.StorageManager.store_path` for shuttling + directories or streams; use these instead of manual ``shutil`` calls. + +By following these patterns, ``ProtocolUnit`` implementations remain portable +across execution backends while benefiting from consistent lifecycle and +storage management. + + +Migrating from legacy Context usage +----------------------------------- + +Before ``Context`` was rewritten in :mod:`gufe 0.15`, it was a simple data +class with two ``pathlib.Path`` handles: ``scratch`` and ``shared``. Existing +protocols adopted a variety of implicit conventions around those attributes. +The current implementation keeps backwards compatibility at the interface +level—``ctx.scratch`` and ``ctx.shared`` still exist—but their behavior is more +structured because they are ``StorageManager`` instances backed by +``ExternalStorage``. + +Follow this checklist when migrating old protocols: + +1. **Swap file paths to StorageManager APIs.** Calls like ``ctx.shared / + "filename"`` should be replaced with the helper methods offered by + :class:`StorageManager`. For example:: + + path = ctx.shared.scratch_dir / "myfile.dat" + ctx.shared.register("myfile.dat") + + or use :meth:`StorageManager.store_path` when moving directories. Register + every artifact that must be transferred off the worker; unregistered files + stay local and will be deleted. + +2. **Avoid storing heavy objects in Python outputs.** Older protocols often + returned raw ``Path`` objects pointing at scratch files. Instead, register + the file with the storage manager and return the storage key (a string) from + ``_execute`` as shown in the example above. Downstream units can then call + ``ctx.shared.storage.load_stream`` or ``StorageManager.load``. + +3. **Handle ``ctx.permanent``.** There was no equivalent in the legacy API. + Decide which results must persist between DAG executions and write them via + ``ctx.permanent``. For migration you can start by mirroring whatever used + to live in ``ctx.shared`` and refine later. + +4. **Expect automatic cleanup.** Old contexts typically left stdout/stderr + directories around. The new context removes these when the unit finishes. + If your code tried to re-read the capture directories after ``_execute`` + returned, move that logic earlier or rely on the logged data captured in the + ``ProtocolUnitResult``. + +5. **Stop constructing Context manually.** Some bespoke execution scripts once + instantiated ``Context(scratch=..., shared=...)`` by hand. That pattern is + obsolete because the constructor now requires ``ExternalStorage`` objects. + Instead, rely on the execution backend (``gufe-client``, custom schedulers, + etc.) to build contexts. For unit tests use the helpers in + ``gufe.storage.externalresource`` (e.g., ``MemoryStorage``) to create the + necessary storage instances. + +During the transition it is safe to read ``ctx.shared`` as a ``Path`` so long as +you treat it as read-only and only for temporary files; the attribute still +implements ``pathlib.Path`` thanks to the embedded ``StorageManager``'s +``scratch_dir``. Prefer using the explicit ``scratch_dir`` attribute to make +the intent clear when touching raw filesystem paths. diff --git a/docs/concepts/index.rst b/docs/concepts/index.rst index 6e8007e0..53692be4 100644 --- a/docs/concepts/index.rst +++ b/docs/concepts/index.rst @@ -13,3 +13,4 @@ utilize **gufe** APIs. included_models serialization storage + context From edd1d8bfcb5306a7e651d30d4a9d9df357418d4e Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Wed, 21 Jan 2026 16:51:45 -0700 Subject: [PATCH 03/15] docs: improve context docs --- docs/concepts/context.rst | 112 +++++++++++--------------------------- 1 file changed, 32 insertions(+), 80 deletions(-) diff --git a/docs/concepts/context.rst b/docs/concepts/context.rst index 34b41cfd..62bbc0a0 100644 --- a/docs/concepts/context.rst +++ b/docs/concepts/context.rst @@ -1,12 +1,9 @@ Context ======= -``Context`` instances carry the execution environment for individual -``ProtocolUnit`` executions. They are created by the execution engine just -before a unit's ``_execute`` method is called and then discarded once the unit -returns. The class lives in ``gufe.protocols.protocolunit`` and acts as a thin -wrapper around two :class:`~gufe.storage.storagemanager.StorageManager` -objects and a scratch directory. +:class:`.Context` instances carry the execution environment for individual :class:`.ProtocolUnit` executions. +They are created by the execution engine just before a unit is excuted and discarded once the unit returns. +The class acts as a thin wrapper around two :class:`.StorageManager` objects and a scratch directory. Why Context exists @@ -20,14 +17,14 @@ Why Context exists soon as the unit finishes. ``shared`` - A :class:`~gufe.storage.storagemanager.StorageManager` backed by - short–lived :class:`~gufe.storage.externalresource.ExternalStorage`. Use + A :class:`.StorageManager` backed by + short–lived :class:`.ExternalStorage`. Use this to hand large files to downstream units without serializing the payloads through Python return values. ``permanent`` - Another ``StorageManager`` targeting long–term storage. Results saved here - survive beyond the life of the ``ProtocolDAG`` run (for example for + Another :class:`.StorageManager` targeting long–term storage. Results saved here + survive beyond the life of the :class:`.ProtocolDAG` run (for example for inspection or for reuse in future extensions). ``stdout`` / ``stderr`` @@ -43,22 +40,19 @@ storage gets flushed and temporary directories are cleaned. Lifecycle --------- -``Context`` implements ``__enter__``/``__exit__`` so the execution engine can -rely on ``with Context(...) as ctx`` semantics internally. When ``__exit__`` -triggers the ``shared`` and ``permanent`` storage managers flush pending -transfers back to their underlying ``ExternalStorage`` instances. Any -``stdout`` or ``stderr`` capture directories are also removed. +``Context`` implements a `Python context manager`_. +When the context is exited, ``shared`` and ``permanent`` storage managers flushed tracked files back to their underlying :class:``ExternalStorage``. +Any ``stdout`` or ``stderr`` capture directories are also removed. -This means a ``ProtocolUnit`` should treat ``ctx.shared`` and -``ctx.permanent`` as handles that stay valid only for the duration of -``_execute``; once the method returns, the engine is responsible for moving the -files into their durable homes. +.. _Python context manager: https://docs.python.org/3/reference/datamodel.html#context-managers + +This means a ``ProtocolUnit`` should treat ``ctx.shared`` and ``ctx.permanent`` as handles that stay valid only for the duration of execution; once the method returns, the engine is responsible for moving the files into their durable homes. Using Context inside ProtocolUnits ---------------------------------- -Every ``ProtocolUnit._execute`` definition must accept ``ctx: Context`` as its +Every :meth:`.ProtocolUnit._execute` definition must accept ``ctx: Context`` as its first argument. Typical usage looks like the example below. .. code-block:: python @@ -91,7 +85,7 @@ first argument. Typical usage looks like the example below. Choosing between shared and permanent storage --------------------------------------------- -Both ``ctx.shared`` and ``ctx.permanent`` expose the same ``StorageManager`` +Both ``ctx.shared`` and ``ctx.permanent`` expose the same :class:`.StorageManager` API but they serve different audiences: ``ctx.shared`` @@ -101,7 +95,7 @@ API but they serve different audiences: ``ctx.permanent`` Intended for outputs that should survive beyond the immediate DAG, such as - user-facing reports or artifacts that will seed future ``extends`` runs. + user-facing reports or artifacts that will seed future runs. As a rule of thumb, prefer ``ctx.shared`` unless you have a clear requirement to keep the data after the ``Protocol`` run concludes. Small scalar values or @@ -112,71 +106,37 @@ they become part of the ``ProtocolUnitResult`` record. Interaction with Protocols -------------------------- -``Protocol`` instances do not instantiate ``Context`` directly; they declare -``ProtocolUnit`` objects via ``Protocol._create``. When an execution backend -(for example, ``gufe-client`` or a workflow manager) walks the resulting -``ProtocolDAG`` it constructs a ``Context`` for each unit using the DAG label, -unit label, scratch directory, and the configured -``ExternalStorage`` implementations. The backend might provide different -``ExternalStorage`` implementations (e.g., local filesystem, object store, -cluster scratch) depending on where the work runs, but the ``Context`` API seen -by ``ProtocolUnit`` authors stays consistent. - -Because the execution backend is in charge of creating the contexts, protocol -authors can rely on ``ctx`` always being populated with valid storage managers -and paths that are safe to write to from distributed workers. - - -Practical tips --------------- - -* Keep return values small. Record numeric summaries, filenames, and status in - the ``dict`` returned by ``_execute``; stream heavy data through - ``ctx.shared`` or ``ctx.permanent``. -* Clean up unit-specific scratch files as you go if they consume lots of disk; - the execution engine only guarantees cleanup once ``_execute`` finishes. -* If you launch subprocesses, direct their output into ``ctx.stdout`` and - ``ctx.stderr`` (when provided) to make debugging easier. -* ``StorageManager`` exposes convenience helpers such as - :meth:`~gufe.storage.storagemanager.StorageManager.store_path` for shuttling - directories or streams; use these instead of manual ``shutil`` calls. +``Protocol`` instances do not instantiate ``Context`` directly; they declare ``ProtocolUnit`` objects via ``Protocol._create``. +When an execution backend walks the resulting +``ProtocolDAG`` it constructs a ``Context`` for each unit using the DAG label, unit label, scratch directory, and the configured ``ExternalStorage`` implementations. +The backend might provide different ``ExternalStorage`` implementations (e.g., local filesystem, object store, cluster scratch) depending on where the work runs, but the ``Context`` API seen by ``ProtocolUnit`` authors stays consistent. -By following these patterns, ``ProtocolUnit`` implementations remain portable -across execution backends while benefiting from consistent lifecycle and -storage management. +Because the execution backend is in charge of creating the contexts, protocol authors can rely on ``ctx`` always being populated with valid storage managers and paths that are safe to write to from distributed workers. Migrating from legacy Context usage ----------------------------------- -Before ``Context`` was rewritten in :mod:`gufe 0.15`, it was a simple data -class with two ``pathlib.Path`` handles: ``scratch`` and ``shared``. Existing -protocols adopted a variety of implicit conventions around those attributes. -The current implementation keeps backwards compatibility at the interface -level—``ctx.scratch`` and ``ctx.shared`` still exist—but their behavior is more -structured because they are ``StorageManager`` instances backed by -``ExternalStorage``. - +Before ``Context`` was rewritten, it was a simple data class with two ``pathlib.Path`` handles: ``scratch`` and ``shared``. +Existing protocols adopted a variety of implicit conventions around those attributes. Follow this checklist when migrating old protocols: 1. **Swap file paths to StorageManager APIs.** Calls like ``ctx.shared / "filename"`` should be replaced with the helper methods offered by :class:`StorageManager`. For example:: +.. code-block:: python + path = ctx.shared.scratch_dir / "myfile.dat" ctx.shared.register("myfile.dat") - or use :meth:`StorageManager.store_path` when moving directories. Register - every artifact that must be transferred off the worker; unregistered files - stay local and will be deleted. - 2. **Avoid storing heavy objects in Python outputs.** Older protocols often returned raw ``Path`` objects pointing at scratch files. Instead, register the file with the storage manager and return the storage key (a string) from ``_execute`` as shown in the example above. Downstream units can then call - ``ctx.shared.storage.load_stream`` or ``StorageManager.load``. + :meth:`.StorageManager.load`. -3. **Handle ``ctx.permanent``.** There was no equivalent in the legacy API. +3. **Handle ctx.permanent.** There was no equivalent in the legacy API. Decide which results must persist between DAG executions and write them via ``ctx.permanent``. For migration you can start by mirroring whatever used to live in ``ctx.shared`` and refine later. @@ -187,16 +147,8 @@ Follow this checklist when migrating old protocols: returned, move that logic earlier or rely on the logged data captured in the ``ProtocolUnitResult``. -5. **Stop constructing Context manually.** Some bespoke execution scripts once - instantiated ``Context(scratch=..., shared=...)`` by hand. That pattern is - obsolete because the constructor now requires ``ExternalStorage`` objects. - Instead, rely on the execution backend (``gufe-client``, custom schedulers, - etc.) to build contexts. For unit tests use the helpers in - ``gufe.storage.externalresource`` (e.g., ``MemoryStorage``) to create the - necessary storage instances. - -During the transition it is safe to read ``ctx.shared`` as a ``Path`` so long as -you treat it as read-only and only for temporary files; the attribute still -implements ``pathlib.Path`` thanks to the embedded ``StorageManager``'s -``scratch_dir``. Prefer using the explicit ``scratch_dir`` attribute to make -the intent clear when touching raw filesystem paths. +5. **Stop constructing Context manually.** + Some bespoke execution scripts once instantiated ``Context(scratch=..., shared=...)`` by hand. + That pattern is obsolete because the constructor now requires ``ExternalStorage`` objects. + Instead, rely on the execution backend to build contexts. + For unit tests use the helpers in ``gufe.storage.externalresource`` (e.g., ``MemoryStorage``) to create the necessary storage instances. From 66a40fba3dae5f59717d81186a060b1548a709f4 Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Wed, 21 Jan 2026 16:54:57 -0700 Subject: [PATCH 04/15] docs: migrate how to write a custom implementation --- docs/concepts/storage.rst | 73 ++++++++++++++++++++------------------- 1 file changed, 37 insertions(+), 36 deletions(-) diff --git a/docs/concepts/storage.rst b/docs/concepts/storage.rst index a9454eae..bffa057e 100644 --- a/docs/concepts/storage.rst +++ b/docs/concepts/storage.rst @@ -90,42 +90,6 @@ The :class:`.MemoryStorage` implementation stores data in a Python dictionary. T .. warning:: ``MemoryStorage`` is not intended for production use and all data is lost when the Python process exits. -StorageManager --------------- - -The :class:`.StorageManager` class provides a higher-level interface for managing storage operations within a computational workflow. -It handles the transfer of files between a scratch directory and external storage: - -.. code-block:: python - - from pathlib import Path - from gufe.storage import StorageManager - from gufe.storage.externalresource import FileStorage - - # Set up storage - storage = FileStorage(root_dir=Path("/path/to/storage")) - scratch_dir = Path("/path/to/scratch") - - # Create a storage manager for a specific DAG and unit - manager = StorageManager( - scratch_dir=scratch_dir, - storage=storage, - dag_label="my_experiment", - unit_label="transformation_1" - ) - - # Register files for later transfer - manager.register("trajectory.dcd") - manager.register("results.json") - - # Transfer all registered files to external storage - manager._transfer() - - # Load files from external storage - trajectory_data = manager.load("my_experiment/transformation_1/trajectory.dcd") - -The ``StorageManager`` uses a namespace combining the ``dag_label`` and ``unit_label`` to organize files in the external storage backend. - Implementing Custom Storage Backends ------------------------------------- @@ -182,6 +146,43 @@ To create a custom storage backend, subclass :class:`.ExternalStorage` and imple .. note:: All storage methods should be blocking operations, even if the underlying storage backend supports asynchronous operations. +StorageManager +-------------- + +The :class:`.StorageManager` class provides a higher-level interface for managing storage operations within a computational workflow. +It handles the transfer of files between a scratch directory and external storage: + +.. code-block:: python + + from pathlib import Path + from gufe.storage import StorageManager + from gufe.storage.externalresource import FileStorage + + # Set up storage + storage = FileStorage(root_dir=Path("/path/to/storage")) + scratch_dir = Path("/path/to/scratch") + + # Create a storage manager for a specific DAG and unit + manager = StorageManager( + scratch_dir=scratch_dir, + storage=storage, + dag_label="my_experiment", + unit_label="transformation_1" + ) + + # Register files for later transfer + manager.register("trajectory.dcd") + manager.register("results.json") + + # Transfer all registered files to external storage + manager._transfer() + + # Load files from external storage + trajectory_data = manager.load("my_experiment/transformation_1/trajectory.dcd") + +The ``StorageManager`` uses a namespace combining the ``dag_label`` and ``unit_label`` to organize files in the external storage backend. + + Error Handling -------------- From adbad03dfb6b854f3132808abb6db172e64028c9 Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Fri, 23 Jan 2026 10:12:37 -0700 Subject: [PATCH 05/15] doc: update with comments from review Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> --- docs/concepts/context.rst | 6 +++--- docs/concepts/storage.rst | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/concepts/context.rst b/docs/concepts/context.rst index 62bbc0a0..93f15635 100644 --- a/docs/concepts/context.rst +++ b/docs/concepts/context.rst @@ -1,9 +1,9 @@ -Context +``ProtocolUnit`` execution ``Context`` ======= :class:`.Context` instances carry the execution environment for individual :class:`.ProtocolUnit` executions. They are created by the execution engine just before a unit is excuted and discarded once the unit returns. -The class acts as a thin wrapper around two :class:`.StorageManager` objects and a scratch directory. +The class acts as a thin wrapper around two :class:`.StorageManager` objects (shared and permanent) and a scratch directory. Why Context exists @@ -29,7 +29,7 @@ Why Context exists ``stdout`` / ``stderr`` Optional directories where the engine captures subprocess output triggered - by the unit. The directories are removed automatically when the context + by the ``ProtocolUnit``. The directories are removed automatically when the context closes. Keeping these handles bundled together and managed by a context manager lets diff --git a/docs/concepts/storage.rst b/docs/concepts/storage.rst index bffa057e..92d658cf 100644 --- a/docs/concepts/storage.rst +++ b/docs/concepts/storage.rst @@ -3,13 +3,13 @@ How storage is handled in **gufe** ================================== -**gufe** abstracts storage into a reusable storage interface using the :class:`.ExternalStorage` abstract base class. +**gufe** abstracts storage into a reusable interface using the :class:`.ExternalStorage` abstract base class. This abstraction enables the storage of any file or byte stream using various storage backends without changing application code. Overview -------- -The storage system is designed to handle (files or byte data) that need to be stored in some location. +The storage system is designed to handle files (or byte data) that need to be stored in some location. Instead of embedding the data, objects can store a reference (a *location* string) to where the data is stored externally. This approach provides several benefits: * **Efficiency**: Large objects don't need to be serialized multiple times From 94469fd934ffac53acb2aa3c61aa128b7798237c Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Fri, 23 Jan 2026 09:21:16 -0800 Subject: [PATCH 06/15] fix docs build --- docs/concepts/context.rst | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/concepts/context.rst b/docs/concepts/context.rst index 93f15635..4284d095 100644 --- a/docs/concepts/context.rst +++ b/docs/concepts/context.rst @@ -123,12 +123,13 @@ Follow this checklist when migrating old protocols: 1. **Swap file paths to StorageManager APIs.** Calls like ``ctx.shared / "filename"`` should be replaced with the helper methods offered by - :class:`StorageManager`. For example:: + :class:`StorageManager`. For example: .. code-block:: python - path = ctx.shared.scratch_dir / "myfile.dat" - ctx.shared.register("myfile.dat") + path = ctx.shared.scratch_dir / "myfile.dat" + ctx.shared.register("myfile.dat") + 2. **Avoid storing heavy objects in Python outputs.** Older protocols often returned raw ``Path`` objects pointing at scratch files. Instead, register From 355a4ceb97da9ca2a8c915117a56076312fdede9 Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Fri, 23 Jan 2026 17:15:12 -0700 Subject: [PATCH 07/15] docs: Add information on changes in feat/return-storage-handles This adds new information on how storage is handled, and how storage gets passed around by units. --- docs/concepts/context.rst | 30 ++++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/docs/concepts/context.rst b/docs/concepts/context.rst index 4284d095..7637dd59 100644 --- a/docs/concepts/context.rst +++ b/docs/concepts/context.rst @@ -18,12 +18,12 @@ Why Context exists ``shared`` A :class:`.StorageManager` backed by - short–lived :class:`.ExternalStorage`. Use + short–lived :class:`.ExternalStorage`. Use this to hand large files to downstream units without serializing the payloads through Python return values. ``permanent`` - Another :class:`.StorageManager` targeting long–term storage. Results saved here + Another :class:`.StorageManager` targeting long–term storage. Results saved here survive beyond the life of the :class:`.ProtocolDAG` run (for example for inspection or for reuse in future extensions). @@ -46,7 +46,11 @@ Any ``stdout`` or ``stderr`` capture directories are also removed. .. _Python context manager: https://docs.python.org/3/reference/datamodel.html#context-managers -This means a ``ProtocolUnit`` should treat ``ctx.shared`` and ``ctx.permanent`` as handles that stay valid only for the duration of execution; once the method returns, the engine is responsible for moving the files into their durable homes. +This means each ``ProtocolUnit``'s ``shared`` and ``permanent`` object are not paths, and should not be treated as such. +Both of these are registries that track if a file should be transferred from its location in ``scratch`` to its final location after completing a unit. + +If you want to use some from ``shared`` or ``permanent``, you can use ``ctx.shared.load`` or ``ctx.permanent.load``. +This will allow your unit to fetch those objects from their storage for use. Using Context inside ProtocolUnits @@ -67,20 +71,26 @@ first argument. Typical usage looks like the example below. scratch_path.mkdir(exist_ok=True) # Read upstream artifacts from ctx.shared - system_file = setup_result.outputs["system_file"] - topology_file = setup_result.outputs["topology_file"] + system_file = ctx.shared.load(setup_result.outputs["system_file"]) + topology_file = ctx.shared.load(setup_result.outputs["topology_file"]) + - # Produce large payloads into ctx.shared so downstream units can - # access them without serialization - result_path = ctx.shared / f"window_{lambda_window}.npz" - save_simulation_outputs(result_path) + result_path = ctx.scratch / "some_output.pdb" + # When you register the filename doesn't matter, + # just as long as you do it before you return + result_path_final_location = ctx.permanent.register(result_path) + # This is an example of running something that you want to save + simulate(output=result_path) # Return only lightweight metadata return { "lambda_window": lambda_window, - "result_path": str(result_path), + # We use this because it is already namespaced and can be used between units. + "result_path": result_path_final_location, } +The example above showcases how + Choosing between shared and permanent storage --------------------------------------------- From 497eb8ca632e0e9cb05ed7c4d2c2a90d0963d08b Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Fri, 23 Jan 2026 17:16:43 -0700 Subject: [PATCH 08/15] docs: add information on pre-namespaced objects from other branch --- docs/concepts/storage.rst | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/concepts/storage.rst b/docs/concepts/storage.rst index 92d658cf..04b0d6a5 100644 --- a/docs/concepts/storage.rst +++ b/docs/concepts/storage.rst @@ -171,8 +171,9 @@ It handles the transfer of files between a scratch directory and external storag ) # Register files for later transfer - manager.register("trajectory.dcd") - manager.register("results.json") + out = manager.register("trajectory.dcd") + out2 = manager.register("results.json") + # Note: out and out2 are pre-namespaced values that allow storage items to be passed around # Transfer all registered files to external storage manager._transfer() From 200973a0f53085cda098cafd89818193e4036e01 Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Wed, 28 Jan 2026 13:30:59 -0700 Subject: [PATCH 09/15] docs: add suggestion to back link to Context docs --- docs/concepts/storage.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/concepts/storage.rst b/docs/concepts/storage.rst index 04b0d6a5..0696c7b6 100644 --- a/docs/concepts/storage.rst +++ b/docs/concepts/storage.rst @@ -182,6 +182,7 @@ It handles the transfer of files between a scratch directory and external storag trajectory_data = manager.load("my_experiment/transformation_1/trajectory.dcd") The ``StorageManager`` uses a namespace combining the ``dag_label`` and ``unit_label`` to organize files in the external storage backend. +To see how these work in practice see our documentation on :doc:`context`. Error Handling From ba7657973529d19d85f1037c81df7bafb768f173 Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Wed, 28 Jan 2026 13:36:40 -0700 Subject: [PATCH 10/15] docs: add becomes in migration guide per suggestion --- docs/concepts/context.rst | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/concepts/context.rst b/docs/concepts/context.rst index 7637dd59..4c58b8c6 100644 --- a/docs/concepts/context.rst +++ b/docs/concepts/context.rst @@ -1,5 +1,5 @@ ``ProtocolUnit`` execution ``Context`` -======= +====================================== :class:`.Context` instances carry the execution environment for individual :class:`.ProtocolUnit` executions. They are created by the execution engine just before a unit is excuted and discarded once the unit returns. @@ -138,6 +138,11 @@ Follow this checklist when migrating old protocols: .. code-block:: python path = ctx.shared.scratch_dir / "myfile.dat" + +becomes: + +.. code-block:: python + ctx.shared.register("myfile.dat") From 3fbe18f864d89f30847ad11d6826e2188a544f0c Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Wed, 28 Jan 2026 13:39:42 -0700 Subject: [PATCH 11/15] docs: add link to protocol how to --- docs/concepts/context.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/concepts/context.rst b/docs/concepts/context.rst index 4c58b8c6..d1b1e1a2 100644 --- a/docs/concepts/context.rst +++ b/docs/concepts/context.rst @@ -168,3 +168,5 @@ becomes: That pattern is obsolete because the constructor now requires ``ExternalStorage`` objects. Instead, rely on the execution backend to build contexts. For unit tests use the helpers in ``gufe.storage.externalresource`` (e.g., ``MemoryStorage``) to create the necessary storage instances. + +If you need more information on how to use these concepts, checkout out: :doc:`../how-tos/protocol`. From 8f1ce78dbc7304a97722aae70f2c20fcb3101577 Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Wed, 28 Jan 2026 17:09:07 -0700 Subject: [PATCH 12/15] docs: add more informaton on StorageManagers --- docs/concepts/storage.rst | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/concepts/storage.rst b/docs/concepts/storage.rst index 0696c7b6..462a8165 100644 --- a/docs/concepts/storage.rst +++ b/docs/concepts/storage.rst @@ -150,7 +150,9 @@ StorageManager -------------- The :class:`.StorageManager` class provides a higher-level interface for managing storage operations within a computational workflow. -It handles the transfer of files between a scratch directory and external storage: +This class is largely used by the :class:`.Context` class and should not be instantiated in protocols. +In general, protocol developers will only use the register and load functions. +It handles the transfer of files between a scratch directory and external storage (such as shared or permenant storage): .. code-block:: python From 83ee9295340ed0a351e862181003d329689049f5 Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Wed, 28 Jan 2026 17:09:52 -0700 Subject: [PATCH 13/15] docs: clean up storage code example --- docs/concepts/storage.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/concepts/storage.rst b/docs/concepts/storage.rst index 462a8165..52633782 100644 --- a/docs/concepts/storage.rst +++ b/docs/concepts/storage.rst @@ -181,7 +181,8 @@ It handles the transfer of files between a scratch directory and external storag manager._transfer() # Load files from external storage - trajectory_data = manager.load("my_experiment/transformation_1/trajectory.dcd") + trajectory_data = manager.load(out) + results_json = manager.load(out2) The ``StorageManager`` uses a namespace combining the ``dag_label`` and ``unit_label`` to organize files in the external storage backend. To see how these work in practice see our documentation on :doc:`context`. From be26dfc8866fea1ed17d52d86f2c399bf4ffa80c Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Thu, 29 Jan 2026 14:23:57 -0700 Subject: [PATCH 14/15] Apply suggestions from code review Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> --- docs/concepts/storage.rst | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/docs/concepts/storage.rst b/docs/concepts/storage.rst index 52633782..4dcdec36 100644 --- a/docs/concepts/storage.rst +++ b/docs/concepts/storage.rst @@ -10,7 +10,7 @@ Overview -------- The storage system is designed to handle files (or byte data) that need to be stored in some location. -Instead of embedding the data, objects can store a reference (a *location* string) to where the data is stored externally. This approach provides several benefits: +Instead of embedding the data, objects can store a reference (a unique string indicating the object's location, such as a path) to where the data is stored externally. This approach provides several benefits: * **Efficiency**: Large objects don't need to be serialized multiple times * **Flexibility**: Different storage backends (local filesystem, cloud storage, in-memory) can be used interchangeably @@ -150,8 +150,11 @@ StorageManager -------------- The :class:`.StorageManager` class provides a higher-level interface for managing storage operations within a computational workflow. -This class is largely used by the :class:`.Context` class and should not be instantiated in protocols. -In general, protocol developers will only use the register and load functions. +.. note:: + + ``StorageManager`` is largely used by the :class:`.Context` class and should not be instantiated in protocols. + In general, protocol developers will only use the ``register`` and ``load`` functions. + It handles the transfer of files between a scratch directory and external storage (such as shared or permenant storage): .. code-block:: python From bdc284a615ca0a06e652829ebf9b25860fa6063c Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Thu, 29 Jan 2026 15:16:06 -0800 Subject: [PATCH 15/15] fix typo --- docs/concepts/storage.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/concepts/storage.rst b/docs/concepts/storage.rst index 4dcdec36..59e4ed0b 100644 --- a/docs/concepts/storage.rst +++ b/docs/concepts/storage.rst @@ -155,7 +155,7 @@ The :class:`.StorageManager` class provides a higher-level interface for managin ``StorageManager`` is largely used by the :class:`.Context` class and should not be instantiated in protocols. In general, protocol developers will only use the ``register`` and ``load`` functions. -It handles the transfer of files between a scratch directory and external storage (such as shared or permenant storage): +It handles the transfer of files between a scratch directory and external storage (such as shared or permanent storage): .. code-block:: python