check if the pid is running before attaching

Signed-off-by: Keming <kemingy94@gmail.com>

Author: kemingy

Lib/profiling/sampling/cli.py

@@ -6,7 +6,7 @@
 import subprocess
 import sys
 
-from .sample import sample, sample_live
+from .sample import sample, sample_live, _is_process_running
 from .pstats_collector import PstatsCollector
 from .stack_collector import CollapsedStackCollector, FlamegraphCollector
 from .heatmap_collector import HeatmapCollector
@@ -596,6 +596,8 @@ def main():
 
 def _handle_attach(args):
     """Handle the 'attach' command."""
+    if not _is_process_running(args.pid):
+        raise sys.exit(f"Process with PID {args.pid} is not running.")
     # Check if live mode is requested
     if args.live:
         _handle_live_attach(args, args.pid)

Lib/profiling/sampling/sample.py

@@ -89,7 +89,7 @@ def sample(self, collector, duration_sec=10, *, async_aware=False):
                         collector.collect_failed_sample()
                         errors += 1
                     except Exception as e:
-                        if not self._is_process_running():
+                        if not _is_process_running(self.pid):
                             break
                         raise e from None
 
@@ -151,22 +151,6 @@ def sample(self, collector, duration_sec=10, *, async_aware=False):
                 f"({(expected_samples - num_samples) / expected_samples * 100:.2f}%)"
             )
 
-    def _is_process_running(self):
-        if sys.platform == "linux" or sys.platform == "darwin":
-            try:
-                os.kill(self.pid, 0)
-                return True
-            except ProcessLookupError:
-                return False
-        elif sys.platform == "win32":
-            try:
-                _remote_debugging.RemoteUnwinder(self.pid)
-            except Exception:
-                return False
-            return True
-        else:
-            raise ValueError(f"Unsupported platform: {sys.platform}")
-
     def _print_realtime_stats(self):
         """Print real-time sampling statistics."""
         if len(self.sample_intervals) < 2:
@@ -282,6 +266,25 @@ def _print_unwinder_stats(self):
             print(f"  {ANSIColors.YELLOW}Stale cache invalidations: {stale_invalidations}{ANSIColors.RESET}")
 
 
+def _is_process_running(pid):
+    if pid <= 0:
+        return False
+    if sys.platform == "linux" or sys.platform == "darwin":
+        try:
+            os.kill(pid, 0)
+            return True
+        except ProcessLookupError:
+            return False
+    elif sys.platform == "win32":
+        try:
+            _remote_debugging.RemoteUnwinder(pid)
+        except Exception:
+            return False
+        return True
+    else:
+        raise ValueError(f"Unsupported platform: {sys.platform}")
+
+
 def sample(
     pid,
     collector,

Lib/test/test_profiling/test_sampling_profiler/test_cli.py

@@ -434,6 +434,7 @@ def test_cli_default_collapsed_filename(self):
 
         with (
             mock.patch("sys.argv", test_args),
+            mock.patch("profiling.sampling.sample._is_process_running", return_value=True),
             mock.patch("profiling.sampling.cli.sample") as mock_sample,
         ):
             from profiling.sampling.cli import main
@@ -476,6 +477,7 @@ def test_cli_custom_output_filenames(self):
         for test_args, expected_filename, expected_format in test_cases:
             with (
                 mock.patch("sys.argv", test_args),
+                mock.patch("profiling.sampling.sample._is_process_running", return_value=True),
                 mock.patch("profiling.sampling.cli.sample") as mock_sample,
             ):
                 main()
@@ -516,6 +518,7 @@ def test_argument_parsing_basic(self):
 
         with (
             mock.patch("sys.argv", test_args),
+            mock.patch("profiling.sampling.sample._is_process_running", return_value=True),
             mock.patch("profiling.sampling.cli.sample") as mock_sample,
         ):
             from profiling.sampling.cli import main
@@ -540,6 +543,7 @@ def test_sort_options(self):
 
             with (
                 mock.patch("sys.argv", test_args),
+                mock.patch("profiling.sampling.sample._is_process_running", return_value=True),
                 mock.patch("profiling.sampling.cli.sample") as mock_sample,
             ):
                 from profiling.sampling.cli import main
@@ -554,6 +558,7 @@ def test_async_aware_flag_defaults_to_running(self):
 
         with (
             mock.patch("sys.argv", test_args),
+            mock.patch("profiling.sampling.sample._is_process_running", return_value=True),
             mock.patch("profiling.sampling.cli.sample") as mock_sample,
         ):
             from profiling.sampling.cli import main
@@ -570,6 +575,7 @@ def test_async_aware_with_async_mode_all(self):
 
         with (
             mock.patch("sys.argv", test_args),
+            mock.patch("profiling.sampling.sample._is_process_running", return_value=True),
             mock.patch("profiling.sampling.cli.sample") as mock_sample,
         ):
             from profiling.sampling.cli import main
@@ -585,6 +591,7 @@ def test_async_aware_default_is_none(self):
 
         with (
             mock.patch("sys.argv", test_args),
+            mock.patch("profiling.sampling.sample._is_process_running", return_value=True),
             mock.patch("profiling.sampling.cli.sample") as mock_sample,
         ):
             from profiling.sampling.cli import main

Lib/test/test_profiling/test_sampling_profiler/test_integration.py

@@ -17,7 +17,7 @@
     import profiling.sampling.sample
     from profiling.sampling.pstats_collector import PstatsCollector
     from profiling.sampling.stack_collector import CollapsedStackCollector
-    from profiling.sampling.sample import SampleProfiler
+    from profiling.sampling.sample import SampleProfiler, _is_process_running
 except ImportError:
     raise unittest.SkipTest(
         "Test only runs when _remote_debugging is available"
@@ -681,7 +681,7 @@ def test_is_process_running(self):
                 self.skipTest(
                     "Insufficient permissions to read the stack trace"
                 )
-            self.assertTrue(profiler._is_process_running())
+            self.assertTrue(_is_process_running(profiler.pid))
             self.assertIsNotNone(profiler.unwinder.get_stack_trace())
             subproc.process.kill()
             subproc.process.wait()
@@ -690,7 +690,7 @@ def test_is_process_running(self):
             )
 
         # Exit the context manager to ensure the process is terminated
-        self.assertFalse(profiler._is_process_running())
+        self.assertFalse(_is_process_running(profiler.pid))
         self.assertRaises(
             ProcessLookupError, profiler.unwinder.get_stack_trace
         )