From 334101536d7d1a54e4ab236006a03a851822ef83 Mon Sep 17 00:00:00 2001 From: Stuart Berg Date: Fri, 4 Sep 2020 00:03:24 -0400 Subject: [PATCH] TensorStoreVolumeService: Use @auto_retry to catch issues during TensorStore unpickling --- .../volumes/tensorstore_volume_service.py | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) 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