From d308262fedffa65f53ea4d2ae1268cc2d17e8de0 Mon Sep 17 00:00:00 2001 From: Valentyn Tymofieiev Date: Thu, 2 Oct 2025 17:57:39 -0700 Subject: [PATCH 1/4] Mark instructions as cleaned up in the GRPC data channel if processing an instruction fails. --- .../python/apache_beam/runners/worker/data_plane.py | 13 ++++++++++++- .../python/apache_beam/runners/worker/sdk_worker.py | 1 + 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/runners/worker/data_plane.py b/sdks/python/apache_beam/runners/worker/data_plane.py index e4cf4f185ad4..e524cba62254 100644 --- a/sdks/python/apache_beam/runners/worker/data_plane.py +++ b/sdks/python/apache_beam/runners/worker/data_plane.py @@ -502,7 +502,7 @@ def _clean_receiving_queue(self, instruction_id): instruction_id cannot be reused for new queue. """ with self._receive_lock: - self._received.pop(instruction_id) + self._received.pop(instruction_id, None) self._cleaned_instruction_ids[instruction_id] = True while len(self._cleaned_instruction_ids) > _MAX_CLEANED_INSTRUCTIONS: self._cleaned_instruction_ids.popitem(last=False) @@ -787,6 +787,12 @@ def close(self): """Close all channels that this factory owns.""" raise NotImplementedError(type(self)) + def cleanup(self, instruction_id): + # type: (str) -> None + + """Clean up resources for a given instruction.""" + pass + class GrpcClientDataChannelFactory(DataChannelFactory): """A factory for ``GrpcClientDataChannel``. @@ -855,6 +861,11 @@ def close(self): channel.close() self._data_channel_cache.clear() + def cleanup(self, instruction_id): + # type: (str) -> None + for channel in self._data_channel_cache.values(): + channel._clean_receiving_queue(instruction_id) + class InMemoryDataChannelFactory(DataChannelFactory): """A singleton factory for ``InMemoryDataChannel``.""" diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker.py b/sdks/python/apache_beam/runners/worker/sdk_worker.py index 0b4c236d6b37..d3810800c644 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker.py @@ -568,6 +568,7 @@ def discard(self, instruction_id, exception): # Perform the shutdown while not holding the lock. processor.shutdown() + self.data_channel_factory.cleanup(instruction_id) def release(self, instruction_id): # type: (str) -> None From 265ce8e90860df3b7ca4ed044bf9f0e14e915b57 Mon Sep 17 00:00:00 2001 From: Valentyn Tymofieiev Date: Thu, 2 Oct 2025 18:47:27 -0700 Subject: [PATCH 2/4] Invoke cleanup even if BP failed to create. --- sdks/python/apache_beam/runners/worker/sdk_worker.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker.py b/sdks/python/apache_beam/runners/worker/sdk_worker.py index d3810800c644..b0d47f9607d7 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker.py @@ -559,15 +559,17 @@ def discard(self, instruction_id, exception): """ Marks the instruction id as failed shutting down the ``BundleProcessor``. """ + processor = None with self._lock: self.failed_instruction_ids[instruction_id] = exception while len(self.failed_instruction_ids) > MAX_FAILED_INSTRUCTIONS: self.failed_instruction_ids.popitem(last=False) - processor = self.active_bundle_processors[instruction_id][1] - del self.active_bundle_processors[instruction_id] + if instruction_id in self.active_bundle_processors: + processor = self.active_bundle_processors.pop(instruction_id)[1] # Perform the shutdown while not holding the lock. - processor.shutdown() + if processor: + processor.shutdown() self.data_channel_factory.cleanup(instruction_id) def release(self, instruction_id): @@ -691,9 +693,9 @@ def process_bundle( instruction_id # type: str ): # type: (...) -> beam_fn_api_pb2.InstructionResponse - bundle_processor = self.bundle_processor_cache.get( - instruction_id, request.process_bundle_descriptor_id) try: + bundle_processor = self.bundle_processor_cache.get( + instruction_id, request.process_bundle_descriptor_id) with bundle_processor.state_handler.process_instruction_id( instruction_id, request.cache_tokens): with self.maybe_profile(instruction_id): From 0597b73941ac6e5203227612f5d7afdf25d5b58d Mon Sep 17 00:00:00 2001 From: Valentyn Tymofieiev Date: Tue, 14 Oct 2025 18:45:37 -0700 Subject: [PATCH 3/4] Address feedback. --- sdks/python/apache_beam/runners/worker/data_plane.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/sdks/python/apache_beam/runners/worker/data_plane.py b/sdks/python/apache_beam/runners/worker/data_plane.py index e524cba62254..cbd28f8b0a3f 100644 --- a/sdks/python/apache_beam/runners/worker/data_plane.py +++ b/sdks/python/apache_beam/runners/worker/data_plane.py @@ -502,6 +502,10 @@ def _clean_receiving_queue(self, instruction_id): instruction_id cannot be reused for new queue. """ with self._receive_lock: + # Per-instruction read queue may or may not be created yet when + # we mark an instruction as 'cleaned up' when creating + # a bundle processor failed, e.g. due to a flake in DoFn.setup(). + # We want to mark an instruction as cleaned up regardless. self._received.pop(instruction_id, None) self._cleaned_instruction_ids[instruction_id] = True while len(self._cleaned_instruction_ids) > _MAX_CLEANED_INSTRUCTIONS: @@ -857,13 +861,13 @@ def create_data_channel(self, remote_grpc_port): def close(self): # type: () -> None _LOGGER.info('Closing all cached grpc data channels.') - for _, channel in self._data_channel_cache.items(): + for channel in list(self._data_channel_cache.values()): channel.close() self._data_channel_cache.clear() def cleanup(self, instruction_id): # type: (str) -> None - for channel in self._data_channel_cache.values(): + for channel in list(self._data_channel_cache.values()): channel._clean_receiving_queue(instruction_id) From bf1625016ea36432834d2891a92c6a08f750cb49 Mon Sep 17 00:00:00 2001 From: Valentyn Tymofieiev Date: Wed, 15 Oct 2025 13:29:48 -0700 Subject: [PATCH 4/4] Add a test --- .../runners/worker/sdk_worker_test.py | 44 +++++++++++++++++-- 1 file changed, 40 insertions(+), 4 deletions(-) diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py index 0ab04ff256cd..7b53f274cac2 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py @@ -37,6 +37,7 @@ from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.portability.api import metrics_pb2 +from apache_beam.runners.worker import data_plane from apache_beam.runners.worker import sdk_worker from apache_beam.runners.worker import statecache from apache_beam.runners.worker.sdk_worker import BundleProcessorCache @@ -126,7 +127,10 @@ def test_fn_registration(self): def test_inactive_bundle_processor_returns_empty_progress_response(self): bundle_processor = mock.MagicMock() - bundle_processor_cache = BundleProcessorCache(None, None, None, {}) + data_channel_factory = mock.create_autospec( + data_plane.GrpcClientDataChannelFactory) + bundle_processor_cache = BundleProcessorCache( + None, None, data_channel_factory, {}) bundle_processor_cache.activate('instruction_id') worker = SdkWorker(bundle_processor_cache) split_request = beam_fn_api_pb2.InstructionRequest( @@ -153,7 +157,10 @@ def test_inactive_bundle_processor_returns_empty_progress_response(self): def test_failed_bundle_processor_returns_failed_progress_response(self): bundle_processor = mock.MagicMock() - bundle_processor_cache = BundleProcessorCache(None, None, None, {}) + data_channel_factory = mock.create_autospec( + data_plane.GrpcClientDataChannelFactory) + bundle_processor_cache = BundleProcessorCache( + None, None, data_channel_factory, {}) bundle_processor_cache.activate('instruction_id') worker = SdkWorker(bundle_processor_cache) @@ -176,7 +183,10 @@ def test_failed_bundle_processor_returns_failed_progress_response(self): def test_inactive_bundle_processor_returns_empty_split_response(self): bundle_processor = mock.MagicMock() - bundle_processor_cache = BundleProcessorCache(None, None, None, {}) + data_channel_factory = mock.create_autospec( + data_plane.GrpcClientDataChannelFactory) + bundle_processor_cache = BundleProcessorCache( + None, None, data_channel_factory, {}) bundle_processor_cache.activate('instruction_id') worker = SdkWorker(bundle_processor_cache) split_request = beam_fn_api_pb2.InstructionRequest( @@ -262,7 +272,10 @@ def test_harness_monitoring_infos_and_metadata(self): def test_failed_bundle_processor_returns_failed_split_response(self): bundle_processor = mock.MagicMock() - bundle_processor_cache = BundleProcessorCache(None, None, None, {}) + data_channel_factory = mock.create_autospec( + data_plane.GrpcClientDataChannelFactory) + bundle_processor_cache = BundleProcessorCache( + None, None, data_channel_factory, {}) bundle_processor_cache.activate('instruction_id') worker = SdkWorker(bundle_processor_cache) @@ -338,6 +351,29 @@ def stop(self): self.assertEqual(response, expected_response) + def test_bundle_processor_creation_failure_cleans_up_grpc_data_channel(self): + data_channel_factory = data_plane.GrpcClientDataChannelFactory() + channel = data_channel_factory.create_data_channel_from_url('some_url') + state_handler_factory = mock.create_autospec( + sdk_worker.GrpcStateHandlerFactory) + bundle_processor_cache = BundleProcessorCache( + frozenset(), state_handler_factory, data_channel_factory, {}) + if bundle_processor_cache.periodic_shutdown: + bundle_processor_cache.periodic_shutdown.cancel() + + bundle_processor_cache.get = mock.MagicMock( + side_effect=RuntimeError('test error')) + + worker = SdkWorker(bundle_processor_cache) + instruction_id = 'instruction_id' + request = beam_fn_api_pb2.ProcessBundleRequest( + process_bundle_descriptor_id='descriptor_id') + + with self.assertRaises(RuntimeError): + worker.process_bundle(request, instruction_id) + + self.assertIn(instruction_id, channel._cleaned_instruction_ids) + class CachingStateHandlerTest(unittest.TestCase): def test_caching(self):