From 5d023728cd3d308b1751430d4bc6c37e63eb35b3 Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Mon, 26 Oct 2026 18:40:24 +0000 Subject: [PATCH] Support closures and clean hard-deletion in MultiProcessShared spawn_process (#40673) --- .../apache_beam/utils/multi_process_shared.py | 41 +++- .../utils/multi_process_shared_test.py | 183 ++++++++++++++++++ 2 files changed, 221 insertions(+), 3 deletions(-) diff --git a/sdks/python/apache_beam/utils/multi_process_shared.py b/sdks/python/apache_beam/utils/multi_process_shared.py index 13a0f13e617b..8a1bce9519ef 100644 --- a/sdks/python/apache_beam/utils/multi_process_shared.py +++ b/sdks/python/apache_beam/utils/multi_process_shared.py @@ -26,6 +26,7 @@ import logging import multiprocessing.managers import os +import sys import tempfile import threading import time @@ -38,6 +39,7 @@ import fasteners +from apache_beam.internal import pickler from apache_beam.utils import retry # In some python versions, there is a bug where AutoProxy doesn't handle @@ -150,8 +152,22 @@ def release(self, proxy): def unsafe_hard_delete(self): with self.lock: if self.initialied: + obj = self.obj del self.obj self.initialied = False + # When hard-deleting in a spawned server process that exits via + # os._exit(0), an in-flight RPC thread (e.g. a timed-out call) may + # still hold a stack reference to `obj`, preventing CPython from + # invoking `obj.__del__()` before `os._exit(0)`. + if (self._hard_delete_callback is not None and + sys.getrefcount(obj) > 2 and hasattr(type(obj), '__del__')): + try: + type(obj).__del__(obj) + except Exception: + logging.warning( + 'Exception while finalizing object during unsafe_hard_delete', + exc_info=True) + del obj if self._hard_delete_callback: self._hard_delete_callback() @@ -232,7 +248,8 @@ def unsafe_hard_delete(self): self._proxyObject.unsafe_hard_delete() -def _run_server_process(address_file, tag, constructor, authkey, life_line): +def _run_server_process( + address_file, tag, serialized_constructor, authkey, life_line): """ Runs in a separate process. Includes a 'Suicide Pact' monitor: If parent dies, I die. @@ -276,6 +293,7 @@ def _monitor_parent(): logging.getLogger().setLevel(logging.INFO) multiprocessing.current_process().authkey = authkey + constructor = pickler.loads(serialized_constructor) serving_manager = _SingletonRegistrar( address=('localhost', 0), authkey=authkey) _process_level_singleton_manager.set_hard_delete_callback( @@ -439,7 +457,19 @@ def unsafe_hard_delete(self): to this object exist, or (b) you are ok with all existing references to this object throwing strange errors when derefrenced. """ - self._get_manager().unsafe_hard_delete_singleton(self._tag) + address_file = os.path.join(self._path, self._tag) + ".address" + if (self._manager is None and + not _process_level_singleton_manager.has_singleton(self._tag) and + not os.path.exists(address_file)): + return + try: + self._get_manager().unsafe_hard_delete_singleton(self._tag) + except (EOFError, OSError, BrokenPipeError): + # Expected when a spawned server process exits via os._exit(0) before + # sending an RPC reply. + pass + finally: + self._manager = None def _create_server(self, address_file): if self._spawn_process: @@ -459,7 +489,12 @@ def _create_server(self, address_file): ctx = multiprocessing.get_context('spawn') p = ctx.Process( target=_run_server_process, - args=(address_file, self._tag, self._constructor, AUTH_KEY, reader), + args=( + address_file, + self._tag, + pickler.dumps(self._constructor), + AUTH_KEY, + reader), daemon=False # Must be False for nested proxies ) p.start() diff --git a/sdks/python/apache_beam/utils/multi_process_shared_test.py b/sdks/python/apache_beam/utils/multi_process_shared_test.py index c597e459e1a1..3256ec00ab79 100644 --- a/sdks/python/apache_beam/utils/multi_process_shared_test.py +++ b/sdks/python/apache_beam/utils/multi_process_shared_test.py @@ -16,11 +16,14 @@ # # pytype: skip-file +import gc import logging import multiprocessing +import multiprocessing.connection import os import tempfile import threading +import time import unittest from typing import Any from unittest.mock import patch @@ -106,6 +109,39 @@ def __getattribute__(self, __name: str) -> Any: return object.__getattribute__(self, __name) +class PidCounter(object): + def __init__(self, start=0): + self.running = start + self.lock = threading.Lock() + + def get_pid(self): + return os.getpid() + + def get(self): + return self.running + + def increment(self, value=1): + with self.lock: + self.running += value + return self.running + + +class SlowFinalizingObject(object): + def __init__(self, entered_path, finalized_path): + self._entered_path = entered_path + self._finalized_path = finalized_path + + def block_call(self): + with open(self._entered_path, 'w') as f: + f.write('entered') + time.sleep(30) + return 'done' + + def __del__(self): + with open(self._finalized_path, 'w') as f: + f.write('finalized') + + class MultiProcessSharedTest(unittest.TestCase): @classmethod def setUpClass(cls): @@ -492,6 +528,153 @@ def side_effect_connect(self_mgr, *args, **kwargs): self.assertEqual(counter1.increment(), 1) self.assertEqual(counter2.increment(), 2) + def test_spawn_process_runs_in_dedicated_pid(self): + parent_pid = os.getpid() + proxies = [ + multi_process_shared.MultiProcessShared( + PidCounter, + tag=f'dedicated_pid_{i}', + always_proxy=True, + spawn_process=True).acquire() for i in range(3) + ] + pids = [proxy.get_pid() for proxy in proxies] + self.assertEqual(len(set(pids)), 3) + for pid in pids: + self.assertNotEqual(pid, parent_pid) + + same_tag_proxy = multi_process_shared.MultiProcessShared( + PidCounter, + tag='dedicated_pid_0', + always_proxy=True, + spawn_process=True).acquire() + self.assertEqual(same_tag_proxy.get_pid(), pids[0]) + self.assertEqual(proxies[0].increment(5), 5) + self.assertEqual(same_tag_proxy.get(), 5) + + def test_no_cyclic_gc_self_deadlock_issue_40673(self): + """Regression test for https://github.com/apache/beam/issues/40673. + + When always_proxy=True and spawn_process=False, cyclic GC collecting an + unreachable BaseProxy on the in-process Server.accepter thread deadlocks + the server against itself. With spawn_process=True, the server runs in a + dedicated child process with no BaseProxy references, preventing deadlock. + """ + orig_accept = multiprocessing.connection.Listener.accept + + def gc_on_accepter_accept(listener_self): + if 'accepter' in threading.current_thread().name: + gc.collect() + return orig_accept(listener_self) + + stop_bg_gc = threading.Event() + + def background_gc(): + while not stop_bg_gc.is_set(): + gc.collect() + time.sleep(0.005) + + errors = [] + + def workload(): + try: + for i in range(15): + handle = multi_process_shared.MultiProcessShared( + Counter, + tag='issue_40673_tag', + always_proxy=True, + spawn_process=True) + proxy = handle.acquire() + self.assertEqual(proxy.increment(), i + 1) + cycle = [] + cycle.append(cycle) + cycle.append(proxy) + cycle.append(handle) + del proxy + del handle + del cycle + if i % 3 == 0: + gc.collect() + except BaseException as e: # pylint: disable=broad-exception-caught + errors.append(e) + + with patch.object(multiprocessing.connection.Listener, + 'accept', + gc_on_accepter_accept): + gc_thread = threading.Thread(target=background_gc, daemon=True) + worker_thread = threading.Thread(target=workload, daemon=True) + gc_thread.start() + worker_thread.start() + worker_thread.join(timeout=30) + stop_bg_gc.set() + gc_thread.join(timeout=5) + + self.assertFalse( + worker_thread.is_alive(), + 'Deadlock detected while acquiring/collecting cyclic BaseProxy ' + 'instances with spawn_process=True') + self.assertEqual(errors, []) + + final_proxy = multi_process_shared.MultiProcessShared( + Counter, tag='issue_40673_tag', always_proxy=True, + spawn_process=True).acquire() + self.assertEqual(final_proxy.get(), 15) + + def test_spawn_process_supports_lambda_and_closure_constructors(self): + captured_offset = 42 + + def local_factory(extra): + return PidCounter(start=captured_offset + extra) + + proxy = multi_process_shared.MultiProcessShared( + lambda: local_factory(8), + tag='closure_constructor', + always_proxy=True, + spawn_process=True).acquire() + self.assertNotEqual(proxy.get_pid(), os.getpid()) + self.assertEqual(proxy.get(), 50) + self.assertEqual(proxy.increment(5), 55) + + def test_unsafe_hard_delete_missing_tag_is_noop(self): + handle = multi_process_shared.MultiProcessShared( + Counter, tag='nonexistent_tag', always_proxy=True, spawn_process=True) + handle.unsafe_hard_delete() + self.assertEqual(multiprocessing.active_children(), []) + + def test_unsafe_hard_delete_runs_del_during_inflight_blocked_call(self): + with tempfile.TemporaryDirectory() as tmp_dir: + entered_path = os.path.join(tmp_dir, 'entered') + finalized_path = os.path.join(tmp_dir, 'finalized') + handle = multi_process_shared.MultiProcessShared( + lambda: SlowFinalizingObject(entered_path, finalized_path), + tag='inflight_hard_delete', + always_proxy=True, + spawn_process=True) + proxy = handle.acquire() + server_proc = next( + p for p in multiprocessing.active_children() if p.is_alive()) + + def call_ignoring_disconnect(): + try: + proxy.block_call() + except (EOFError, OSError): + pass + + call_thread = threading.Thread( + target=call_ignoring_disconnect, daemon=True) + call_thread.start() + + deadline = time.time() + 5 + while not os.path.exists(entered_path) and time.time() < deadline: + time.sleep(0.02) + self.assertTrue(os.path.exists(entered_path)) + + handle.unsafe_hard_delete() + server_proc.join(timeout=5) + self.assertFalse(server_proc.is_alive()) + self.assertTrue(os.path.exists(finalized_path)) + with open(finalized_path) as f: + self.assertEqual(f.read(), 'finalized') + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO)