google-deepmind / google-deepmind/bagz
Useful `__repr__` for `Reader`
- Dominant language
- C++
- Stars
- 70
- Forks
- 13
- Avg merge
- 1d 3h
- Merged PRs (30d)
- 10
Description
When utilizing `grain` and `orbax`, the `__repr__` of the data source is used for validation when restoring checkpoints. Currently, this causes errors when using `bagz.Reader` for the data source, as it utilizes the default `__repr__` implementation that indicates the pointer location.
This results in the following error when resuming checkpoints that restore the data iterator:
```
ValueError: DataSource in checkpoint does not match datasource in dataloader.
data source in checkpoint:
data source in dataloader:
Grain uses `repr(data_source)` to validate the source, so you may need to implement a custom `__repr__`.
```
This can be patched currently with a trivial subclass that indicates the location of the underlying file:
```python
class ReprReader(bagz.Reader):
def __init__(self, fp, *args):
self.fp = fp
super().__init__(fp, *args)
def __repr__(self):
return f"Reader(fp={self.fp})"
```
Minimal example:
```python
import grain
import orbax.checkpoint as ocp
import bagz
import tempfile
import os
with tempfile.TemporaryDirectory() as temp_dir:
bagz_file = os.path.join(temp_dir, "demo.bagz")
with bagz.Writer(bagz_file) as writer:
writer.write(b"123")
def make_iter():
reader = bagz.Reader(bagz_file)
dataset = grain.DataLoader(data_source=reader, sampler=grain.samplers.SequentialSampler(1))
return iter(dataset)
mgr = ocp.CheckpointManager(os.path.join(temp_dir, "ckp"), item_names=("data_iter",))
# Write a checkpoint containing a data iterator sourced from bagz
data_iter = make_iter()
mgr.save(0, args=ocp.args.Composite(
data_iter=grain.checkpoint.CheckpointSave(data_iter)
))
mgr.wait_until_finished()
# Reinstantiate, to emulate restoring to an instance in a new process
data_iter = make_iter()
mgr.restore(0, args=ocp.args.Composite(
data_iter=grain.checkpoint.CheckpointRestore(data_iter)
))
```
Contributor guide
Assessment
This issue has not been assessed yet.