Add AlchemicalArchive extraction

CLAUDE.md

@@ -75,6 +75,7 @@ The system has four main services communicating through Neo4j and S3:
 - **Frozen Pydantic models**: Core models use `frozen=True` for immutability.
 - **Module docstrings**: Follow the pattern `:mod:\`alchemiscale.module_name\` --- description`.
 - **GUFE integration**: Core chemistry types (`AlchemicalNetwork`, `Transformation`, `ChemicalSystem`, `Protocol`) come from the `gufe` library and are stored/retrieved via their tokenization system (`GufeKey`, `GufeTokenizable`).
+- **`GufeKey` is not stable across `gufe` versions**: A `ScopedKey` is stable once created (it carries the *ingestion-time* `gufe_key`), but a `GufeTokenizable` deserialized later may recompute a *different* `GufeKey` if the `gufe` tokenization has changed since ingestion. Never derive a `ScopedKey` from a freshly-deserialized object's `.key` (e.g. `ScopedKey(gufe_key=obj.key, **scope.to_dict())`) to join against stored data — such a join can silently miss instead of erroring. Instead, obtain the authoritative `ScopedKey` from the server (e.g. `get_network_transformations`, or the keys of `get_network_results`) and fetch the object *by that `ScopedKey`* (e.g. `get_transformation(sk)`). Comparisons purely among objects deserialized in the same process (e.g. "is this transformation in `network.edges`?") are safe, since both sides use the current-stack key consistently.
 
 ## Testing
 

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,174 @@ 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, ValueError) 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)
+
+        # retrieve all successful ProtocolDAGResults for the network's
+        # Transformations in parallel, keyed by their authoritative
+        # (server-side) Transformation ScopedKey
+        results = self.get_network_results(
+            network,
+            return_as=ResultFormat.PROTOCOL_DAG_RESULTS,
+            compress=compress,
+            visualize=visualize,
+        )
+
+        # pair each Transformation with its results by fetching the
+        # Transformation via its authoritative ScopedKey. We deliberately do
+        # not derive a ScopedKey from a deserialized Transformation's GufeKey:
+        # a `Transformation` deserialized now may not reproduce the GufeKey it
+        # had when ingested (gufe tokenization can change across versions), so
+        # its key is not a reliable join to the stored ScopedKey.
+        transformation_results = []
+        for transformation_sk, pdrs in results.items():
+            transformation = self.get_transformation(
+                transformation_sk, compress=compress, visualize=False
+            )
+            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,126 @@ 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 string ScopedKey is accepted and produces an equivalent archive
+        archive_from_str = user_client.get_network_archive(
+            str(network_sk), visualize=False
+        )
+        assert archive_from_str.network == an
+        str_results = {
+            tf.key: pdrs for tf, pdrs in archive_from_str.transformation_results
+        }
+        assert set(str_results[transformation.key]) == set(protocoldagresults)
+
+        # the compress=False path produces an equivalent archive
+        archive_uncompressed = user_client.get_network_archive(
+            network_sk, compress=False, visualize=False
+        )
+        assert archive_uncompressed.network == an
+        uncompressed_results = {
+            tf.key: pdrs for tf, pdrs in archive_uncompressed.transformation_results
+        }
+        assert set(uncompressed_results[transformation.key]) == set(protocoldagresults)
+
+        # 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(),
+        )
+
+        # per-network metadata is attached in order, and a not-found network
+        # yields None in its slot
+        archives = user_client.get_network_archives(
+            [network_sk, fake_network_sk],
+            metadata=[metadata, {"unused": "for missing network"}],
+            visualize=False,
+        )
+        assert len(archives) == 2
+        assert isinstance(archives[0], AlchemicalArchive)
+        assert archives[0].network == an
+        assert archives[0].metadata == metadata
+        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