[Bug]: Python sdk data plane can become stuck on instructions that had an exception while calling BundleProcessor setup

[scwhittle]

What happened?

When receiving elements over the grpc stream, the data plane dispatches to per-bundle queues which have a maximum size. Thus if nothing consumes from the queue the data plane can become blocked forever.

216f0d9 fixed this issue when an exception was triggered during processing of data for a instruction by removing the channel and remembering that it had been cleaned up.
https://github.com/apache/beam/blob/master/sdks/python/apache_beam/runners/worker/data_plane.py#L578

However if an exception occurs during setup of the bundle processor, _clean_receiving_queue is not invoked.
https://github.com/apache/beam/blob/master/sdks/python/apache_beam/runners/worker/sdk_worker.py#L687
just appears to propagate the failure to the response on the control stream:
https://github.com/apache/beam/blob/master/sdks/python/apache_beam/runners/worker/sdk_worker.py#L312

This stuckness was observed in Python 2.38 but it seems like it could still trigger by exceptions during bundle processor setup (which currently includes user DoFn setup method).

Some possible fixes:

• defer invoking setup until first data has arrived so that it is beneath the data_plane try block. However this wouldn't work if we add support to python sdk to inline data with the control request as the data plane may never be used for a request.
• expose the cleanup method on the data plane and invoke it from process_bundle error paths in sdk_worker. This could work but it theoretically is possible to still get stuck if the set size of cleaned up instructions in the data_plane is exceeded.

Issue Priority

Priority: 2 (default / most bugs should be filed as P2)

Issue Components

• Component: Python SDK
• Component: Java SDK
• Component: Go SDK
• Component: Typescript SDK
• Component: IO connector
• Component: Beam YAML
• Component: Beam examples
• Component: Beam playground
• Component: Beam katas
• Component: Website
• Component: Spark Runner
• Component: Flink Runner
• Component: Samza Runner
• Component: Twister2 Runner
• Component: Hazelcast Jet Runner
• Component: Google Cloud Dataflow Runner

[scwhittle]

@tvalentyn what are your thoughts on how to best fix this?

[tvalentyn]

It should be possible, but seems to require some coordination between the thread that populates the data plane and the bundle processor. I think I have a repro. The repro has:

• many elements for a single instruction
• waits in DoFn.setup() to saturate input queue
• fails in setup 1 in 3 times.

class FlakySetup(beam.DoFn):
  import collections

  counter = collections.defaultdict(int)

  def process(self, x):
    yield x

  def setup(self):
    self.counter['setup'] = self.counter['setup']+1
    logging.warning("In setup.")
    if self.counter['setup'] % 3 == 0:
      logging.warning("Will fail in 60 seconds")
      time.sleep(60)
      raise RuntimeError

  def start_bundle(self):
    self.counter['start_bundle'] = self.counter['start_bundle']+1
    logging.warning("In start_bundle.")

def flaky_setup(p):
  for i in range(20):
    p | beam.Create(list(range(1000000))) | beam.Reshuffle() | beam.ParDo(FlakySetup()) #| beam.LogElements(level=logging.WARNING)


def run(argv=None):
  parser = argparse.ArgumentParser()

  known_args, pipeline_args = parser.parse_known_args(argv)

  with beam.Pipeline(argv=pipeline_args) as p:
    flaky_setup(p)

[tvalentyn]

Was also using --no_use_multiple_sdk_containers --number_of_worker_harness_threads=12 --auto_unique_labels --num_workers=1 --max_num_workers=1

[tvalentyn]

stuck job (with some extra logging: https://console.cloud.google.com/dataflow/jobs/us-central1/2025-09-30_16_16_12-13233345857016967751)