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

@@ -19,6 +19,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/path/gcs_utils.py

@@ -15,6 +15,7 @@
 """Utils for interacting with GCS paths."""
 
 import functools
+import os
 from urllib import parse
 from etils import epath
 
@@ -26,19 +27,55 @@ def is_gcs_path(path: epath.Path) -> bool:
   return path.as_posix().startswith(_GCS_PATH_PREFIX)
 
 
-def parse_gcs_path(path: epath.PathLike) -> tuple[str, str]:
+def parse_gcs_path(
+    path: epath.PathLike, add_trailing_slash: bool = True
+) -> tuple[str, str]:
+  """Parses a GCS path into bucket and path within the bucket.
+
+  Args:
+    path: The GCS path to parse (e.g., "gs://my-bucket/path/to/object").
+    add_trailing_slash: Whether to ensure the returned path has a trailing
+      slash.
+
+  Returns:
+    A tuple containing the bucket name and the path within the bucket.
+  """
   parsed = parse.urlparse(str(path))
   assert parsed.scheme == 'gs', f'Unsupported scheme for GCS: {parsed.scheme}'
   # Strip the leading slash from the path.
   standardized_path = parsed.path
   if standardized_path.startswith('/'):
     standardized_path = standardized_path[1:]
   # Add a trailing slash if it's missing.
-  if not standardized_path.endswith('/'):
+  if add_trailing_slash and not standardized_path.endswith('/'):
     standardized_path = standardized_path + '/'
   return parsed.netloc, standardized_path
 
 
+def get_kvstore_for_gcs(ckpt_path: str):
+  """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.
+  """
+  gcs_bucket, path_without_bucket = parse_gcs_path(
+      ckpt_path, add_trailing_slash=False
+  )
+  # 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')
+  spec = {'driver': gcs_backend, 'bucket': gcs_bucket}
+  if path_without_bucket:
+    spec['path'] = path_without_bucket
+  return spec
+
+
 @functools.lru_cache(maxsize=32)
 def get_bucket(bucket_name: str):
   # pylint: disable=g-import-not-at-top

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

@@ -19,7 +19,6 @@
 from collections.abc import Mapping
 import contextlib
 import os
-import re
 from typing import Any, AsyncIterator, Dict, Optional, Protocol, Sequence, Union
 
 from absl import logging
@@ -32,6 +31,7 @@
 from orbax.checkpoint._src.arrays import numpy_utils as np_utils
 from orbax.checkpoint._src.arrays import types
 from orbax.checkpoint._src.multihost import multihost
+from orbax.checkpoint._src.path import gcs_utils
 from orbax.checkpoint._src.serialization import replica_slices
 from orbax.checkpoint._src.serialization import tensorstore_utils as ts_utils
 import tensorstore as ts
@@ -65,36 +65,6 @@ def _spec_has_metadata(tree):
   )
 
 
-def _get_kvstore_for_gcs(ckpt_path: str):
-  """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('^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)
-  # 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,
-  }
-
-
 def get_tensorstore_spec(ckpt_path: str, ocdbt: bool = False):
   """Constructs a TensorStore spec for the given checkpoint path."""
   # Normalize path to exclude trailing '/'. In GCS path case, we will need to
@@ -107,7 +77,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
+        gcs_utils.get_kvstore_for_gcs(base_path)
         if is_gcs_path
         else {'driver': ts_utils.DEFAULT_DRIVER, 'path': base_path}
     )
@@ -118,7 +88,7 @@ def get_tensorstore_spec(ckpt_path: str, ocdbt: bool = False):
     }
   else:
     if is_gcs_path:
-      spec['kvstore'] = _get_kvstore_for_gcs(ckpt_path)
+      spec['kvstore'] = gcs_utils.get_kvstore_for_gcs(ckpt_path)
     else:
       spec['kvstore'] = {'driver': ts_utils.DEFAULT_DRIVER, 'path': ckpt_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

@@ -25,6 +25,7 @@
 from orbax.checkpoint._src.arrays import subchunking
 from orbax.checkpoint._src.arrays import types
 from orbax.checkpoint._src.metadata import array_metadata
+from orbax.checkpoint._src.path import gcs_utils
 import tensorstore as ts
 
 JsonSpec: TypeAlias = dict[str, Any]
@@ -43,8 +44,6 @@
 ZARR_VER2 = 'zarr'
 ZARR_VER3 = 'zarr3'
 
-_GCS_PATH_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
 # cases, but it allows us to verify that data was actually written when
@@ -111,18 +110,6 @@ def get_ts_context(
 ### Building KvStore specs.
 
 
-def _get_kvstore_for_gcs(ckpt_path: str) -> JsonSpec:
-  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}'
-    )
-  gcs_bucket = m.group(1)
-  path_without_bucket = m.group(2)
-  return {'driver': 'gcs', 'bucket': gcs_bucket, 'path': path_without_bucket}
-
-
 def build_kvstore_tspec(
     directory: str,
     name: str | None = None,
@@ -165,7 +152,7 @@ def build_kvstore_tspec(
           directory, f'{PROCESS_SUBDIR_PREFIX}{process_id}'
       )
     base_driver_spec = (
-        directory
+        gcs_utils.get_kvstore_for_gcs(str(directory))
         if is_gcs_path
         else {'driver': default_driver, 'path': str(directory)}
     )
@@ -196,7 +183,7 @@ def build_kvstore_tspec(
     else:
       path = os.path.join(directory, name)
     if is_gcs_path:
-      kv_spec = _get_kvstore_for_gcs(path)
+      kv_spec = gcs_utils.get_kvstore_for_gcs(path)
     else:
       kv_spec = {'driver': default_driver, 'path': path}
 

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(