Skip to content

Commit e04181e

Browse files
ianmkenneydotsdlclaudemikemhenry
authored
Add AlchemicalArchive extraction (#486)
* Add unimplemented functions for archive extraction * Move archive construction to API layer * Allow user to specify metadata * Validate metadata types * Extract PDR from bytes to function * Update requests for archives * Implement client-side AlchemicalNetwork archival extracts Add AlchemiscaleClient.get_network_archive / get_network_archives, which bundle an AlchemicalNetwork with all successful ProtocolDAGResults for its Transformations into a gufe AlchemicalArchive, with optional user-supplied metadata. Not-found networks yield None. Implemented client-side atop the existing get_network and get_transformation_results machinery, replacing the abandoned server-side /bulk/networks/archive endpoint (which referenced undefined names and a non-existent statestore method). Also fixes a missing `json` import in client.py and restores pdr_from_bytes in utils.py (zstd import + decode fallback). Closes #246 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Join archive results via authoritative ScopedKeys; broaden tests Rework AlchemiscaleClient._get_network_archive to pair each Transformation with its results by fetching results in parallel via get_network_results and then retrieving each Transformation by its authoritative (server-side) ScopedKey. A Transformation deserialized now may not reproduce the GufeKey it had at ingestion (gufe tokenization can change across versions), so a reconstructed ScopedKey is not a reliable join and could silently drop results; the server-side ScopedKey is stable. Also broaden metadata-serialization validation to catch ValueError (e.g. circular references), and extend tests to cover string ScopedKey input, the compress=False path, and per-network metadata ordering in the bulk method. Document the GufeKey-instability gotcha in CLAUDE.md. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Fix docs build for AlchemicalArchive The docs build failed to import the package: client.py now imports gufe.archival, absent from the docs env's pinned gufe=1.3.0. Bump it to 1.10.0 to match the runtime pin. That bump surfaced a latent conf.py issue: gufe>=1.10.0 evaluates `PositiveFloat | None` at class-definition time in its settings models, which raises "unsupported operand type(s) for |" when pydantic is mocked via autodoc_mock_imports. Drop pydantic from the mock list (it is a real, installed gufe dependency); autodoc then imports gufe cleanly. Finally, render AlchemicalArchive as an inline literal in getting_started.rst rather than an intersphinx cross-reference: the gufe inventory is pinned to v1.2.0, which predates gufe.archival, so the reference cannot resolve. Verified: `sphinx -W` builds with zero warnings in both the docs.yml env and the full runtime env. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Added more checks for metadata validation --------- Co-authored-by: David L. Dotson <dotsdl@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Mike Henry <11765982+mikemhenry@users.noreply.github.com>
1 parent 829532e commit e04181e

9 files changed

Lines changed: 398 additions & 16 deletions

File tree

CLAUDE.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ The system has four main services communicating through Neo4j and S3:
7575
- **Frozen Pydantic models**: Core models use `frozen=True` for immutability.
7676
- **Module docstrings**: Follow the pattern `:mod:\`alchemiscale.module_name\` --- description`.
7777
- **GUFE integration**: Core chemistry types (`AlchemicalNetwork`, `Transformation`, `ChemicalSystem`, `Protocol`) come from the `gufe` library and are stored/retrieved via their tokenization system (`GufeKey`, `GufeTokenizable`).
78+
- **`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.
7879

7980
## Testing
8081

alchemiscale/interface/client.py

Lines changed: 185 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,16 +7,18 @@
77
from __future__ import annotations
88

99
import asyncio
10+
import json
1011
from enum import StrEnum
1112
from typing import Any, Literal
12-
from collections.abc import Iterable
13+
from collections.abc import Iterable, Mapping, Sequence
1314
from itertools import chain
1415
from functools import lru_cache
1516

1617
from async_lru import alru_cache
1718
import networkx as nx
1819
from gufe import AlchemicalNetwork, Transformation, ChemicalSystem
19-
from gufe.tokenization import GufeTokenizable, KeyedChain
20+
from gufe.archival import AlchemicalArchive
21+
from gufe.tokenization import GufeTokenizable, KeyedChain, JSON_HANDLER
2022
from gufe.protocols import ProtocolResult, ProtocolDAGResult
2123
import zstandard as zstd
2224

@@ -34,6 +36,7 @@
3436
StrategyState,
3537
)
3638
from stratocaster.base import Strategy
39+
from ..utils import pdr_from_bytes
3740
from ..validators import validate_network_nonself
3841

3942
from warnings import warn
@@ -1536,12 +1539,7 @@ async def _async_get_protocoldagresult(
15361539
pdr_bytes,
15371540
)
15381541

1539-
try:
1540-
# Attempt to decompress the ProtocolDAGResult object
1541-
pdr = decompress_gufe_zstd(pdr_bytes)
1542-
except zstd.ZstdError:
1543-
# If decompress fails, assume it's a UTF-8 encoded JSON string
1544-
pdr = json_to_gufe(pdr_bytes.decode("utf-8"))
1542+
pdr = pdr_from_bytes(pdr_bytes)
15451543

15461544
return pdr
15471545

@@ -1767,6 +1765,185 @@ def get_network_failures(
17671765
network=network, ok=False, compress=compress, visualize=visualize
17681766
)
17691767

1768+
def get_network_archives(
1769+
self,
1770+
networks: list[ScopedKey | str],
1771+
metadata: list[dict | None] | None = None,
1772+
compress: bool = True,
1773+
visualize: bool = True,
1774+
) -> list[AlchemicalArchive | None]:
1775+
r"""Produce archival-quality extracts for the given ``AlchemicalNetwork``\s.
1776+
1777+
Each returned ``AlchemicalArchive`` bundles an
1778+
``AlchemicalNetwork`` together with all successful
1779+
``ProtocolDAGResult``\s currently available for its
1780+
``Transformation``\s, in a form suitable for long-term storage,
1781+
sharing, and downstream analysis.
1782+
1783+
Parameters
1784+
----------
1785+
networks
1786+
A list of ``AlchemicalNetwork`` ``ScopedKey`` values. The
1787+
list must not contain duplicate entries.
1788+
metadata
1789+
Metadata to attach to the produced ``AlchemicalArchive``
1790+
objects. This must be a list of dictionaries that are
1791+
compatible with ``GufeTokenizable`` serialization, in the
1792+
same order as ``networks``. A ``None`` entry in the list
1793+
attaches no metadata to the corresponding
1794+
``AlchemicalArchive``. Passing ``None`` in place of the
1795+
list is interpreted as a list of ``None``, which is the
1796+
default.
1797+
compress
1798+
If ``True``, compress objects server-side before shipping
1799+
them to the client. This is a performance optimization; it
1800+
has no bearing on the result of this method call.
1801+
visualize
1802+
If ``True``, show retrieval progress indicators.
1803+
1804+
Returns
1805+
-------
1806+
A list of ``AlchemicalArchive`` instances matching the order of
1807+
``networks``. If a network was not found, ``None`` is returned
1808+
in its place.
1809+
1810+
Raises
1811+
------
1812+
ValueError
1813+
If the provided metadata is not serializable, if the
1814+
lengths of the ``metadata`` and ``networks`` lists differ,
1815+
or if ``networks`` contains duplicate entries.
1816+
1817+
"""
1818+
network_sks = [
1819+
ScopedKey.from_str(network) if isinstance(network, str) else network
1820+
for network in networks
1821+
]
1822+
1823+
if len(set(network_sks)) != len(network_sks):
1824+
raise ValueError("`networks` list must not contain duplicate entries")
1825+
1826+
if metadata is None:
1827+
metadata = [None] * len(network_sks)
1828+
elif isinstance(metadata, Mapping) or not isinstance(metadata, Sequence):
1829+
raise ValueError(
1830+
"`metadata` must be a list/sequence of dictionaries or None"
1831+
)
1832+
1833+
if len(metadata) != len(network_sks):
1834+
raise ValueError("`metadata` and `networks` lists must be the same length")
1835+
1836+
# validate that all metadata is serializable up-front, before
1837+
# performing any (potentially expensive) retrieval
1838+
for network_sk, meta in zip(network_sks, metadata):
1839+
if meta is None:
1840+
continue
1841+
1842+
if not isinstance(meta, Mapping):
1843+
raise ValueError(
1844+
f"Metadata for '{network_sk}' must be a dictionary/mapping or None"
1845+
)
1846+
1847+
try:
1848+
json.dumps(meta, cls=JSON_HANDLER.encoder)
1849+
except (TypeError, ValueError) as e:
1850+
raise ValueError(
1851+
f"Unable to serialize metadata for '{network_sk}': {e}"
1852+
)
1853+
1854+
return [
1855+
self._get_network_archive(
1856+
network_sk, meta, compress=compress, visualize=visualize
1857+
)
1858+
for network_sk, meta in zip(network_sks, metadata)
1859+
]
1860+
1861+
def get_network_archive(
1862+
self,
1863+
network: ScopedKey | str,
1864+
metadata: dict | None = None,
1865+
compress: bool = True,
1866+
visualize: bool = True,
1867+
) -> AlchemicalArchive | None:
1868+
r"""Produce an archival-quality extract for a given ``AlchemicalNetwork``.
1869+
1870+
The returned ``AlchemicalArchive`` bundles the
1871+
``AlchemicalNetwork`` together with all successful
1872+
``ProtocolDAGResult``\s currently available for its
1873+
``Transformation``\s, in a form suitable for long-term storage,
1874+
sharing, and downstream analysis.
1875+
1876+
Parameters
1877+
----------
1878+
network
1879+
The ``ScopedKey`` of the ``AlchemicalNetwork`` to archive.
1880+
metadata
1881+
Metadata to attach to the produced ``AlchemicalArchive``.
1882+
This must be a dictionary that is compatible with
1883+
``GufeTokenizable`` serialization.
1884+
compress
1885+
If ``True``, compress objects server-side before shipping
1886+
them to the client. This is a performance optimization; it
1887+
has no bearing on the result of this method call.
1888+
visualize
1889+
If ``True``, show retrieval progress indicators.
1890+
1891+
Returns
1892+
-------
1893+
An ``AlchemicalArchive`` for the provided ``AlchemicalNetwork``.
1894+
If the network was not found, ``None`` is returned.
1895+
1896+
Raises
1897+
------
1898+
ValueError
1899+
If the provided metadata is not serializable.
1900+
1901+
"""
1902+
return self.get_network_archives(
1903+
[network],
1904+
metadata=None if metadata is None else [metadata],
1905+
compress=compress,
1906+
visualize=visualize,
1907+
)[0]
1908+
1909+
def _get_network_archive(
1910+
self,
1911+
network: ScopedKey,
1912+
metadata: dict | None,
1913+
compress: bool = True,
1914+
visualize: bool = True,
1915+
) -> AlchemicalArchive | None:
1916+
# returns None if the network does not exist in the given Scope
1917+
if not self.check_exists(network):
1918+
return None
1919+
1920+
an = self.get_network(network, compress=compress, visualize=visualize)
1921+
1922+
# retrieve all successful ProtocolDAGResults for the network's
1923+
# Transformations in parallel, keyed by their authoritative
1924+
# (server-side) Transformation ScopedKey
1925+
results = self.get_network_results(
1926+
network,
1927+
return_as=ResultFormat.PROTOCOL_DAG_RESULTS,
1928+
compress=compress,
1929+
visualize=visualize,
1930+
)
1931+
1932+
# pair each Transformation with its results by fetching the
1933+
# Transformation via its authoritative ScopedKey. We deliberately do
1934+
# not derive a ScopedKey from a deserialized Transformation's GufeKey:
1935+
# a `Transformation` deserialized now may not reproduce the GufeKey it
1936+
# had when ingested (gufe tokenization can change across versions), so
1937+
# its key is not a reliable join to the stored ScopedKey.
1938+
transformation_results = []
1939+
for transformation_sk, pdrs in results.items():
1940+
transformation = self.get_transformation(
1941+
transformation_sk, compress=compress, visualize=False
1942+
)
1943+
transformation_results.append((transformation, pdrs))
1944+
1945+
return AlchemicalArchive(an, transformation_results, metadata=metadata)
1946+
17701947
def get_transformation_results(
17711948
self,
17721949
transformation: ScopedKey,

alchemiscale/strategist/service.py

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
from ..compression import compress_keyed_chain_zstd, decompress_gufe_zstd, json_to_gufe
3838
from ..sleep import InterruptableSleep, SleepInterrupted
3939
from .settings import StrategistSettings
40+
from ..utils import pdr_from_bytes
4041

4142

4243
def execute_strategy_worker(
@@ -173,12 +174,7 @@ def _get_protocoldagresult_cached(
173174
location=result_ref.location, ok=result_ref.ok
174175
)
175176

176-
# Decompress the raw bytes to get ProtocolDAGResult
177-
try:
178-
pdr = decompress_gufe_zstd(pdr_bytes)
179-
except zstd.ZstdError:
180-
# Fallback to JSON deserialization for uncompressed data
181-
pdr = json_to_gufe(pdr_bytes.decode("utf-8"))
177+
pdr = pdr_from_bytes(pdr_bytes)
182178

183179
if self._cache_enabled:
184180
try:

0 commit comments

Comments
 (0)