google / google/orbax

Tensorstore spec configuration

Open
#872 7 comments 0 reactions 0 assignees View on GitHub
checkpoint type:feature
Dominant language
Python
Stars
535
Forks
101
Avg merge
1d 13h
Merged PRs (30d)
23

Description

Hi Orbax community,

Under the hood, Orbax uses [TensorStore](https://google.github.io/tensorstore) for tensor IO, TensorStore integration is a part of [type_handlers.py](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/type_handlers.py).

TensorStore comes with KVStore implementations for [File](https://google.github.io/tensorstore/kvstore/file/index.html), [GCS](https://google.github.io/tensorstore/kvstore/gcs/index.html), [S3](https://google.github.io/tensorstore/kvstore/s3/index.html), [GRPC](https://google.github.io/tensorstore/kvstore/tsgrpc/index.html), etc.

Unfortunately, TensorStore integration part is quite rigid and does not support any form of extension from client code.
Namely, it only supports 'gcs' and 'file' kvstores, their spec is hard-coded and can't be configured.

I guess Orbax should provide a way to configure custom kvstore spec based on `directory`. I.e. make `typed_handlers.py` flexible and extendable and avoid hard-coded if/else. So that, as a user, I can support alternative "directory -> kvstore spec" mapping when needed and patch other parts of tsspec too.

For instance, there could be a class like:
```python
class TsSpecStrategyBase(ABC):
def supported_for(self, directory: str) -> bool:
...

def get_spec(self, directory: str) -> dict:
...

class FileTsSpecStrategy(TsSpecStrategyBase):
...

class GrpcTsSpecStrategy(TsSpecStrategyBase):
...

class TsSpecStrategyResolver:
def register_strategy(self, strategy: TsSpecStrategyBase):
...

def resolve(self, directory: str) -> TsSpecStrategyBase:
...
```

---

**Motivation**. My colleagues have implemented tsgrpc-compatible storage that we want to use as a checkpoint storage. Unfortunately, we can't use it without custom patches to orbax code. Namely:
- we've added another `if/else` to [_get_tensorstore_spec](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/type_handlers.py#L782) to activate `grpc` driver for paths like `yt://*`.
- supported patches via env variables to `tspec['metadata']`, `tspec['kvstore']['config']`, `spec['kvstore'] (ocdbt)` in order to configure `experimental_read_coalescing_interval` and disable compression.

---

I will be happy to assist and submit a PR

Regards,
Simon

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.