Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 20 additions & 10 deletions Lib/profiling/sampling/_sync_coordinator.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
the startup of target processes. It should not be called directly by users.
"""

import importlib.util
import os
import sys
import socket
Expand Down Expand Up @@ -138,7 +139,7 @@ def _execute_module(module_name: str, module_args: List[str]) -> None:
"""
# Replace sys.argv to match how Python normally runs modules
# When running 'python -m module args', sys.argv is ["__main__.py", "args"]
sys.argv = [f"__main__.py"] + module_args
sys.argv = ["__main__.py"] + module_args

try:
runpy.run_module(module_name, run_name="__main__", alter_sys=True)
Expand Down Expand Up @@ -215,22 +216,31 @@ def main() -> NoReturn:
# Set up execution environment
_setup_environment(cwd)

# Signal readiness to profiler
_signal_readiness(sync_port)

# Execute the target
if target_args[0] == "-m":
# Module execution
# Determine execution type and validate target exists
is_module = target_args[0] == "-m"
if is_module:
if len(target_args) < 2:
raise ArgumentError("Module name required after -m")

module_name = target_args[1]
module_args = target_args[2:]
_execute_module(module_name, module_args)

if importlib.util.find_spec(module_name) is None:
raise TargetError(f"Module not found: {module_name}")
else:
# Script execution
script_path = target_args[0]
script_args = target_args[1:]
# Match the path resolution logic in _execute_script
check_path = script_path if os.path.isabs(script_path) else os.path.join(cwd, script_path)
if not os.path.isfile(check_path):
raise TargetError(f"Script not found: {script_path}")

# Signal readiness to profiler
_signal_readiness(sync_port)

# Execute the target
if is_module:
_execute_module(module_name, module_args)
else:
_execute_script(script_path, script_args, cwd)

except CoordinatorError as e:
Expand Down
98 changes: 87 additions & 11 deletions Lib/profiling/sampling/cli.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,13 @@
"""Command-line interface for the sampling profiler."""

import argparse
import importlib.util
import os
import selectors
import socket
import subprocess
import sys
import time

from .sample import sample, sample_live
from .pstats_collector import PstatsCollector
Expand Down Expand Up @@ -92,6 +95,54 @@ def _parse_mode(mode_string):
return mode_map[mode_string]


def _check_process_died(process):
"""Check if process died and raise an error with stderr if available."""
if process.poll() is None:
return # Process still running

# Process died - try to get stderr for error message
stderr_msg = ""
if process.stderr:
try:
stderr_msg = process.stderr.read().decode().strip()
except (OSError, UnicodeDecodeError):
pass

if stderr_msg:
raise RuntimeError(stderr_msg)
raise RuntimeError(f"Process exited with code {process.returncode}")


def _wait_for_ready_signal(sync_sock, process, timeout):
"""Wait for the ready signal from the subprocess, checking for early death."""
deadline = time.monotonic() + timeout
sel = selectors.DefaultSelector()
sel.register(sync_sock, selectors.EVENT_READ)

try:
while True:
_check_process_died(process)

remaining = deadline - time.monotonic()
if remaining <= 0:
raise socket.timeout("timed out")

if not sel.select(timeout=min(0.1, remaining)):
continue

conn, _ = sync_sock.accept()
try:
ready_signal = conn.recv(_RECV_BUFFER_SIZE)
finally:
conn.close()

if ready_signal != _READY_MESSAGE:
raise RuntimeError(f"Invalid ready signal received: {ready_signal!r}")
return
finally:
sel.close()


def _run_with_sync(original_cmd, suppress_output=False):
"""Run a command with socket-based synchronization and return the process."""
# Create a TCP socket for synchronization with better socket options
Expand All @@ -117,24 +168,24 @@ def _run_with_sync(original_cmd, suppress_output=False):
) + tuple(target_args)

# Start the process with coordinator
# Suppress stdout/stderr if requested (for live mode)
# When suppress_output=True (live mode), capture stderr so we can
# report errors if the process dies before signaling ready.
# When suppress_output=False (normal mode), let stderr inherit so
# script errors print to the terminal.
popen_kwargs = {}
if suppress_output:
popen_kwargs["stdin"] = subprocess.DEVNULL
popen_kwargs["stdout"] = subprocess.DEVNULL
popen_kwargs["stderr"] = subprocess.DEVNULL
popen_kwargs["stderr"] = subprocess.PIPE

process = subprocess.Popen(cmd, **popen_kwargs)

try:
# Wait for ready signal with timeout
with sync_sock.accept()[0] as conn:
ready_signal = conn.recv(_RECV_BUFFER_SIZE)
_wait_for_ready_signal(sync_sock, process, _SYNC_TIMEOUT)

if ready_signal != _READY_MESSAGE:
raise RuntimeError(
f"Invalid ready signal received: {ready_signal!r}"
)
# Close stderr pipe if we were capturing it
if process.stderr:
process.stderr.close()

except socket.timeout:
# If we timeout, kill the process and raise an error
Expand Down Expand Up @@ -632,6 +683,25 @@ def _handle_attach(args):

def _handle_run(args):
"""Handle the 'run' command."""
# Validate target exists before launching subprocess
if args.module:
# Temporarily add cwd to sys.path so we can find modules in the
# current directory, matching the coordinator's behavior
cwd = os.getcwd()
added_cwd = False
if cwd not in sys.path:
sys.path.insert(0, cwd)
added_cwd = True
try:
if importlib.util.find_spec(args.target) is None:
sys.exit(f"Error: Module not found: {args.target}")
finally:
if added_cwd:
sys.path.remove(cwd)
else:
if not os.path.exists(args.target):
sys.exit(f"Error: Script not found: {args.target}")

# Check if live mode is requested
if args.live:
_handle_live_run(args)
Expand All @@ -644,7 +714,10 @@ def _handle_run(args):
cmd = (sys.executable, args.target, *args.args)

# Run with synchronization
process = _run_with_sync(cmd, suppress_output=False)
try:
process = _run_with_sync(cmd, suppress_output=False)
except RuntimeError as e:
sys.exit(f"Error: {e}")

# Use PROFILING_MODE_ALL for gecko format
mode = (
Expand Down Expand Up @@ -732,7 +805,10 @@ def _handle_live_run(args):
cmd = (sys.executable, args.target, *args.args)

# Run with synchronization, suppressing output for live mode
process = _run_with_sync(cmd, suppress_output=True)
try:
process = _run_with_sync(cmd, suppress_output=True)
except RuntimeError as e:
sys.exit(f"Error: {e}")

mode = _parse_mode(args.mode)

Expand Down
Loading
Loading