[Python SDK] Add support for named data streams to provide bundle isolation

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,48 @@ 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)
+
 
 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 +198,14 @@ 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 +214,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 +1266,9 @@ 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 +1277,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 +1305,27 @@ 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(
-            list(self.timers_info.keys()))
+      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 +1672,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 +1693,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)

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,80 @@ 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):
+      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, None)
+    data_channel_factory.create_data_channel.assert_any_call(
+        sink_port, 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()