Skip to content
11 changes: 9 additions & 2 deletions Lib/profiling/sampling/_sync_coordinator.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import socket
import runpy
import time
import types
from typing import List, NoReturn


Expand Down Expand Up @@ -175,15 +176,21 @@ def _execute_script(script_path: str, script_args: List[str], cwd: str) -> None:
try:
with open(script_path, 'rb') as f:
source_code = f.read()

except FileNotFoundError as e:
raise TargetError(f"Script file not found: {script_path}") from e
except PermissionError as e:
raise TargetError(f"Permission denied reading script: {script_path}") from e

try:
# Compile and execute the script
# gh-140729: Create a __mp_main__ module to allow pickling
main_module = types.ModuleType("__main__")
main_module.__file__ = script_path
main_module.__builtins__ = __builtins__
sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module

code = compile(source_code, script_path, 'exec')
exec(code, {'__name__': '__main__', '__file__': script_path})
exec(code, main_module.__dict__)
except SyntaxError as e:
raise TargetError(f"Syntax error in script {script_path}: {e}") from e
except SystemExit:
Expand Down
52 changes: 51 additions & 1 deletion Lib/test/test_profiling/test_sampling_profiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,13 @@
from profiling.sampling.gecko_collector import GeckoCollector

from test.support.os_helper import unlink
from test.support import force_not_colorized_test_class, SHORT_TIMEOUT
from test.support import (
force_not_colorized_test_class,
SHORT_TIMEOUT,
script_helper,
os_helper,
SuppressCrashReport,
)
from test.support.socket_helper import find_unused_port
from test.support import requires_subprocess, is_emscripten
from test.support import captured_stdout, captured_stderr
Expand Down Expand Up @@ -3007,5 +3013,49 @@ def test_parse_mode_function(self):
profiling.sampling.sample._parse_mode("invalid")


@requires_subprocess()
@skip_if_not_supported
class TestProcessPoolExecutorSupport(unittest.TestCase):
"""
Test that ProcessPoolExecutor works correctly with profiling.sampling.
"""

def test_process_pool_executor_pickle(self):
# gh-140729: test use ProcessPoolExecutor.map() can sampling
test_script = '''
import concurrent.futures

def worker(x):
return x * 2

if __name__ == "__main__":
with concurrent.futures.ProcessPoolExecutor() as executor:
results = list(executor.map(worker, [1, 2, 3]))
print(f"Results: {results}")
'''
with os_helper.temp_dir() as temp_dir:
script = script_helper.make_script(
temp_dir, 'test_process_pool_executor_pickle', test_script
)
with SuppressCrashReport():
with script_helper.spawn_python(
"-m", "profiling.sampling.sample",
"-d", "5",
"-i", "100000",
script,
stderr=subprocess.PIPE,
text=True
) as proc:
proc.wait(timeout=SHORT_TIMEOUT)
stdout = proc.stdout.read()
stderr = proc.stderr.read()

if "PermissionError" in stderr:
self.skipTest("Insufficient permissions for remote profiling")

self.assertIn("Results: [2, 4, 6]", stdout)
self.assertNotIn("Can't pickle", stderr)


if __name__ == "__main__":
unittest.main()
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Fix pickling error in the sampling profiler when using ``concurrent.futures.ProcessPoolExecutor``
script can not be properly pickled and executed in worker processes.
Loading