Use regex for parsing step_dir #739
Conversation
| @@ -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) | |||
There was a problem hiding this comment.
Thanks! I do think we should be restrictive about the patterns that we accept, specifically that it ends with STEP_NUM_DIGITS digits. Should we update to a pattern that matches that and also add a test case?
There was a problem hiding this comment.
We run into an issue quite often where we'll have some checkpoint suffixed (i.e. checkpoints/step_00001000-ext) where the only way to run this with axlearn is to copy this directory to step_00001000. How would you suggest we incorporate a pattern like this? Happy to add test cases to this PR as well.
There was a problem hiding this comment.
I'll probably need to understand more about your use-case and why we need to suffix checkpoints. Can you open an internal PR for discussion?
Instead of assuming the end of a step directory will be a series of digits to be converted to a step number, uses regex to find the last group of
STEP_NUM_DIGITSfrom the directory passed into the function.