Add additional support for TENSORSTORE_GCS_BACKEND environment variable

This allows the user to configure the Tensorstore GCS backend (i.e. with `gcs` for http or `gcs_grpc` for grpc) for additional places in the code that did not previously have support, including when using ocdbt.

PiperOrigin-RevId: 809285436

Author: Orbax Authors

checkpoint/CHANGELOG.md

@@ -16,6 +16,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - #v1 Modify LeafHandler definitions so that `AbstractLeaf` or
 `Type[AbstractLeaf]` are always accepted as valid abstract values.
+- Configuring the `TENSORSTORE_GCS_BACKEND` environment variable is now
+supported for additional locations in the code, notably when using ocdbt.
 
 ## [0.11.25] - 2025-09-10
 

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

@@ -79,20 +79,22 @@ def _get_kvstore_for_gcs(ckpt_path: str):
   """
   m = re.fullmatch('^gs://([^/]*)/(.*)$', ckpt_path, re.DOTALL)
   if m is None:
-    raise ValueError(
-        'The ckpt_path should contain the bucket name and the '
-        f'file path inside the bucket. Got: {ckpt_path}'
-    )
+    # The path might only be a bucket name.
+    m = re.fullmatch('^gs://([^/]*)$', ckpt_path, re.DOTALL)
+    if m is None:
+      raise ValueError(
+          'The ckpt_path should contain the bucket name and the '
+          f'file path inside the bucket. Got: {ckpt_path}'
+      )
   gcs_bucket = m.group(1)
-  path_without_bucket = m.group(2)
+  path_without_bucket = m.group(2) if m.lastindex == 2 else None
   # TODO(stoelinga): Switch to gcs_grpc by default.
   # gcs_grpc performs roughly twice as fast as gcs backend.
   gcs_backend = os.environ.get('TENSORSTORE_GCS_BACKEND', 'gcs')
-  return {
-      'driver': f'{gcs_backend}',
-      'bucket': gcs_bucket,
-      'path': path_without_bucket,
-  }
+  spec = {'driver': gcs_backend, 'bucket': gcs_bucket}
+  if path_without_bucket:
+    spec['path'] = path_without_bucket
+  return spec
 
 
 def get_tensorstore_spec(ckpt_path: str, ocdbt: bool = False):
@@ -107,7 +109,7 @@ def get_tensorstore_spec(ckpt_path: str, ocdbt: bool = False):
       raise ValueError(f'Checkpoint path should be absolute. Got {ckpt_path}')
     base_path = os.path.dirname(ckpt_path)
     base_driver_spec = (
-        base_path
+        _get_kvstore_for_gcs(base_path)
         if is_gcs_path
         else {'driver': ts_utils.DEFAULT_DRIVER, 'path': base_path}
     )

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

@@ -475,7 +475,14 @@ def test_get_tensorstore_spec_ocdbt(self, path):
     spec = serialization.get_tensorstore_spec(path, ocdbt=True)
     is_gcs_path = path.startswith('gs://')
     if is_gcs_path:
-      self.assertEqual(spec['kvstore']['base'], os.path.dirname(path))
+      self.assertEqual(
+          spec['kvstore']['base'],
+          {
+              'driver': 'gcs',
+              'bucket': 'my',
+              'path': 'ckpt/dir',
+          },
+      )
     else:
       self.assertEqual(
           spec['kvstore']['base'],
@@ -493,6 +500,52 @@ def test_get_tensorstore_spec_not_absolute_path(self):
     ):
       serialization.get_tensorstore_spec(path, ocdbt=True)
 
+  @parameterized.named_parameters(
+      dict(testcase_name='none', backend=None, target_driver='gcs'),
+      dict(testcase_name='gcs', backend='gcs', target_driver='gcs'),
+      dict(
+          testcase_name='gcs_grpc', backend='gcs_grpc', target_driver='gcs_grpc'
+      ),
+  )
+  def test_get_tensorstore_spec_ocdbt_grpc(self, backend, target_driver):
+    if backend:
+      os.environ['TENSORSTORE_GCS_BACKEND'] = backend
+      self.addCleanup(lambda: os.environ.pop('TENSORSTORE_GCS_BACKEND'))
+    spec = serialization.get_tensorstore_spec(
+        'gs://my/ckpt/dir/path', ocdbt=True
+    )
+    self.assertEqual(
+        spec['kvstore']['base'],
+        {
+            'driver': target_driver,
+            'bucket': 'my',
+            'path': 'ckpt/dir',
+        },
+    )
+
+  @parameterized.named_parameters(
+      dict(testcase_name='none', backend=None, target_driver='gcs'),
+      dict(testcase_name='gcs', backend='gcs', target_driver='gcs'),
+      dict(
+          testcase_name='gcs_grpc', backend='gcs_grpc', target_driver='gcs_grpc'
+      ),
+  )
+  def test_get_tensorstore_spec_grpc(self, backend, target_driver):
+    if backend:
+      os.environ['TENSORSTORE_GCS_BACKEND'] = backend
+      self.addCleanup(lambda: os.environ.pop('TENSORSTORE_GCS_BACKEND'))
+    spec = serialization.get_tensorstore_spec(
+        'gs://my/ckpt/dir/path', ocdbt=False
+    )
+    self.assertEqual(
+        spec['kvstore'],
+        {
+            'driver': target_driver,
+            'bucket': 'my',
+            'path': 'ckpt/dir/path',
+        },
+    )
+
   def test_deserialization_with_int4(self):
     dtype = jnp.int4
     shape = (8, 2)

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

@@ -44,6 +44,7 @@
 ZARR_VER3 = 'zarr3'
 
 _GCS_PATH_RE = r'^gs://([^/]*)/(.*)$'
+_GCS_BUCKET_RE = r'^gs://([^/]*)$'
 
 # Even if the data is equal to the fill value, we still want to write it
 # to the checkpoint. This results in unnecessary writes in some edge
@@ -112,15 +113,35 @@ def get_ts_context(
 
 
 def _get_kvstore_for_gcs(ckpt_path: str) -> JsonSpec:
+  """Constructs a TensorStore kvstore spec for a GCS path.
+
+  Args:
+    ckpt_path: A GCS path of the form gs://<bucket>/<path>.
+
+  Returns:
+    A dictionary containing the TensorStore kvstore spec.
+
+  Raises:
+    ValueError: if ckpt_path is not a valid GCS path.
+  """
   m = re.fullmatch(_GCS_PATH_RE, ckpt_path, re.DOTALL)
   if m is None:
-    raise ValueError(
-        'The ckpt_path should contain the bucket name and the '
-        f'file path inside the bucket. Got: {ckpt_path}'
-    )
+    # The path might only be a bucket name.
+    m = re.fullmatch(_GCS_BUCKET_RE, ckpt_path, re.DOTALL)
+    if m is None:
+      raise ValueError(
+          'The ckpt_path should contain the bucket name and the '
+          f'file path inside the bucket. Got: {ckpt_path}'
+      )
   gcs_bucket = m.group(1)
-  path_without_bucket = m.group(2)
-  return {'driver': 'gcs', 'bucket': gcs_bucket, 'path': path_without_bucket}
+  path_without_bucket = m.group(2) if m.lastindex == 2 else None
+  # TODO(mirvine): Switch to gcs_grpc by default.
+  # gcs_grpc performs roughly twice as fast as gcs backend.
+  gcs_backend = os.environ.get('TENSORSTORE_GCS_BACKEND', 'gcs')
+  spec = {'driver': gcs_backend, 'bucket': gcs_bucket}
+  if path_without_bucket:
+    spec['path'] = path_without_bucket
+  return spec
 
 
 def build_kvstore_tspec(
@@ -165,7 +186,7 @@ def build_kvstore_tspec(
           directory, f'{PROCESS_SUBDIR_PREFIX}{process_id}'
       )
     base_driver_spec = (
-        directory
+        _get_kvstore_for_gcs(str(directory))
         if is_gcs_path
         else {'driver': default_driver, 'path': str(directory)}
     )

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

@@ -203,12 +203,12 @@ def test_ocdbt_kvstore(
       dict(
           testcase_name='regular_path',
           directory='gs://gcs_bucket/object_path',
-          expected_directory=None,
+          expected_directory='object_path',
       ),
       dict(
           testcase_name='path_with_single_slash',
           directory='gs:/gcs_bucket/object_path',
-          expected_directory='gs://gcs_bucket/object_path',
+          expected_directory='object_path',
       ),
   )
   def test_ocdbt_kvstore_with_gcs_path(
@@ -228,10 +228,60 @@ def test_ocdbt_kvstore_with_gcs_path(
     self.assertEqual(kvstore_tspec['driver'], 'ocdbt')
     self.assertEqual(
         kvstore_tspec['base'],
-        os.path.join(expected_directory or directory, 'ocdbt.process_0'),
+        {
+            'driver': 'gcs',
+            'bucket': 'gcs_bucket',
+            'path': os.path.join(
+                expected_directory or directory, 'ocdbt.process_0'
+            ),
+        },
     )
     self.assertEqual(kvstore_tspec['path'], self.param_name)
 
+  @parameterized.named_parameters(
+      dict(testcase_name='none', backend=None, target_driver='gcs'),
+      dict(testcase_name='gcs', backend='gcs', target_driver='gcs'),
+      dict(
+          testcase_name='gcs_grpc', backend='gcs_grpc', target_driver='gcs_grpc'
+      ),
+  )
+  def test_get_tensorstore_spec_ocdbt_grpc(self, backend, target_driver):
+    if backend:
+      os.environ['TENSORSTORE_GCS_BACKEND'] = backend
+      self.addCleanup(lambda: os.environ.pop('TENSORSTORE_GCS_BACKEND'))
+    spec = ts_utils.build_kvstore_tspec('gs://my/ckpt/dir/path', use_ocdbt=True)
+    self.assertEqual(
+        spec['base'],
+        {
+            'driver': target_driver,
+            'bucket': 'my',
+            'path': 'ckpt/dir/path',
+        },
+    )
+
+  @parameterized.named_parameters(
+      dict(testcase_name='none', backend=None, target_driver='gcs'),
+      dict(testcase_name='gcs', backend='gcs', target_driver='gcs'),
+      dict(
+          testcase_name='gcs_grpc', backend='gcs_grpc', target_driver='gcs_grpc'
+      ),
+  )
+  def test_get_tensorstore_spec_grpc(self, backend, target_driver):
+    if backend:
+      os.environ['TENSORSTORE_GCS_BACKEND'] = backend
+      self.addCleanup(lambda: os.environ.pop('TENSORSTORE_GCS_BACKEND'))
+    spec = ts_utils.build_kvstore_tspec(
+        'gs://my/ckpt/dir/path', use_ocdbt=False
+    )
+    self.assertEqual(
+        spec,
+        {
+            'driver': target_driver,
+            'bucket': 'my',
+            'path': 'ckpt/dir/path',
+        },
+    )
+
   @parameterized.product(use_zarr3=(True, False))
   def test_ocdbt_kvstore_default_target_data_file_size(self, use_zarr3: bool):
     tspec = self.array_write_spec_constructor(