Fix step_prefix option to allow for arbitrary naming

PiperOrigin-RevId: 832504226

Author: Orbax Authors

checkpoint/CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed
+
+- Fix `step_from_checkpoint_name` to allow the passed in checkpoint name to
+  include an arbitrary `step_prefix` with any character(s) such as underscores.
+
 ### Changed
 
 - Validate checkpoints before writing merged OCDBT database using in-memory

checkpoint/orbax/checkpoint/_src/path/step.py

@@ -38,6 +38,8 @@
 from orbax.checkpoint._src.path import gcs_utils
 from orbax.checkpoint._src.path import temporary_paths
 
+# Allowed checkpoint step naming using any non empty `step_prefix`.
+ALLOWED_STEP_NAME_PATTERN = r'(.+)_(\d+)'
 
 TMP_DIR_SUFFIX = temporary_paths.TMP_DIR_SUFFIX
 # prefix_1000.orbax-checkpoint-tmp-1010101
@@ -739,10 +741,8 @@ def step_from_checkpoint_name(name: str) -> int:
   """Returns the step from a checkpoint name. Also works for tmp checkpoints."""
   if name.isdigit():
     return int(name)
-  elif name.split('_')[-1].isdigit():
-    split = name.split('_')
-    if len(split) == 2 and split[0]:
-      return int(split[-1])
+  elif m := re.fullmatch(ALLOWED_STEP_NAME_PATTERN, name):
+    return int(m.group(2))
   elif tmp_match := re.match(TMP_DIR_STEP_PATTERN, name):
     return int(tmp_match.group(1))
   raise ValueError(f'Unrecognized name format: {name}.')

checkpoint/orbax/checkpoint/_src/path/step_test.py

@@ -471,11 +471,13 @@ def test_is_path_temporary(self, step_prefix):
       ('checkpoint_0000', 0),
       ('checkpoint_003400', 3400),
       ('foobar_1000', 1000),
+      ('foo_bar_1000', 1000),
       ('0.orbax-checkpoint-tmp-1010101', 0),
       ('0000.orbax-checkpoint-tmp-12323232', 0),
       ('foobar_1.orbax-checkpoint-tmp-12424424', 1),
       ('foobar_000505.orbax-checkpoint-tmp-13124', 505),
       ('checkpoint_16.orbax-checkpoint-tmp-123214324', 16),
+      ('foo_bar_1000.orbax-checkpoint-tmp-123123', 1000),
   )
   def test_step_from_checkpoint_name(self, name, step):
     self.assertEqual(step_lib.step_from_checkpoint_name(name), step)