mirror of
https://github.com/python/cpython.git
synced 2025-12-08 06:10:17 +00:00
gh-140729: Add __mp_main__ as a duplicate for __main__ for pickle to work (#140735)
This commit is contained in:
parent
20b64bdf23
commit
994ab5c922
3 changed files with 62 additions and 3 deletions
|
|
@ -10,6 +10,7 @@
|
||||||
import socket
|
import socket
|
||||||
import runpy
|
import runpy
|
||||||
import time
|
import time
|
||||||
|
import types
|
||||||
from typing import List, NoReturn
|
from typing import List, NoReturn
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -175,15 +176,21 @@ def _execute_script(script_path: str, script_args: List[str], cwd: str) -> None:
|
||||||
try:
|
try:
|
||||||
with open(script_path, 'rb') as f:
|
with open(script_path, 'rb') as f:
|
||||||
source_code = f.read()
|
source_code = f.read()
|
||||||
|
|
||||||
except FileNotFoundError as e:
|
except FileNotFoundError as e:
|
||||||
raise TargetError(f"Script file not found: {script_path}") from e
|
raise TargetError(f"Script file not found: {script_path}") from e
|
||||||
except PermissionError as e:
|
except PermissionError as e:
|
||||||
raise TargetError(f"Permission denied reading script: {script_path}") from e
|
raise TargetError(f"Permission denied reading script: {script_path}") from e
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# Compile and execute the script
|
main_module = types.ModuleType("__main__")
|
||||||
|
main_module.__file__ = script_path
|
||||||
|
main_module.__builtins__ = __builtins__
|
||||||
|
# gh-140729: Create a __mp_main__ module to allow pickling
|
||||||
|
sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module
|
||||||
|
|
||||||
code = compile(source_code, script_path, 'exec', module='__main__')
|
code = compile(source_code, script_path, 'exec', module='__main__')
|
||||||
exec(code, {'__name__': '__main__', '__file__': script_path})
|
exec(code, main_module.__dict__)
|
||||||
except SyntaxError as e:
|
except SyntaxError as e:
|
||||||
raise TargetError(f"Syntax error in script {script_path}: {e}") from e
|
raise TargetError(f"Syntax error in script {script_path}: {e}") from e
|
||||||
except SystemExit:
|
except SystemExit:
|
||||||
|
|
|
||||||
|
|
@ -22,7 +22,13 @@
|
||||||
from profiling.sampling.gecko_collector import GeckoCollector
|
from profiling.sampling.gecko_collector import GeckoCollector
|
||||||
|
|
||||||
from test.support.os_helper import unlink
|
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.socket_helper import find_unused_port
|
||||||
from test.support import requires_subprocess, is_emscripten
|
from test.support import requires_subprocess, is_emscripten
|
||||||
from test.support import captured_stdout, captured_stderr
|
from test.support import captured_stdout, captured_stderr
|
||||||
|
|
@ -3009,5 +3015,49 @@ def test_parse_mode_function(self):
|
||||||
profiling.sampling.sample._parse_mode("invalid")
|
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__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|
|
||||||
|
|
@ -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…
Add table
Add a link
Reference in a new issue