Enable OCDBT read coalescing for remote storage

PiperOrigin-RevId: 725709173

Author: ChromeHearts

checkpoint/CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.11.5] - 2025-02-10
+
+### Fixed
+
+- Enable OCDBT read coalescing for remote storage
+
 ### Added
 
 - `ocp.metadata.get_step_metadata(path)` to public api.

checkpoint/orbax/checkpoint/_src/serialization/BUILD

@@ -49,6 +49,7 @@ py_test(
     name = "tensorstore_utils_test",
     srcs = ["tensorstore_utils_test.py"],
     deps = [
+        ":serialization",
         ":tensorstore_utils",
         "//checkpoint/orbax/checkpoint/_src/arrays:subchunking",
         "//checkpoint/orbax/checkpoint/_src/arrays:types",

checkpoint/orbax/checkpoint/_src/serialization/serialization.py

@@ -40,12 +40,6 @@
 TS_CONTEXT = ts.Context({'file_io_concurrency': {'limit': 128}})
 _REMOVED_VALUE = 'Value removed'
 _CHECKPOINT_SUCCESS = 'checkpoint_write_success'
-_REMOTE_URL_PREFIXES = ['gs://', 's3://']
-_REMOTE_DRIVER_VALIDATIONS = [
-    {'driver': 'gcs', 'path_regex': None},
-    {'driver': 's3', 'path_regex': None},
-]
-
 
 Index = types.Index
 Layout = layout.Layout
@@ -110,42 +104,6 @@ def get_tensorstore_spec(ckpt_path: str, ocdbt: bool = False):
   return spec
 
 
-def is_remote_storage(tspec: Union[Dict[str, Any], str]) -> bool:
-  """Detect if user is using remote storages.
-
-  This can detect common defines and unable to detect some corner cases such as
-  using gcsfuse.
-
-  Args:
-    tspec: Tensorstore spec.
-
-  Returns:
-    True if the spec is using remote storage.
-  """
-  if isinstance(tspec, str):
-    # KvStoreUrl
-    if re.match(rf'^({"|".join(_REMOTE_URL_PREFIXES)})', tspec):
-      return True
-    else:
-      return False
-
-  for key in ('base', 'kvstore'):
-    if key in tspec:
-      return is_remote_storage(tspec[key])
-
-  if 'driver' in tspec:
-    for rule in _REMOTE_DRIVER_VALIDATIONS:
-      if tspec['driver'] == rule['driver']:
-        if rule['path_regex'] is None:
-          return True
-
-        # check if path matches the regex.
-        if re.match(rule['path_regex'], tspec['path']):
-          return True
-
-  return False
-
-
 class ByteLimiter(Protocol):
 
   async def wait_for_bytes(self, requested_bytes: int):

checkpoint/orbax/checkpoint/_src/serialization/serialization_test.py

@@ -472,25 +472,6 @@ def test_get_tensorstore_spec_not_absolute_path(self):
     ):
       serialization.get_tensorstore_spec(path, ocdbt=True)
 
-  def test_maybe_cloud_storage(self):
-    gs_path = 'gs://some-buck/path'
-    gs_spec = serialization.get_tensorstore_spec(gs_path, ocdbt=True)
-    self.assertTrue(serialization.is_remote_storage(gs_spec))
-
-    local_path = '/tmp/checkpoint'
-    local_spec = serialization.get_tensorstore_spec(local_path, ocdbt=True)
-    self.assertFalse(serialization.is_remote_storage(local_spec))
-
-    nested_tspec = {
-        'driver': 'cast',
-        'dtype': 'int32',
-        'base': {
-            'driver': 'zarr',
-            'kvstore': {'driver': 'ocdbt', 'base': 's3://some-bucket/path'},
-        },
-    }
-    self.assertTrue(serialization.is_remote_storage(nested_tspec))
-
   def test_deserialization_with_int4(self):
     dtype = jnp.int4
     shape = (8, 2)

checkpoint/orbax/checkpoint/_src/serialization/tensorstore_utils.py

@@ -17,7 +17,7 @@
 import math
 import os
 import re
-from typing import Any, TypeAlias
+from typing import Any, Dict, TypeAlias, Union
 
 from absl import logging
 from jax import numpy as jnp
@@ -62,6 +62,13 @@
     **{'cache_pool#ocdbt': {'total_bytes_limit': 100000000}},
 }
 
+_REMOTE_URL_PREFIXES = ['gs://', 's3://']
+_REMOTE_DRIVER_VALIDATIONS = [
+    {'driver': 'gcs', 'path_regex': None},
+    {'driver': 's3', 'path_regex': None},
+]
+
+
 
 def get_ts_context(*, use_ocdbt: bool = True) -> ts.Context:
   """Creates a TensorStore context object.
@@ -155,7 +162,7 @@ def build_kvstore_tspec(
         'cache_pool': 'cache_pool#ocdbt',
     })
 
-    if default_driver != FILE_DRIVER:
+    if is_remote_storage(kv_spec):
       kv_spec.update({  # pytype: disable=attribute-error
           # Enable read coalescing.  This feature merges adjacent read_ops into
           # one, which could reduce I/O ops by a factor of 10. This is
@@ -415,3 +422,39 @@ def json(self) -> JsonSpec:
   def metadata(self) -> ArrayMetadata:
     """Checkpoint-relevant TensorStore metadata of the array."""
     return self._metadata
+
+
+def is_remote_storage(tspec: Union[Dict[str, Any], str]) -> bool:
+  """Detect if user is using remote storages.
+
+  This can detect common defines and unable to detect some corner cases such as
+  using gcsfuse.
+
+  Args:
+    tspec: Tensorstore spec.
+
+  Returns:
+    True if the spec is using remote storage.
+  """
+  if isinstance(tspec, str):
+    # KvStoreUrl
+    if re.match(rf'^({"|".join(_REMOTE_URL_PREFIXES)})', tspec):
+      return True
+    else:
+      return False
+
+  for key in ('base', 'kvstore'):
+    if key in tspec:
+      return is_remote_storage(tspec[key])
+
+  if 'driver' in tspec:
+    for rule in _REMOTE_DRIVER_VALIDATIONS:
+      if tspec['driver'] == rule['driver']:
+        if rule['path_regex'] is None:
+          return True
+
+        # check if path matches the regex.
+        if re.match(rule['path_regex'], tspec['path']):
+          return True
+
+  return False

checkpoint/orbax/checkpoint/_src/serialization/tensorstore_utils_test.py

@@ -21,6 +21,7 @@
 import numpy as np
 from orbax.checkpoint._src.arrays import subchunking
 from orbax.checkpoint._src.arrays import types
+from orbax.checkpoint._src.serialization import serialization
 from orbax.checkpoint._src.serialization import tensorstore_utils as ts_utils
 
 
@@ -599,6 +600,25 @@ def test_chunk_byte_size_is_adjusted_for_target_data_file_size(
           expected_chunk_byte_size_limit,
       )
 
+  def test_maybe_cloud_storage(self):
+    gs_path = 'gs://some-buck/path'
+    gs_spec = serialization.get_tensorstore_spec(gs_path, ocdbt=True)
+    self.assertTrue(ts_utils.is_remote_storage(gs_spec))
+
+    local_path = '/tmp/checkpoint'
+    local_spec = serialization.get_tensorstore_spec(local_path, ocdbt=True)
+    self.assertFalse(ts_utils.is_remote_storage(local_spec))
+
+    nested_tspec = {
+        'driver': 'cast',
+        'dtype': 'int32',
+        'base': {
+            'driver': 'zarr',
+            'kvstore': {'driver': 'ocdbt', 'base': 's3://some-bucket/path'},
+        },
+    }
+    self.assertTrue(ts_utils.is_remote_storage(nested_tspec))
+
 
 if __name__ == '__main__':
   absltest.main()

checkpoint/orbax/checkpoint/version.py

@@ -16,7 +16,7 @@
 
 # A new PyPI release will be pushed everytime `__version__` is increased.
 # Also modify version and date in CHANGELOG.
-__version__ = '0.11.4'
+__version__ = '0.11.5'
 
 
 # TODO: b/362813406 - Add latest change timestamp and commit number.