Tensorstore spec configuration
- 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
Assessment
This issue has not been assessed yet.