Skip to content

Commit 994ab5c

Browse files
authored
pythongh-140729: Add __mp_main__ as a duplicate for __main__ for pickle to work (python#140735)
1 parent 20b64bd commit 994ab5c

File tree

3 files changed

+62
-3
lines changed

3 files changed

+62
-3
lines changed

Lib/profiling/sampling/_sync_coordinator.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import socket
1111
import runpy
1212
import time
13+
import types
1314
from typing import List, NoReturn
1415

1516

@@ -175,15 +176,21 @@ def _execute_script(script_path: str, script_args: List[str], cwd: str) -> None:
175176
try:
176177
with open(script_path, 'rb') as f:
177178
source_code = f.read()
179+
178180
except FileNotFoundError as e:
179181
raise TargetError(f"Script file not found: {script_path}") from e
180182
except PermissionError as e:
181183
raise TargetError(f"Permission denied reading script: {script_path}") from e
182184

183185
try:
184-
# Compile and execute the script
186+
main_module = types.ModuleType("__main__")
187+
main_module.__file__ = script_path
188+
main_module.__builtins__ = __builtins__
189+
# gh-140729: Create a __mp_main__ module to allow pickling
190+
sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module
191+
185192
code = compile(source_code, script_path, 'exec', module='__main__')
186-
exec(code, {'__name__': '__main__', '__file__': script_path})
193+
exec(code, main_module.__dict__)
187194
except SyntaxError as e:
188195
raise TargetError(f"Syntax error in script {script_path}: {e}") from e
189196
except SystemExit:

Lib/test/test_profiling/test_sampling_profiler.py

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,13 @@
2222
from profiling.sampling.gecko_collector import GeckoCollector
2323

2424
from test.support.os_helper import unlink
25-
from test.support import force_not_colorized_test_class, SHORT_TIMEOUT
25+
from test.support import (
26+
force_not_colorized_test_class,
27+
SHORT_TIMEOUT,
28+
script_helper,
29+
os_helper,
30+
SuppressCrashReport,
31+
)
2632
from test.support.socket_helper import find_unused_port
2733
from test.support import requires_subprocess, is_emscripten
2834
from test.support import captured_stdout, captured_stderr
@@ -3009,5 +3015,49 @@ def test_parse_mode_function(self):
30093015
profiling.sampling.sample._parse_mode("invalid")
30103016

30113017

3018+
@requires_subprocess()
3019+
@skip_if_not_supported
3020+
class TestProcessPoolExecutorSupport(unittest.TestCase):
3021+
"""
3022+
Test that ProcessPoolExecutor works correctly with profiling.sampling.
3023+
"""
3024+
3025+
def test_process_pool_executor_pickle(self):
3026+
# gh-140729: test use ProcessPoolExecutor.map() can sampling
3027+
test_script = '''
3028+
import concurrent.futures
3029+
3030+
def worker(x):
3031+
return x * 2
3032+
3033+
if __name__ == "__main__":
3034+
with concurrent.futures.ProcessPoolExecutor() as executor:
3035+
results = list(executor.map(worker, [1, 2, 3]))
3036+
print(f"Results: {results}")
3037+
'''
3038+
with os_helper.temp_dir() as temp_dir:
3039+
script = script_helper.make_script(
3040+
temp_dir, 'test_process_pool_executor_pickle', test_script
3041+
)
3042+
with SuppressCrashReport():
3043+
with script_helper.spawn_python(
3044+
"-m", "profiling.sampling.sample",
3045+
"-d", "5",
3046+
"-i", "100000",
3047+
script,
3048+
stderr=subprocess.PIPE,
3049+
text=True
3050+
) as proc:
3051+
proc.wait(timeout=SHORT_TIMEOUT)
3052+
stdout = proc.stdout.read()
3053+
stderr = proc.stderr.read()
3054+
3055+
if "PermissionError" in stderr:
3056+
self.skipTest("Insufficient permissions for remote profiling")
3057+
3058+
self.assertIn("Results: [2, 4, 6]", stdout)
3059+
self.assertNotIn("Can't pickle", stderr)
3060+
3061+
30123062
if __name__ == "__main__":
30133063
unittest.main()
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix pickling error in the sampling profiler when using ``concurrent.futures.ProcessPoolExecutor``
2+
script can not be properly pickled and executed in worker processes.

0 commit comments

Comments
 (0)