From 58f23d21613b30abdef224748f76908974c0c332 Mon Sep 17 00:00:00 2001 From: Gary Pavlis Date: Mon, 21 Sep 2026 07:54:25 -0400 Subject: [PATCH 1/5] Add a generic module for handling s3 data access in mspass --- python/mspasspy/io/s3_worker_plugin.py | 300 +++++++++++++++++++++++++ 1 file changed, 300 insertions(+) create mode 100644 python/mspasspy/io/s3_worker_plugin.py diff --git a/python/mspasspy/io/s3_worker_plugin.py b/python/mspasspy/io/s3_worker_plugin.py new file mode 100644 index 0000000000..a2d1ff9dd4 --- /dev/null +++ b/python/mspasspy/io/s3_worker_plugin.py @@ -0,0 +1,300 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Prototype worker plugin for using dask on AWS to allow workers to +not have to instantiate an s3 client on each submit. Modeled after +the mongodb worker plugin. + +Specific to the supplied EarthScope/GeoLab workflow. Live deployment validation +is still required; the tests exercise credential lifecycle without AWS access. + +Created on Mon Mar 23 09:26:00 2026 + +@author: pavlis +""" + +import boto3 +from contextlib import ExitStack +from botocore.config import Config +from botocore import UNSIGNED +from functools import lru_cache +from earthscope_sdk import EarthScopeClient +from dask.distributed import WorkerPlugin, get_worker + +# this decorator is necessary to avoid the cost of instantiating a +# client on serial tasks every time it is called +@lru_cache(maxsize=1) +def fetch_s3_client(parallel=True, + session=None, + session_config=None, + worker_data_key="geolab_s3client"): + """ + Generic tool to fetch s3 client access S3 data. + + This function provides a standard tool to get an s3 client efficiently + in either a parallel (default) or serial operation. It behaves + totally differently depending on the value or the "parallel" argument. + + Case 1: parallel == True (default) + In this case the function assumes an instance of a versions of + the S3Worker class (defined in this same module) has been pushed to + all workers with the "register_plugin" method of the dask.distributed.Client. + Different clients with different configurations can be loaded onto + workers with different values of a key. To have this function + access an alternative to the default client run the functionw with + an appropriate value for `worker_data_key` (see below for supported options). + + Case 2: parallel==False + Set the parallel argument False if you are using this function in + a serial task. If you are running serial and you do not set parallel + False function will throw a KeyError exception. For efficiency + this function is pushed to a global cache using the `functools.lru_cache` + decorator. That improves performance as otherwise every time the function + is called a new client would have to instantiated from a Session object. + Not as bad as starting from nothing but not zero either. The cache + approach makes subsequent access times tiny. + + :param parallel: boolean that when set True (default) assumes a client + has already been pushed to dask workers with register_plugin. In that + mode it will try to fetch that client from the worker with a key + defined by the worker_data_key argument. When False the first + time the function is called a client is created but subsequent + calls fetch the client cached in memory. + :param session: boto3 Session object that should be used to + instiate a client. This should only be used for initializing + a client in serial mode. It will be silently ignored if + parallel is True. Default is None which causes a default + Session object to be created for the cached serial client. + Pass a valid value if need to use an alternative. + :param session_config: boto3 Config object to specify alternative + configurations to default. Use this only if you know what you are doing. + It is used only for serial mode and will be silently ignored if used + with parallel set True. + :param worker_data_key: string key used only when parallel is set True. + This parameter defines the key used to access the appropriate + instance of s3 client on each worker. Multiple keys can be used to + allow using clients with different configurations for different access + credentials. There are currently three distinct plugin implementations + with different default key values. + + 1. GeoLabS3Worker uses the key "geolab_s3client" by default. + 2. AnonymousS3Worker uses the key "anonymous_s3client" by default + 3. StockS3Worker uses the key "stock_s3client" as default. + + Those keys, however, are not frozen. When any of the worker plugins + are instatiated you can set the value for that key with the + common argument "key=". If you do, the driver code will need to + change worker_data_key appropriately as well. You will need to + set worker_data_key running parallel on anywhere but on geolab. + """ + if parallel: + try: + worker = get_worker() + except Exception as e: + raise ValueError( + "fetch_s3_client: this function must be " + "executed within a Dask worker context so that get_worker() succeeds and an " + "s3 client is available via a worker plugin." + ) from e + try: + s3_client = worker.data[worker_data_key] + + except KeyError as e: + message = "fetch_s3_client: dask worker has no S3 plugin registered with name={}\n".format( + worker_data_key + ) + message += "Register S3Worker before submitting tasks to the dask cluster" + raise ValueError(message) from e + else: + if session is None: + # This should work as this would be the use for a stock AWS + # login. There the credentials are supposed to be cached so this + # should work unless proven othewise + session = boto3.Session() + if session_config is None: + # default config frozen for now + # may need a more flexible way to do this than hard code this + s3_client = session.client( + "s3", + config=Config( + request_checksum_calculation="when_required", + response_checksum_validation="when_required", + ), + ) + else: + s3_client = session.client("s3", config=session_config) + + return s3_client + + +class GeoLabS3Worker(WorkerPlugin): + """ + Dask worker plugin to create a worker resident client on each dask worker. + This instance works only with earthscope_sdk used on GeoLab. + + An S3 client requires significant time to construct for a variety of reasons. + It is known to be a really serious bottleneck if s3 clients are instantiated + inside a parallel workflow to access s3. This class can be used to + create what is called a worker plugin for dask. The standard use in a + python script is: + ``` + dask_client = mspass_client.get_scheduler() + s3plugin = GeoLabS3Worker() + dask_client.register_plugin(s3plugin) + ``` + That creates and loads a memory resident s3 client in each worker's memory space. + Processing functions that need access to s3 should then use the function + `fetch_s3_client` to get a reference to the memory resident s3 client. + That produces near zero overhead in a worker process particularly compared to + the time needed to instantiate a new instance of an s3 client. + + Note this plugin differs from it's siblings in this module as it caches the + s3 client as an attribute of this class. That is a necessary evil to deal with + timeout issues in GeoLab that complicate this code. Since this plugin is + only of use on GeoLab this should not matter as one would not expect multiple + s3 clients with different credentialsl run on GeoLab simultaneously that + would require distinguishing the multiple instances. + """ + + def __init__(self, key="geolab_s3client"): + """ + Standard constructor. + + Normal use requires no arguments. + """ + # probably unnecessary initializations but makes clear these are two attributes + # managed by this class + self.worker_key = key + self.s3_client = None + self._cleanup = None + + + def setup(self, worker): + """ + Required method for a worker plugin. This acts like a secondary constructor. + It is involked on each worker when this object is pushed to + each worker with the register_plugin method of the dask client. + """ + # Gemini's comment on why this complication is needed here: + # Keep the SDK's refresh provider and real expiration intact. A frozen + # credential snapshot has no expiration, and a new SDK session may reuse + # cached credentials; inventing a new expiry would extend their lifetime + # locally without extending their actual AWS validity. + with ExitStack() as cleanup: + esclient = EarthScopeClient() + cleanup.callback(esclient.close) + session = esclient.user.get_boto3_session() + # use parallel=False because this method is executed on each + # worker and to that python instance it is a serial thing + # confusing though as this always done in a parallel context + s3_client = fetch_s3_client(parallel=False,session=session) + # this is a callback to ExitStack which means it is executed only when the + # with block exits + cleanup.callback(s3_client.close) + # Gemini suggests we should use the construct commented out for reason here. + # A client is a worker resource, not task data: keeping it in + # worker.data could ask Dask to pickle/spill its credentials. + #self.s3_client = s3_client + # Testing showed that did work running in isolation on GeoLab but I (glp) don't think + # what gemini suggests would work if multiple worker clients were defined + # this use makes this consistent with other instances of the WorkerPlugin for + # using s3 + worker.data[self.worker_key] = s3_client + self._cleanup = cleanup.pop_all() + + def teardown(self, worker): + """ + Required method for a worker plugin. This method is effectively a destructor + called when the class goes out of scope. It is essential in this case to + avoid a resource leak as it properly closes the clients connections to s3. + """ + s3_client = worker.data.get(self.worker_key) + s3_client.close() + cleanup = getattr(self, "_cleanup", None) + if cleanup is not None: + self._cleanup = None + #self.s3_client = None + cleanup.close() + +class AnonymousS3Client(WorkerPlugin): + """ + Dask worker plugin to create a worker resident client on each dask worker. + This instance works only for anonymous access. + + An S3 client requires significant time to construct for a variety of reasons. + It is known to be a really serious bottleneck if s3 clients are instantiated + inside a parallel workflow to access s3. This class can be used to + create what is called a worker plugin for dask. The standard use in a + python script is: + ``` + dask_client = mspass_client.get_scheduler() + s3plugin = GeoLabS3Worker() + dask_client.register_plugin(s3plugin) + ``` + That creates and loads a memory resident s3 client in each worker's memory space. + Processing functions that need access to s3 should then use the function + `fetch_s3_client` to get a reference to the memory resident s3 client. + That produces near zero overhead in a worker process particularly compared to + the time needed to instantiate a new instance of an s3 client. + + Note the default key the fetch_s3_client function needs to use + this client is "anonymous_s3client". It can be changed if desired but + you will need to change calls to fetch_s3_client appropriately if you do. + """ + def __init__(self,key="anonymous_s3client",region="us-east-2"): + self.worker_key = key + self.region = region + def setup(self,worker): + s3_client = boto3.client( + 's3', + region_name=self.region, + config=Config(signature_version=UNSIGNED) + ) + + worker.data[self.worker_key] = s3_client + def teardown(self,worker): + s3_client = worker.data.get(self.worker_key) + s3_client.close() + +class StockS3Client(WorkerPlugin): + """ + Dask worker plugin to create a worker resident client on each dask worker. + This instance should be used for conventional AWS login accounts where + the user has already valid credentials cached in their run environment. + + An S3 client requires significant time to construct for a variety of reasons. + It is known to be a really serious bottleneck if s3 clients are instantiated + inside a parallel workflow to access s3. This class can be used to + create what is called a worker plugin for dask. The standard use in a + python script is: + ``` + dask_client = mspass_client.get_scheduler() + s3plugin = GeoLabS3Worker() + dask_client.register_plugin(s3plugin) + ``` + That creates and loads a memory resident s3 client in each worker's memory space. + Processing functions that need access to s3 should then use the function + `fetch_s3_client` to get a reference to the memory resident s3 client. + That produces near zero overhead in a worker process particularly compared to + the time needed to instantiate a new instance of an s3 client. + + Note the default key the fetch_s3_client function needs to use + this client is "stock_s3client". It can be changed if desired but + you will need to change calls to fetch_s3_client appropriately if you do. + """ + def __init__(self,key="stock_s3client",region="us-east-2"): + self.worker_key = key + self.region = region + def setup(self,worker): + # this simple construct assumes credentials are set using + # the default AWS credentials chain. In that case the + # access keys are cached and this should work according to Gemini + s3_client = boto3.client( + 's3', + region_name=self.region, + ) + + worker.data[self.worker_key] = s3_client + def teardown(self,worker): + s3_client = worker.data.get(self.worker_key) + s3_client.close() \ No newline at end of file From 3697334aff31618315be73eb72187cf555001254 Mon Sep 17 00:00:00 2001 From: pavlis <184871+pavlis@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:09:05 +0000 Subject: [PATCH 2/5] Apply Black 25.1.0 formatting --- python/mspasspy/io/s3_worker_plugin.py | 187 +++++++++++++------------ 1 file changed, 96 insertions(+), 91 deletions(-) diff --git a/python/mspasspy/io/s3_worker_plugin.py b/python/mspasspy/io/s3_worker_plugin.py index a2d1ff9dd4..6a9f0fca5c 100644 --- a/python/mspasspy/io/s3_worker_plugin.py +++ b/python/mspasspy/io/s3_worker_plugin.py @@ -21,71 +21,71 @@ from earthscope_sdk import EarthScopeClient from dask.distributed import WorkerPlugin, get_worker -# this decorator is necessary to avoid the cost of instantiating a + +# this decorator is necessary to avoid the cost of instantiating a # client on serial tasks every time it is called @lru_cache(maxsize=1) -def fetch_s3_client(parallel=True, - session=None, - session_config=None, - worker_data_key="geolab_s3client"): +def fetch_s3_client( + parallel=True, session=None, session_config=None, worker_data_key="geolab_s3client" +): """ Generic tool to fetch s3 client access S3 data. - - This function provides a standard tool to get an s3 client efficiently - in either a parallel (default) or serial operation. It behaves - totally differently depending on the value or the "parallel" argument. - + + This function provides a standard tool to get an s3 client efficiently + in either a parallel (default) or serial operation. It behaves + totally differently depending on the value or the "parallel" argument. + Case 1: parallel == True (default) - In this case the function assumes an instance of a versions of - the S3Worker class (defined in this same module) has been pushed to + In this case the function assumes an instance of a versions of + the S3Worker class (defined in this same module) has been pushed to all workers with the "register_plugin" method of the dask.distributed.Client. - Different clients with different configurations can be loaded onto - workers with different values of a key. To have this function - access an alternative to the default client run the functionw with + Different clients with different configurations can be loaded onto + workers with different values of a key. To have this function + access an alternative to the default client run the functionw with an appropriate value for `worker_data_key` (see below for supported options). - + Case 2: parallel==False - Set the parallel argument False if you are using this function in + Set the parallel argument False if you are using this function in a serial task. If you are running serial and you do not set parallel - False function will throw a KeyError exception. For efficiency - this function is pushed to a global cache using the `functools.lru_cache` - decorator. That improves performance as otherwise every time the function - is called a new client would have to instantiated from a Session object. - Not as bad as starting from nothing but not zero either. The cache - approach makes subsequent access times tiny. - - :param parallel: boolean that when set True (default) assumes a client - has already been pushed to dask workers with register_plugin. In that - mode it will try to fetch that client from the worker with a key - defined by the worker_data_key argument. When False the first - time the function is called a client is created but subsequent + False function will throw a KeyError exception. For efficiency + this function is pushed to a global cache using the `functools.lru_cache` + decorator. That improves performance as otherwise every time the function + is called a new client would have to instantiated from a Session object. + Not as bad as starting from nothing but not zero either. The cache + approach makes subsequent access times tiny. + + :param parallel: boolean that when set True (default) assumes a client + has already been pushed to dask workers with register_plugin. In that + mode it will try to fetch that client from the worker with a key + defined by the worker_data_key argument. When False the first + time the function is called a client is created but subsequent calls fetch the client cached in memory. - :param session: boto3 Session object that should be used to - instiate a client. This should only be used for initializing - a client in serial mode. It will be silently ignored if - parallel is True. Default is None which causes a default - Session object to be created for the cached serial client. + :param session: boto3 Session object that should be used to + instiate a client. This should only be used for initializing + a client in serial mode. It will be silently ignored if + parallel is True. Default is None which causes a default + Session object to be created for the cached serial client. Pass a valid value if need to use an alternative. - :param session_config: boto3 Config object to specify alternative + :param session_config: boto3 Config object to specify alternative configurations to default. Use this only if you know what you are doing. - It is used only for serial mode and will be silently ignored if used + It is used only for serial mode and will be silently ignored if used with parallel set True. - :param worker_data_key: string key used only when parallel is set True. - This parameter defines the key used to access the appropriate - instance of s3 client on each worker. Multiple keys can be used to + :param worker_data_key: string key used only when parallel is set True. + This parameter defines the key used to access the appropriate + instance of s3 client on each worker. Multiple keys can be used to allow using clients with different configurations for different access - credentials. There are currently three distinct plugin implementations - with different default key values. - + credentials. There are currently three distinct plugin implementations + with different default key values. + 1. GeoLabS3Worker uses the key "geolab_s3client" by default. 2. AnonymousS3Worker uses the key "anonymous_s3client" by default - 3. StockS3Worker uses the key "stock_s3client" as default. - - Those keys, however, are not frozen. When any of the worker plugins - are instatiated you can set the value for that key with the - common argument "key=". If you do, the driver code will need to - change worker_data_key appropriately as well. You will need to - set worker_data_key running parallel on anywhere but on geolab. + 3. StockS3Worker uses the key "stock_s3client" as default. + + Those keys, however, are not frozen. When any of the worker plugins + are instatiated you can set the value for that key with the + common argument "key=". If you do, the driver code will need to + change worker_data_key appropriately as well. You will need to + set worker_data_key running parallel on anywhere but on geolab. """ if parallel: try: @@ -108,7 +108,7 @@ def fetch_s3_client(parallel=True, else: if session is None: # This should work as this would be the use for a stock AWS - # login. There the credentials are supposed to be cached so this + # login. There the credentials are supposed to be cached so this # should work unless proven othewise session = boto3.Session() if session_config is None: @@ -119,8 +119,8 @@ def fetch_s3_client(parallel=True, config=Config( request_checksum_calculation="when_required", response_checksum_validation="when_required", - ), - ) + ), + ) else: s3_client = session.client("s3", config=session_config) @@ -147,20 +147,20 @@ class GeoLabS3Worker(WorkerPlugin): `fetch_s3_client` to get a reference to the memory resident s3 client. That produces near zero overhead in a worker process particularly compared to the time needed to instantiate a new instance of an s3 client. - - Note this plugin differs from it's siblings in this module as it caches the - s3 client as an attribute of this class. That is a necessary evil to deal with - timeout issues in GeoLab that complicate this code. Since this plugin is - only of use on GeoLab this should not matter as one would not expect multiple - s3 clients with different credentialsl run on GeoLab simultaneously that - would require distinguishing the multiple instances. + + Note this plugin differs from it's siblings in this module as it caches the + s3 client as an attribute of this class. That is a necessary evil to deal with + timeout issues in GeoLab that complicate this code. Since this plugin is + only of use on GeoLab this should not matter as one would not expect multiple + s3 clients with different credentialsl run on GeoLab simultaneously that + would require distinguishing the multiple instances. """ def __init__(self, key="geolab_s3client"): """ Standard constructor. - Normal use requires no arguments. + Normal use requires no arguments. """ # probably unnecessary initializations but makes clear these are two attributes # managed by this class @@ -168,7 +168,6 @@ def __init__(self, key="geolab_s3client"): self.s3_client = None self._cleanup = None - def setup(self, worker): """ Required method for a worker plugin. This acts like a secondary constructor. @@ -184,19 +183,19 @@ def setup(self, worker): esclient = EarthScopeClient() cleanup.callback(esclient.close) session = esclient.user.get_boto3_session() - # use parallel=False because this method is executed on each + # use parallel=False because this method is executed on each # worker and to that python instance it is a serial thing # confusing though as this always done in a parallel context - s3_client = fetch_s3_client(parallel=False,session=session) - # this is a callback to ExitStack which means it is executed only when the + s3_client = fetch_s3_client(parallel=False, session=session) + # this is a callback to ExitStack which means it is executed only when the # with block exits cleanup.callback(s3_client.close) # Gemini suggests we should use the construct commented out for reason here. # A client is a worker resource, not task data: keeping it in # worker.data could ask Dask to pickle/spill its credentials. - #self.s3_client = s3_client - # Testing showed that did work running in isolation on GeoLab but I (glp) don't think - # what gemini suggests would work if multiple worker clients were defined + # self.s3_client = s3_client + # Testing showed that did work running in isolation on GeoLab but I (glp) don't think + # what gemini suggests would work if multiple worker clients were defined # this use makes this consistent with other instances of the WorkerPlugin for # using s3 worker.data[self.worker_key] = s3_client @@ -213,13 +212,14 @@ def teardown(self, worker): cleanup = getattr(self, "_cleanup", None) if cleanup is not None: self._cleanup = None - #self.s3_client = None + # self.s3_client = None cleanup.close() + class AnonymousS3Client(WorkerPlugin): """ Dask worker plugin to create a worker resident client on each dask worker. - This instance works only for anonymous access. + This instance works only for anonymous access. An S3 client requires significant time to construct for a variety of reasons. It is known to be a really serious bottleneck if s3 clients are instantiated @@ -236,31 +236,33 @@ class AnonymousS3Client(WorkerPlugin): `fetch_s3_client` to get a reference to the memory resident s3 client. That produces near zero overhead in a worker process particularly compared to the time needed to instantiate a new instance of an s3 client. - - Note the default key the fetch_s3_client function needs to use + + Note the default key the fetch_s3_client function needs to use this client is "anonymous_s3client". It can be changed if desired but you will need to change calls to fetch_s3_client appropriately if you do. """ - def __init__(self,key="anonymous_s3client",region="us-east-2"): + + def __init__(self, key="anonymous_s3client", region="us-east-2"): self.worker_key = key self.region = region - def setup(self,worker): + + def setup(self, worker): s3_client = boto3.client( - 's3', - region_name=self.region, - config=Config(signature_version=UNSIGNED) - ) + "s3", region_name=self.region, config=Config(signature_version=UNSIGNED) + ) worker.data[self.worker_key] = s3_client - def teardown(self,worker): + + def teardown(self, worker): s3_client = worker.data.get(self.worker_key) s3_client.close() - + + class StockS3Client(WorkerPlugin): """ Dask worker plugin to create a worker resident client on each dask worker. - This instance should be used for conventional AWS login accounts where - the user has already valid credentials cached in their run environment. + This instance should be used for conventional AWS login accounts where + the user has already valid credentials cached in their run environment. An S3 client requires significant time to construct for a variety of reasons. It is known to be a really serious bottleneck if s3 clients are instantiated @@ -277,24 +279,27 @@ class StockS3Client(WorkerPlugin): `fetch_s3_client` to get a reference to the memory resident s3 client. That produces near zero overhead in a worker process particularly compared to the time needed to instantiate a new instance of an s3 client. - - Note the default key the fetch_s3_client function needs to use + + Note the default key the fetch_s3_client function needs to use this client is "stock_s3client". It can be changed if desired but you will need to change calls to fetch_s3_client appropriately if you do. """ - def __init__(self,key="stock_s3client",region="us-east-2"): + + def __init__(self, key="stock_s3client", region="us-east-2"): self.worker_key = key self.region = region - def setup(self,worker): + + def setup(self, worker): # this simple construct assumes credentials are set using - # the default AWS credentials chain. In that case the + # the default AWS credentials chain. In that case the # access keys are cached and this should work according to Gemini s3_client = boto3.client( - 's3', - region_name=self.region, - ) + "s3", + region_name=self.region, + ) worker.data[self.worker_key] = s3_client - def teardown(self,worker): + + def teardown(self, worker): s3_client = worker.data.get(self.worker_key) - s3_client.close() \ No newline at end of file + s3_client.close() From 9799cb8e461d61f57c68a63c7bad13bf60f55d83 Mon Sep 17 00:00:00 2001 From: wangyinz Date: Sat, 26 Sep 2026 01:19:26 -0500 Subject: [PATCH 3/5] test: cover S3 worker lifecycle and fix client ownership (#1039) --- python/mspasspy/io/s3_worker_plugin.py | 117 +++++----- python/tests/io/test_s3_worker_plugin.py | 264 +++++++++++++++++++++++ 2 files changed, 330 insertions(+), 51 deletions(-) create mode 100644 python/tests/io/test_s3_worker_plugin.py diff --git a/python/mspasspy/io/s3_worker_plugin.py b/python/mspasspy/io/s3_worker_plugin.py index 6a9f0fca5c..c38ce9f187 100644 --- a/python/mspasspy/io/s3_worker_plugin.py +++ b/python/mspasspy/io/s3_worker_plugin.py @@ -22,9 +22,6 @@ from dask.distributed import WorkerPlugin, get_worker -# this decorator is necessary to avoid the cost of instantiating a -# client on serial tasks every time it is called -@lru_cache(maxsize=1) def fetch_s3_client( parallel=True, session=None, session_config=None, worker_data_key="geolab_s3client" ): @@ -47,10 +44,9 @@ def fetch_s3_client( Case 2: parallel==False Set the parallel argument False if you are using this function in a serial task. If you are running serial and you do not set parallel - False function will throw a KeyError exception. For efficiency - this function is pushed to a global cache using the `functools.lru_cache` - decorator. That improves performance as otherwise every time the function - is called a new client would have to instantiated from a Session object. + False the function will raise a ValueError exception. For efficiency + serial client creation is cached using `functools.lru_cache`. Otherwise + every call would instantiate a new client from a Session object. Not as bad as starting from nothing but not zero either. The cache approach makes subsequent access times tiny. @@ -78,8 +74,8 @@ def fetch_s3_client( with different default key values. 1. GeoLabS3Worker uses the key "geolab_s3client" by default. - 2. AnonymousS3Worker uses the key "anonymous_s3client" by default - 3. StockS3Worker uses the key "stock_s3client" as default. + 2. AnonymousS3Client uses the key "anonymous_s3client" by default + 3. StockS3Client uses the key "stock_s3client" as default. Those keys, however, are not frozen. When any of the worker plugins are instatiated you can set the value for that key with the @@ -106,27 +102,33 @@ def fetch_s3_client( message += "Register S3Worker before submitting tasks to the dask cluster" raise ValueError(message) from e else: - if session is None: - # This should work as this would be the use for a stock AWS - # login. There the credentials are supposed to be cached so this - # should work unless proven othewise - session = boto3.Session() - if session_config is None: - # default config frozen for now - # may need a more flexible way to do this than hard code this - s3_client = session.client( - "s3", - config=Config( - request_checksum_calculation="when_required", - response_checksum_validation="when_required", - ), - ) - else: - s3_client = session.client("s3", config=session_config) + s3_client = _serial_s3_client(session, session_config) return s3_client +def _create_s3_client(session=None, session_config=None): + if session is None: + session = boto3.Session() + if session_config is None: + session_config = Config( + request_checksum_calculation="when_required", + response_checksum_validation="when_required", + ) + return session.client("s3", config=session_config) + + +@lru_cache(maxsize=1) +def _serial_s3_client(session=None, session_config=None): + return _create_s3_client(session, session_config) + + +# Keep the public cache controls while caching only serial clients. Worker +# lookups must reflect the current worker and its current plugin registration. +fetch_s3_client.cache_clear = _serial_s3_client.cache_clear +fetch_s3_client.cache_info = _serial_s3_client.cache_info + + class GeoLabS3Worker(WorkerPlugin): """ Dask worker plugin to create a worker resident client on each dask worker. @@ -183,22 +185,15 @@ def setup(self, worker): esclient = EarthScopeClient() cleanup.callback(esclient.close) session = esclient.user.get_boto3_session() - # use parallel=False because this method is executed on each - # worker and to that python instance it is a serial thing - # confusing though as this always done in a parallel context - s3_client = fetch_s3_client(parallel=False, session=session) - # this is a callback to ExitStack which means it is executed only when the - # with block exits + # Plugin clients have their own lifetime; never reuse a serial + # cached client that another owner may already have closed. + s3_client = _create_s3_client(session=session) + # Retain cleanup after successful registration, or release resources + # immediately if setup fails. cleanup.callback(s3_client.close) - # Gemini suggests we should use the construct commented out for reason here. - # A client is a worker resource, not task data: keeping it in - # worker.data could ask Dask to pickle/spill its credentials. - # self.s3_client = s3_client - # Testing showed that did work running in isolation on GeoLab but I (glp) don't think - # what gemini suggests would work if multiple worker clients were defined - # this use makes this consistent with other instances of the WorkerPlugin for - # using s3 + # Publish under the configurable key, retaining ownership for teardown. worker.data[self.worker_key] = s3_client + self.s3_client = s3_client self._cleanup = cleanup.pop_all() def teardown(self, worker): @@ -207,12 +202,12 @@ def teardown(self, worker): called when the class goes out of scope. It is essential in this case to avoid a resource leak as it properly closes the clients connections to s3. """ - s3_client = worker.data.get(self.worker_key) - s3_client.close() - cleanup = getattr(self, "_cleanup", None) + if worker.data.get(self.worker_key) is self.s3_client: + worker.data.pop(self.worker_key, None) + cleanup = self._cleanup + self._cleanup = None + self.s3_client = None if cleanup is not None: - self._cleanup = None - # self.s3_client = None cleanup.close() @@ -245,17 +240,27 @@ class AnonymousS3Client(WorkerPlugin): def __init__(self, key="anonymous_s3client", region="us-east-2"): self.worker_key = key self.region = region + self.s3_client = None def setup(self, worker): s3_client = boto3.client( "s3", region_name=self.region, config=Config(signature_version=UNSIGNED) ) - worker.data[self.worker_key] = s3_client + try: + worker.data[self.worker_key] = s3_client + except Exception: + s3_client.close() + raise + self.s3_client = s3_client def teardown(self, worker): - s3_client = worker.data.get(self.worker_key) - s3_client.close() + s3_client = self.s3_client + self.s3_client = None + if s3_client is not None: + if worker.data.get(self.worker_key) is s3_client: + worker.data.pop(self.worker_key, None) + s3_client.close() class StockS3Client(WorkerPlugin): @@ -288,6 +293,7 @@ class StockS3Client(WorkerPlugin): def __init__(self, key="stock_s3client", region="us-east-2"): self.worker_key = key self.region = region + self.s3_client = None def setup(self, worker): # this simple construct assumes credentials are set using @@ -298,8 +304,17 @@ def setup(self, worker): region_name=self.region, ) - worker.data[self.worker_key] = s3_client + try: + worker.data[self.worker_key] = s3_client + except Exception: + s3_client.close() + raise + self.s3_client = s3_client def teardown(self, worker): - s3_client = worker.data.get(self.worker_key) - s3_client.close() + s3_client = self.s3_client + self.s3_client = None + if s3_client is not None: + if worker.data.get(self.worker_key) is s3_client: + worker.data.pop(self.worker_key, None) + s3_client.close() diff --git a/python/tests/io/test_s3_worker_plugin.py b/python/tests/io/test_s3_worker_plugin.py new file mode 100644 index 0000000000..3f28da78dd --- /dev/null +++ b/python/tests/io/test_s3_worker_plugin.py @@ -0,0 +1,264 @@ +"""Offline S3 worker lifecycle tests; no AWS credentials or network are used.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from botocore.config import Config +from botocore import UNSIGNED + +import mspasspy.io.s3_worker_plugin as s3 + + +@pytest.fixture(autouse=True) +def offline_clients(monkeypatch): + s3.fetch_s3_client.cache_clear() + monkeypatch.setattr( + s3.boto3, "Session", Mock(side_effect=AssertionError("unmocked Session")) + ) + monkeypatch.setattr( + s3.boto3, "client", Mock(side_effect=AssertionError("unmocked client")) + ) + monkeypatch.setattr( + s3, "EarthScopeClient", Mock(side_effect=AssertionError("unmocked SDK")) + ) + yield + s3.fetch_s3_client.cache_clear() + + +def test_serial_default_session_cached(monkeypatch): + session = Mock() + factory = Mock(return_value=session) + monkeypatch.setattr(s3.boto3, "Session", factory) + assert s3.fetch_s3_client(parallel=False) is session.client.return_value + assert s3.fetch_s3_client(parallel=False) is session.client.return_value + factory.assert_called_once_with() + session.client.assert_called_once() + args, kwargs = session.client.call_args + assert args == ("s3",) + assert kwargs["config"].request_checksum_calculation == "when_required" + assert kwargs["config"].response_checksum_validation == "when_required" + + +def test_serial_custom_session_and_config(): + session = Mock() + config = Config(region_name="us-west-2") + first = s3.fetch_s3_client(False, session, config) + assert s3.fetch_s3_client(False, session, config) is first + session.client.assert_called_once_with("s3", config=config) + alternate = Mock() + assert s3.fetch_s3_client(False, alternate, config) is alternate.client.return_value + alternate.client.assert_called_once_with("s3", config=config) + + +def test_parallel_looks_up_current_worker_and_registration(monkeypatch): + first, second, replacement = Mock(), Mock(), Mock() + workers = [ + SimpleNamespace(data={"geolab_s3client": first}), + SimpleNamespace(data={"geolab_s3client": second}), + ] + get_worker = Mock(side_effect=workers) + monkeypatch.setattr(s3, "get_worker", get_worker) + assert s3.fetch_s3_client() is first + assert s3.fetch_s3_client() is second + monkeypatch.setattr(s3, "get_worker", lambda: workers[1]) + workers[1].data["geolab_s3client"] = replacement + assert s3.fetch_s3_client() is replacement + workers[1].data["custom"] = first + # Ignored serial parameters need not be hashable in worker mode. + assert ( + s3.fetch_s3_client(session={}, session_config={}, worker_data_key="custom") + is first + ) + + +def test_parallel_errors_include_cause_and_key(monkeypatch): + failure = RuntimeError("outside worker") + monkeypatch.setattr(s3, "get_worker", Mock(side_effect=failure)) + with pytest.raises(ValueError, match="Dask worker context") as error: + s3.fetch_s3_client() + assert error.value.__cause__ is failure + monkeypatch.setattr(s3, "get_worker", lambda: SimpleNamespace(data={})) + with pytest.raises(ValueError, match="name=missing") as error: + s3.fetch_s3_client(worker_data_key="missing") + assert isinstance(error.value.__cause__, KeyError) + + +@pytest.mark.parametrize("plugin_type", [s3.AnonymousS3Client, s3.StockS3Client]) +def test_standard_plugin_configuration_and_lifecycle(monkeypatch, plugin_type): + client = Mock() + factory = Mock(return_value=client) + monkeypatch.setattr(s3.boto3, "client", factory) + plugin = plugin_type(key="custom", region="us-west-1") + worker = SimpleNamespace(data={"unrelated": object()}) + plugin.teardown(worker) + plugin.setup(worker) + assert worker.data["custom"] is client + args, kwargs = factory.call_args + assert args == ("s3",) + assert kwargs["region_name"] == "us-west-1" + if plugin_type is s3.AnonymousS3Client: + assert kwargs["config"].signature_version is UNSIGNED + else: + assert "config" not in kwargs + plugin.teardown(worker) + plugin.teardown(worker) + client.close.assert_called_once_with() + assert "custom" not in worker.data + assert "unrelated" in worker.data + + +class RejectWrites(dict): + def __setitem__(self, key, value): + raise RuntimeError("registration failed") + + +@pytest.mark.parametrize("plugin_type", [s3.AnonymousS3Client, s3.StockS3Client]) +def test_standard_plugin_setup_failures(monkeypatch, plugin_type): + plugin = plugin_type() + worker = SimpleNamespace(data=RejectWrites()) + client = Mock() + factory = Mock(return_value=client) + monkeypatch.setattr(s3.boto3, "client", factory) + with pytest.raises(RuntimeError, match="registration failed"): + plugin.setup(worker) + client.close.assert_called_once_with() + plugin.teardown(worker) + factory.side_effect = RuntimeError("creation failed") + with pytest.raises(RuntimeError, match="creation failed"): + plugin.setup(worker) + plugin.teardown(worker) + + +def geolab_mocks(monkeypatch): + sdk = Mock() + session = sdk.user.get_boto3_session.return_value + client = session.client.return_value + monkeypatch.setattr(s3, "EarthScopeClient", Mock(return_value=sdk)) + return sdk, session, client + + +def test_geolab_lifecycle_and_refresh_provider(monkeypatch): + sdk, session, client = geolab_mocks(monkeypatch) + plugin = s3.GeoLabS3Worker(key="custom") + worker = SimpleNamespace(data={}) + plugin.teardown(worker) + plugin.setup(worker) + assert worker.data["custom"] is client + sdk.user.get_boto3_session.assert_called_once_with() + session.client.assert_called_once() + sdk.close.assert_not_called() + client.close.assert_not_called() + plugin.teardown(worker) + plugin.teardown(worker) + sdk.close.assert_called_once_with() + client.close.assert_called_once_with() + assert "custom" not in worker.data + + +@pytest.mark.parametrize("stage", ["session", "client", "registration"]) +def test_geolab_setup_failure_releases_acquired_resources(monkeypatch, stage): + sdk, session, client = geolab_mocks(monkeypatch) + worker = SimpleNamespace(data=RejectWrites() if stage == "registration" else {}) + if stage == "session": + sdk.user.get_boto3_session.side_effect = RuntimeError("failed") + elif stage == "client": + session.client.side_effect = RuntimeError("failed") + plugin = s3.GeoLabS3Worker() + with pytest.raises(RuntimeError): + plugin.setup(worker) + sdk.close.assert_called_once_with() + if stage == "registration": + client.close.assert_called_once_with() + else: + client.close.assert_not_called() + plugin.teardown(worker) + assert not worker.data + + +def test_geolab_cleanup_continues_when_s3_close_fails(monkeypatch): + sdk, session, client = geolab_mocks(monkeypatch) + plugin = s3.GeoLabS3Worker() + worker = SimpleNamespace(data={}) + plugin.setup(worker) + client.close.side_effect = RuntimeError("close failed") + with pytest.raises(RuntimeError, match="close failed"): + plugin.teardown(worker) + sdk.close.assert_called_once_with() + plugin.teardown(worker) + assert not worker.data + + +def test_geolab_does_not_reuse_closed_serial_cached_client(monkeypatch): + sdk, session, client = geolab_mocks(monkeypatch) + closed = s3.fetch_s3_client(False, session) + closed.close() + fresh = Mock() + session.client.return_value = fresh + plugin = s3.GeoLabS3Worker() + worker = SimpleNamespace(data={}) + plugin.setup(worker) + assert worker.data["geolab_s3client"] is fresh + plugin.teardown(worker) + fresh.close.assert_called_once_with() + assert session.client.call_count == 2 + + +@pytest.mark.parametrize( + "plugin_type", [s3.AnonymousS3Client, s3.StockS3Client, s3.GeoLabS3Worker] +) +def test_teardown_preserves_replacement_client(monkeypatch, plugin_type): + if plugin_type is s3.GeoLabS3Worker: + _, _, client = geolab_mocks(monkeypatch) + else: + client = Mock() + monkeypatch.setattr(s3.boto3, "client", Mock(return_value=client)) + plugin = plugin_type() + worker = SimpleNamespace(data={}) + plugin.setup(worker) + replacement = Mock() + worker.data[plugin.worker_key] = replacement + plugin.teardown(worker) + assert worker.data[plugin.worker_key] is replacement + replacement.close.assert_not_called() + client.close.assert_called_once_with() + + +def test_serial_client_reads_existing_miniseed_via_stubbed_s3(monkeypatch): + from io import BytesIO + from pathlib import Path + + import boto3 + import numpy as np + from botocore.response import StreamingBody + from botocore.stub import Stubber + from obspy import read + + payload = (Path(__file__).parents[1] / "data" / "3channels.mseed").read_bytes() + expected = read(BytesIO(payload), format="MSEED") + # Explicit test credentials keep boto3 away from the credential-provider chain. + session = boto3.session.Session( + aws_access_key_id="testing", + aws_secret_access_key="testing", + region_name="us-east-1", + ) + client = s3.fetch_s3_client(False, session) + with Stubber(client) as stubber: + stubber.add_response( + "get_object", + { + "Body": StreamingBody(BytesIO(payload), len(payload)), + "ContentLength": len(payload), + }, + {"Bucket": "test-waveforms", "Key": "3channels.mseed"}, + ) + response = s3.fetch_s3_client(False, session).get_object( + Bucket="test-waveforms", Key="3channels.mseed" + ) + actual = read(BytesIO(response["Body"].read()), format="MSEED") + stubber.assert_no_pending_responses() + client.close() + assert len(actual) == len(expected) == 3 + for result, reference in zip(actual, expected): + assert result.stats == reference.stats + np.testing.assert_array_equal(result.data, reference.data) From 015670691ded29f2fd35e14f37595ac55e4403d6 Mon Sep 17 00:00:00 2001 From: Gary Pavlis Date: Wed, 30 Sep 2026 05:19:12 -0400 Subject: [PATCH 4/5] change module name The old name of s3_worker_plugin for the module was misleading. s3client is more descriptive because the module contains the fetch_s3_client function and not just the plugin code. --- python/mspasspy/io/{s3_worker_plugin.py => s3client.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename python/mspasspy/io/{s3_worker_plugin.py => s3client.py} (100%) diff --git a/python/mspasspy/io/s3_worker_plugin.py b/python/mspasspy/io/s3client.py similarity index 100% rename from python/mspasspy/io/s3_worker_plugin.py rename to python/mspasspy/io/s3client.py From e028e07360cc3fd8dc87261a8926172e4ecc6b1d Mon Sep 17 00:00:00 2001 From: Gary Pavlis Date: Wed, 30 Sep 2026 05:47:24 -0400 Subject: [PATCH 5/5] fix import in test program Fix a dumb mistake. After I changed the module name I neglected to change the import line for the test code for that module. --- python/tests/io/test_s3_worker_plugin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tests/io/test_s3_worker_plugin.py b/python/tests/io/test_s3_worker_plugin.py index 3f28da78dd..55b34e3b00 100644 --- a/python/tests/io/test_s3_worker_plugin.py +++ b/python/tests/io/test_s3_worker_plugin.py @@ -7,7 +7,7 @@ from botocore.config import Config from botocore import UNSIGNED -import mspasspy.io.s3_worker_plugin as s3 +import mspasspy.io.s3client as s3 @pytest.fixture(autouse=True)