google-deepmind / google-deepmind/bagz

Useful `__repr__` for `Reader`

Open
#3 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.