Skip to content

Region batch jobs

batch

Build a frame cube with independent region jobs (PRD C6, plan task T37.4).

Modeled on earth-mover's serverless-datacube-demo. A coordinator:

  1. builds a lazy template of the output (for example dxr.chain(dxr.open_disp(granules))), which reads no pixel;
  2. writes the store's metadata once with :func:init_store;
  3. splits the frame into regions aligned to the shards with :func:plan_jobs; each :class:RegionJob serializes to JSON, one per AWS Batch array index or dask task;

and each worker runs :func:run_region: it opens the granules on its region, applies a processing function, overwrites its own shards and returns a :class:RegionResult. Shards are the unit of safe parallel writing, and a job owns whole shards, so a retry overwrites them cleanly. The results form a manifest, and :func:failed_jobs lists the jobs to resubmit.

The granule list is resolved once by the coordinator and travels in the jobs. Workers reading s3:// granules get temporary ASF credentials, cached and renewed after 50 minutes, so each job starts with valid ones.

Processing must be local in space: masking, corrections, chaining and per-pixel fits work region by region; a spatial reference or a deramp needs the whole frame, so apply it after the cube is assembled (lazily, on the store).

Examples:

>>> template = dxr.chain(dxr.open_disp(granules))
>>> shards = init_store(template, "frame.zarr")
>>> jobs = plan_jobs(template, "frame.zarr", granules, shards)
>>> results = [run_region(job, chain) for job in jobs]
>>> retry = failed_jobs(jobs, results)

RegionJob dataclass

One region of a frame cube: what to read, where to write it.

Examples:

>>> job = RegionJob(
...     job_id="F11116-y0-x0",
...     frame=11116,
...     store="frame.zarr",
...     region={"y": [0, 1024], "x": [0, 1024]},
...     granules=["g1.nc", "g2.nc"],
... )
>>> RegionJob.from_json(job.to_json()) == job
True
>>> job.slices["y"]
slice(0, 1024, None)
region instance-attribute 
region: dict[str, list[int]]

{"y": [start, stop], "x": [start, stop]} in pixels.

slices property 
slices: dict[str, slice]

The region as slices.

to_json 
to_json() -> str

Return the job as JSON (one line).

from_json classmethod 
from_json(text: str) -> RegionJob

Return the job encoded by :meth:to_json.

RegionResult dataclass

What a region job reports: success or the error, and how long it took.

Examples:

>>> result = run_region(job, chain)
>>> result.status, result.seconds
status instance-attribute 
status: str

"ok" or "failed".

init_store 

init_store(template: Dataset, store: Union[str, Path], *, preset: Preset = 'timeseries', target_mb: Optional[float] = None, compression_level: int = 3) -> dict[str, int]

Write the metadata of template to a new Zarr v3 store, no data.

Coordinates are written; data variables are created empty with the encoding :func:disp_xr.to_zarr would use. Returns the shard sizes by dimension, which :func:plan_jobs aligns the regions to.

Examples:

>>> template = dxr.chain(dxr.open_disp(granules))
>>> shards = init_store(template, "frame.zarr")

plan_regions 

plan_regions(template: Dataset, shards: dict[str, int]) -> list[dict[str, list[int]]]

Return regions of whole shards in y and x (all epochs each).

Examples:

>>> import numpy as np, xarray as xr
>>> template = xr.Dataset(coords={"y": np.arange(300), "x": np.arange(500)})
>>> len(plan_regions(template, {"y": 256, "x": 256}))
4

plan_jobs 

plan_jobs(template: Dataset, store: Union[str, Path], granules: Sequence[Union[str, Path]], shards: dict[str, int], *, parameters: Optional[dict[str, Any]] = None) -> list[RegionJob]

Return one :class:RegionJob per region of :func:plan_regions.

Examples:

>>> jobs = plan_jobs(template, "frame.zarr", granules, shards)
>>> lines = [job.to_json() + "\n" for job in jobs]
>>> open("jobs.jsonl", "w").writelines(lines)

run_region 

run_region(job: RegionJob, process: Callable[[Dataset, RegionJob], Dataset]) -> RegionResult

Open the job's region, process it and overwrite its shards.

process(ds, job) gets the region of open_disp(job.granules) and returns the variables of the template on that region. Errors are caught and reported in the result, so a coordinator can collect every job.

Examples:

>>> def chain(ds, job):
...     return dxr.chain(dxr.mask(ds, "recommended"))
>>> result = run_region(job, chain)

failed_jobs 

failed_jobs(jobs: Sequence[RegionJob], results: Sequence[RegionResult]) -> list[RegionJob]

Return the jobs without an "ok" result, to resubmit.

Examples:

>>> retry = failed_jobs(jobs, read_manifest("manifest.jsonl"))

write_manifest 

write_manifest(results: Sequence[RegionResult], path: Union[str, Path]) -> None

Append results to a JSON-lines manifest.

Examples:

>>> write_manifest([result], "manifest.jsonl")

read_manifest 

read_manifest(path: Union[str, Path]) -> list[RegionResult]

Return the results in a JSON-lines manifest (later lines win per job).

Examples:

>>> results = read_manifest("manifest.jsonl")