Attempt to decompress object store objects in a try block

If a decompression error is raised, assume that the original data was
never compressed.

alchemiscale/compute/client.py

@@ -115,20 +115,25 @@ def get_task_transformation(self, task: ScopedKey) -> ScopedKey:
 
     def retrieve_task_transformation(
         self, task: ScopedKey
-    ) -> Tuple[Transformation, Optional[ProtocolDAGResult]]:
+    ) -> tuple[Transformation, ProtocolDAGResult | None]:
         transformation, protocoldagresult = self._get_resource(
             f"/tasks/{task}/transformation/gufe"
         )
 
         if protocoldagresult is not None:
-            protocoldagresult = decompress_gufe_zstd(
-                protocoldagresult.encode("latin-1")
-            )
 
-        return (
-            json_to_gufe(transformation),
-            protocoldagresult,
-        )
+            protocoldagresult_bytes = protocoldagresult.encode("latin-1")
+
+            try:
+                # Attempt to decompress the ProtocolDAGResult object
+                protocoldagresult = decompress_gufe_zstd(
+                    protocoldagresult_bytes
+                )
+            except zstd.ZstdError:
+                # If decompression fails, assume it's a UTF-8 encoded JSON string
+                protocoldagresult = json_to_gufe(protocoldagresult_bytes.decode('utf-8'))
+
+        return json_to_gufe(transformation), protocoldagresult
 
     def set_task_result(
         self,

alchemiscale/interface/client.py

@@ -15,11 +15,13 @@
 from gufe import AlchemicalNetwork, Transformation, ChemicalSystem
 from gufe.tokenization import GufeTokenizable, JSON_HANDLER, KeyedChain
 from gufe.protocols import ProtocolResult, ProtocolDAGResult
+import zstandard as zstd
 
 
 from ..base.client import (
     AlchemiscaleBaseClient,
     AlchemiscaleBaseClientError,
+    json_to_gufe,
     use_session,
 )
 from ..compression import decompress_gufe_zstd
@@ -1353,11 +1355,18 @@ def get_tasks_priority(
     async def _async_get_protocoldagresult(
         self, protocoldagresultref, transformation, route, compress
     ):
-        pdr_compressed_latin1 = await self._get_resource_async(
+        pdr_latin1_decoded = await self._get_resource_async(
             f"/transformations/{transformation}/{route}/{protocoldagresultref}",
             compress=compress,
         )
-        pdr = decompress_gufe_zstd(pdr_compressed_latin1[0].encode("latin-1"))
+
+        try:
+            # Attempt to decompress the ProtocolDAGResult object
+            pdr_bytes = pdr_latin1_decoded[0].encode("latin-1")
+            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'))
 
         return pdr
 

alchemiscale/storage/objectstore.py

@@ -228,7 +228,7 @@ def push_protocoldagresult(
             transformation.gufe_key,
             route,
             pdr.key,
-            "obj.json",
+            "obj",
         )
 
         response = self._store_bytes(location, protocoldagresult)
@@ -289,7 +289,7 @@ def pull_protocoldagresult(
                 transformation.gufe_key,
                 route,
                 protocoldagresult.gufe_key,
-                "obj.json",
+                "obj",
             )
 
         ## TODO: want organization alongside `obj.json` of `ProtocolUnit` gufe_keys

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

@@ -1,10 +1,12 @@
 import pytest
 from time import sleep
+import os
 from pathlib import Path
 from itertools import chain
+import json
 
 from gufe import AlchemicalNetwork
-from gufe.tokenization import TOKENIZABLE_REGISTRY, GufeKey
+from gufe.tokenization import TOKENIZABLE_REGISTRY, GufeKey, JSON_HANDLER
 from gufe.protocols.protocoldag import execute_DAG
 import networkx as nx
 
@@ -1860,6 +1862,66 @@ def _execute_tasks(tasks, n4js, s3os_server):
 
         return protocoldagresults
 
+    def test_get_transformation_and_network_results_json(
+            self,
+            scope_test,
+            n4js_preloaded,
+            s3os_server,
+            user_client: client.AlchemiscaleClient,
+            network_tyk2,
+            tmpdir,
+    ):
+        n4js = n4js_preloaded
+
+        # select the transformation we want to compute
+        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)
+
+        # user client : create three independent tasks for the transformation
+        user_client.create_tasks(transformation_sk, count=3)
+
+        # user client : action the tasks for execution
+        all_tasks = user_client.get_transformation_tasks(transformation_sk)
+        actioned_tasks = user_client.action_tasks(all_tasks, network_sk)
+
+        # execute the actioned tasks and push results directly using statestore and object store
+        with tmpdir.as_cwd():
+            protocoldagresults = self._execute_tasks(actioned_tasks, n4js, s3os_server)
+            # overwrite what's in the object store
+            for protocoldagresult in protocoldagresults:
+                pdr_jb = json.dumps(
+                    protocoldagresult.to_dict(), cls=JSON_HANDLER.encoder
+                ).encode("utf-8")
+
+                location = os.path.join(
+                    "protocoldagresult",
+                    *transformation_sk.scope.to_tuple(),
+                    transformation_sk.gufe_key,
+                    "results",
+                    protocoldagresult.key,
+                    "obj",
+                )
+
+                s3os_server._store_bytes(location, pdr_jb)
+
+        # clear local gufe registry of pdr objects
+        # not critical, but ensures we see the objects that are deserialized
+        # instead of our instances already in memory post-pull
+        for pdr in protocoldagresults:
+            TOKENIZABLE_REGISTRY.pop(pdr.key, None)
+
+        # get back protocoldagresults instead
+        protocoldagresults_r = user_client.get_transformation_results(
+            transformation_sk, return_protocoldagresults=True
+        )
+
+        assert set(protocoldagresults_r) == set(protocoldagresults)
+
+        pass
+
     def test_get_transformation_and_network_results(
         self,
         scope_test,