[DNM] Add AlchemicalArchive extraction

alchemiscale/interface/client.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import asyncio
+import json
 from enum import StrEnum
 from typing import Any, Literal
 from collections.abc import Iterable
@@ -16,7 +17,8 @@
 from async_lru import alru_cache
 import networkx as nx
 from gufe import AlchemicalNetwork, Transformation, ChemicalSystem
-from gufe.tokenization import GufeTokenizable, KeyedChain
+from gufe.archival import AlchemicalArchive
+from gufe.tokenization import GufeTokenizable, KeyedChain, JSON_HANDLER
 from gufe.protocols import ProtocolResult, ProtocolDAGResult
 import zstandard as zstd
 
@@ -34,6 +36,7 @@
     StrategyState,
 )
 from stratocaster.base import Strategy
+from ..utils import pdr_from_bytes
 from ..validators import validate_network_nonself
 
 from warnings import warn
@@ -1536,12 +1539,7 @@ async def _async_get_protocoldagresult(
                     pdr_bytes,
                 )
 
-        try:
-            # Attempt to decompress the ProtocolDAGResult object
-            pdr = decompress_gufe_zstd(pdr_bytes)
-        except zstd.ZstdError:
-            # If decompress fails, assume it's a UTF-8 encoded JSON string
-            pdr = json_to_gufe(pdr_bytes.decode("utf-8"))
+        pdr = pdr_from_bytes(pdr_bytes)
 
         return pdr
 
@@ -1767,6 +1765,164 @@ def get_network_failures(
             network=network, ok=False, compress=compress, visualize=visualize
         )
 
+    def get_network_archives(
+        self,
+        networks: list[ScopedKey | str],
+        metadata: list[dict | None] | None = None,
+        compress: bool = True,
+        visualize: bool = True,
+    ) -> list[AlchemicalArchive | None]:
+        r"""Produce archival-quality extracts for the given ``AlchemicalNetwork``\s.
+
+        Each returned ``AlchemicalArchive`` bundles an
+        ``AlchemicalNetwork`` together with all successful
+        ``ProtocolDAGResult``\s currently available for its
+        ``Transformation``\s, in a form suitable for long-term storage,
+        sharing, and downstream analysis.
+
+        Parameters
+        ----------
+        networks
+            A list of ``AlchemicalNetwork`` ``ScopedKey`` values. The
+            list must not contain duplicate entries.
+        metadata
+            Metadata to attach to the produced ``AlchemicalArchive``
+            objects. This must be a list of dictionaries that are
+            compatible with ``GufeTokenizable`` serialization, in the
+            same order as ``networks``. A ``None`` entry in the list
+            attaches no metadata to the corresponding
+            ``AlchemicalArchive``. Passing ``None`` in place of the
+            list is interpreted as a list of ``None``, which is the
+            default.
+        compress
+            If ``True``, compress objects server-side before shipping
+            them to the client. This is a performance optimization; it
+            has no bearing on the result of this method call.
+        visualize
+            If ``True``, show retrieval progress indicators.
+
+        Returns
+        -------
+        A list of ``AlchemicalArchive`` instances matching the order of
+        ``networks``. If a network was not found, ``None`` is returned
+        in its place.
+
+        Raises
+        ------
+        ValueError
+            If the provided metadata is not serializable, if the
+            lengths of the ``metadata`` and ``networks`` lists differ,
+            or if ``networks`` contains duplicate entries.
+
+        """
+        network_sks = [
+            ScopedKey.from_str(network) if isinstance(network, str) else network
+            for network in networks
+        ]
+
+        if len(set(network_sks)) != len(network_sks):
+            raise ValueError("`networks` list must not contain duplicate entries")
+
+        metadata = [None] * len(network_sks) if metadata is None else metadata
+
+        if len(metadata) != len(network_sks):
+            raise ValueError("`metadata` and `networks` lists must be the same length")
+
+        # validate that all metadata is serializable up-front, before
+        # performing any (potentially expensive) retrieval
+        for network_sk, meta in zip(network_sks, metadata):
+            if meta is None:
+                continue
+            try:
+                json.dumps(meta, cls=JSON_HANDLER.encoder)
+            except TypeError as e:
+                raise ValueError(
+                    f"Unable to serialize metadata for '{network_sk}': {e}"
+                )
+
+        return [
+            self._get_network_archive(
+                network_sk, meta, compress=compress, visualize=visualize
+            )
+            for network_sk, meta in zip(network_sks, metadata)
+        ]
+
+    def get_network_archive(
+        self,
+        network: ScopedKey | str,
+        metadata: dict | None = None,
+        compress: bool = True,
+        visualize: bool = True,
+    ) -> AlchemicalArchive | None:
+        r"""Produce an archival-quality extract for a given ``AlchemicalNetwork``.
+
+        The returned ``AlchemicalArchive`` bundles the
+        ``AlchemicalNetwork`` together with all successful
+        ``ProtocolDAGResult``\s currently available for its
+        ``Transformation``\s, in a form suitable for long-term storage,
+        sharing, and downstream analysis.
+
+        Parameters
+        ----------
+        network
+            The ``ScopedKey`` of the ``AlchemicalNetwork`` to archive.
+        metadata
+            Metadata to attach to the produced ``AlchemicalArchive``.
+            This must be a dictionary that is compatible with
+            ``GufeTokenizable`` serialization.
+        compress
+            If ``True``, compress objects server-side before shipping
+            them to the client. This is a performance optimization; it
+            has no bearing on the result of this method call.
+        visualize
+            If ``True``, show retrieval progress indicators.
+
+        Returns
+        -------
+        An ``AlchemicalArchive`` for the provided ``AlchemicalNetwork``.
+        If the network was not found, ``None`` is returned.
+
+        Raises
+        ------
+        ValueError
+            If the provided metadata is not serializable.
+
+        """
+        return self.get_network_archives(
+            [network],
+            metadata=None if metadata is None else [metadata],
+            compress=compress,
+            visualize=visualize,
+        )[0]
+
+    def _get_network_archive(
+        self,
+        network: ScopedKey,
+        metadata: dict | None,
+        compress: bool = True,
+        visualize: bool = True,
+    ) -> AlchemicalArchive | None:
+        # returns None if the network does not exist in the given Scope
+        if not self.check_exists(network):
+            return None
+
+        an = self.get_network(network, compress=compress, visualize=visualize)
+
+        transformation_results = []
+        for transformation in an.edges:
+            transformation_sk = ScopedKey(
+                gufe_key=transformation.key, **network.scope.to_dict()
+            )
+            pdrs = self.get_transformation_results(
+                transformation_sk,
+                return_as=ResultFormat.PROTOCOL_DAG_RESULTS,
+                compress=compress,
+                visualize=visualize,
+            )
+            transformation_results.append((transformation, pdrs))
+
+        return AlchemicalArchive(an, transformation_results, metadata=metadata)
+
     def get_transformation_results(
         self,
         transformation: ScopedKey,

alchemiscale/strategist/service.py

@@ -37,6 +37,7 @@
 from ..compression import compress_keyed_chain_zstd, decompress_gufe_zstd, json_to_gufe
 from ..sleep import InterruptableSleep, SleepInterrupted
 from .settings import StrategistSettings
+from ..utils import pdr_from_bytes
 
 
 def execute_strategy_worker(
@@ -173,12 +174,7 @@ def _get_protocoldagresult_cached(
                 location=result_ref.location, ok=result_ref.ok
             )
 
-        # Decompress the raw bytes to get ProtocolDAGResult
-        try:
-            pdr = decompress_gufe_zstd(pdr_bytes)
-        except zstd.ZstdError:
-            # Fallback to JSON deserialization for uncompressed data
-            pdr = json_to_gufe(pdr_bytes.decode("utf-8"))
+        pdr = pdr_from_bytes(pdr_bytes)
 
         if self._cache_enabled:
             try:

alchemiscale/tests/integration/interface/client/test_client.py

@@ -7,6 +7,7 @@
 import json
 
 from gufe import AlchemicalNetwork
+from gufe.archival import AlchemicalArchive
 from gufe.tokenization import TOKENIZABLE_REGISTRY, GufeKey, JSON_HANDLER
 from gufe.protocols import ProtocolResult
 from gufe.protocols.protocoldag import execute_DAG
@@ -2633,6 +2634,101 @@ def test_get_task_failures(
         # TODO: can we mix in a success in here somewhere?
         # not possible with current BrokenProtocol, unfortunately
 
+    def test_get_network_archives(
+        self,
+        scope_test,
+        n4js_preloaded,
+        s3os_server_fresh,
+        user_client_no_cache: client.AlchemiscaleClient,
+        network_tyk2,
+        tmpdir,
+    ):
+        user_client = user_client_no_cache
+        n4js = n4js_preloaded
+        s3os_server = s3os_server_fresh
+
+        an = network_tyk2
+        transformation = list(t for t in an.edges if "_solvent" in t.name)[0]
+
+        network_sk = user_client.get_scoped_key(an, scope_test)
+        transformation_sk = user_client.get_scoped_key(transformation, scope_test)
+
+        # create, action, execute, and push results for three tasks on a
+        # single transformation
+        all_tasks = user_client.create_tasks(transformation_sk, count=3)
+        actioned_tasks = user_client.action_tasks(all_tasks, network_sk)
+
+        with tmpdir.as_cwd():
+            protocoldagresults = self._execute_tasks(actioned_tasks, n4js, s3os_server)
+            self._push_results(actioned_tasks, protocoldagresults, n4js, s3os_server)
+
+        # clear local gufe registry of pdr objects so we exercise
+        # deserialization of the objects pulled back from the server
+        for pdr in protocoldagresults:
+            TOKENIZABLE_REGISTRY.pop(pdr.key, None)
+
+        # produce a single-network archive with metadata attached
+        metadata = {"project": "tyk2", "note": "archival extract"}
+        archive = user_client.get_network_archive(
+            network_sk, metadata=metadata, visualize=False
+        )
+
+        assert isinstance(archive, AlchemicalArchive)
+        assert archive.network == an
+        assert archive.metadata == metadata
+
+        # every Transformation in the network should be represented, but only
+        # the one we computed results for should carry ProtocolDAGResults
+        results_by_key = {tf.key: pdrs for tf, pdrs in archive.transformation_results}
+        assert set(results_by_key) == {tf.key for tf in an.edges}
+
+        for tf in an.edges:
+            if tf.key == transformation.key:
+                assert set(results_by_key[tf.key]) == set(protocoldagresults)
+                for pdr in results_by_key[tf.key]:
+                    assert pdr.ok()
+                    assert pdr.transformation_key == transformation.key
+            else:
+                assert results_by_key[tf.key] == []
+
+        # the archive should round-trip through its serialized form
+        roundtrip = AlchemicalArchive.from_json(content=archive.to_json())
+        assert roundtrip == archive
+        assert roundtrip.metadata == metadata
+
+        # a nonexistent network yields None, both in bulk and singly
+        fake_network_sk = ScopedKey(
+            gufe_key=GufeKey(f"AlchemicalNetwork-{'0' * 32}"),
+            **scope_test.to_dict(),
+        )
+
+        archives = user_client.get_network_archives(
+            [network_sk, fake_network_sk], visualize=False
+        )
+        assert len(archives) == 2
+        assert isinstance(archives[0], AlchemicalArchive)
+        assert archives[0].network == an
+        assert archives[1] is None
+
+        assert user_client.get_network_archive(fake_network_sk, visualize=False) is None
+
+    def test_get_network_archives_validation(
+        self, scope_test, n4js_preloaded, user_client, network_tyk2
+    ):
+        network_sk = user_client.get_scoped_key(network_tyk2, scope_test)
+
+        # duplicate networks are rejected
+        with pytest.raises(ValueError, match="duplicate"):
+            user_client.get_network_archives([network_sk, network_sk])
+
+        # metadata and networks lists must be the same length
+        with pytest.raises(ValueError, match="same length"):
+            user_client.get_network_archives([network_sk], metadata=[{}, {}])
+
+        # unserializable metadata is rejected before any retrieval
+        with pytest.raises(ValueError, match="serialize"):
+            user_client.get_network_archive(network_sk, metadata={"bad": object()})
+
 
 class TestTaskRestartPolicy:
     default_max_retries = 3

alchemiscale/utils.py

@@ -1,3 +1,4 @@
+from gufe.protocols import ProtocolDAGResult
 from gufe.tokenization import (
     GufeTokenizable,
     TOKENIZABLE_REGISTRY,
@@ -7,6 +8,9 @@
 )
 from itertools import chain
 import networkx as nx
+import zstandard as zstd
+
+from .compression import decompress_gufe_zstd, json_to_gufe
 
 
 class RegistryBackup:
@@ -112,3 +116,13 @@ def add_edges(o):
     _ = modify_dependencies(sd, add_edges, is_gufe_obj, mode="encode")
 
     return graph
+
+
+def pdr_from_bytes(buffer: bytes) -> ProtocolDAGResult:
+    try:
+        # Attempt to decompress the ProtocolDAGresult object
+        pdr = decompress_gufe_zstd(buffer)
+    except zstd.ZstdError:
+        # If decompress fails, assume it's a UTF-8 encoded JSON string
+        pdr = json_to_gufe(buffer.decode("utf-8"))
+    return pdr

docs/user_guide/getting_started.rst

@@ -457,3 +457,29 @@ If you wish to pull results for only a single :py:class:`~alchemiscale.storage.m
     [<ProtocolDAGResult-54a3ed32cbd3e3d60d87b2a17519e399>]
 
 You can then iteratively create and action new :py:class:`~alchemiscale.storage.models.Task`\s on your desired :external+gufe:py:class:`~gufe.transformations.transformation.Transformation`\s based on their current estimate and uncertainty, allocating effort where it will be most beneficial.
+
+
+****************************
+Producing archival extracts
+****************************
+
+``alchemiscale`` is an execution system, not an archival system.
+When you want a durable, self-contained snapshot of an :external+gufe:py:class:`~gufe.network.AlchemicalNetwork` together with its existing results — for example to deposit on a `FAIR <https://en.wikipedia.org/wiki/FAIR_data>`_ data repository such as Zenodo, to share with collaborators, or to feed into downstream analysis tooling — you can produce an :external+gufe:py:class:`~gufe.archival.AlchemicalArchive` with::
+
+    >>> an_sk: ScopedKey
+    >>> archive = asc.get_network_archive(an_sk)
+    >>> archive
+    <AlchemicalArchive-6f3c...>
+
+The returned :external+gufe:py:class:`~gufe.archival.AlchemicalArchive` bundles the :external+gufe:py:class:`~gufe.network.AlchemicalNetwork` with all successful :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAGResult`\s currently available for its :external+gufe:py:class:`~gufe.transformations.transformation.Transformation`\s.
+If the network is not found in the given :py:class:`~alchemiscale.models.Scope`, ``None`` is returned.
+
+You can attach arbitrary, serializable metadata to the archive::
+
+    >>> archive = asc.get_network_archive(an_sk, metadata={"doi": "10.5281/zenodo.1234567"})
+
+To produce archives for several networks at once, use :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_network_archives`, optionally passing a list of per-network metadata dictionaries in the same order.
+
+The archive is a :external+gufe:py:class:`~gufe.tokenization.GufeTokenizable`, so it can be written to and read back from disk with its :external+gufe:py:meth:`~gufe.tokenization.GufeTokenizable.to_json` and :external+gufe:py:meth:`~gufe.tokenization.GufeTokenizable.from_json` methods::
+
+    >>> archive.to_json(file="my_network_archive.json")