fix: profiling or tracing multiprocessing can cause error

Signed-off-by: yihong0618 <zouzou0208@gmail.com>
Co-Authored-By: YvesDup <yduprat@gmail.com>

Author: yihong0618

Lib/profile.py

@@ -602,13 +602,16 @@ def main():
                 code = compile(fp.read(), progname, 'exec')
             spec = importlib.machinery.ModuleSpec(name='__main__', loader=None,
                                                   origin=progname)
-            globs = {
-                '__spec__': spec,
+            module = importlib.util.module_from_spec(spec)
+            sys.modules['__main__'] = module
+            globs = module.__dict__
+            globs.update({
+                '__spec__': None,
                 '__file__': spec.origin,
                 '__name__': spec.name,
                 '__package__': None,
                 '__cached__': None,
-            }
+            })
         try:
             runctx(code, globs, None, options.outfile, options.sort)
         except BrokenPipeError as exc:

Lib/profiling/tracing/__init__.py

@@ -197,7 +197,7 @@ def main():
             # in the module's namespace.
             globs = module.__dict__
             globs.update({
-                '__spec__': spec,
+                '__spec__': None,
                 '__file__': spec.origin,
                 '__name__': spec.name,
                 '__package__': None,

Lib/test/test_profile.py

@@ -4,10 +4,12 @@
 import pstats
 import unittest
 import os
+import subprocess
 import warnings
 from difflib import unified_diff
 from io import StringIO
 from test.support.os_helper import TESTFN, unlink, temp_dir, change_cwd
+from test.support import script_helper, os_helper, SHORT_TIMEOUT
 from contextlib import contextmanager, redirect_stdout
 
 # Suppress deprecation warning for profile module (PEP 799)
@@ -134,6 +136,39 @@ def test_output_file_when_changing_directory(self):
 
             self.assertTrue(os.path.exists('out.pstats'))
 
+    def test_profile_multiprocessing(self):
+        test_script = '''
+import multiprocessing
+def worker_proc(x):
+    return x * 42
+
+def main_proc():
+    p = multiprocessing.Process(target=worker_proc, args=(10,))
+    p.start()
+    p.join()
+
+if __name__ == "__main__":
+    main_proc()
+'''
+        with os_helper.temp_dir() as temp_dir:
+            script = script_helper.make_script(
+                temp_dir, 'test_profile_multiprocessing', test_script
+            )
+            with script_helper.spawn_python(
+                "-m", self.profilermodule.__name__,
+                script,
+                stderr=subprocess.PIPE,
+                text=True
+            ) as proc:
+                proc.wait(timeout=SHORT_TIMEOUT)
+                stdout = proc.stdout.read()
+                stderr = proc.stderr.read()
+
+        self.assertIn("main_proc", stdout)
+        self.assertNotIn("Can't pickle", stderr)
+        self.assertNotIn("ModuleSpec", stderr)
+        self.assertEqual(proc.returncode, 0)
+
 
 def regenerate_expected_output(filename, cls):
     filename = filename.rstrip('co')

Lib/test/test_profiling/test_tracing_profiler.py

@@ -2,12 +2,14 @@
 
 import sys
 import unittest
+import subprocess
 
 # rip off all interesting stuff from test_profile
 import profiling.tracing as cProfile
 import tempfile
 import textwrap
 from test.test_profile import ProfileTest, regenerate_expected_output
+from test.support import script_helper, os_helper, SHORT_TIMEOUT
 from test.support.script_helper import assert_python_failure, assert_python_ok
 from test import support
 
@@ -170,6 +172,39 @@ class Foo:
             f.close()
             assert_python_ok('-m', "cProfile", f.name)
 
+    def test_profile_multiprocessing(self):
+        test_script = '''
+import multiprocessing
+
+def worker_proc(x):
+    return x * 42
+
+def main_proc():
+    p = multiprocessing.Process(target=worker_proc, args=(10,))
+    p.start()
+    p.join()
+    print("SUCCESS")
+
+if __name__ == "__main__":
+    main_proc()
+'''
+        with os_helper.temp_dir() as temp_dir:
+            script = script_helper.make_script(
+                temp_dir, 'test_cprofile_multiprocessing', test_script
+            )
+            with script_helper.spawn_python(
+                "-m", "cProfile",
+                script,
+                stderr=subprocess.PIPE,
+                text=True
+            ) as proc:
+                proc.wait(timeout=SHORT_TIMEOUT)
+                stdout = proc.stdout.read()
+                stderr = proc.stderr.read()
+
+        self.assertIn("SUCCESS", stdout)
+        self.assertNotIn("has no attribute \'worker_proc\'", stderr)
+
 
 def main():
     if '-r' not in sys.argv: