diff --git a/flyemflows/volumes/tensorstore_volume_service.py b/flyemflows/volumes/tensorstore_volume_service.py index bd693c20..79990ec8 100644 --- a/flyemflows/volumes/tensorstore_volume_service.py +++ b/flyemflows/volumes/tensorstore_volume_service.py @@ -1,4 +1,5 @@ import copy +import pickle import numpy as np from confiddler import validate, emit_defaults @@ -220,6 +221,31 @@ def store(self, scale): self._stores[scale] = store return store + def __getstate__(self): + """ + Pickle representation. + """ + d = self.__dict__.copy() + + # Apparently unpickling can fail intermittently due to network issues, so we + # pickle/unpickle it explicitly so we can control the process and retry if necessary. + d['_stores'] = pickle.dumps(d['_stores']) + return d + + @auto_retry(3, pause_between_tries=5.0, logging_name=__name__) + def __setstate__(self, d): + # We unpickle the TensorStore(s) explicitly here, to ensure + # that errors are caught by the auto_retry decorator. + # + # Here's an example error we want to catch: + # + # ValueError: Error opening "neuroglancer_precomputed" driver: + # Error reading "gs://vnc-v3-seg/rc4_wsexp/info": + # CURL error[16] Error in the HTTP2 framing layer + # + d['_stores'] = pickle.loads(d['_stores']) + self.__dict__ = d + @property def dtype(self): return self._dtype