diff --git a/src/lib.rs b/src/lib.rs index c013cc7..daa03ab 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -44,24 +44,33 @@ impl Default for ParseConfig { } fn maybe_remove_convert_frame_suffixes(frames: &mut Vec) { - let target_frames = [ - ("torch/_dynamo/convert_frame.py", "catch_errors"), - ("torch/_dynamo/convert_frame.py", "_convert_frame"), - ("torch/_dynamo/convert_frame.py", "_convert_frame_assert"), + let all_target_frames = [ + [ + ("torch/_dynamo/convert_frame.py", "catch_errors"), + ("torch/_dynamo/convert_frame.py", "_convert_frame"), + ("torch/_dynamo/convert_frame.py", "_convert_frame_assert"), + ], + [ + ("torch/_dynamo/convert_frame.py", "__call__"), + ("torch/_dynamo/convert_frame.py", "__call__"), + ("torch/_dynamo/convert_frame.py", "__call__"), + ], ]; let len = frames.len(); - if len >= target_frames.len() { - let suffix = &frames[len - target_frames.len()..]; - if suffix - .iter() - .zip(target_frames.iter()) - .all(|(frame, target)| { - simplify_filename(unintern_str(frame.filename).as_ref()) == target.0 - && frame.name == target.1 - }) - { - frames.truncate(len - target_frames.len()); + for target_frames in all_target_frames { + if len >= target_frames.len() { + let suffix = &frames[len - target_frames.len()..]; + if suffix + .iter() + .zip(target_frames.iter()) + .all(|(frame, target)| { + simplify_filename(unintern_str(frame.filename).as_ref()) == target.0 + && frame.name == target.1 + }) + { + frames.truncate(len - target_frames.len()); + } } } } diff --git a/src/types.rs b/src/types.rs index 2d1bb3d..8f91e11 100644 --- a/src/types.rs +++ b/src/types.rs @@ -160,7 +160,7 @@ pub fn simplify_filename<'a>(filename: &'a str) -> &'a str { if parts.len() > 1 { return parts[1]; } - let re = Regex::new(r"\d+e\d+-seed-nspid\d+_cgpid\d+-ns-\d+/").unwrap(); + let re = Regex::new(r"[^/]+-seed-nspid[^/]+/").unwrap(); if let Some(captures) = re.captures(filename) { if let Some(capture) = captures.get(0) { return &filename[capture.end()..];