Parallel updates

[aldenks]

Make dataset updates parallel across batches of data variables, similar to how we parallelize backfills.

We have machinery to get_jobs that would give multiple region jobs to do an update.

The challenge is that certain steps need to be done before and after the main region job processing, namely

• For a standard zarr v3 dataset: only write the updated template metadata after all region jobs finish processing
• For an icechunk zarr:
  • Before any region jobs: create & checkout new temp branch, resize commit, fresh icechunk session, fork and serialize (pickle) it to object storage that each region job can access.
  • Each (parallel) region job loads the pickled session, does its normal writing of chunk data, and then writes the updated session back to a key in object storage that includes its worker index.
  • Finally something waits for all region jobs to finish, reads and merges all their seasons and does an icechunk commit.

In short we need to add a way to do some work before the region jobs main work, and other work after every parallel process has completed its main region job work.

I’d like to not add new infrastructure dependencies in addition to kubernetes / docker.

[aldenks]

From earthmover:

Deepak:

1. create & checkout new temp branch;
2. resize commit;
3. fork/write/merge/commit;
4. reset main to head of new temp branch;
5. delete branch

See https://github.com/earth-mover/forecast-datacube-demo/blob/24a258d21ff9cb54df5e84d731d7a09dd9beaa01/src/modal_app.py#L175 for an example

Sebastián Galkin
Quick comment: to make this safer it would be great to use the from_snapshot_id: str argument to reset_branch

that would guarantee you don't lose any data if somebody else made a commit to main you don't know about

[aldenks]

approach that fits in the k8s index job paradigm is to have an abstraction that lets you define work to be done if first that runs on worker index 0, or waits until that work is done on workers 1+.

And then similarly at the end, exits immediately if worker index isn’t the last, else if last worker waits until all other workers are done and then does its work.