diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 5ee65afa0615..8bc04e3cea4a 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -2073,7 +2073,10 @@ def load_model_status( tag = f'{model_tag}_model_status' if share_across_processes: return multi_process_shared.MultiProcessShared( - lambda: _ModelStatus(True), tag=tag, always_proxy=True).acquire() + lambda: _ModelStatus(True), + tag=tag, + always_proxy=True, + spawn_process=True).acquire() return _model_statuses.setdefault(tag, _ModelStatus(False)) @@ -2135,9 +2138,10 @@ def __init__( self.loader_func = loader_func if not self.use_model_manager and len(models) > 1: self.model_router = multi_process_shared.MultiProcessShared( - lambda: _ModelRoutingStrategy(), + _ModelRoutingStrategy, tag=f'{model_tag}_counter', - always_proxy=True).acquire() + always_proxy=True, + spawn_process=True).acquire() def next_model(self): if self.use_model_manager: @@ -2250,19 +2254,24 @@ def load(): model_manager = multi_process_shared.MultiProcessShared( lambda: ModelManager(**self._model_manager_args), tag='model_manager', - always_proxy=True).acquire() + always_proxy=False).acquire() model_wrapper = _SharedModelWrapper( model_manager, self._cur_tag, self._model_handler.load_model, metric_tag=self._model_tag) elif self._model_handler.share_model_across_processes(): + if isinstance(side_input_model_path, str): + self._model_handler.update_model_path(side_input_model_path) models = [] for copy_tag in _get_tags_for_copies(self._cur_tag, self._model_handler.model_copies()): models.append( multi_process_shared.MultiProcessShared( - load, tag=copy_tag, always_proxy=True).acquire()) + OOMProtectedFn(self._model_handler.load_model), + tag=copy_tag, + always_proxy=True, + spawn_process=True).acquire()) model_wrapper = _SharedModelWrapper(models, self._cur_tag) else: model = self._shared_model_handle.acquire(load, tag=self._cur_tag) @@ -2331,7 +2340,8 @@ def _run_inference(self, batch, inference_args): # ModelManager with MultiProcessShared returns the model tag unique_tag = model model = multi_process_shared.MultiProcessShared( - lambda: None, tag=model, always_proxy=True).acquire() + lambda: None, tag=model, always_proxy=True, + spawn_process=True).acquire() try: result_generator = (OOMProtectedFn(self._model_handler.run_inference))( batch, model, inference_args) @@ -2410,8 +2420,8 @@ def finish_bundle(self): for unprefixed_tag in tags_to_gc: for tag in _get_tags_for_copies(unprefixed_tag, self._model_handler.model_copies()): - multi_process_shared.MultiProcessShared(lambda: None, - tag).unsafe_hard_delete() + multi_process_shared.MultiProcessShared( + lambda: None, tag, spawn_process=True).unsafe_hard_delete() self._model_metadata.mark_tags_deleted(tags_to_gc) diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index dec1c2906cd2..d2c6a4f27441 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -20,6 +20,7 @@ import inspect import math import multiprocessing +import multiprocessing.connection import os import pickle import random @@ -3153,5 +3154,234 @@ def test_run_inference_dofn_model_update_acquires_and_releases_matching_tag( dofn._model.force_reset() +class PidReportingModel: + """Model that reports the OS PIDs where it was loaded and invoked.""" + def __init__(self, model_path: str, env_var_name: Optional[str] = None): + self._model_path = model_path + self._env_var_name = env_var_name + self._load_pid = os.getpid() + self._load_env = os.environ.get(env_var_name) if env_var_name else None + + def predict(self, example: int) -> dict[str, Any]: + return { + 'example': example, + 'load_pid': self._load_pid, + 'infer_pid': os.getpid(), + 'model_path': self._model_path, + 'load_env': self._load_env, + 'infer_env': ( + os.environ.get(self._env_var_name) if self._env_var_name else None), + } + + +class PidReportingModelHandler(base.ModelHandler[int, + dict[str, Any], + PidReportingModel]): + def __init__( + self, + *, + model_path: str = 'default_path', + model_copies: int = 1, + env_vars: Optional[dict[str, Any]] = None, + env_var_name: Optional[str] = None, + inference_fn=None): + super().__init__(env_vars=env_vars or {}) + self._model_path = model_path + self._model_copies = model_copies + self._env_var_name = env_var_name + self._inference_fn = inference_fn + + def load_model(self) -> PidReportingModel: + return PidReportingModel( + model_path=self._model_path, env_var_name=self._env_var_name) + + def update_model_path(self, model_path: Optional[str] = None): + if model_path: + self._model_path = model_path + + def run_inference( + self, + batch: Sequence[int], + model: PidReportingModel, + inference_args=None) -> Iterable[dict[str, Any]]: + worker_pid = os.getpid() + for example in batch: + res = model.predict(example) + res['worker_pid'] = worker_pid + if self._inference_fn is not None: + res['fn_out'] = self._inference_fn(example) + yield res + + def share_model_across_processes(self) -> bool: + return True + + def model_copies(self) -> int: + return self._model_copies + + def batch_elements_kwargs(self): + return {'min_batch_size': 1, 'max_batch_size': 1} + + +class SpawnedModelProcessTest(unittest.TestCase): + """Tests verifying shared models spawn in dedicated processes (#40673).""" + def tearDown(self): + for p in multiprocessing.active_children(): + if p.is_alive(): + p.terminate() + p.join(timeout=2) + + def test_shared_model_spawns_in_dedicated_process(self): + with tempfile.TemporaryDirectory() as tmp_dir: + results_path = os.path.join(tmp_dir, 'results.csv') + mh = PidReportingModelHandler(model_path='single_v1') + with TestPipeline('FnApiRunner') as pipeline: + _ = ( + pipeline + | 'Create' >> beam.Create([10, 20, 30]) + | 'Infer' >> base.RunInference( + mh, model_identifier='dedicated_single_' + uuid.uuid4().hex) + | 'Write' >> beam.Map( + lambda r: open(results_path, 'a').write( + f'{r["worker_pid"]},{r["load_pid"]},{r["infer_pid"]},{r["model_path"]}\n' + ) # pylint: disable=line-too-long + )) + + with open(results_path) as f: + rows = [line.strip().split(',') for line in f if line.strip()] + self.assertEqual(len(rows), 3) + load_pids = {int(r[1]) for r in rows} + self.assertEqual(len(load_pids), 1) + for worker_pid_s, load_pid_s, infer_pid_s, model_path in rows: + self.assertEqual(int(load_pid_s), int(infer_pid_s)) + self.assertNotEqual(int(load_pid_s), int(worker_pid_s)) + self.assertEqual(model_path, 'single_v1') + + def test_shared_model_multiple_copies_spawn_distinct_processes(self): + with tempfile.TemporaryDirectory() as tmp_dir: + results_path = os.path.join(tmp_dir, 'results.csv') + mh = PidReportingModelHandler(model_path='multi_copy', model_copies=3) + with TestPipeline('FnApiRunner') as pipeline: + _ = ( + pipeline + | 'Create' >> beam.Create([1, 2, 3, 4, 5, 6], reshuffle=False) + | 'Infer' >> base.RunInference( + mh, model_identifier='dedicated_copies_' + uuid.uuid4().hex) + | 'Write' >> beam.Map( + lambda r: open(results_path, 'a').write( + f'{r["worker_pid"]},{r["load_pid"]},{r["infer_pid"]}\n') # pylint: disable=line-too-long + )) + + with open(results_path) as f: + rows = [line.strip().split(',') for line in f if line.strip()] + self.assertEqual(len(rows), 6) + worker_pids = {int(r[0]) for r in rows} + model_pids = {int(r[1]) for r in rows} + self.assertEqual( + len(model_pids), 3, f'Expected 3 distinct model PIDs: {rows}') + self.assertTrue(model_pids.isdisjoint(worker_pids)) + for _, load_pid_s, infer_pid_s in rows: + self.assertEqual(int(load_pid_s), int(infer_pid_s)) + + def test_shared_model_env_vars_propagated_to_spawned_process(self): + env_key = 'BEAM_TEST_SPAWNED_MODEL_ENV_' + uuid.uuid4().hex[:8] + try: + mh = PidReportingModelHandler( + model_path='env_test', + env_vars={env_key: 'expected_env_value'}, + env_var_name=env_key) + dofn = base._RunInferenceDoFn( + mh, + FakeClock(), + metrics_namespace=None, + model_tag='env_tag_' + uuid.uuid4().hex) + dofn.setup() + preds = dofn.process([1], None, None) + self.assertEqual(len(preds), 1) + self.assertNotEqual(preds[0]['load_pid'], os.getpid()) + self.assertEqual(preds[0]['load_env'], 'expected_env_value') + self.assertEqual(preds[0]['infer_env'], 'expected_env_value') + finally: + os.environ.pop(env_key, None) + + def test_shared_model_local_closure_and_side_input_updates_no_gc_deadlock( + self): + """End-to-end regression test for https://github.com/apache/beam/issues/40673.""" + factor = 10 + + def local_closure_fn(x: int) -> int: + return x * factor + + 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 = [] + seen_pids = [] + + def run_updates(): + try: + mh = PidReportingModelHandler( + model_path='init', inference_fn=local_closure_fn) + dofn = base._RunInferenceDoFn( + mh, + FakeClock(), + metrics_namespace='gc_test', + load_model_at_runtime=True, + model_tag='gc_dofn_' + uuid.uuid4().hex) + dofn.setup() + + for i in range(5): + model_id = f'side_model_{i}_{uuid.uuid4().hex[:6]}' + si = base.ModelMetadata(model_id=model_id, model_name=f'm{i}') + out = dofn.process([i + 1], None, si) + self.assertEqual(out[0]['model_path'], model_id) + self.assertEqual(out[0]['fn_out'], (i + 1) * 10) + self.assertNotEqual(out[0]['load_pid'], os.getpid()) + seen_pids.append(out[0]['load_pid']) + + cycle = [] + cycle.append(cycle) + cycle.append(dofn._model) + del cycle + gc.collect() + + last_pid = seen_pids[-1] + last_proc = next( + p for p in multiprocessing.active_children() if p.pid == last_pid) + dofn._model_metadata.try_mark_current_model_invalid(0) + dofn.finish_bundle() + last_proc.join(timeout=5) + self.assertFalse(last_proc.is_alive()) + except BaseException as e: # pylint: disable=broad-exception-caught + errors.append(e) + + with unittest.mock.patch.object(multiprocessing.connection.Listener, + 'accept', + gc_on_accepter_accept): + gc_thread = threading.Thread(target=background_gc, daemon=True) + worker_thread = threading.Thread(target=run_updates, daemon=True) + gc_thread.start() + worker_thread.start() + worker_thread.join(timeout=45) + stop_bg_gc.set() + gc_thread.join(timeout=5) + + self.assertFalse( + worker_thread.is_alive(), + 'Deadlock detected during RunInference side-input model updates') + self.assertEqual(errors, []) + self.assertEqual(len(set(seen_pids)), 5) + + if __name__ == '__main__': unittest.main() 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)