From 69304c5763cc66343d7c15aa74eafa2cd7dd7cc8 Mon Sep 17 00:00:00 2001 From: Wouter Zwerink <35296208+wouterzwerink@users.noreply.github.com> Date: Mon, 9 Dec 2024 21:53:16 +0100 Subject: [PATCH] Prevent _s3_client from being serialized (#847) Co-authored-by: Saaketh Narayan Co-authored-by: Saaketh Narayan --- streaming/base/storage/download.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/streaming/base/storage/download.py b/streaming/base/storage/download.py index d92df6a1e..b2dff253d 100644 --- a/streaming/base/storage/download.py +++ b/streaming/base/storage/download.py @@ -261,6 +261,15 @@ def _create_s3_client(self, unsigned: bool = False, timeout: float = DEFAULT_TIM config=config, endpoint_url=os.environ.get('S3_ENDPOINT_URL')) + def __getstate__(self) -> dict: + state = self.__dict__.copy() + state['_s3_client'] = None # Exclude _s3_client from being pickled + return state + + def __setstate__(self, state: dict): + self.__dict__.update(state) + self._s3_client = None # Ensure _s3_client is reset after unpickling + class SFTPDownloader(CloudDownloader): """Download files from SFTP to local filesystem."""