Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
Prev Previous commit
Next Next commit
Fail fast
  • Loading branch information
savannahostrowski committed Dec 11, 2025
commit 2ee9bf94fe276436416f069ce87a86db75fda2be
28 changes: 18 additions & 10 deletions Lib/profiling/sampling/_sync_coordinator.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,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 +215,30 @@ 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)

import importlib.util
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:]
if not os.path.isfile(os.path.join(cwd, script_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
61 changes: 44 additions & 17 deletions Lib/profiling/sampling/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,11 @@

import argparse
import os
import select
import socket
import subprocess
import sys
import time

from .sample import sample, sample_live
from .pstats_collector import PstatsCollector
Expand Down Expand Up @@ -119,25 +121,50 @@ 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)
# Suppress stdout if requested (for live mode), but capture stderr
# so we can report errors if the process fails to start
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)

if ready_signal != _READY_MESSAGE:
# Wait for ready signal, but also check if process dies early
deadline = time.monotonic() + _SYNC_TIMEOUT
while True:
# Check if process died
if process.poll() is not None:
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"Invalid ready signal received: {ready_signal!r}"
f"Process exited with code {process.returncode}"
)

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

# Wait for socket with short timeout to allow process checks
ready, _, _ = select.select([sync_sock], [], [], min(0.1, remaining))
if ready:
with sync_sock.accept()[0] as conn:
ready_signal = conn.recv(_RECV_BUFFER_SIZE)
if ready_signal != _READY_MESSAGE:
raise RuntimeError(
f"Invalid ready signal received: {ready_signal!r}"
)
break

except socket.timeout:
# If we timeout, kill the process and raise an error
if process.poll() is None:
Expand Down Expand Up @@ -614,15 +641,6 @@ def _handle_attach(args):

def _handle_run(args):
"""Handle the 'run' command."""
# Validate target exists before launching
if args.module:
import importlib.util
if importlib.util.find_spec(args.target) is None:
sys.exit(f"Error: module not found: {args.target}")
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 Down Expand Up @@ -681,6 +699,14 @@ def _handle_run(args):

def _handle_live_attach(args, pid):
"""Handle live mode for an existing process."""
# Check if process exists
try:
os.kill(pid, 0)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure this works on windows. On Windows, os.kill() only supports signals SIGTERM, CTRL_C_EVENT, and CTRL_BREAK_EVENT. Signal 0 (which is the Unix idiom for checking if a process exists without sending a signal) is not supported on Windows.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This also has a TOCTOU problem: the process may die after you check and before we attach

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The fix probably needs to be done more on the line of https://github.com/python/cpython/pull/142655/changes (we should handle this there)

except ProcessLookupError:
sys.exit(f"Error: process not found: {pid}")
except PermissionError:
pass # Process exists, permission error will be handled later

mode = _parse_mode(args.mode)

# Determine skip_idle based on mode
Expand Down Expand Up @@ -720,6 +746,7 @@ def _handle_live_run(args):
cmd = (sys.executable, args.target, *args.args)

# Run with synchronization, suppressing output for live mode
# Note: _run_with_sync will raise if the process dies before signaling ready
process = _run_with_sync(cmd, suppress_output=True)

mode = _parse_mode(args.mode)
Expand Down
7 changes: 7 additions & 0 deletions Lib/test/test_profiling/test_sampling_profiler/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ def test_cli_module_argument_parsing(self):
mock.patch("profiling.sampling.cli.sample") as mock_sample,
mock.patch("subprocess.Popen") as mock_popen,
mock.patch("socket.socket") as mock_socket,
mock.patch("profiling.sampling.cli.select.select", return_value=([True], [], [])),
):
from profiling.sampling.cli import main
self._setup_sync_mocks(mock_socket, mock_popen)
Expand All @@ -96,6 +97,7 @@ def test_cli_module_with_arguments(self):
mock.patch("profiling.sampling.cli.sample") as mock_sample,
mock.patch("subprocess.Popen") as mock_popen,
mock.patch("socket.socket") as mock_socket,
mock.patch("profiling.sampling.cli.select.select", return_value=([True], [], [])),
):
self._setup_sync_mocks(mock_socket, mock_popen)
from profiling.sampling.cli import main
Expand All @@ -115,6 +117,7 @@ def test_cli_script_argument_parsing(self):
mock.patch("profiling.sampling.cli.sample") as mock_sample,
mock.patch("subprocess.Popen") as mock_popen,
mock.patch("socket.socket") as mock_socket,
mock.patch("profiling.sampling.cli.select.select", return_value=([True], [], [])),
):
self._setup_sync_mocks(mock_socket, mock_popen)
from profiling.sampling.cli import main
Expand All @@ -139,6 +142,7 @@ def test_cli_script_with_arguments(self):
mock.patch("profiling.sampling.cli.sample") as mock_sample,
mock.patch("subprocess.Popen") as mock_popen,
mock.patch("socket.socket") as mock_socket,
mock.patch("profiling.sampling.cli.select.select", return_value=([True], [], [])),
):
# Use the helper to set up mocks consistently
mock_process = self._setup_sync_mocks(mock_socket, mock_popen)
Expand Down Expand Up @@ -248,6 +252,7 @@ def test_cli_module_with_profiler_options(self):
mock.patch("profiling.sampling.cli.sample") as mock_sample,
mock.patch("subprocess.Popen") as mock_popen,
mock.patch("socket.socket") as mock_socket,
mock.patch("profiling.sampling.cli.select.select", return_value=([True], [], [])),
):
self._setup_sync_mocks(mock_socket, mock_popen)
from profiling.sampling.cli import main
Expand Down Expand Up @@ -278,6 +283,7 @@ def test_cli_script_with_profiler_options(self):
mock.patch("profiling.sampling.cli.sample") as mock_sample,
mock.patch("subprocess.Popen") as mock_popen,
mock.patch("socket.socket") as mock_socket,
mock.patch("profiling.sampling.cli.select.select", return_value=([True], [], [])),
):
self._setup_sync_mocks(mock_socket, mock_popen)
from profiling.sampling.cli import main
Expand Down Expand Up @@ -319,6 +325,7 @@ def test_cli_long_module_option(self):
mock.patch("profiling.sampling.cli.sample") as mock_sample,
mock.patch("subprocess.Popen") as mock_popen,
mock.patch("socket.socket") as mock_socket,
mock.patch("profiling.sampling.cli.select.select", return_value=([True], [], [])),
):
self._setup_sync_mocks(mock_socket, mock_popen)
from profiling.sampling.cli import main
Expand Down
Loading