From 1220f56821531ddb92bb1be49407d9c1dcd07862 Mon Sep 17 00:00:00 2001 From: Sam Whittle Date: Mon, 31 Aug 2026 15:47:42 +0200 Subject: [PATCH 1/3] Initial support for named data streams for python sdk. --- .../runners/worker/bundle_processor.py | 66 ++++++++++--- .../runners/worker/bundle_processor_test.py | 74 +++++++++++++++ .../apache_beam/runners/worker/data_plane.py | 93 ++++++++++++------- .../runners/worker/data_plane_test.py | 39 +++++++- .../apache_beam/runners/worker/sdk_worker.py | 3 +- .../runners/worker/sdk_worker_test.py | 19 ++++ .../runners/worker/worker_id_interceptor.py | 40 +++++++- .../worker/worker_id_interceptor_test.py | 45 ++++++++- .../apache_beam/transforms/environments.py | 1 + .../transforms/environments_test.py | 1 + 10 files changed, 327 insertions(+), 54 deletions(-) diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor.py b/sdks/python/apache_beam/runners/worker/bundle_processor.py index 851efc81221d..a254eca52dbf 100644 --- a/sdks/python/apache_beam/runners/worker/bundle_processor.py +++ b/sdks/python/apache_beam/runners/worker/bundle_processor.py @@ -26,6 +26,7 @@ import collections import concurrent.futures import copy +import functools import heapq import itertools import json @@ -136,22 +137,46 @@ def __init__( state_sampler: statesampler.StateSampler, windowed_coder: coders.Coder, transform_id: str, - data_channel: data_plane.DataChannel) -> None: + data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel]) -> None: super().__init__(name_context, None, counter_factory, state_sampler) self.windowed_coder = windowed_coder self.windowed_coder_impl = windowed_coder.get_impl() # transform_id represents the consumer for the bytes in the data plane for a # DataInputOperation or a producer of these bytes for a DataOutputOperation. self.transform_id = transform_id - self.data_channel = data_channel + self.data_channel_factory = data_channel_factory for _, consumer_ops in consumers.items(): for consumer in consumer_ops: self.add_receiver(consumer, 0) + def get_data_channel( + self, data_stream_id: Optional[str] = None) -> data_plane.DataChannel: + return self.data_channel_factory(data_stream_id=data_stream_id) + class DataOutputOperation(RunnerIOOperation): """A sink-like operation that gathers outputs to be sent back to the runner. """ + def __init__( + self, + operation_name: common.NameContext, + step_name: Any, + consumers: Mapping[Any, list[operations.Operation]], + counter_factory: counters.CounterFactory, + state_sampler: statesampler.StateSampler, + windowed_coder: coders.Coder, + transform_id: str, + data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel]) -> None: + super().__init__( + operation_name, + step_name, + consumers, + counter_factory, + state_sampler, + windowed_coder, + transform_id=transform_id, + data_channel_factory=data_channel_factory) + def set_output_stream( self, output_stream: data_plane.ClosableOutputStream) -> None: self.output_stream = output_stream @@ -171,13 +196,13 @@ class DataInputOperation(RunnerIOOperation): def __init__( self, operation_name: common.NameContext, - step_name, + step_name: Any, consumers: Mapping[Any, list[operations.Operation]], counter_factory: counters.CounterFactory, state_sampler: statesampler.StateSampler, windowed_coder: coders.Coder, - transform_id, - data_channel: data_plane.GrpcClientDataChannel) -> None: + transform_id: str, + data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel]) -> None: super().__init__( operation_name, step_name, @@ -186,7 +211,7 @@ def __init__( state_sampler, windowed_coder, transform_id=transform_id, - data_channel=data_channel) + data_channel_factory=data_channel_factory) self.consumer = next(iter(consumers.values())) self.splitting_lock = threading.Lock() @@ -1238,7 +1263,7 @@ def reset(self) -> None: op.reset() def process_bundle( - self, instruction_id: str + self, instruction_id: str, data_stream_id: Optional[str] = None ) -> tuple[list[beam_fn_api_pb2.DelayedBundleApplication], bool]: expected_input_ops: list[DataInputOperation] = [] @@ -1247,8 +1272,9 @@ def process_bundle( if isinstance(op, DataOutputOperation): # TODO(robertwb): Is there a better way to pass the instruction id to # the operation? + data_channel = op.get_data_channel(data_stream_id) op.set_output_stream( - op.data_channel.output_stream(instruction_id, op.transform_id)) + data_channel.output_stream(instruction_id, op.transform_id)) elif isinstance(op, DataInputOperation): # We must wait until we receive "end of stream" for each of these ops. expected_input_ops.append(op) @@ -1274,18 +1300,28 @@ def process_bundle( # Add expected data inputs for each data channel. input_op_by_transform_id = {} for input_op in expected_input_ops: - data_channels[input_op.data_channel].append(input_op.transform_id) + data_channel = input_op.get_data_channel(data_stream_id) + data_channels[data_channel].append(input_op.transform_id) input_op_by_transform_id[input_op.transform_id] = input_op # Update timer_data channel with expected timer inputs. - if self.timer_data_channel: - data_channels[self.timer_data_channel].extend( + timer_data_channel = None + if self.process_bundle_descriptor.timer_api_service_descriptor.url: + timer_data_channel = ( + self.data_channel_factory.create_data_channel_from_url( + self.process_bundle_descriptor.timer_api_service_descriptor.url, + data_stream_id=data_stream_id)) + elif self.timer_data_channel: + timer_data_channel = self.timer_data_channel + + if timer_data_channel: + data_channels[timer_data_channel].extend( list(self.timers_info.keys())) # Set up timer output stream for DoOperation. for ((transform_id, timer_family_id), timer_info) in self.timers_info.items(): - output_stream = self.timer_data_channel.output_timer_stream( + output_stream = timer_data_channel.output_timer_stream( instruction_id, transform_id, timer_family_id) timer_info.output_stream = output_stream self.ops[transform_id].add_timer_info(timer_family_id, timer_info) @@ -1632,7 +1668,8 @@ def create_source_runner( factory.state_sampler, output_coder, transform_id=transform_id, - data_channel=factory.data_channel_factory.create_data_channel(grpc_port)) + data_channel_factory=functools.partial( + factory.data_channel_factory.create_data_channel, grpc_port)) @BeamTransformFactory.register_urn( @@ -1652,7 +1689,8 @@ def create_sink_runner( factory.state_sampler, output_coder, transform_id=transform_id, - data_channel=factory.data_channel_factory.create_data_channel(grpc_port)) + data_channel_factory=functools.partial( + factory.data_channel_factory.create_data_channel, grpc_port)) @BeamTransformFactory.register_urn(OLD_DATAFLOW_RUNNER_HARNESS_READ_URN, None) diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor_test.py b/sdks/python/apache_beam/runners/worker/bundle_processor_test.py index 0eb4dd9485fd..8615ae3d7f1b 100644 --- a/sdks/python/apache_beam/runners/worker/bundle_processor_test.py +++ b/sdks/python/apache_beam/runners/worker/bundle_processor_test.py @@ -20,6 +20,7 @@ import random import unittest +from unittest import mock import apache_beam as beam from apache_beam.coders import StrUtf8Coder @@ -736,5 +737,78 @@ def test_continuation_token(self): self.assertEqual([A1, A2, A7, B7, A8], list(self.state.read())) +class NamedDataStreamsTest(unittest.TestCase): + def test_named_data_streams_routing(self): + descriptor = beam_fn_api_pb2.ProcessBundleDescriptor(id='descriptor_id') + + # Coders + CODER_ID = 'coder' + descriptor.coders[CODER_ID].spec.urn = common_urns.StandardCoders.Enum.BYTES.urn + + # PCollections + PCOLLECTION_IN = 'pcoll_in' + descriptor.pcollections[PCOLLECTION_IN].unique_name = PCOLLECTION_IN + descriptor.pcollections[PCOLLECTION_IN].coder_id = CODER_ID + + PCOLLECTION_OUT = 'pcoll_out' + descriptor.pcollections[PCOLLECTION_OUT].unique_name = PCOLLECTION_OUT + descriptor.pcollections[PCOLLECTION_OUT].coder_id = CODER_ID + + # Source transform + SOURCE_ID = 'source' + source_transform = descriptor.transforms[SOURCE_ID] + source_transform.spec.urn = bundle_processor.DATA_INPUT_URN + source_port = beam_fn_api_pb2.RemoteGrpcPort(coder_id=CODER_ID) + source_port.api_service_descriptor.url = 'localhost:123' + source_transform.spec.payload = source_port.SerializeToString() + source_transform.outputs['None'] = PCOLLECTION_IN + + # Sink transform + SINK_ID = 'sink' + sink_transform = descriptor.transforms[SINK_ID] + sink_transform.spec.urn = bundle_processor.DATA_OUTPUT_URN + sink_port = beam_fn_api_pb2.RemoteGrpcPort(coder_id=CODER_ID) + sink_port.api_service_descriptor.url = 'localhost:123' + sink_transform.spec.payload = sink_port.SerializeToString() + sink_transform.inputs['None'] = PCOLLECTION_IN + sink_transform.outputs['None'] = PCOLLECTION_OUT + + data_channel_factory = mock.MagicMock() + mock_channel_default = mock.MagicMock() + mock_channel_named = mock.MagicMock() + + def get_channel(port, data_stream_id=None): + if data_stream_id == 'named_stream': + return mock_channel_named + return mock_channel_default + + data_channel_factory.create_data_channel.side_effect = get_channel + + mock_channel_default.input_elements.return_value = [] + mock_channel_named.input_elements.return_value = [] + + processor = BundleProcessor( + frozenset(), descriptor, None, data_channel_factory) + + # Process on default stream + processor.process_bundle('inst_1') + data_channel_factory.create_data_channel.assert_any_call( + source_port, data_stream_id=None) + data_channel_factory.create_data_channel.assert_any_call( + sink_port, data_stream_id=None) + mock_channel_default.output_stream.assert_called_once_with('inst_1', SINK_ID) + + processor.reset() + + # Process on named stream + processor.process_bundle('inst_2', data_stream_id='named_stream') + + data_channel_factory.create_data_channel.assert_any_call( + source_port, data_stream_id='named_stream') + data_channel_factory.create_data_channel.assert_any_call( + sink_port, data_stream_id='named_stream') + mock_channel_named.output_stream.assert_called_once_with('inst_2', SINK_ID) + + if __name__ == '__main__': unittest.main() diff --git a/sdks/python/apache_beam/runners/worker/data_plane.py b/sdks/python/apache_beam/runners/worker/data_plane.py index cfefa37d76b6..a591f3b22f8b 100644 --- a/sdks/python/apache_beam/runners/worker/data_plane.py +++ b/sdks/python/apache_beam/runners/worker/data_plane.py @@ -48,6 +48,7 @@ from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.runners.worker.channel_factory import GRPCChannelFactory +from apache_beam.runners.worker.worker_id_interceptor import DataStreamIdInterceptor from apache_beam.runners.worker.worker_id_interceptor import WorkerIdInterceptor from apache_beam.utils.byte_limited_queue import ByteLimitedQueue @@ -804,12 +805,15 @@ def __init__( self._lock = threading.Lock() self._connections_by_worker_id = collections.defaultdict( lambda: _GrpcDataChannel(data_buffer_time_limit_ms) - ) # type: DefaultDict[str, _GrpcDataChannel] + ) # type: DefaultDict[Tuple[str, str], _GrpcDataChannel] - def get_conn_by_worker_id(self, worker_id): - # type: (str) -> _GrpcDataChannel + def get_conn_by_worker_id( + self, + worker_id: str, + data_stream_id: Optional[str] = None + ) -> _GrpcDataChannel: with self._lock: - return self._connections_by_worker_id[worker_id] + return self._connections_by_worker_id[(worker_id, data_stream_id or '')] def Data( self, @@ -817,8 +821,10 @@ def Data( context # type: Any ): # type: (...) -> Iterator[beam_fn_api_pb2.Elements] - worker_id = dict(context.invocation_metadata())['worker_id'] - data_conn = self.get_conn_by_worker_id(worker_id) + metadata = dict(context.invocation_metadata()) + worker_id = metadata['worker_id'] + data_stream_id = metadata.get('data_stream_id', '') + data_conn = self.get_conn_by_worker_id(worker_id, data_stream_id) data_conn.set_inputs(elements_iterator) for elements in data_conn._write_outputs(): yield elements @@ -827,16 +833,20 @@ def Data( class DataChannelFactory(metaclass=abc.ABCMeta): """An abstract factory for creating ``DataChannel``.""" @abc.abstractmethod - def create_data_channel(self, remote_grpc_port): - # type: (beam_fn_api_pb2.RemoteGrpcPort) -> GrpcClientDataChannel - + def create_data_channel( + self, + remote_grpc_port: beam_fn_api_pb2.RemoteGrpcPort, + data_stream_id: Optional[str] = None + ) -> DataChannel: """Returns a ``DataChannel`` from the given RemoteGrpcPort.""" raise NotImplementedError(type(self)) @abc.abstractmethod - def create_data_channel_from_url(self, url): - # type: (str) -> Optional[GrpcClientDataChannel] - + def create_data_channel_from_url( + self, + url: str, + data_stream_id: Optional[str] = None + ) -> Optional[DataChannel]: """Returns a ``DataChannel`` from the given url.""" raise NotImplementedError(type(self)) @@ -857,7 +867,7 @@ def cleanup(self, instruction_id): class GrpcClientDataChannelFactory(DataChannelFactory): """A factory for ``GrpcClientDataChannel``. - Caches the created channels by ``data descriptor url``. + Caches the created channels by ``(data descriptor url, data_stream_id)``. """ def __init__( self, @@ -866,7 +876,7 @@ def __init__( data_buffer_time_limit_ms=0 # type: int ): # type: (...) -> None - self._data_channel_cache = {} # type: Dict[str, GrpcClientDataChannel] + self._data_channel_cache = {} # type: Dict[Tuple[str, str], GrpcClientDataChannel] self._lock = threading.Lock() self._credentials = None self._worker_id = worker_id @@ -875,14 +885,22 @@ def __init__( _LOGGER.info('Using secure channel creds.') self._credentials = credentials - def create_data_channel_from_url(self, url): - # type: (str) -> Optional[GrpcClientDataChannel] + def create_data_channel_from_url( + self, + url: str, + data_stream_id: Optional[str] = None + ) -> Optional[GrpcClientDataChannel]: if not url: return None - if url not in self._data_channel_cache: + data_stream_id = data_stream_id or '' + cache_key = (url, data_stream_id) + if cache_key not in self._data_channel_cache: with self._lock: - if url not in self._data_channel_cache: - _LOGGER.info('Creating client data channel for %s', url) + if cache_key not in self._data_channel_cache: + _LOGGER.info( + 'Creating client data channel for %s (data_stream_id: %s)', + url, + data_stream_id) # Options to have no limits (-1) on the size of the messages # received or sent over the data plane. The actual buffer size # is controlled in a layer above. @@ -897,22 +915,27 @@ def create_data_channel_from_url(self, url): grpc_channel = GRPCChannelFactory.secure_channel( url, self._credentials, options=channel_options) _LOGGER.info('Data channel established.') - # Add workerId to the grpc channel - grpc_channel = grpc.intercept_channel( - grpc_channel, WorkerIdInterceptor(self._worker_id)) - self._data_channel_cache[url] = GrpcClientDataChannel( + # Add workerId and optional data_stream_id to the grpc channel + interceptors = [WorkerIdInterceptor(self._worker_id)] + if data_stream_id: + interceptors.append(DataStreamIdInterceptor(data_stream_id)) + grpc_channel = grpc.intercept_channel(grpc_channel, *interceptors) + self._data_channel_cache[cache_key] = GrpcClientDataChannel( beam_fn_api_pb2_grpc.BeamFnDataStub(grpc_channel), self._data_buffer_time_limit_ms) - return self._data_channel_cache[url] + return self._data_channel_cache[cache_key] - def create_data_channel(self, remote_grpc_port): - # type: (beam_fn_api_pb2.RemoteGrpcPort) -> GrpcClientDataChannel + def create_data_channel( + self, + remote_grpc_port: beam_fn_api_pb2.RemoteGrpcPort, + data_stream_id: Optional[str] = None + ) -> GrpcClientDataChannel: url = remote_grpc_port.api_service_descriptor.url # TODO(https://github.com/apache/beam/issues/19737): this can return None # if url is falsey, but this seems incorrect, as code that calls this # method seems to always expect non-Optional values. - return self.create_data_channel_from_url(url) # type: ignore[return-value] + return self.create_data_channel_from_url(url, data_stream_id=data_stream_id) # type: ignore[return-value] def close(self): # type: () -> None @@ -930,15 +953,21 @@ def cleanup(self, instruction_id): class InMemoryDataChannelFactory(DataChannelFactory): """A singleton factory for ``InMemoryDataChannel``.""" def __init__(self, in_memory_data_channel): - # type: (GrpcClientDataChannel) -> None + # type: (DataChannel) -> None self._in_memory_data_channel = in_memory_data_channel - def create_data_channel(self, unused_remote_grpc_port): - # type: (beam_fn_api_pb2.RemoteGrpcPort) -> GrpcClientDataChannel + def create_data_channel( + self, + unused_remote_grpc_port: beam_fn_api_pb2.RemoteGrpcPort, + data_stream_id: Optional[str] = None + ) -> DataChannel: return self._in_memory_data_channel - def create_data_channel_from_url(self, url): - # type: (Any) -> GrpcClientDataChannel + def create_data_channel_from_url( + self, + url: Any, + data_stream_id: Optional[str] = None + ) -> DataChannel: return self._in_memory_data_channel def close(self): diff --git a/sdks/python/apache_beam/runners/worker/data_plane_test.py b/sdks/python/apache_beam/runners/worker/data_plane_test.py index 5124bb69e6c8..534f3f91ad4e 100644 --- a/sdks/python/apache_beam/runners/worker/data_plane_test.py +++ b/sdks/python/apache_beam/runners/worker/data_plane_test.py @@ -29,6 +29,7 @@ from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.runners.worker import data_plane +from apache_beam.runners.worker.worker_id_interceptor import DataStreamIdInterceptor from apache_beam.runners.worker.worker_id_interceptor import WorkerIdInterceptor from apache_beam.utils import thread_pool_executor @@ -40,7 +41,10 @@ def test_grpc_data_channel(self): def test_time_based_flush_grpc_data_channel(self): self._grpc_data_channel_test(True) - def _grpc_data_channel_test(self, time_based_flush=False): + def test_named_grpc_data_channel(self): + self._grpc_data_channel_test(data_stream_id='stream_1') + + def _grpc_data_channel_test(self, time_based_flush=False, data_stream_id=''): if time_based_flush: data_servicer = data_plane.BeamFnDataServicer( data_buffer_time_limit_ms=100) @@ -48,7 +52,7 @@ def _grpc_data_channel_test(self, time_based_flush=False): data_servicer = data_plane.BeamFnDataServicer() worker_id = 'worker_0' data_channel_service = \ - data_servicer.get_conn_by_worker_id(worker_id) + data_servicer.get_conn_by_worker_id(worker_id, data_stream_id=data_stream_id) server = grpc.server(thread_pool_executor.shared_unbounded_instance()) beam_fn_api_pb2_grpc.add_BeamFnDataServicer_to_server(data_servicer, server) @@ -56,9 +60,10 @@ def _grpc_data_channel_test(self, time_based_flush=False): server.start() grpc_channel = grpc.insecure_channel('localhost:%s' % test_port) - # Add workerId to the grpc channel - grpc_channel = grpc.intercept_channel( - grpc_channel, WorkerIdInterceptor(worker_id)) + interceptors = [WorkerIdInterceptor(worker_id)] + if data_stream_id: + interceptors.append(DataStreamIdInterceptor(data_stream_id)) + grpc_channel = grpc.intercept_channel(grpc_channel, *interceptors) data_channel_stub = beam_fn_api_pb2_grpc.BeamFnDataStub(grpc_channel) if time_based_flush: data_channel_client = data_plane.GrpcClientDataChannel( @@ -75,6 +80,30 @@ def _grpc_data_channel_test(self, time_based_flush=False): data_channel_client.wait() data_channel_service.wait() + def test_grpc_client_data_channel_factory_named_streams(self): + data_servicer = data_plane.BeamFnDataServicer() + server = grpc.server(thread_pool_executor.shared_unbounded_instance()) + beam_fn_api_pb2_grpc.add_BeamFnDataServicer_to_server(data_servicer, server) + test_port = server.add_insecure_port('[::]:0') + server.start() + + try: + factory = data_plane.GrpcClientDataChannelFactory(worker_id='worker_0') + url = 'localhost:%s' % test_port + ch_default = factory.create_data_channel_from_url(url) + ch_default2 = factory.create_data_channel_from_url(url, data_stream_id='') + ch_stream1 = factory.create_data_channel_from_url(url, data_stream_id='stream_1') + ch_stream1_dup = factory.create_data_channel_from_url(url, data_stream_id='stream_1') + ch_stream2 = factory.create_data_channel_from_url(url, data_stream_id='stream_2') + + self.assertIs(ch_default, ch_default2) + self.assertIs(ch_stream1, ch_stream1_dup) + self.assertIsNot(ch_default, ch_stream1) + self.assertIsNot(ch_stream1, ch_stream2) + factory.close() + finally: + server.stop(0) + def test_in_memory_data_channel(self): channel = data_plane.InMemoryDataChannel() self._data_channel_test(channel, channel.inverse()) diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker.py b/sdks/python/apache_beam/runners/worker/sdk_worker.py index db25a40a405c..84b4ae1bf231 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker.py @@ -706,7 +706,8 @@ def process_bundle( instruction_id, request.cache_tokens): with self.maybe_profile(instruction_id): delayed_applications, requests_finalization = ( - bundle_processor.process_bundle(instruction_id)) + bundle_processor.process_bundle( + instruction_id, getattr(request, 'data_stream_id', ''))) monitoring_infos = bundle_processor.monitoring_infos() response = beam_fn_api_pb2.InstructionResponse( instruction_id=instruction_id, 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 bea313a4d2fd..25d5ca96c346 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py @@ -412,6 +412,25 @@ def test_bundle_processor_creation_failure_cleans_up_grpc_data_channel(self): self.assertIn(instruction_id, channel._cleaned_instruction_ids) + def test_process_bundle_passes_data_stream_id(self): + mock_bundle_processor = mock.MagicMock() + mock_bundle_processor.process_bundle.return_value = ([], False) + mock_bundle_processor.monitoring_infos.return_value = [] + mock_bundle_processor.state_handler.process_instruction_id.return_value = contextlib.nullcontext() + + bundle_processor_cache = mock.MagicMock() + bundle_processor_cache.get.return_value = mock_bundle_processor + + worker = SdkWorker(bundle_processor_cache) + instruction_id = 'instruction_id' + request = beam_fn_api_pb2.ProcessBundleRequest( + process_bundle_descriptor_id='descriptor_id', + data_stream_id='stream_xyz') + + worker.process_bundle(request, instruction_id) + mock_bundle_processor.process_bundle.assert_called_once_with( + instruction_id, 'stream_xyz') + class CachingStateHandlerTest(unittest.TestCase): def test_caching(self): diff --git a/sdks/python/apache_beam/runners/worker/worker_id_interceptor.py b/sdks/python/apache_beam/runners/worker/worker_id_interceptor.py index 1db2b5f4a151..bd38528e6d1b 100644 --- a/sdks/python/apache_beam/runners/worker/worker_id_interceptor.py +++ b/sdks/python/apache_beam/runners/worker/worker_id_interceptor.py @@ -61,12 +61,50 @@ def _intercept(self, continuation, client_call_details, request): metadata = [] if client_call_details.metadata is not None: metadata = list(client_call_details.metadata) - if 'worker_id' in metadata: + if any(k == 'worker_id' for k, _ in metadata): raise RuntimeError('Header metadata already has a worker_id.') metadata.append(('worker_id', self._worker_id)) + new_client_details = _ClientCallDetails( client_call_details.method, client_call_details.timeout, metadata, client_call_details.credentials) return continuation(new_client_details, request) + + +class DataStreamIdInterceptor(grpc.UnaryUnaryClientInterceptor, + grpc.StreamStreamClientInterceptor): + """Client Interceptor to inject data_stream_id into metadata.""" + def __init__(self, data_stream_id: Optional[str] = None) -> None: + self._data_stream_id = data_stream_id + + def intercept_unary_unary(self, continuation, client_call_details, request): + return self._intercept(continuation, client_call_details, request) + + def intercept_unary_stream(self, continuation, client_call_details, request): + return self._intercept(continuation, client_call_details, request) + + def intercept_stream_unary(self, continuation, client_call_details, request): + return self._intercept(continuation, client_call_details, request) + + def intercept_stream_stream( + self, continuation, client_call_details, request_iterator): + return self._intercept(continuation, client_call_details, request_iterator) + + def _intercept(self, continuation, client_call_details, request): + if self._data_stream_id: + metadata = [] + if client_call_details.metadata is not None: + metadata = list(client_call_details.metadata) + if any(k == 'data_stream_id' for k, _ in metadata): + raise RuntimeError('Header metadata already has a data_stream_id.') + metadata.append(('data_stream_id', self._data_stream_id)) + new_client_details = _ClientCallDetails( + client_call_details.method, + client_call_details.timeout, + metadata, + client_call_details.credentials) + return continuation(new_client_details, request) + return continuation(client_call_details, request) + diff --git a/sdks/python/apache_beam/runners/worker/worker_id_interceptor_test.py b/sdks/python/apache_beam/runners/worker/worker_id_interceptor_test.py index 0db9c1b4ddc0..1f42f97d419d 100644 --- a/sdks/python/apache_beam/runners/worker/worker_id_interceptor_test.py +++ b/sdks/python/apache_beam/runners/worker/worker_id_interceptor_test.py @@ -24,6 +24,7 @@ import grpc +from apache_beam.runners.worker.worker_id_interceptor import DataStreamIdInterceptor from apache_beam.runners.worker.worker_id_interceptor import WorkerIdInterceptor @@ -62,9 +63,51 @@ def continuation(client_details, request_iterator): with self.assertRaises(RuntimeError): WorkerIdInterceptor().intercept_stream_stream( continuation, - _ClientCallDetails(None, None, {'worker_id': '1'}, None), []) + _ClientCallDetails(None, None, [('worker_id', '1')], None), []) + + +class DataStreamIdInterceptorTest(unittest.TestCase): + def test_data_stream_id_insertion(self): + data_stream_id_key = 'data_stream_id' + headers_holder = {} + + def continuation(client_details, request_iterator): + headers_holder.update( + {data_stream_id_key: dict(client_details.metadata).get(data_stream_id_key)}) + + DataStreamIdInterceptor('stream_123').intercept_stream_stream( + continuation, _ClientCallDetails(None, None, None, None), []) + self.assertEqual( + headers_holder[data_stream_id_key], + 'stream_123', + 'data_stream_id not set') + + def test_no_data_stream_id_when_empty(self): + headers_holder = {} + + def continuation(client_details, request_iterator): + headers_holder['metadata'] = client_details.metadata + + DataStreamIdInterceptor('').intercept_stream_stream( + continuation, _ClientCallDetails(None, None, None, None), []) + self.assertIsNone(headers_holder['metadata']) + + def test_failure_when_data_stream_id_exists(self): + data_stream_id_key = 'data_stream_id' + headers_holder = {} + + def continuation(client_details, request_iterator): + headers_holder.update( + {data_stream_id_key: dict(client_details.metadata).get(data_stream_id_key)}) + + with self.assertRaises(RuntimeError): + DataStreamIdInterceptor('stream_123').intercept_stream_stream( + continuation, + _ClientCallDetails(None, None, [('data_stream_id', 'existing')], None), + []) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() + diff --git a/sdks/python/apache_beam/transforms/environments.py b/sdks/python/apache_beam/transforms/environments.py index 6f17ca270ebd..9d5e6a48cc03 100644 --- a/sdks/python/apache_beam/transforms/environments.py +++ b/sdks/python/apache_beam/transforms/environments.py @@ -917,6 +917,7 @@ def _python_sdk_capabilities_iter(): yield common_urns.protocols.DATA_SAMPLING.urn yield common_urns.protocols.SDK_CONSUMING_RECEIVED_DATA.urn yield common_urns.protocols.ORDERED_LIST_STATE.urn + yield common_urns.protocols.NAMED_DATA_STREAMS.urn def python_sdk_dependencies(options, tmp_dir=None): diff --git a/sdks/python/apache_beam/transforms/environments_test.py b/sdks/python/apache_beam/transforms/environments_test.py index c32a85579fcb..ebc9a9f7bbc8 100644 --- a/sdks/python/apache_beam/transforms/environments_test.py +++ b/sdks/python/apache_beam/transforms/environments_test.py @@ -72,6 +72,7 @@ def test_sdk_capabilities(self): common_urns.sdf_components.TRUNCATE_SIZED_RESTRICTION.urn, sdk_capabilities) self.assertIn(common_urns.primitives.TO_STRING.urn, sdk_capabilities) + self.assertIn(common_urns.protocols.NAMED_DATA_STREAMS.urn, sdk_capabilities) def test_default_capabilities(self): environment = DockerEnvironment.from_options( From ca39ed4c42f3e562b3ee99ac6816f1ffd43efd32 Mon Sep 17 00:00:00 2001 From: Sam Whittle Date: Thu, 10 Sep 2026 18:37:20 +0200 Subject: [PATCH 2/3] fix format and lint:: --- .../runners/worker/bundle_processor.py | 18 ++++++++------ .../runners/worker/bundle_processor_test.py | 6 +++-- .../apache_beam/runners/worker/data_plane.py | 24 +++++++------------ .../runners/worker/data_plane_test.py | 9 ++++--- .../runners/worker/sdk_worker_test.py | 3 ++- .../runners/worker/worker_id_interceptor.py | 1 - .../worker/worker_id_interceptor_test.py | 17 +++++++------ .../transforms/environments_test.py | 3 ++- 8 files changed, 44 insertions(+), 37 deletions(-) diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor.py b/sdks/python/apache_beam/runners/worker/bundle_processor.py index a254eca52dbf..d7cf0af4bb0b 100644 --- a/sdks/python/apache_beam/runners/worker/bundle_processor.py +++ b/sdks/python/apache_beam/runners/worker/bundle_processor.py @@ -137,7 +137,8 @@ def __init__( state_sampler: statesampler.StateSampler, windowed_coder: coders.Coder, transform_id: str, - data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel]) -> None: + data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel]) + -> None: super().__init__(name_context, None, counter_factory, state_sampler) self.windowed_coder = windowed_coder self.windowed_coder_impl = windowed_coder.get_impl() @@ -151,7 +152,7 @@ def __init__( def get_data_channel( self, data_stream_id: Optional[str] = None) -> data_plane.DataChannel: - return self.data_channel_factory(data_stream_id=data_stream_id) + return self.data_channel_factory(data_stream_id) class DataOutputOperation(RunnerIOOperation): @@ -166,7 +167,8 @@ def __init__( state_sampler: statesampler.StateSampler, windowed_coder: coders.Coder, transform_id: str, - data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel]) -> None: + data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel]) + -> None: super().__init__( operation_name, step_name, @@ -202,7 +204,8 @@ def __init__( state_sampler: statesampler.StateSampler, windowed_coder: coders.Coder, transform_id: str, - data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel]) -> None: + data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel]) + -> None: super().__init__( operation_name, step_name, @@ -1263,7 +1266,9 @@ def reset(self) -> None: op.reset() def process_bundle( - self, instruction_id: str, data_stream_id: Optional[str] = None + self, + instruction_id: str, + data_stream_id: Optional[str] = None ) -> tuple[list[beam_fn_api_pb2.DelayedBundleApplication], bool]: expected_input_ops: list[DataInputOperation] = [] @@ -1315,8 +1320,7 @@ def process_bundle( timer_data_channel = self.timer_data_channel if timer_data_channel: - data_channels[timer_data_channel].extend( - list(self.timers_info.keys())) + data_channels[timer_data_channel].extend(list(self.timers_info.keys())) # Set up timer output stream for DoOperation. for ((transform_id, timer_family_id), diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor_test.py b/sdks/python/apache_beam/runners/worker/bundle_processor_test.py index 8615ae3d7f1b..196be8a3864a 100644 --- a/sdks/python/apache_beam/runners/worker/bundle_processor_test.py +++ b/sdks/python/apache_beam/runners/worker/bundle_processor_test.py @@ -743,7 +743,8 @@ def test_named_data_streams_routing(self): # Coders CODER_ID = 'coder' - descriptor.coders[CODER_ID].spec.urn = common_urns.StandardCoders.Enum.BYTES.urn + descriptor.coders[ + CODER_ID].spec.urn = common_urns.StandardCoders.Enum.BYTES.urn # PCollections PCOLLECTION_IN = 'pcoll_in' @@ -796,7 +797,8 @@ def get_channel(port, data_stream_id=None): source_port, data_stream_id=None) data_channel_factory.create_data_channel.assert_any_call( sink_port, data_stream_id=None) - mock_channel_default.output_stream.assert_called_once_with('inst_1', SINK_ID) + mock_channel_default.output_stream.assert_called_once_with( + 'inst_1', SINK_ID) processor.reset() diff --git a/sdks/python/apache_beam/runners/worker/data_plane.py b/sdks/python/apache_beam/runners/worker/data_plane.py index a591f3b22f8b..84f9b16f3670 100644 --- a/sdks/python/apache_beam/runners/worker/data_plane.py +++ b/sdks/python/apache_beam/runners/worker/data_plane.py @@ -810,8 +810,7 @@ def __init__( def get_conn_by_worker_id( self, worker_id: str, - data_stream_id: Optional[str] = None - ) -> _GrpcDataChannel: + data_stream_id: Optional[str] = None) -> _GrpcDataChannel: with self._lock: return self._connections_by_worker_id[(worker_id, data_stream_id or '')] @@ -836,8 +835,7 @@ class DataChannelFactory(metaclass=abc.ABCMeta): def create_data_channel( self, remote_grpc_port: beam_fn_api_pb2.RemoteGrpcPort, - data_stream_id: Optional[str] = None - ) -> DataChannel: + data_stream_id: Optional[str] = None) -> DataChannel: """Returns a ``DataChannel`` from the given RemoteGrpcPort.""" raise NotImplementedError(type(self)) @@ -845,8 +843,7 @@ def create_data_channel( def create_data_channel_from_url( self, url: str, - data_stream_id: Optional[str] = None - ) -> Optional[DataChannel]: + data_stream_id: Optional[str] = None) -> Optional[DataChannel]: """Returns a ``DataChannel`` from the given url.""" raise NotImplementedError(type(self)) @@ -876,7 +873,8 @@ def __init__( data_buffer_time_limit_ms=0 # type: int ): # type: (...) -> None - self._data_channel_cache = {} # type: Dict[Tuple[str, str], GrpcClientDataChannel] + self._data_channel_cache = { + } # type: Dict[Tuple[str, str], GrpcClientDataChannel] self._lock = threading.Lock() self._credentials = None self._worker_id = worker_id @@ -888,8 +886,7 @@ def __init__( def create_data_channel_from_url( self, url: str, - data_stream_id: Optional[str] = None - ) -> Optional[GrpcClientDataChannel]: + data_stream_id: Optional[str] = None) -> Optional[GrpcClientDataChannel]: if not url: return None data_stream_id = data_stream_id or '' @@ -929,8 +926,7 @@ def create_data_channel_from_url( def create_data_channel( self, remote_grpc_port: beam_fn_api_pb2.RemoteGrpcPort, - data_stream_id: Optional[str] = None - ) -> GrpcClientDataChannel: + data_stream_id: Optional[str] = None) -> GrpcClientDataChannel: url = remote_grpc_port.api_service_descriptor.url # TODO(https://github.com/apache/beam/issues/19737): this can return None # if url is falsey, but this seems incorrect, as code that calls this @@ -959,15 +955,13 @@ def __init__(self, in_memory_data_channel): def create_data_channel( self, unused_remote_grpc_port: beam_fn_api_pb2.RemoteGrpcPort, - data_stream_id: Optional[str] = None - ) -> DataChannel: + data_stream_id: Optional[str] = None) -> DataChannel: return self._in_memory_data_channel def create_data_channel_from_url( self, url: Any, - data_stream_id: Optional[str] = None - ) -> DataChannel: + data_stream_id: Optional[str] = None) -> DataChannel: return self._in_memory_data_channel def close(self): diff --git a/sdks/python/apache_beam/runners/worker/data_plane_test.py b/sdks/python/apache_beam/runners/worker/data_plane_test.py index 534f3f91ad4e..16648f23f5f1 100644 --- a/sdks/python/apache_beam/runners/worker/data_plane_test.py +++ b/sdks/python/apache_beam/runners/worker/data_plane_test.py @@ -92,9 +92,12 @@ def test_grpc_client_data_channel_factory_named_streams(self): url = 'localhost:%s' % test_port ch_default = factory.create_data_channel_from_url(url) ch_default2 = factory.create_data_channel_from_url(url, data_stream_id='') - ch_stream1 = factory.create_data_channel_from_url(url, data_stream_id='stream_1') - ch_stream1_dup = factory.create_data_channel_from_url(url, data_stream_id='stream_1') - ch_stream2 = factory.create_data_channel_from_url(url, data_stream_id='stream_2') + ch_stream1 = factory.create_data_channel_from_url( + url, data_stream_id='stream_1') + ch_stream1_dup = factory.create_data_channel_from_url( + url, data_stream_id='stream_1') + ch_stream2 = factory.create_data_channel_from_url( + url, data_stream_id='stream_2') self.assertIs(ch_default, ch_default2) self.assertIs(ch_stream1, ch_stream1_dup) 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 25d5ca96c346..4b00051a3370 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py @@ -416,7 +416,8 @@ def test_process_bundle_passes_data_stream_id(self): mock_bundle_processor = mock.MagicMock() mock_bundle_processor.process_bundle.return_value = ([], False) mock_bundle_processor.monitoring_infos.return_value = [] - mock_bundle_processor.state_handler.process_instruction_id.return_value = contextlib.nullcontext() + mock_bundle_processor.state_handler.process_instruction_id.return_value = contextlib.nullcontext( + ) bundle_processor_cache = mock.MagicMock() bundle_processor_cache.get.return_value = mock_bundle_processor diff --git a/sdks/python/apache_beam/runners/worker/worker_id_interceptor.py b/sdks/python/apache_beam/runners/worker/worker_id_interceptor.py index bd38528e6d1b..ca9476f0c2fd 100644 --- a/sdks/python/apache_beam/runners/worker/worker_id_interceptor.py +++ b/sdks/python/apache_beam/runners/worker/worker_id_interceptor.py @@ -107,4 +107,3 @@ def _intercept(self, continuation, client_call_details, request): client_call_details.credentials) return continuation(new_client_details, request) return continuation(client_call_details, request) - diff --git a/sdks/python/apache_beam/runners/worker/worker_id_interceptor_test.py b/sdks/python/apache_beam/runners/worker/worker_id_interceptor_test.py index 1f42f97d419d..06546ab3ce2c 100644 --- a/sdks/python/apache_beam/runners/worker/worker_id_interceptor_test.py +++ b/sdks/python/apache_beam/runners/worker/worker_id_interceptor_test.py @@ -72,8 +72,10 @@ def test_data_stream_id_insertion(self): headers_holder = {} def continuation(client_details, request_iterator): - headers_holder.update( - {data_stream_id_key: dict(client_details.metadata).get(data_stream_id_key)}) + headers_holder.update({ + data_stream_id_key: dict( + client_details.metadata).get(data_stream_id_key) + }) DataStreamIdInterceptor('stream_123').intercept_stream_stream( continuation, _ClientCallDetails(None, None, None, None), []) @@ -97,17 +99,18 @@ def test_failure_when_data_stream_id_exists(self): headers_holder = {} def continuation(client_details, request_iterator): - headers_holder.update( - {data_stream_id_key: dict(client_details.metadata).get(data_stream_id_key)}) + headers_holder.update({ + data_stream_id_key: dict( + client_details.metadata).get(data_stream_id_key) + }) with self.assertRaises(RuntimeError): DataStreamIdInterceptor('stream_123').intercept_stream_stream( continuation, - _ClientCallDetails(None, None, [('data_stream_id', 'existing')], None), - []) + _ClientCallDetails( + None, None, [('data_stream_id', 'existing')], None), []) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() - diff --git a/sdks/python/apache_beam/transforms/environments_test.py b/sdks/python/apache_beam/transforms/environments_test.py index ebc9a9f7bbc8..21ede9e0d075 100644 --- a/sdks/python/apache_beam/transforms/environments_test.py +++ b/sdks/python/apache_beam/transforms/environments_test.py @@ -72,7 +72,8 @@ def test_sdk_capabilities(self): common_urns.sdf_components.TRUNCATE_SIZED_RESTRICTION.urn, sdk_capabilities) self.assertIn(common_urns.primitives.TO_STRING.urn, sdk_capabilities) - self.assertIn(common_urns.protocols.NAMED_DATA_STREAMS.urn, sdk_capabilities) + self.assertIn( + common_urns.protocols.NAMED_DATA_STREAMS.urn, sdk_capabilities) def test_default_capabilities(self): environment = DockerEnvironment.from_options( From b19c1ebd886f0bd9601bec3fc984acebee380e64 Mon Sep 17 00:00:00 2001 From: Sam Whittle Date: Thu, 10 Sep 2026 18:45:31 +0200 Subject: [PATCH 3/3] fix format and lint --- .../apache_beam/runners/worker/bundle_processor.py | 12 ++++++------ .../runners/worker/bundle_processor_test.py | 6 +++--- sdks/python/apache_beam/runners/worker/data_plane.py | 4 +--- 3 files changed, 10 insertions(+), 12 deletions(-) diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor.py b/sdks/python/apache_beam/runners/worker/bundle_processor.py index d7cf0af4bb0b..e62e522c00cf 100644 --- a/sdks/python/apache_beam/runners/worker/bundle_processor.py +++ b/sdks/python/apache_beam/runners/worker/bundle_processor.py @@ -137,8 +137,8 @@ def __init__( state_sampler: statesampler.StateSampler, windowed_coder: coders.Coder, transform_id: str, - data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel]) - -> None: + data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel] + ) -> None: super().__init__(name_context, None, counter_factory, state_sampler) self.windowed_coder = windowed_coder self.windowed_coder_impl = windowed_coder.get_impl() @@ -167,8 +167,8 @@ def __init__( state_sampler: statesampler.StateSampler, windowed_coder: coders.Coder, transform_id: str, - data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel]) - -> None: + data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel] + ) -> None: super().__init__( operation_name, step_name, @@ -204,8 +204,8 @@ def __init__( state_sampler: statesampler.StateSampler, windowed_coder: coders.Coder, transform_id: str, - data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel]) - -> None: + data_channel_factory: Callable[[Optional[str]], data_plane.DataChannel] + ) -> None: super().__init__( operation_name, step_name, diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor_test.py b/sdks/python/apache_beam/runners/worker/bundle_processor_test.py index 196be8a3864a..3651cea2e600 100644 --- a/sdks/python/apache_beam/runners/worker/bundle_processor_test.py +++ b/sdks/python/apache_beam/runners/worker/bundle_processor_test.py @@ -778,7 +778,7 @@ def test_named_data_streams_routing(self): mock_channel_default = mock.MagicMock() mock_channel_named = mock.MagicMock() - def get_channel(port, data_stream_id=None): + def get_channel(port, data_stream_id): if data_stream_id == 'named_stream': return mock_channel_named return mock_channel_default @@ -794,9 +794,9 @@ def get_channel(port, data_stream_id=None): # Process on default stream processor.process_bundle('inst_1') data_channel_factory.create_data_channel.assert_any_call( - source_port, data_stream_id=None) + source_port, None) data_channel_factory.create_data_channel.assert_any_call( - sink_port, data_stream_id=None) + sink_port, None) mock_channel_default.output_stream.assert_called_once_with( 'inst_1', SINK_ID) diff --git a/sdks/python/apache_beam/runners/worker/data_plane.py b/sdks/python/apache_beam/runners/worker/data_plane.py index 84f9b16f3670..2a626b163a62 100644 --- a/sdks/python/apache_beam/runners/worker/data_plane.py +++ b/sdks/python/apache_beam/runners/worker/data_plane.py @@ -959,9 +959,7 @@ def create_data_channel( return self._in_memory_data_channel def create_data_channel_from_url( - self, - url: Any, - data_stream_id: Optional[str] = None) -> DataChannel: + self, url: Any, data_stream_id: Optional[str] = None) -> DataChannel: return self._in_memory_data_channel def close(self):