From f3b25287908fbd23072469c220611c9c3755ec63 Mon Sep 17 00:00:00 2001 From: Jake Gimenes Date: Mon, 8 Jun 2026 14:00:12 -0600 Subject: [PATCH 01/14] Add workflow simplification plan --- PLAN.md | 118 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 PLAN.md diff --git a/PLAN.md b/PLAN.md new file mode 100644 index 0000000..bd0f2f5 --- /dev/null +++ b/PLAN.md @@ -0,0 +1,118 @@ +# 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. +- **Out of scope (frozen):** `geoglows_ecflow/workflow/comfies/*`. This is + vendored ECMWF framework code (Apache 2.0, ~4,700 lines). Refactoring it would + break future upstream syncs. + +## Decisions + +- **Tests:** start with unit tests on pure functions only. ecFlow-server / + suite-definition tests are deferred to a later task. +- **HRES member:** high-resolution is always ensemble member **52**. Keep this + as a single named constant (not a configurable value). The goal is to replace + scattered magic `52`/`51`/`53` literals with that one named constant. + +--- + +## Phase 1 — Pure cleanup (no behavior change) + +Small, fast, independently reviewable. Builds confidence before touching logic. + +- [ ] `builder.py`: fix `Task("dimmy")` typo (should be `"dummy"`) in the + non-osuite `run_en` branch. +- [ ] `helper_functions.py`: fix `create_logger` — when `log_file` is set, the + handler's `setLevel` / `setFormatter` / `addHandler` calls live in the + `else` branch, so file logging never attaches. Move them out so both + branches configure and register the handler. +- [ ] `builder.py`: remove duplicate imports (`complete`, `Family`, `Task` + imported multiple times) and the duplicated `nodes` import line. +- [ ] `builder.py`: remove read-but-unused config vars (`with_flood_hazard`, + `wb_days`, `ens_range`) — or wire them in if they were meant to be used + (confirm intent first). +- [ ] `builder.py`: fix stale docstring ("GLOFAS suite" → GEOGloWS). +- [ ] `generate_esri_table.py`: fix `int or str` type hint (evaluates to `int`). +- [ ] Standardize the `argparse(nargs=1)` + `args.x[0]` antipattern to plain + positional args (`day_one_forecast.py`, `netcdf_to_zarr.py`, + `archive_to_aws.py`). + +## Phase 2 — Test harness + CI + +The safety net that makes every later refactor safe. Pin **current** behavior. + +- [ ] Add `pytest` (+ `pytest-cov`) as a dev dependency in `pyproject.toml`. +- [ ] Create `tests/` with fixtures (tiny synthetic `forecast_run.json`, + a small return-period table, a minimal `Qout` netCDF). +- [ ] Unit tests for the pure functions: + - [ ] `helper_functions.get_ensemble_number_from_forecast` (incl. the + `.205.runoff.grib.runoff.netcdf` special case) + - [ ] `helper_functions.get_valid_vpucode_list` (3-digit dir filtering) + - [ ] `helper_functions.get_date_from_forecast_dir` + - [ ] `run_river_route_forecast._find_state_init` (24/48/72h lookback + + seasonal fallback) + - [ ] `prep_river_route_forecast.forecast_preprocess` (manifest shape / + job keys / HRES-first ordering) + - [ ] `day_one_forecast.check_for_return_period_flow` and + `get_time_of_first_exceedance` + - [ ] `compute_init_flows` time-index selection (`INIT_TIME_INDEX`) +- [ ] Add a GitHub Actions workflow (`.github/workflows/tests.yml`) running the + suite on the supported Python range (3.11–3.13). + +## Phase 3 — Centralize duplication + +Guarded by Phase 2 tests. + +- [ ] Extract the duplicated dask `config.set({...})` + Blosc/zstd compressor + + encoding block (currently copy-pasted in `day_one_forecast.py` and + `netcdf_to_zarr.py`) into one shared zarr-writing helper. +- [ ] Add a small `forecast_run.json` loader to remove the open-and-parse + boilerplate repeated across `run_river_route_forecast.py`, + `prep_river_route_forecast.py`, `compute_init_flows.py`, + `netcdf_to_zarr.py`, and `archive_to_aws.py`. +- [ ] Define return periods `[2, 5, 10, 25, 50, 100]` once and drive the + currently hand-unrolled ladders from it + (`day_one_forecast.check_for_return_period_flow` and the two blocks in + `generate_esri_table.py`). +- [ ] Standardize `logging` setup (formats currently differ per module). + +## Phase 4 — Make implicit explicit + +Highest value for "make implicit explicit," but touches the most files, so it +goes last. + +- [ ] Introduce a single named `HRES_ENSEMBLE_MEMBER = 52` constant and replace + the scattered literals: `grep -v Qout_..._52.nc` (3× in `nco_calc.ecf`), + `Qout_*_52.nc` in `netcdf_to_zarr.py`, and `range(1, 53)` / + `np.arange(1, 52)` / `ensemble=52` in `builder.py` and the scripts. +- [ ] Name the `EMOS_BASE != "12"` gate (repeated 3× in `builder.py`) with a + meaningful variable. +- [ ] Name the remaining magic numbers: timer offsets (`hours=7`, `hours=9`, + `"14:15"`), `MEM` values (6000/4000), stream-order threshold (`>= 3`), + flow-thickness thresholds (20/250/1500/10000/30000), and the 10-day + windows. (`compute_init_flows.INIT_TIME_INDEX` is the model to follow.) +- [ ] Consolidate the config keys currently read ad hoc via `self.config.get(...)` + in `builder.py` into one documented place, so what's tunable is visible + at a glance. + +## Follow-ups (later tasks) + +- ecFlow-server / suite-definition smoke tests (e.g. building the def in + `--dry` mode). +- README refresh (carried over from PR #27 review). From 7d2a1a590218036eeef9a11a47137c6bfc83435d Mon Sep 17 00:00:00 2001 From: Jake Gimenes Date: Mon, 8 Jun 2026 15:01:46 -0600 Subject: [PATCH 02/14] Phase 1 cleanup: fix logger bug, typo, dead imports/vars --- geoglows_ecflow/resources/archive_to_aws.py | 6 ++---- geoglows_ecflow/resources/day_one_forecast.py | 9 +++------ geoglows_ecflow/resources/generate_esri_table.py | 2 +- geoglows_ecflow/resources/helper_functions.py | 12 ++++++------ geoglows_ecflow/resources/netcdf_to_zarr.py | 3 +-- geoglows_ecflow/workflow/builders/builder.py | 15 ++++----------- 6 files changed, 17 insertions(+), 30 deletions(-) diff --git a/geoglows_ecflow/resources/archive_to_aws.py b/geoglows_ecflow/resources/archive_to_aws.py index 71560bd..4301e19 100644 --- a/geoglows_ecflow/resources/archive_to_aws.py +++ b/geoglows_ecflow/resources/archive_to_aws.py @@ -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/day_one_forecast.py b/geoglows_ecflow/resources/day_one_forecast.py index 02ec49e..a557fb8 100644 --- a/geoglows_ecflow/resources/day_one_forecast.py +++ b/geoglows_ecflow/resources/day_one_forecast.py @@ -337,27 +337,24 @@ 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 diff --git a/geoglows_ecflow/resources/generate_esri_table.py b/geoglows_ecflow/resources/generate_esri_table.py index 1047701..31f3b3c 100644 --- a/geoglows_ecflow/resources/generate_esri_table.py +++ b/geoglows_ecflow/resources/generate_esri_table.py @@ -11,7 +11,7 @@ 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( diff --git a/geoglows_ecflow/resources/helper_functions.py b/geoglows_ecflow/resources/helper_functions.py index b33e3e4..ebd3fdf 100644 --- a/geoglows_ecflow/resources/helper_functions.py +++ b/geoglows_ecflow/resources/helper_functions.py @@ -21,13 +21,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..8a540b2 100644 --- a/geoglows_ecflow/resources/netcdf_to_zarr.py +++ b/geoglows_ecflow/resources/netcdf_to_zarr.py @@ -95,8 +95,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/workflow/builders/builder.py b/geoglows_ecflow/workflow/builders/builder.py index bd27ed5..b20fdfa 100644 --- a/geoglows_ecflow/workflow/builders/builder.py +++ b/geoglows_ecflow/workflow/builders/builder.py @@ -1,9 +1,9 @@ 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, @@ -18,9 +18,6 @@ from geoglows_ecflow.workflow.comfies.partition import partition -from geoglows_ecflow.workflow.parts.nodes import Family, Task, NominalTime - - class Builder(GEOGLOWSBaseBuilder): comfies_minimum_version = "1.6.2" @@ -39,7 +36,7 @@ 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 @@ -57,14 +54,10 @@ def build(self): 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") @@ -219,7 +212,7 @@ def build(self): 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_run_en.add(Task("dummy"), Timer(tnom + TimeDelta(hours=9))) n_barrier_epilog = DummyEpilog(done=Timer("14:15")) barrier_hh.add(n_run_hr, n_run_en) From 11caefdf9d080629093f86db31ac24cb1c997eba Mon Sep 17 00:00:00 2001 From: Jake Gimenes Date: Tue, 9 Jun 2026 09:58:15 -0600 Subject: [PATCH 03/14] Phase 2: add pytest harness, unit tests, and CI --- pyproject.toml | 9 +++ tests/conftest.py | 55 ++++++++++++++ tests/test_compute_init_flows.py | 41 +++++++++++ tests/test_day_one_forecast.py | 49 +++++++++++++ tests/test_helper_functions.py | 49 +++++++++++++ tests/test_prep_river_route_forecast.py | 95 +++++++++++++++++++++++++ tests/test_run_river_route_forecast.py | 41 +++++++++++ 7 files changed, 339 insertions(+) create mode 100644 tests/conftest.py create mode 100644 tests/test_compute_init_flows.py create mode 100644 tests/test_day_one_forecast.py create mode 100644 tests/test_helper_functions.py create mode 100644 tests/test_prep_river_route_forecast.py create mode 100644 tests/test_run_river_route_forecast.py diff --git a/pyproject.toml b/pyproject.toml index 806e3c2..a400466 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,15 @@ requires-python = ">=3.11,<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..8916a10 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,55 @@ +"""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. They deliberately avoid the heavy optional +dependency ``river_route`` (not needed by the functions under test) by +injecting a Dummy module into ``sys.modules`` before any test imports +``run_river_route_forecast``. +""" + +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. Inject a +# Dummy module so the import succeeds without the real dependency installed. +if "river_route" not in sys.modules: + 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_helper_functions.py b/tests/test_helper_functions.py new file mode 100644 index 0000000..3e3f342 --- /dev/null +++ b/tests/test_helper_functions.py @@ -0,0 +1,49 @@ +"""Unit tests pinning current behavior of helper_functions pure functions.""" + +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, +) + + +@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") 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 From ec2a164dd333a2eae8c14aa50d7eba787b2efbdb Mon Sep 17 00:00:00 2001 From: Jake Gimenes Date: Wed, 10 Jun 2026 10:18:39 -0600 Subject: [PATCH 04/14] Phase 3: centralize duplicated zarr writing, json loading, return periods, logging --- .github/workflows/tests.yml | 30 ++++ geoglows_ecflow/resources/archive_to_aws.py | 10 +- .../resources/combine_esri_tables.py | 9 +- .../resources/compute_init_flows.py | 6 +- .../resources/concat_forecast_warnings.py | 9 +- geoglows_ecflow/resources/day_one_forecast.py | 136 ++++++------------ .../resources/generate_esri_table.py | 29 ++-- geoglows_ecflow/resources/helper_functions.py | 32 +++++ geoglows_ecflow/resources/netcdf_to_zarr.py | 86 ++++------- .../resources/prep_river_route_forecast.py | 8 +- .../resources/run_river_route_forecast.py | 5 +- geoglows_ecflow/resources/zarr_io.py | 67 +++++++++ tests/test_generate_esri_table.py | 62 ++++++++ tests/test_helper_functions.py | 9 ++ tests/test_zarr_io.py | 35 +++++ 15 files changed, 336 insertions(+), 197 deletions(-) create mode 100644 .github/workflows/tests.yml create mode 100644 geoglows_ecflow/resources/zarr_io.py create mode 100644 tests/test_generate_esri_table.py create mode 100644 tests/test_zarr_io.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..d4c9824 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,30 @@ +name: tests + +on: + push: + branches: [main] + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.12", "3.13"] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install package with dev dependencies + run: | + python -m pip install --upgrade pip + pip install -e ".[dev]" + + - name: Run tests + run: pytest -q --cov=geoglows_ecflow.resources --cov-report=term-missing diff --git a/geoglows_ecflow/resources/archive_to_aws.py b/geoglows_ecflow/resources/archive_to_aws.py index 4301e19..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( 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 a557fb8..69d6294 100644 --- a/geoglows_ecflow/resources/day_one_forecast.py +++ b/geoglows_ecflow/resources/day_one_forecast.py @@ -2,14 +2,17 @@ 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 def check_for_return_period_flow( @@ -17,59 +20,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): @@ -288,48 +270,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") @@ -358,7 +308,7 @@ def netcdf_forecast_record_to_zarr(record_path) -> None: 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 31f3b3c..07edc72 100644 --- a/geoglows_ecflow/resources/generate_esri_table.py +++ b/geoglows_ecflow/resources/generate_esri_table.py @@ -1,12 +1,16 @@ 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, +) + def postprocess_vpu_forecast_directory( output_dir: str, @@ -43,12 +47,8 @@ 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"][:], ) @@ -63,12 +63,8 @@ def postprocess_vpu_forecast_directory( 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 +114,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 ebd3fdf..b14a365 100644 --- a/geoglows_ecflow/resources/helper_functions.py +++ b/geoglows_ecflow/resources/helper_functions.py @@ -2,11 +2,43 @@ # 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] + +# 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 diff --git a/geoglows_ecflow/resources/netcdf_to_zarr.py b/geoglows_ecflow/resources/netcdf_to_zarr.py index 8a540b2..5fde924 100644 --- a/geoglows_ecflow/resources/netcdf_to_zarr.py +++ b/geoglows_ecflow/resources/netcdf_to_zarr.py @@ -1,21 +1,19 @@ 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 ( + 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,10 +23,9 @@ 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"))]) @@ -41,53 +38,30 @@ def netcdf_forecasts_to_zarr(workspace: str) -> None: 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, 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") - 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=52) - 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__": 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/tests/test_generate_esri_table.py b/tests/test_generate_esri_table.py new file mode 100644 index 0000000..eb2909f --- /dev/null +++ b/tests/test_generate_esri_table.py @@ -0,0 +1,62 @@ +"""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] diff --git a/tests/test_helper_functions.py b/tests/test_helper_functions.py index 3e3f342..f67aba6 100644 --- a/tests/test_helper_functions.py +++ b/tests/test_helper_functions.py @@ -1,5 +1,6 @@ """Unit tests pinning current behavior of helper_functions pure functions.""" +import json import os import pytest @@ -8,6 +9,7 @@ get_date_from_forecast_dir, get_ensemble_number_from_forecast, get_valid_vpucode_list, + load_forecast_run, ) @@ -47,3 +49,10 @@ def test_get_date_from_forecast_dir_parses_timestamp(): 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_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() From 1ccab32e69235c71bdafb14c1ba03ecef3d42c6e Mon Sep 17 00:00:00 2001 From: Jake Gimenes Date: Wed, 10 Jun 2026 18:27:13 -0600 Subject: [PATCH 05/14] Add Phase 4 handoff notes to PLAN.md --- PLAN.md | 52 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/PLAN.md b/PLAN.md index bd0f2f5..bd53da0 100644 --- a/PLAN.md +++ b/PLAN.md @@ -116,3 +116,55 @@ goes last. - ecFlow-server / suite-definition smoke tests (e.g. building the def in `--dry` mode). - README refresh (carried over from PR #27 review). + +--- + +## Handoff — current status (resume Phase 4 on Ubuntu) + +**Done & committed (branch `workflow-simplification`, on fork `JakeGimenes`):** +Phases 1–3 complete; 24 pytest tests pass. Phase 3 added `resources/zarr_io.py` +(shared `write_dataset_to_zarr` + `DASK_ZARR_CONFIG`), `helper_functions` +gained `load_forecast_run`, `RETURN_PERIODS`, and `configure_logging`. + +**Why Phase 4 moves to Ubuntu:** on the Windows dev box, `ecflow` is not +pip-installable (no wheel) and `river-route` is not on PyPI, so `builder.py` +can't even be imported and the `.ecf` scripts have no harness. The Ubuntu box +has geoglows-ecflow installed, so the suite definition *can* be built and +tested there. + +### Phase 4a — resources/ constants (verifiable; do first) +- `HRES_ENSEMBLE_MEMBER = 52` in `helper_functions.py`. In `netcdf_to_zarr.py`: + `np.arange(1, HRES_ENSEMBLE_MEMBER)`, `f"Qout_*_{HRES_ENSEMBLE_MEMBER}.nc"`, + `ensemble=HRES_ENSEMBLE_MEMBER`. +- `THICKNESS_THRESHOLDS = [20, 250, 1500, 10000, 30000]` in + `generate_esri_table.py`; drive the thickness ladder via `enumerate` + (levels 2..6). Covered by `tests/test_generate_esri_table.py` — extend it to + assert thickness too. +- `MIN_STREAM_ORDER = 3` (`day_one_forecast.py:114`), + `FORECAST_WINDOW_DAYS = 10` (`generate_esri_table.py` 10-day filter). + +### Phase 4b — builder.py + nco_calc.ecf (verify on Ubuntu, expression-preserving) +- `range(1, 53)` → `range(1, HRES_ENSEMBLE_MEMBER + 1)` (builder imports the + constant cross-package from `resources.helper_functions`). +- `EMOS_BASE != "12"` gate (×3) → helper `is_00z_cycle(node)` (gates the + full-ensemble build; `"12"` = 12Z cycle, so `!= "12"` = 00Z). +- Timers: `HRES_RUN_OFFSET_HOURS = 7`, `ENS_RUN_OFFSET_HOURS = 9`, + `BARRIER_DONE_TIME = "14:15"`. MEM: `ENS_TASK_MEM_MB = 6000`, + `ARCHIVE_QINIT_MEM_MB = 4000`. +- Consolidate `self.config.get(...)` reads into one documented block. +- `nco_calc.ecf` `grep -v ..._52.nc` (×3): wire an ecflow + `Variable("HRES_MEMBER", HRES_ENSEMBLE_MEMBER)` and reference + `%HRES_MEMBER%` — DECISION PENDING (vs. leaving `52` + a comment). Highest + risk: wrong wiring silently changes which member is excluded from the + ensemble mean. + +### To do on Ubuntu (before/with Phase 4b) +1. Pull the deferred suite-definition smoke tests forward: build the def + in-memory and assert structure (e.g. 00Z cycle builds 52 ensemble tasks, + `HRES_MEMBER` variable resolves, run timers are +7h/+9h). These guard 4b. +2. Make the `conftest.py` `river_route` shim conditional — only inject the + dummy when the real import fails — so the real package is used where present. +3. Fix CI: `pip install -e ".[dev]"` cannot resolve `river-route` on PyPI, so + the workflow will fail at install. Run `pip show river-route ecflow` to find + their real source, then either point CI at that index or install test-only + deps. (`ecflow` is also an undeclared dependency — not in `pyproject.toml`.) From c66258a48174696a36e634530e077bdfcec5ffc8 Mon Sep 17 00:00:00 2001 From: Jake Gimenes Date: Wed, 10 Jun 2026 19:39:36 -0600 Subject: [PATCH 06/14] Phase 4a: name HRES member, thickness, stream-order, and window constants; use real river_route in tests --- geoglows_ecflow/resources/day_one_forecast.py | 7 ++++++- .../resources/generate_esri_table.py | 18 +++++++++++------ geoglows_ecflow/resources/helper_functions.py | 5 +++++ geoglows_ecflow/resources/netcdf_to_zarr.py | 11 ++++++---- tests/conftest.py | 19 +++++++++++------- tests/test_generate_esri_table.py | 20 +++++++++++++++++++ 6 files changed, 62 insertions(+), 18 deletions(-) diff --git a/geoglows_ecflow/resources/day_one_forecast.py b/geoglows_ecflow/resources/day_one_forecast.py index 69d6294..2ed7cda 100644 --- a/geoglows_ecflow/resources/day_one_forecast.py +++ b/geoglows_ecflow/resources/day_one_forecast.py @@ -14,6 +14,10 @@ ) 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( largeflows_df, forecasted_flows_df, stream_order, rp_data @@ -111,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 diff --git a/geoglows_ecflow/resources/generate_esri_table.py b/geoglows_ecflow/resources/generate_esri_table.py index 07edc72..c9f628b 100644 --- a/geoglows_ecflow/resources/generate_esri_table.py +++ b/geoglows_ecflow/resources/generate_esri_table.py @@ -11,6 +11,14 @@ 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, @@ -38,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 @@ -55,11 +64,8 @@ def postprocess_vpu_forecast_directory( 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 diff --git a/geoglows_ecflow/resources/helper_functions.py b/geoglows_ecflow/resources/helper_functions.py index b14a365..1ebdcab 100644 --- a/geoglows_ecflow/resources/helper_functions.py +++ b/geoglows_ecflow/resources/helper_functions.py @@ -12,6 +12,11 @@ # 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" diff --git a/geoglows_ecflow/resources/netcdf_to_zarr.py b/geoglows_ecflow/resources/netcdf_to_zarr.py index 5fde924..934ef8f 100644 --- a/geoglows_ecflow/resources/netcdf_to_zarr.py +++ b/geoglows_ecflow/resources/netcdf_to_zarr.py @@ -8,6 +8,7 @@ import xarray as xr from geoglows_ecflow.resources.helper_functions import ( + HRES_ENSEMBLE_MEMBER, configure_logging, load_forecast_run, ) @@ -28,11 +29,11 @@ def netcdf_forecasts_to_zarr(workspace: str) -> None: 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): @@ -43,13 +44,15 @@ def netcdf_forecasts_to_zarr(workspace: str) -> None: 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, 52)) + 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") - 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") diff --git a/tests/conftest.py b/tests/conftest.py index 8916a10..711c7b1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,10 +1,11 @@ """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. They deliberately avoid the heavy optional -dependency ``river_route`` (not needed by the functions under test) by -injecting a Dummy module into ``sys.modules`` before any test imports -``run_river_route_forecast``. +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 @@ -14,10 +15,14 @@ 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. Inject a -# Dummy module so the import succeeds without the real dependency installed. +# 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: - sys.modules["river_route"] = types.ModuleType("river_route") + try: + import river_route # noqa: F401 (use the real package when present) + except ModuleNotFoundError: + sys.modules["river_route"] = types.ModuleType("river_route") @pytest.fixture diff --git a/tests/test_generate_esri_table.py b/tests/test_generate_esri_table.py index eb2909f..ace68f5 100644 --- a/tests/test_generate_esri_table.py +++ b/tests/test_generate_esri_table.py @@ -60,3 +60,23 @@ def test_return_period_ladder_assigns_expected_levels(tmp_path): 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] From e34a865079beac4aea319edb7e4352cf20627249 Mon Sep 17 00:00:00 2001 From: Jake Gimenes Date: Thu, 11 Jun 2026 14:57:59 -0600 Subject: [PATCH 07/14] Make comfies importable on Python 3.12+ and add suite definition smoke tests --- geoglows_ecflow/workflow/comfies/config.py | 25 +++- geoglows_ecflow/workflow/comfies/ooflow.py | 2 +- geoglows_ecflow/workflow/comfies/sdeploy.py | 25 ++-- tests/test_suite_definition.py | 145 ++++++++++++++++++++ 4 files changed, 186 insertions(+), 11 deletions(-) create mode 100644 tests/test_suite_definition.py 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/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()) From 078effb8d7cf76187f363caf6e36ea3df1208b0c Mon Sep 17 00:00:00 2001 From: Jake Gimenes Date: Thu, 11 Jun 2026 15:24:45 -0600 Subject: [PATCH 08/14] Replace builder magic numbers with named constants and an is_00z_cycle helper --- PLAN.md | 8 ++++++ geoglows_ecflow/workflow/builders/builder.py | 27 +++++++++++++++----- 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/PLAN.md b/PLAN.md index bd53da0..d763aff 100644 --- a/PLAN.md +++ b/PLAN.md @@ -116,6 +116,14 @@ goes last. - ecFlow-server / suite-definition smoke tests (e.g. building the def in `--dry` mode). - README refresh (carried over from PR #27 review). +- **Research (`rd`) mode is broken.** `mode='rd'` (documented in `README.md`) + is the only path with `follow_osuite=False`, and it crashes unconditionally + at `builder.py:324` (`barrier_hh.ymd`; `barrier_hh` is a `NominalTime`, which + has no `ymd`) — broken since the original 2024-07-25 authoring. That branch + is also the only place the `+7h`/`+9h`/`14:15` run timers exist, so the timer + constants can't be extracted/tested until this is resolved. Decide later: + fix research mode (needs the intended barrier-repeat wiring) or remove it (and + the timers + the `rd` choice) if unused. --- diff --git a/geoglows_ecflow/workflow/builders/builder.py b/geoglows_ecflow/workflow/builders/builder.py index b20fdfa..3ed02fd 100644 --- a/geoglows_ecflow/workflow/builders/builder.py +++ b/geoglows_ecflow/workflow/builders/builder.py @@ -16,6 +16,21 @@ 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 + + +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): @@ -251,14 +266,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}"), @@ -293,7 +308,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 @@ -323,7 +338,7 @@ def build(self): 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, @@ -348,7 +363,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) From ac780dd90eea6d2454f9804b00cacec7ad67bae4 Mon Sep 17 00:00:00 2001 From: Jake Gimenes Date: Thu, 11 Jun 2026 16:28:04 -0600 Subject: [PATCH 09/14] Run CI on a conda environment so ecflow suite tests execute; declare packaging dependency --- .github/workflows/tests.yml | 20 ++++++++++++++------ pyproject.toml | 1 + 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index d4c9824..d394a3c 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -13,18 +13,26 @@ jobs: matrix: python-version: ["3.11", "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 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 + # 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: - python-version: ${{ matrix.python-version }} + environment-file: environment.yml + create-args: >- + python=${{ matrix.python-version }} + cache-environment: true - name: Install package with dev dependencies - run: | - python -m pip install --upgrade pip - pip install -e ".[dev]" + run: pip install -e ".[dev]" - name: Run tests run: pytest -q --cov=geoglows_ecflow.resources --cov-report=term-missing diff --git a/pyproject.toml b/pyproject.toml index a400466..89a69fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ dependencies = [ "zarr>=2.16.1", "boto3>=1.28.65", "river-route>=2.1.1", + "packaging>=21.0", ] requires-python = ">=3.11,<3.14" readme = "README.md" From 5ed23858fc6e0214a34d5ae8294f1f20105fea8a Mon Sep 17 00:00:00 2001 From: Jake Gimenes Date: Fri, 12 Jun 2026 11:40:39 -0600 Subject: [PATCH 10/14] Fix CI: drop Python 3.11 (river-route needs >=3.12) and pin ecflow<5.17 --- .github/workflows/tests.yml | 2 +- PLAN.md | 5 +++++ environment.yml | 6 ++++-- pyproject.toml | 2 +- 4 files changed, 11 insertions(+), 4 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index d394a3c..0203922 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -11,7 +11,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.11", "3.12", "3.13"] + python-version: ["3.12", "3.13"] defaults: run: diff --git a/PLAN.md b/PLAN.md index d763aff..e1d97cf 100644 --- a/PLAN.md +++ b/PLAN.md @@ -116,6 +116,11 @@ goes last. - ecFlow-server / suite-definition smoke tests (e.g. building the def in `--dry` mode). - 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 is broken.** `mode='rd'` (documented in `README.md`) is the only path with `follow_osuite=False`, and it crashes unconditionally at `builder.py:324` (`barrier_hh.ymd`; `barrier_hh` is a `NominalTime`, which 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/pyproject.toml b/pyproject.toml index 89a69fa..f8f1677 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,7 @@ dependencies = [ "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" } From 15064a2ebe858591636566d61dbc87bc899266e6 Mon Sep 17 00:00:00 2001 From: Jake Gimenes Date: Fri, 12 Jun 2026 12:08:21 -0600 Subject: [PATCH 11/14] Consolidate builder config reads into one documented block; refresh PLAN --- PLAN.md | 206 +++++++------------ geoglows_ecflow/workflow/builders/builder.py | 36 ++-- 2 files changed, 85 insertions(+), 157 deletions(-) diff --git a/PLAN.md b/PLAN.md index e1d97cf..9a150c8 100644 --- a/PLAN.md +++ b/PLAN.md @@ -18,103 +18,69 @@ river-route codebase, not the RAPID `main`. - **In scope:** `geoglows_ecflow/resources/*.py`, `geoglows_ecflow/workflow/builders/builder.py`, `geoglows_ecflow/workflow/parts/*`, and the `.ecf` task scripts. -- **Out of scope (frozen):** `geoglows_ecflow/workflow/comfies/*`. This is - vendored ECMWF framework code (Apache 2.0, ~4,700 lines). Refactoring it would - break future upstream syncs. +- **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:** start with unit tests on pure functions only. ecFlow-server / - suite-definition tests are deferred to a later task. -- **HRES member:** high-resolution is always ensemble member **52**. Keep this - as a single named constant (not a configurable value). The goal is to replace - scattered magic `52`/`51`/`53` literals with that one named constant. +- **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) - -Small, fast, independently reviewable. Builds confidence before touching logic. - -- [ ] `builder.py`: fix `Task("dimmy")` typo (should be `"dummy"`) in the - non-osuite `run_en` branch. -- [ ] `helper_functions.py`: fix `create_logger` — when `log_file` is set, the - handler's `setLevel` / `setFormatter` / `addHandler` calls live in the - `else` branch, so file logging never attaches. Move them out so both - branches configure and register the handler. -- [ ] `builder.py`: remove duplicate imports (`complete`, `Family`, `Task` - imported multiple times) and the duplicated `nodes` import line. -- [ ] `builder.py`: remove read-but-unused config vars (`with_flood_hazard`, - `wb_days`, `ens_range`) — or wire them in if they were meant to be used - (confirm intent first). -- [ ] `builder.py`: fix stale docstring ("GLOFAS suite" → GEOGloWS). -- [ ] `generate_esri_table.py`: fix `int or str` type hint (evaluates to `int`). -- [ ] Standardize the `argparse(nargs=1)` + `args.x[0]` antipattern to plain - positional args (`day_one_forecast.py`, `netcdf_to_zarr.py`, - `archive_to_aws.py`). - -## Phase 2 — Test harness + CI - -The safety net that makes every later refactor safe. Pin **current** behavior. - -- [ ] Add `pytest` (+ `pytest-cov`) as a dev dependency in `pyproject.toml`. -- [ ] Create `tests/` with fixtures (tiny synthetic `forecast_run.json`, - a small return-period table, a minimal `Qout` netCDF). -- [ ] Unit tests for the pure functions: - - [ ] `helper_functions.get_ensemble_number_from_forecast` (incl. the - `.205.runoff.grib.runoff.netcdf` special case) - - [ ] `helper_functions.get_valid_vpucode_list` (3-digit dir filtering) - - [ ] `helper_functions.get_date_from_forecast_dir` - - [ ] `run_river_route_forecast._find_state_init` (24/48/72h lookback + - seasonal fallback) - - [ ] `prep_river_route_forecast.forecast_preprocess` (manifest shape / - job keys / HRES-first ordering) - - [ ] `day_one_forecast.check_for_return_period_flow` and - `get_time_of_first_exceedance` - - [ ] `compute_init_flows` time-index selection (`INIT_TIME_INDEX`) -- [ ] Add a GitHub Actions workflow (`.github/workflows/tests.yml`) running the - suite on the supported Python range (3.11–3.13). - -## Phase 3 — Centralize duplication - -Guarded by Phase 2 tests. - -- [ ] Extract the duplicated dask `config.set({...})` + Blosc/zstd compressor + - encoding block (currently copy-pasted in `day_one_forecast.py` and - `netcdf_to_zarr.py`) into one shared zarr-writing helper. -- [ ] Add a small `forecast_run.json` loader to remove the open-and-parse - boilerplate repeated across `run_river_route_forecast.py`, - `prep_river_route_forecast.py`, `compute_init_flows.py`, - `netcdf_to_zarr.py`, and `archive_to_aws.py`. -- [ ] Define return periods `[2, 5, 10, 25, 50, 100]` once and drive the - currently hand-unrolled ladders from it - (`day_one_forecast.check_for_return_period_flow` and the two blocks in - `generate_esri_table.py`). -- [ ] Standardize `logging` setup (formats currently differ per module). - -## Phase 4 — Make implicit explicit - -Highest value for "make implicit explicit," but touches the most files, so it -goes last. - -- [ ] Introduce a single named `HRES_ENSEMBLE_MEMBER = 52` constant and replace - the scattered literals: `grep -v Qout_..._52.nc` (3× in `nco_calc.ecf`), - `Qout_*_52.nc` in `netcdf_to_zarr.py`, and `range(1, 53)` / - `np.arange(1, 52)` / `ensemble=52` in `builder.py` and the scripts. -- [ ] Name the `EMOS_BASE != "12"` gate (repeated 3× in `builder.py`) with a - meaningful variable. -- [ ] Name the remaining magic numbers: timer offsets (`hours=7`, `hours=9`, - `"14:15"`), `MEM` values (6000/4000), stream-order threshold (`>= 3`), - flow-thickness thresholds (20/250/1500/10000/30000), and the 10-day - windows. (`compute_init_flows.INIT_TIME_INDEX` is the model to follow.) -- [ ] Consolidate the config keys currently read ad hoc via `self.config.get(...)` - in `builder.py` into one documented place, so what's tunable is visible - at a glance. +## 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`). +- [ ] Timer offsets (`hours=7`/`hours=9`/`"14:15"`) — **blocked**: they live + only in the broken `rd`/research-mode branch (see follow-ups). +- [x] Consolidate `self.config.get(...)` reads in `builder.py` into one + documented block. ## Follow-ups (later tasks) -- ecFlow-server / suite-definition smoke tests (e.g. building the def in - `--dry` mode). +- 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 @@ -132,52 +98,18 @@ goes last. --- -## Handoff — current status (resume Phase 4 on Ubuntu) - -**Done & committed (branch `workflow-simplification`, on fork `JakeGimenes`):** -Phases 1–3 complete; 24 pytest tests pass. Phase 3 added `resources/zarr_io.py` -(shared `write_dataset_to_zarr` + `DASK_ZARR_CONFIG`), `helper_functions` -gained `load_forecast_run`, `RETURN_PERIODS`, and `configure_logging`. - -**Why Phase 4 moves to Ubuntu:** on the Windows dev box, `ecflow` is not -pip-installable (no wheel) and `river-route` is not on PyPI, so `builder.py` -can't even be imported and the `.ecf` scripts have no harness. The Ubuntu box -has geoglows-ecflow installed, so the suite definition *can* be built and -tested there. - -### Phase 4a — resources/ constants (verifiable; do first) -- `HRES_ENSEMBLE_MEMBER = 52` in `helper_functions.py`. In `netcdf_to_zarr.py`: - `np.arange(1, HRES_ENSEMBLE_MEMBER)`, `f"Qout_*_{HRES_ENSEMBLE_MEMBER}.nc"`, - `ensemble=HRES_ENSEMBLE_MEMBER`. -- `THICKNESS_THRESHOLDS = [20, 250, 1500, 10000, 30000]` in - `generate_esri_table.py`; drive the thickness ladder via `enumerate` - (levels 2..6). Covered by `tests/test_generate_esri_table.py` — extend it to - assert thickness too. -- `MIN_STREAM_ORDER = 3` (`day_one_forecast.py:114`), - `FORECAST_WINDOW_DAYS = 10` (`generate_esri_table.py` 10-day filter). - -### Phase 4b — builder.py + nco_calc.ecf (verify on Ubuntu, expression-preserving) -- `range(1, 53)` → `range(1, HRES_ENSEMBLE_MEMBER + 1)` (builder imports the - constant cross-package from `resources.helper_functions`). -- `EMOS_BASE != "12"` gate (×3) → helper `is_00z_cycle(node)` (gates the - full-ensemble build; `"12"` = 12Z cycle, so `!= "12"` = 00Z). -- Timers: `HRES_RUN_OFFSET_HOURS = 7`, `ENS_RUN_OFFSET_HOURS = 9`, - `BARRIER_DONE_TIME = "14:15"`. MEM: `ENS_TASK_MEM_MB = 6000`, - `ARCHIVE_QINIT_MEM_MB = 4000`. -- Consolidate `self.config.get(...)` reads into one documented block. -- `nco_calc.ecf` `grep -v ..._52.nc` (×3): wire an ecflow - `Variable("HRES_MEMBER", HRES_ENSEMBLE_MEMBER)` and reference - `%HRES_MEMBER%` — DECISION PENDING (vs. leaving `52` + a comment). Highest - risk: wrong wiring silently changes which member is excluded from the - ensemble mean. - -### To do on Ubuntu (before/with Phase 4b) -1. Pull the deferred suite-definition smoke tests forward: build the def - in-memory and assert structure (e.g. 00Z cycle builds 52 ensemble tasks, - `HRES_MEMBER` variable resolves, run timers are +7h/+9h). These guard 4b. -2. Make the `conftest.py` `river_route` shim conditional — only inject the - dummy when the real import fails — so the real package is used where present. -3. Fix CI: `pip install -e ".[dev]"` cannot resolve `river-route` on PyPI, so - the workflow will fail at install. Run `pip show river-route ecflow` to find - their real source, then either point CI at that index or install test-only - deps. (`ecflow` is also an undeclared dependency — not in `pyproject.toml`.) +## 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 two +items flagged above (`nco_calc.ecf` decision, and the timer constants which are +blocked on the `rd`-mode decision). **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, the +`rd`-mode fix-or-remove decision (which unblocks the timer constants), and the +README refresh. diff --git a/geoglows_ecflow/workflow/builders/builder.py b/geoglows_ecflow/workflow/builders/builder.py index 3ed02fd..cce45e3 100644 --- a/geoglows_ecflow/workflow/builders/builder.py +++ b/geoglows_ecflow/workflow/builders/builder.py @@ -58,25 +58,21 @@ def build(self): """ 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") - - # 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) - 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") + mode = cfg.get("mode", choices=["prod", "test", "rd"]) + 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] != "/": @@ -84,6 +80,8 @@ def build(self): if mc_suite[0] != "/": mc_suite = f"/{mc_suite}" + suite = self.suite + # these flags are not user-configurable but # depend on other flags @@ -146,8 +144,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) From adee6daeeafc1fcf21f3b8be71a7eddaa71b5735 Mon Sep 17 00:00:00 2001 From: Jake Gimenes Date: Mon, 15 Jun 2026 10:09:26 -0600 Subject: [PATCH 12/14] Remove dead rd (research) mode and fix stale README entries --- PLAN.md | 33 +++---- README.md | 6 +- geoglows_ecflow/workflow/builders/builder.py | 94 +++++++------------- 3 files changed, 54 insertions(+), 79 deletions(-) diff --git a/PLAN.md b/PLAN.md index 9a150c8..943ace6 100644 --- a/PLAN.md +++ b/PLAN.md @@ -72,8 +72,9 @@ river-route codebase, not the RAPID `main`. - [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`). -- [ ] Timer offsets (`hours=7`/`hours=9`/`"14:15"`) — **blocked**: they live - only in the broken `rd`/research-mode branch (see follow-ups). +- [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. @@ -87,29 +88,29 @@ river-route codebase, not the RAPID `main`. 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 is broken.** `mode='rd'` (documented in `README.md`) - is the only path with `follow_osuite=False`, and it crashes unconditionally - at `builder.py:324` (`barrier_hh.ymd`; `barrier_hh` is a `NominalTime`, which - has no `ymd`) — broken since the original 2024-07-25 authoring. That branch - is also the only place the `+7h`/`+9h`/`14:15` run timers exist, so the timer - constants can't be extracted/tested until this is resolved. Decide later: - fix research mode (needs the intended barrier-repeat wiring) or remove it (and - the timers + the `rd` choice) if unused. +- **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 two -items flagged above (`nco_calc.ecf` decision, and the timer constants which are -blocked on the `rd`-mode decision). **33 pytest tests pass** — the resources -tests run anywhere; the suite-definition tests require `ecflow` (conda-forge). +`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, the -`rd`-mode fix-or-remove decision (which unblocks the timer constants), and the +**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/geoglows_ecflow/workflow/builders/builder.py b/geoglows_ecflow/workflow/builders/builder.py index cce45e3..630f6f8 100644 --- a/geoglows_ecflow/workflow/builders/builder.py +++ b/geoglows_ecflow/workflow/builders/builder.py @@ -6,11 +6,10 @@ 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 @@ -64,7 +63,9 @@ def build(self): # (exparch/workroot are consumed by the task scripts via templating, # so they are intentionally not read here.) suite_name = cfg.get("name") - mode = cfg.get("mode", choices=["prod", "test", "rd"]) + # 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) @@ -82,9 +83,6 @@ def build(self): suite = self.suite - # these flags are not user-configurable but - # depend on other flags - # admin family n_admin = Family("admin") n_admin_toggles = Task("toggles") @@ -105,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") @@ -184,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("dummy"), 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) @@ -240,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") @@ -331,9 +308,6 @@ 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 is_00z_cycle(main_hh): main_hh.add( n_initialize, From db193ad5fe940b7501a990ad898d43d613572983 Mon Sep 17 00:00:00 2001 From: Jake Gimenes Date: Thu, 25 Jun 2026 17:29:42 -0600 Subject: [PATCH 13/14] Replace Python .cfg deploy config with YAML gdeploy now reads a plain YAML config (yaml.safe_load) instead of executing a Python source file as config. Removes PythonConfigFile/PythonConfigPath/ _load_source, repairs the broken YAMLConfigFile, and adds YAMLConfigPath; DeployConfigFile/DeployConfigPath now subclass the YAML loaders. The Config wrapper and builder.py are unchanged (YAML loads into the same nested-dict shape). Computed values the .cfg format allowed (f-strings, variable reuse) must now be written as literals. Adds config.example.yaml template and loader tests; updates README, PLAN, and .gitignore. Touches the otherwise-frozen comfies loader by design (agreed clean break, no Python-config fallback). Verified: pytest 40 passed. Reviewed via writer-reviewer (correct/merge-ready; all actionable findings addressed). Co-Authored-By: Claude Opus 4.8 (1M context) --- .gitignore | 4 +- PLAN.md | 13 +- README.md | 130 ++++++++++---------- config.example.yaml | 64 ++++++++++ geoglows_ecflow/workflow/comfies/config.py | 90 +++++--------- geoglows_ecflow/workflow/comfies/sdeploy.py | 8 +- tests/test_config_yaml.py | 85 +++++++++++++ 7 files changed, 259 insertions(+), 135 deletions(-) create mode 100644 config.example.yaml create mode 100644 tests/test_config_yaml.py diff --git a/.gitignore b/.gitignore index 46f071f..bd2be20 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,8 @@ # geoglows_ecflow +# Deployment configs hold secrets and stay untracked; the example is tracked. *config.yml +*config.yaml +!config.example.yaml *ecflow_start.sh *data/ *log/ @@ -7,7 +10,6 @@ ecflow_home/ *geoglows_forecast.def *.ecf.* -*.cfg # pytest and coverage .coverage diff --git a/PLAN.md b/PLAN.md index 943ace6..8e99197 100644 --- a/PLAN.md +++ b/PLAN.md @@ -21,7 +21,8 @@ river-route codebase, not the RAPID `main`. - **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). + imported at all without them) and the deploy-config format switch (see + Decisions: config format). ## Decisions @@ -30,6 +31,16 @@ river-route codebase, not the RAPID `main`. 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. +- **Config format:** the deploy config is now plain **YAML** instead of the old + `.cfg` (Python source executed via `SourceFileLoader`). `gdeploy` reads + `config.yaml`; `config.example.yaml` is the committed template. This is a + clean break — the Python-config path was removed, not kept as a fallback. It + required touching the otherwise-frozen comfies loader (`config.py`, + `sdeploy.py`): the broken `YAMLConfigFile` was repaired and a `YAMLConfigPath` + added; `PythonConfigFile`/`PythonConfigPath` and `_load_source` were deleted. + The `Config` wrapper and `builder.py` are unchanged — YAML loads into the same + nested-dict shape. Computed values the Python format allowed (f-strings, + variable reuse) must now be written out as literals. --- diff --git a/README.md b/README.md index f9c8722..0d85d29 100644 --- a/README.md +++ b/README.md @@ -22,72 +22,68 @@ pip install -e . - nco>=5.1.8 - ksh>=2020.0.0 -## geoglows_ecflow configuration file (config.cfg) +## geoglows_ecflow configuration file (config.yaml) -```python - name = 'suite_name' - srcroot = "/path/to/source" - first_date = first_barrier = 'YYYYMMDD' - vpu_list = [] - mars_bond_id='251' - staticdata = '/path/to/assets' - workroot = f'/path/to/workroot' - mode = 'test' # suite mode ('test':test, 'prod':production) - expver = 'geoglows' - exparch = '/path/to/archive' - iniexparch = '/path/to/init_archive' - mars_workers = '3' - script_extension = '.ecf' - - # suite's source code - source = dict( - root = srcroot, - builder = 'geoglows_ecflow.workflow.builders.builder', - includes = 'scripts/troika:suites/scripts/tems:{includes}', - scripts = 'scripts/tems:{scripts}' - ) - - # deploy location - target = dict( - root = "/path/to/deploy_location", - ) - - # where to run computations - jobs = dict( - manager = dict( - name='troika', - ), - root = '/path/to/job_root', - limit = 26, - destinations = dict( - default = dict( - host = '%SCHOST:ab%', - bkup_host = '%SCHOST_BKUP%', - user = 'user_name', - queue = 'nf', - account = 'ECACCOUNT', - sthost = 'sthost', - ), - parallel = dict( - host = '%SCHOST:ab%', - bkup_host = '%SCHOST_BKUP%', - user = user, - queue = 'nf', - ncpus = '12', - mem = '1000', - ) - ) - ) - - # -------------------------------------------- - # Configuration of GEOGloWS software packages - # which are installed together with the suite. - # -------------------------------------------- - packages = dict( - scripts = dict( - srcdir = srcroot + 'scripts', - ), - ) +The deployment configuration is a plain YAML file. Copy +[`config.example.yaml`](config.example.yaml) to `config.yaml` and edit the +values for your environment (`config.yaml` is gitignored so secrets stay out +of version control). Values like `%SCHOST:ab%` are ecFlow variables passed +through verbatim, and `{includes}`/`{scripts}` are sdeploy search-path +placeholders. + +```yaml +name: suite_name +mode: test # 'test' or 'prod' +first_date: "YYYYMMDD" +first_barrier: "YYYYMMDD" +vpu_list: [] +ens_members: 51 +mars_workers: 3 +script_extension: ".ecf" + +expver: geoglows +exparch: /path/to/archive +iniexparch: /path/to/init_archive +staticdata: /path/to/assets +workroot: /path/to/workroot + +# suite's source code +source: + root: /path/to/source + builder: geoglows_ecflow.workflow.builders.builder + includes: "scripts/troika:suites/scripts/tems:{includes}" + scripts: "scripts/tems:{scripts}" + +# deploy location +target: + root: /path/to/deploy_location + +# where to run computations +jobs: + manager: + name: troika + root: /path/to/job_root + limit: 26 + destinations: + default: + host: "%SCHOST:ab%" + bkup_host: "%SCHOST_BKUP%" + user: user_name + queue: nf + account: ECACCOUNT + sthost: sthost + parallel: + host: "%SCHOST:ab%" + bkup_host: "%SCHOST_BKUP%" + user: user_name + queue: nf + ncpus: "12" + mem: "1000" + +# GEOGloWS software packages installed alongside the suite +packages: + scripts: + srcdir: /path/to/source/scripts ``` ## AWS configuration file (aws_config.yml) @@ -117,12 +113,12 @@ ecflow_start.sh -d /path/to/ecflow_home Generate the suite definition (via CLI or Python): ```bash -gdeploy --config /path/to/config.cfg +gdeploy --config /path/to/config.yaml ``` ```python from geoglows_ecflow.workflow.create import main -main("/path/to/config.cfg") +main("/path/to/config.yaml") ``` Start a local ecflow server, then load and begin the suite: diff --git a/config.example.yaml b/config.example.yaml new file mode 100644 index 0000000..48c5883 --- /dev/null +++ b/config.example.yaml @@ -0,0 +1,64 @@ +# Example deployment configuration for gdeploy. +# +# Copy this file to "config.yaml" (kept out of version control) and edit the +# values for your environment, then deploy the suite with: +# +# gdeploy --config /path/to/config.yaml +# +# This is plain data: every value is a literal. Unlike the old Python ".cfg" +# format there is no variable reuse or f-strings, so paths are written out in +# full. Values like "%SCHOST:ab%" are ecFlow variables and are passed through +# verbatim; "{includes}"/"{scripts}" are sdeploy search-path placeholders. + +name: suite_name +mode: test # 'test' or 'prod' +first_date: "YYYYMMDD" +first_barrier: "YYYYMMDD" +vpu_list: [] # list at least one VPU code for a real run +ens_members: 51 +mars_workers: 3 +script_extension: ".ecf" + +expver: geoglows +exparch: /path/to/archive +iniexparch: /path/to/init_archive +staticdata: /path/to/assets +workroot: /path/to/workroot + +# Suite source code +source: + root: /path/to/source + builder: geoglows_ecflow.workflow.builders.builder + includes: "scripts/troika:suites/scripts/tems:{includes}" + scripts: "scripts/tems:{scripts}" + +# Deploy location +target: + root: /path/to/deploy_location + +# Where to run computations +jobs: + manager: + name: troika + root: /path/to/job_root + limit: 26 + destinations: + default: + host: "%SCHOST:ab%" + bkup_host: "%SCHOST_BKUP%" + user: user_name + queue: nf + account: ECACCOUNT + sthost: sthost + parallel: + host: "%SCHOST:ab%" + bkup_host: "%SCHOST_BKUP%" + user: user_name + queue: nf + ncpus: "12" + mem: "1000" + +# GEOGloWS software packages installed alongside the suite +packages: + scripts: + srcdir: /path/to/source/scripts diff --git a/geoglows_ecflow/workflow/comfies/config.py b/geoglows_ecflow/workflow/comfies/config.py index ce48dbc..f1e7ca1 100644 --- a/geoglows_ecflow/workflow/comfies/config.py +++ b/geoglows_ecflow/workflow/comfies/config.py @@ -13,13 +13,11 @@ import os -import sys -import importlib.util -from importlib.machinery import SourceFileLoader import collections.abc as collections import datetime import copy +import yaml from .py2 import basestring # ----------------------------------- @@ -55,26 +53,6 @@ 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. @@ -106,6 +84,10 @@ def __init__(self, name): def _find(self, name): for dir_ in self.search_path: + if dir_ is None: + # e.g. the default search path is [os.getenv('PWD')] and PWD + # is unset; skip rather than crash in os.path.join. + continue path = os.path.join(dir_, name+self.extension) if os.path.isfile(path): return path @@ -117,57 +99,41 @@ def _load(self, path): -class PythonConfigFile(ConfigFile): +class YAMLConfigFile(ConfigFile): """ - Python config file + YAML config file. + + The file is parsed as plain data with ``yaml.safe_load``; the top level + must be a mapping of config keys (the same nested-dict shape the rest of + the config layer expects). """ - extension = '.cfg' + extension = '.yaml' + def _load(self, path): try: - old_dont_write_bytecode = sys.dont_write_bytecode - sys.dont_write_bytecode = True - data = _load_source('_sdeploy_config_'+path, path).__dict__ - sys.dont_write_bytecode = old_dont_write_bytecode + with open(path) as f: + data = yaml.safe_load(f) except IOError as e: - msg = path + ": " + e.strerror - raise ConfigLoadingError(msg) - except SyntaxError as e: - msg = path + ": " + str(e) + "\n" + e.text - raise ConfigLoadingError(msg) - except NameError as e: - msg = path + ": " + str(e) - raise ConfigLoadingError(msg) + raise ConfigLoadingError(path + ": " + e.strerror) + except yaml.YAMLError as e: + raise ConfigLoadingError(path + ": " + str(e)) + if not isinstance(data, dict): + raise ConfigLoadingError( + path + ": top-level YAML must be a mapping of config keys" + ) return data - -class PythonConfigPath(PythonConfigFile): +class YAMLConfigPath(YAMLConfigFile): """ - Like PythonConfigFile but 'name' argument - in the constructor is treated as explicit - path to config file rather than config ID. + Like YAMLConfigFile but the 'name' argument in the constructor is treated + as an explicit path to the config file rather than a config ID. """ + def _find(self, name): if os.path.isfile(name): - return os.path.join(os.environ['PWD'],name) - msg = 'Cannot find "{}"'.format(name) - raise ConfigNotFoundError(msg) - - - -class YAMLConfigFile(ConfigFile): - """ - YAML config file (not tested much..) - """ - extension = '.yaml' - def _load(self, path): - try: - f = open(path) - except IOError as e: - msg = path + ": " + e.strerror - raise ConfigLoadingError(msg) - config_data = yaml.load(f) - return data + return os.path.abspath(name) + raise ConfigNotFoundError('Cannot find "{}"'.format(name)) # --------------------------------------------------------------------- diff --git a/geoglows_ecflow/workflow/comfies/sdeploy.py b/geoglows_ecflow/workflow/comfies/sdeploy.py index 1193d19..3bc5588 100644 --- a/geoglows_ecflow/workflow/comfies/sdeploy.py +++ b/geoglows_ecflow/workflow/comfies/sdeploy.py @@ -42,8 +42,8 @@ ) from geoglows_ecflow.workflow.comfies.config import ( Config, - PythonConfigFile, - PythonConfigPath, + YAMLConfigFile, + YAMLConfigPath, ) from geoglows_ecflow.workflow.comfies.config import ConfigNotFoundError from geoglows_ecflow.workflow.comfies.config import ( @@ -412,14 +412,14 @@ def prepend_root_dirs(paths, roots): return new_paths -class DeployConfigFile(PythonConfigFile): +class DeployConfigFile(YAMLConfigFile): search_path = [ os.getcwd(), os.path.join(os.environ["HOME"], ".comfies", "sdeploy"), ] -class DeployConfigPath(PythonConfigPath): +class DeployConfigPath(YAMLConfigPath): pass diff --git a/tests/test_config_yaml.py b/tests/test_config_yaml.py new file mode 100644 index 0000000..79688d6 --- /dev/null +++ b/tests/test_config_yaml.py @@ -0,0 +1,85 @@ +"""Tests for the YAML deployment-config loader. + +These exercise the comfies config layer directly (no ecflow, no server) so they +run anywhere. They guard the switch from the old Python ".cfg" format to plain +YAML: the loader must turn a YAML mapping into the same nested-dict shape +the ``Config`` wrapper navigates, and reject input that is missing or is +not a mapping. +""" + +import pytest + +from geoglows_ecflow.workflow.comfies.config import ( + Config, + YAMLConfigFile, + YAMLConfigPath, + ConfigNotFoundError, + ConfigLoadingError, +) + +SAMPLE = """\ +name: geoglows_test +mode: prod +mars_workers: 3 +vpu_list: + - "101" + - "102" +jobs: + root: /tmp/jobs + destinations: + default: + name: localhost +""" + + +def _write(tmp_path, text, filename="config.yaml"): + path = tmp_path / filename + path.write_text(text) + return str(path) + + +def test_yaml_path_loads_scalars_and_nested_sections(tmp_path): + cfg = Config(YAMLConfigPath(_write(tmp_path, SAMPLE))) + assert cfg.get("name") == "geoglows_test" + assert cfg.get("mars_workers", type=int) == 3 + assert cfg.get("vpu_list", type=list) == ["101", "102"] + # dotted path descends into nested mappings + assert cfg.get("jobs.destinations.default.name") == "localhost" + + +def test_yaml_section_returns_subtree(tmp_path): + cfg = Config(YAMLConfigPath(_write(tmp_path, SAMPLE))) + jobs = cfg.section("jobs") + assert jobs.get("root") == "/tmp/jobs" + + +def test_config_file_finds_by_id_on_search_path(tmp_path): + # YAMLConfigFile resolves a config ID to ".yaml" on its search_path. + _write(tmp_path, SAMPLE, filename="mysuite.yaml") + + class _LocalConfig(YAMLConfigFile): + search_path = [str(tmp_path)] + + cfg = Config(_LocalConfig("mysuite")) + assert cfg.get("name") == "geoglows_test" + + +def test_missing_path_raises_not_found(tmp_path): + with pytest.raises(ConfigNotFoundError): + YAMLConfigPath(str(tmp_path / "does_not_exist.yaml")) + + +def test_empty_yaml_is_rejected(tmp_path): + # safe_load("") returns None, which is not a config mapping. + with pytest.raises(ConfigLoadingError): + YAMLConfigPath(_write(tmp_path, "")) + + +def test_non_mapping_yaml_is_rejected(tmp_path): + with pytest.raises(ConfigLoadingError): + YAMLConfigPath(_write(tmp_path, "- just\n- a\n- list\n")) + + +def test_malformed_yaml_is_rejected(tmp_path): + with pytest.raises(ConfigLoadingError): + YAMLConfigPath(_write(tmp_path, "key: [unclosed\n")) From b5a39e3562f6171b033e500462efe21a1059a25a Mon Sep 17 00:00:00 2001 From: Jake Gimenes Date: Fri, 24 Jul 2026 18:34:12 -0600 Subject: [PATCH 14/14] Document local troika job submission Add a README section and troika.example.yml describing how to run the suite locally via troika, plus an optional troika dependency. On Atos troika is provided by the system; local users install it with pip install .[troika] and point jobs.manager at a direct/local troika config. Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 42 ++++++++++++++++++++++++++++++++++++++++++ pyproject.toml | 3 +++ troika.example.yml | 10 ++++++++++ 3 files changed, 55 insertions(+) create mode 100644 troika.example.yml diff --git a/README.md b/README.md index 0d85d29..f5634a0 100644 --- a/README.md +++ b/README.md @@ -86,6 +86,48 @@ packages: srcdir: /path/to/source/scripts ``` +## Troika job submission + +The suite submits jobs through [troika](https://github.com/ecmwf/troika), ECMWF's +job-submission tool. On Atos, troika and its site configuration are provided by the +system. To run **locally**, install the optional `troika` dependency and point the +suite at a small local troika config. + +Install with the troika extra: + +```bash +pip install .[troika] +``` + +Create a local troika config (copy [`troika.example.yml`](troika.example.yml) to +`troika.yml`) that runs jobs as plain local processes. The **site name must match the +`host`** used in the config's job destinations: + +```yaml +sites: + localhost: + type: direct # run the job directly (no SLURM/PBS) + connection: local # on this machine (no ssh) +``` + +Then add `executable` and `config` to the `jobs.manager` block of your `config.yaml`: + +```yaml +jobs: + manager: + name: troika + executable: /path/to/troika # output of `which troika` + config: /path/to/troika.yml + # ... + destinations: + default: + host: localhost # must match the site name in troika.yml + user: your_user +``` + +With that, deploying and running the suite (see *Local run example*) submits every task +through troika as a local process. + ## AWS configuration file (aws_config.yml) ```yaml diff --git a/pyproject.toml b/pyproject.toml index f8f1677..57f67bd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,9 @@ readme = "README.md" license = { text = "BSD-3-Clause" } [project.optional-dependencies] +troika = [ + "troika>=0.2.0", +] dev = [ "pytest>=7.0", "pytest-cov>=4.0", diff --git a/troika.example.yml b/troika.example.yml new file mode 100644 index 0000000..c506148 --- /dev/null +++ b/troika.example.yml @@ -0,0 +1,10 @@ +# troika.example.yml -- minimal local job submission for troika. +# +# Copy this to "troika.yml" and point config.yaml's jobs.manager.config at it. +# This runs every job as a plain local process (no SLURM/PBS, no ssh), which is +# useful for testing the suite off-Atos. The site name must match the `host` +# used in the config's job destinations (see README). +sites: + localhost: + type: direct # run the job directly (no scheduler) + connection: local # on this machine (no ssh)