diff --git a/CLAUDE.md b/CLAUDE.md index 832a26be..c9307441 100644 --- a/CLAUDE.md +++ b/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 diff --git a/alchemiscale/interface/client.py b/alchemiscale/interface/client.py index f53859ba..f4fb7e0d 100644 --- a/alchemiscale/interface/client.py +++ b/alchemiscale/interface/client.py @@ -7,16 +7,18 @@ from __future__ import annotations import asyncio +import json from enum import StrEnum from typing import Any, Literal -from collections.abc import Iterable +from collections.abc import Iterable, Mapping, Sequence from itertools import chain from functools import lru_cache 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,185 @@ 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") + + if metadata is None: + metadata = [None] * len(network_sks) + elif isinstance(metadata, Mapping) or not isinstance(metadata, Sequence): + raise ValueError( + "`metadata` must be a list/sequence of dictionaries or None" + ) + + 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 + + if not isinstance(meta, Mapping): + raise ValueError( + f"Metadata for '{network_sk}' must be a dictionary/mapping or None" + ) + + 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, diff --git a/alchemiscale/strategist/service.py b/alchemiscale/strategist/service.py index a5cc149b..6a9824cd 100644 --- a/alchemiscale/strategist/service.py +++ b/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: diff --git a/alchemiscale/tests/integration/interface/client/test_client.py b/alchemiscale/tests/integration/interface/client/test_client.py index 390acb71..041deb15 100644 --- a/alchemiscale/tests/integration/interface/client/test_client.py +++ b/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,151 @@ 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()}) + + # bulk metadata must be a sequence aligned with networks, not a single dict. + # This is a subtle case because a one-key dict has len == len(networks), + # so a length-only check would pass and iteration would yield the key string. + with pytest.raises(ValueError, match="Metadata.*list|metadata.*sequence"): + user_client.get_network_archives( + [network_sk], + metadata={"project": "tyk2"}, + ) + + # each per-network metadata entry must be a dict or None, even if the value + # is otherwise JSON-serializable + for bad_meta in ["project", ["project", "tyk2"], 1, 1.5, True]: + with pytest.raises(ValueError, match="Metadata.*dict|metadata.*mapping"): + user_client.get_network_archives( + [network_sk], + metadata=[bad_meta], + ) + + # None remains valid as a per-network metadata entry + user_client.get_network_archives( + [network_sk], + metadata=[None], + visualize=False, + ) + class TestTaskRestartPolicy: default_max_retries = 3 diff --git a/alchemiscale/utils.py b/alchemiscale/utils.py index 37cc0b51..45d7ebe4 100644 --- a/alchemiscale/utils.py +++ b/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 diff --git a/devtools/conda-envs/docs.yml b/devtools/conda-envs/docs.yml index f080a0af..86b7b004 100644 --- a/devtools/conda-envs/docs.yml +++ b/devtools/conda-envs/docs.yml @@ -10,6 +10,6 @@ dependencies: - myst-parser>=0.14 - docutils - sphinx-notfound-page - - gufe=1.3.0 + - gufe=1.10.0 - py2neo - stratocaster diff --git a/docs/conf.py b/docs/conf.py index 16eea6ca..00a3cdcb 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -42,7 +42,6 @@ "httpx", "jose", "neo4j", - "pydantic", "pydantic_settings", "starlette", "yaml", diff --git a/docs/user_guide/getting_started.rst b/docs/user_guide/getting_started.rst index f4460887..2fa2313b 100644 --- a/docs/user_guide/getting_started.rst +++ b/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 [] 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 `_ data repository such as Zenodo, to share with collaborators, or to feed into downstream analysis tooling — you can produce an ``AlchemicalArchive`` with:: + + >>> an_sk: ScopedKey + >>> archive = asc.get_network_archive(an_sk) + >>> archive + + +The returned ``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") diff --git a/news/issue-246.rst b/news/issue-246.rst new file mode 100644 index 00000000..86eab94e --- /dev/null +++ b/news/issue-246.rst @@ -0,0 +1,23 @@ +**Added:** + +* ``AlchemiscaleClient.get_network_archive`` and ``AlchemiscaleClient.get_network_archives`` produce archival-quality ``AlchemicalArchive`` extracts, bundling an ``AlchemicalNetwork`` with all successful ``ProtocolDAGResult``\s for its ``Transformation``\s, with optional user-supplied metadata. Networks that are not found yield ``None``. + +**Changed:** + +* + +**Deprecated:** + +* + +**Removed:** + +* + +**Fixed:** + +* + +**Security:** + +*