[Prism] Fix DEADLINE_EXCEEDED errors caused by worker failures (#38523)

* Set b.BundleErr correctly when dynamic split happens.

* Fix deadlock on worker failures by introducing a bundle-level Done signal

The cancellation signal `b.Done` is to abort bundle data streaming immediately
if the worker fails or completes the bundle early.

Without this signal, a deadlock scenario can occur.
- Runner sends elements through data channel
- SDK worker reads and processes the first element, raises an exception
  and sends InstructionResponse through control channel. However,
  at the same time, Runner keeps sending data and is blocked until
  SDK worker reads more data.
- The existed two signals are not sufficient to do an immediate break
  - ctx.Done() is closed when the pipeline is done or timeout. Given
    the bundle failure, timeout is the only option but it is too late.
  - wk.StoppedChan() is closed when the worker is stopped by the worker
    pool, which is also not happening when the runner is waiting to
    send data.

* Fix a problem of prism log not relaying correctly.

* Trigger more python tests.

* Add a mutex and safe getter/setter methods for bundleErr

* Rename b.Done to b.DataAbort.

* Set bundleErr before closing DataAbort to ensure error will be propagated correctly.

* Refactor Respond to prevent race condition on BundleErr setting

* Remove debug msg

* Formatting.

* Address reviewer comments.

Author: shunping

.github/trigger_files/beam_PostCommit_Python_Versions.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "revision": 1
+  "revision": 2
 }

sdks/go/pkg/beam/runners/prism/internal/execute_test.go

@@ -541,6 +541,27 @@ func TestFailureHang(t *testing.T) {
 	}
 }
 
+func TestFailure_SplitError(t *testing.T) {
+	initRunner(t)
+
+	p, s := beam.NewPipelineWithRoot()
+	configs := beam.Create(s, SourceConfig{NumElements: 100, InitialSplits: 1})
+	col := beam.ParDo(s, &slowFailSDF{}, configs)
+	beam.ParDo(s, &int64Check{Name: "sdf_fail", Want: []int{}}, col)
+
+	// Set a short timeout so the test fails quickly if the pipeline hangs due to split error handling bug
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	_, err := executeWithT(ctx, t, p)
+	if err == nil {
+		t.Fatalf("expected pipeline failure, but got a success")
+	}
+	if want := "intentional split error from tracker"; !strings.Contains(err.Error(), want) {
+		t.Fatalf("expected pipeline failure with %q, but was %v", want, err)
+	}
+}
+
 func TestRunner_Passert(t *testing.T) {
 	initRunner(t)
 	tests := []struct {

sdks/go/pkg/beam/runners/prism/internal/stage.go

@@ -207,8 +207,9 @@ progress:
 			return context.Cause(ctx)
 		case resp = <-b.Resp:
 			bundleFinished = true
-			if b.BundleErr != nil {
-				return b.BundleErr
+			if err := b.GetErr(); err != nil {
+				slog.Error("stage.Execute aborting due to bundle error", "stage", s.ID, "bundle", rb.BundleID)
+				return err
 			}
 			if dataFinished && bundleFinished {
 				break progress // exit progress loop on close.
@@ -240,10 +241,16 @@ progress:
 				sr, err := b.Split(ctx, wk, 0.5 /* fraction of remainder */, nil /* allowed splits */)
 				if err != nil {
 					slog.Warn("SDK Error from split, aborting splits and failing bundle", "bundle", rb, "error", err.Error())
-					if b.BundleErr != nil {
-						b.BundleErr = err
+					// Safely set the split error if no primary bundle error has been set yet.
+					// Both SetErr and GetErr are synchronized under the same mutex, guaranteeing
+					// no memory races occur. The write-read inconsistency is explicitly guarded
+					// by the nil/null check inside SetErr(): if a concurrent primary error (e.g.,
+					// worker crash) was set first, it will not be overwritten and b.GetErr() will
+					// correctly preserve and return it instead of the secondary split error.
+					if !b.SetErr(err) {
+						slog.Debug("Error for bundle already set, logging dropped split error", "bundle", rb, "error", err)
 					}
-					return b.BundleErr
+					return b.GetErr()
 				}
 				if sr.GetChannelSplits() == nil {
 					slog.Debug("SDK returned no splits", "bundle", rb)

sdks/go/pkg/beam/runners/prism/internal/testdofns_test.go

@@ -64,6 +64,7 @@ func init() {
 	register.Function2x1(combineIntSum)
 
 	register.DoFn3x1[*sdf.LockRTracker, SourceConfig, func(int64), error]((*intRangeFn)(nil))
+	register.DoFn4x1[context.Context, *sdf.LockRTracker, SourceConfig, func(int64), error]((*slowFailSDF)(nil))
 	register.Emitter1[int64]()
 	register.Emitter2[int64, int64]()
 }
@@ -404,3 +405,36 @@ func (fn *selfCheckpointingDoFn) ProcessElement(rt *sdf.LockRTracker, _ []byte,
 		}
 	}
 }
+
+type errorSplitTracker struct {
+	*offsetrange.Tracker
+}
+
+func (t *errorSplitTracker) TrySplit(fraction float64) (any, any, error) {
+	return nil, nil, fmt.Errorf("intentional split error from tracker")
+}
+
+type slowFailSDF struct{}
+
+func (fn *slowFailSDF) CreateInitialRestriction(config SourceConfig) offsetrange.Restriction {
+	return offsetrange.Restriction{Start: 0, End: config.NumElements}
+}
+
+func (fn *slowFailSDF) SplitRestriction(config SourceConfig, rest offsetrange.Restriction) []offsetrange.Restriction {
+	return rest.EvenSplits(config.InitialSplits)
+}
+
+func (fn *slowFailSDF) RestrictionSize(_ SourceConfig, rest offsetrange.Restriction) float64 {
+	return rest.Size()
+}
+
+func (fn *slowFailSDF) CreateTracker(rest offsetrange.Restriction) *sdf.LockRTracker {
+	return sdf.NewLockRTracker(&errorSplitTracker{Tracker: offsetrange.NewTracker(rest)})
+}
+
+func (fn *slowFailSDF) ProcessElement(ctx context.Context, rt *sdf.LockRTracker, config SourceConfig, emit func(int64)) error {
+	for i := rt.GetRestriction().(offsetrange.Restriction).Start; rt.TryClaim(i); i++ {
+		<-ctx.Done()
+	}
+	return nil
+}

sdks/go/pkg/beam/runners/prism/internal/worker/bundle.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"sync"
 	"sync/atomic"
 
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
@@ -61,8 +62,19 @@ type B struct {
 	dataSema   atomic.Int32
 	OutputData engine.TentativeData
 
-	Resp      chan *fnpb.ProcessBundleResponse
-	BundleErr error
+	Resp chan *fnpb.ProcessBundleResponse
+	// DataAbort is closed when the worker responds to the bundle instruction
+	// (with success or failure), signaling ProcessOn to stop streaming data.
+	//
+	// This prevents a deadlock where a worker fails mid-bundle and stops reading
+	// from the data channel while the runner blocks indefinitely attempting to
+	// write remaining elements. Other signals are insufficient to abort immediately:
+	// - ctx.Done() only triggers on global timeouts/cancellations, which is too late.
+	// - wk.StoppedChan is only closed when tearing down the worker pool, which does
+	//   not happen while the runner is waiting on the current bundle to finish.
+	DataAbort chan struct{}
+	mu        sync.Mutex
+	bundleErr error
 	responded bool
 
 	SinkToPCollection map[string]string
@@ -80,6 +92,7 @@ func (b *B) Init() {
 		close(b.DataWait) // Can happen if there are no outputs for the bundle.
 	}
 	b.Resp = make(chan *fnpb.ProcessBundleResponse, 1)
+	b.DataAbort = make(chan struct{})
 }
 
 // DataOrTimerDone indicates a final element has been received from a Data or Timer output.
@@ -96,14 +109,40 @@ func (b *B) LogValue() slog.Value {
 		slog.String("stage", b.PBDID))
 }
 
+// SetErr sets the bundle error if it is not already set, returning true if it was set, and false otherwise.
+func (b *B) SetErr(err error) bool {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if b.bundleErr == nil {
+		b.bundleErr = err
+		return true
+	}
+	return false
+}
+
+// GetErr gets the current bundle error.
+func (b *B) GetErr() error {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.bundleErr
+}
+
 func (b *B) Respond(resp *fnpb.InstructionResponse) {
 	if b.responded {
 		slog.Warn("additional bundle response", "bundle", b, "resp", resp)
 		return
 	}
 	b.responded = true
+	if b.DataAbort != nil {
+		// Defer closing DataAbort to guarantee that the abort signal is sent
+		// when this function returns. This ensures it is always executed after
+		// any error has been safely written and synchronized via b.SetErr() or,
+		// in the happy path, after the successful response is sent to b.Resp.
+		defer close(b.DataAbort)
+	}
 	if resp.GetError() != "" {
-		b.BundleErr = fmt.Errorf("bundle %v %v failed:%v", resp.GetInstructionId(), b.PBDID, resp.GetError())
+		slog.Error("Prism received bundle error from worker response", "bundle", resp.GetInstructionId())
+		b.SetErr(fmt.Errorf("bundle %v %v failed:%v", resp.GetInstructionId(), b.PBDID, resp.GetError()))
 		close(b.Resp)
 		return
 	}
@@ -143,6 +182,13 @@ func (b *B) ProcessOn(ctx context.Context, wk *W) <-chan struct{} {
 			b.DataOrTimerDone()
 		}
 		return b.DataWait
+	case <-b.DataAbort:
+		// The bundle completed/failed before req was sent.
+		outCap := b.OutputCount + len(b.HasTimers)
+		for i := 0; i < outCap; i++ {
+			b.DataOrTimerDone()
+		}
+		return b.DataWait
 	case wk.InstReqs <- req:
 		// desired outcome
 	}
@@ -181,6 +227,9 @@ func (b *B) ProcessOn(ctx context.Context, wk *W) <-chan struct{} {
 		case <-ctx.Done():
 			b.DataOrTimerDone()
 			return b.DataWait
+		case <-b.DataAbort:
+			b.DataOrTimerDone()
+			return b.DataWait
 		case wk.DataReqs <- elms:
 		}
 	}
@@ -202,6 +251,9 @@ func (b *B) ProcessOn(ctx context.Context, wk *W) <-chan struct{} {
 	case <-ctx.Done():
 		b.DataOrTimerDone()
 		return b.DataWait
+	case <-b.DataAbort:
+		b.DataOrTimerDone()
+		return b.DataWait
 	case wk.DataReqs <- &fnpb.Elements{
 		Timers: timers,
 		Data: []*fnpb.Elements_Data{

sdks/go/pkg/beam/runners/prism/internal/worker/worker_test.go

@@ -522,6 +522,73 @@ func TestWorker_State_MultimapSideInput(t *testing.T) {
 	}
 }
 
+// TestBundle_ProcessOn_WorkerFailure verifies that the runner does not deadlock when
+// a worker fails mid-bundle and stops reading elements from the Data plane stream.
+func TestBundle_ProcessOn_WorkerFailure(t *testing.T) {
+	ctx, wk, clientConn := serveTestWorker(t)
+
+	dataCli := fnpb.NewBeamFnDataClient(clientConn)
+	dataStream, err := dataCli.Data(ctx)
+	if err != nil {
+		t.Fatal("couldn't create data client:", err)
+	}
+
+	instID := wk.NextInst()
+
+	// Create 15 large input blocks (512 KB each) to saturate the 10-slot channel buffer
+	// and the gRPC flow control window, forcing the Data sender inside worker.go to block.
+	largeBytes := make([]byte, 512*1024)
+	var inputBlocks []*engine.Block
+	for i := 0; i < 15; i++ {
+		inputBlocks = append(inputBlocks, &engine.Block{
+			Kind:  engine.BlockData,
+			Bytes: [][]byte{largeBytes},
+		})
+	}
+
+	b := &B{
+		InstID:      instID,
+		PBDID:       "teststageID",
+		Input:       inputBlocks,
+		OutputCount: 1,
+	}
+	b.Init()
+	wk.activeInstructions[instID] = b
+
+	processOnDone := make(chan struct{})
+	go func() {
+		b.ProcessOn(ctx, wk)
+		close(processOnDone)
+	}()
+
+	// Send the initial process bundle request trigger.
+	wk.InstReqs <- &fnpb.InstructionRequest{
+		InstructionId: instID,
+	}
+
+	// Read only the first block to simulate worker processing start.
+	_, err = dataStream.Recv()
+	if err != nil {
+		t.Fatal("couldn't receive first data element:", err)
+	}
+
+	// Simulate worker failure by responding with an error on the Control channel.
+	// Without the fix, ProcessOn's background goroutine deadlocks at `wk.DataReqs <- elms`
+	// because the client stopped reading and the buffer/flow-control is saturated.
+	wk.activeInstructions[instID].Respond(&fnpb.InstructionResponse{
+		InstructionId: instID,
+		Error:         "Intentional worker failure",
+	})
+
+	// Verify that ProcessOn exits cleanly and does not deadlock/hang.
+	select {
+	case <-processOnDone:
+		// Test passed: ProcessOn exited successfully!
+	case <-time.After(10 * time.Second):
+		t.Fatal("ProcessOn deadlocked / hung after worker failure!")
+	}
+}
+
 func newWorker() *W {
 	mw := &MultiplexW{
 		pool: map[string]*W{},

sdks/python/apache_beam/runners/portability/job_server.py

@@ -121,6 +121,11 @@ def start(self):
       logger = logging.getLogger(f"{self.__class__.__name__}")
       if self._log_filter is not None:
         logger.addFilter(self._log_filter)
+        # Explicitly set logger level to INFO so logger.info(...) calls in
+        # SubprocessServer._really_start_process's log_stdout pass the initial
+        # isEnabledFor check, allowing the filter to dynamically elevate log
+        # levels to WARNING or ERROR.
+        logger.setLevel(logging.INFO)
       self._server = subprocess_server.SubprocessServer(
           beam_job_api_pb2_grpc.JobServiceStub, cmd, port=port, logger=logger)
     return self._server.start()

sdks/python/apache_beam/runners/portability/prism_runner.py

@@ -154,9 +154,13 @@ def filter(self, record):
           record.msg = json_record["sdk"]["msg"]
         else:
           record.name = "PrismRunner"
-          record.msg = (
-              f"{json_record['msg']} "
-              f"({', '.join(f'{k}={v!r}' for k, v in extras.items())})")
+          formatted_extras = []
+          for k, v in extras.items():
+            if isinstance(v, str) and '\n' in v:
+              formatted_extras.append(f"\n{k}:\n{v}")
+            else:
+              formatted_extras.append(f"{k}={v!r}")
+          record.msg = f"{json_record['msg']} ({', '.join(formatted_extras)})"
       except (json.JSONDecodeError,
               KeyError,
               ValueError,

sdks/python/apache_beam/runners/portability/prism_runner_test.py

@@ -300,6 +300,17 @@ def test_after_count_trigger_streaming(self):
                   ('B-3', {10, 15, 16}),
               ])))
 
+  def test_failing_dofn(self):
+    # Prism interprets all bundle failures as RuntimeError
+    with self.assertRaisesRegex(
+        RuntimeError, "Intentional DoFn failure for prism logging test"):
+      with self.create_pipeline() as p:
+
+        def failing_fn(x):
+          raise ValueError("Intentional DoFn failure for prism logging test")
+
+        _ = p | beam.Create([1, 2, 3]) | beam.Map(failing_fn)
+
 
 class PrismJobServerTest(unittest.TestCase):
   def setUp(self) -> None: