From 3d0e44aa5313811000f5dfa3be6edf1f00cc14be Mon Sep 17 00:00:00 2001
From: Nick Lusskin <nlusskin@apple.com>
Date: Mon, 7 Oct 2024 15:39:38 -0700
Subject: [PATCH] use regex for parsing step_dir

---
 axlearn/common/checkpointer.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/axlearn/common/checkpointer.py b/axlearn/common/checkpointer.py
index 979a3c840..b80931f5b 100644
--- a/axlearn/common/checkpointer.py
+++ b/axlearn/common/checkpointer.py
@@ -7,6 +7,7 @@
 import enum
 import json
 import os.path
+import re
 import tempfile
 import threading
 import time
@@ -81,8 +82,10 @@ class CheckpointValidationType(str, enum.Enum):
 
 
 def parse_step_from_dir(step_dir: str) -> int:
-    # TODO(markblee): use regex.
-    return int(step_dir[-STEP_NUM_DIGITS:])
+    step = re.findall(r'(\d{' + str(STEP_NUM_DIGITS) + r'})', step_dir)
+    if len(step) < 1:
+        raise ValueError(f"Could not find step in '{step_dir}'")
+    return int(step[-1])
 
 
 def check_state_structure(