google / google/orbax

Make GCS-style checkpointing configurable instead deciding it on is_gcs_path function output

Open
#710 4 comments 6 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, i'm using s3fs which has very similar restrictions as GCS (no atomic moves, for example).
To enable GCS-style checkpointing I'm patching orbax which seems very ugly

```
def patch_with_gcs_style_checkpointing_for_path_prefix(prefix):
def is_gcs_path(path: str) -> bool:
return os.fspath(path).startswith(prefix)

ocp.utils.is_gcs_path = is_gcs_path

patch_with_gcs_style_checkpointing_for_path_prefix("/mnt/s3")
```

Instead, I would prefer to have a parameter in `ocp.CheckpointManagerOptions` that would enable this checkpointing style.
Its default value may be False and in `ocp.CheckpointManager.__init__` there could be `self._gcs_style_checkpointing = gcs_style_checkpointing or is_gcs_path(directory)` which would be totally backward compatible.

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.