From a71f6c82e0db42be01196e8578e231c7c064ab91 Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Mon, 19 Oct 2026 10:45:10 +0200 Subject: [PATCH 1/3] orchestration: run the pipeline on Argo, one tile per pod An optional `orchestration` package with one subcommand per workflow task (plan, ingest-tile, compute, trace), and an example Argo WorkflowTemplate that runs each tile's ingest -> compute as a DAG of its own, so a failed tile stops only itself. jdluc runs unmodified in the pods; every cache-key argument passes straight through, and an import-linter contract keeps jdluc from importing orchestration. trace writes the emissions-factor table to a new EXPORT_ROOT, through a small `export.write_table` utility. The package is built to be wrapped rather than edited: `build_parser` and `run` let a wrapper add subcommands, `plan.get_task_to_items` returns the lists for it to add to before `write_task_to_items`, and the template's tasks can be called by templateRef. See docs/orchestration.md. Also adds the Dockerfile the workflow's image is built from. --- .dockerignore | 23 ++ .env.example | 4 + .pre-commit-config.yaml | 2 +- Dockerfile | 25 ++ README.md | 1 + docs/orchestration.md | 95 ++++++ jdluc/__tests__/trace_test.py | 36 ++ jdluc/config.py | 1 + jdluc/export.py | 62 ++++ jdluc/trace.py | 36 +- orchestration/__init__.py | 22 ++ orchestration/__main__.py | 255 ++++++++++++++ orchestration/__tests__/__init__.py | 0 orchestration/__tests__/conftest.py | 2 + orchestration/__tests__/plan_test.py | 313 ++++++++++++++++++ orchestration/plan.py | 230 +++++++++++++ .../workflows/cornerstone.example.yaml | 238 +++++++++++++ pyproject.toml | 13 +- 18 files changed, 1355 insertions(+), 3 deletions(-) create mode 100644 .dockerignore create mode 100644 Dockerfile create mode 100644 docs/orchestration.md create mode 100644 jdluc/export.py create mode 100644 orchestration/__init__.py create mode 100644 orchestration/__main__.py create mode 100644 orchestration/__tests__/__init__.py create mode 100644 orchestration/__tests__/conftest.py create mode 100644 orchestration/__tests__/plan_test.py create mode 100644 orchestration/plan.py create mode 100644 orchestration/workflows/cornerstone.example.yaml diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..40c24c3 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,23 @@ +# Claude worktrees and other local tooling state. +.claude/ + +# A host venv, at any depth, must not shadow the one built inside the image. +**/.venv/ + +# Exact-patch pin (e.g. 3.14.5) a base image's uv may not fetch; requires-python (>=3.14,<3.15) +# is satisfied by any 3.14.x, so let uv use the base image's own. +.python-version + +# Secrets: supply .env at runtime, never bake it into the image. `*.env` matches any env file +# but not the `*.env.example` templates, which stay. +.env +**/*.env + +# Caches / build noise. +**/__pycache__/ +*.pyc +.mypy_cache/ +.pytest_cache/ +.ruff_cache/ +.cache/ +.DS_Store diff --git a/.env.example b/.env.example index dff0bf2..b53768a 100644 --- a/.env.example +++ b/.env.example @@ -1,3 +1,7 @@ +# EXPORT_ROOT is where terminal deliverables land (the emissions-factor table): not a cache +# the pipeline reads back. Any fsspec URI, like the roots below. +# EXPORT_ROOT = '/absolute/path/to/export' +EXPORT_ROOT = 'gs://cornerstone-emissions-us-central1' HARVARD_DATAVERSE_GUESTBOOK_JSON = '{"email":Replace with your email address,"institution":Replace with your institution}' # INGEST_ROOT and SCRATCH_ROOT take any fsspec URI. A local root must be an absolute path # that already exists; a gs:// root needs application-default credentials. diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ade003c..4c69651 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -18,7 +18,7 @@ repos: - id: mypy name: mypy language: system - entry: uv run mypy jdluc validation + entry: uv run mypy jdluc orchestration validation pass_filenames: false - repo: https://github.com/astral-sh/uv-pre-commit rev: 0.12.7 diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..0ccee1b --- /dev/null +++ b/Dockerfile @@ -0,0 +1,25 @@ +# syntax=docker/dockerfile:1.19 +# Build from a full clone, not a git worktree: a worktree's .git is a file pointing outside the +# build context, and the image needs the repository for git provenance (utils.get_git_version). + +FROM ghcr.io/astral-sh/uv:python3.14-bookworm-slim + +WORKDIR /app + +# The rasterio/GDAL wheels vendor GDAL but still link libexpat, which the slim image omits. +# safe.directory lets git read /app when the container runs as a user other than its owner. +RUN apt-get update \ + && apt-get install -y --no-install-recommends git libexpat1 \ + && rm -rf /var/lib/apt/lists/* \ + && git config --system --add safe.directory /app + +COPY pyproject.toml uv.lock ./ +RUN uv sync --frozen --no-install-project --no-dev + +COPY --exclude=.git . . +RUN uv sync --frozen --no-dev + +# .git in its own last layer, so a commit that changes no source rebuilds only this layer. +COPY .git .git + +ENV PATH="/app/.venv/bin:${PATH}" diff --git a/README.md b/README.md index f60f5fe..fbe4e2c 100644 --- a/README.md +++ b/README.md @@ -45,6 +45,7 @@ Once you're ready to look under the hood: - `docs/data.md` — how to get the data products, grids, and full schemas for the harmonized inputs, per-pixel emissions, and emissions-factor table. - `docs/coverage.md` — the full list of countries and crops the pipeline produces emissions factors for, by ISO 3166-1 alpha-3 code. - `docs/validation.md` — how the pipeline is measured against external datasets: how the (country, crop) targets are chosen, which relationships are held as regression controls, and what such a comparison can and cannot conclude. The tooling lives in `validation/` and reports against the gaps in `docs/further_research.md`. +- `docs/orchestration.md` — an optional way to run the pipeline on Kubernetes with Argo Workflows, one tile per pod, for country runs and global backfills in hours rather than days. ## Running and contributing diff --git a/docs/orchestration.md b/docs/orchestration.md new file mode 100644 index 0000000..b4ddb43 --- /dev/null +++ b/docs/orchestration.md @@ -0,0 +1,95 @@ +# Running the pipeline on a cluster + +The pipeline runs on a single host by design (see [`architecture.md`](architecture.md#single-host-by-design)), and nothing here changes that. The `orchestration` package is an optional layer for running it on Kubernetes with [Argo Workflows](https://argoproj.github.io/workflows/), one 10° tile per pod, so that a country or a global backfill takes hours instead of days. `jdluc` runs unmodified inside the pods. Argo takes over only the driver's job: sequencing, fan-out, retries and visibility. + +## Design + +**One tile per pod, and no Dask cluster.** The pipeline processes each 10° tile independently, and one tile's working set fits one high-memory node. So the unit of distribution is a whole tile: a pod runs a tile's stages, each with its own in-process `LocalCluster`, exactly as a laptop run does. Nothing crosses between tiles except trace's final merge, over kilobyte-scale parquets. A shared Dask scheduler would add version skew between driver and workers and a scheduler that a heavy tile can destabilise, and a Dask failure would then look like a pipeline failure. Tasks are hour-long with tens of GB working sets, and they're already parallel inside the pod, so a distributed Dask cluster is the wrong granularity anyway. + +**The cache is the resume mechanism.** Every stage is wrapped in `@storage.cache_to_zarr` / `@storage.cache_to_parquet`, writing idempotently to `SCRATCH_ROOT`. Retrying anything skips finished work, so the orchestration can stay simple. A zarr counts as cached only once its `.complete` marker is written, so a pod killed mid-write leaves nothing a later pod would mistake for a result. + +**`orchestration` may not change a result.** It is one subcommand per workflow task, and every cache-key argument goes straight through to `jdluc`, with `jdluc`'s own defaults. An import-linter contract keeps `jdluc` from importing it. + +**Ingest is separate from compute.** They are different workloads: + +| | Ingest | Compute | +| --------------------- | --------------------------------- | ------------------------- | +| Bound by | network and disk | memory and CPU | +| Needs source-API keys | **yes** (USDA, Harvard Dataverse) | **no** (reads warm cache) | +| Cadence | once per country × source version | every code change | + +So only the ingest pods carry the keys, a flaky third-party source fails an ingest task rather than a compute run, and one warm `INGEST_ROOT` serves many compute runs. A laptop run ingests lazily, from inside harmonize; a workflow ingests up front instead, and `compute` runs harmonize with `skip_ingest=True`. That needs the list of datasets a run reads, which no `jdluc` constant declares, so `orchestration/plan.py` keeps one by hand. When `jdluc` starts reading a new dataset, add it there. + +## The DAG + +[`orchestration/workflows/cornerstone.example.yaml`](../orchestration/workflows/cornerstone.example.yaml) is an Argo `WorkflowTemplate`, with placeholders for the image, service account and storage class. Each run covers its `scope`: one country, or `world` for every country and GNW tile, as `--backfill` does locally. + +``` +plan -+-> ingest-world[dataset] -+-> tile[tile] --> trace + +--------------------------+ + +tile: ingest[dataset] -------------------+-> compute + ingest-allowed-failures[dataset] --+ +``` + +| Task | Pods | Runs | Keys | +| ------------------------- | -------------------- | --------------------------------------------------------------------------------------------------------- | ---- | +| `plan` | 1 | `python -m orchestration plan`: ingests the boundaries if absent, then writes the JSON lists to loop over | no | +| `ingest-world` | 1 per dataset | `jdluc/ingest.py NAME --backfill`, for each whole-world dataset | yes | +| `ingest` | 1 per dataset × tile | `python -m orchestration ingest-tile`: one dataset on one tile | yes | +| `ingest-allowed-failures` | 1 per dataset × tile | The same, for the datasets in `allow-failures` | yes | +| `compute` | 1 per tile | `python -m orchestration compute`: harmonize → emit → attribute, for each country the tile touches | no | +| `trace` | 1 | `python -m orchestration trace`: the cross-tile merge, written as the emissions-factor table | no | + +- **The tile list is derived, never typed in.** `plan` lists the tiles the scope's boundaries touch, the same derivation the `jdluc` CLIs use. +- **A failed tile fails only itself.** Each tile is a DAG of its own. A failed required ingest leaves that tile's compute unscheduled, and the other tiles carry on. +- **Ingest fails loudly, with exceptions.** A missing tile would otherwise reach harmonize's `ignore_missing_tiles` and become an all-nodata band without complaint. The datasets in `allow-failures` are the exception: by default only `DESCALS_OIL_PALM`, which is published only where oil palm was found, so most tiles have none. Compute runs whether or not those ingests fail. +- **`compute` checks its inputs.** With `--check-ingests` it raises before harmonize if a required ingest is missing from the ingest root, as a second line of defence behind the DAG. +- **`trace` needs every tile.** A table missing tiles would pass for a whole one, so with any tile failed `trace` doesn't run, and the run ends Failed. Once the failed tiles are fixed, `argo retry` reruns them, then `trace`. +- **Tiles wait for `ingest-world`.** That takes minutes once the whole-world datasets are cached. + +The flat alternative, one fan-out per stage across every tile, is simpler to read in the Argo UI, but one failed (dataset, tile) ingest then stops compute on every tile. `plan` writes the flat lists too (`ingest-tiles`, `tiles`, …), so a flat workflow needs no further code. + +## Setup + +**The image.** The repo-root `Dockerfile` builds it. Build from a full clone, not a git worktree, since the image carries `.git` for provenance (`utils.get_git_version`). Tag the image with the commit's SHA, and set the template's `image` to it. + +**Identity.** The pods need read/write access to `INGEST_ROOT`, `SCRATCH_ROOT` and `EXPORT_ROOT`. On a cloud provider, bind the template's Kubernetes service account to a cloud identity (on GKE, Workload Identity), so that no storage keys go in the image or the Secrets. + +**The two Secrets.** `Config.from_dot_env()` reads a `.env` *file*, so each pod mounts a Secret at `/app/.env`. Start from `.env.example`, keep the real file outside the repo, and fill in every field. Then make a copy with the two API-key fields present but blank, so that only the ingest pods get the keys: + +```bash +kubectl -n argo create secret generic cornerstone-env --from-file=.env=cloud.env +kubectl -n argo create secret generic cornerstone-env-nokeys --from-file=.env=cloud-nokeys.env +``` + +**Disk is what limits speed.** Ingest writes tens of GB of uncompressed GeoTIFF locally before converting it to a COG. On a node's shared boot disk, those writes can throttle the whole node to a crawl. So each pod gets its own SSD-backed generic ephemeral volume at `/localtmp`, with `TMPDIR` and `CPL_TMPDIR` pointed at it, created and deleted with the pod (on GKE, the built-in `premium-rwo` class). Don't shrink these volumes: SSD throughput scales with size. Each container also has a small `ephemeral-storage` request and limit, which covers only its own overlay, so a stray writer evicts its own pod and not the node. + +**Pod sizing.** `compute` requests 56 GiB, with peaks near 40 GiB observed: one tile per 64 GiB node, with room for system pods. The other tasks fit in 24 GiB. Spot or preemptible nodes work well, since a reclaimed node only fails its pod and `argo retry` resumes from the cache. Add their `nodeSelector` and `tolerations` under `templateDefaults`. + +## Running + +```bash +argo template create orchestration/workflows/cornerstone.example.yaml +argo submit --from workflowtemplate/cornerstone -p scope=HND --generate-name cornerstone-hnd- +``` + +Parameters: `scope` (one ISO 3166 alpha-3 code, or `world`), `methodology` (`STATISTICAL`, the default, or `JURISDICTIONAL_DIRECT`, which is US-only) and `allow-failures` (comma-separated tiled datasets, by default `DESCALS_OIL_PALM`). The template's `parallelism` and `activeDeadlineSeconds` are sized for one country; raise them for `world`. + +`trace` writes `{EXPORT_ROOT}/emissions-factors/{scope}-{methodology}.parquet`. Intermediate results stay under `INGEST_ROOT` and `SCRATCH_ROOT`. + +**Reruns are cheap.** Rerunning a country, or `argo retry` after a failure, skips everything already cached. A pod killed mid-write leaves a zarr without its `.complete` marker, and the next attempt to write it raises rather than reading back missing chunks. Once no pod is writing that zarr, delete it by hand to let it be recomputed. + +**Gotchas:** + +- **`compute` and `trace` must use the same methodology.** It is a cache-key argument, so a mismatch doesn't fail: it silently recomputes every tile inside the single `trace` pod. The template passes one `methodology` to both, so this bites only when running `python -m orchestration` by hand. +- **Storage, not compute, is the cost that recurs.** Per-tile intermediates in `SCRATCH_ROOT` are hundreds of GiB. Before a large run, decide their retention, for example with bucket lifecycle rules. +- **A warm cache can fake a fast run.** When measuring timings, make sure the run does real work. + +## Extending it + +To add tasks of your own, such as a per-tile export after `compute`, wrap the package rather than editing it: + +- **Subcommands.** `orchestration.__main__.build_parser()` returns the parser and its subparsers. Add a subcommand with `add_subparser` and the `add_scope` / `add_methodology` / `add_tile_id` helpers, then call `run(parser)`. A subcommand's flags are exactly its function's parameters. +- **Plan lists.** `orchestration.plan.get_task_to_items()` returns the lists without writing them. Add lists of your own, or keys to each `tile-units` item, then write them with `write_task_to_items()`. Call `check_allow_failures()` and `ingest_boundaries()` first, as `write_plan` does. +- **The workflow.** Call the template's tasks with `templateRef: {name: cornerstone, template: compute}` (and likewise `ingest-tile`, `trace`, …) from DAGs of your own, instead of copying the file. diff --git a/jdluc/__tests__/trace_test.py b/jdluc/__tests__/trace_test.py index aa02b59..56050e5 100644 --- a/jdluc/__tests__/trace_test.py +++ b/jdluc/__tests__/trace_test.py @@ -1,6 +1,7 @@ import numpy import pandas +from jdluc import trace from jdluc.attribute import Methodology from jdluc.datasets.worldbank_jurisdictions import AdminLevel from jdluc.trace import ( @@ -314,3 +315,38 @@ def test_methodologies_share_a_schema() -> None: assert set(jurisdictional.reset_index().columns) == set( statistical.reset_index().columns ) + + +def test_export_workflow_writes_the_table_by_aoi_and_methodology( + tmp_path, monkeypatch +) -> None: + table = attach_ratios( + df=derive_statistical_production_kg(emissions=STATISTICAL_EMISSIONS) + ) + monkeypatch.setattr(trace, "workflow", lambda **kwargs: table) + uri = trace.export_workflow( + commodity_names=("MAIZE",), + iso_3166s=("GTM", "HND"), + methodology=Methodology.STATISTICAL, + output_root=str(tmp_path), + ) + assert uri == str(tmp_path / "emissions-factors" / "GTM-HND-STATISTICAL.parquet") + written = pandas.read_parquet(uri) + pandas.testing.assert_frame_equal(written, table) + assert written.attrs["cornerstone-product-name"] == "emissions-factors" + assert written.attrs["cornerstone-source-name"] == "jdluc-trace" + + +def test_export_workflow_takes_a_name_for_its_countries(tmp_path, monkeypatch) -> None: + table = attach_ratios( + df=derive_statistical_production_kg(emissions=STATISTICAL_EMISSIONS) + ) + monkeypatch.setattr(trace, "workflow", lambda **kwargs: table) + uri = trace.export_workflow( + commodity_names=("MAIZE",), + iso_3166s=("GTM", "HND"), + methodology=Methodology.STATISTICAL, + name="world", + output_root=str(tmp_path), + ) + assert uri == str(tmp_path / "emissions-factors" / "world-STATISTICAL.parquet") diff --git a/jdluc/config.py b/jdluc/config.py index 5eab326..e455673 100644 --- a/jdluc/config.py +++ b/jdluc/config.py @@ -11,6 +11,7 @@ @dataclasses.dataclass class Config: + export_root: str harvard_dataverse_guestbook_json: str ingest_root: str number_of_dask_workers: int diff --git a/jdluc/export.py b/jdluc/export.py new file mode 100644 index 0000000..5860d5e --- /dev/null +++ b/jdluc/export.py @@ -0,0 +1,62 @@ +"""Write a pipeline step's output as a deliverable under EXPORT_ROOT. + +A utility, like `storage`: it knows nothing about the steps. A tabular deliverable (see +`trace.export_workflow`) is one parquet per set of countries, written by `write_table` to +`{EXPORT_ROOT}/{deliverable}/{name}.parquet`, with provenance in its `attrs`. + +Unlike SCRATCH_ROOT, EXPORT_ROOT is not a cache that the pipeline reads back: what lands +there is a terminal output. +""" + +import logging +import os +import tempfile + +import pandas + +from jdluc import config, storage, utils + +logger = logging.getLogger(__name__) + + +def get_metadata(product_name: str, source_name: str) -> dict[str, str]: + """Provenance, as `datasets.base` stamps on ingested COGs.""" + return { + "cornerstone-processing-time": utils.get_utc_timestamp(), + "cornerstone-processing-version": utils.get_git_version(), + "cornerstone-product-name": product_name, + "cornerstone-remote-url": utils.get_git_remote_url(), + "cornerstone-source-name": source_name, + } + + +def get_table_output_uri( + deliverable_name: str, name: str, output_root: str | None = None +) -> str: + """`{output_root}/{deliverable}/{name}.parquet`; `output_root` defaults to EXPORT_ROOT.""" + if output_root is None: + output_root = config.Config.from_dot_env().export_root + return storage.join_uri( + root=output_root, prefix=f"{deliverable_name:s}/{name:s}.parquet" + ) + + +def write_table( + df: pandas.DataFrame, + deliverable_name: str, + name: str, + source_name: str, + output_root: str | None = None, +) -> str: + """Write a tabular deliverable as one parquet, provenance in its `attrs`; return the URI.""" + uri = get_table_output_uri( + deliverable_name=deliverable_name, name=name, output_root=output_root + ) + df = df.copy() + df.attrs |= get_metadata(product_name=deliverable_name, source_name=source_name) + with tempfile.TemporaryDirectory() as tmpdir: + local_path = os.path.join(tmpdir, os.path.basename(uri)) + df.to_parquet(path=local_path) + storage.put_file(local_path=local_path, uri=uri) + logger.info(f"Wrote {deliverable_name=:s} ({len(df):d} rows) to {uri=:s}") + return uri diff --git a/jdluc/trace.py b/jdluc/trace.py index bb414b4..289a70b 100644 --- a/jdluc/trace.py +++ b/jdluc/trace.py @@ -9,6 +9,7 @@ It then derives, identically for both, each commodity's yield (kg per hectare) and emissions factor (kgCO2e per kg), and rolls the provincials up to national totals. Returns a cached pandas.DataFrame indexed by (admin_level, admin_id, commodity_name, methodology). +`export_workflow` writes that table as the `emissions-factors` deliverable under EXPORT_ROOT. The countries come from the positional ISO 3166 alpha-3 codes, or -- with `--backfill` -- from every country in the World Bank admin-0 layer. @@ -24,11 +25,14 @@ import pandas -from jdluc import attribute, emit, storage +from jdluc import attribute, emit, export, storage from jdluc.datasets import usda_nass_quickstats, worldbank_jurisdictions logger = logging.getLogger(__name__) +# Named after the published `emissions-factors.parquet` +EMISSIONS_FACTORS = "emissions-factors" + NASS_YIELD_YEARS = (2017, 2018, 2019, 2020) KG_PER_TONNE = 1000 @@ -170,6 +174,36 @@ def workflow( ) +def export_workflow( + commodity_names: tuple[str, ...], + iso_3166s: tuple[str, ...], + methodology: attribute.Methodology, + name: str | None = None, + output_root: str | None = None, +) -> str: + """Write the countries' table from the cached `workflow` to + `{EXPORT_ROOT}/emissions-factors/{name}-{METHODOLOGY}.parquet`; return the URI. + + `name` defaults to the countries' codes joined, `iso[-iso...]`; pass one for a set too large + to spell out. The methodology is in the name so the two legs over one set of countries don't + overwrite each other. + """ + df = workflow( + # NB: not a cache-key argument; with the attribute partials warm, this only reads them + concurrency=1, + commodity_names=commodity_names, + iso_3166s=iso_3166s, + methodology=methodology, + ) + return export.write_table( + df=df, + deliverable_name=EMISSIONS_FACTORS, + name=f"{name or '-'.join(iso_3166s):s}-{methodology.name:s}", + output_root=output_root, + source_name="jdluc-trace", + ) + + def main() -> int: logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", diff --git a/orchestration/__init__.py b/orchestration/__init__.py new file mode 100644 index 0000000..0d8371b --- /dev/null +++ b/orchestration/__init__.py @@ -0,0 +1,22 @@ +"""The glue between a workflow engine (Argo) and jdluc: one subcommand per workflow task. + +What may live here: listing the units of work a run fans out over, and running exactly one +unit per call. What must not: anything that changes a result. Every argument that reaches a +cache key is passed straight through to jdluc, with jdluc's own defaults. + + python -m orchestration plan HND --methodology=STATISTICAL --output-dir=/plan + python -m orchestration ingest-tile --dataset=GNW_TREE_COVER_LOSS --tile-id=10N_090W + python -m orchestration compute HND --methodology=STATISTICAL --tile-id=10N_090W + python -m orchestration trace HND --methodology=STATISTICAL + +Each subcommand that takes countries takes `world` in their place, for every country and GNW +tile, as jdluc's `--backfill` covers. + +Whole-world ingest needs no wrapper: the workflow runs jdluc's own CLI for it. A tile's ingest +has one, `ingest-tile`, so that its pod names the tile rather than the run's countries. + +To add tasks of your own, wrap rather than edit: `__main__.build_parser` returns the parser and +its subparsers to add subcommands to, and `plan.get_task_to_items` the lists to add to before +`plan.write_task_to_items`. The example in `workflows/` shows how a workflow calls each +subcommand. See docs/orchestration.md. +""" diff --git a/orchestration/__main__.py b/orchestration/__main__.py new file mode 100644 index 0000000..6daf089 --- /dev/null +++ b/orchestration/__main__.py @@ -0,0 +1,255 @@ +import argparse +import collections.abc +import logging + +import orchestration +from jdluc import ( + attribute, + config, + emit, + harmonize, + ingest, + jurisdictional_direct, + statistical, + storage, + tiling, + trace, +) +from jdluc.datasets import NAME_TO_CLS, DatasetName, worldbank_jurisdictions +from orchestration import plan + +logger = logging.getLogger(__name__) + + +def ingest_tile(dataset: DatasetName, tile_id: str) -> None: + """One ten-degree dataset on one tile, into the ingest root. Raises if the tile fails.""" + (path_or_exc,) = ingest.workflow( + concurrency=1, + dataset=NAME_TO_CLS[dataset], + overwrite=False, + root=config.Config.from_dot_env().ingest_root, + tile_ids=(tile_id,), + ).values() + if isinstance(path_or_exc, Exception): + raise path_or_exc + logger.info(f"ingested {dataset.name:s} {tile_id=:s} to {path_or_exc:s}") + + +def get_missing_ingests( + allow_failures: tuple[str, ...], methodology: attribute.Methodology, tile_id: str +) -> list[str]: + """The URIs of the tile's inputs absent from the ingest root: its ten-degree datasets + other than ``allow_failures``, and every whole-world dataset.""" + root = config.Config.from_dot_env().ingest_root + name_and_tile_ids = [ + *( + (name, tile_id) + for name in plan.get_dataset_names( + methodology=methodology, + partitioning=tiling.Partitioning.TEN_DEGREE_TILE, + ) + if name.name not in allow_failures + ), + *( + (name, tiling.WHOLE_WORLD_TILE_ID) + for name in plan.get_dataset_names( + methodology=methodology, partitioning=tiling.Partitioning.WHOLE_WORLD + ) + ), + ] + uris = [ + storage.join_uri( + prefix=NAME_TO_CLS[name].get_prefix(tile_id=tile_id), root=root + ) + for name, tile_id in name_and_tile_ids + ] + return [uri for uri in uris if not storage.path_exists(uri=uri)] + + +def compute( + allow_failures: tuple[str, ...], + check_ingests: bool, + methodology: attribute.Methodology, + scope: tuple[str, ...], + tile_id: str, +) -> None: + """harmonize -> emit -> attribute for one tile, in one process. + + harmonize goes first with ``skip_ingest=True``: the later stages call it again with + ``skip_ingest=False``, and the cache key ignores the flag, so warming it here is what + keeps source credentials off compute pods. + + With ``check_ingests``, raises first if any input other than ``allow_failures`` is + missing. harmonize would otherwise read it as an all-nodata band, and cache that. + """ + if check_ingests: + missing = get_missing_ingests( + allow_failures=allow_failures, methodology=methodology, tile_id=tile_id + ) + if missing: + raise FileNotFoundError(f"{tile_id=:s} is missing ingests: {missing}") + for stack in harmonize.Stack: + logger.info(f"harmonize {stack.name:s} {tile_id=:s}") + harmonize.workflow( + dataset_names=stack.value, + ignore_missing_tiles=True, + skip_ingest=True, + tile_id=tile_id, + tile_resolution=tiling.TileResolution.GLAD, + ) + logger.info(f"emit {tile_id=:s}") + emit.workflow(tile_id=tile_id) + + workflow_for_tile = { + attribute.Methodology.JURISDICTIONAL_DIRECT: jurisdictional_direct.workflow, + attribute.Methodology.STATISTICAL: statistical.workflow, + }[methodology] + # A tile can straddle several of the run's countries; trace sums exactly these + # per-(tile, country) partials, so compute has to produce all of them + for iso_3166 in get_iso_3166s_touching_tile( + iso_3166s=plan.get_iso_3166s(scope=scope), tile_id=tile_id + ): + logger.info(f"attribute {methodology.name:s} {tile_id=:s} {iso_3166=:s}") + workflow_for_tile( + commodity_names=attribute.get_commodity_names(methodology=methodology), + iso_3166=iso_3166, + tile_id=tile_id, + ) + + +def get_iso_3166s_touching_tile( + iso_3166s: collections.abc.Iterable[str], tile_id: str +) -> list[str]: + return sorted( + iso_3166 + for iso_3166 in iso_3166s + if tile_id + in worldbank_jurisdictions.get_ten_degree_tile_ids_for_iso_3166s( + iso_3166s=(iso_3166,) + ) + ) + + +def trace_countries(methodology: attribute.Methodology, scope: tuple[str, ...]) -> None: + """The cross-tile merge, written as the emissions-factor table. Cheap only if compute ran + with the same methodology.""" + uri = trace.export_workflow( + commodity_names=attribute.get_commodity_names(methodology=methodology), + iso_3166s=plan.get_iso_3166s(scope=scope), + methodology=methodology, + name=plan.get_name(scope=scope), + ) + logger.info(f"exported emissions factors to {uri:s}") + + +def add_subparser( + subparsers: argparse._SubParsersAction[argparse.ArgumentParser], + name: str, + func: collections.abc.Callable[..., None], +) -> argparse.ArgumentParser: + """A subcommand that calls ``func`` with its flags, which must be exactly ``func``'s + parameters.""" + subparser = subparsers.add_parser(name, help=(func.__doc__ or "").split("\n")[0]) + subparser.set_defaults(func=func) + return subparser + + +def add_scope(subparser: argparse.ArgumentParser) -> None: + subparser.add_argument( + "scope", + help=f"ISO 3166 alpha-3 codes, or {plan.WORLD:s} for every country and GNW tile", + nargs=argparse.ONE_OR_MORE, + type=plan.scope_str, + ) + + +def add_methodology(subparser: argparse.ArgumentParser) -> None: + subparser.add_argument( + "--methodology", + choices=list(attribute.Methodology), + metavar="{" + ",".join(m.name for m in attribute.Methodology) + "}", + required=True, + type=attribute.Methodology.__getitem__, + ) + + +def add_tile_id(subparser: argparse.ArgumentParser) -> None: + subparser.add_argument("--tile-id", required=True) + + +def add_allow_failures(subparser: argparse.ArgumentParser) -> None: + subparser.add_argument( + "--allow-failures", + default=(), + help="comma-separated ten-degree datasets whose failed tiles must not stop the run", + type=lambda s: tuple(name for name in s.split(",") if name), + ) + + +def build_parser() -> tuple[ + argparse.ArgumentParser, argparse._SubParsersAction[argparse.ArgumentParser] +]: + """The CLI, and its subparsers for a wrapper to add subcommands to before `run`.""" + parser = argparse.ArgumentParser( + description=orchestration.__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + subparsers = parser.add_subparsers(required=True) + + subparser = add_subparser(subparsers, "plan", plan.write_plan) + add_scope(subparser) + add_methodology(subparser) + add_allow_failures(subparser) + subparser.add_argument("--output-dir", required=True) + + subparser = add_subparser(subparsers, "ingest-tile", ingest_tile) + subparser.add_argument( + "--dataset", + choices=list(DatasetName), + metavar="DATASET", + required=True, + # NB: not `DatasetName.__getitem__`, which a StrEnum resolves to `str.__getitem__` + type=lambda s: DatasetName[s], + ) + add_tile_id(subparser) + + subparser = add_subparser(subparsers, "compute", compute) + add_scope(subparser) + add_methodology(subparser) + add_tile_id(subparser) + add_allow_failures(subparser) + subparser.add_argument( + "--check-ingests", + action="store_true", + help="fail if an input other than --allow-failures is missing from the ingest root", + ) + + subparser = add_subparser(subparsers, "trace", trace_countries) + add_scope(subparser) + add_methodology(subparser) + + return parser, subparsers + + +def run( + parser: argparse.ArgumentParser, args: collections.abc.Sequence[str] | None = None +) -> int: + """Parse ``args`` (by default the command line) and call the chosen subcommand.""" + logging.basicConfig( + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + level=logging.INFO, + ) + # Each subcommand's flags are exactly its function's parameters + kwargs = vars(parser.parse_args(args)) + func = kwargs.pop("func") + func(**{k: tuple(v) if isinstance(v, list) else v for k, v in kwargs.items()}) + return 0 + + +def main() -> int: + parser, _ = build_parser() + return run(parser) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/orchestration/__tests__/__init__.py b/orchestration/__tests__/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/orchestration/__tests__/conftest.py b/orchestration/__tests__/conftest.py new file mode 100644 index 0000000..139f566 --- /dev/null +++ b/orchestration/__tests__/conftest.py @@ -0,0 +1,2 @@ +# Unit tests here, as in jdluc, need no .env and must never touch real storage +from jdluc.__tests__.conftest import config_without_dot_env # noqa: F401 diff --git a/orchestration/__tests__/plan_test.py b/orchestration/__tests__/plan_test.py new file mode 100644 index 0000000..3773942 --- /dev/null +++ b/orchestration/__tests__/plan_test.py @@ -0,0 +1,313 @@ +import json +import pathlib +import re +import sys +import types + +import pytest + +from jdluc import attribute, datasets, tiling, trace +from jdluc.datasets import worldbank_jurisdictions +from orchestration import __main__ as cli +from orchestration import plan + +TILE_IDS = {"10N_010W", "20N_010W", "30N_010W"} +# A stand-in for every GNW tile: more than the countries' tiles, as the real list is +WORLD_TILE_IDS = ("00N_010W", *sorted(TILE_IDS), "40N_010W") + + +@pytest.fixture +def ingested(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, str]]: + """Stub the boundaries and every dataset; return the (dataset, tile) pairs ingested.""" + calls: list[tuple[str, str]] = [] + monkeypatch.setattr( + worldbank_jurisdictions, + "get_ten_degree_tile_ids_for_iso_3166s", + lambda iso_3166s: TILE_IDS, + ) + monkeypatch.setattr( + worldbank_jurisdictions, "get_all_iso_3166s", lambda: {"AAA", "BBB"} + ) + monkeypatch.setattr(tiling, "GLOBAL_NATURE_WATCH_TILE_IDS", WORLD_TILE_IDS) + for name, cls in list(datasets.NAME_TO_CLS.items()): + + def ingest_a_tile( + overwrite: bool, root: str, tile_id: str, name: str = name.name + ) -> str: + calls.append((name, tile_id)) + return f"{tile_id:s}.tif" + + monkeypatch.setitem( + datasets.NAME_TO_CLS, + name, + types.SimpleNamespace( + get_prefix=lambda tile_id, name=name.name: f"{name:s}/{tile_id:s}", + ingest_a_tile=ingest_a_tile, + partitioning=cls.partitioning, + ), + ) + return calls + + +def write_plan( + tmp_path: pathlib.Path, + allow_failures: tuple[str, ...], + scope: tuple[str, ...] = ("AAA",), +) -> dict: + plan.write_plan( + allow_failures=allow_failures, + methodology=attribute.Methodology.STATISTICAL, + output_dir=str(tmp_path), + scope=scope, + ) + return {p.stem: json.loads(p.read_text()) for p in tmp_path.glob("*.json")} + + +def test_plan_ingests_the_boundaries_it_reads( + ingested: list[tuple[str, str]], tmp_path: pathlib.Path +) -> None: + task_to_items = write_plan(tmp_path=tmp_path, allow_failures=()) + assert sorted(ingested) == [ + ("WORLD_BANK_ADMIN_0", tiling.WHOLE_WORLD_TILE_ID), + ("WORLD_BANK_ADMIN_1", tiling.WHOLE_WORLD_TILE_ID), + ] + assert "WORLD_BANK_ADMIN_0" not in task_to_items["ingest-world"] + + +def test_allowed_failures_get_their_own_list( + ingested: list[tuple[str, str]], tmp_path: pathlib.Path +) -> None: + task_to_items = write_plan(tmp_path=tmp_path, allow_failures=("DESCALS_OIL_PALM",)) + allowed = task_to_items["ingest-tiles-allowed-failures"] + assert {u["dataset"] for u in allowed} == {"DESCALS_OIL_PALM"} + assert {u["tile_id"] for u in allowed} == TILE_IDS + assert "DESCALS_OIL_PALM" not in { + u["dataset"] for u in task_to_items["ingest-tiles"] + } + + +def test_an_unknown_allowed_failure_fails( + ingested: list[tuple[str, str]], tmp_path: pathlib.Path +) -> None: + with pytest.raises(AssertionError, match="DESCALS_OIL_PLAM"): + write_plan(tmp_path=tmp_path, allow_failures=("DESCALS_OIL_PLAM",)) + + +@pytest.mark.parametrize("scope", [("AAA",), (plan.WORLD,)]) +def test_each_planned_unit_ingests_exactly_its_tile( + ingested: list[tuple[str, str]], + monkeypatch: pytest.MonkeyPatch, + scope: tuple[str, ...], + tmp_path: pathlib.Path, +) -> None: + task_to_items = write_plan(tmp_path=tmp_path, allow_failures=(), scope=scope) + for unit in task_to_items["ingest-tiles"]: + ingested.clear() + monkeypatch.setattr( + sys, + "argv", + [ + "orchestration", + "ingest-tile", + f"--dataset={unit['dataset']:s}", + f"--tile-id={unit['tile_id']:s}", + ], + ) + assert cli.main() == 0 + assert ingested == [(unit["dataset"], unit["tile_id"])] + + +def test_a_failed_tile_fails_its_ingest(monkeypatch: pytest.MonkeyPatch) -> None: + def ingest_a_tile(overwrite: bool, root: str, tile_id: str) -> str: + raise RuntimeError(f"no {tile_id:s}") + + monkeypatch.setitem( + datasets.NAME_TO_CLS, + datasets.DatasetName.DESCALS_OIL_PALM, + types.SimpleNamespace( + ingest_a_tile=ingest_a_tile, + partitioning=tiling.Partitioning.TEN_DEGREE_TILE, + ), + ) + with pytest.raises(RuntimeError, match="no 10N_010W"): + cli.ingest_tile( + dataset=datasets.DatasetName.DESCALS_OIL_PALM, tile_id="10N_010W" + ) + + +def test_world_plans_every_gnw_tile( + ingested: list[tuple[str, str]], tmp_path: pathlib.Path +) -> None: + task_to_items = write_plan( + tmp_path=tmp_path, allow_failures=(), scope=(plan.WORLD,) + ) + assert [u["tile_id"] for u in task_to_items["tiles"]] == list(WORLD_TILE_IDS) + assert {u["tile_id"] for u in task_to_items["ingest-tiles"]} == set(WORLD_TILE_IDS) + + +@pytest.mark.parametrize("scope", [(), (plan.WORLD, "AAA")]) +def test_a_scope_is_countries_or_world_alone(scope: tuple[str, ...]) -> None: + with pytest.raises(AssertionError): + plan.get_tile_ids(scope=scope) + + +def test_world_names_its_run_wide_outputs_world( + ingested: list[tuple[str, str]], monkeypatch: pytest.MonkeyPatch +) -> None: + calls: list[dict] = [] + + def record(**kwargs: object) -> str: + calls.append(kwargs) + return "uri" + + monkeypatch.setattr(trace, "export_workflow", record) + monkeypatch.setattr( + sys, "argv", ["orchestration", "trace", "world", "--methodology=STATISTICAL"] + ) + assert cli.main() == 0 + + (trace_kwargs,) = calls + assert trace_kwargs["name"] == "world" + assert trace_kwargs["iso_3166s"] == ("AAA", "BBB") + + +def test_a_country_run_names_its_outputs_by_its_countries() -> None: + assert plan.get_name(scope=("HND", "GTM")) == "GTM-HND" + + +def test_get_task_to_items_ingests_nothing_and_matches_the_plan( + ingested: list[tuple[str, str]], tmp_path: pathlib.Path +) -> None: + task_to_items = plan.get_task_to_items( + allow_failures=(), + methodology=attribute.Methodology.STATISTICAL, + scope=("AAA",), + ) + assert ingested == [] + assert task_to_items == write_plan(tmp_path=tmp_path, allow_failures=()) + + +def test_a_wrapper_adds_a_list_and_a_subcommand( + ingested: list[tuple[str, str]], tmp_path: pathlib.Path +) -> None: + """The extension points a wrapper builds on: a list added to the plan, and a subcommand + added to the CLI, with no edit to this package.""" + task_to_items = plan.get_task_to_items( + allow_failures=(), + methodology=attribute.Methodology.STATISTICAL, + scope=("AAA",), + ) + task_to_items["extra-tiles"] = [u["tile_id"] for u in task_to_items["tiles"]] + plan.write_task_to_items(output_dir=str(tmp_path), task_to_items=task_to_items) + assert json.loads((tmp_path / "extra-tiles.json").read_text()) == sorted(TILE_IDS) + + calls: list[dict] = [] + + def extra(scope: tuple[str, ...], tile_id: str) -> None: + calls.append({"scope": scope, "tile_id": tile_id}) + + parser, subparsers = cli.build_parser() + subparser = cli.add_subparser(subparsers, "extra", extra) + cli.add_scope(subparser) + cli.add_tile_id(subparser) + assert cli.run(parser, ["extra", "AAA", "--tile-id=10N_010W"]) == 0 + assert calls == [{"scope": ("AAA",), "tile_id": "10N_010W"}] + + +@pytest.mark.parametrize("scope", [("AAA",), (plan.WORLD,)]) +def test_tile_units_hold_the_flat_lists_by_tile( + ingested: list[tuple[str, str]], scope: tuple[str, ...], tmp_path: pathlib.Path +) -> None: + task_to_items = write_plan( + tmp_path=tmp_path, + allow_failures=("DESCALS_OIL_PALM",), + scope=scope, + ) + tile_units = task_to_items["tile-units"] + assert [u["tile_id"] for u in tile_units] == [ + u["tile_id"] for u in task_to_items["tiles"] + ] + for key, task in ( + ("ingests", "ingest-tiles"), + ("ingests_allowed_failures", "ingest-tiles-allowed-failures"), + ): + assert sorted( + (dataset, u["tile_id"]) for u in tile_units for dataset in u[key] + ) == sorted((u["dataset"], u["tile_id"]) for u in task_to_items[task]) + + +def test_each_tile_lists_each_tiled_dataset_once( + ingested: list[tuple[str, str]], tmp_path: pathlib.Path +) -> None: + task_to_items = write_plan(tmp_path=tmp_path, allow_failures=("DESCALS_OIL_PALM",)) + tiled = { + name.name + for name in plan.get_dataset_names( + methodology=attribute.Methodology.STATISTICAL, + partitioning=tiling.Partitioning.TEN_DEGREE_TILE, + ) + } + for u in task_to_items["tile-units"]: + assert u["ingests_allowed_failures"] == ["DESCALS_OIL_PALM"] + assert sorted([*u["ingests"], *u["ingests_allowed_failures"]]) == sorted(tiled) + + +@pytest.fixture +def present_uris( + ingested: list[tuple[str, str]], monkeypatch: pytest.MonkeyPatch +) -> set[str]: + """Every input of tile 10N_010W, as the ingest root would hold them; remove some to + make them missing.""" + root = cli.config.Config.from_dot_env().ingest_root + uris = { + cli.storage.join_uri(prefix=f"{name.name:s}/{tile_id:s}", root=root) + for partitioning, tile_id in ( + (tiling.Partitioning.TEN_DEGREE_TILE, "10N_010W"), + (tiling.Partitioning.WHOLE_WORLD, tiling.WHOLE_WORLD_TILE_ID), + ) + for name in plan.get_dataset_names( + methodology=attribute.Methodology.STATISTICAL, partitioning=partitioning + ) + } + monkeypatch.setattr(cli.storage, "path_exists", lambda uri: uri in uris) + return uris + + +def test_compute_checks_ingests_before_harmonize( + monkeypatch: pytest.MonkeyPatch, present_uris: set[str] +) -> None: + (missing,) = [uri for uri in present_uris if "GPW_GRASSLAND/" in uri] + present_uris.remove(missing) + + def harmonize(**kwargs: object) -> None: + raise AssertionError("harmonize ran on a tile missing an ingest") + + monkeypatch.setattr(cli.harmonize, "workflow", harmonize) + with pytest.raises(FileNotFoundError, match=re.escape(missing)): + cli.compute( + allow_failures=("DESCALS_OIL_PALM",), + check_ingests=True, + methodology=attribute.Methodology.STATISTICAL, + scope=("AAA",), + tile_id="10N_010W", + ) + + +def test_an_allowed_failure_is_not_missing(present_uris: set[str]) -> None: + present_uris -= {uri for uri in present_uris if "DESCALS_OIL_PALM/" in uri} + assert ( + cli.get_missing_ingests( + allow_failures=("DESCALS_OIL_PALM",), + methodology=attribute.Methodology.STATISTICAL, + tile_id="10N_010W", + ) + == [] + ) + assert ( + cli.get_missing_ingests( + allow_failures=(), + methodology=attribute.Methodology.STATISTICAL, + tile_id="10N_010W", + ) + != [] + ) diff --git a/orchestration/plan.py b/orchestration/plan.py new file mode 100644 index 0000000..964a17e --- /dev/null +++ b/orchestration/plan.py @@ -0,0 +1,230 @@ +"""What a run consists of, as the JSON lists the workflow loops over. + +The dataset lists are a hand-kept claim about what jdluc reads: no jdluc constant declares it +all. A laptop run ingests lazily, from inside harmonize, so it never needs them. A workflow +ingests up front instead, so that compute runs with ``skip_ingest`` and a missing dataset fails +at ingest rather than surfacing as an all-nodata band (``ignore_missing_tiles``). In particular +``faostat.load`` and ``usda_nass_quickstats.load`` read the ingest root directly and never +ingest, so those tables appear here by name. When jdluc starts reading a new dataset, add it +here. +""" + +import collections.abc +import json +import logging +import os + +from jdluc import ( + attribute, + config, + harmonize, + jurisdictional_direct, + statistical, + tiling, +) +from jdluc.datasets import NAME_TO_CLS, DatasetName, worldbank_jurisdictions + +logger = logging.getLogger(__name__) + +# Read by every tile lookup and every per-jurisdiction rollup, but named by no workflow +BOUNDARY_DATASET_NAMES = ( + DatasetName.WORLD_BANK_ADMIN_0, + DatasetName.WORLD_BANK_ADMIN_1, +) + +# Read with `load()` straight from the ingest root, never ingested on demand +METHODOLOGY_TO_LOADED_DATASET_NAMES = { + attribute.Methodology.JURISDICTIONAL_DIRECT: (DatasetName.USDA_NASS_QUICKSTATS,), + attribute.Methodology.STATISTICAL: ( + DatasetName.FAOSTAT_LIVESTOCK_PATTERNS, + DatasetName.FAOSTAT_PRODUCTION_LIVESTOCK, + ), +} + +METHODOLOGY_TO_HARMONIZED_DATASET_NAMES = { + attribute.Methodology.JURISDICTIONAL_DIRECT: jurisdictional_direct.DATASET_NAMES, + attribute.Methodology.STATISTICAL: statistical.DATASET_NAMES, +} + + +# The scope of a run over every country and GNW tile, as jdluc's `--backfill`; also what its +# run-wide outputs are named, in place of every country's code joined +WORLD = "world" + + +def get_dataset_names( + methodology: attribute.Methodology, partitioning: tiling.Partitioning +) -> list[DatasetName]: + """Every dataset of one partitioning that a run of ``methodology`` reads.""" + dataset_names = { + *BOUNDARY_DATASET_NAMES, + *(name for stack in harmonize.Stack for name in stack.value), + *METHODOLOGY_TO_HARMONIZED_DATASET_NAMES[methodology], + *METHODOLOGY_TO_LOADED_DATASET_NAMES[methodology], + } + return sorted( + ( + name + for name in dataset_names + if NAME_TO_CLS[name].partitioning == partitioning + ), + key=lambda name: name.name, + ) + + +def scope_str(s: str) -> str: + """One element of a run's scope: an ISO 3166 alpha-3 code, or ``world``.""" + return s if s == WORLD else worldbank_jurisdictions.iso_3166_str(s=s) + + +def is_world(scope: tuple[str, ...]) -> bool: + assert scope, "the scope names no country" + assert scope == (WORLD,) or WORLD not in scope, ( + f"{scope=} mixes {WORLD} with countries" + ) + return scope == (WORLD,) + + +def get_iso_3166s(scope: tuple[str, ...]) -> tuple[str, ...]: + """The scope's countries: those named, or for ``world`` every country, as jdluc's + ``--backfill`` reads them.""" + return tuple( + sorted( + worldbank_jurisdictions.get_all_iso_3166s() if is_world(scope) else scope + ) + ) + + +def get_name(scope: tuple[str, ...]) -> str: + """What the run's run-wide outputs (its mosaics and emissions-factor table) are named.""" + return WORLD if is_world(scope) else "-".join(sorted(scope)) + + +def get_tile_ids(scope: tuple[str, ...]) -> list[str]: + """The tiles the scope's countries touch, or for ``world`` every GNW tile, as jdluc's + ``--backfill`` reads them: tiles no country touches included.""" + return sorted( + tiling.GLOBAL_NATURE_WATCH_TILE_IDS + if is_world(scope) + else worldbank_jurisdictions.get_ten_degree_tile_ids_for_iso_3166s( + iso_3166s=get_iso_3166s(scope=scope) + ) + ) + + +def ingest_boundaries() -> None: + """The tile list is read off the boundaries, so the plan ingests them first. + + A no-op when they are already in the ingest root, as after any earlier run. + """ + root = config.Config.from_dot_env().ingest_root + for name in BOUNDARY_DATASET_NAMES: + logger.info(f"Ingesting {name.name:s} if absent") + NAME_TO_CLS[name].ingest_a_tile( + overwrite=False, root=root, tile_id=tiling.WHOLE_WORLD_TILE_ID + ) + + +def check_allow_failures( + allow_failures: tuple[str, ...], methodology: attribute.Methodology +) -> None: + tiled = get_dataset_names( + methodology=methodology, partitioning=tiling.Partitioning.TEN_DEGREE_TILE + ) + unknown = sorted(set(allow_failures) - {name.name for name in tiled}) + assert not unknown, f"{unknown} are not ten-degree datasets this run ingests" + + +def get_task_to_items( + allow_failures: tuple[str, ...], + methodology: attribute.Methodology, + scope: tuple[str, ...], +) -> dict[str, list]: + """Each list the workflow loops over, by task name. Writes nothing, so that a wrapper can + add lists of its own, or keys to each tile's unit, before `write_task_to_items`. + + Reads the boundaries for the tile list: `ingest_boundaries` first, on a fresh ingest root. + + The flat workflow loops over each stage's list across every tile. ``tile-units`` holds + the same units grouped by tile, for the per-tile workflow, which runs each tile's ingest + and compute as a DAG of its own. + """ + check_allow_failures(allow_failures=allow_failures, methodology=methodology) + tiled = get_dataset_names( + methodology=methodology, partitioning=tiling.Partitioning.TEN_DEGREE_TILE + ) + tile_ids = get_tile_ids(scope=scope) + # One flat list of (dataset, tile), so that the workflow needs no nested loop + ingest_units = [ + {"dataset": name.name, "tile_id": tile_id} + for name in tiled + for tile_id in tile_ids + ] + return { + "ingest-world": [ + name.name + for name in get_dataset_names( + methodology=methodology, partitioning=tiling.Partitioning.WHOLE_WORLD + ) + if name not in BOUNDARY_DATASET_NAMES + ], + "ingest-tiles": [u for u in ingest_units if u["dataset"] not in allow_failures], + "ingest-tiles-allowed-failures": [ + u for u in ingest_units if u["dataset"] in allow_failures + ], + "tiles": [{"tile_id": tile_id} for tile_id in tile_ids], + "tile-units": [ + { + "tile_id": tile_id, + "ingests": [ + u["dataset"] + for u in ingest_units + if u["tile_id"] == tile_id and u["dataset"] not in allow_failures + ], + "ingests_allowed_failures": [ + u["dataset"] + for u in ingest_units + if u["tile_id"] == tile_id and u["dataset"] in allow_failures + ], + } + for tile_id in tile_ids + ], + } + + +def write_task_to_items( + output_dir: str, task_to_items: collections.abc.Mapping[str, list] +) -> None: + """Write each list to ``output_dir``, as ``.json``.""" + os.makedirs(output_dir, exist_ok=True) + for task, items in task_to_items.items(): + path = os.path.join(output_dir, f"{task:s}.json") + logger.info(f"Writing {len(items):d} items to {path=:s}") + with open(path, "w") as fp: + json.dump(items, fp) + + +def write_plan( + allow_failures: tuple[str, ...], + methodology: attribute.Methodology, + output_dir: str, + scope: tuple[str, ...], +) -> None: + """Write each list the workflow loops over to ``output_dir``, as ``.json``. + + The run covers ``scope``: one or more countries, or ``world`` for every country and GNW + tile. + + ``allow_failures`` names ten-degree datasets whose failed tiles must not stop the run. + Their units go to their own list, so that the workflow can run them as a task that + compute does not wait to succeed. + """ + # Before the boundaries' ingest, so that a typo fails fast + check_allow_failures(allow_failures=allow_failures, methodology=methodology) + ingest_boundaries() + write_task_to_items( + output_dir=output_dir, + task_to_items=get_task_to_items( + allow_failures=allow_failures, methodology=methodology, scope=scope + ), + ) diff --git a/orchestration/workflows/cornerstone.example.yaml b/orchestration/workflows/cornerstone.example.yaml new file mode 100644 index 0000000..015143a --- /dev/null +++ b/orchestration/workflows/cornerstone.example.yaml @@ -0,0 +1,238 @@ +# The cornerstone pipeline as an Argo WorkflowTemplate, one DAG per tile: an example with +# placeholder names. Replace every UPPER-CASE placeholder. docs/orchestration.md describes the +# setup: the image, the two Secrets and the service account. +# +# argo template create cornerstone.example.yaml +# argo submit --from workflowtemplate/cornerstone -p scope=HND --generate-name cornerstone-hnd- +# argo submit --from workflowtemplate/cornerstone -p scope=world --generate-name cornerstone-world- +# +# Each tile runs its own ingest -> compute, so a failed (dataset, tile) ingest stops only its +# own tile, and every other tile goes on to compute. trace needs every tile, so with any tile +# failed it doesn't run, and the run ends Failed. `argo retry` resumes from the cache. +# +# `plan` writes, as `tile-units`, one item per tile listing its ingests, which the tile's own +# DAG loops over. A tile's tasks start once the whole-world datasets are in. +# +# plan ─┬─ ingest-world ─┐ +# └────────────────┴─ tile[tile] ── trace +# +# tile: ingest ──────────────────┐ +# ingest-allowed-failures ─┴─ compute +# +# Each task template is reusable on its own. A workflow that adds tasks of its own (say, a +# per-tile export after compute) calls them by `templateRef: {name: cornerstone, template: ...}` +# from DAGs of its own, rather than copying this file. +# +# Sizing: compute requests 56 GiB, with peaks near 40 GiB observed; the other tasks fit in +# 24 GiB. Every pod writes its temp files to a volume of its own at /localtmp: on a cloud +# node's boot disk, GDAL's temp files can throttle the whole node. + +apiVersion: argoproj.io/v1alpha1 +kind: WorkflowTemplate +metadata: + name: cornerstone +spec: + entrypoint: cornerstone + # Needs read/write access to INGEST_ROOT, SCRATCH_ROOT and EXPORT_ROOT, for instance through + # workload identity rather than keys + serviceAccountName: SERVICE-ACCOUNT + # Backstops against a runaway fan-out: at most 25 pods at once, and 6 hours for the whole run. + # Sized for one country: raise both for a scope=world run + parallelism: 25 + activeDeadlineSeconds: 21600 + arguments: + parameters: + # The run's country, as one ISO 3166 alpha-3 code, or world for every country and GNW tile + - name: scope + # STATISTICAL or JURISDICTIONAL_DIRECT (the latter is US-only) + - name: methodology + value: STATISTICAL + # Ten-degree datasets whose failed tiles must not stop the run, comma-separated. Descals + # is published only where there is oil palm, and its ingest fails on tiles without any. + - name: allow-failures + value: DESCALS_OIL_PALM + + # To run on preemptible (spot) nodes, add a nodeSelector and tolerations here. A reclaimed + # node fails its pod; `argo retry` resumes from the cache, as a zarr counts only once written. + # templateDefaults: + # nodeSelector: { NODE-POOL-LABEL: NODE-POOL } + # tolerations: [{ key: NODE-POOL-TAINT, operator: Equal, value: "true", effect: NoSchedule }] + + templates: + - name: cornerstone + dag: + # A failed task must not stop the scheduling of tasks that wait on its failure, nor a + # failed tile the other tiles + failFast: false + tasks: + - name: plan + template: plan + + - name: ingest-world + template: ingest-world + depends: plan + withParam: "{{tasks.plan.outputs.parameters.ingest-world}}" + arguments: { parameters: [{ name: dataset, value: "{{item}}" }] } + + # Each list in an item reaches the tile's DAG as a JSON string, for it to loop over + - name: tile + template: tile + depends: plan && ingest-world + withParam: "{{tasks.plan.outputs.parameters.tile-units}}" + arguments: + parameters: + - { name: tile-id, value: "{{item.tile_id}}" } + - { name: ingests, value: "{{item.ingests}}" } + - { name: ingests-allowed-failures, value: "{{item.ingests_allowed_failures}}" } + + # Runs only if every tile succeeded: a table missing tiles would pass for a whole one + - name: trace + template: trace + depends: tile + + # One tile's ingest -> compute. A failed required ingest leaves compute unscheduled, which + # fails this tile alone + - name: tile + inputs: + parameters: + - name: tile-id + - name: ingests + - name: ingests-allowed-failures + dag: + failFast: false + tasks: + # Ingest exits non-zero if its tile fails, which fails the task, and so the tile + - name: ingest + template: ingest-tile + withParam: "{{inputs.parameters.ingests}}" + arguments: &ingest-tile-arguments + parameters: + - { name: dataset, value: "{{item}}" } + - { name: tile-id, value: "{{inputs.parameters.tile-id}}" } + + # Failed tiles here stay visible as failed pods, but compute runs regardless + - name: ingest-allowed-failures + template: ingest-tile + withParam: "{{inputs.parameters.ingests-allowed-failures}}" + arguments: *ingest-tile-arguments + + - name: compute + template: compute + depends: ingest && (ingest-allowed-failures || ingest-allowed-failures.Failed) + arguments: { parameters: [{ name: tile-id, value: "{{inputs.parameters.tile-id}}" }] } + + # Ingests the boundaries if absent, then writes one JSON list per fan-out, and the per-tile list + - name: plan + activeDeadlineSeconds: 1800 + container: + image: &image IMAGE-REGISTRY/cornerstone:COMMIT-SHA + command: [python, -m, orchestration, plan] + args: + - "{{workflow.parameters.scope}}" + - "--methodology={{workflow.parameters.methodology}}" + - --output-dir=/plan + - "--allow-failures={{workflow.parameters.allow-failures}}" + env: &tmpdir-env + - { name: TMPDIR, value: /localtmp } + - { name: CPL_TMPDIR, value: /localtmp } + resources: &light-resources + requests: { cpu: "2", memory: 24Gi, ephemeral-storage: 2Gi } + limits: { cpu: "4", memory: 24Gi, ephemeral-storage: 8Gi } + volumeMounts: &mounts + - { name: cloud-env, mountPath: /app/.env, subPath: .env } + - { name: localtmp, mountPath: /localtmp } + volumes: + - &nokeys { name: cloud-env, secret: { secretName: cornerstone-env-nokeys } } + - &localtmp-20 + name: localtmp + ephemeral: + volumeClaimTemplate: + spec: + accessModes: [ReadWriteOnce] + # An SSD-backed class: throughput scales with size, so don't shrink these + storageClassName: SSD-STORAGE-CLASS + resources: { requests: { storage: 20Gi } } + outputs: + parameters: + - { name: ingest-world, valueFrom: { path: /plan/ingest-world.json } } + - { name: tile-units, valueFrom: { path: /plan/tile-units.json } } + + # One whole-world dataset, with jdluc's own CLI + - name: ingest-world + inputs: { parameters: [{ name: dataset }] } + activeDeadlineSeconds: 5400 + container: + image: *image + command: [python, jdluc/ingest.py] + args: ["{{inputs.parameters.dataset}}", --backfill] + env: *tmpdir-env + resources: *light-resources + volumeMounts: *mounts + volumes: + - &keys { name: cloud-env, secret: { secretName: cornerstone-env } } + - name: localtmp + ephemeral: + volumeClaimTemplate: + spec: + accessModes: [ReadWriteOnce] + storageClassName: SSD-STORAGE-CLASS + resources: { requests: { storage: 100Gi } } + + # One dataset on one tile. Room for a global source download (GPW livestock, ~5 GB) plus + # one tile staged uncompressed for its COG (Liao, ~55 GB) + - name: ingest-tile + inputs: { parameters: [{ name: dataset }, { name: tile-id }] } + activeDeadlineSeconds: 10800 + container: + image: *image + command: [python, -m, orchestration, ingest-tile] + args: + - "--dataset={{inputs.parameters.dataset}}" + - "--tile-id={{inputs.parameters.tile-id}}" + env: *tmpdir-env + resources: *light-resources + volumeMounts: *mounts + volumes: + - *keys + - name: localtmp + ephemeral: + volumeClaimTemplate: + spec: + accessModes: [ReadWriteOnce] + storageClassName: SSD-STORAGE-CLASS + resources: { requests: { storage: 150Gi } } + + # harmonize -> emit -> attribute for one tile, from the warm ingest cache, without keys + - name: compute + inputs: { parameters: [{ name: tile-id }] } + activeDeadlineSeconds: 14400 + container: + image: *image + command: [python, -m, orchestration, compute] + args: + - "{{workflow.parameters.scope}}" + - "--methodology={{workflow.parameters.methodology}}" + - "--tile-id={{inputs.parameters.tile-id}}" + # A second line of defence: harmonize would read a missing ingest as nodata + - --check-ingests + - "--allow-failures={{workflow.parameters.allow-failures}}" + env: *tmpdir-env + resources: + requests: { cpu: "6", memory: 56Gi, ephemeral-storage: 2Gi } + limits: { cpu: "8", memory: 56Gi, ephemeral-storage: 8Gi } + volumeMounts: *mounts + volumes: [*nokeys, *localtmp-20] + + # The cross-tile merge over the scope; writes emissions-factors/{scope}-{methodology}.parquet + - name: trace + activeDeadlineSeconds: 3600 + container: + image: *image + command: [python, -m, orchestration, trace] + args: + - "{{workflow.parameters.scope}}" + - "--methodology={{workflow.parameters.methodology}}" + env: *tmpdir-env + resources: *light-resources + volumeMounts: *mounts + volumes: [*nokeys, *localtmp-20] diff --git a/pyproject.toml b/pyproject.toml index 6e37fee..a2ee865 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,7 +59,7 @@ follow_imports = "skip" implicit_reexport = true [tool.importlinter] -root_packages = ["jdluc", "validation"] +root_packages = ["jdluc", "orchestration", "validation"] [[tool.importlinter.contracts]] name = "The pipeline never imports the validator" @@ -71,6 +71,16 @@ forbidden_modules = [ "validation", ] +[[tool.importlinter.contracts]] +name = "The pipeline never imports its orchestration" +type = "forbidden" +source_modules = [ + "jdluc", +] +forbidden_modules = [ + "orchestration", +] + [[tool.importlinter.contracts]] name = "ETL" type = "layers" @@ -88,6 +98,7 @@ name = "Utilities are foundational" type = "forbidden" source_modules = [ "jdluc.config", + "jdluc.export", "jdluc.geo", "jdluc.storage", "jdluc.tiling", From c6e3b0a3c5e0e6e23026d18908c90c6a4492edfb Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Mon, 19 Oct 2026 12:16:11 +0200 Subject: [PATCH 2/3] orchestration: docstrings name nothing that lives only in the fork get_name no longer mentions mosaics, get_task_to_items describes the flat lists without a flat workflow to point at, and trace's --help is a whole sentence. --- orchestration/__main__.py | 5 +++-- orchestration/plan.py | 8 ++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/orchestration/__main__.py b/orchestration/__main__.py index 6daf089..e4eeed4 100644 --- a/orchestration/__main__.py +++ b/orchestration/__main__.py @@ -131,8 +131,9 @@ def get_iso_3166s_touching_tile( def trace_countries(methodology: attribute.Methodology, scope: tuple[str, ...]) -> None: - """The cross-tile merge, written as the emissions-factor table. Cheap only if compute ran - with the same methodology.""" + """The cross-tile merge, written as the emissions-factor table. + + Cheap only if compute ran with the same methodology.""" uri = trace.export_workflow( commodity_names=attribute.get_commodity_names(methodology=methodology), iso_3166s=plan.get_iso_3166s(scope=scope), diff --git a/orchestration/plan.py b/orchestration/plan.py index 964a17e..70141a8 100644 --- a/orchestration/plan.py +++ b/orchestration/plan.py @@ -96,7 +96,7 @@ def get_iso_3166s(scope: tuple[str, ...]) -> tuple[str, ...]: def get_name(scope: tuple[str, ...]) -> str: - """What the run's run-wide outputs (its mosaics and emissions-factor table) are named.""" + """What the run's run-wide outputs, such as its emissions-factor table, are named.""" return WORLD if is_world(scope) else "-".join(sorted(scope)) @@ -145,9 +145,9 @@ def get_task_to_items( Reads the boundaries for the tile list: `ingest_boundaries` first, on a fresh ingest root. - The flat workflow loops over each stage's list across every tile. ``tile-units`` holds - the same units grouped by tile, for the per-tile workflow, which runs each tile's ingest - and compute as a DAG of its own. + ``tile-units`` groups the units by tile, for the example workflow, which runs each tile's + ingest and compute as a DAG of its own. The other lists hold the same units flat, one per + stage across every tile, for a workflow that fans out stage by stage instead. """ check_allow_failures(allow_failures=allow_failures, methodology=methodology) tiled = get_dataset_names( From e5611e79b45b5e7d41e27e4e7712ad40acb6a05a Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Mon, 19 Oct 2026 23:05:12 +0200 Subject: [PATCH 3/3] export: write_table names its product, as get_metadata stamps it deliverable_name becomes product_name in write_table and get_table_output_uri, matching the cornerstone-product-name tag that get_metadata already writes from it. --- jdluc/export.py | 22 +++++++++++----------- jdluc/trace.py | 4 ++-- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/jdluc/export.py b/jdluc/export.py index 5860d5e..ef1d6fc 100644 --- a/jdluc/export.py +++ b/jdluc/export.py @@ -1,8 +1,8 @@ -"""Write a pipeline step's output as a deliverable under EXPORT_ROOT. +"""Write a pipeline step's output as a product under EXPORT_ROOT. -A utility, like `storage`: it knows nothing about the steps. A tabular deliverable (see +A utility, like `storage`: it knows nothing about the steps. A tabular product (see `trace.export_workflow`) is one parquet per set of countries, written by `write_table` to -`{EXPORT_ROOT}/{deliverable}/{name}.parquet`, with provenance in its `attrs`. +`{EXPORT_ROOT}/{product}/{name}.parquet`, with provenance in its `attrs`. Unlike SCRATCH_ROOT, EXPORT_ROOT is not a cache that the pipeline reads back: what lands there is a terminal output. @@ -31,32 +31,32 @@ def get_metadata(product_name: str, source_name: str) -> dict[str, str]: def get_table_output_uri( - deliverable_name: str, name: str, output_root: str | None = None + product_name: str, name: str, output_root: str | None = None ) -> str: - """`{output_root}/{deliverable}/{name}.parquet`; `output_root` defaults to EXPORT_ROOT.""" + """`{output_root}/{product}/{name}.parquet`; `output_root` defaults to EXPORT_ROOT.""" if output_root is None: output_root = config.Config.from_dot_env().export_root return storage.join_uri( - root=output_root, prefix=f"{deliverable_name:s}/{name:s}.parquet" + root=output_root, prefix=f"{product_name:s}/{name:s}.parquet" ) def write_table( df: pandas.DataFrame, - deliverable_name: str, + product_name: str, name: str, source_name: str, output_root: str | None = None, ) -> str: - """Write a tabular deliverable as one parquet, provenance in its `attrs`; return the URI.""" + """Write a tabular product as one parquet, provenance in its `attrs`; return the URI.""" uri = get_table_output_uri( - deliverable_name=deliverable_name, name=name, output_root=output_root + name=name, output_root=output_root, product_name=product_name ) df = df.copy() - df.attrs |= get_metadata(product_name=deliverable_name, source_name=source_name) + df.attrs |= get_metadata(product_name=product_name, source_name=source_name) with tempfile.TemporaryDirectory() as tmpdir: local_path = os.path.join(tmpdir, os.path.basename(uri)) df.to_parquet(path=local_path) storage.put_file(local_path=local_path, uri=uri) - logger.info(f"Wrote {deliverable_name=:s} ({len(df):d} rows) to {uri=:s}") + logger.info(f"Wrote {product_name=:s} ({len(df):d} rows) to {uri=:s}") return uri diff --git a/jdluc/trace.py b/jdluc/trace.py index 289a70b..916844c 100644 --- a/jdluc/trace.py +++ b/jdluc/trace.py @@ -9,7 +9,7 @@ It then derives, identically for both, each commodity's yield (kg per hectare) and emissions factor (kgCO2e per kg), and rolls the provincials up to national totals. Returns a cached pandas.DataFrame indexed by (admin_level, admin_id, commodity_name, methodology). -`export_workflow` writes that table as the `emissions-factors` deliverable under EXPORT_ROOT. +`export_workflow` writes that table as the `emissions-factors` product under EXPORT_ROOT. The countries come from the positional ISO 3166 alpha-3 codes, or -- with `--backfill` -- from every country in the World Bank admin-0 layer. @@ -197,9 +197,9 @@ def export_workflow( ) return export.write_table( df=df, - deliverable_name=EMISSIONS_FACTORS, name=f"{name or '-'.join(iso_3166s):s}-{methodology.name:s}", output_root=output_root, + product_name=EMISSIONS_FACTORS, source_name="jdluc-trace", )