diff --git a/.github/workflows/earthscope-s3.yml b/.github/workflows/earthscope-s3.yml new file mode 100644 index 0000000..71a3020 --- /dev/null +++ b/.github/workflows/earthscope-s3.yml @@ -0,0 +1,35 @@ +name: EarthScope S3 workflow + +on: + pull_request: + paths: + - 'notebooks/EarthScopeS3/**' + - '.github/workflows/earthscope-s3.yml' + push: + branches: [master] + paths: + - 'notebooks/EarthScopeS3/**' + - '.github/workflows/earthscope-s3.yml' + +permissions: + contents: read + +jobs: + offline-auth-and-notebooks: + runs-on: ubuntu-latest + timeout-minutes: 15 + strategy: + fail-fast: false + matrix: + python: ['3.10', '3.13'] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python }} + - run: python -m pip install 'earthscope-sdk>=1.6.1,<1.8' 'dask[distributed]' pytest + - name: Credential lifecycle, worker process, and sanitized notebooks + working-directory: notebooks/EarthScopeS3 + run: python -m pytest -q test_s3_worker_plugin.py test_plugin_process.py test_notebooks.py + - name: Compile modules without running notebooks + run: python -m compileall -q notebooks/EarthScopeS3 diff --git a/notebooks/EarthScopeS3/.gitignore b/notebooks/EarthScopeS3/.gitignore new file mode 100644 index 0000000..7a083a5 --- /dev/null +++ b/notebooks/EarthScopeS3/.gitignore @@ -0,0 +1,8 @@ +__pycache__/ +.pytest_cache/ +.ipynb_checkpoints/ +css30/ +pf/ +*.pickle +*.bin +ANF48_*.json diff --git a/notebooks/EarthScopeS3/ExportMetadata.ipynb b/notebooks/EarthScopeS3/ExportMetadata.ipynb new file mode 100644 index 0000000..8f0499d --- /dev/null +++ b/notebooks/EarthScopeS3/ExportMetadata.ipynb @@ -0,0 +1,76 @@ +{ + "nbformat": 4, + "nbformat_minor": 5, + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Contributed workflow — live GeoLab validation is still required.**\n", + "\n", + "Read [README.md](README.md) first. Use a dedicated database and a small date subset. These notebooks can modify database collections and write S3 objects. No CSS data or credentials are distributed here." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Exporting Metadata Collections\n", + "This notebook uses pymongo to export database collections common to the entire data set I'm assembling. The json files it produces can be downloaded and the inverse performed on the local system. (a different notebook)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import json\n", + "import pymongo\n", + "from bson import json_util\n", + "\n", + "from mspasspy.client import Client\n", + "mspass_client=Client()\n", + "db = mspass_client.get_database(\"ANF48\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# list of collections to export\n", + "export_list = [\"site\",\"channel\",\"source\",\"arrival_css30\"]\n", + "for collection in export_list:\n", + " cursor = db[collection].find({})\n", + " output_file = f\"{db.name}_{collection}.json\"\n", + " with open(output_file,\"w\") as fh:\n", + " # Stream a JSON array without materializing the full collection.\n", + " fh.write(\"[\\n\")\n", + " for index, document in enumerate(cursor):\n", + " if index:\n", + " fh.write(\",\\n\")\n", + " fh.write(json_util.dumps(document))\n", + " fh.write(\"\\n]\\n\")\n", + " print(\"Wrote data for collection {} to file={}\".format(collection,output_file))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "" + ], + "outputs": [], + "execution_count": null + } + ] +} diff --git a/notebooks/EarthScopeS3/README.md b/notebooks/EarthScopeS3/README.md new file mode 100644 index 0000000..f697580 --- /dev/null +++ b/notebooks/EarthScopeS3/README.md @@ -0,0 +1,175 @@ +# EarthScope S3 daily waveform workflow + +This directory adapts the collaborator-supplied `reexternal2014datastatus.zip` +notebooks and Python modules (August 2026). Original module authorship is retained. +It belongs in the tutorial repository, not the MsPASS library. It is a contributed +workflow requiring local data and authorization, **not a validated GeoLab +production recipe**. No notebook cells have been executed against the supplied +CSS data or live GeoLab/S3 as part of this change. + +## What is fixed + +- `S3Worker` uses the EarthScope SDK's refreshable boto3 session directly. It + does not freeze temporary credentials or fabricate a new expiration for cached + credentials. One SDK/client pair lives until worker-plugin teardown; partial + setup is cleaned up. The client is a named worker-plugin resource, not spillable + `worker.data` or a task argument. +- Read/list/authentication/transport failures retain their original exception. + Only a definite missing-object HEAD response permits `#N` version fallback. + A generic HEAD 403 does not prove that the object exists. Missing/undecodable + waveforms still produce an error log; systemic failures stop processing. +- S3 response bodies close on success and failure. Upload errors, including an + error on file close/commit, propagate instead of returning an ignored `False`. +- Completion consumes its disposable input list as it merges station ensembles, + preserving member order and values while releasing source ensembles before + pickling. The native member vector is pre-reserved. This does not eliminate + native pickle's safe snapshot copies or guarantee a maximum RSS. +- Index listing is paginated. Day groups use the union of required days, and an + arrival is emitted once rather than once per joined holding. Year filters are + half-open and the index includes neighboring days needed by padded windows. +- Multi-rate data are converted once; requested detrending acts on the decoded + data. The metadata loader preserves unmatched station names and excludes `AK` + by value, not by the order MongoDB returns network codes. +- Saved notebook outputs, credential-printing cells and personal scratch paths + are removed. Metadata export streams a JSON array instead of loading the entire + collection into a second list. + +## What is deliberately unchanged + +Output is **one ordinary `TimeSeriesEnsemble` pickle per day**, not a new stream +format or serialization option. Normal keys remain +`//_.pickle`, including the complete user +prefix from `SCRATCH_BUCKET`. Existing readers using `pickle.load` keep working +with a compatible MsPASS installation. A waveform may still require a large +whole-day worker result, a driver gather, native copies, and upload buffering. + +The input access point and `s3-miniseed` SDK role are retained from the supplied +workflow. This is not an automatic migration to EarthScope's newer S3 interfaces. +Confirm that this legacy endpoint and your intended networks remain available +and authorized. See the [EarthScope S3 documentation](https://docs.earthscope.org/sdk/s3-direct-access-tutorial) +for current access rules. A role/bucket migration needs its own validation. + +The reader plugin and output `s3fs.S3FileSystem()` have **separate credential +providers**. SDK refresh on the reader does not renew copied/static credentials +used by the output filesystem, and does not grant scratch write permission. + +## Prerequisites and setup + +Use a dedicated working directory and database. The default database name in the +notebooks is `ANF48`; change it consistently if this is not a disposable test +database. Initialization drops/replaces `source`, `netmag` and `wf_s3` collections +and may duplicate arrivals on repeated runs. Do not run it on a valuable existing +database or rerun all initialization cells as a recovery step. + +1. Install a compatible current MsPASS build in the notebook **and workers**. + Required behavior includes `sliding_window_pipeline(..., retain_results=False)` + and the native ensemble pickle fixes through mspass-team/mspass#1025. + Updating Python source does not rebuild an already loaded native extension. +2. In that environment install `python -m pip install -r requirements.txt`. + The SDK range covers the session API used/tested here. Resolve s3fs/aiobotocore + dependencies together in a clean environment; do not overwrite a running + shared environment. pandas, ObsPy, PyMongo, Dask and boto3 are also required. +3. Authenticate EarthScope using the supported GeoLab/SDK mechanism on every + worker. Separately configure a refreshable identity authorized to write your + GeoLab scratch prefix. Never put keys/tokens into task arguments or notebook + output. No authentication setup is performed by the tests. +4. Obtain the CSS/Antelope tables separately, including the required `snetsta` + mapping, and place the supplied `usarray48.*` tables under `css30/`. The large + tables are not distributed here. Do not assume the public bulletin archive + contains the complete station cross-reference table used by the contributor. +5. Copy `data/pf/DatascopeDatabase.pf` from your MsPASS distribution into + `pf/DatascopeDatabase.pf`, as expected by `load_anfdata.ipynb`. +6. Open the notebooks with this directory as the working directory, so module + imports and `dask_client.upload_file(...)` find the adjacent Python files. + +## Notebook order + +1. `load_anfdata.ipynb`: load the CSS catalog and retrieve station metadata. + Review the contributor's geographic/phase selection before accepting it. +2. `s3indexing_2014.ipynb`: set `year` and build that year's index plus adjacent + days. Review `BUCKET` if your supported input endpoint differs. +3. `load_waveforms_s3_2014.ipynb`: set the same `year`/`BUCKET`, then choose a + populated subset with `first_julday` and `days_to_process`. The initial defaults + process three days with a sliding window of **1**. Missing holdings are exposed + before task submission. The scalar accumulator reports successful day files; + an empty/dead day is not a successful write. +4. Optional: `ExportMetadata.ipynb` exports metadata as JSON arrays; + `SetupDataTransfer.ipynb` describes copying output with an independently + authorized CLI, without printing credentials. + +`save_jday_outputs` runs on the driver, so `s3_day_workflow.py` stays there. +The notebook uploads the reader/plugin modules to workers. Do not enable +`completion_on_worker=True` with live database/filesystem handles as task +arguments; worker-side saving is a separate workflow change, not required here. + +Restart the kernel/workers when replacing modules from the old ZIP, then register +`S3Worker()` again. Its default registration name is `s3client`; with a custom +`S3Worker(key="archive")`, readers must use `fetch_s3_client(worker_data_key="archive")`. +Fetch no longer looks in `worker.data`. Do not override the registration name +independently of the key. Unregister the plugin/close the cluster when finished. + +If a systemic read/write error stops processing, already completed files and +failure records are not rolled back. Check actual output keys and read failures +before resuming. Tagged day outputs use the same keys on rerun and may overwrite +previous objects. This workflow is not a transactional or automatic resume system. + +## Offline checks + +From this directory in a compatible MsPASS environment: + +```bash +python -m pip install pytest +python -m pytest -q test_s3_worker_plugin.py test_notebooks.py test_plugin_process.py test_workflow.py +``` + +The credential tests use the real SDK session builder and botocore, with a fake +credential endpoint/cache and a controllable clock. They exercise a cached token +at minutes 41/55 and refresh beyond minutes 60/120, near-expired initial tokens, +refresh failures, and resource cleanup. They do not contact AWS or log keys. +The separate-process Dask probe uses an intentionally unpickleable fake client to +verify registration, lookup outside task storage, and teardown. + +The native tests use real ObsPy miniSEED and MsPASS objects with fake S3/database +boundaries. They check pagination, cross-year grouping, error propagation, body +cleanup, conversion/detrending, member order and sample values, ordinary pickle +round trips, and weak references **before serialization**, not just after return. +The Dask probe has a 45-second timeout and cleans up its process group. + +For a small memory comparison, run each mode in a fresh process (requires psutil): + +```bash +python benchmark_completion_memory.py legacy +python benchmark_completion_memory.py fixed +``` + +A local Python 3.13/protocol-4 run with 64 MiB of samples and the #1025 native +pickle implementation measured approximately 323 MiB versus 259 MiB RSS above +baseline during serialization. Both produced 64.16 MiB pickle streams. This is a +synthetic, no-S3 comparison, not a prediction for GeoLab or a peak-memory bound. + +GitHub Actions runs only credential, notebook, and worker-process checks on +Python 3.10/3.13; it does not install/build native MsPASS. `test_workflow.py` must +also run in a compatible native environment. No fake-backed test proves that +live authentication, scratch policies or a complete year's data work correctly. + +## Required GeoLab acceptance checks + +- First run `import os; print(os.getpid())` in the kernel and correlate it with + `ps`/`top`; do not infer the driver solely from a command named `python`. +- Verify loaded native behavior on the driver and workers. For example, + `type(TimeSeriesEnsemble().__getstate__()[3]).__name__` should be `list`, not + legacy `bytes`. This checks #1025 behavior, not every installed commit. +- Run a low-volume authentication check across the *real* token expiration + (e.g. 75–90 minutes), including the input and output providers independently. + Preserve operation, Error.Code, HTTP status and RequestId on the first failure; + never include secret keys, session tokens or signed headers in a report. +- Run a small day, representative day and heavy day. Record driver RSS before + building request lists, at completion entry, after merge and after upload. + Check actual object counts, member counts and samples with downstream readers. +- Increase concurrency only after measuring total pod memory. A window bounds + task count, not bytes. `del ens` immediately before return or forced garbage + collection is not a substitute for reducing live copies, task size and + allocator high-water usage. + +Full CSS/GeoLab execution, live credential renewal, real scratch upload throughput +and the collaborator's 12–18 GiB RSS remain explicitly unverified here. diff --git a/notebooks/EarthScopeS3/SetupDataTransfer.ipynb b/notebooks/EarthScopeS3/SetupDataTransfer.ipynb new file mode 100644 index 0000000..ab6a52e --- /dev/null +++ b/notebooks/EarthScopeS3/SetupDataTransfer.ipynb @@ -0,0 +1,30 @@ +{ + "nbformat": 4, + "nbformat_minor": 4, + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Transfer daily pickle files\n", + "\n", + "Use your institution's approved AWS authentication mechanism on the destination host. Ask the GeoLab administrators how to obtain a refreshable identity authorized for your scratch prefix. Do not print/export access keys, secret keys or session tokens in a notebook: saved outputs retain them and copied temporary credentials expire.\n", + "\n", + "The scratch value is a full URI including your user prefix; do not replace it with the bare bucket. Once the destination CLI is authenticated and authorized, substitute your own URI, year and region in:\n", + "\n", + "```bash\n", + "aws s3 sync 's3://YOUR-SCRATCH-BUCKET/YOUR-PREFIX/2014/' ./2014/ --region YOUR-REGION\n", + "```\n", + "\n", + "Verify object counts and sizes before treating a transfer as complete. Keep the source data until the destination is verified. No credentials or personal scratch URI are distributed in this notebook." + ] + } + ] +} diff --git a/notebooks/EarthScopeS3/benchmark_completion_memory.py b/notebooks/EarthScopeS3/benchmark_completion_memory.py new file mode 100644 index 0000000..fd91a0c --- /dev/null +++ b/notebooks/EarthScopeS3/benchmark_completion_memory.py @@ -0,0 +1,91 @@ +"""Synthetic 64 MiB completion comparison; run each mode in a fresh process. + +Usage: python benchmark_completion_memory.py legacy|fixed +No AWS, MongoDB, or Dask calls. RSS is sampled during pickle writes, so shorter +peaks may be missed. This is not a GeoLab benchmark or a bound on real day sizes. +""" + +import json +import pickle +import sys +from contextlib import contextmanager + +import psutil +from mspasspy.ccore.seismic import TimeSeries, TimeSeriesEnsemble + +import s3_day_workflow as workflow + +process = psutil.Process() +baseline = process.memory_info().rss +peak = baseline +written = 0 + + +class Sink: + def write(self, data): + global peak, written + count = memoryview(data).nbytes + written += count + peak = max(peak, process.memory_info().rss) + return count + + +@contextmanager +def open_sink(key, mode): + yield Sink() + + +def legacy_save(reader_output): + merged = TimeSeriesEnsemble() + for ensemble in reader_output: + for trace in ensemble.member: + merged.member.append(trace) + merged["jday_tag"] = "2014_1" + merged.set_live() + del reader_output + pickle.dump(merged, Sink()) + return True + + +def build(): + outputs = [] + for _ in range(32): + ensemble = TimeSeriesEnsemble(16) + for _ in range(16): + trace = TimeSeries(16384) + trace.set_live() + ensemble.member.append(trace) + ensemble["jday_tag"] = "2014_1" + ensemble.set_live() + outputs.append(ensemble) + return outputs + + +result = build() +mode = sys.argv[1] +if mode not in ("legacy", "fixed"): + raise ValueError("mode must be legacy or fixed") +if mode == "legacy": + result = legacy_save(result) +else: + workflow.fetch_dbhandle = lambda handle: {} + result = workflow.save_jday_outputs( + result, + None, + type("FS", (), {"open": staticmethod(open_sink)})(), + "s3://test/u", + 2014, + ) +assert result is True +print( + json.dumps( + { + "mode": mode, + "baseline_mib": baseline / 2**20, + "observed_peak_mib": peak / 2**20, + "peak_over_baseline_mib": (peak - baseline) / 2**20, + "pickle_mib": written / 2**20, + }, + sort_keys=True, + ) +) diff --git a/notebooks/EarthScopeS3/load_anfdata.ipynb b/notebooks/EarthScopeS3/load_anfdata.ipynb new file mode 100644 index 0000000..f9a8151 --- /dev/null +++ b/notebooks/EarthScopeS3/load_anfdata.ipynb @@ -0,0 +1,636 @@ +{ + "nbformat": 4, + "nbformat_minor": 5, + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Contributed workflow — live GeoLab validation is still required.**\n", + "\n", + "Read [README.md](README.md) first. Use a dedicated database and a small date subset. These notebooks can modify database collections and write S3 objects. No CSS data or credentials are distributed here." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Load ANF data\n", + "This workflow is driven by a set of CSS3.0 relational database tables created with great effort by analysts at the Earthscope Array Network Facility during the operation of the Earthscope Transportable Array (TA). The permanent archive for these data can be found [here](https://doi.org/10.17611/DP/EB.1). From that top-level you browse to \"anf\" directory. That large web page has arrival database tables by month of TA operation. To build a master arrival database for building a usarray data set, however, it was far more convenient to download the merge of all these bundled into a single tar file. That tar file is called \"TA_Event_DB.tar\". Note again it is found in \"anf\" directory one in the doi link above. A more volatile link is [this one](https://gage-data.earthscope.org/archive/seismology/products/event_bulletins/anf/). I did some external work on my desktop to reduce the size of that database. That reduction is just limiting the time range to when the TA was operating in the lower 48. Deployment pages say the last official TA station in the lower 48 ended operation in September 2015. The first monthly is June 2004 so I kept all arrival data from June 2004 through the end of September 2015. That operation is a simple set of shell commands easily recreated. The downloaded tar file is an image of that entire anf directory tree sans the tar file itself. I took apart the tar file into a scratch work area and then ran this in that directory:\n", + "```\n", + "tables=(arrival assoc event netmag origin stamag)\n", + "for t in \"${tables[@]}\"; do\n", + " echo Working on relation=$t\n", + " cat 200*/*.$t > usarray48.$t\n", + " cat 201[0-4]_??/*.$t >> usarray48.$t\n", + " cat 2015_0?/*.$t >> usarray48.$t\n", + "done\n", + "```\n", + "Noting I explicitly dropped wfmeas as it wasn't consistently used and we don't need it for our study. I also dropped origerr as we don't need that here either and it would just be baggage. I downloaded that data to the css3.0 directory found in this work area.\n", + "\n", + "There is one additional gotcha. I had to obtain the final \"snetsta\" table from Frank Vernon who was the PI for the Earthscope ANF. Earthscope, for some unknown reason, has not archived that table. It is a bit of a detail but they are needed in this workflow to sort out the mismatch in concept between how css3.0 defines a station name and seed defines a station name. If anyone is reading this who needs to reproduce this workflow you will need that table. If this work ever results in publicationw we will need to supply that file. The data handled by that shell script above is not trivial as it totals around 5 G. \n", + "\n", + "The data are tabulated in a set of text files that are used for data storage in Antelope's relational database system that ANF used for their operation. MsPASS has a special python class to handle that type of data. Hence, we first have to instantiate an instance of the what we call the DatascopeDatabase class. The %env magic is needed to tell the DatascopeDatabase constructor how to load its dictionary of what attributes are in what table. The need for that is an oddity of the setup here on GeoLab. Oh, actually that file had to be copied from the master. i.e. DatascopeDatabase.pf is in the mspass distribution data/pf directory. Here I copied that to a local \".pf\" directory." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from mspasspy.preprocessing.css30.datascope import DatascopeDatabase\n", + "anfdb = DatascopeDatabase(\"css30/usarray48\",pffile=\"./pf/DatascopeDatabase.pf\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This does all the dirty wrk to build a single DataFrame of what we need except the orid==prefor subsetting. That is done later in this workflow." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "df = anfdb.CSS30Catalog2df()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(len(df))\n", + "print(df.keys())" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Unassociated arrival entries turn out to leave NaN values in the evid and orid fields. That causes pandas to change the type from int to float. I think the reason is there is no natural undefined value in a dataframe for a int. The next line fixes that by deleting all tuples for which evid or orid are NaN and then forcing evid and orid to be integers." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Drop the nan tuples and then convert evid to int\n", + "df = df.dropna(subset=['evid','orid'])\n", + "df['evid'] = df['evid'].astype(int)\n", + "df['orid'] = df['orid'].astype(int)\n", + "print(\"Size of DataFrame with unassociated arrivals removed=\",len(df))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## station metadata\n", + "We will need a list of stations eventually to fetch proper metadata. For now this just validates an algorithm." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "stalist = df['sta'].unique()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# examine result for validity\n", + "print(stalist)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This month doesn't seem to have a net-sta issue but to make this more general one should handle that. This little block builds a dictionary to cross-reference Antelope naming defined by their netsta table and net:sta used by Earthscope. Note a complication is the method we called to create the giant dataframe doesn't merge netsta.\n", + "\n", + "There is a strange issue with the nan station name. That may cause problems downstream.\n", + "\n", + "IMPORTANT: I hacked in the snetsta table currently in the css30 directory. The master data repository has arrival data in month directories with the dbmaster data somewhere else. I'll need to figure that out eventually. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "df_netdata=anfdb.get_table('snetsta')\n", + "# useful to show this\n", + "print(df_netdata.columns)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Not obvious and explained only by digging into the antelope schema documents is that fsta is the seed sta name while \"sta\" for antelope is the composite string used to handle stations with duplicate station names and different net codes. This little block demonstrates how many of them we have in the snetsta table we loaded." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(df_netdata)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for row in df_netdata.itertuples():\n", + " if row.fsta!=row.sta:\n", + " print(row.fsta,row.sta)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now that we know the difference between fsta and sta we can build this cross reference. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "netsta_xref=dict()\n", + "for row in df_netdata.itertuples():\n", + " netsta_xref[row.sta] = [row.snet,row.fsta]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# print a couple examples and verify the length\n", + "print(\"size of cross-reference table=\",len(netsta_xref))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By combining the contents of stalist and netsta_xref we can build a bombproof algorithm to fetch station metadata with now standard web servcies. \n", + "\n", + "## Build subsetted arrivals DataFrame \n", + "We next need to build a subset of the catalog dataframe that includes only teleseismic events. For this case because the anf data are clean we can use the column of data with the key \"delta\", which in css3.0 means the distance in degrees. That is a relatively standard DataFrame operation done in the next frame." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# there are multiple ways to do the following operation. This is the clearest to humans\n", + "# this query uses stock limits for teleseismic P wave processing for receiver function deconvolution \n", + "# which the tutorial workflow will be doing\n", + "# note prefor clause is needed because the method that created df doesn't do that according to the docstring\n", + "dftele = df.query('delta>30.0 and delta<100.0 and iphase==\"P\" and orid==prefor')\n", + "print(\"Size of original catalog = \",len(df))\n", + "print(\"Size of P wave subset = \",len(dftele))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "That is a reasonably large data set for a tutorial exercise with a group. It would seem a monthly will work fine. I will probably change the month eventually to be sure the data set contains some useful data, but this will certainly do for prototyping. \n", + "\n", + "I think the next step is pulling out only columns from the DataFrame that are essential to save as metadata. This will also require using the cross reference even if we don't need it for this month of data. First, here are the keys printed in a more readable form:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for k in dftele.keys():\n", + " print(k)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "copylist=['evid','orid','arid','phase','delta','seaz','esaz','timeres','iphase','fm']\n", + "rename_list = { 'lat' : 'source_lat',\n", + " 'lon' : 'source_lon',\n", + " 'depth' : 'source_depth',\n", + " 'time' : 'source_time',\n", + " 'time_arrival' : 'Ptime',\n", + " 'chan' : 'pick_channel',\n", + " }\n", + "magkeys=['ms','ml','mb']\n", + "# sta and net are handled specially using the cross reference\n", + "# note iterrows is not the fastest way to do this but good enough for \n", + "# this application - cleaner syntax for sure for these data\n", + "doclist=list()\n", + "for index, row in dftele.iterrows():\n", + " doc=dict()\n", + " for k in copylist:\n", + " doc[k] = row[k]\n", + " for k in rename_list.keys():\n", + " newname = rename_list[k]\n", + " val=row[k]\n", + " doc[newname] = val\n", + " for k in magkeys:\n", + " val = row[k]\n", + " # null magnitudes are given a large negative value in antelope \n", + " # this needs to be more generic if reused in different context\n", + " if val>4.0: # could be 0 but appropriate for teleseism dataset\n", + " doc[k] = val\n", + " # handle neT:AR code\n", + " css_sta = row.sta\n", + " if css_sta in netsta_xref:\n", + " x = netsta_xref[css_sta]\n", + " net = x[0]\n", + " sta = x[1]\n", + " doc['net'] = net\n", + " doc['sta'] = sta\n", + " else:\n", + " print(f\"Warning: {css_sta} not found in cross reference dictionary created from snetsta\")\n", + " print(\"sta is left unaltered and net is set to TA\")\n", + " doc[\"net\"] = \"TA\"\n", + " doc[\"sta\"] = css_sta\n", + " doclist.append(doc)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# clean up memory \n", + "del df\n", + "del dftele" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# QC results\n", + "from mspasspy.util.seismic import print_metadata\n", + "print(\"Size of arrival document list=\",len(doclist))\n", + "print(\"Example document with pretty printing\")\n", + "print_metadata(doclist[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Save to MongoDB\n", + "I think the right approach here is to save these data to special mongodb collection. We can read that later to drive waveform extraction from s3." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from mspasspy.client import Client\n", + "mspass_client=Client()\n", + "db = mspass_client.get_database(\"ANF48\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "out = db.arrival_css30.insert_many(doclist)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# verify \n", + "n=db.arrival_css30.count_documents({})\n", + "print(\"size of new collection created=\",n)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "del doclist" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Get Station Metadata\n", + "We know from lots of previous experience that downloading station metadata with obspy is best done as needed. The data volume is relatively lightweight, the service runs fast, and there is a simple path to load the result to MongoDB found in all previous Earthscope tutorials. \n", + "\n", + "This version differs a bit from previously only in the way we select what metadata to download. In this case we need to download station data for all networks that have picks in the arrival_css30 collection we just created." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "netlist=db.arrival_css30.distinct('net')\n", + "print(netlist)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "There are extras here from when the TA operated stations in Alaska at the same times as the lower 48. Because I selected only by distance there was nothing to exclude that data. A good way to keep them out of our data set is to not download the station metadata. The get_stations run below will do most of that, but I do want to exclude AK (Alaska monitoring network) immediately. The extra code box before the main one that runs a long time does that.\n", + "\n", + "I'm using a generous time window but limiting the search region to a generous box around the lower 48 stations of the US. That choice is because the focus here is stations the ANF used to produce their catalog. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "netlist = sorted(net for net in netlist if isinstance(net, str) and net and net != \"AK\")\n", + "print(netlist)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from obspy import UTCDateTime\n", + "from obspy.clients.fdsn import Client\n", + "from obspy.clients.fdsn.header import FDSNNoDataException\n", + "client=Client(\"Earthscope\")\n", + "\n", + "# time range all of 2004 through end of 2015 \n", + "ts=UTCDateTime('2004-01-01T00:00:00.0')\n", + "starttime=ts\n", + "te=UTCDateTime('2016-01-01T00:00:00.0')\n", + "latitude_range=[20.0,55.0]\n", + "longitude_range=[-130.0,-60]\n", + "for net in netlist:\n", + " try:\n", + " inv=client.get_stations(network=net,\n", + " starttime=ts,\n", + " endtime=te,\n", + " minlatitude=latitude_range[0],\n", + " maxlatitude=latitude_range[1],\n", + " minlongitude=longitude_range[0],\n", + " maxlongitude=longitude_range[1],\n", + " format='xml',\n", + " channel='BH?',\n", + " level='response',\n", + " )\n", + " print(\"Saving metadata for network=\",net)\n", + " db.save_inventory(inv)\n", + " except FDSNNoDataException as ex:\n", + " print(\"client threw an FDSNNoDataException. This is the message it posted\")\n", + " print(str(ex))\n", + " print(\"Cannot obtain station metadata for network code=\",net)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finally, get a summary of how many documents were saved in channel and site." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "nsite=db.site.count_documents({})\n", + "nchan=db.channel.count_documents({})\n", + "print(\"Number of documents in channel collection=\",nchan)\n", + "print(\"Number of documents in site collection=\",nsite)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Save source data\n", + "Saving source data is a bit different because I don't wish to save all the source data in the Datascope database. We could and it wouldn't matter a lot but this is cleaner. There is also a serious complication in handling magnitude data. That turns out to be problematic with the anf data. Earlier tests showed I could not use the netmag table, which exists in this database, with mspass join. What happens is some origin tuples have multiple netmag entries. That caused duplicates when I had the join command commented out below. Revision below handles that differently." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "df_src = anfdb.get_table(\"event\")\n", + "df_src = anfdb.join(df_src,\"origin\",join_keys=[\"evid\"])\n", + "df_src = df_src[df_src[\"orid\"] == df_src[\"prefor\"]]\n", + "# these had to be excluded as they caused duplicates\n", + "#df_src = anfdb.join(df_src,\"netmag\",join_keys=[\"evid\",\"orid\"])\n", + "#df_src = df_src[df_src[\"orid\"] == df_src[\"prefor\"]]\n", + "print(df_src)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# similar to above drop rows with prefor null and convert prefor to an int\n", + "df_src = df_src.dropna(subset=['prefor'])\n", + "df_src['prefor'] = df_src['prefor'].astype(int)\n", + "print(\"Size of source DataFrame with null prefor tuples removed=\",len(df_src))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# use this list to select only orid values of data used \n", + "oridlist=db.arrival_css30.distinct(\"orid\")\n", + "print(\"Number of events with arrivals saved=\",len(oridlist))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we need to handle the earlier problem I had with netmag. I'm going to create a scratch netmag collection in MongoDB. Note I always clear that collection since it is scratch." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import numpy as np\n", + "db.drop_collection(\"netmag\")\n", + "df_nm = anfdb.get_table(\"netmag\")\n", + "# this is needed to make mongodb cleaner with any NaN values\n", + "df_nm = df_nm.replace({np.nan: None})\n", + "records = df_nm.to_dict(orient=\"records\")\n", + "save_ret = db.netmag.insert_many(records)\n", + "print(\"Number of records saved=\",len(save_ret.inserted_ids))\n", + "del save_ret" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# This index will speed the queries below\n", + "db.netmag.create_index([(\"evid\",1)])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This next block that saves data is not at all generic. It intentionally discards attributes from the css table that are not needed for this tutorial." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "df_src.keys()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "db.drop_collection(\"source\")\n", + "doclist=list()\n", + "nmags=0\n", + "n=0\n", + "for index,row in df_src.iterrows():\n", + " if row.orid in oridlist:\n", + " n += 1\n", + " doc = { \"lat\" : row.lat,\n", + " \"lon\" : row.lon, \n", + " \"depth\" : row.depth,\n", + " \"time\" : row.time, \n", + " \"orid\" : row.orid,\n", + " \"evid\" : row.evid,\n", + " \"nass\" : row.nass,\n", + " \"ndef\" : row.ndef,\n", + " \"auth\" : row.auth,\n", + " }\n", + " evid = row.evid\n", + " magdoc = db.netmag.find_one({\"evid\" : evid})\n", + " if magdoc is not None:\n", + " nmags += 1\n", + " if \"magnitude\" in magdoc:\n", + " doc[\"magnitude\"] = magdoc[\"magnitude\"]\n", + " if \"magtype\" in magdoc:\n", + " doc[\"magtype\"] = magdoc[\"magtype\"]\n", + "\n", + " if row.ml>0.0:\n", + " doc[\"ml\"] = row.ml\n", + " if row.mb>0.0:\n", + " doc[\"mb\"] = row.mb\n", + " if row.ms>0.0:\n", + " doc[\"ms\"] = row.ms\n", + " doclist.append(doc)\n", + "insert_result=db.source.insert_many(doclist) \n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(n,nmags)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from mspasspy.util.seismic import print_metadata\n", + "n=db.source.count_documents({})\n", + "print(\"Number of source documents saved=\",n)\n", + "print(\"Typical content\")\n", + "doc=db.source.find_one()\n", + "print_metadata(doc)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "" + ], + "outputs": [], + "execution_count": null + } + ] +} diff --git a/notebooks/EarthScopeS3/load_waveforms_s3_2014.ipynb b/notebooks/EarthScopeS3/load_waveforms_s3_2014.ipynb new file mode 100644 index 0000000..b74d4ba --- /dev/null +++ b/notebooks/EarthScopeS3/load_waveforms_s3_2014.ipynb @@ -0,0 +1,371 @@ +{ + "nbformat": 4, + "nbformat_minor": 5, + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Contributed workflow — live GeoLab validation is still required.**\n", + "\n", + "Read [README.md](README.md) first. Use a dedicated database and a small date subset. These notebooks can modify database collections and write S3 objects. No CSS data or credentials are distributed here." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Read Waveform Data\n", + "As the name suggests this notebook constructs a working data set for a year set in box 2. This notebook needs to be run after running load_anfdata and s3indexing. The output is two things: (1) a set of day pickle files written to Earthscope's scratch bucket, (2) a new collection called \"s3_read_failures\". The pickle files can be transferred to local store and loaded easily into mongoDB with a read - save_data loop over all files (or map driven by the list of pickle files). Warning: those pickle files are day files, not event files. To sort the data into common source ensembles is best done after loading with appropriate queries. (2) is designed to make it possible to retry failures. That is, each document in that collection has the original document sent to the s3 reader to pull data that failed. \n", + "\n", + "First, the stock starting incantation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from mspasspy.client import Client\n", + "mspass_client=Client()\n", + "db=mspass_client.get_database(\"ANF48\")\n", + "dask_client = mspass_client.get_scheduler()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# this workflow is run in yearly blocks. Change this to change year being handled this run\n", + "year = 2014\n", + "first_julday = 1\n", + "days_to_process = 3 # Set to None only after validating a representative subset." + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# this isn't essential but convenient for monitoring a long run with dask diagnostics\n", + "dask_client" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fetch only data suitable for teleseismic P wave processing for receiver functions. Use the \"delta\" attribute from assoc to limit data to epicentral distances between 30 and 95. Note in this notebook that does nothing as we already filtered input that way, but shows an alternative way to do tha with MongoDB." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from obspy import UTCDateTime\n", + "starttime = UTCDateTime(year=year,julday=first_julday)\n", + "endtime = UTCDateTime(year=year+1,month=1,day=1)\n", + "if days_to_process is not None:\n", + " endtime = min(endtime, starttime + 86400 * days_to_process)\n", + "query = {\"$and\" : [\n", + " {\"delta\" : {\"$gte\" : 30.0, \"$lte\" : 95.0}},\n", + " {\"Ptime\" : {\"$gte\" : starttime.timestamp, \"$lt\" : endtime.timestamp}},\n", + " ] }\n", + "\n", + "cursor=db.arrival_css30.find(query)\n", + "doclist = list(cursor)\n", + "if not doclist:\n", + " raise ValueError(\"No arrivals in this date/distance subset; choose a populated interval.\")\n", + "print(\"Size of arrival table loaded = \",len(doclist)) # validate this is reasonable" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# reader we use requires a DataFrame as input \n", + "import pandas as pd\n", + "df = pd.DataFrame(doclist)\n", + "print(df)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The reader function we will use below requires the DataFrame key \"time\" to define the reference time it is to use to define waveform segments. For clarity I chose to set the arrival time from the ANF pick table to \"Ptime\". Hence, we need this trivial patch from the pandas API." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "df = df.rename(columns={\"Ptime\" : \"time\"})" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "del doclist" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Build Parallel Input\n", + "We need to define a series of functions needed to drive this workflow. These may be moved to a module but for now the notebook is preferable as there are likely residual bugs that are way easier to fix if they are inside this notebook.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import s3segment_reader as geos3" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "With those functions we can now create a (perhaps overcomplicated) data structure to drive parallel processing. This function is run serial because it runs fast enough that a parallel construct is unnecessary." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data_list,error_list = geos3.build_day_processing_blocks(df,\n", + " -240,\n", + " 300,\n", + " db,\n", + " auxkeys=[\"evid\",\"orid\",\"arid\",\"iphase\",\"delta\",\"pick_channel\"],\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The outputs (data_list and error_list) are rather complicated. Be aware the v1 copies of this notebook have a debug section that clarifies that content. For this production version that has been removed.\n", + "\n", + "We are done with the DataFrame so release it to reduce memory use." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "del df" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This plugin is required for the parallel job below to work. It instantiates a memory resident client on each worker." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dask_client.upload_file(\"s3_worker_plugin.py\")\n", + "\n", + "from s3_worker_plugin import S3Worker\n", + "s3worker = S3Worker()\n", + "dask_client.register_plugin(s3worker)\n", + "# This registers the named \"s3client\" plugin. Restart workers when updating modules." + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This box is needed for both serial and parallel processing. The parallel version needs the session object for the completion function that writes pickle files. Kind of confusing but we need this client for the completion function because it runs under the control of the python process connected to this notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "BUCKET = \"earthscope-mseed-res-na3mtd4fq5kz7pntcyr1uh46use2a--ol-s3\"\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The parallel job below had a serial test found in the version 1 notebooks. I deleted it here as it was baggage for production.\n", + "\n", + "This notebook does the writer as a completion function with sliding_window_pipeline. We give it a dbname and have it fetch db with fetch_db_handle for more flexibility but the expectation is database activity will be limited to the process running the master script - norm for a completion function. For GeoLab this doesn't seem to be a bottleneck because only there is currently a limit of 4 cpus and the work done to crack the miniseed files and extract segments seems to keep that many cpus busy.\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from s3_day_workflow import save_jday_outputs, count_saved\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from mspasspy.workflow import sliding_window_pipeline\n", + "import os\n", + "\n", + "# Fetch the scratch bucket path\n", + "scratch_bucket = os.environ[\"SCRATCH_BUCKET\"]\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# The way I did the above requires I strip component 1 from the day list to drive the processing. \n", + "process_docs=list()\n", + "for i in range(len(data_list)):\n", + " process_docs.append(data_list[i][1])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# worth recording this before starting\n", + "len(process_docs)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# we need to merge these to run sliding_window_pipeline\n", + "def read_and_merge(doclist,bucket=BUCKET):\n", + " enslist=list()\n", + " for doc in doclist:\n", + " ens = geos3.get_s3_segments(doc,bucket=bucket,detrend_type=\"simple\")\n", + " enslist.append(ens)\n", + " return enslist\n", + " " + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# this is needed to give workers access to functions in the local s3segement_reader.py file\n", + "dask_client.upload_file(\"s3segment_reader.py\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import time\n", + "import s3fs\n", + "s3fshandle = s3fs.S3FileSystem()\n", + "t0=time.time()\n", + "print(\"Submitting retrieve to dask cluster\")\n", + "out = sliding_window_pipeline(process_docs,\n", + " read_and_merge,\n", + " dask_client,\n", + " pfunc_kwargs={\"bucket\" : BUCKET},\n", + " completion_function=save_jday_outputs,\n", + " cfunc_args=[db,s3fshandle,scratch_bucket,year],\n", + " sliding_window_size=1,\n", + " accumulator=count_saved,\n", + " retain_results=False,\n", + " verbose=True,\n", + " progress_report_interval=10,\n", + " )\n", + "t=time.time()\n", + "print(\"Elapsed time=\",t-t0)\n", + "print(\"Number of items handled=\",len(process_docs))\n", + "print(\"Number of day files successfully saved=\", out)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# important check on data loss\n", + "n=db.s3_read_failures.count_documents({})\n", + "print(\"Number of failures=\",n)\n", + "print(\"Check contents of s3_read_failures collection if it is not zero\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "" + ], + "outputs": [], + "execution_count": null + } + ] +} diff --git a/notebooks/EarthScopeS3/plugin_process_probe.py b/notebooks/EarthScopeS3/plugin_process_probe.py new file mode 100644 index 0000000..d204c67 --- /dev/null +++ b/notebooks/EarthScopeS3/plugin_process_probe.py @@ -0,0 +1,66 @@ +"""Offline worker-process probe, launched with a timeout by its pytest test.""" + +from types import SimpleNamespace + +from distributed import Client, LocalCluster + +import s3_worker_plugin as plugin + + +class UnpickleableClient: + def __init__(self, events): + self.events = events + + def __reduce__(self): + raise RuntimeError("An S3 client must never be serialized as task data") + + def close(self): + self.events.append("s3") + + +class ProbePlugin(plugin.S3Worker): + def setup(self, worker): + worker.auth_probe_events = [] + events = worker.auth_probe_events + session = SimpleNamespace(client=lambda *a, **kw: UnpickleableClient(events)) + sdk = SimpleNamespace( + user=SimpleNamespace(get_boto3_session=lambda: session), + close=lambda: events.append("sdk"), + ) + original = plugin.EarthScopeClient + plugin.EarthScopeClient = lambda: sdk + try: + super().setup(worker) + finally: + plugin.EarthScopeClient = original + + +def check_client(): + from distributed import get_worker + from s3_worker_plugin import fetch_s3_client + + worker = get_worker() + client = fetch_s3_client(worker_data_key="archive") + return ( + client is worker.plugins["archive"].s3_client and "archive" not in worker.data + ) + + +def close_events(dask_worker=None): + return dask_worker.auth_probe_events + + +if __name__ == "__main__": + with LocalCluster( + n_workers=1, + threads_per_worker=1, + processes=True, + memory_limit=0, + dashboard_address=None, + ) as cluster: + with Client(cluster) as client: + client.register_plugin(ProbePlugin(key="archive")) + assert client.submit(check_client, pure=False).result(timeout=10) + client.unregister_worker_plugin("archive") + assert list(client.run(close_events).values()) == [["s3", "sdk"]] + print("worker-process registration, fetch, and teardown: PASS") diff --git a/notebooks/EarthScopeS3/requirements.txt b/notebooks/EarthScopeS3/requirements.txt new file mode 100644 index 0000000..806802f --- /dev/null +++ b/notebooks/EarthScopeS3/requirements.txt @@ -0,0 +1,3 @@ +# Install into an existing, compatible MsPASS environment, not in place of it. +earthscope-sdk>=1.6.1,<1.8 +s3fs diff --git a/notebooks/EarthScopeS3/s3_day_workflow.py b/notebooks/EarthScopeS3/s3_day_workflow.py new file mode 100644 index 0000000..95a87ee --- /dev/null +++ b/notebooks/EarthScopeS3/s3_day_workflow.py @@ -0,0 +1,91 @@ +"""Driver completion for the contributed daily-pickle workflow. + +This preserves one ordinary TimeSeriesEnsemble pickle per day, including the +scratch URI's user prefix. It does not make a whole day a bounded-size task. +""" + +import pickle + +from bson.objectid import ObjectId +from mspasspy.ccore.seismic import TimeSeriesEnsemble +from mspasspy.db.database import elog2doc +from mspasspy.util.db_utils import fetch_dbhandle + + +def merge_ensembles(enslist, yearday_key="jday_tag"): + """Merge in order and consume the input list, replacing entries with None. + + Called only for a disposable worker result. A Python parameter deletion + cannot release the pipeline's reference to that list; clearing its entries + releases source ensembles incrementally. Members are still copied into the + output, but the complete input need not coexist with serialization state. + """ + count = sum(len(ens.member) for ens in enslist if ens.live) + merged = TimeSeriesEnsemble(count) + failures = [] + for index in range(len(enslist)): + ens = enslist[index] + try: + if ens.live: + if yearday_key not in merged and yearday_key in ens: + merged[yearday_key] = ens[yearday_key] + for member in ens.member: + merged.member.append(member) + # A pybind member wrapper can keep its source vector alive. + if len(ens.member): + del member + else: + document = dict(ens) + document["elog_content"] = elog2doc(ens.elog) + failures.append(document) + finally: + enslist[index] = None + del ens + if len(merged.member): + merged.set_live() + return [merged, failures] + + +def save_jday_outputs( + reader_output, + dbname_or_handle, + s3fshandle, + bucket, + year, + yearday_key="jday_tag", + s3_fail_collection="s3_read_failures", + verbose=False, +): + """Consume reader_output and save a daily ensemble; return whether saved. + + bucket is the full SCRATCH_BUCKET URI, potentially including a user prefix. + An empty/dead day returns False after recording read failures. Database and + upload errors propagate unchanged so a failed write cannot look successful. + The caller owns the filesystem and database handles. + """ + if not isinstance(bucket, str) or not bucket.strip(): + raise ValueError("SCRATCH_BUCKET must be a nonempty bucket/prefix URI") + db = fetch_dbhandle(dbname_or_handle) + ens, failures = merge_ensembles(reader_output, yearday_key=yearday_key) + nfailed = len(failures) + if failures: + db[s3_fail_collection].insert_many(failures) + del failures + if ens.dead(): + return False + if yearday_key in ens: + basename = ens[yearday_key] + else: + basename = str(ObjectId()) + print(f"Missing {yearday_key}; using unique output name {basename}") + object_key = f"{bucket.rstrip('/')}/{year}/{basename}.pickle" + if verbose: + print(f"Saving {len(ens.member)} members to {object_key}; failures={nfailed}") + with s3fshandle.open(object_key, "wb") as output: + pickle.dump(ens, output) + return True + + +def count_saved(total, saved): + """Count successful day files with constant-size pipeline state.""" + return (0 if total is None else total) + int(saved) diff --git a/notebooks/EarthScopeS3/s3_worker_plugin.py b/notebooks/EarthScopeS3/s3_worker_plugin.py new file mode 100644 index 0000000..228f426 --- /dev/null +++ b/notebooks/EarthScopeS3/s3_worker_plugin.py @@ -0,0 +1,134 @@ +#!/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 +""" + +from contextlib import ExitStack +from botocore.config import Config +from earthscope_sdk import EarthScopeClient +from dask.distributed import WorkerPlugin, get_worker + + +def fetch_s3_client(session=None, worker_data_key="s3client"): + """ + Generic tool to fetch s3 client for Earthscope s3 archive. + + When session is None (default) the funcion assumes it is running in a parallel environment with dask. + In that situation it assumes workers have been previously initialized with the WorkerPlugin "S3Worker" + and it fetches the client from worker.plugins using "worker_data_key" as the plugin name. + + When session is defined a client is instantiated with the client method of the session object. + Note on GeoLab the session should be constructed using this construct: + ``` + client = EarthScopeClient() + session = client.user.get_boto3_session() + ``` + There are other ways to do that, but Earthscope advises that construct to avoid a problem + timeout of credentials after an hour that can happen otherwise. + """ + if session is None: + 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.plugins[worker_data_key].s3_client + + 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 = session.client( + "s3", + config=Config( + request_checksum_calculation="when_required", + response_checksum_validation="when_required", + ), + ) + + return s3_client + + +class S3Worker(WorkerPlugin): + """ + Dask worker plugin to create a worker resident client on each dask worker. + + 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 = S3Worker() + 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. + """ + + def __init__(self, key="s3client"): + """ + Standard constructor. + + Normal use requires no arguments. One could change the key argument + but that is not advised. The S3 client is held by a named worker plugin, outside the spillable + worker.data dictionary. The key selects that plugin name. + The only reason to change it would be if some other worker plugin + used the same key, which should never happen. + """ + self.worker_key = key + self.name = key + + 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. + """ + # 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() + s3_client = fetch_s3_client(session) + cleanup.callback(s3_client.close) + # 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 + 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. + """ + cleanup = getattr(self, "_cleanup", None) + if cleanup is not None: + self._cleanup = None + self.s3_client = None + cleanup.close() diff --git a/notebooks/EarthScopeS3/s3indexing_2014.ipynb b/notebooks/EarthScopeS3/s3indexing_2014.ipynb new file mode 100644 index 0000000..55789f9 --- /dev/null +++ b/notebooks/EarthScopeS3/s3indexing_2014.ipynb @@ -0,0 +1,270 @@ +{ + "nbformat": 4, + "nbformat_minor": 5, + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Contributed workflow — live GeoLab validation is still required.**\n", + "\n", + "Read [README.md](README.md) first. Use a dedicated database and a small date subset. These notebooks can modify database collections and write S3 objects. No CSS data or credentials are distributed here." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Build index to s3 objects\n", + "s3 objects are equivalent to files. They are just stored and managed in a different way than file systems. This notebook demonstrates how to build an index of the holdings of data by Earthscope for the test TA data. \n", + "\n", + "Earthscope organizes their waveform in miniseed day volumes. The day volumes are held in the s3 equivalent of a directory which they call a \"prefix\". This notebook creates a MongoDB collection with the name \"wf_s3\" \n", + "that will be used to drive processing to extract event data. This notebook is indepent of that and more-or-less builds a directory of all file-like (s3 objects) linked to TA data for all of the year 2010. Our test data does not cover the entire year, but we do the entire year to demonstrate that creating the index is very lightweight. Fetching waveforms using it is not and is done in a different notebook," + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from mspasspy.client import Client\n", + "mspass_client=Client()\n", + "dask_client = mspass_client.get_scheduler()\n", + "db = mspass_client.get_database(\"ANF48\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This notebook will work one year at a time. To run this for different years one need only change the year symbol in the next box." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "year=2014" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from earthscope_sdk import EarthScopeClient\n", + "from s3_worker_plugin import fetch_s3_client\n", + "\n", + "# Keep the SDK client alive for the refresh callback throughout indexing.\n", + "client = EarthScopeClient()\n", + "session = client.user.get_boto3_session()\n", + "s3_client = fetch_s3_client(session)\n", + "BUCKET = \"earthscope-mseed-res-na3mtd4fq5kz7pntcyr1uh46use2a--ol-s3\"\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This module contains special functions used to run this workflow and the related load_waveform workflows. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import s3segment_reader as geos3" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We need to get a list of networks that have data in the year we are processing. The entire list is unnecessarily large and would create a bit inefficiency below. There may be a more elegant way to do this with a mongodb aggregation but this is simpler with pymongo from my perspective." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from obspy import UTCDateTime\n", + "starttime = UTCDateTime(year=year,month=1,day=1)\n", + "endtime = UTCDateTime(year=year+1,month=1,day=1)\n", + "query=dict()\n", + "query['$and'] = [\n", + " {'Ptime' : {'$gte' : starttime.timestamp} },\n", + " {'Ptime' : {'$lt' : endtime.timestamp} }\n", + "]\n", + "n=db.arrival_css30.count_documents(query)\n", + "print(f\"Number of arrivals for {year}={n}\")\n", + "net_set = set()\n", + "cursor=db.arrival_css30.find(query)\n", + "for doc in cursor:\n", + " net=doc.get(\"net\")\n", + " if isinstance(net, str) and net:\n", + " net_set.add(net)\n", + "print(f\"Distinct networks in year {year}\")\n", + "netlist=list()\n", + "for net in net_set:\n", + " print(net)\n", + " netlist.append(net)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Include neighboring days needed by padded windows at year boundaries.\n", + "from datetime import date, timedelta\n", + "first_day = date(year, 1, 1) - timedelta(days=1)\n", + "last_day = date(year + 1, 1, 1)\n", + "index_days = []\n", + "day = first_day\n", + "while day <= last_day:\n", + " index_days.append((day.year, day.timetuple().tm_yday))\n", + " day += timedelta(days=1)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now use the functions we defined above to create a document for each s3 object with auxmd added. Save the resuls to MongoDB with the collection name \"wf_s3\". " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# always clear wf_s3 here to allow this to be run for a new year easily\n", + "db.drop_collection(\"wf_s3\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import time \n", + "t0=time.time()\n", + "prefix_list=list()\n", + "for net in netlist:\n", + " for index_year, day in index_days:\n", + " s3names,auxmd = geos3.fetch_net_day_list(s3_client,net,index_year,day,bucket=BUCKET)\n", + " doclist=geos3.objlist2doclist(s3names,auxmd)\n", + " if len(doclist)>0:\n", + " db.wf_s3.insert_many(doclist)\n", + " else:\n", + " print(f\"No data found in s3 for net={net} for day={day} of {year}\")\n", + "t=time.time()\n", + "n=db.wf_s3.count_documents({})\n", + "print(\"Total number of documents now in wf_s3=\",n)\n", + "print(\"Elapsed time to build index=\",t-t0)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Note that process is fast enough there is no need to add the complexity of running it in parallel. It looks like one could build an index for the entire TA data set in less than 30 minutes running serial. \n", + "\n", + "Before leaving, verify what a typical document looks like with the mspass pretty print function." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from mspasspy.util.seismic import print_metadata\n", + "doc=db.wf_s3.find_one()\n", + "print_metadata(doc)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For this algorithm this index will speed processing enormously." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from pymongo import ASCENDING\n", + "db.wf_s3.create_index([\n", + " (\"year\",ASCENDING),\n", + " (\"jday\",ASCENDING),\n", + "])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# important summary to print \n", + "ns3 = db.wf_s3.count_documents({})\n", + "nfailures = db.s3_read_failures.count_documents({})\n", + "print(\"Total number of wf_s3 documents created=\",ns3)\n", + "print(\"Current size of s3_read_failures collection (should be zero)=\",nfailures)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Comments relevant if you want to utilize this prototype for a larger data set assembly:\n", + "1. The wf_s3 collection can get very large very fast. Choose the bounds carefully.\n", + "2. Creating wf_s3 is fast. Don't create it until you need to for multiple reasons.\n", + "3. When data set is assembled (your version of load_waveforms_s3 completes) you can run `db.drop_collection(\"wf_s3\") to reduce storage. wf_s3 is more-or-less a scratch collection.\n", + "4. Don't run that last box until wf_s3 is fully populated. Each index you add to a collection slows write speed because the index has to be updated after the document is saved. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "# Close only after all indexing requests have finished.\n", + "s3_client.close()\n", + "client.close()\n" + ] + } + ] +} diff --git a/notebooks/EarthScopeS3/s3segment_reader.py b/notebooks/EarthScopeS3/s3segment_reader.py new file mode 100644 index 0000000..c00045b --- /dev/null +++ b/notebooks/EarthScopeS3/s3segment_reader.py @@ -0,0 +1,546 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +This module is something around verions 3 or 4 of a set of utility +functions for working with s3 on geolab. Their purpose is to standardize +componnts needed in workflows to extract waveform segments from the +continuous archives. + +Created on Sat Aug 22 06:04:50 2026 + +@author: pavlis +""" +import io +from contextlib import closing +import obspy +from obspy import UTCDateTime +import pandas as pd +import calendar +from mspasspy.util.db_utils import fetch_dbhandle +from mspasspy.ccore.seismic import TimeSeriesEnsemble +from mspasspy.util.converter import Stream2TimeSeriesEnsemble +from mspasspy.util.seismic import has_live_data +from mspasspy.algorithms.window import WindowData +from mspasspy.algorithms.signals import detrend +from mspasspy.ccore.utility import ErrorLogger, ErrorSeverity +from botocore.exceptions import ClientError +from s3_worker_plugin import fetch_s3_client + + +def fetch_net_day_list( + s3_client, + net, + year, + day, + bucket="earthscope-mseed-res-na3mtd4fq5kz7pntcyr1uh46use2a--ol-s3", + server_data_key="server_data", +) -> tuple: + """ + Fetch the raw list of object names from s3 for network net for year and julian day day. + Returns a tuple/array of length 2. Comp 0 is a list of s3 object names for that + net, year, day combination. Comp 1 is a dictionary returned by aws of the attribute + they call "HPPTHeader". The main use of that is to define the server and aws region + from which the data can be fetched. That is contant for this tutorial using Earthscope + data only, but it is useful to include it as access from multiple servers would need that data. + Default bucket argument is that for earthscope s3 miniseed archives. + """ + base_prefix = "miniseed" + prefix = f"{base_prefix}/{net}/{year}/{day:03d}/" + keys = set() + server_data = {} + # One ListObjectsV2 response is capped at 1,000 keys. Propagate listing + # failures rather than turning an authentication outage into an empty day. + for response in s3_client.get_paginator("list_objects_v2").paginate( + Bucket=bucket, Prefix=prefix, Delimiter="/" + ): + keys.update(obj["Key"].split("#")[0] for obj in response.get("Contents", [])) + server_data = response.get("ResponseMetadata", {}) + return [sorted(keys), {server_data_key: server_data}] + + +def objlist2doclist(objlist, auxmd=None): + """ + Splits s3 object names into dictionary entries to create a document to be saved in MongoDB. + Use dictionary like object *auxmd* to load additional, constant metadata to each doc returned. + """ + doclist = list() + for objname in objlist: + doc = dict() + doc["s3key"] = objname + s = objname.split("/") + doc["net"] = s[1] + yr = int(s[2]) + doc["year"] = yr + jday = int(s[3]) + doc["jday"] = jday + starttime = UTCDateTime(year=yr, julday=jday) + doc["starttime"] = starttime.timestamp + endtime = starttime.timestamp + 86400.0 + doc["endtime"] = endtime + s2 = s[4].split(".") + doc["sta"] = s2[0] + if auxmd is not None: + for k in auxmd.keys(): + # this maybe should warn if k overwrites an existing key + doc[k] = auxmd[k] + doclist.append(doc) + return doclist + + +def days_needed(t0, start, end, pad=100.0) -> list: + """ + Computes a range of days from a time interval t0+start to t0+tend. + Optional pad in seconds. That is, normally the function returns a list of + year-day pairs defined as the range t0+start-pad to t0+end+pad. + That is necessary because miniseed day files usually have inexact start times + due to compresssion. + + The idea of the function is that t0 is either a measured or theoretical + phase arrival time and start:end is the time window around that time + that we aim to extract. Downstream code has to take this output + and design an efficient read strategy for groups of arrivals + + Return as a list of (year,day) tuples. + """ + if end < start: + message = "days_needed: start time={} is less than end time={}".format( + start, end + ) + raise ValueError(message) + tsutc = UTCDateTime(t0 + start - pad) + teutc = UTCDateTime(t0 + end + pad) + result = list() + if tsutc.year == teutc.year: + yr = tsutc.year + for jday in range(tsutc.julday, teutc.julday + 1): + result.append([yr, jday]) + else: + # complexity required to cross year boundary + # first a sanity check + if (teutc.year - tsutc.year) != 1: + message = "days_needed: received irrational input.\n" + message += "Time window of {} -> {} spans multiple years\n".format( + tsutc, teutc + ) + message += "Irrational input as the volume of data to download is too large" + raise ValueError(message) + if calendar.isleap(tsutc.year): + days_in_year = 366 + else: + days_in_year = 365 + for jday in range(tsutc.julday, days_in_year + 1): + result.append([tsutc.year, jday]) + for jday in range(1, teutc.julday + 1): + result.append([teutc.year, jday]) + return result + + +def build_day_processing_blocks( + df, + start, + end, + dbname_or_handle, + pad=100.0, + channel_select="B*", + maxdays=2, + collection="wf_s3", + auxkeys=["source_id"], +) -> list: + """ + This function takes a (potentially large) DataFrame df of arrival + time data and returns a grouped list of documents that can be used to + efficiently extract waveform data segments containing those arrival + times. + + The data structure returned is a bit complex because it provides a + clean, if not obvious, way to work efficiently within constraints of + GeoLab and the current implementation of the MsPASS Database. + That is, the only scalable method of output on GeoLab at present is to + push data to an output s3 bucket on their scratch bucket. The only + we we have a present to do that is to serialize data with pickle and + push the bits to s3 objects. The idea of the output of this function + is to organize the processing to produce ensembles that are essentially + segmented day files. i.e. the groupings define all data that have a + starttime on a particular day. A subtle element of this approach is + that some events will be split across day volumes in the rare situation + where the set of arrival times for that event cross a day boundary. + + To handle that complexity the output is a list of lists. The + outer list of the return is of the form: + [data_list, error_list] + The content of "data_list" is itself a list of pairs (2 element array/list) + with conetent: + [yearday_string,document_list] + where "yearday_string" is a like it sound (e.g. 2004122 for julian + day 122 of 2004). document_list is a list of dictionaries that are used + downstream in parallel by a function called `get_s3_segments`. + + The complexity of the outer list component I called `error_list` above is + needed as a necessary evil as an error return. Because the input to + this function is normally huge the list of failures can also be huge. + Simply printing every failure can overwhelm the jupyter server when this + is run on GeoLab (the expectation). The list is a list of MongoDB + queries that yield no data. The caller should handle it appropriately. + The expectation at present is to save the content as a pickle file + in the run directory for later inspection. + + This function uses the DataFrame api in combination with pymongo + bo build the output. It assumes an index to the earthscope s3 + archive already existss for the time period to be handled. + That index is expected to live in a collection defined by the + "collction" argument (default is "wf_s3"). + + The input DataFrame must contain the following required attributes: + 'time' - either a theoretical or measured arrival time used as a basis for + what segments are generated. The functio aims to extract a + window of size time+start-pad to time+end+pad for each row of df. + 'net' - SEED station net code + 'sta' - SEED station code + + Optional: the keys defined by "auxkeys" are also expected to be in the + input DataFrame but will be silently ignored if they are not present. + The default expects "source_id" which is one way to provide a + cross reference to source data. Something of that sort is + normally required to make the results easier to use since + arrivals by definition are associated with some source. + The idea is this function is generic, however, and it does not + need to be dogmatic about source data. The main concept is to + extract event segments. Use data in auxkeys as metadata that define + what those segments are. + + A perspective on this function is it takes an input DataFrame with + the required metadata listed above and convert to an input that + can be used to implement processing by day groups. The day groups + provide a reasonably managable size package to push data as pickled + objects to an output s3 bucket. The primary content of each document + in the output structure is: + + 's3objects' - list of s3 object name to be retrieved as an atomic operation + Note when more than 1 reader needs to assume the pieces + need to be glued together. Currently the list can be no longer + than 2. + 'net' - SEED net code + 'sta' - SEED station code + 'channel_select' - echo of the channel_select argument + (used in s3 reader to select matching channel codes only) + 'arrivals' - list of dictionaries defining data to extract + 'start' - relative start time for windowing + 'end' - relative end time for windowing + + As noted any valid data match key defined by "auxdata" will also be + posted to those documents. + + Finally, note the 'arrivals' dictionaries each contain a 'time' attribute + that defines the arrival time from which a window is computed. + Any other data in tha dictionary is treated as auxiliary Metadata that + is to be copied to the output. Note downstream uses of that aux + data must handle it carefully to make sure the contents do not contain + metadata that conflict with metadata used inside children of + BasicTimeSeries data (e.g. starttime). Also note the start end + range will be modified from input by pad value. + + """ + db = fetch_dbhandle(dbname_or_handle) + df = df.sort_values("time") + # Add two columns to the DataFrame that are used below to generate + # dayfiles to load to handle which arrivals. + yrdaylists = list() + yrdaystr = list() # string version to be hashable - needed by groupby below + for row in df.itertuples(): + ydlist = days_needed(row.time, start, end, pad=pad) + if len(ydlist) > maxdays: + # this should only happen if the user makes an input error for start and/or end + message = ( + "Number of days for time={} and range={}-{} is too large\n".format( + row.time, start, end + ) + ) + message += ( + "Number of days this would try to return={} but maxdays={}".format( + len(ydlist), maxdays + ) + ) + raise ValueError(message) + # ydstr = f"{ydlist[0]}{ydlist[1]}" + # ydstr = str(ydlist) + ydstr = f"{ydlist[0][0]}{ydlist[0][1]}" + # print(ydlist[0][0]) + # print(ydlist[0][1]) + yrdaylists.append(ydlist) + yrdaystr.append(ydstr) + df["yrdaylist"] = yrdaylists + df["yearday"] = yrdaystr + + result = [] + fail_list = [] + for _, gdf in df.groupby("yearday", sort=False): + # The first window need not span all the days used by later windows. + required_days = sorted( + {tuple(day) for days in gdf["yrdaylist"] for day in days} + ) + query = {"$or": [{"year": yr, "jday": day} for yr, day in required_days]} + holdings = {} + for holding in db[collection].find(query): + holdings.setdefault((holding["net"], holding["sta"]), []).append(holding) + if not holdings: + fail_list.append(query) + continue + first_day = gdf["yrdaylist"].iloc[0][0] + jday_tag = f"{first_day[0]}_{first_day[1]}" + daylist = [] + for (net, sta), station_rows in gdf.groupby(["net", "sta"]): + objects = set() + arrivals, starts, ends = [], [], [] + available = holdings.get((net, sta), []) + for row in station_rows.to_dict("records"): + days = {tuple(day) for day in row["yrdaylist"]} + matching = [h for h in available if (h["year"], h["jday"]) in days] + if not matching: + fail_list.append( + { + "net": net, + "sta": sta, + "$or": [ + {"year": yr, "jday": day} for yr, day in sorted(days) + ], + } + ) + continue + # An arrival appears once, even when it needs two day objects. + # A DataFrame join with the holdings would multiply arrivals. + arrival = {"time": row["time"]} + arrival.update({key: row[key] for key in auxkeys if key in row}) + arrivals.append(arrival) + starts.append(row["time"] + start - pad) + ends.append(row["time"] + end + pad) + objects.update(h["s3key"] for h in matching) + if arrivals: + daylist.append( + { + "net": net, + "sta": sta, + "channel_select": channel_select, + "s3objects": sorted(objects), + "arrival": arrivals, + "start": starts, + "end": ends, + "jday_tag": jday_tag, + } + ) + if daylist: + result.append([jday_tag, daylist]) + return [result, fail_list] + + +def get_s3_segments( + doc, + session=None, + bucket="earthscope-mseed-res-na3mtd4fq5kz7pntcyr1uh46use2a--ol-s3", + detrend_type="None", + short_segment_handling="kill", + jday_tag_key="jday_tag", +) -> TimeSeriesEnsemble: + """ + Low level function to take content of doc, which is assumed to + be a componet of the output of run_by_days, and return a + TimeSeries ensemble of waveform segments the doc defines. + """ + alg = "get_s3_segments" + IMMUTABLE_METADATA = [ + "starttime", + "endtime", + "delta", + "npts", + "utc_convertible", + "time_standard", + ] + # defined here so if the strm is empty we return a default + # constructed ensemble with error log content + result = TimeSeriesEnsemble() + s3_client = None + # Always add the bucket name to the doc so it will be saved for both success and failure + doc["Bucket"] = bucket + # always put this in the ensemble metadata so there is always a way to identify a problem child + result["s3_request_document"] = doc + # Only serial calls create an owned client; a worker's plugin keeps its + # client open for subsequent tasks. + s3_client = fetch_s3_client() if session is None else fetch_s3_client(session) + try: + s3objects = doc["s3objects"] + if len(s3objects) == 0: + message = f"Request document had no entries in the list defined by the s3objects key" + result.elog.log_error(alg, message, ErrorSeverity.Invalid) + return result + strm = None + for s3key in s3objects: + # Missing/undecodable data return None; infrastructure errors raise. + # print("Debug: trying to fetch s3key=",s3key) + stread, elog = get_object_with_iris_versions( + s3_client, s3key, bucket=bucket + ) + if stread is None: + result.elog = elog + message = f"No data retrieved for request with s3object={s3objects}" + result.elog.log_error(alg, message, ErrorSeverity.Invalid) + # temporary for debugging + # print(bad_result.elog.get_error_log()) + return result + else: + if strm is None: + strm = stread + else: + strm += stread + del stread + + try: + channel_select = doc["channel_select"] + if len(channel_select) > 0: + strm = strm.select(channel=channel_select) + # print("stream of size after select=",len(strm)) + if len(strm) > 0: + # This handles a rare data problem where the same seed net code + # will have multiple sample rate data. Happens because SEED + # only specifies a range for a a channel code + rates = set(tr.stats.sampling_rate for tr in strm) + # strm.sort() + if len(rates) == 1: + strm.merge() + alldata = Stream2TimeSeriesEnsemble(strm) + del strm + else: + # should not need to handle case with len 0 + # here we split and merge pieces + merged_streams = obspy.Stream() + for rate in rates: + sub_st = strm.select(sampling_rate=rate).copy() + sub_st.merge(method=1) + merged_streams += sub_st + del sub_st + alldata = Stream2TimeSeriesEnsemble(merged_streams) + del strm + del merged_streams + + # elog contains error messages posted by get_object_with_iris_versions + # always save them to sort out data store problems by sifting through elogs + result.elog += elog + if detrend_type not in (None, "None"): + alldata = detrend(alldata, type=detrend_type) + # workaround for bug github issue 710 + if alldata.dead(): + if has_live_data(alldata): + alldata.set_live() + else: + # print("Ensemble with ",len(alldata.member)," members has no data marked live - returning empty result") + result.elog.log_error( + alg, + "converted TimeSeriesEnsemble following detrend has no live data", + ErrorSeverity.Invalid, + ) + return result + stlist = doc["start"] + etlist = doc["end"] + arrival_auxdata = doc["arrival"] + # We assume the three items indexed here have the same length + for i in range(len(stlist)): + stime = stlist[i] + etime = etlist[i] + arrival_doc = arrival_auxdata[i] + # print("Running window data") + ens = WindowData( + alldata, + stime, + etime, + short_segment_handling=short_segment_handling, + ) + # print("window output state=",ens.live) + # print("Size of output=",len(ens.member)) + for i in range(len(ens.member)): + d = ens.member[i] + for k in arrival_doc: + if k not in IMMUTABLE_METADATA: + if k == "time": + d["arrival_time"] = arrival_doc["time"] + else: + d[k] = arrival_doc[k] + ens.member[i] = d + # print("Size ens after metadata copy=",len(ens.member)) + for d in ens.member: + result.member.append(d) + # print("Size of result = ",len(result.member)) + if has_live_data(result) > 0: + # print("Setting result live") + result.set_live() + if jday_tag_key in doc: + result[jday_tag_key] = doc[jday_tag_key] + else: + message = f"No value found for key={jday_tag_key} - may cause problems downstream" + result.elog.log_error(alg, message, ErrorSeverity.Complaint) + else: + result.kill() # not necessary but sure kill make this more robust + message = "No live data extracted by windowing loop of this code" + result.elog.log_error(alg, message, ErrorSeverity.Invalid) + del alldata + else: + result.elog = elog + message = f"No data satisfied Stream.select({channel_select})" + result.elog.log_error(alg, message, ErrorSeverity.Invalid) + except MemoryError: + raise + except Exception as e: + message = ( + "get_s3_segments failed with when the following exception was thrown:\n" + ) + message += str(e) + result.elog.log_error(alg, message, ErrorSeverity.Invalid) + return result + finally: + if session is not None: + s3_client.close() + + +def get_object_with_iris_versions( + s3_client, + base_name, + bucket="earthscope-mseed-res-na3mtd4fq5kz7pntcyr1uh46use2a--ol-s3", + maxversion=10, +): + """Read a base object or its legacy #N version. + + Try the base name, then #maxversion down to #1. Only a definite missing + object (404/NoSuchKey/NotFound) permits fallback. HEAD's generic 403 cannot + establish whether an object exists. Authentication, transport, body-read, + and other infrastructure failures retain their original exception and + stop the workflow. Missing/undecodable data return [None, elog]. + """ + elog = ErrorLogger() + alg = "get_object_with_iris_versions" + names = [base_name] + [ + f"{base_name}#{version}" for version in range(maxversion, 0, -1) + ] + for name in names: + try: + s3_client.head_object(Bucket=bucket, Key=name) + except ClientError as error: + if error.response["Error"]["Code"] in ("404", "NoSuchKey", "NotFound"): + continue + raise + response = s3_client.get_object(Bucket=bucket, Key=name) + with closing(response["Body"]) as body: + data = body.read() + try: + with io.BytesIO(data) as buffer: + stream = obspy.read(buffer, format="mseed") + return [stream, elog] + except Exception as error: + elog.log_error( + alg, + f"Cannot decode miniSEED object {name}: {error}", + ErrorSeverity.Invalid, + ) + return [None, elog] + elog.log_error( + alg, + f"No valid version of object {base_name} was found", + ErrorSeverity.Invalid, + ) + return [None, elog] diff --git a/notebooks/EarthScopeS3/test_notebooks.py b/notebooks/EarthScopeS3/test_notebooks.py new file mode 100644 index 0000000..9877259 --- /dev/null +++ b/notebooks/EarthScopeS3/test_notebooks.py @@ -0,0 +1,46 @@ +import ast +import json +from pathlib import Path + +import pytest + +ROOT = Path(__file__).parent + + +@pytest.mark.parametrize("path", sorted(ROOT.glob("*.ipynb")), ids=lambda p: p.name) +def test_notebooks_are_output_free_and_valid_python(path): + notebook = json.loads(path.read_text()) + assert notebook["nbformat"] == 4 + for index, cell in enumerate(notebook["cells"]): + assert not cell["metadata"] + if cell["cell_type"] == "code": + assert cell["execution_count"] is None + assert cell["outputs"] == [] + source = "".join(cell["source"]) + ast.parse(source, filename=f"{path.name}:cell{index}") + assert "get_frozen_credentials" not in source + assert "AWS_SECRET_ACCESS_KEY" not in source + assert "AWS_SESSION_TOKEN" not in source + + +def test_index_includes_adjacent_days_and_uses_half_open_year(): + notebook = json.loads((ROOT / "s3indexing_2014.ipynb").read_text()) + sources = [ + "".join(c["source"]) for c in notebook["cells"] if c["cell_type"] == "code" + ] + index_source = next(s for s in sources if "first_day =" in s) + for year, previous_last_day in ((2013, 366), (2014, 365)): + namespace = {"year": year} + exec(index_source, namespace) + assert namespace["index_days"][0] == (year - 1, previous_last_day) + assert namespace["index_days"][-1] == (year + 1, 1) + query = next(s for s in sources if "query['$and']" in s) + assert "'$lt' : endtime.timestamp" in query + + +def test_loader_preserves_unmatched_station_and_filters_network_by_value(): + notebook = json.loads((ROOT / "load_anfdata.ipynb").read_text()) + source = "\n".join("".join(c["source"]) for c in notebook["cells"]) + assert 'doc["sta"] = css_sta' in source + assert 'net != "AK"' in source + assert "del netlist[0]" not in source diff --git a/notebooks/EarthScopeS3/test_plugin_process.py b/notebooks/EarthScopeS3/test_plugin_process.py new file mode 100644 index 0000000..2c0f587 --- /dev/null +++ b/notebooks/EarthScopeS3/test_plugin_process.py @@ -0,0 +1,27 @@ +import os +from pathlib import Path +import signal +import subprocess +import sys + + +def test_plugin_in_separate_worker_process(): + process = subprocess.Popen( + [sys.executable, str(Path(__file__).with_name("plugin_process_probe.py"))], + start_new_session=True, + ) + try: + assert process.wait(timeout=45) == 0 + finally: + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + pass + try: + process.wait(timeout=5) + finally: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait() diff --git a/notebooks/EarthScopeS3/test_s3_worker_plugin.py b/notebooks/EarthScopeS3/test_s3_worker_plugin.py new file mode 100644 index 0000000..c86bfde --- /dev/null +++ b/notebooks/EarthScopeS3/test_s3_worker_plugin.py @@ -0,0 +1,212 @@ +"""Offline tests; no EarthScope login, AWS requests, or worker processes. + +Use the real EarthScope SDK session builder and real botocore credentials, +replacing only the credential service with a deterministic cache and clock. +""" + +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from earthscope_sdk.client.user._service import UserService +from earthscope_sdk.client.user.models import AwsTemporaryCredentials + +import s3_worker_plugin as plugin + + +class CachedCredentialService: + def __init__(self, initial_ttl=60): + self.now = datetime(2026, 9, 9, tzinfo=timezone.utc) + self.expiry = self.now + timedelta(minutes=initial_ttl) + self.generation = 0 + self.failure = None + self.closed = False + self.calls = 0 + + def get_aws_credentials(self, *, role): + assert role == "s3-miniseed" # Do not change the existing archive role. + assert not self.closed + self.calls += 1 + if self.failure: + raise self.failure + if self.expiry - self.now <= timedelta(seconds=30): + self.generation += 1 + self.expiry = self.now + timedelta(minutes=60) + return AwsTemporaryCredentials( + aws_access_key_id=f"FAKE_ACCESS_{self.generation}", + aws_secret_access_key="NOT_A_REAL_SECRET", + aws_session_token=f"FAKE_TOKEN_{self.generation}", + expiration=self.expiry, + ) + + +@pytest.fixture +def harness(monkeypatch): + # Do not consult the operator's AWS configuration or instance metadata. + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + monkeypatch.setenv("AWS_DEFAULT_REGION", "us-west-2") + monkeypatch.setenv("AWS_CONFIG_FILE", "/dev/null") + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", "/dev/null") + + def make(initial_ttl=60): + cache = CachedCredentialService(initial_ttl) + # No SdkContext: its network/credential endpoint is the one mocked seam. + service = object.__new__(UserService) + service.get_aws_credentials = cache.get_aws_credentials + events = [] + sessions = [] + + def session_factory(): + session = service.get_boto3_session() + credentials = session.get_credentials() + credentials._time_fetcher = lambda: cache.now + sessions.append(session) + return session + + def close_sdk(): + cache.closed = True + events.append("sdk") + + sdk = SimpleNamespace( + user=SimpleNamespace(get_boto3_session=Mock(side_effect=session_factory)), + close=Mock(side_effect=close_sdk), + ) + factory = Mock(return_value=sdk) + monkeypatch.setattr(plugin, "EarthScopeClient", factory) + return SimpleNamespace( + cache=cache, + sdk=sdk, + sessions=sessions, + events=events, + factory=factory, + worker=SimpleNamespace(data={}, plugins={}), + ) + + return make + + +def test_cached_credentials_keep_real_expiry_and_refresh_past_one_hour(harness): + h = harness() + worker_plugin = plugin.S3Worker() + worker_plugin.setup(h.worker) + try: + client = worker_plugin.s3_client + credentials = client._request_signer._credentials + # No second credential provider or frozen snapshot may replace the SDK's. + assert credentials is h.sessions[0].get_credentials() + assert credentials._expiry_time == h.cache.expiry + assert not h.cache.closed + start = h.cache.now + first = credentials.get_frozen_credentials().access_key + + for minute in (41, 55): + h.cache.now = start + timedelta(minutes=minute) + assert credentials.get_frozen_credentials().access_key == first + assert credentials._expiry_time == start + timedelta(minutes=60) + + h.cache.now = start + timedelta(minutes=61) + assert credentials.get_frozen_credentials().access_key != first + assert credentials._expiry_time == h.cache.expiry > h.cache.now + assert h.cache.generation == 1 + h.cache.now = start + timedelta(minutes=122) + credentials.get_frozen_credentials() + assert h.cache.generation == 2 + assert credentials._expiry_time == h.cache.expiry > h.cache.now + assert h.factory.call_count == 1 + assert h.sdk.user.get_boto3_session.call_count == 1 + assert client.meta.config.request_checksum_calculation == "when_required" + assert client.meta.config.response_checksum_validation == "when_required" + finally: + worker_plugin.teardown(h.worker) + + +def test_cached_credentials_already_near_expiry_are_not_extended(harness): + h = harness(initial_ttl=8) + worker_plugin = plugin.S3Worker() + worker_plugin.setup(h.worker) + try: + credentials = worker_plugin.s3_client._request_signer._credentials + start = h.cache.now + assert credentials._expiry_time == start + timedelta(minutes=8) + first = credentials.get_frozen_credentials().access_key + h.cache.now = start + timedelta(minutes=9) + assert credentials.get_frozen_credentials().access_key != first + assert credentials._expiry_time == h.cache.expiry > h.cache.now + finally: + worker_plugin.teardown(h.worker) + + +def test_mandatory_refresh_failure_is_not_hidden(harness): + h = harness() + worker_plugin = plugin.S3Worker() + worker_plugin.setup(h.worker) + try: + credentials = worker_plugin.s3_client._request_signer._credentials + h.cache.now += timedelta(minutes=61) + error = RuntimeError("simulated credential endpoint failure") + h.cache.failure = error + with pytest.raises(RuntimeError) as raised: + credentials.get_frozen_credentials() + assert raised.value is error + h.cache.failure = None + credentials.get_frozen_credentials() + assert credentials._expiry_time == h.cache.expiry > h.cache.now + finally: + worker_plugin.teardown(h.worker) + + +@pytest.mark.parametrize("failure_stage", ["session", "client"]) +def test_failed_setup_closes_resources(harness, monkeypatch, failure_stage): + h = harness() + error = RuntimeError(f"failed {failure_stage}") + if failure_stage == "session": + h.sdk.user.get_boto3_session.side_effect = error + elif failure_stage == "client": + monkeypatch.setattr(plugin, "fetch_s3_client", Mock(side_effect=error)) + worker_plugin = plugin.S3Worker() + with pytest.raises(RuntimeError) as raised: + worker_plugin.setup(h.worker) + assert raised.value is error + h.sdk.close.assert_called_once_with() + assert "s3client" not in h.worker.data + worker_plugin.teardown(h.worker) + h.sdk.close.assert_called_once_with() + + +def test_custom_key_fetch_and_idempotent_teardown(harness, monkeypatch): + h = harness() + h.worker.data["unrelated"] = 42 + worker_plugin = plugin.S3Worker(key="archive") + real_fetch = plugin.fetch_s3_client + + def instrument_client(session): + client = real_fetch(session) + real_close = client.close + + def close_client(): + h.events.append("s3") + real_close() + + client.close = Mock(side_effect=close_client) + return client + + with monkeypatch.context() as setup_patch: + setup_patch.setattr(plugin, "fetch_s3_client", instrument_client) + worker_plugin.setup(h.worker) + client = worker_plugin.s3_client + h.worker.plugins["archive"] = worker_plugin + monkeypatch.setattr(plugin, "get_worker", lambda: h.worker) + assert plugin.fetch_s3_client(worker_data_key="archive") is client + worker_plugin.teardown(h.worker) + worker_plugin.teardown(h.worker) + assert h.worker.data == {"unrelated": 42} + assert h.events == ["s3", "sdk"] + client.close.assert_called_once_with() + h.sdk.close.assert_called_once_with() + + +def test_teardown_before_setup_is_safe(): + worker = SimpleNamespace(data={}, plugins={}) + plugin.S3Worker().teardown(worker) + assert worker.data == {} diff --git a/notebooks/EarthScopeS3/test_workflow.py b/notebooks/EarthScopeS3/test_workflow.py new file mode 100644 index 0000000..4ab2d41 --- /dev/null +++ b/notebooks/EarthScopeS3/test_workflow.py @@ -0,0 +1,353 @@ +"""Native waveform tests with fake network/database boundaries, not live S3.""" + +import io +import pickle +import weakref +from contextlib import contextmanager +from types import SimpleNamespace + +import numpy as np +import obspy +import pandas as pd +import pytest +from botocore.exceptions import ClientError +from mspasspy.ccore.seismic import TimeSeries, TimeSeriesEnsemble +from mspasspy.ccore.utility import ErrorSeverity + +import s3_day_workflow as workflow +import s3segment_reader as reader + + +def aws_error(code, operation="HeadObject"): + return ClientError( + { + "Error": {"Code": code, "Message": "test failure"}, + "ResponseMetadata": {"HTTPStatusCode": 403, "RequestId": "TEST_REQUEST"}, + }, + operation, + ) + + +class FakeS3: + def __init__(self, payload=b"", missing=(), error=None): + self.payload, self.missing, self.error = payload, missing, error + self.heads = [] + self.body = None + + def head_object(self, *, Bucket, Key): + self.heads.append(Key) + if self.error: + raise self.error + if Key in self.missing: + raise aws_error("404") + + def get_object(self, **kwargs): + self.body = io.BytesIO(self.payload) + return {"Body": self.body} + + +def miniseed(rates=(10,)): + stream = obspy.Stream() + for index, rate in enumerate(rates): + stream += obspy.Trace( + np.arange(100, dtype=np.int32), + header={ + "network": "TA", + "station": "ABC", + "channel": "BHZ", + "location": f"{index:02d}", + "sampling_rate": rate, + "starttime": obspy.UTCDateTime(2014, 1, 1), + }, + ) + buffer = io.BytesIO() + stream.write(buffer, format="MSEED") + return buffer.getvalue() + + +def request_document(): + start = obspy.UTCDateTime(2014, 1, 1).timestamp + return { + "s3objects": ["day"], + "channel_select": "B*", + "start": [start + 1], + "end": [start + 2], + "arrival": [{"time": start + 1, "arid": 7}], + "jday_tag": "2014_1", + } + + +@pytest.mark.parametrize("code", ["403", "AccessDenied", "ExpiredToken", "SlowDown"]) +def test_authentication_and_service_errors_do_not_try_other_versions(code, monkeypatch): + error = aws_error(code) + client = FakeS3(error=error) + monkeypatch.setattr(reader, "fetch_s3_client", lambda: client) + with pytest.raises(ClientError) as raised: + reader.get_s3_segments(request_document()) + assert raised.value is error + assert client.heads == ["day"] + + +def test_refresh_runtime_error_is_not_converted_to_dead_data(monkeypatch): + error = RuntimeError("credential refresh failed") + client = FakeS3(error=error) + monkeypatch.setattr(reader, "fetch_s3_client", lambda: client) + with pytest.raises(RuntimeError) as raised: + reader.get_s3_segments(request_document()) + assert raised.value is error + + +def test_native_memory_error_stops_processing(monkeypatch): + error = MemoryError("test allocation failure") + client = FakeS3(miniseed()) + monkeypatch.setattr(reader, "fetch_s3_client", lambda: client) + + def fail(stream): + raise error + + monkeypatch.setattr(reader, "Stream2TimeSeriesEnsemble", fail) + with pytest.raises(MemoryError) as raised: + reader.get_s3_segments(request_document()) + assert raised.value is error + + +@pytest.mark.parametrize("fails", [False, True]) +def test_serial_reader_closes_owned_client(fails, monkeypatch): + error = aws_error("403") if fails else None + client = FakeS3(miniseed(), error=error) + closed = [] + client.close = lambda: closed.append(True) + session = object() + + def fetch(actual_session): + assert actual_session is session + return client + + monkeypatch.setattr(reader, "fetch_s3_client", fetch) + if fails: + with pytest.raises(ClientError) as raised: + reader.get_s3_segments(request_document(), session=session) + assert raised.value is error + else: + assert reader.get_s3_segments(request_document(), session=session).live + assert closed == [True] + + +def test_missing_base_falls_back_and_closes_body(): + client = FakeS3(miniseed(), missing={"day", "day#3"}) + stream, errors = reader.get_object_with_iris_versions(client, "day", maxversion=3) + assert client.heads == ["day", "day#3", "day#2"] + assert len(stream) == 1 + assert client.body.closed + + +def test_all_versions_missing_and_invalid_miniseed_are_data_failures(): + client = FakeS3(missing={"day", "day#1"}) + stream, errors = reader.get_object_with_iris_versions(client, "day", maxversion=1) + assert stream is None + assert errors.size() > 0 + invalid = FakeS3(b"not miniSEED") + stream, errors = reader.get_object_with_iris_versions(invalid, "day") + assert stream is None + assert invalid.body.closed + + +def test_body_read_error_propagates_and_closes(): + error = OSError("truncated S3 response") + + class FailedBody(io.BytesIO): + def read(self, *args): + raise error + + body = FailedBody() + client = FakeS3() + client.get_object = lambda **kwargs: {"Body": body} + with pytest.raises(OSError) as raised: + reader.get_object_with_iris_versions(client, "day") + assert raised.value is error + assert body.closed + + +def test_get_error_propagates_without_decode_or_version_fallback(): + client = FakeS3() + error = aws_error("ExpiredToken", "GetObject") + + def fail(**kwargs): + raise error + + client.get_object = fail + with pytest.raises(ClientError) as raised: + reader.get_object_with_iris_versions(client, "day") + assert raised.value is error + assert client.heads == ["day"] + + +@pytest.mark.parametrize("rates", [(10,), (10, 20)]) +def test_real_miniseed_decode_detrend_and_window(rates, monkeypatch): + client = FakeS3(miniseed(rates)) + monkeypatch.setattr(reader, "fetch_s3_client", lambda: client) + converted = [] + converter = reader.Stream2TimeSeriesEnsemble + + def convert(stream): + converted.append(len(stream)) + return converter(stream) + + monkeypatch.setattr(reader, "Stream2TimeSeriesEnsemble", convert) + result = reader.get_s3_segments(request_document(), detrend_type="simple") + assert result.live, str(result.elog) + assert converted == [len(rates)] + assert len(result.member) == len(rates) + for member in result.member: + assert member["arid"] == 7 + np.testing.assert_allclose(np.asarray(member.data), 0, atol=1e-10) + + +def test_paginated_index_deduplicates_versioned_names(): + def pages(**kwargs): + assert kwargs["Prefix"] == "miniseed/TA/2014/001/" + yield {"Contents": [{"Key": f"object{i}"} for i in range(1000)]} + yield {"Contents": [{"Key": "object1000"}, {"Key": "object0#2"}]} + + client = SimpleNamespace(get_paginator=lambda name: SimpleNamespace(paginate=pages)) + names, metadata = reader.fetch_net_day_list(client, "TA", 2014, 1) + assert len(names) == 1001 + assert names == sorted(set(names)) + assert metadata == {"server_data": {}} + + +def test_index_error_propagates(): + error = aws_error("ExpiredToken", "ListObjectsV2") + + def pages(**kwargs): + raise error + + client = SimpleNamespace(get_paginator=lambda name: SimpleNamespace(paginate=pages)) + with pytest.raises(ClientError) as raised: + reader.fetch_net_day_list(client, "TA", 2014, 1) + assert raised.value is error + + +def test_day_group_uses_union_without_multiplying_arrivals(monkeypatch): + beginning = obspy.UTCDateTime(2014, 12, 31).timestamp + frame = pd.DataFrame( + [ + {"time": beginning + 100, "net": "TA", "sta": "ABC", "arid": 1}, + {"time": beginning + 86390, "net": "TA", "sta": "ABC", "arid": 2}, + ] + ) + queries = [] + + def find(query): + queries.append(query) + return [ + {"year": yr, "jday": day, "net": "TA", "sta": "ABC", "s3key": name} + for yr, day, name in [(2014, 365, "a"), (2015, 1, "b")] + ] + + db = {"wf_s3": SimpleNamespace(find=find)} + monkeypatch.setattr(reader, "fetch_dbhandle", lambda handle: db) + days, errors = reader.build_day_processing_blocks( + frame, -10, 20, db, pad=0, auxkeys=["arid"] + ) + assert errors == [] + assert len(days) == 1 + assert queries == [ + {"$or": [{"year": 2014, "jday": 365}, {"year": 2015, "jday": 1}]} + ] + doc = days[0][1][0] + assert doc["s3objects"] == ["a", "b"] + assert [a["arid"] for a in doc["arrival"]] == [1, 2] + assert len(doc["start"]) == len(doc["end"]) == 2 + + +def make_ensemble(value=1, live=True): + ens = TimeSeriesEnsemble() + if live: + trace = TimeSeries(8) + trace.set_live() + trace["arid"] = value + trace.data[0] = value + ens.member.append(trace) + ens.set_live() + else: + ens.elog.log_error("read", "missing test data", ErrorSeverity.Invalid) + ens["jday_tag"] = "2014_1" + return ens + + +def test_completion_releases_sources_before_pickle_and_preserves_output(monkeypatch): + inputs = [make_ensemble(1), make_ensemble(2)] + references = [weakref.ref(ens) for ens in inputs] + saved = {} + + @contextmanager + def open_file(key, mode): + assert mode == "wb" + assert inputs == [None, None] + assert all(reference() is None for reference in references) + buffer = io.BytesIO() + yield buffer + saved[key] = buffer.getvalue() + + monkeypatch.setattr(workflow, "fetch_dbhandle", lambda handle: {}) + assert ( + workflow.save_jday_outputs( + inputs, + None, + SimpleNamespace(open=open_file), + "s3://scratch/user-prefix/", + 2014, + verbose=True, + ) + is True + ) + assert list(saved) == ["s3://scratch/user-prefix/2014/2014_1.pickle"] + result = pickle.loads(next(iter(saved.values()))) + assert isinstance(result, TimeSeriesEnsemble) + assert result.live and result["jday_tag"] == "2014_1" + assert [t["arid"] for t in result.member] == [1, 2] + assert [t.data[0] for t in result.member] == [1, 2] + + +@pytest.mark.parametrize("failure_stage", ["open", "dump", "close"]) +def test_write_errors_escape_completion(failure_stage, monkeypatch): + error = OSError(f"test {failure_stage} error") + + @contextmanager + def open_file(key, mode): + if failure_stage == "open": + raise error + yield io.BytesIO() + if failure_stage == "close": + raise error + + def fail_dump(*args): + raise error + + if failure_stage == "dump": + monkeypatch.setattr(workflow.pickle, "dump", fail_dump) + monkeypatch.setattr(workflow, "fetch_dbhandle", lambda handle: {}) + with pytest.raises(OSError) as raised: + workflow.save_jday_outputs( + [make_ensemble()], + None, + SimpleNamespace(open=open_file), + "s3://scratch/u", + 2014, + ) + assert raised.value is error + + +def test_dead_and_empty_days_record_failures_without_writing(monkeypatch): + documents = [] + db = {"s3_read_failures": SimpleNamespace(insert_many=documents.extend)} + monkeypatch.setattr(workflow, "fetch_dbhandle", lambda handle: db) + assert not workflow.save_jday_outputs( + [make_ensemble(live=False)], db, None, "scratch/u", 2014 + ) + assert len(documents) == 1 + assert documents[0]["elog_content"] + assert not workflow.save_jday_outputs([], db, None, "scratch/u", 2014) + assert workflow.count_saved(None, True) == 1 + assert workflow.count_saved(1, False) == 1