diff --git a/tests/profiling/collector/pprof_utils.py b/tests/profiling/collector/pprof_utils.py index 03937e81573..43bacfe9dc8 100644 --- a/tests/profiling/collector/pprof_utils.py +++ b/tests/profiling/collector/pprof_utils.py @@ -143,7 +143,9 @@ def __init__(self, *args, **kwargs): super().__init__(event_type=LockEventType.RELEASE, *args, **kwargs) -def parse_newest_profile(filename_prefix: str, assert_samples: bool = True) -> pprof_pb2.Profile: +def parse_newest_profile( + filename_prefix: str, assert_samples: bool = True, allow_penultimate: bool = False +) -> pprof_pb2.Profile: """Parse the newest profile that has given filename prefix. The profiler outputs profile file with following naming convention: ...pprof, and in tests, we'd want to parse @@ -160,6 +162,18 @@ def parse_newest_profile(filename_prefix: str, assert_samples: bool = True) -> p profile = pprof_pb2.Profile() profile.ParseFromString(serialized_data) + if allow_penultimate and len(profile.sample) == 0 and len(files) > 1: + # The newest profile file may be empty if it was just created and has not yet accumulated samples. + # In this case, we try to parse the (second-to-last) file instead, which is more likely + # to contain complete data. We temporarily rename the newest file so parse_newest_profile picks + # the previous one. + temp_filename = filename + ".temp" + os.rename(filename, temp_filename) + try: + return parse_newest_profile(filename_prefix, assert_samples, allow_penultimate=False) + finally: + os.rename(temp_filename, filename) + if assert_samples: assert len(profile.sample) > 0, "No samples found in profile" diff --git a/tests/profiling/test_main.py b/tests/profiling/test_main.py index 345d01fd08e..ca272f44cf0 100644 --- a/tests/profiling/test_main.py +++ b/tests/profiling/test_main.py @@ -60,7 +60,7 @@ def test_call_script_pprof_output(tmp_path: Path) -> None: stdout = stdout.decode() if isinstance(stdout, bytes) else stdout _, pid = list(s.strip() for s in stdout.strip().split("\n")) - profile = pprof_utils.parse_newest_profile(f"{filename}.{pid}") + profile = pprof_utils.parse_newest_profile(f"{filename}.{pid}", allow_penultimate=True) samples = pprof_utils.get_samples_with_value_type(profile, "cpu-time") assert len(samples) > 0 @@ -77,7 +77,7 @@ def test_fork(tmp_path: Path) -> None: assert exitcode == 0 stdout = stdout.decode() if isinstance(stdout, bytes) else stdout child_pid = stdout.strip() - profile = pprof_utils.parse_newest_profile(f"{filename}.{pid}") + profile = pprof_utils.parse_newest_profile(f"{filename}.{pid}", allow_penultimate=True) parent_expected_acquire_events = [ pprof_utils.LockAcquireEvent( caller_name="", @@ -99,7 +99,7 @@ def test_fork(tmp_path: Path) -> None: expected_acquire_events=parent_expected_acquire_events, expected_release_events=parent_expected_release_events, ) - child_profile = pprof_utils.parse_newest_profile(filename + "." + str(child_pid)) + child_profile = pprof_utils.parse_newest_profile(f"{filename}.{child_pid}") # We expect the child profile to not have lock events from the parent process # Note that assert_lock_events function only checks that the given events # exists, and doesn't assert that other events don't exist.