Support grain data checkpoint for elastic training

[aireenmei]

Description

• migrate RemoteIterator to colocated python class
• Add checkpointing logic to RemoteIterator, so data iterator in the colocated sidecar writes checkpoint to the checkpoint path, prevent sending data iterator state to the controller
• Add grain.ElasticIterator support controlled by flag grain_use_elastic_iterator, can be used with or without elastic training. This class allows recovering checkpoint with a dynamic scale (up or down), with the following limitations:
  (1). Only arrayrecord files are supported, parquet or tfrecord are not supported
  (2). Does not support many-to-one transformations, including packing, filtering
  (3). Does not support mixing datasets

Tests

Tested on Pathways saving and restoring data iterator checkpoints with different # of v5e-32 slices
jobset.yaml
Test scaling down:

• Start with 2 slices: log (Num devices: 64)
• Scale down to 1 slice: log (Num devices: 32)
  Test scaling up:
• Start with 1 slice: log(Num devices: 32)
• Scale up to 2 slices: log(Num devices: 64)

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 23.21429% with 129 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...rc/maxtext/input_pipeline/multihost_dataloading.py	20.28%	55 Missing ⚠️
src/maxtext/common/checkpointing.py	20.93%	25 Missing and 9 partials ⚠️
...rc/maxtext/input_pipeline/grain_data_processing.py	31.11%	26 Missing and 5 partials ⚠️
...rc/maxtext/input_pipeline/data_processing_utils.py	22.22%	4 Missing and 3 partials ⚠️
src/maxtext/input_pipeline/tfds_data_processing.py	0.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!