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
45 changes: 44 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,42 @@ 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(script, stderr=subprocess.PIPE) as proc:
proc.wait()
stdout = proc.stdout.read()
stderr = proc.stderr.read()

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

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


if __name__ == "__main__":
unittest.main()
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Fix: Add __mp_main__ as a duplicate for __main__ for pickle to work in
Copy link
Member

Choose a reason for hiding this comment

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

This needs to refer the profiler otherwise is not possible to undestand what this is fixing. Also this should say WHAT is fixed not HOW

Copy link
Contributor Author

Choose a reason for hiding this comment

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

addressed

sampling
Loading