diff --git a/.github/workflows/run-weather-backfill.yml b/.github/workflows/run-weather-backfill.yml index 0f0b9fba..1e8a8e97 100644 --- a/.github/workflows/run-weather-backfill.yml +++ b/.github/workflows/run-weather-backfill.yml @@ -40,15 +40,16 @@ on: default: "latest" type: string mode: - description: "pilot = 1 station (cheap, default). full = the station×year fleet (65 stations × N years ≈ the ~28 TB reduce; needs cost sign-off; default GOES satellites cover Americas/E-Pacific stations — non-GOES shards no-op, native ring is 28-26)." + description: "pilot = 1 station, 1 monolithic task (cheap smoke test, default). station = 1+ stations sharded across ALL years (task_count = stations × years; the fast full-single-station cost/speed run, e.g. KNYC 2017..now ACMC+DSRF — NOT cost-gated, bounded by station count). full = the 65-station×year fleet (~28 TB reduce; needs cost sign-off; non-GOES shards no-op, native ring is 28-26)." required: true default: "pilot" type: choice options: - pilot + - station - full pilot_station: - description: "ICAO for the pilot (mode=pilot only), e.g. KNYC." + description: "Station ICAO(s). mode=pilot: a single ICAO (e.g. KNYC). mode=station: one OR a comma-separated list (e.g. KNYC or KNYC,KDEN) — each station is sharded across the year window." required: true default: "KNYC" type: string @@ -58,12 +59,12 @@ on: default: "ABI-L2-ACMC" type: string year_start: - description: "First year (inclusive) of the pilot backfill window (mode=pilot only)." + description: "First year (inclusive) of the backfill window (mode=pilot and mode=station)." required: true default: "2017" type: string year_end: - description: "Last year (inclusive) of the pilot backfill window (mode=pilot only). Leave empty to use the current year." + description: "Last year (inclusive) of the backfill window (mode=pilot and mode=station). Leave empty to use the current year." required: false default: "" type: string @@ -164,6 +165,60 @@ jobs: --arg ys "$FULL_YEAR_START" --arg ye "$FULL_YEAR_END" \ '["--mirror","gcp","--roster","kalshi,polymarket","--year-start",$ys,"--year-end",$ye,"--max-workers","16","--progress-bucket",$pb,"--r2-bucket",$rb]') fi + elif [ "$MODE" = "station" ]; then + # FAST FULL-SINGLE-STATION path (28-29): shard ONE (or a few) + # station(s) across ALL years so each array task is a bounded + # station-year unit that fits the 6h Spot ceiling — instead of the + # monolithic pilot (task_count=1) that blows maxRunDuration on a + # multi-year single-station run. TASK_COUNT = num_stations × num_years; + # the CLI's explicit-mode station×year sharding decomposes + # BATCH_TASK_INDEX into exactly one (station, year) per task. This is + # the path for a full-station cost/speed MEASUREMENT (e.g. KNYC + # 2017..now, ABI-L2-ACMC,ABI-L2-DSRF). Not the 28 TB fleet, so NOT + # cost-gated; it is bounded by the station count the operator passes. + YEAR_START_EFF="$YEAR_START" + YEAR_END_EFF="${YEAR_END:-$(date -u +%Y)}" + # Validate the year window BEFORE the arithmetic: reversed + # (year_end < year_start), empty, or non-numeric inputs would make + # NUM_YEARS <= 0 and yield a negative/zero TASK_COUNT + PARALLELISM + # that jq --argjson would bake into an invalid Batch job. Fail loud + # here with a clear operator message instead. + # EXACTLY four digits (not just "numeric"): a short-year typo like + # year_start=17 would otherwise make NUM_YEARS ≈ 2010 and submit + # thousands of n2-standard-8 tasks for one station (Codex P2). + case "$YEAR_START_EFF" in [0-9][0-9][0-9][0-9]) ;; *) + echo "::error::year_start must be a 4-digit year, got '$YEAR_START_EFF'"; exit 1 ;; + esac + case "$YEAR_END_EFF" in [0-9][0-9][0-9][0-9]) ;; *) + echo "::error::year_end must be a 4-digit year (or empty), got '$YEAR_END_EFF'"; exit 1 ;; + esac + if [ "$YEAR_END_EFF" -lt "$YEAR_START_EFF" ]; then + echo "::error::year_end ($YEAR_END_EFF) must be >= year_start ($YEAR_START_EFF)"; exit 1 + fi + NUM_YEARS=$(( YEAR_END_EFF - YEAR_START_EFF + 1 )) + # Count non-empty stations in the comma list (pilot_station accepts a + # comma-separated list here, e.g. "KNYC" or "KNYC,KDEN"). + NUM_STATIONS=$(printf '%s' "$PILOT_STATION" | tr ',' '\n' | grep -c '[^[:space:]]' || true) + if [ "$NUM_STATIONS" -lt 1 ]; then + echo "::error::station mode needs at least one station in pilot_station"; exit 1 + fi + TASK_COUNT=$(( NUM_STATIONS * NUM_YEARS )) + # Mid-tier machine: n2-standard-8 (8 vCPU / 32 GiB) — DSRF full-disk + # decode is CPU/mem heavy (--executor process), so bigger than the + # pilot's n2-standard-4, but one station is far cheaper than the + # 65-station fleet. Run all year-shards concurrently (parallelism + # capped at 48 to match the fleet's Spot-availability ceiling). + MACHINE_TYPE="n2-standard-8" + PARALLELISM=$(( TASK_COUNT < 48 ? TASK_COUNT : 48 )) + CPU_MILLI=8000 + MEM_MIB=32768 + # Explicit mode (satellites/products/year-window/out all passed). The + # CLI decomposes BATCH_TASK_INDEX/COUNT into this task's (station, + # year). --executor process for DSRF full-disk; --max-workers 8 matches + # the 8-vCPU machine. GOES-East (goes16) covers the US pilot stations. + COMMANDS=$(jq -nc --arg st "$PILOT_STATION" --arg ys "$YEAR_START_EFF" --arg ye "$YEAR_END_EFF" \ + --arg pb "$PROGRESS_BUCKET" --arg rb "$R2_BUCKET" --arg pr "$PRODUCTS" \ + '["--mirror","gcp","--satellites","goes16","--products",$pr,"--stations",$st,"--year-start",$ys,"--year-end",$ye,"--out","/tmp/derived","--r2-target","--r2-bucket",$rb,"--progress-bucket",$pb,"--max-workers","8","--executor","process"]') else TASK_COUNT=1 # Pilot machine — UNCHANGED (n2-standard-4, parallelism 16): a single diff --git a/packages/weather/src/mostlyright/weather/_fetchers/_goes_s3.py b/packages/weather/src/mostlyright/weather/_fetchers/_goes_s3.py index 678df7ae..df777d61 100644 --- a/packages/weather/src/mostlyright/weather/_fetchers/_goes_s3.py +++ b/packages/weather/src/mostlyright/weather/_fetchers/_goes_s3.py @@ -142,10 +142,13 @@ p: (64 * 1024 * 1024 if p == "ABI-L2-DSRF" else _DEFAULT_SIZE_CAP_BYTES) for (p, _v) in PRODUCTS } -# Accepted dataset shapes for the goes16-coarse vs goes19-full-disk DSRF split -# (D6 §A.6). The registry carries the goes19 full-disk ABI shape (5424x5424); -# goes16 DSRF ships on a coarse 326x326 lat/lon grid (see _goes_extract lat/lon -# projection branch). This split appears on the GCS mirror too. +# The coarse DSRF grid shape (D6 §A.6). The registry carries the full-disk ABI +# shape (5424x5424); goes16 DSRF ALSO shipped on this coarse 326x326 lat/lon grid +# (see the _goes_extract lat/lon projection branch) through early 2024, then +# switched to the full-disk ABI grid from ~2024 day-108 — so goes16 DSRF is +# bimodal and _validate_dataset_shape keys the accepted shape on the file's own +# projection variable, NOT on the satellite. This split appears on the GCS mirror +# too. _DSRF_GOES16_COARSE_SHAPE: tuple[int, int] = (326, 326) @@ -414,10 +417,27 @@ def _validate_dataset_shape(ds: Any, product: str, satellite: str) -> None: accepted = {expected} if product == "ABI-L2-DSRF": - # goes16 publishes DSRF on the coarse lat/lon grid; goes19 on the - # full-disk ABI grid. Accept ONLY the variant matching this file's - # satellite so a swapped grid still fails loudly. - accepted = {_DSRF_GOES16_COARSE_SHAPE} if satellite == "goes16" else {expected} + if satellite == "goes16": + # goes16 DSRF is BIMODAL: it shipped on the coarse 326x326 lat/lon + # grid (goes_lat_lon_projection) through early 2024, then switched + # to the full-disk ABI fixed grid (goes_imager_projection, + # 5424x5424) from ~2024 day-108. Key the accepted shape on the + # file's OWN projection variable — the SAME per-file branch + # _extract_from_dataset uses — so each grid must match its own + # expected shape and a mislabeled/swapped grid still fails loudly + # (no silent accept). Before this, goes16 DSRF was hard-coded to + # the coarse shape and every post-cutover full-disk file was + # rejected "before pixel read", stalling the DSRF >= mid-2024 + # backfill. + if "goes_lat_lon_projection" in ds.variables: + accepted = {_DSRF_GOES16_COARSE_SHAPE} + else: + accepted = {expected} + else: + # Other satellites (goes19) ship DSRF ONLY on the full-disk ABI + # grid — a coarse/lat-lon file here is swapped/corrupt, so keep the + # registry full-disk shape as the sole accepted variant. + accepted = {expected} if actual not in accepted: raise GoesDataCorruptError( f"{product}/{var_name} on {satellite}: shape {actual} does not " diff --git a/packages/weather/src/mostlyright/weather/satellite/__main__.py b/packages/weather/src/mostlyright/weather/satellite/__main__.py index 904d006d..0f54322a 100644 --- a/packages/weather/src/mostlyright/weather/satellite/__main__.py +++ b/packages/weather/src/mostlyright/weather/satellite/__main__.py @@ -140,14 +140,19 @@ def _build_parser() -> argparse.ArgumentParser: type=int, default=None, dest="shard_index", - help="0-based shard index for --roster (default: env BATCH_TASK_INDEX, else 0).", + help="0-based shard index for the array-job stationxyear decomposition " + "(default: env BATCH_TASK_INDEX, else 0). Applies to BOTH --roster mode " + "and explicit --stations mode (the fast single-station year-sharded path).", ) bf.add_argument( "--shard-count", type=int, default=None, dest="shard_count", - help="Total shard count for --roster (default: env BATCH_TASK_COUNT, else 1).", + help="Total shard count for the array-job stationxyear decomposition " + "(default: env BATCH_TASK_COUNT, else 1). count == stationsxyears triggers " + "the one-station-one-year decomposition; count 1 is a no-op. Applies to " + "both --roster and explicit --stations mode.", ) bf.add_argument( "--progress-bucket", @@ -412,6 +417,23 @@ def _run_backfill(args: argparse.Namespace) -> int: products = args.products stations = args.stations out = args.out + # 28-29: array-job stationxyear sharding for the EXPLICIT path too. When + # this task is one shard of an array job (shard count > 1, from + # --shard-count or the Cloud Batch BATCH_TASK_COUNT env), decompose the + # (stations x year-range) grid to THIS task's shard using the SAME + # shard_roster mapping roster mode uses. This is the fast + # full-single-station path: one station x N years submits as N parallel + # year-shards (each a bounded < 6h Spot unit) instead of one monolithic + # task that blows the maxRunDuration ceiling. Shard count == 1 (the + # classic single-task pilot / local run) is a NO-OP — shard_roster with + # count 1 returns the whole station set over the full range — so the + # explicit single-task contract stays byte-identical. + index, count = _resolve_shard_index_count(args) + if count > 1: + shard_stations, year_start, year_end = shard_roster( + tuple(stations), index, count, year_start=year_start, year_end=year_end + ) + stations = list(shard_stations) kwargs: dict = { "satellites": satellites, diff --git a/packages/weather/tests/satellite/test_cli_roster.py b/packages/weather/tests/satellite/test_cli_roster.py index bfe3d938..15104a35 100644 --- a/packages/weather/tests/satellite/test_cli_roster.py +++ b/packages/weather/tests/satellite/test_cli_roster.py @@ -387,3 +387,81 @@ def test_incremental_explicit_mode(captured, tmp_path): assert kw["year_start"] == year assert kw["year_end"] == year assert kw["resume"] is True + + +# --------------------------------------------------------------------------- +# 28-29: EXPLICIT-mode stationxyear sharding — the fast full-single-station +# path. Explicit --stations mode now honors --shard-index/--shard-count (and +# the Cloud Batch BATCH_TASK_INDEX/COUNT env) with the SAME shard_roster +# decomposition roster mode uses, so one station x N years runs as N parallel +# year-shards. Shard count == 1 stays byte-identical to the classic pilot. +# --------------------------------------------------------------------------- +def _explicit_args(station: str, year_start: int, year_end: int, *extra: str) -> list[str]: + """The required explicit-mode backfill args (bulk_backfill is mocked).""" + return [ + "backfill", + "--mirror", + "gcp", + "--satellites", + "goes16", + "--products", + "ABI-L2-ACMC,ABI-L2-DSRF", + "--stations", + station, + "--year-start", + str(year_start), + "--year-end", + str(year_end), + "--out", + "/tmp/derived", + *extra, + ] + + +def test_explicit_single_station_year_shard_decomposes(captured): + """One station x 10 years with count 10 -> shard i is (station, year_start+i).""" + rc = cli.main(_explicit_args("KNYC", 2017, 2026, "--shard-index", "3", "--shard-count", "10")) + assert rc == 0 + kw = captured[0] + # num_years = 10, count = 10 == len({KNYC}) * 10 -> stationxyear decomposition. + assert kw["stations"] == ["KNYC"] + assert kw["year_start"] == 2020 # 2017 + (3 % 10) + assert kw["year_end"] == 2020 + # Products/satellites still flow through explicit mode unchanged. + assert kw["satellites"] == ["goes16"] + assert kw["products"] == ["ABI-L2-ACMC", "ABI-L2-DSRF"] + + +def test_explicit_shard_count_one_is_noop(captured): + """No shard flags (count defaults to 1) -> full station set + full year range.""" + rc = cli.main(_explicit_args("KNYC", 2017, 2026)) + assert rc == 0 + kw = captured[0] + assert kw["stations"] == ["KNYC"] + assert kw["year_start"] == 2017 + assert kw["year_end"] == 2026 + + +def test_explicit_shard_from_batch_env(captured, monkeypatch): + """BATCH_TASK_INDEX/COUNT drive explicit-mode sharding when flags are absent.""" + monkeypatch.setenv("BATCH_TASK_INDEX", "0") + monkeypatch.setenv("BATCH_TASK_COUNT", "10") + rc = cli.main(_explicit_args("KNYC", 2017, 2026)) + assert rc == 0 + kw = captured[0] + assert kw["stations"] == ["KNYC"] + assert kw["year_start"] == 2017 # shard 0 -> year_start + assert kw["year_end"] == 2017 + + +def test_explicit_multi_station_shard(captured): + """2 stations x 2 years, count 4: shard 2 -> (second station, first year).""" + rc = cli.main( + _explicit_args("KNYC,KDEN", 2017, 2018, "--shard-index", "2", "--shard-count", "4") + ) + assert rc == 0 + kw = captured[0] + # index 2 -> station_index = 2 // 2 = 1 (KDEN), year = 2017 + (2 % 2) = 2017. + assert kw["stations"] == ["KDEN"] + assert kw["year_start"] == 2017 + assert kw["year_end"] == 2017 diff --git a/packages/weather/tests/test_satellite_s3.py b/packages/weather/tests/test_satellite_s3.py index 3f393a32..087a14dd 100644 --- a/packages/weather/tests/test_satellite_s3.py +++ b/packages/weather/tests/test_satellite_s3.py @@ -385,6 +385,46 @@ def test_shape_validation_dsrf_swapped_grids_rejected() -> None: _goes_s3._validate_dataset_shape(ds_coarse, "ABI-L2-DSRF", "goes19") +def test_shape_validation_dsrf_goes16_fulldisk_accepted() -> None: + """goes16 DSRF switched coarse-latlon -> full-disk ABI grid (~2024 day-108). + + Regression for the backfill wall: a real ``_G16_`` DSRF file on the + ``goes_imager_projection`` fixed grid at 5424x5424 (``_ds_with_var`` builds + exactly that) must PASS the gate — the extractor already routes it through + the fixed-grid branch. Before the fix the gate hard-coded goes16 DSRF to the + coarse (326, 326) shape and rejected every post-cutover file "before pixel + read", so the whole DSRF >= mid-2024 backfill failed. + """ + ds = _ds_with_var("DSR", ("y", "x"), (5424, 5424)) + _goes_s3._validate_dataset_shape(ds, "ABI-L2-DSRF", "goes16") + + +def test_shape_validation_dsrf_goes16_fulldisk_wrong_shape_rejected() -> None: + """A goes16 DSRF file on the fixed grid with a NON-full-disk shape is loud.""" + ds = _ds_with_var("DSR", ("y", "x"), (999, 999)) + with pytest.raises(GoesDataCorruptError): + _goes_s3._validate_dataset_shape(ds, "ABI-L2-DSRF", "goes16") + + +def test_shape_validation_dsrf_goes16_fulldisk_shape_on_latlon_rejected() -> None: + """A full-disk shape on a goes16 lat/lon-projection file is a mismatch -> loud. + + The accepted shape is keyed on the file's OWN projection variable, mirroring + ``_extract_from_dataset``: a ``goes_lat_lon_projection`` file must be the + coarse (326, 326) grid, so a 5424x5424 array under lat/lon metadata is a + corrupt/mislabeled file and still fails loudly (no silent accept). + """ + ds = mock.MagicMock() + var = mock.MagicMock() + var.dims = ("lat", "lon") + var.shape = (5424, 5424) + ds.variables = {"DSR": var, "goes_lat_lon_projection": mock.MagicMock()} + ds.__contains__ = lambda self, k: k in ds.variables + ds.__getitem__ = lambda self, k: ds.variables[k] + with pytest.raises(GoesDataCorruptError): + _goes_s3._validate_dataset_shape(ds, "ABI-L2-DSRF", "goes16") + + # --------------------------------------------------------------------------- # P2-3: 3D profile products (ABI-L2-LVMPC / ABI-L2-LVTPC) carry a TRAILING # pressure axis (layout ``(y, x, pressure)``) — matching the byte-faithful