diff --git a/deployment/aws/lambda_handler.py b/deployment/aws/lambda_handler.py index 8751d0b5..eebf69d6 100644 --- a/deployment/aws/lambda_handler.py +++ b/deployment/aws/lambda_handler.py @@ -6,9 +6,9 @@ Event payload (default / process mode): { "chunk_idx": int, - "parent_morton": int, - "parent_order": int, - "child_order": int, + "shard_key": int, # grid-agnostic shard identifier + "parent_order": int, # HEALPix only (omit for other grids) + "child_order": int, # HEALPix only (omit for other grids) "granule_urls": [str, ...], "store_path": str, # e.g. "s3://bucket/prefix.zarr" "s3_credentials": { # creds for reading NSIDC source data @@ -182,7 +182,7 @@ def _handle_process(event: Dict[str, Any], context: Any) -> Dict[str, Any]: json.dumps( { "event_type": "lambda_invocation", - "parent_morton": event.get("parent_morton"), + "shard_key": event.get("shard_key"), "granule_count": len(event.get("granule_urls", [])), "child_order": event.get("child_order"), "request_id": context.aws_request_id, @@ -192,11 +192,13 @@ def _handle_process(event: Dict[str, Any], context: Any) -> Dict[str, Any]: ) try: - # Validate required parameters + # Validate required parameters. ``child_order`` is HEALPix-specific and + # only required once the grid is known to be HEALPix (checked below); + # ``parent_order`` is forwarded by the orchestrator for every grid (None + # for non-HEALPix), so its key is always present. required_params = [ - "parent_morton", + "shard_key", "parent_order", - "child_order", "granule_urls", "store_path", "s3_credentials", @@ -228,13 +230,22 @@ def _handle_process(event: Dict[str, Any], context: Any) -> Dict[str, Any]: if config is None: from zagg.config import default_config config = default_config("atl06") - grid = from_config(config, parent_order=event["parent_order"]) + + # child_order is required for HEALPix runs (drives the leaf order); it is + # absent/unused for non-HEALPix grids. + grid_type = config.output.get("grid", {}).get("type", "healpix") + if grid_type == "healpix" and "child_order" not in event: + error_msg = "Missing required parameters: child_order" + logger.error(error_msg) + return {"statusCode": 400, "body": json.dumps({"error": error_msg})} + + grid = from_config(config, parent_order=event.get("parent_order")) # Process the shard using cloud-agnostic function from zagg.processing import process_shard df_out, metadata = process_shard( grid, - event["parent_morton"], + event["shard_key"], event["granule_urls"], s3_credentials=s3_creds, config=config, @@ -274,7 +285,7 @@ def _handle_process(event: Dict[str, Any], context: Any) -> Dict[str, Any]: json.dumps( { "event_type": "processing_complete", - "parent_morton": metadata["parent_morton"], + "shard_key": metadata["shard_key"], "cells_with_data": metadata["cells_with_data"], "total_obs": metadata["total_obs"], "duration_s": metadata["duration_s"], @@ -302,7 +313,7 @@ def _handle_process(event: Dict[str, Any], context: Any) -> Dict[str, Any]: "body": json.dumps( { "error": f"Unhandled exception: {str(e)}", - "parent_morton": event.get("parent_morton"), + "shard_key": event.get("shard_key"), "request_id": context.aws_request_id, } ), diff --git a/docs/deployment/lambda.md b/docs/deployment/lambda.md index fbb833de..bac4bef6 100644 --- a/docs/deployment/lambda.md +++ b/docs/deployment/lambda.md @@ -47,7 +47,7 @@ The Lambda function processes a single morton cell (order 6) by: ```json { - "parent_morton": 123456, + "shard_key": 123456, "parent_order": 6, "child_order": 12, "granule_urls": [ @@ -74,24 +74,24 @@ The Lambda function processes a single morton cell (order 6) by: | Parameter | Type | Required | Description | |-----------|------|----------|-------------| -| `parent_morton` | int | Yes | Morton index of parent cell (order 6) | -| `parent_order` | int | Yes | Order of parent cell (typically 6) | -| `child_order` | int | Yes | Order of child cells for statistics (typically 12) | +| `shard_key` | int | Yes | Grid-agnostic shard identifier (HEALPix: the parent-cell morton index) | +| `parent_order` | int | Yes | Order of parent cell (typically 6); HEALPix-only (`null` for other grids) | +| `child_order` | int | HEALPix only | Order of child cells for statistics (typically 12); omitted for non-HEALPix grids | | `granule_urls` | list | Yes | Pre-computed list of S3 URLs from catalog | | `store_path` | str | Yes | Output Zarr store path (e.g. `s3://bucket/prefix.zarr`) | | `s3_credentials` | dict | Yes | NSIDC S3 credentials for reading source data | | `output_credentials` | dict | No | Explicit credentials for *writing* the output store. Omit to use the execution role (in-account writes). Supply to write an external / S3-compatible target. Keys: `accessKeyId`, `secretAccessKey`, optional `sessionToken`/`endpointUrl`/`region`. | -!!! note "Shard vs. `parent_order` naming" - The unit of work is a **shard** — one parent (order-6) HEALPix cell. The +!!! note "Grid-neutral event fields" + The unit of work is a **shard** — for HEALPix, one parent (order-6) cell. The orchestrator and the catalog use that vocabulary (`python -m zagg.catalog` emits a shard map with `shard_keys` + a `grid_signature`). The Lambda - **event** schema, however, still carries the HEALPix-shaped field names - `parent_morton` / `parent_order` / `child_order` shown above (the handler - requires them — see `deployment/aws/lambda_handler.py`). A grid-neutral - rename of those event fields is tracked in - [#24](https://github.com/englacial/zagg/issues/24), not done here, so the - field names above are the current, accurate ones. + **event** schema uses the grid-neutral field name `shard_key` (the shard + identifier for any grid; for HEALPix it is the parent-cell morton index). + `parent_order`/`child_order` are HEALPix-specific: `parent_order` is + forwarded for every grid (`null` for non-HEALPix), while `child_order` is + only required/sent for HEALPix runs. See `deployment/aws/lambda_handler.py`. + This rename landed via [#24](https://github.com/englacial/zagg/issues/24). ### S3 Credentials @@ -105,7 +105,7 @@ s3_creds = get_nsidc_s3_credentials() # Pass to each Lambda invocation event = { - "parent_morton": -6134114, + "shard_key": -6134114, "parent_order": 6, "child_order": 12, "granule_urls": [...], diff --git a/src/zagg/auth.py b/src/zagg/auth.py index e5336a54..56f13fda 100644 --- a/src/zagg/auth.py +++ b/src/zagg/auth.py @@ -56,7 +56,7 @@ def get_nsidc_s3_credentials() -> dict: # Pass to Lambda invocation event = { - "parent_morton": -6134114, + "shard_key": -6134114, "s3_credentials": creds, # ... other params } diff --git a/src/zagg/processing.py b/src/zagg/processing.py index 7b30f610..49aec1d9 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -487,7 +487,7 @@ def _kernel_aggregate( def _read_group( - h5obj, group: str, data_source: dict, parent_morton: int, grid, arrow: bool = False + h5obj, group: str, data_source: dict, shard_key: int, grid, arrow: bool = False ): """Read and spatially filter one HDF5 group. @@ -513,7 +513,7 @@ def _read_group( # Assign points to leaf cells, then filter to the current shard. leaf_ids = grid.assign(lats, lons) - mask_spatial = grid.shards_of(leaf_ids) == parent_morton + mask_spatial = grid.shards_of(leaf_ids) == shard_key if np.sum(mask_spatial) == 0: return None @@ -629,8 +629,8 @@ def process_shard( raise ValueError(f"handoff must be 'pandas', 'arrow', or 'arrow-kernel', got {handoff!r}") data_source = config.data_source - parent_morton = int(shard_key) - logger.info(f"Processing shard: {parent_morton}") + shard_key = int(shard_key) + logger.info(f"Processing shard: {shard_key}") start_time = datetime.now() # Resolve driver @@ -648,7 +648,7 @@ def process_shard( # Prepare metadata metadata: ProcessingMetadata = { - "parent_morton": parent_morton, + "shard_key": shard_key, "cells_with_data": 0, "total_obs": 0, "granule_count": len(granule_urls), @@ -659,7 +659,7 @@ def process_shard( # Check for granules if not granule_urls: - logger.info(f" No granules provided for morton {parent_morton} - skipping") + logger.info(f" No granules provided for shard {shard_key} - skipping") metadata["error"] = "No granules found" metadata["duration_s"] = (datetime.now() - start_time).total_seconds() return pd.DataFrame(), metadata @@ -701,7 +701,7 @@ def process_shard( for g in data_source["groups"]: try: - chunk = _read_group(h5obj, g, data_source, parent_morton, grid, arrow=use_arrow) + chunk = _read_group(h5obj, g, data_source, shard_key, grid, arrow=use_arrow) if chunk is not None: all_reads.append(chunk) except Exception as e: @@ -718,12 +718,12 @@ def process_shard( metadata["files_processed"] = files_processed if not all_reads: - logger.info(f" No data after filtering for morton {parent_morton} - skipping") + logger.info(f" No data after filtering for shard {shard_key} - skipping") metadata["error"] = "No data after filtering" metadata["duration_s"] = (datetime.now() - start_time).total_seconds() return pd.DataFrame(), metadata - children = grid.children(parent_morton) + children = grid.children(shard_key) data_vars = get_data_vars(config) if handoff == "arrow-kernel": @@ -790,7 +790,7 @@ def process_shard( df_out[col_name] = vals duration = (datetime.now() - start_time).total_seconds() - logger.info(f"Completed shard {parent_morton} in {duration:.1f}s") + logger.info(f"Completed shard {shard_key} in {duration:.1f}s") metadata["cells_with_data"] = cells_with_data metadata["total_obs"] = int(stats_arrays["count"].sum()) diff --git a/src/zagg/runner.py b/src/zagg/runner.py index 95fe86dc..35e5b60a 100644 --- a/src/zagg/runner.py +++ b/src/zagg/runner.py @@ -127,7 +127,9 @@ def agg( if not store_path: raise ValueError("No store path specified (pass store= or set output.store: in config)") - child_order = get_child_order(config) + # child_order is HEALPix-specific (leaf order); other grids don't define it. + grid_type = config.output.get("grid", {}).get("type", "healpix") + child_order = get_child_order(config) if grid_type == "healpix" else None _maybe_warn_dense(get_layout(config)) # Resolve driver: kwarg > config > default @@ -506,7 +508,7 @@ def _run_lambda(config, catalog_data, store_path, child_order, *, cells_with_data += 1 elif error not in ("No granules found", "No data after filtering"): cells_error += 1 - logger.warning(f" [{i}/{len(cells)}] morton {result.get('morton')}: {error}") + logger.warning(f" [{i}/{len(cells)}] shard {result.get('shard_key')}: {error}") if i % 50 == 0: elapsed = time.time() - start_time @@ -631,7 +633,7 @@ def _invoke_lambda_finalize(lambda_client, function_name, store_path, def _invoke_lambda_cell( - lambda_client, chunk_idx, parent_morton, parent_order, child_order, + lambda_client, chunk_idx, shard_key, parent_order, child_order, granule_urls, store_path, s3_credentials, *, function_name, config_dict, output_creds_event=None, max_retries=3, max_workers=None, @@ -645,9 +647,8 @@ def _invoke_lambda_cell( event = { "chunk_idx": chunk_idx, - "parent_morton": parent_morton, + "shard_key": shard_key, "parent_order": parent_order, - "child_order": child_order, "granule_urls": granule_urls, "store_path": store_path, "s3_credentials": { @@ -656,6 +657,10 @@ def _invoke_lambda_cell( "sessionToken": s3_credentials["sessionToken"], }, } + # child_order is HEALPix-specific; only forward it when set (non-HEALPix + # grids leave it None and the handler doesn't require it). + if child_order is not None: + event["child_order"] = child_order if config_dict is not None: event["config"] = config_dict if output_creds_event is not None: @@ -692,7 +697,7 @@ def _invoke_lambda_cell( body = {} return { - "morton": parent_morton, + "shard_key": shard_key, "status_code": result.get("statusCode"), "body": body, "wall_time": time.time() - wall_start, @@ -717,7 +722,7 @@ def _invoke_lambda_cell( break return { - "morton": parent_morton, + "shard_key": shard_key, "status_code": None, "body": {}, "wall_time": time.time() - wall_start, diff --git a/src/zagg/schema.py b/src/zagg/schema.py index 66565e97..545d7bd0 100644 --- a/src/zagg/schema.py +++ b/src/zagg/schema.py @@ -17,7 +17,7 @@ class ProcessingMetadata(TypedDict): - parent_morton: int + shard_key: int cells_with_data: int total_obs: int granule_count: int diff --git a/tests/test_lambda_handler.py b/tests/test_lambda_handler.py new file mode 100644 index 00000000..7a9df133 --- /dev/null +++ b/tests/test_lambda_handler.py @@ -0,0 +1,130 @@ +"""Tests for the AWS Lambda handler's process-mode event contract (#24). + +The handler lives under ``deployment/aws/`` (not an importable package module), +so it is loaded by path. These tests exercise the grid-neutral event schema: +the shard identifier is ``shard_key`` (not ``parent_morton``), and the +HEALPix-specific ``child_order`` is required for HEALPix runs but optional for +other grids. +""" + +import importlib.util +import json +from dataclasses import asdict +from pathlib import Path +from unittest.mock import MagicMock + +import pandas as pd +import pytest + +from zagg.config import default_config + +REPO_ROOT = Path(__file__).parent.parent +HANDLER_PATH = REPO_ROOT / "deployment" / "aws" / "lambda_handler.py" + + +@pytest.fixture(scope="module") +def handler_mod(): + spec = importlib.util.spec_from_file_location("zagg_lambda_handler", HANDLER_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def _context(): + ctx = MagicMock() + ctx.aws_request_id = "req-1" + ctx.function_name = "process-shard" + ctx.memory_limit_in_mb = 2048 + ctx.get_remaining_time_in_millis.return_value = 900_000 + return ctx + + +def _healpix_config_dict(): + return asdict(default_config("atl06")) + + +def _rectilinear_config_dict(): + return asdict(default_config("atl06_polar")) + + +_CREDS = {"accessKeyId": "a", "secretAccessKey": "s", "sessionToken": "t"} + + +def _base_event(config_dict): + return { + "shard_key": 12345, + "parent_order": 6, + "granule_urls": ["s3://b/g.h5"], + "store_path": "s3://out/x.zarr", + "s3_credentials": _CREDS, + "config": config_dict, + } + + +class TestProcessEventGate: + def test_missing_shard_key_rejected(self, handler_mod): + event = _base_event(_healpix_config_dict()) + del event["shard_key"] + resp = handler_mod._handle_process(event, _context()) + assert resp["statusCode"] == 400 + assert "shard_key" in json.loads(resp["body"])["error"] + + def test_legacy_parent_morton_not_accepted(self, handler_mod): + # Hard rename: the old field name is no longer a valid shard identifier. + event = _base_event(_healpix_config_dict()) + del event["shard_key"] + event["parent_morton"] = 12345 + resp = handler_mod._handle_process(event, _context()) + assert resp["statusCode"] == 400 + assert "shard_key" in json.loads(resp["body"])["error"] + + def test_healpix_requires_child_order(self, handler_mod): + # child_order omitted on a HEALPix run -> rejected. + event = _base_event(_healpix_config_dict()) + resp = handler_mod._handle_process(event, _context()) + assert resp["statusCode"] == 400 + assert "child_order" in json.loads(resp["body"])["error"] + + +class TestProcessEventDispatch: + """The gate passes and the shard key flows into ``process_shard`` for both + a HEALPix event (with child_order) and a rectilinear event (without it).""" + + def _run(self, handler_mod, monkeypatch, event): + import zagg.grids as grids + import zagg.processing as processing + + captured = {} + + def fake_process_shard(grid, shard_key, granule_urls, **kwargs): + captured["shard_key"] = shard_key + meta = { + "shard_key": shard_key, + "cells_with_data": 0, + "total_obs": 0, + "granule_count": len(granule_urls), + "files_processed": 0, + "duration_s": 0.0, + "error": None, + } + return pd.DataFrame(), meta + + monkeypatch.setattr(processing, "process_shard", fake_process_shard) + monkeypatch.setattr(grids, "from_config", lambda *a, **k: MagicMock()) + resp = handler_mod._handle_process(event, _context()) + return resp, captured + + def test_healpix_dispatch(self, handler_mod, monkeypatch): + event = _base_event(_healpix_config_dict()) + event["child_order"] = 12 + resp, captured = self._run(handler_mod, monkeypatch, event) + assert resp["statusCode"] == 200 + assert captured["shard_key"] == 12345 + + def test_rectilinear_dispatch_without_child_order(self, handler_mod, monkeypatch): + event = _base_event(_rectilinear_config_dict()) + event["parent_order"] = None # rectilinear has no parent_order + assert "child_order" not in event + resp, captured = self._run(handler_mod, monkeypatch, event) + assert resp["statusCode"] == 200 + assert captured["shard_key"] == 12345 diff --git a/tests/test_runner.py b/tests/test_runner.py index 9d0312bc..b7d4649d 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -174,3 +174,42 @@ def test_endpoint_and_region_override(self): assert block["endpointUrl"] == "https://r2.example" assert block["region"] == "eu-west-1" assert "sessionToken" not in block + + +class TestInvokeLambdaCellEvent: + """The per-cell Lambda event uses the grid-neutral ``shard_key`` field, and + only forwards the HEALPix-specific ``child_order`` when it is set (#24).""" + + _CREDS = {"accessKeyId": "a", "secretAccessKey": "s", "sessionToken": "t"} + + def _captured_event(self, *, child_order): + from unittest.mock import MagicMock + + from zagg.runner import _invoke_lambda_cell + + client = MagicMock() + payload = MagicMock() + payload.read.return_value = json.dumps( + {"statusCode": 200, "body": json.dumps({"total_obs": 0, "duration_s": 0.0})} + ).encode() + client.invoke.return_value = {"Payload": payload, "FunctionError": None} + _invoke_lambda_cell( + client, (0,), 12345, 6, child_order, + ["s3://b/g.h5"], "s3://out/x.zarr", self._CREDS, + function_name="process-shard", config_dict=None, max_workers=4, + ) + return json.loads(client.invoke.call_args.kwargs["Payload"]) + + def test_healpix_event_uses_shard_key_and_keeps_child_order(self): + event = self._captured_event(child_order=12) + assert event["shard_key"] == 12345 + assert "parent_morton" not in event + assert event["parent_order"] == 6 + assert event["child_order"] == 12 + + def test_non_healpix_event_omits_child_order(self): + # Rectilinear runs pass child_order=None; the field is dropped. + event = self._captured_event(child_order=None) + assert event["shard_key"] == 12345 + assert "child_order" not in event + assert "parent_morton" not in event diff --git a/tests/test_runner_concurrency.py b/tests/test_runner_concurrency.py index 6e489955..fc39d4b9 100644 --- a/tests/test_runner_concurrency.py +++ b/tests/test_runner_concurrency.py @@ -92,7 +92,7 @@ def test_pool_sized_to_clamped_workers(self, lambda_env, monkeypatch): "body": {"total_obs": 5}, "error": None, "lambda_duration": 1.0, - "morton": 0, + "shard_key": 0, }, ) summary = runner._run_lambda(