diff --git a/python/mspasspy/io/s3client.py b/python/mspasspy/io/s3client.py new file mode 100644 index 0000000000..c38ce9f187 --- /dev/null +++ b/python/mspasspy/io/s3client.py @@ -0,0 +1,320 @@ +#!/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 + + +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 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. + + :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. 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 + 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: + 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. + 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() + # 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) + # 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): + """ + 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. + """ + 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: + 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 + self.s3_client = None + + def setup(self, worker): + s3_client = boto3.client( + "s3", region_name=self.region, config=Config(signature_version=UNSIGNED) + ) + + 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 = 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): + """ + 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 + self.s3_client = None + + 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, + ) + + 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 = 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..55b34e3b00 --- /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.s3client 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)