11"""Command-line interface for the sampling profiler."""
22
33import argparse
4+ import importlib .util
45import os
6+ import selectors
57import socket
68import subprocess
79import sys
10+ import time
811
912from .sample import sample , sample_live
1013from .pstats_collector import PstatsCollector
@@ -92,6 +95,54 @@ def _parse_mode(mode_string):
9295 return mode_map [mode_string ]
9396
9497
98+ def _check_process_died (process ):
99+ """Check if process died and raise an error with stderr if available."""
100+ if process .poll () is None :
101+ return # Process still running
102+
103+ # Process died - try to get stderr for error message
104+ stderr_msg = ""
105+ if process .stderr :
106+ try :
107+ stderr_msg = process .stderr .read ().decode ().strip ()
108+ except (OSError , UnicodeDecodeError ):
109+ pass
110+
111+ if stderr_msg :
112+ raise RuntimeError (stderr_msg )
113+ raise RuntimeError (f"Process exited with code { process .returncode } " )
114+
115+
116+ def _wait_for_ready_signal (sync_sock , process , timeout ):
117+ """Wait for the ready signal from the subprocess, checking for early death."""
118+ deadline = time .monotonic () + timeout
119+ sel = selectors .DefaultSelector ()
120+ sel .register (sync_sock , selectors .EVENT_READ )
121+
122+ try :
123+ while True :
124+ _check_process_died (process )
125+
126+ remaining = deadline - time .monotonic ()
127+ if remaining <= 0 :
128+ raise socket .timeout ("timed out" )
129+
130+ if not sel .select (timeout = min (0.1 , remaining )):
131+ continue
132+
133+ conn , _ = sync_sock .accept ()
134+ try :
135+ ready_signal = conn .recv (_RECV_BUFFER_SIZE )
136+ finally :
137+ conn .close ()
138+
139+ if ready_signal != _READY_MESSAGE :
140+ raise RuntimeError (f"Invalid ready signal received: { ready_signal !r} " )
141+ return
142+ finally :
143+ sel .close ()
144+
145+
95146def _run_with_sync (original_cmd , suppress_output = False ):
96147 """Run a command with socket-based synchronization and return the process."""
97148 # Create a TCP socket for synchronization with better socket options
@@ -117,24 +168,24 @@ def _run_with_sync(original_cmd, suppress_output=False):
117168 ) + tuple (target_args )
118169
119170 # Start the process with coordinator
120- # Suppress stdout/stderr if requested (for live mode)
171+ # When suppress_output=True (live mode), capture stderr so we can
172+ # report errors if the process dies before signaling ready.
173+ # When suppress_output=False (normal mode), let stderr inherit so
174+ # script errors print to the terminal.
121175 popen_kwargs = {}
122176 if suppress_output :
123177 popen_kwargs ["stdin" ] = subprocess .DEVNULL
124178 popen_kwargs ["stdout" ] = subprocess .DEVNULL
125- popen_kwargs ["stderr" ] = subprocess .DEVNULL
179+ popen_kwargs ["stderr" ] = subprocess .PIPE
126180
127181 process = subprocess .Popen (cmd , ** popen_kwargs )
128182
129183 try :
130- # Wait for ready signal with timeout
131- with sync_sock .accept ()[0 ] as conn :
132- ready_signal = conn .recv (_RECV_BUFFER_SIZE )
184+ _wait_for_ready_signal (sync_sock , process , _SYNC_TIMEOUT )
133185
134- if ready_signal != _READY_MESSAGE :
135- raise RuntimeError (
136- f"Invalid ready signal received: { ready_signal !r} "
137- )
186+ # Close stderr pipe if we were capturing it
187+ if process .stderr :
188+ process .stderr .close ()
138189
139190 except socket .timeout :
140191 # If we timeout, kill the process and raise an error
@@ -632,6 +683,25 @@ def _handle_attach(args):
632683
633684def _handle_run (args ):
634685 """Handle the 'run' command."""
686+ # Validate target exists before launching subprocess
687+ if args .module :
688+ # Temporarily add cwd to sys.path so we can find modules in the
689+ # current directory, matching the coordinator's behavior
690+ cwd = os .getcwd ()
691+ added_cwd = False
692+ if cwd not in sys .path :
693+ sys .path .insert (0 , cwd )
694+ added_cwd = True
695+ try :
696+ if importlib .util .find_spec (args .target ) is None :
697+ sys .exit (f"Error: Module not found: { args .target } " )
698+ finally :
699+ if added_cwd :
700+ sys .path .remove (cwd )
701+ else :
702+ if not os .path .exists (args .target ):
703+ sys .exit (f"Error: Script not found: { args .target } " )
704+
635705 # Check if live mode is requested
636706 if args .live :
637707 _handle_live_run (args )
@@ -644,7 +714,10 @@ def _handle_run(args):
644714 cmd = (sys .executable , args .target , * args .args )
645715
646716 # Run with synchronization
647- process = _run_with_sync (cmd , suppress_output = False )
717+ try :
718+ process = _run_with_sync (cmd , suppress_output = False )
719+ except RuntimeError as e :
720+ sys .exit (f"Error: { e } " )
648721
649722 # Use PROFILING_MODE_ALL for gecko format
650723 mode = (
@@ -732,7 +805,10 @@ def _handle_live_run(args):
732805 cmd = (sys .executable , args .target , * args .args )
733806
734807 # Run with synchronization, suppressing output for live mode
735- process = _run_with_sync (cmd , suppress_output = True )
808+ try :
809+ process = _run_with_sync (cmd , suppress_output = True )
810+ except RuntimeError as e :
811+ sys .exit (f"Error: { e } " )
736812
737813 mode = _parse_mode (args .mode )
738814
0 commit comments