[Bug]: Timers fire incorrectly in Prism when bundle splitting happens

[shunping]

What happened?

Running the following code can result in a prism error.

import logging

import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.utils.timestamp import Timestamp
from apache_beam.utils.timestamp import MIN_TIMESTAMP, MAX_TIMESTAMP
from apache_beam.transforms.timeutil import TimeDomain
from apache_beam.transforms.userstate import TimerSpec, on_timer


logging.basicConfig(level=logging.INFO)

options = PipelineOptions([
    "--environment_type=LOOPBACK",
    "--runner=PrismRunner",
])


data = [(0, 0), (1, 1), (2, 0)]

class MyDoFn(beam.DoFn):
  WINDOW_TIMER = TimerSpec('window_timer', TimeDomain.WATERMARK)

  def __init__(self):
    pass

  def process(
      self,
      element,
      window=beam.DoFn.WindowParam,
      window_timer=beam.DoFn.TimerParam(WINDOW_TIMER),
      ):
    print("process is called")
    window_timer.set(window.end)


  @on_timer(WINDOW_TIMER)
  def on_timer(self, key=beam.DoFn.KeyParam):
    import time
    time.sleep(1)
    yield "on_timer called on key %s" % str(key)


with beam.Pipeline(options=options) as p:
  _ = (
      p | beam.Create(data)
      | beam.ParDo(MyDoFn())
      | beam.LogElements(
          level=logging.WARNING,
          ))

Error message:

RuntimeError: Pipeline job-001 failed in state FAILED: nothing in progress and no refreshes with non zero pending elements: 2
ElementManager Now: 1754103888458 processingTimeEvents: map[] injectedBundles: []
stage-000 watermark in -inf out -inf upstream -inf from IMPULSE   pending [] byKey map[] inprogressKeys map[] byBundle map[] holds [] holdCounts map[] holdsInBundle map[] pttEvents map[] bundlesToInject []
        sideInputs: [] outputCols: [ref_PCollection_PCollection_4] outputConsumers: [stage-001] sideConsumers: []
stage-001 watermark in -inf out -inf upstream -inf from stage-000 pending [] byKey map[] inprogressKeys map[] byBundle map[] holds [] holdCounts map[] holdsInBundle map[] pttEvents map[] bundlesToInject []
        sideInputs: [] outputCols: [ref_PCollection_PCollection_5] outputConsumers: [stage-002] sideConsumers: []
stage-002 watermark in -inf out -inf upstream -inf from stage-001 pending [] byKey map[] inprogressKeys map[] byBundle map[] holds [] holdCounts map[] holdsInBundle map[] pttEvents map[] bundlesToInject []
        sideInputs: [] outputCols: [] outputConsumers: [] sideConsumers: []
stage-003 watermark in +inf out +inf upstream +inf from IMPULSE   pending [] byKey map[] inprogressKeys map[] byBundle map[] holds [] holdCounts map[] holdsInBundle map[] pttEvents map[] bundlesToInject []
        sideInputs: [] outputCols: [ref_PCollection_PCollection_1] outputConsumers: [stage-004] sideConsumers: []
stage-004 watermark in +inf out +inf upstream +inf from stage-003 pending [] byKey map[] inprogressKeys map[] byBundle map[] holds [] holdCounts map[] holdsInBundle map[] pttEvents map[] bundlesToInject []
        sideInputs: [] outputCols: [ref_PCollection_PCollection_2] outputConsumers: [stage-005] sideConsumers: []
stage-005 watermark in +inf out +inf upstream +inf from stage-004 pending [] byKey map[] inprogressKeys map[] byBundle map[] holds [] holdCounts map[] holdsInBundle map[] pttEvents map[] bundlesToInject []
        sideInputs: [] outputCols: [ref_PCollection_PCollection_8] outputConsumers: [stage-006] sideConsumers: []
stage-006 watermark in glo out glo upstream +inf from stage-005 pending [{Timer - Window [*], EventTime glo, Hold glo, "ref_AppliedPTransform_ParDo-MyDoFn-_14" "ts-window_timer" "" "\x02"} {Timer - Window [*], EventTime glo, Hold glo, "ref_AppliedPTransform_ParDo-MyDoFn-_14" "ts-window_timer" "" "\x01"}] byKey map[] inprogressKeys map[] byBundle map[] holds [] holdCounts map[] holdsInBundle map[] pttEvents map[] bundlesToInject []
        sideInputs: [] outputCols: [] outputConsumers: [] sideConsumers: []

If we comment out the time.sleep() in on_timer(), the pipeline can run successfully. Turning on the debug log level, and see there are bundle splits reported.

The same pipeline can run successfully in FnApiRunner or DataflowRunner.

---

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: Infrastructure
• Component: Spark Runner
• Component: Flink Runner
• Component: Samza Runner
• Component: Twister2 Runner
• Component: Hazelcast Jet Runner
• Component: Google Cloud Dataflow Runner

[shunping]

After discussing with @lostluck, the problem is confirmed to be related to bundle splitting in prism. If we set OneKeyPerBundle in element manager to true, the pipeline can run fine.

Originally, I thought changing the code in SplitBundle to address the change of stage state when a split happens is sufficient.

However, after applying the changes, I soon discovered some deeper problem in Python SDK.

---

Let's fast forward to the time when 3 timers are sent to a SDK worker to execute.

1. When Prism sends out 3 timers, they are sent as a block and there is a TransformId associated with it. In this case, the TransformId is from the stateful dofn. (

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

   Line 168 in 5395c9d

   TransformId: block.Transform,

   )

2. Then when Prism sees it runs slowly (due to the sleep), it sends out a split request to a SDK worker. The split request is for the input transform of that stage (

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

   Line 260 in 5395c9d

   b.InputTransformID: {

   ), e.g. stage-006-source in our above example. In the request, the runner lets SDK worker know that there are about three timers, and SDK worker determines that they can be split (without knowing they are data or timer elements) (

   beam/sdks/python/apache_beam/runners/worker/bundle_processor.py

   Line 1320 in 5395c9d

   if isinstance(op, DataInputOperation):

   ).

3. However, when the worker is processing timers, it is actually not going through DataInputOperation of "stage-006-source", but DoOperation of the DoFn (

   beam/sdks/python/apache_beam/runners/worker/bundle_processor.py

   Lines 1265 to 1272 in 5395c9d

   	if isinstance(element, beam_fn_api_pb2.Elements.Timers):
   	timer_coder_impl = (
   	self.timers_info[(
   	element.transform_id,
   	element.timer_family_id)].timer_coder_impl)
   	for timer_data in timer_coder_impl.decode_all(element.timers):
   	self.ops[element.transform_id].process_timer(
   	element.timer_family_id, timer_data)

   ). The latter doesn't know anything about splitting. Therefore, it just processes every single timer in the data channel.

[shunping]

There are three approaches to fix this:

1. 35770: fix both prism and python sdk (1a36379).
2. Don't allow bundle splitting on bundles full of timers
3. Only include one timer per bundle.

Thoughts? @lostluck

[lostluck]

There's definitely a bug in Prism's split return handling. That should be fixed anyway.

As for whether python should be fixed, i say yes, as it makes the SDK more robust to correct implementations of the model, vs just what Dataflow does.

WRT the model, timers for different keys can be sent to SDKs by a runner within the same bundle. So Prism's implementation of timer bundles is correct.

In practice, Prism is the only one that does this, as Dataflow and Flink (and probably the others) simply never mix keys in bundles. Similarly, they never split such bundles since they only have one element in them. This is also a correct implementation choice in the model. (And akin to the 2 and 3 solutions there). It's important to recognize that single element/key bundle approach is largely a performance based choice to reduce latency across independent keys.

My view on this is that the SDK has a bug that needs to be fixed, as it's not robust to a correct runner implementation.

Ideally, we run the same test pipeline implemented in Java and Go, and see what they do. If they have the same kind of bug, we should probably just change Prism's behavior then.

[shunping]

Ideally, we run the same test pipeline implemented in Java and Go, and see what they do. If they have the same kind of bug, we should probably just change Prism's behavior then.

I made a similar pipeline using go sdk.

package main

import (
	"context"
	"fmt"
	"log"
	"time"

	"github.com/apache/beam/sdks/v2/go/pkg/beam"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/timers"
	beamLog "github.com/apache/beam/sdks/v2/go/pkg/beam/log"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/x/beamx"
)

func init() {
	register.DoFn7x0[context.Context, beam.Window, beam.EventTime, timers.Provider, string, int, func(string)](&eventTimerDoFn{})
	register.Function3x0(add_key)
}

func add_key(ctx context.Context, e int, emit func(string, int)) {
	key := fmt.Sprint(e)
	// beamLog.Infof(ctx, "%d -> (%s, %d)", e, key, e+1)
	emit(key, e+1)
}

type eventTimerDoFn struct {
	Timer timers.EventTime
}

func (fn *eventTimerDoFn) ProcessElement(ctx context.Context, w beam.Window, ts beam.EventTime, tp timers.Provider, key string, value int, emitWords func(string)) {
	beamLog.Infof(ctx, "ProcessElement called: received %s,%d", key, value)

	// Set an event-time timer to the end of window
	fn.Timer.Set(tp, w.MaxTimestamp().ToTime())
}

func (fn *eventTimerDoFn) OnTimer(ctx context.Context, w beam.Window, ts beam.EventTime, tp timers.Provider, key string, timer timers.Context, emitWords func(string)) {
	switch timer.Family {
	case fn.Timer.Family:
		beamLog.Infof(ctx, "OnTimer called: for key %s", key)

		// sleep to trigger split bundle
		time.Sleep(1 * time.Second)
	}
}

func AddEventTimeDoFn(s beam.Scope, in beam.PCollection) beam.PCollection {
	return beam.ParDo(s, &eventTimerDoFn{
		Timer: timers.InEventTime("processWatermark"),
	}, in)
}

func main() {
	beam.Init()

	ctx := context.Background()
	p := beam.NewPipeline()
	s := p.Root()

	a := beam.CreateList(s, []int{1, 2, 3})
	b := beam.ParDo(s, add_key, a)
	AddEventTimeDoFn(s, b)

	if err := beamx.Run(ctx, p); err != nil {
		log.Fatalf("Failed to execute job: %v", err)
	}
}

I ran it before my PR change in #35770, and I was able to see a similar error as mentioned in #35771 (comment): two timers in pending in the last stage.

stage-000 watermark in +inf out +inf upstream +inf from IMPULSE   pending [] byKey map[] inprogressKeys map[] byBundle map[] holds [] holdCounts map[] holdsInBundle map[] pttEvents map[] bundlesToInject []
        sideInputs: [] outputCols: [n1] outputConsumers: [stage-001] sideConsumers: []
stage-001 watermark in +inf out +inf upstream +inf from stage-000 pending [] byKey map[] inprogressKeys map[] byBundle map[] holds [] holdCounts map[] holdsInBundle map[] pttEvents map[] bundlesToInject []
        sideInputs: [] outputCols: [n3] outputConsumers: [stage-002] sideConsumers: []
stage-002 watermark in glo out glo upstream +inf from stage-001 pending [{Timer - Window [*], EventTime glo, Hold glo, "e4" "processWatermark" "" "\x011"} {Timer - Window [*], EventTime glo, Hold glo, "e4" "processWatermark" "" "\x013"}] byKey map[] inprogressKeys map[] byBundle map[] holds [] holdCounts map[] holdsInBundle map[] pttEvents map[] bundlesToInject []
        sideInputs: [] outputCols: [] outputConsumers: [] sideConsumers: []

When I ran the above pipeline against #35770, which includes the fix of putting split residual to the correct place, I see the timers are fired more than I expected (similar to the python pipeline prior to Python SDK change).

2025/08/02 22:00:04 INFO log from SDK worker worker.ID=job-001[go-job-1-1754186404643657000]_go worker.endpoint=localhost:62176 sdk.transformID=e4 sdk.location=/Users/shunping/Projects/beam-dev-go-prism/sdks/go/examples/periodic_bug.go:41 sdk.time=2025-08-03T02:00:04.654Z sdk.msg="OnTimer called: for key 3"
2025/08/02 22:00:04 INFO log from SDK worker worker.ID=job-001[go-job-1-1754186404643657000]_go worker.endpoint=localhost:62176 sdk.transformID=e4 sdk.location=/Users/shunping/Projects/beam-dev-go-prism/sdks/go/examples/periodic_bug.go:41 sdk.time=2025-08-03T02:00:04.855Z sdk.msg="OnTimer called: for key 2"
2025/08/02 22:00:05 INFO log from SDK worker worker.ID=job-001[go-job-1-1754186404643657000]_go worker.endpoint=localhost:62176 sdk.transformID=e4 sdk.location=/Users/shunping/Projects/beam-dev-go-prism/sdks/go/examples/periodic_bug.go:41 sdk.time=2025-08-03T02:00:05.258Z sdk.msg="OnTimer called: for key 1"
2025/08/02 22:00:05 INFO log from SDK worker worker.ID=job-001[go-job-1-1754186404643657000]_go worker.endpoint=localhost:62176 sdk.transformID=e4 sdk.location=/Users/shunping/Projects/beam-dev-go-prism/sdks/go/examples/periodic_bug.go:41 sdk.time=2025-08-03T02:00:05.655Z sdk.msg="OnTimer called: for key 1"
2025/08/02 22:00:06 INFO log from SDK worker worker.ID=job-001[go-job-1-1754186404643657000]_go worker.endpoint=localhost:62176 sdk.transformID=e4 sdk.location=/Users/shunping/Projects/beam-dev-go-prism/sdks/go/examples/periodic_bug.go:41 sdk.time=2025-08-03T02:00:06.655Z sdk.msg="OnTimer called: for key 2"

I think a bundle with multiple timers is valid, but we probably should not try to split it ( Approach 2 in #35771 (comment))

[shunping]

Basically the current way of sending split request does not work for bundles full of timers.

beam/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go

Lines 1019 to 1024 in 1386408

	// ReturnResiduals is called after a successful split, so the remaining work
	// can be re-assigned to a new bundle.
	func (em *ElementManager) ReturnResiduals(rb RunBundle, firstRsIndex int, inputInfo PColInfo, residuals Residuals) {
	stage := em.stages[rb.StageID]
	stage.splitBundle(rb, firstRsIndex)

That's because, although timer elements are sent through the same data channel as data elements, the split handling in SDK worker does not consider them the same as data elements which are splittable.

I think one way to fix that is to change estimatedInputElements.

https://github.com/apache/beam/blob/138640827e00d6f2e139dfcc38a41ab21d4bb79a/sdks/go/pkg/beam/runners/prism/internal/worker/bundle.go#L254-L274

We send the estimatedInputElements to the SDK worker for split request. Currently, we include both the number of timers and data elements in that number.

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

Lines 144 to 152 in 1386408

	input, estimatedElements := em.DataAndTimerInputForBundle(rb, s.inputInfo)
	b = &worker.B{
	PBDID: s.ID,
	InstID: rb.BundleID,
	InputTransformID: s.inputTransformID,
	Input: input,
	EstimatedInputElements: estimatedElements,

Maybe we don't count the number of timers there? (See my latest change at #35770)

Also a quote from the offline discussion with @lostluck:

since Timers don't target the source transform, they target the actual callback transform. The source transforms are what split

[lostluck]

This is ultimately an edge case where the model and reality don't mix well. The model also allows a transform to have multiple sources generally, but in practice the switching/ handling of those sources is also never implemented by any sdk. Great find.