Reduce default GCS target size to 400MiB for OCDBT in GCS

The default target data file size for GCS OCDBT checkpoints is reduced from 2 GiB to 400 MiB for improved write and read performance. This also reduces the impact of a single GCS write taking too long.

PiperOrigin-RevId: 908509104

Author: samos123

checkpoint/CHANGELOG.md

@@ -20,6 +20,10 @@ merging.
 leaves can also be saved as individual checkpointables.
 - Move MTC files to multi_tier_checkpointing and use local checkpoint engine
 
+### Changed
+
+- Reduced default OCDBT target file size from 2GiB to 400MiB for GCS paths.
+
 ## [0.11.36] - 2026-04-14
 
 ### Added

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

@@ -34,10 +34,10 @@
 from orbax.checkpoint._src.metadata import sharding as sharding_metadata
 from orbax.checkpoint._src.metadata import value as value_metadata
 from orbax.checkpoint._src.path import async_path
+from orbax.checkpoint._src.path import gcs_utils
 from orbax.checkpoint._src.serialization import types
 import tensorstore as ts
 
-
 JsonSpec: TypeAlias = dict[str, Any]
 Shape: TypeAlias = arrays_types.Shape
 DType: TypeAlias = arrays_types.DType
@@ -51,6 +51,7 @@
 REPLICA_SUBDIR_SUFFIX = 'replica_'
 _OCDBT_PROCESS_ID_RE = r'[A-Za-z0-9]+'
 _DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE = 2**31  # 2 GiB
+_GCS_OCDBT_TARGET_DATA_FILE_SIZE = 400 * 2**20  # 400 MiB
 
 ZARR_VER2 = 'zarr'
 ZARR_VER3 = 'zarr3'
@@ -240,19 +241,45 @@ def build_kvstore_tspec_for_merge(
   )
 
 
+def _get_backend_ocdbt_target_data_file_size(
+    kvstore_spec: JsonSpec | None,
+) -> int:
+  """Gets OCDBT target data file size based on kvstore spec."""
+  if kvstore_spec is None:
+    return _DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE
+  base = kvstore_spec.get('base')
+
+  if isinstance(base, str):
+    # OCDBT base is generally a string when it's a GCS path.
+    if gcs_utils.is_gcs_path(epath.Path(base)):
+      return _GCS_OCDBT_TARGET_DATA_FILE_SIZE
+  elif isinstance(base, dict):
+    # OCDBT base can also be a dict with 'driver' and 'path' keys.
+    if base.get('driver') == 'gcs':
+      return _GCS_OCDBT_TARGET_DATA_FILE_SIZE
+    path_str = base.get('path')
+    if path_str and gcs_utils.is_gcs_path(epath.Path(path_str)):
+      return _GCS_OCDBT_TARGET_DATA_FILE_SIZE
+
+  return _DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE
+
+
 def add_ocdbt_write_options(
     kvstore_tspec: JsonSpec,
     target_data_file_size: int | None = None,
 ) -> None:
   """Adds write-specific options to a TensorStore OCDBT KVStore spec."""
-  if target_data_file_size is not None:
-    # TODO: b/354139177 - disallow too small values, too.
-    if target_data_file_size < 0:
-      raise ValueError(
-          'OCDBT target_data_file_size must be >= 0, where 0 means no limit'
-          f'; got {target_data_file_size}'
-      )
-    kvstore_tspec['target_data_file_size'] = target_data_file_size
+  if target_data_file_size is None:
+    target_data_file_size = _get_backend_ocdbt_target_data_file_size(
+        kvstore_tspec
+    )
+  # TODO: b/354139177 - Disallow too small values, too.
+  if target_data_file_size < 0:
+    raise ValueError(
+        'OCDBT target_data_file_size must be >= 0, where 0 means no limit'
+        f'; got {target_data_file_size}'
+    )
+  kvstore_tspec['target_data_file_size'] = target_data_file_size
 
   kvstore_tspec['config'] = {
       # Store .zarray metadata inline but not large chunks.
@@ -342,13 +369,14 @@ def calculate_chunk_byte_size(
     *,
     chunk_byte_size: int | None,
     ocdbt_target_data_file_size: int | None = None,
+    kvstore_spec: JsonSpec | None = None,
 ) -> int | None:
   """Selects chunk byte size to fit both target data file and chunk sizes."""
   # Check if the chunk size would exceed ocdbt target file size.
   if ocdbt_target_data_file_size is None:
-    # Set to default used by TensorStore
-    # (from https://google.github.io/tensorstore/kvstore/ocdbt/index.html).
-    ocdbt_target_data_file_size = _DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE
+    ocdbt_target_data_file_size = _get_backend_ocdbt_target_data_file_size(
+        kvstore_spec
+    )
 
   if ocdbt_target_data_file_size == 0:
     # No limit.
@@ -509,6 +537,7 @@ def __init__(
           target_storage_dtype,
           chunk_byte_size=chunk_byte_size,
           ocdbt_target_data_file_size=ocdbt_target_data_file_size,
+          kvstore_spec=tspec['kvstore'],
       )
     # Choose chunk shape.
     chunk_shape = subchunking.choose_chunk_shape(
@@ -719,6 +748,7 @@ def get_json_tspec_write(
         dtype,
         chunk_byte_size=chunk_byte_size,
         ocdbt_target_data_file_size=ocdbt_target_data_file_size,
+        kvstore_spec=tspec['kvstore'],
     )
 
   chunk_shape = subchunking.choose_chunk_shape(

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

@@ -33,12 +33,20 @@ class AddOcdbtWriteOptionsTest(parameterized.TestCase):
   def test_ocdbt_target_data_file_size_none(self):
     kvstore_tspec = {}
     ts_utils.add_ocdbt_write_options(kvstore_tspec, target_data_file_size=None)
-    self.assertNotIn('target_data_file_size', kvstore_tspec)
+    self.assertIn('target_data_file_size', kvstore_tspec)
+    self.assertEqual(
+        kvstore_tspec['target_data_file_size'],
+        ts_utils._DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE,
+    )
 
   def test_ocdbt_target_data_file_size_none_is_default(self):
     kvstore_tspec = {}
     ts_utils.add_ocdbt_write_options(kvstore_tspec)
-    self.assertNotIn('target_data_file_size', kvstore_tspec)
+    self.assertIn('target_data_file_size', kvstore_tspec)
+    self.assertEqual(
+        kvstore_tspec['target_data_file_size'],
+        ts_utils._DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE,
+    )
 
   def test_ocdbt_target_data_file_size_rejects_negative_value(self):
     with self.assertRaises(ValueError):
@@ -58,6 +66,56 @@ def test_ocdbt_target_data_file_size_sets_value(
     )
 
 
+class GetBackendOcdbtTargetDataFileSizeTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='none_spec',
+          kvstore_spec=None,
+          expected_size=ts_utils._DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE,
+      ),
+      dict(
+          testcase_name='base_is_gcs_str',
+          kvstore_spec={'base': 'gs://bucket/path'},
+          expected_size=ts_utils._GCS_OCDBT_TARGET_DATA_FILE_SIZE,
+      ),
+      dict(
+          testcase_name='base_is_local_str',
+          kvstore_spec={'base': '/tmp/path'},
+          expected_size=ts_utils._DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE,
+      ),
+      dict(
+          testcase_name='base_is_gcs_driver_dict',
+          kvstore_spec={'base': {'driver': 'gcs', 'bucket': 'bucket'}},
+          expected_size=ts_utils._GCS_OCDBT_TARGET_DATA_FILE_SIZE,
+      ),
+      dict(
+          testcase_name='base_is_dict_with_gcs_path',
+          kvstore_spec={'base': {'driver': 'file', 'path': 'gs://bucket/path'}},
+          expected_size=ts_utils._GCS_OCDBT_TARGET_DATA_FILE_SIZE,
+      ),
+      dict(
+          testcase_name='base_is_dict_with_local_path',
+          kvstore_spec={'base': {'driver': 'file', 'path': '/tmp/path'}},
+          expected_size=ts_utils._DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE,
+      ),
+      dict(
+          testcase_name='spec_without_base',
+          kvstore_spec={},
+          expected_size=ts_utils._DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE,
+      ),
+  )
+  def test_get_backend_ocdbt_target_data_file_size(
+      self,
+      kvstore_spec: ts_utils.JsonSpec | None,
+      expected_size: int,
+  ):
+    self.assertEqual(
+        ts_utils._get_backend_ocdbt_target_data_file_size(kvstore_spec),
+        expected_size,
+    )
+
+
 class BuildArrayTSpecForWriteTest(parameterized.TestCase):
 
   def setUp(self):
@@ -231,11 +289,41 @@ def test_ocdbt_kvstore_with_gcs_path(
         os.path.join(expected_directory or directory, 'ocdbt.process_0'),
     )
     self.assertEqual(kvstore_tspec['path'], self.param_name)
+    self.assertEqual(
+        kvstore_tspec['target_data_file_size'],
+        ts_utils._GCS_OCDBT_TARGET_DATA_FILE_SIZE,
+    )
 
-  @parameterized.product(use_zarr3=(True, False))
-  def test_ocdbt_kvstore_default_target_data_file_size(self, use_zarr3: bool):
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='zarr3_None',
+          use_zarr3=True,
+          directory_override=None,
+      ),
+      dict(
+          testcase_name='zarr3_tfhub',
+          use_zarr3=True,
+          directory_override='/tfhub/prod/model',
+      ),
+      dict(
+          testcase_name='no_zarr3_None',
+          use_zarr3=False,
+          directory_override=None,
+      ),
+      dict(
+          testcase_name='no_zarr3_tfhub',
+          use_zarr3=False,
+          directory_override='/tfhub/prod/model',
+      ),
+  )
+  def test_ocdbt_kvstore_default_target_data_file_size(
+      self,
+      use_zarr3: bool,
+      directory_override: str | None,
+  ):
+    directory = directory_override or self.directory
     tspec = self.array_write_spec_constructor(
-        directory=self.directory,
+        directory=directory,
         relative_array_filename=self.param_name,
         use_zarr3=use_zarr3,
         use_ocdbt=True,
@@ -244,7 +332,11 @@ def test_ocdbt_kvstore_default_target_data_file_size(self, use_zarr3: bool):
     self.assertEqual(tspec.metadata.use_zarr3, use_zarr3)
     self.assertTrue(tspec.metadata.use_ocdbt)
     self.assertEqual(tspec.json['kvstore']['driver'], 'ocdbt')
-    self.assertNotIn('target_data_file_size', tspec.json['kvstore'])
+    self.assertIn('target_data_file_size', tspec.json['kvstore'])
+    self.assertEqual(
+        tspec.json['kvstore']['target_data_file_size'],
+        ts_utils._DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE,
+    )
 
   @parameterized.named_parameters(
       dict(testcase_name='none', target_data_file_size=None),
@@ -267,7 +359,11 @@ def test_ocdbt_kvstore_target_data_file_size(
     kvstore_tspec = tspec.json['kvstore']
     self.assertEqual(kvstore_tspec['driver'], 'ocdbt')
     if target_data_file_size is None:
-      self.assertNotIn('target_data_file_size', kvstore_tspec)
+      self.assertIn('target_data_file_size', kvstore_tspec)
+      self.assertEqual(
+          kvstore_tspec['target_data_file_size'],
+          ts_utils._DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE,
+      )
     else:
       self.assertEqual(
           kvstore_tspec['target_data_file_size'], target_data_file_size
@@ -593,6 +689,50 @@ def test_chunk_byte_size_is_adjusted_for_target_data_file_size(
           expected_chunk_byte_size_limit,
       )
 
+  @parameterized.product(
+      chunk_byte_size=[None, 512 * 2**20],
+      use_zarr3=[True, False],
+  )
+  def test_gcs_chunk_byte_size_is_adjusted_for_target_data_file_size(
+      self,
+      chunk_byte_size: int | None,
+      use_zarr3: bool,
+  ):
+    gcs_dir = 'gs://gcs_bucket/object_path'
+    self.shape = (8 * 1024, 2 * 1024, 4 * 1024)
+    self.write_shape = (2 * 1024, 1024, 2 * 1024)
+    storage_dtype = self.dtype
+
+    tspec = ts_utils.ArrayWriteSpec(
+        directory=gcs_dir,
+        relative_array_filename=self.param_name,
+        global_shape=self.shape,
+        write_shape=self.write_shape,
+        dtype=self.dtype,
+        target_dtype=None,
+        chunk_byte_size=chunk_byte_size,
+        use_zarr3=use_zarr3,
+        use_ocdbt=True,
+        process_id='w13',
+        ocdbt_target_data_file_size=None,
+    )
+    self.assertEqual(tspec.metadata.dtype, storage_dtype)
+    chunk_shape = self._get_chunk_shape_from_tspec(
+        tspec.json,
+        use_zarr3=use_zarr3,
+    )
+    np.testing.assert_array_equal(chunk_shape, tspec.metadata.chunk_shape)
+    self.assertTrue(
+        subchunking.validate_divisible_shapes(self.write_shape, chunk_shape),
+        f'Write shape {self.write_shape} is not divisible by chunk shape'
+        f' {chunk_shape}.',
+    )
+
+    self.assertLessEqual(
+        math.prod(chunk_shape) * storage_dtype.itemsize,
+        ts_utils._GCS_OCDBT_TARGET_DATA_FILE_SIZE,
+    )
+
   def test_maybe_cloud_storage(self):
     gs_path = 'gs://some-buck/path'
     gs_spec = serialization.get_tensorstore_spec(gs_path, ocdbt=True)