Code duplication around tensorstore_spec logic in orbax-checkpoint
- Dominant language
- Python
- Stars
- 535
- Forks
- 101
- Avg merge
- 1d 13h
- Merged PRs (30d)
- 23
Description
Hi Orbax team,
I was looking at Orbax code at the latest version==0.7.0 and found that pieces of code with quite heavy logic around tensorstore_spec creation seem to contain duplicates.
I'd like to know if this code duplication intended by design or I am welcome to submit a PR.
Here `get_tensorstore_spec` is a part of public API, and I can't find any usage of `get_tensorstore_spec` by `orbax-checkpoint` itself
https://github.com/google/orbax/blob/8b4e90d573082a5c7caa5f99c51db376f62a6995/checkpoint/orbax/checkpoint/serialization.py#L97C5-L124
And here is a very similar piece of code in `build_kvstore_tspec` in `_internal` package, and `build_kvstore_tspec` is used heavily by `type_handlers.py`
https://github.com/google/orbax/blob/8b4e90d573082a5c7caa5f99c51db376f62a6995/checkpoint/orbax/checkpoint/_src/serialization/tensorstore_utils.py#L62-L135
Would you consider `get_tensorstore_spec` to reuse `build_kvstore_tspec` under the hood?
---
Also, there seems to be a bit of obscurity with default ts_context value.
- In `orbax/checkpoint/serialization.py`, there is [TS_CONTEXT](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/serialization.py#L39) in public `serialization.py` that is used as a default value of `context` in [async_serialize](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/serialization.py#L352) (actually, `orbax` does not use `async_serialize` anywhere and recommends using `async_serialize_shards`) , [async_serialize_shards](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/serialization.py#L405), [async_deserialize](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/serialization.py#L573) and by [StringHandler](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/type_handlers.py#L1700).
- At the same time, in`type_handlers.py`, there is [get_ts_context()](https://github.com/google/orbax/blob/8b4e90d573082a5c7caa5f99c51db376f62a6995/checkpoint/orbax/checkpoint/type_handlers.py#L735-L737) (it references `_DEFAULT_OCDBT_TS_CONTEXT`), and `get_ts_context` is used by all other handler implementations.
So, `TS_CONTEXT` from `serialization.py` seems to be never used by common checkpoint IO code.
Should we somehow leave only 1 source of truth for default ts_context values?
Contributor guide
Assessment
This issue has not been assessed yet.