diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..0203922 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,38 @@ +name: tests + +on: + push: + branches: [main] + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.12", "3.13"] + + defaults: + run: + # Login shell so the micromamba environment stays activated across steps. + shell: bash -el {0} + + steps: + - uses: actions/checkout@v4 + + # ecflow is not on PyPI; it is installed from conda-forge via the + # environment file. river-route and the rest come along with it. + - name: Create environment + uses: mamba-org/setup-micromamba@v2 + with: + environment-file: environment.yml + create-args: >- + python=${{ matrix.python-version }} + cache-environment: true + + - name: Install package with dev dependencies + run: pip install -e ".[dev]" + + - name: Run tests + run: pytest -q --cov=geoglows_ecflow.resources --cov-report=term-missing diff --git a/PLAN.md b/PLAN.md new file mode 100644 index 0000000..943ace6 --- /dev/null +++ b/PLAN.md @@ -0,0 +1,116 @@ +# Workflow Simplification Plan + +A living plan for simplifying the `geoglows_ecflow` workflow. Branched off the +`rapid-to-river-route` work (PR #27), so all references below assume the +river-route codebase, not the RAPID `main`. + +## Goals + +1. **Reduce complexity** — eliminate duplication, simplify functions, separate + inputs from logic. +2. **Make implicit explicit** — name constants and configuration, split mixed + functions, document non-obvious behavior. +3. **Improve maintainability** — make different configurations easy to run and + add a unit-test safety net. + +## Scope + +- **In scope:** `geoglows_ecflow/resources/*.py`, + `geoglows_ecflow/workflow/builders/builder.py`, + `geoglows_ecflow/workflow/parts/*`, and the `.ecf` task scripts. +- **Mostly frozen:** `geoglows_ecflow/workflow/comfies/*` — vendored ECMWF + framework code (Apache 2.0, ~4,700 lines). Left untouched except for the + minimal Python-3.12+ compatibility fixes noted below (the suite could not be + imported at all without them). + +## Decisions + +- **Tests:** unit tests on the pure functions, plus in-memory + suite-definition smoke tests (these were originally deferred but pulled + forward to guard the builder refactor). +- **HRES member:** high-resolution is always ensemble member **52**, kept as a + single named constant `HRES_ENSEMBLE_MEMBER` rather than a configurable value. + +--- + +## Phase 1 — Pure cleanup (no behavior change) — DONE + +- [x] `builder.py`: fix `Task("dimmy")` typo. +- [x] `helper_functions.py`: fix `create_logger` so file logging attaches. +- [x] `builder.py`: remove duplicate imports and the duplicated `nodes` import. +- [x] `builder.py`: remove read-but-unused config vars. +- [x] `builder.py`: fix stale docstring ("GLOFAS suite" → GEOGloWS). +- [x] `generate_esri_table.py`: fix `int or str` type hint. +- [x] Standardize the `argparse(nargs=1)` + `args.x[0]` antipattern. + +## Phase 2 — Test harness + CI — DONE + +- [x] Add `pytest` (+ `pytest-cov`) as a dev dependency. +- [x] Create `tests/` with fixtures. +- [x] Unit tests for the pure functions (ensemble parsing, VPU listing, date + parsing, state-init lookback, forecast preprocess, return-period / + exceedance, init-flow time index, zarr round-trip). +- [x] GitHub Actions workflow (`.github/workflows/tests.yml`). Now builds a + conda env from `environment.yml` (ecflow has no PyPI wheel) on a + Python **3.12–3.13** matrix and runs the whole suite. + +## Phase 3 — Centralize duplication — DONE + +- [x] Shared zarr-writing helper (`resources/zarr_io.py`). +- [x] `helper_functions.load_forecast_run` loader. +- [x] `RETURN_PERIODS` defined once and used for the ladders. +- [x] Standardized `logging` setup (`configure_logging`). + +## Phase 4 — Make implicit explicit — MOSTLY DONE + +- [x] `HRES_ENSEMBLE_MEMBER = 52` — used in `netcdf_to_zarr.py` and the + `range(1, HRES_ENSEMBLE_MEMBER + 1)` ensemble loop in `builder.py`. +- [ ] `nco_calc.ecf` `grep -v ..._52.nc` (×3) — **decision pending**: wire an + ecflow `%HRES_MEMBER%` variable vs. leave `52` + a comment. Highest risk: + wrong wiring silently changes which member is excluded from the mean. +- [x] `EMOS_BASE != "12"` gate (×3) → `is_00z_cycle()` helper. +- [x] Magic numbers named: thickness ladder (`THICKNESS_THRESHOLDS`), + stream-order (`MIN_STREAM_ORDER`), 10-day window (`FORECAST_WINDOW_DAYS`), + `MEM` values (`ENS_TASK_MEM_MB` / `ARCHIVE_QINIT_MEM_MB`). +- [x] Timer offsets (`hours=7`/`hours=9`/`"14:15"`) — **resolved by deletion**: + they lived only in the broken `rd`/research-mode branch, which has been + removed (see follow-ups). Nothing to extract; revisit if `rd` returns. +- [x] Consolidate `self.config.get(...)` reads in `builder.py` into one + documented block. + +## Follow-ups (later tasks) + +- ecFlow-**server** tests (building the def against a live server / `--dry`); + the in-memory structural tests are done, this is the heavier version. +- README refresh (carried over from PR #27 review). +- **comfies is incompatible with ecflow 5.17+.** Its node wrappers set + `Variable.parent` (`ooflow.py:1925`), which ecflow 5.17 made a read-only + built-in, so building any suite raises `AttributeError`. Worked around by + pinning `ecflow<5.17` in `environment.yml`; the real fix is to rename + comfies' parent-tracking attribute so it no longer collides. +- **Research (`rd`) mode removed.** `mode='rd'` was the only path with + `follow_osuite=False` and crashed unconditionally at `barrier_hh.ymd` + (`barrier_hh` is a `NominalTime`, which has no `ymd`) — broken since the + original 2024-07-25 authoring, so never usable. Removed the `rd` choice, the + `follow_osuite`/`in_production`/`in_test` flags, the non-`follow_osuite` + branch (the `+7h`/`+9h`/`14:15` run timers), and the dead crash line. `prod` + and `test` are the remaining modes. If research mode is wanted again, it + should be reintroduced correctly (with the intended barrier-repeat wiring). + +--- + +## Current status + +**Branch `workflow-simplification`** (fork `JakeGimenes`), open as a PR against +`rapid-to-river-route`. Phases 1–3 complete; Phase 4 complete except the +`nco_calc.ecf` decision. The timer constants are resolved by deletion (the +`rd`/research mode that owned them has been removed). **33 pytest tests pass** — +the resources tests run anywhere; the suite-definition tests require `ecflow` +(conda-forge). + +The vendored `comfies` framework got the minimum Python-3.12+ compatibility +fixes needed to import it at all (`imp` → `importlib`, `pkg_resources` → +`packaging`); everything else in `comfies/*` is unchanged. + +**Remaining actionable work:** the `nco_calc.ecf` `HRES_MEMBER` decision and the +README refresh. diff --git a/README.md b/README.md index 84d5a93..f9c8722 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,7 @@ pip install -e . ## Non-Python Dependencies -- ecflow>=5.11.3 +- ecflow>=5.11.3,<5.17 - nco>=5.1.8 - ksh>=2020.0.0 @@ -32,7 +32,7 @@ pip install -e . mars_bond_id='251' staticdata = '/path/to/assets' workroot = f'/path/to/workroot' - mode = 'test' # suite mode ('rd':research, 'test':test, 'prod':production) + mode = 'test' # suite mode ('test':test, 'prod':production) expver = 'geoglows' exparch = '/path/to/archive' iniexparch = '/path/to/init_archive' @@ -80,7 +80,7 @@ pip install -e . ) # -------------------------------------------- - # Configuration of EFAS software packages + # Configuration of GEOGloWS software packages # which are installed together with the suite. # -------------------------------------------- packages = dict( diff --git a/environment.yml b/environment.yml index ce56eb3..7710e78 100644 --- a/environment.yml +++ b/environment.yml @@ -3,9 +3,11 @@ channels: - conda-forge - defaults dependencies: - - python>=3.11,<3.14 + - python>=3.12,<3.14 - pip - - ecflow + # comfies' node wrappers set Variable.parent, which ecflow 5.17 made + # read-only; pin below that until comfies is updated. + - ecflow<5.17 - nco - pyyaml - numpy diff --git a/geoglows_ecflow/resources/archive_to_aws.py b/geoglows_ecflow/resources/archive_to_aws.py index 71560bd..fcaae55 100644 --- a/geoglows_ecflow/resources/archive_to_aws.py +++ b/geoglows_ecflow/resources/archive_to_aws.py @@ -1,11 +1,12 @@ import argparse import glob -import json import os import boto3 import yaml +from geoglows_ecflow.resources.helper_functions import load_forecast_run + def upload_to_s3(workspace: str, aws_config_file: str): """ @@ -22,10 +23,9 @@ def upload_to_s3(workspace: str, aws_config_file: str): forecast_bucket_uri = config["bucket_forecast_archive"] mapstyletable_bucket_uri = config["bucket_maptable_archive"] - with open(os.path.join(workspace, "forecast_run.json"), "r") as f: - data = json.load(f) - date = data["date"] - output_dir = data["output_dir"] + data = load_forecast_run(workspace) + date = data["date"] + output_dir = data["output_dir"] # Create an S3 client s3 = boto3.client( @@ -60,17 +60,15 @@ def upload_to_s3(workspace: str, aws_config_file: str): argparser = argparse.ArgumentParser() argparser.add_argument( "workspace", - nargs=1, help="Path to suite home directory", ) argparser.add_argument( "aws_config_file", - nargs=1, help="Path to AWS config file", ) args = argparser.parse_args() - workspace = args.workspace[0] - aws_config_file = args.aws_config_file[0] + workspace = args.workspace + aws_config_file = args.aws_config_file upload_to_s3(workspace, aws_config_file) diff --git a/geoglows_ecflow/resources/combine_esri_tables.py b/geoglows_ecflow/resources/combine_esri_tables.py index 7ade203..bb99e74 100644 --- a/geoglows_ecflow/resources/combine_esri_tables.py +++ b/geoglows_ecflow/resources/combine_esri_tables.py @@ -1,12 +1,13 @@ import argparse import logging -import sys import glob import os import pandas as pd +from geoglows_ecflow.resources.helper_functions import configure_logging + def combine_esri_tables(workspace: str): """Combines the map_style_tables from each VPU into 1 CSV file per time @@ -56,10 +57,6 @@ def combine_esri_tables(workspace: str): args = argparser.parse_args() workspace = args.workspace[0] - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s %(levelname)s %(message)s", - stream=sys.stdout, - ) + configure_logging() combine_esri_tables(workspace) diff --git a/geoglows_ecflow/resources/compute_init_flows.py b/geoglows_ecflow/resources/compute_init_flows.py index f777bc7..b10cc04 100644 --- a/geoglows_ecflow/resources/compute_init_flows.py +++ b/geoglows_ecflow/resources/compute_init_flows.py @@ -1,10 +1,11 @@ import argparse -import json import os import pandas as pd import xarray as xr +from geoglows_ecflow.resources.helper_functions import load_forecast_run + # Time index 7 of the ensemble-mean Q corresponds to t+24h on the ENS # 3h-resolution grid (the HRES member is excluded from the average upstream by @@ -14,8 +15,7 @@ def main(workspace: str, vpu: str) -> None: - with open(os.path.join(workspace, "forecast_run.json"), "r") as f: - ymd = json.load(f)["date"] + ymd = load_forecast_run(workspace)["date"] avg_path = os.path.join(workspace, "output", f"nces_avg_{vpu}.nc") out_path = os.path.join(workspace, "input", vpu, f"Qinit_{ymd}.parquet") diff --git a/geoglows_ecflow/resources/concat_forecast_warnings.py b/geoglows_ecflow/resources/concat_forecast_warnings.py index 8468e50..3070272 100644 --- a/geoglows_ecflow/resources/concat_forecast_warnings.py +++ b/geoglows_ecflow/resources/concat_forecast_warnings.py @@ -1,10 +1,11 @@ import os import glob -import sys import logging import pandas as pd import argparse +from geoglows_ecflow.resources.helper_functions import configure_logging + def concat_warnings(workdir: str) -> None: """ @@ -51,10 +52,6 @@ def concat_warnings(workdir: str) -> None: args = parser.parse_args() workspace = args.workspace[0] - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", - stream=sys.stdout, - ) + configure_logging() concat_warnings(workspace) diff --git a/geoglows_ecflow/resources/day_one_forecast.py b/geoglows_ecflow/resources/day_one_forecast.py index 02ec49e..2ed7cda 100644 --- a/geoglows_ecflow/resources/day_one_forecast.py +++ b/geoglows_ecflow/resources/day_one_forecast.py @@ -2,14 +2,21 @@ import glob import logging import os -import sys import argparse import numpy as np import pandas as pd import xarray as xr import netCDF4 as nc -import dask -from numcodecs import Blosc + +from geoglows_ecflow.resources.helper_functions import ( + RETURN_PERIODS, + configure_logging, +) +from geoglows_ecflow.resources.zarr_io import write_dataset_to_zarr + +# Only streams of at least this Strahler order are checked against return +# periods for the warnings summary (smaller headwater streams are skipped). +MIN_STREAM_ORDER = 3 def check_for_return_period_flow( @@ -17,59 +24,38 @@ def check_for_return_period_flow( ): max_flow = max(forecasted_flows_df["means"]) - # temporary dates - date_r5 = "" - date_r10 = "" - date_r25 = "" - date_r50 = "" - date_r100 = "" - - # retrieve return period flow levels from dataframe - r2 = float(rp_data["rp2"].values[0]) - r5 = float(rp_data["rp5"].values[0]) - r10 = float(rp_data["rp10"].values[0]) - r25 = float(rp_data["rp25"].values[0]) - r50 = float(rp_data["rp50"].values[0]) - r100 = float(rp_data["rp100"].values[0]) - - # then compare the timeseries to the return period thresholds - if max_flow >= r2: - date_r2 = get_time_of_first_exceedance(forecasted_flows_df, r2) + # retrieve return period flow levels from the dataframe + thresholds = { + rp: float(rp_data[f"rp{rp}"].values[0]) for rp in RETURN_PERIODS + } + # if the flow is not larger than the smallest return period, return the # dataframe without appending anything - else: + if max_flow < thresholds[RETURN_PERIODS[0]]: return largeflows_df - # check the rest of the return period flow levels - if max_flow >= r5: - date_r5 = get_time_of_first_exceedance(forecasted_flows_df, r5) - if max_flow >= r10: - date_r10 = get_time_of_first_exceedance(forecasted_flows_df, r10) - if max_flow >= r25: - date_r25 = get_time_of_first_exceedance(forecasted_flows_df, r25) - if max_flow >= r50: - date_r50 = get_time_of_first_exceedance(forecasted_flows_df, r50) - if max_flow >= r100: - date_r100 = get_time_of_first_exceedance(forecasted_flows_df, r100) - - new_row = pd.DataFrame( - { - "comid": rp_data.index[0], - "stream_order": stream_order, - "max_forecasted_flow": round(max_flow, 2), - "date_exceeds_return_period_2": date_r2, - "date_exceeds_return_period_5": date_r5, - "date_exceeds_return_period_10": date_r10, - "date_exceeds_return_period_25": date_r25, - "date_exceeds_return_period_50": date_r50, - "date_exceeds_return_period_100": date_r100, - }, - index=[0], - ) + # compare the timeseries to each return period threshold, ascending, so the + # progressive masking inside get_time_of_first_exceedance is preserved + exceedance_dates = {} + for rp in RETURN_PERIODS: + if max_flow >= thresholds[rp]: + exceedance_dates[rp] = get_time_of_first_exceedance( + forecasted_flows_df, thresholds[rp] + ) + else: + exceedance_dates[rp] = "" + + row = { + "comid": rp_data.index[0], + "stream_order": stream_order, + "max_forecasted_flow": round(max_flow, 2), + } + for rp in RETURN_PERIODS: + row[f"date_exceeds_return_period_{rp}"] = exceedance_dates[rp] - largeflows_df = pd.concat([largeflows_df, new_row], ignore_index=True) + new_row = pd.DataFrame(row, index=[0]) - return largeflows_df + return pd.concat([largeflows_df, new_row], ignore_index=True) def get_time_of_first_exceedance(forecasted_flows_df, flow): @@ -129,7 +115,8 @@ def postprocess_vpu( streams_file_path = os.path.join(input_dir, "master_table.parquet") streams_df = pd.read_parquet(streams_file_path) large_vpu_streams_df = streams_df[ - (streams_df["VPUCode"] == int(vpu)) & ((streams_df["strmOrder"] >= 3)) + (streams_df["VPUCode"] == int(vpu)) + & (streams_df["strmOrder"] >= MIN_STREAM_ORDER) ] # get the list of comids @@ -288,48 +275,16 @@ def netcdf_forecast_record_to_zarr(record_path) -> None: zarr_path = record_path.replace(".nc", ".zarr") record_nc = xr.open_dataset(record_path) - with dask.config.set(**{ - 'array.slicing.split_large_chunks': False, - # set the max chunk size to 5MB - 'array.chunk-size': '40MB', - # use the threads scheduler - 'scheduler': 'threads', - # set the maximum memory target usage to 90% of total memory - 'distributed.worker.memory.target': 0.80, - # do not allow spilling to disk - 'distributed.worker.memory.spill': False, - # specify the amount of resources to allocate to dask workers - 'distributed.worker.resources': { - 'memory': 3e9, # 1e9=1GB, this is the amount per worker - 'cpu': os.cpu_count(), # num CPU per worker - } - }): - #set compressing information - logging.info("Configuring compression") - - #if we get rid of dask, we can get rid of the compressor - #the compressor throws an error for version 3 so specify version 2 - compressor = Blosc(cname="zstd", clevel=3, shuffle=Blosc.BITSHUFFLE) - encoding = {'Q': {"compressor": compressor}} - - logging.info("Writing to zarr") - ( - record_nc - .drop_vars(["lat", "lon"], errors="ignore") - .chunk({ - "time": -1, - "river_id": "auto" - }) - .to_zarr( - zarr_path, - consolidated=True, - encoding=encoding, - mode = 'w', - zarr_version=2 - ) - ) - - record_nc.close() + logging.info("Writing to zarr") + write_dataset_to_zarr( + record_nc, + zarr_path, + {"time": -1, "river_id": "auto"}, + drop_vars=["lat", "lon"], + mode="w", + zarr_version=2, + ) + record_nc.close() logging.info("Done") @@ -337,31 +292,28 @@ def netcdf_forecast_record_to_zarr(record_path) -> None: parser = argparse.ArgumentParser() parser.add_argument( "workspace", - nargs=1, help="path to the daily workspace directory", ) parser.add_argument( "vpu", - nargs=1, help="VPU number", ) parser.add_argument( "output_dir", - nargs=1, help="path to the forecast records output directory", ) args = parser.parse_args() - workspace = args.workspace[0] - vpu = args.vpu[0] + workspace = args.workspace + vpu = args.vpu input_dir = os.path.join(workspace, "input") output_dir = os.path.join(workspace, "output") returnperiods = os.path.join(workspace, "return_periods_dir") - forecast_records = args.output_dir[0] + forecast_records = args.output_dir output_dir = os.path.join(workspace, "output") # start logging - logging.basicConfig(stream=sys.stdout, level=logging.INFO) + configure_logging() postprocess_vpu( vpu, input_dir, output_dir, returnperiods, forecast_records diff --git a/geoglows_ecflow/resources/generate_esri_table.py b/geoglows_ecflow/resources/generate_esri_table.py index 1047701..c9f628b 100644 --- a/geoglows_ecflow/resources/generate_esri_table.py +++ b/geoglows_ecflow/resources/generate_esri_table.py @@ -1,17 +1,29 @@ import argparse import logging import os -import sys import netCDF4 as nc import pandas as pd import xarray as xr +from geoglows_ecflow.resources.helper_functions import ( + RETURN_PERIODS, + configure_logging, +) + +# Only the first 10 days of the forecast are summarized in the style table. +FORECAST_WINDOW_DAYS = 10 + +# Mean-flow thresholds (m^3/s) that drive the map line-thickness ladder. Flows +# below the first threshold get thickness 1; each threshold crossed bumps the +# thickness by one (levels 2..6). +THICKNESS_THRESHOLDS = [20, 250, 1500, 10000, 30000] + def postprocess_vpu_forecast_directory( output_dir: str, returnperiods: str, - vpu: int or str, + vpu: int | str, ): # creates file name for the csv file date_string = os.path.basename( @@ -34,7 +46,8 @@ def postprocess_vpu_forecast_directory( # limit both dataframes to the first 10 days mean_flow_df = mean_flow_df[ - mean_flow_df.index <= mean_flow_df.index[0] + pd.Timedelta(days=10) + mean_flow_df.index + <= mean_flow_df.index[0] + pd.Timedelta(days=FORECAST_WINDOW_DAYS) ] # creating pandas dataframe with return periods @@ -43,32 +56,21 @@ def postprocess_vpu_forecast_directory( with nc.Dataset(rp_path, "r") as rp_ncfile: rp_df = pd.DataFrame( { - "return_2": rp_ncfile.variables["rp2"][:], - "return_5": rp_ncfile.variables["rp5"][:], - "return_10": rp_ncfile.variables["rp10"][:], - "return_25": rp_ncfile.variables["rp25"][:], - "return_50": rp_ncfile.variables["rp50"][:], - "return_100": rp_ncfile.variables["rp100"][:], + f"return_{rp}": rp_ncfile.variables[f"rp{rp}"][:] + for rp in RETURN_PERIODS }, index=rp_ncfile.variables["river_id"][:], ) mean_thickness_df = pd.DataFrame(columns=comids, index=dates, dtype=int) mean_thickness_df[:] = 1 - mean_thickness_df[mean_flow_df >= 20] = 2 - mean_thickness_df[mean_flow_df >= 250] = 3 - mean_thickness_df[mean_flow_df >= 1500] = 4 - mean_thickness_df[mean_flow_df >= 10000] = 5 - mean_thickness_df[mean_flow_df >= 30000] = 6 + for level, threshold in enumerate(THICKNESS_THRESHOLDS, start=2): + mean_thickness_df[mean_flow_df >= threshold] = level mean_ret_per_df = pd.DataFrame(columns=comids, index=dates, dtype=int) mean_ret_per_df[:] = 0 - mean_ret_per_df[mean_flow_df.gt(rp_df["return_2"], axis=1)] = 2 - mean_ret_per_df[mean_flow_df.gt(rp_df["return_5"], axis=1)] = 5 - mean_ret_per_df[mean_flow_df.gt(rp_df["return_10"], axis=1)] = 10 - mean_ret_per_df[mean_flow_df.gt(rp_df["return_25"], axis=1)] = 25 - mean_ret_per_df[mean_flow_df.gt(rp_df["return_50"], axis=1)] = 50 - mean_ret_per_df[mean_flow_df.gt(rp_df["return_100"], axis=1)] = 100 + for rp in RETURN_PERIODS: + mean_ret_per_df[mean_flow_df.gt(rp_df[f"return_{rp}"], axis=1)] = rp mean_flow_df = mean_flow_df.stack().to_frame().rename(columns={0: "mean"}) mean_thickness_df = ( @@ -118,12 +120,7 @@ def postprocess_vpu_forecast_directory( returnperiods = os.path.join(workspace, "return_periods_dir") vpu = args.vpu[0] - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s %(levelname)s %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", - stream=sys.stdout, - ) + configure_logging() params = [output_dir, returnperiods, vpu] diff --git a/geoglows_ecflow/resources/helper_functions.py b/geoglows_ecflow/resources/helper_functions.py index b33e3e4..1ebdcab 100644 --- a/geoglows_ecflow/resources/helper_functions.py +++ b/geoglows_ecflow/resources/helper_functions.py @@ -2,11 +2,48 @@ # See: spt_compute (https://github.com/erdc/spt_compute) # Updated by Michael Souffront, 2023 +import json import os import sys import re import logging as log +# Return periods (years) used throughout the forecast post-processing. Defined +# once here and imported by the modules that build return-period ladders. +RETURN_PERIODS = [2, 5, 10, 25, 50, 100] + +# The high-resolution (HRES) forecast is always ensemble member 52. Members +# 1-51 are the ensemble perturbations. Defined once here and imported wherever +# the workflow needs to single out (or exclude) the HRES member. +HRES_ENSEMBLE_MEMBER = 52 + +# Shared logging format so every resource module logs identically. +LOG_FORMAT = "%(asctime)s %(levelname)s %(message)s" +LOG_DATEFMT = "%Y-%m-%d %H:%M:%S" + + +def configure_logging(level: str = "INFO") -> None: + """Configure root logging with the shared format, writing to stdout.""" + log.basicConfig( + level=level, + format=LOG_FORMAT, + datefmt=LOG_DATEFMT, + stream=sys.stdout, + ) + + +def load_forecast_run(workspace: str) -> dict: + """Load and parse /forecast_run.json. + + Args: + workspace (str): Directory containing forecast_run.json. + + Returns: + dict: The parsed forecast-run manifest. + """ + with open(os.path.join(workspace, "forecast_run.json"), "r") as f: + return json.load(f) + def create_logger( name: str, level: str = "INFO", log_file: str | None = None @@ -21,13 +58,13 @@ def create_logger( else: handler = log.StreamHandler(sys.stdout) - handler.setLevel(level) - handler.setFormatter( - log.Formatter("%(asctime)s - %(levelname)s - %(message)s") - ) + handler.setLevel(level) + handler.setFormatter( + log.Formatter("%(asctime)s - %(levelname)s - %(message)s") + ) - # Add the handler to the logger - logger.addHandler(handler) + # Add the handler to the logger + logger.addHandler(handler) return logger diff --git a/geoglows_ecflow/resources/netcdf_to_zarr.py b/geoglows_ecflow/resources/netcdf_to_zarr.py index 1fd2ffb..934ef8f 100644 --- a/geoglows_ecflow/resources/netcdf_to_zarr.py +++ b/geoglows_ecflow/resources/netcdf_to_zarr.py @@ -1,21 +1,20 @@ import argparse import glob -import json import logging import os import shutil -import sys -import dask import numpy as np import xarray as xr -from numcodecs import Blosc -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s %(message)s", - stream=sys.stdout, +from geoglows_ecflow.resources.helper_functions import ( + HRES_ENSEMBLE_MEMBER, + configure_logging, + load_forecast_run, ) +from geoglows_ecflow.resources.zarr_io import write_dataset_to_zarr + +configure_logging() def netcdf_forecasts_to_zarr(workspace: str) -> None: @@ -25,69 +24,47 @@ def netcdf_forecasts_to_zarr(workspace: str) -> None: Args: workspace (str): Path to forecast_run.json base directory. """ - with open(os.path.join(workspace, "forecast_run.json"), "r") as f: - data = json.load(f) - output_dir = data["output_dir"] - date = data["date"] + data = load_forecast_run(workspace) + output_dir = data["output_dir"] + date = data["date"] vpu_nums = sorted( - set([os.path.basename(x).split("_")[1] for x in glob.glob(os.path.join(output_dir, f"Qout_*_52.nc"))]) + set([os.path.basename(x).split("_")[1] for x in glob.glob(os.path.join(output_dir, f"Qout_*_{HRES_ENSEMBLE_MEMBER}.nc"))]) ) qout_1_51_files = sorted([os.path.join(output_dir, f"Qout_{vpu}.nc") for vpu in vpu_nums]) - qout_52_files = sorted(glob.glob(os.path.join(output_dir, f"Qout_*_52.nc"))) + qout_52_files = sorted(glob.glob(os.path.join(output_dir, f"Qout_*_{HRES_ENSEMBLE_MEMBER}.nc"))) zarr_file_path = os.path.join(output_dir, f"Qout_{date}.zarr") if os.path.exists(zarr_file_path): shutil.rmtree(zarr_file_path) - with dask.config.set(**{ - 'array.slicing.split_large_chunks': False, - # set the max chunk size to 5MB - 'array.chunk-size': '40MB', - # use the threads scheduler - 'scheduler': 'threads', - # set the maximum memory target usage to 90% of total memory - 'distributed.worker.memory.target': 0.80, - # do not allow spilling to disk - 'distributed.worker.memory.spill': False, - # specify the amount of resources to allocate to dask workers - 'distributed.worker.resources': { - 'memory': 3e9, # 1e9=1GB, this is the amount per worker - 'cpu': os.cpu_count(), # num CPU per worker - } - }): - logging.info("Opening ensembles 1-51 datasets") - with xr.open_mfdataset(qout_1_51_files, combine="nested", concat_dim="river_id") as ds151: + logging.info("Opening ensembles 1-51 datasets") + with xr.open_mfdataset( + qout_1_51_files, combine="nested", concat_dim="river_id" + ) as ds151: + logging.info("Assigning the ensemble coordinate variable") + ds151 = ds151.assign_coords( + ensemble=np.arange(1, HRES_ENSEMBLE_MEMBER) + ) + logging.info("Opening ensemble 52 dataset") + with xr.open_mfdataset( + qout_52_files, combine="nested", concat_dim="river_id" + ) as ds52: logging.info("Assigning the ensemble coordinate variable") - ds151 = ds151.assign_coords(ensemble=np.arange(1, 52)) - logging.info("Opening ensemble 52 dataset") - with xr.open_mfdataset(qout_52_files, combine="nested", concat_dim="river_id") as ds52: - logging.info("Assigning the ensemble coordinate variable") - ds52 = ds52.assign_coords(ensemble=52) + ds52 = ds52.assign_coords(ensemble=HRES_ENSEMBLE_MEMBER) - logging.info("Concatenating 1-51 and 52 datasets") - ds = xr.concat([ds151, ds52], dim="ensemble") + logging.info("Concatenating 1-51 and 52 datasets") + ds = xr.concat([ds151, ds52], dim="ensemble") - logging.info("Configuring compression") - compressor = Blosc(cname="zstd", clevel=3, shuffle=Blosc.BITSHUFFLE) - encoding = {'Q': {"compressor": compressor}} - logging.info("Writing to zarr") - ( - ds - .drop_vars(["crs", "lat", "lon", "time_bnds"], errors="ignore") - .chunk({ - "time": -1, - "river_id": "auto", - "ensemble": -1 - }) - .to_zarr( - zarr_file_path, - consolidated=True, - encoding=encoding, - ) - ) - logging.info("Done") + logging.info("Writing to zarr") + write_dataset_to_zarr( + ds, + zarr_file_path, + {"time": -1, "river_id": "auto", "ensemble": -1}, + drop_vars=["crs", "lat", "lon", "time_bnds"], + ) + logging.info("Done") if __name__ == "__main__": @@ -95,8 +72,7 @@ def netcdf_forecasts_to_zarr(workspace: str) -> None: parser.add_argument( "workspace", - nargs=1, help="Path to the suite home directory.", ) args = parser.parse_args() - netcdf_forecasts_to_zarr(workspace=args.workspace[0]) + netcdf_forecasts_to_zarr(workspace=args.workspace) diff --git a/geoglows_ecflow/resources/prep_river_route_forecast.py b/geoglows_ecflow/resources/prep_river_route_forecast.py index 9f2e154..d1e9b6b 100644 --- a/geoglows_ecflow/resources/prep_river_route_forecast.py +++ b/geoglows_ecflow/resources/prep_river_route_forecast.py @@ -6,17 +6,13 @@ from glob import glob from geoglows_ecflow.resources.helper_functions import ( + configure_logging, get_ensemble_number_from_forecast, get_valid_vpucode_list, ) -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s %(levelname)s %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", - stream=sys.stdout, -) +configure_logging() def forecast_preprocess( diff --git a/geoglows_ecflow/resources/run_river_route_forecast.py b/geoglows_ecflow/resources/run_river_route_forecast.py index e88cf13..6c8424a 100644 --- a/geoglows_ecflow/resources/run_river_route_forecast.py +++ b/geoglows_ecflow/resources/run_river_route_forecast.py @@ -1,6 +1,5 @@ import argparse import datetime -import json import os import sys from glob import glob @@ -10,6 +9,7 @@ from geoglows_ecflow.resources.helper_functions import ( create_logger, get_ensemble_number_from_forecast, + load_forecast_run, ) @@ -30,8 +30,7 @@ def _find_state_init(vpu_input_dir: str, date: str) -> str | None: def river_route_forecast_exec(workspace: str, job_id: str, log_dir: str) -> None: - with open(os.path.join(workspace, "forecast_run.json"), "r") as f: - data = json.load(f) + data = load_forecast_run(workspace) job = data.get(job_id) if not job: diff --git a/geoglows_ecflow/resources/zarr_io.py b/geoglows_ecflow/resources/zarr_io.py new file mode 100644 index 0000000..c86b1af --- /dev/null +++ b/geoglows_ecflow/resources/zarr_io.py @@ -0,0 +1,67 @@ +"""Shared helpers for writing xarray datasets to zarr. + +The dask configuration and Blosc/zstd compression settings were previously +copy-pasted into both ``netcdf_to_zarr`` and ``day_one_forecast``. They live +here so the heavy ``dask``/``numcodecs`` imports stay out of the lightweight +``helper_functions`` module. +""" + +import os + +import dask +from numcodecs import Blosc + +# Dask settings shared by every zarr write in the workflow. +DASK_ZARR_CONFIG = { + "array.slicing.split_large_chunks": False, + # cap the max chunk size + "array.chunk-size": "40MB", + # use the threads scheduler + "scheduler": "threads", + # set the maximum memory target usage to 80% of total memory + "distributed.worker.memory.target": 0.80, + # do not allow spilling to disk + "distributed.worker.memory.spill": False, + # resources to allocate to dask workers + "distributed.worker.resources": { + "memory": 3e9, # 1e9=1GB, per worker + "cpu": os.cpu_count(), # num CPU per worker + }, +} + + +def write_dataset_to_zarr( + ds, + zarr_path: str, + chunks: dict, + drop_vars=(), + **to_zarr_kwargs, +) -> None: + """Write an open xarray dataset to a consolidated, Blosc-compressed zarr. + + The caller is responsible for opening (and closing) ``ds``; this helper + only applies the shared dask config, drops the requested vars, rechunks, + and writes. Per-call differences (chunk dims, dropped vars, ``mode``, + ``zarr_version``) are passed in rather than hardcoded. + + Args: + ds: An open xarray Dataset. + zarr_path (str): Destination zarr store path. + chunks (dict): Chunk spec passed to ``ds.chunk``. + drop_vars: Variable names to drop (ignored if absent). + **to_zarr_kwargs: Extra keyword args forwarded to ``to_zarr`` + (e.g. ``mode``, ``zarr_version``). + """ + with dask.config.set(**DASK_ZARR_CONFIG): + compressor = Blosc(cname="zstd", clevel=3, shuffle=Blosc.BITSHUFFLE) + encoding = {"Q": {"compressor": compressor}} + ( + ds.drop_vars(list(drop_vars), errors="ignore") + .chunk(chunks) + .to_zarr( + zarr_path, + consolidated=True, + encoding=encoding, + **to_zarr_kwargs, + ) + ) diff --git a/geoglows_ecflow/workflow/builders/builder.py b/geoglows_ecflow/workflow/builders/builder.py index bd27ed5..630f6f8 100644 --- a/geoglows_ecflow/workflow/builders/builder.py +++ b/geoglows_ecflow/workflow/builders/builder.py @@ -1,24 +1,35 @@ from geoglows_ecflow.workflow.builders.base import GEOGLOWSBaseBuilder from geoglows_ecflow.workflow.comfies.ooflow import Trigger, Defuser from geoglows_ecflow.workflow.comfies.ooflow import all_complete, Event, complete -from geoglows_ecflow.workflow.comfies.ooflow import complete, Limit, InLimit, Variable +from geoglows_ecflow.workflow.comfies.ooflow import Limit, InLimit, Variable from geoglows_ecflow.workflow.comfies.ooflow import RepeatDate, Defstatus -from geoglows_ecflow.workflow.parts.nodes import Family, Task +from geoglows_ecflow.workflow.parts.nodes import Family, Task, NominalTime from geoglows_ecflow.workflow.parts.times import ( t2t, - Timer, CronDateRefresh, CronDataAvailability, ) -from geoglows_ecflow.workflow.comfies.dateandtime import Date, TimeDelta, CalSeq +from geoglows_ecflow.workflow.comfies.dateandtime import Date, CalSeq from geoglows_ecflow.workflow.parts.admin import AdminFamily from geoglows_ecflow.workflow.parts.epilogs import DummyEpilog from geoglows_ecflow.workflow.parts.repeats import calseq_repeat from geoglows_ecflow.workflow.parts.packages import PackageInstallers from geoglows_ecflow.workflow.comfies.partition import partition +from geoglows_ecflow.resources.helper_functions import HRES_ENSEMBLE_MEMBER +# Scheduler memory reservations (MB) for the heavier tasks. +ENS_TASK_MEM_MB = 6000 +ARCHIVE_QINIT_MEM_MB = 4000 -from geoglows_ecflow.workflow.parts.nodes import Family, Task, NominalTime + +def is_00z_cycle(nominal_time): + """Whether a nominal-time family is the 00Z cycle. + + The 00Z cycle runs the full ensemble pipeline; the 12Z cycle runs a + reduced tree. The cycle is identified by the EMOS_BASE variable the + NominalTime family carries ("00" or "12"). + """ + return nominal_time.get_variable("EMOS_BASE").value() != "12" class Builder(GEOGLOWSBaseBuilder): @@ -39,36 +50,30 @@ class Builder(GEOGLOWSBaseBuilder): def build(self): """ - Create parts and wire them together into an GLOFAS suite. + Create parts and wire them together into a GEOGloWS suite. Naming conventions: n_* -- ecFlow Node object e_* -- ecFlow Event object """ super(Builder, self).build() cfg = self.config - # get suite configuration parameters from the deployment config file - suite_name = self.config.get("name") - mode = self.config.get("mode", choices=["prod", "test", "rd"]) - first_date = self.config.get("first_date", type=int) - last_date = self.config.get("last_date", type=int, default="20300101") - first_barrier = self.config.get( - "first_barrier", type=int, default=first_date - ) - archive_path = self.config.get("exparch") - suite_dir = self.config.get("workroot") - - with_flood_hazard = self.config.get("with_floodhazard", default=False) - wb_days = self.config.get("wb_days", type=int, default=10) - # initially empty suite, provided by parent - # class will be filled up with content here. - mars_nworkers = self.config.get("mars_workers", type=int, default=1) - ens_members = self.config.get("ens_members", type=int, default=51) - ens_range = self.config.get("ens_range", type=int, default=30) - suite = self.suite - par_jobvars = self.jobvars.dest("parallel", fallback="PARENT") - - # Selectable Trigger suites + # All tunable parameters read from the deployment config file are + # gathered here so what the suite exposes is visible at a glance. + # (exparch/workroot are consumed by the task scripts via templating, + # so they are intentionally not read here.) + suite_name = cfg.get("name") + # Validate the run mode (consumed by suite.h via templating); the + # builder no longer branches on it, so the return is discarded. + cfg.get("mode", choices=["prod", "test"]) + first_date = cfg.get("first_date", type=int) + last_date = cfg.get("last_date", type=int, default="20300101") + first_barrier = cfg.get("first_barrier", type=int, default=first_date) + mars_nworkers = cfg.get("mars_workers", type=int, default=1) + ens_members = cfg.get("ens_members", type=int, default=51) + vpu_list = cfg.get("vpu_list", type=list, default=[]) + + # Operational suites this suite triggers off (normalized to lead "/"). o_suite = cfg.get("o_suite", default="/o") mc_suite = cfg.get("mc_suite", default="/mc") if o_suite[0] != "/": @@ -76,8 +81,7 @@ def build(self): if mc_suite[0] != "/": mc_suite = f"/{mc_suite}" - # these flags are not user-configurable but - # depend on other flags + suite = self.suite # admin family n_admin = Family("admin") @@ -99,17 +103,6 @@ def build(self): with_webpush = False with_diss = False - follow_osuite = False - in_production = False - in_test = False - - if mode == "prod": - in_production = True - if mode == "test": - in_test = True - - if in_production or in_test: - follow_osuite = True # make family n_make = Family("make") @@ -138,8 +131,6 @@ def build(self): n_make.add_inlimit("make") - vpu_list = self.config.get("vpu_list", type=list, default=[]) - n_make.add(Variable("YMD", first_date)) suite.add(n_make, n_admin) @@ -180,47 +171,38 @@ def build(self): (barrier_00, main_00, lag_00), ): cycle = str(main_hh.time.hh) - tnom = main_hh.time - - if follow_osuite: - self.defs.add_extern(f"{mc_suite}/main:YMD") - self.defs.add_extern(f"{mc_suite}/main/{cycle}/fc0015d/fc") - self.defs.add_extern(f"{o_suite}/main:YMD") - self.defs.add_extern(f"{o_suite}/main/{cycle}/fc/model") - n_run_hr = Family("run_hr").add( - Trigger( - f"({o_suite}/main:YMD == /{suite_name}/barrier/daily:YMD " - f"and {o_suite}/main/{cycle}/fc/model == complete) " - f"or ({o_suite}/main:YMD > /{suite_name}/barrier/daily:YMD)" - ) - ) - n_run_hr.add( - Task("dummy").add(Trigger("0==1")).add(Defuser("1==1")) - ) - n_run_en = Family("run_en").add( - Trigger( - f"({mc_suite}/main:YMD == /{suite_name}/barrier/daily:YMD " - f"and {mc_suite}/main/{cycle}/fc0015d/fc == complete) " - f"or ({mc_suite}/main:YMD > /{suite_name}/barrier/daily:YMD)" - ) - ) - n_run_en.add( - Task("dummy").add(Trigger("0==1")).add(Defuser("1==1")) - ) - n_barrier_epilog = Family("last").add( - Trigger( - f"{o_suite}/main:YMD > /{suite_name}/barrier/daily:YMD" - ), - Task("sleep").add(Trigger("0==1"), Defuser("1==1")), + self.defs.add_extern(f"{mc_suite}/main:YMD") + self.defs.add_extern(f"{mc_suite}/main/{cycle}/fc0015d/fc") + self.defs.add_extern(f"{o_suite}/main:YMD") + self.defs.add_extern(f"{o_suite}/main/{cycle}/fc/model") + n_run_hr = Family("run_hr").add( + Trigger( + f"({o_suite}/main:YMD == /{suite_name}/barrier/daily:YMD " + f"and {o_suite}/main/{cycle}/fc/model == complete) " + f"or ({o_suite}/main:YMD > /{suite_name}/barrier/daily:YMD)" ) + ) + n_run_hr.add( + Task("dummy").add(Trigger("0==1")).add(Defuser("1==1")) + ) + n_run_en = Family("run_en").add( + Trigger( + f"({mc_suite}/main:YMD == /{suite_name}/barrier/daily:YMD " + f"and {mc_suite}/main/{cycle}/fc0015d/fc == complete) " + f"or ({mc_suite}/main:YMD > /{suite_name}/barrier/daily:YMD)" + ) + ) + n_run_en.add( + Task("dummy").add(Trigger("0==1")).add(Defuser("1==1")) + ) - else: - n_run_hr = Family("run_hr") - n_run_en = Family("run_en") - n_run_hr.add(Task("dummy"), Timer(tnom + TimeDelta(hours=7))) - n_run_en.add(Task("dimmy"), Timer(tnom + TimeDelta(hours=9))) - n_barrier_epilog = DummyEpilog(done=Timer("14:15")) + n_barrier_epilog = Family("last").add( + Trigger( + f"{o_suite}/main:YMD > /{suite_name}/barrier/daily:YMD" + ), + Task("sleep").add(Trigger("0==1"), Defuser("1==1")), + ) barrier_hh.add(n_run_hr, n_run_en) n_barrier_daily.add(barrier_hh) @@ -236,8 +218,7 @@ def build(self): n_hr.add(n_ret_hr) n_ens = Family("ens") - if follow_osuite: - n_ens.trigger = n_run_en.complete.across("YMD") + n_ens.trigger = n_run_en.complete.across("YMD") n_ens.trigger &= n_initialize.complete n_ens.add(Variable("CONTEXT", "ens")) n_ret_ens = Family("retrieve") @@ -258,14 +239,14 @@ def build(self): n_prep_ens.trigger &= n_ret_hr.complete n_ens_ens = Family("ens_members") n_ens_ens.trigger = n_prep_ens.complete - n_ens_ens.add_variable("MEM", 6000) + n_ens_ens.add_variable("MEM", ENS_TASK_MEM_MB) n_ens.add(n_ret_ens) - if main_hh.get_variable("EMOS_BASE").value() != "12": + if is_00z_cycle(main_hh): n_ens.add(n_prep_ens, n_ens_ens) for vpu in vpu_list: # Create the ensemble tasks - for mem in reversed(range(1, 53)): + for mem in reversed(range(1, HRES_ENSEMBLE_MEMBER + 1)): n_member = Family(f"{vpu}_{mem:02d}").add( Task("ens_member"), Variable("JOB_ID", f"job_{vpu}_{mem}"), @@ -300,7 +281,7 @@ def build(self): n_forecast_warnings.trigger = n_vpus.complete n_archive_qinit = Task("archive_qinit") - n_archive_qinit.add_variable('MEM', 4000) + n_archive_qinit.add_variable('MEM', ARCHIVE_QINIT_MEM_MB) n_archive_qinit.add_variable('NCPUS', 12) n_archive_qinit.trigger = n_vpus.complete @@ -327,10 +308,7 @@ def build(self): n_web_test.defuser = e_no_web_test n_web.add(n_web_prod, n_web_test) - if not follow_osuite: - barrier_ymd = barrier_hh.ymd - - if main_hh.get_variable("EMOS_BASE").value() != "12": + if is_00z_cycle(main_hh): main_hh.add( n_initialize, n_hr, @@ -355,7 +333,7 @@ def build(self): n_lag_arch_init.defuser = e_no_ecfs_archive n_lag_arch_fc = Task("arch_fc") n_lag_arch_fc.defuser = e_no_ecfs_archive - if main_hh.get_variable("EMOS_BASE").value() != "12": + if is_00z_cycle(main_hh): lag_hh.add(n_lag_arch_init, n_lag_arch_fc) lag_hh.trigger = main_hh.complete.across("YMD") n_daily_lag.add(lag_hh) diff --git a/geoglows_ecflow/workflow/comfies/config.py b/geoglows_ecflow/workflow/comfies/config.py index 1075676..ce48dbc 100644 --- a/geoglows_ecflow/workflow/comfies/config.py +++ b/geoglows_ecflow/workflow/comfies/config.py @@ -14,7 +14,8 @@ import os import sys -import imp +import importlib.util +from importlib.machinery import SourceFileLoader import collections.abc as collections import datetime @@ -54,6 +55,26 @@ class ConfigItemTypeError(ConfigError): # ---------------------------------------- +def _load_source(name, path): + """Load a Python source file as a module and return it. + + An explicit ``SourceFileLoader`` is used so the file is parsed as Python + source regardless of its extension (config files use ``.cfg``). + """ + loader = SourceFileLoader(name, path) + spec = importlib.util.spec_from_file_location(name, path, loader=loader) + module = importlib.util.module_from_spec(spec) + # Register before executing so the file's own relative imports resolve; + # drop it again if execution fails, leaving no half-built module behind. + sys.modules[name] = module + try: + loader.exec_module(module) + except BaseException: + sys.modules.pop(name, None) + raise + return module + + class ConfigSource(object): """ Abstract Base Class. @@ -105,7 +126,7 @@ def _load(self, path): try: old_dont_write_bytecode = sys.dont_write_bytecode sys.dont_write_bytecode = True - data = imp.load_source('_sdeploy_config_'+path, path).__dict__ + data = _load_source('_sdeploy_config_'+path, path).__dict__ sys.dont_write_bytecode = old_dont_write_bytecode except IOError as e: msg = path + ": " + e.strerror diff --git a/geoglows_ecflow/workflow/comfies/ooflow.py b/geoglows_ecflow/workflow/comfies/ooflow.py index 0a3c79b..f35e85a 100644 --- a/geoglows_ecflow/workflow/comfies/ooflow.py +++ b/geoglows_ecflow/workflow/comfies/ooflow.py @@ -21,7 +21,7 @@ import os from functools import wraps -from pkg_resources import parse_version +from packaging.version import parse as parse_version import ecflow from .py2 import reduce try: diff --git a/geoglows_ecflow/workflow/comfies/sdeploy.py b/geoglows_ecflow/workflow/comfies/sdeploy.py index c92227d..1193d19 100644 --- a/geoglows_ecflow/workflow/comfies/sdeploy.py +++ b/geoglows_ecflow/workflow/comfies/sdeploy.py @@ -28,8 +28,9 @@ import subprocess import contextlib import logging as log -import imp import importlib +import importlib.util +import importlib.machinery import argparse from geoglows_ecflow.workflow.comfies.memoize import memoize @@ -60,7 +61,7 @@ TemplateError, ) from geoglows_ecflow.workflow.comfies.sjob import SshHost -from pkg_resources import parse_version +from packaging.version import parse as parse_version from geoglows_ecflow.workflow.comfies.version import __version__ from geoglows_ecflow.workflow.comfies.py2 import basestring import ecflow @@ -368,12 +369,20 @@ def module_exists(name, path): """ Test if module exists in path. If module doesn't exist return False. + + Walk the dotted name one component at a time, descending into each + package's directory, and return the final source path (truthy) or False. """ for x in name.split("."): - try: - file, path, descr = imp.find_module(x, [path]) - except ImportError: + spec = importlib.machinery.PathFinder.find_spec(x, [path]) + if spec is None: return False + if spec.submodule_search_locations: + # x is a package; search the next component inside its directory + path = list(spec.submodule_search_locations)[0] + else: + # x is a module; its origin is the source file + path = spec.origin return path @@ -772,9 +781,9 @@ class BaseBuilder(object): def __init__(self, config): globals()["ecflow"] = importlib.import_module(self.ecflow_module) - if parse_version(self.comfies_minimum_version) > parse_version( - __version__ - ): + if self.comfies_minimum_version != "any" and parse_version( + self.comfies_minimum_version + ) > parse_version(__version__): raise ComfiesVersionError( "This suite needs version {}" " or later of comfies package".format( diff --git a/pyproject.toml b/pyproject.toml index 806e3c2..f8f1677 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,11 +18,21 @@ dependencies = [ "zarr>=2.16.1", "boto3>=1.28.65", "river-route>=2.1.1", + "packaging>=21.0", ] -requires-python = ">=3.11,<3.14" +requires-python = ">=3.12,<3.14" readme = "README.md" license = { text = "BSD-3-Clause" } +[project.optional-dependencies] +dev = [ + "pytest>=7.0", + "pytest-cov>=4.0", +] + +[tool.pytest.ini_options] +testpaths = ["tests"] + [tool.coverage.report] show_missing = true diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..711c7b1 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,60 @@ +"""Shared pytest fixtures and import shims for the resources test suite. + +These tests pin the *current* behavior of the pure helper functions used by +the river-route forecast workflow. The functions under test never touch the +heavy optional dependency ``river_route``, so when it is not installed we +inject a Dummy module into ``sys.modules`` before any test imports +``run_river_route_forecast``. When the real package *is* installed (e.g. on the +Ubuntu workflow box), it is used as-is. +""" + +import sys +import types + +import pandas as pd +import pytest + +# ``run_river_route_forecast`` does ``import river_route as rr`` at module top, +# but ``_find_state_init`` (the function under test) never touches it. Prefer +# the real package when installed; only fall back to a Dummy module so the +# import succeeds where the dependency is absent. +if "river_route" not in sys.modules: + try: + import river_route # noqa: F401 (use the real package when present) + except ModuleNotFoundError: + sys.modules["river_route"] = types.ModuleType("river_route") + + +@pytest.fixture +def make_flows(): + """Return a factory building a forecasted-flows DataFrame. + + The frame mirrors what ``day_one_forecast`` passes around: a ``means`` + column of forecast values and a parallel ``times`` column. + """ + + def _make(means, times=None): + if times is None: + times = [f"t{i}" for i in range(len(means))] + return pd.DataFrame({"means": list(means), "times": list(times)}) + + return _make + + +@pytest.fixture +def rp_table(): + """A single-row return-period table indexed by comid. + + Columns match what ``check_for_return_period_flow`` reads: rp2..rp100. + """ + return pd.DataFrame( + { + "rp2": [2.0], + "rp5": [5.0], + "rp10": [10.0], + "rp25": [25.0], + "rp50": [50.0], + "rp100": [100.0], + }, + index=[12345], + ) diff --git a/tests/test_compute_init_flows.py b/tests/test_compute_init_flows.py new file mode 100644 index 0000000..afe8bb6 --- /dev/null +++ b/tests/test_compute_init_flows.py @@ -0,0 +1,41 @@ +"""Test that compute_init_flows selects the configured time index.""" + +import json +import os + +import numpy as np +import pandas as pd +import xarray as xr + +from geoglows_ecflow.resources import compute_init_flows + + +def test_init_time_index_constant(): + # t+24h on the ENS 3h grid; documented contract relied on downstream. + assert compute_init_flows.INIT_TIME_INDEX == 7 + + +def test_qinit_parquet_holds_q_at_init_time_index(tmp_path): + vpu = "101" + ymd = "2023010100" + + (tmp_path / "forecast_run.json").write_text(json.dumps({"date": ymd})) + output_dir = tmp_path / "output" + output_dir.mkdir() + (tmp_path / "input" / vpu).mkdir(parents=True) + + # Q[time, river_id] with distinct values so the selected slice is unique. + data = np.arange(10 * 3).reshape(10, 3) + ds = xr.Dataset( + {"Q": (("time", "river_id"), data)}, + coords={"time": range(10), "river_id": [1, 2, 3]}, + ) + ds.to_netcdf(output_dir / f"nces_avg_{vpu}.nc") + + compute_init_flows.main(str(tmp_path), vpu) + + out = pd.read_parquet( + os.path.join(str(tmp_path), "input", vpu, f"Qinit_{ymd}.parquet") + ) + expected = data[compute_init_flows.INIT_TIME_INDEX] + assert out["Q"].tolist() == expected.tolist() diff --git a/tests/test_day_one_forecast.py b/tests/test_day_one_forecast.py new file mode 100644 index 0000000..79493e3 --- /dev/null +++ b/tests/test_day_one_forecast.py @@ -0,0 +1,49 @@ +"""Tests for the return-period exceedance logic in day_one_forecast.""" + +import pandas as pd + +from geoglows_ecflow.resources.day_one_forecast import ( + check_for_return_period_flow, + get_time_of_first_exceedance, +) + + +def test_first_exceedance_returns_earliest_time_at_or_above_flow(make_flows): + flows = make_flows([1, 5, 10, 2], times=["a", "b", "c", "d"]) + + # Rows below 4 (1 and 2) are dropped; first remaining time is "b" (5). + assert get_time_of_first_exceedance(flows, 4) == "b" + + +def test_below_smallest_return_period_returns_input_unchanged( + make_flows, rp_table +): + largeflows = pd.DataFrame() + flows = make_flows([0.5, 1.0]) # max 1.0 < rp2 (2.0) + + result = check_for_return_period_flow(largeflows, flows, 3, rp_table) + + assert result is largeflows + assert len(result) == 0 + + +def test_appends_row_with_exceeded_thresholds(make_flows, rp_table): + largeflows = pd.DataFrame() + flows = make_flows([1, 3, 6, 8], times=["a", "b", "c", "d"]) # max 8 + + result = check_for_return_period_flow(largeflows, flows, 3, rp_table) + + assert len(result) == 1 + row = result.iloc[0] + assert row["comid"] == 12345 + assert row["stream_order"] == 3 + assert row["max_forecasted_flow"] == 8 + # rp2 (2) and rp5 (5) are exceeded; the input frame is mutated between + # calls, so the first time >= rp5 is "c", not "b". + assert row["date_exceeds_return_period_2"] == "b" + assert row["date_exceeds_return_period_5"] == "c" + # rp10 (10) and above are not reached (max is 8). + assert row["date_exceeds_return_period_10"] == "" + assert row["date_exceeds_return_period_25"] == "" + assert row["date_exceeds_return_period_50"] == "" + assert row["date_exceeds_return_period_100"] == "" diff --git a/tests/test_generate_esri_table.py b/tests/test_generate_esri_table.py new file mode 100644 index 0000000..ace68f5 --- /dev/null +++ b/tests/test_generate_esri_table.py @@ -0,0 +1,82 @@ +"""Return-period ladder tests for generate_esri_table.""" + +import netCDF4 as nc +import numpy as np +import pandas as pd +import xarray as xr + +from geoglows_ecflow.resources.generate_esri_table import ( + postprocess_vpu_forecast_directory, +) + +VPU = "101" +DATE = "20230101.00" + + +def _write_inputs(tmp_path): + """Create the nces-average and return-period netCDFs the function reads.""" + base = tmp_path / DATE # basename(dirname(output_dir)) becomes the date + output_dir = base / "output" + returnperiods = base / "rp" + output_dir.mkdir(parents=True) + returnperiods.mkdir(parents=True) + + # river 1 flows at 300 (exceeds rp5 but not rp10); river 2 stays at 5. + times = pd.date_range("2023-01-01", periods=2, freq="3h") + flows = np.array([[300.0, 5.0], [300.0, 5.0]]) + ds = xr.Dataset( + {"Q": (("time", "river_id"), flows)}, + coords={"time": times, "river_id": [1, 2]}, + ) + ds.to_netcdf(output_dir / f"nces_avg_{VPU}.nc") + + rp = nc.Dataset(str(returnperiods / f"returnperiods_{VPU}.nc"), "w") + rp.createDimension("river_id", 2) + rp.createVariable("river_id", "i4", ("river_id",))[:] = [1, 2] + thresholds = {"rp2": 10, "rp5": 100, "rp10": 500, + "rp25": 1000, "rp50": 2000, "rp100": 5000} + for name, value in thresholds.items(): + rp.createVariable(name, "f8", ("river_id",))[:] = [value, value] + rp.close() + + return output_dir, returnperiods + + +def test_return_period_ladder_assigns_expected_levels(tmp_path): + output_dir, returnperiods = _write_inputs(tmp_path) + + postprocess_vpu_forecast_directory( + str(output_dir), str(returnperiods), VPU + ) + + table = pd.read_parquet( + output_dir + / "map_style_tables" + / f"mapstyletable_{VPU}_{DATE}.parquet" + ) + ret_per_by_comid = table.groupby("comid")["ret_per"].unique() + + # 300 exceeds rp5 (100) but not rp10 (500) -> level 5. + assert ret_per_by_comid[1].tolist() == [5] + # 5 is below rp2 (10) -> level 0. + assert ret_per_by_comid[2].tolist() == [0] + + +def test_thickness_ladder_assigns_expected_levels(tmp_path): + output_dir, returnperiods = _write_inputs(tmp_path) + + postprocess_vpu_forecast_directory( + str(output_dir), str(returnperiods), VPU + ) + + table = pd.read_parquet( + output_dir + / "map_style_tables" + / f"mapstyletable_{VPU}_{DATE}.parquet" + ) + thickness_by_comid = table.groupby("comid")["thickness"].unique() + + # 300 crosses the 250 threshold but not 1500 -> thickness 3. + assert thickness_by_comid[1].tolist() == [3] + # 5 is below the first (20) threshold -> thickness 1. + assert thickness_by_comid[2].tolist() == [1] diff --git a/tests/test_helper_functions.py b/tests/test_helper_functions.py new file mode 100644 index 0000000..f67aba6 --- /dev/null +++ b/tests/test_helper_functions.py @@ -0,0 +1,58 @@ +"""Unit tests pinning current behavior of helper_functions pure functions.""" + +import json +import os + +import pytest + +from geoglows_ecflow.resources.helper_functions import ( + get_date_from_forecast_dir, + get_ensemble_number_from_forecast, + get_valid_vpucode_list, + load_forecast_run, +) + + +@pytest.mark.parametrize( + "forecast_name, expected", + [ + ("1.runoff.nc", 1), + ("52.runoff.nc", 52), + ("/abs/path/to/7.runoff.nc", 7), + # Legacy RAPID-era name: the ".205.runoff.grib.runoff.netcdf" suffix + # routes to split(".")[2] for the ensemble number. + ("20230101.00.52.205.runoff.grib.runoff.netcdf", 52), + ], +) +def test_get_ensemble_number_from_forecast(forecast_name, expected): + assert get_ensemble_number_from_forecast(forecast_name) == expected + + +def test_get_valid_vpucode_list_keeps_three_digit_dirs(tmp_path): + (tmp_path / "101").mkdir() + (tmp_path / "102").mkdir() + # Non-three-digit dir and a file are both skipped. + (tmp_path / "abc").mkdir() + (tmp_path / "103.txt").write_text("not a dir") + + result = get_valid_vpucode_list(str(tmp_path)) + + assert sorted(result) == ["101", "102"] + + +def test_get_date_from_forecast_dir_parses_timestamp(): + forecast_dir = os.path.join("some", "root", "20230101.00") + + assert get_date_from_forecast_dir(forecast_dir) == "20230101.00" + + +def test_get_date_from_forecast_dir_raises_without_match(): + with pytest.raises(AttributeError): + get_date_from_forecast_dir("no-date-here") + + +def test_load_forecast_run_round_trips_manifest(tmp_path): + manifest = {"date": "2023010100", "output_dir": "/out", "job_101_52": {}} + (tmp_path / "forecast_run.json").write_text(json.dumps(manifest)) + + assert load_forecast_run(str(tmp_path)) == manifest diff --git a/tests/test_prep_river_route_forecast.py b/tests/test_prep_river_route_forecast.py new file mode 100644 index 0000000..7b6a4f3 --- /dev/null +++ b/tests/test_prep_river_route_forecast.py @@ -0,0 +1,95 @@ +"""Tests for the forecast job-manifest builder (forecast_preprocess).""" + +import json +import os + +from geoglows_ecflow.resources.prep_river_route_forecast import ( + forecast_preprocess, +) + + +def _build_cycle(tmp_path): + """Create a minimal input/runoff layout and return its paths.""" + input_dir = tmp_path / "input" + for vpu in ("101", "102"): + (input_dir / vpu).mkdir(parents=True) + output_dir = tmp_path / "output" + # The runoff dir's basename is taken verbatim as the cycle date. + runoff_dir = tmp_path / "2023010100" + runoff_dir.mkdir() + for mem in (1, 2, 52): + (runoff_dir / f"{mem}.runoff.nc").write_text("") + return input_dir, output_dir, runoff_dir + + +def test_manifest_top_level_shape(tmp_path): + input_dir, output_dir, runoff_dir = _build_cycle(tmp_path) + + master = forecast_preprocess( + str(tmp_path), str(input_dir), str(output_dir), str(runoff_dir) + ) + + assert master["date"] == "2023010100" + assert master["input_dir"] == str(input_dir) + assert master["output_dir"] == str(output_dir) + assert master["runoff_dir"] == str(runoff_dir) + + +def test_one_job_per_vpu_ensemble_pair(tmp_path): + input_dir, output_dir, runoff_dir = _build_cycle(tmp_path) + + master = forecast_preprocess( + str(tmp_path), str(input_dir), str(output_dir), str(runoff_dir) + ) + + job_keys = {k for k in master if k.startswith("job_")} + assert job_keys == { + "job_101_1", "job_101_2", "job_101_52", + "job_102_1", "job_102_2", "job_102_52", + } + + +def test_jobs_ordered_hres_first(tmp_path): + input_dir, output_dir, runoff_dir = _build_cycle(tmp_path) + + master = forecast_preprocess( + str(tmp_path), str(input_dir), str(output_dir), str(runoff_dir) + ) + + ens_order_101 = [ + int(k.rsplit("_", 1)[1]) for k in master if k.startswith("job_101_") + ] + # Largest ensemble number (HRES = 52) is scheduled first. + assert ens_order_101 == [52, 2, 1] + + +def test_job_entry_fields(tmp_path): + input_dir, output_dir, runoff_dir = _build_cycle(tmp_path) + + master = forecast_preprocess( + str(tmp_path), + str(input_dir), + str(output_dir), + str(runoff_dir), + initialize_flows=False, + ) + + job = master["job_101_52"] + assert job["vpu"] == "101" + assert job["ensemble"] == 52 + assert job["input_dir"] == str(input_dir / "101") + assert job["output_file"] == str(output_dir / "Qout_101_52.nc") + assert job["init_flows"] is False + assert job["runoff"].endswith("52.runoff.nc") + + +def test_manifest_written_to_disk(tmp_path): + input_dir, output_dir, runoff_dir = _build_cycle(tmp_path) + + master = forecast_preprocess( + str(tmp_path), str(input_dir), str(output_dir), str(runoff_dir) + ) + + written_path = os.path.join(str(tmp_path), "forecast_run.json") + with open(written_path) as f: + assert json.load(f) == master diff --git a/tests/test_run_river_route_forecast.py b/tests/test_run_river_route_forecast.py new file mode 100644 index 0000000..7652b14 --- /dev/null +++ b/tests/test_run_river_route_forecast.py @@ -0,0 +1,41 @@ +"""Tests for _find_state_init prior-cycle Qinit lookup. + +``run_river_route_forecast`` imports ``river_route`` at module scope; the +conftest Dummy shim lets this import succeed without the real dependency. +""" + +from geoglows_ecflow.resources.run_river_route_forecast import _find_state_init + +DATE = "2023010100" # base cycle: 2023-01-01 00 UTC + + +def _touch(directory, name): + path = directory / name + path.write_text("") + return str(path) + + +def test_returns_24h_lookback_when_present(tmp_path): + # 24h before the base cycle. + expected = _touch(tmp_path, "Qinit_2022123100.parquet") + + assert _find_state_init(str(tmp_path), DATE) == expected + + +def test_falls_through_to_72h_lookback(tmp_path): + # 24h and 48h missing; only the 72h-prior file exists. + expected = _touch(tmp_path, "Qinit_2022122900.parquet") + + assert _find_state_init(str(tmp_path), DATE) == expected + + +def test_seasonal_fallback_when_no_recent_qinit(tmp_path): + _touch(tmp_path, "seasonal_qinit_b.parquet") + expected = _touch(tmp_path, "seasonal_qinit_a.parquet") + + # Seasonal candidates are sorted; the lexicographically first is chosen. + assert _find_state_init(str(tmp_path), DATE) == expected + + +def test_returns_none_when_nothing_found(tmp_path): + assert _find_state_init(str(tmp_path), DATE) is None diff --git a/tests/test_suite_definition.py b/tests/test_suite_definition.py new file mode 100644 index 0000000..dcc3e66 --- /dev/null +++ b/tests/test_suite_definition.py @@ -0,0 +1,145 @@ +"""Structural smoke tests for the GEOGloWS suite definition. + +The suite definition is built entirely in memory here (no ecflow server, no +files written to disk) and the resulting node tree is asserted directly. These +pin the structure the routing workflow depends on so that changes to the +builder are caught. + +Requires the ``ecflow`` Python bindings; the module is skipped where they are +not installed. +""" + +import pytest + +pytest.importorskip("ecflow") + +from geoglows_ecflow.workflow.comfies.config import Config # noqa: E402 +from geoglows_ecflow.workflow.builders.builder import Builder # noqa: E402 + +SUITE = "geoglows_test" + +# Families that are only attached on the 00Z cycle (the full forecast +# pipeline). The 12Z cycle builds a reduced tree without them. +GATED_FAMILIES = { + "initialize", + "vpu_list", + "nc_to_zarr", + "combine_plain_table", + "combine_forecast_warnings", + "archive_qinit", + "archive_to_aws", + "web_push", +} + + +class _DictConfigSource: + """Minimal in-memory config source, so no ``.cfg`` file is needed.""" + + def __init__(self, data): + self.name = SUITE + self.origin = "" + self.data = data + + +def _make_config(mode="prod", vpu_list=("101",)): + """A minimal in-memory ``Config`` sufficient to construct the builder.""" + data = { + "name": SUITE, + "mode": mode, + "first_date": "20230101", + "last_date": "20230102", + "first_barrier": "20230101", + "exparch": "/tmp/arch", + "workroot": "/tmp/work", + "vpu_list": list(vpu_list), + "ens_members": 51, + "mars_workers": 1, + "target": {"root": "/tmp/target"}, + "jobs": { + "root": "/tmp/jobs", + "destinations": {"default": {"name": "localhost"}}, + }, + } + return Config(_DictConfigSource(data)) + + +def build_defs(mode="prod", vpu_list=("101",)): + """Build the suite definition in memory for the given run mode. + + The returned object is the ecflow ``Defs``. Building runs the builder's + own structural validation (``defs.check()``), so a malformed tree raises + here rather than producing a bad definition. + """ + builder = Builder(_make_config(mode=mode, vpu_list=vpu_list)) + builder() # runs build() and defs.check() + return builder.defs + + +def _member_family_names(defs, cycle): + """Names of the ensemble-member families on a cycle, or None if absent.""" + node = defs.find_abs_node(f"/{SUITE}/main/{cycle}/ens/ens_members") + if node is None: + return None + return sorted(n.name() for n in node.nodes) + + +def _cycle_children(defs, cycle): + node = defs.find_abs_node(f"/{SUITE}/main/{cycle}") + return [n.name() for n in node.nodes] + + +@pytest.fixture(scope="module") +def prod_defs(): + """A production-mode suite definition, built once and treated read-only.""" + return build_defs(mode="prod") + + +def test_00_cycle_builds_one_family_per_ensemble_member(prod_defs): + expected = sorted(f"101_{m:02d}" for m in range(1, 53)) + assert _member_family_names(prod_defs, "00") == expected + + +def test_12_cycle_has_no_ensemble_members(prod_defs): + assert _member_family_names(prod_defs, "12") is None + + +def test_00_cycle_builds_the_full_pipeline(prod_defs): + assert GATED_FAMILIES <= set(_cycle_children(prod_defs, "00")) + + +def test_12_cycle_skips_the_gated_families(prod_defs): + assert _cycle_children(prod_defs, "12") == ["hres", "ens", "diss"] + + +def test_nominal_time_families_carry_their_cycle_in_emos_base(prod_defs): + assert prod_defs.find_abs_node( + f"/{SUITE}/main/00" + ).find_variable("EMOS_BASE").value() == "00" + assert prod_defs.find_abs_node( + f"/{SUITE}/main/12" + ).find_variable("EMOS_BASE").value() == "12" + + +def test_ensemble_and_archive_memory_limits(prod_defs): + ens_members = prod_defs.find_abs_node( + f"/{SUITE}/main/00/ens/ens_members" + ) + archive_qinit = prod_defs.find_abs_node( + f"/{SUITE}/main/00/archive_qinit" + ) + assert ens_members.find_variable("MEM").value() == "6000" + assert archive_qinit.find_variable("MEM").value() == "4000" + + +def test_vpu_count_scales_the_member_families(): + defs = build_defs(mode="prod", vpu_list=("101", "102")) + assert len(_member_family_names(defs, "00")) == 2 * 52 + + +def test_default_minimum_comfies_version_does_not_block_construction(): + # A builder that leaves comfies_minimum_version at its "any" default must + # still construct: the version gate treats "any" as "no minimum". + class _AnyVersionBuilder(Builder): + comfies_minimum_version = "any" + + _AnyVersionBuilder(_make_config()) diff --git a/tests/test_zarr_io.py b/tests/test_zarr_io.py new file mode 100644 index 0000000..1d4c282 --- /dev/null +++ b/tests/test_zarr_io.py @@ -0,0 +1,35 @@ +"""Tests for the shared zarr-writing helper.""" + +import os + +import numpy as np +import xarray as xr + +from geoglows_ecflow.resources.zarr_io import write_dataset_to_zarr + + +def test_write_dataset_to_zarr_round_trips_and_drops_vars(tmp_path): + data = np.arange(5 * 3).reshape(5, 3).astype(float) + ds = xr.Dataset( + { + "Q": (("time", "river_id"), data), + "lat": (("river_id",), [10.0, 11.0, 12.0]), + }, + coords={"time": range(5), "river_id": [1, 2, 3]}, + ) + zarr_path = str(tmp_path / "out.zarr") + + write_dataset_to_zarr( + ds, + zarr_path, + {"time": -1, "river_id": "auto"}, + drop_vars=["lat"], + mode="w", + zarr_version=2, + ) + + # consolidated=True writes a .zmetadata file at the store root. + assert os.path.exists(os.path.join(zarr_path, ".zmetadata")) + with xr.open_zarr(zarr_path) as result: + assert "lat" not in result.variables + assert result["Q"].values.tolist() == data.tolist()