Skip to content

[Prism] Schedule consumers of a self checkpointing source #39572

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Eliaaazzz wants to merge 1 commit into
base: master
Choose a base branch
from
+377 −13
Open

[Prism] Schedule consumers of a self checkpointing source #39572

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 78 additions & 13 deletions sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -235,6 +235,10 @@ type ElementManager struct {
livePending atomic.Int64 // An accessible live pending count. DEBUG USE ONLY
pendingElements sync.WaitGroup // pendingElements counts all unprocessed elements in a job. Jobs with no pending elements terminate successfully.

// Latched once any bundle returns a residual, after which a stage watermark
// can be pinned indefinitely and can't be relied on to schedule consumers.
sawResidual atomic.Bool

processTimeEvents *stageRefreshQueue // Manages sequence of stage updates when interfacing with processing time. Callers must hold refreshCond.L lock.
testStreamHandler *testStreamHandler // Optional test stream handler when a test stream is in the pipeline.
}
Expand Down Expand Up @@ -841,6 +845,17 @@ func reElementResiduals(residuals []Residual, inputInfo PColInfo, rb RunBundle)
// input elements, and the committed output elements.
func (em *ElementManager) PersistBundle(rb RunBundle, col2Coders map[string]PColInfo, d TentativeData, inputInfo PColInfo, residuals Residuals) {
stage := em.stages[rb.StageID]
// Consumers that received data from this bundle, recorded so they can still
// be scheduled when this stage's output watermark is held back. Only needed
// once something self checkpoints, so pipelines that never do keep their
// previous bundle scheduling exactly.
if len(residuals.Data) > 0 {
em.sawResidual.Store(true)
}
var changedConsumers set[string]
if em.sawResidual.Load() {
changedConsumers = set[string]{}
}
var seq int
for output, data := range d.Raw {
info := col2Coders[output]
Expand Down Expand Up @@ -910,6 +925,9 @@ func (em *ElementManager) PersistBundle(rb RunBundle, col2Coders map[string]PCol
count = consumer.AddPending(em, newPending)
}
em.addPending(count)
if changedConsumers != nil && count > 0 {
changedConsumers.insert(sID)
}
}
for _, link := range sideConsumers {
consumer := em.stages[link.Global]
Expand Down Expand Up @@ -942,6 +960,12 @@ func (em *ElementManager) PersistBundle(rb RunBundle, col2Coders map[string]PCol
// even if a panic occurs during `em.addPending`. This prevents potential deadlocks
// if the waitgroup unexpectedly drops below zero due to a runner bug.
defer stage.mu.Unlock()
if len(changedConsumers) > 0 {
if stage.consumersWithNewData == nil {
stage.consumersWithNewData = set[string]{}
}
stage.consumersWithNewData.merge(changedConsumers)
}
completed := stage.inprogress[rb.BundleID]
em.addPending(-len(completed.es))
delete(stage.inprogress, rb.BundleID)
Expand Down Expand Up @@ -1101,6 +1125,9 @@ func (em *ElementManager) ReturnResiduals(rb RunBundle, firstRsIndex int, inputI
stage := em.stages[rb.StageID]

stage.splitBundle(rb, firstRsIndex, em)
if len(residuals.Data) > 0 {
em.sawResidual.Store(true)
}
unprocessedElements := reElementResiduals(residuals.Data, inputInfo, rb)
if len(unprocessedElements) > 0 {
slog.Debug("ReturnResiduals: unprocessed elements", "bundle", rb, "count", len(unprocessedElements))
Expand Down Expand Up @@ -1211,6 +1238,9 @@ type stageState struct {
estimatedOutput mtime.Time // Estimated watermark output from DoFns
previousInput mtime.Time // input watermark before the latest watermark refresh

// Consumers handed data by a bundle, pending delivery to the scheduler.
consumersWithNewData set[string]

pending elementHeap // pending input elements for this stage that are to be processesd
inprogress map[string]elements // inprogress elements by active bundles, keyed by bundle
sideInputs map[LinkID]map[typex.Window][][]byte // side input data for this stage, from {tid, inputID} -> window
Expand Down Expand Up @@ -1822,12 +1852,6 @@ keysPerBundle:
if ss.inprogressKeys.present(k) {
continue
}
newKeys.insert(k)
// Track the min-timestamp for later watermark handling.
if dnt.elements[0].timestamp < minTs {
minTs = dnt.elements[0].timestamp
}

dataInBundle := false

var toProcessForKey []element
Expand Down Expand Up @@ -1874,22 +1898,51 @@ keysPerBundle:
break
}
}
if len(toProcessForKey) > 0 {
newKeys.insert(k)
// Track the min-timestamp for later watermark handling. Elements pop
// in timestamp order, so the first selected one is the earliest.
if ts := toProcessForKey[0].timestamp; ts < minTs {
minTs = ts
}
}
toProcess = append(toProcess, toProcessForKey...)

if dnt.elements.Len() == 0 {
delete(ss.pendingByKeys, k)
}
if OneKeyPerBundle {
// A key that yielded nothing, such as one headed by a timer above the
// watermark, must not consume the single key slot, or the bundle is empty
// and the stage keeps rescheduling on it.
if OneKeyPerBundle && len(toProcessForKey) > 0 {
break keysPerBundle
}
}

// If we're out of data, and timers were not cleared then the watermark is accurate.
stillSchedulable := !(len(ss.pendingByKeys) == 0 && !timerCleared)
// Reschedule only when a later bundle could build something, or a cleared
// timer may have held back the minimum pending timestamp.
stillSchedulable := timerCleared || ss.hasBuildableDataLocked(watermark)

return toProcess, minTs, newKeys, holdsInBundle, nil, stillSchedulable, 0
}

// hasBuildableDataLocked reports whether a key that isn't in progress heads its
// heap with data, or with a timer the watermark has reached. Callers hold ss.mu.
func (ss *stageState) hasBuildableDataLocked(watermark mtime.Time) bool {
for k, dnt := range ss.pendingByKeys {
if ss.inprogressKeys.present(k) {
continue
}
if dnt.elements.Len() == 0 {
continue
}
if e := dnt.elements[0]; e.IsData() || e.timestamp <= watermark {
return true
}
}
return false
}

// buildEventTimeBundle for aggregation stages, processes all elements that are within the watermark for completed windows.
func (*aggregateStageKind) buildEventTimeBundle(ss *stageState, watermark mtime.Time) (toProcess elementHeap, _ mtime.Time, _ set[string], _ map[mtime.Time]int, panesInBundle []bundlePane, schedulable bool, pendingAdjustment int) {
minTs := mtime.MaxTimestamp
Expand Down Expand Up @@ -2286,9 +2339,13 @@ func (ss *stageState) updateWatermarks(em *ElementManager) set[string] {
if minWatermarkHold < newOut {
newOut = minWatermarkHold
}
// If the newOut is smaller, then don't change downstream watermarks.
// If the newOut is smaller, then don't change downstream watermarks. Any
// consumer that received data still needs scheduling, since an unadvancing
// watermark is otherwise the only thing that would surface it.
if newOut <= ss.output {
return nil
refreshes := ss.consumersWithNewData
ss.consumersWithNewData = nil
return refreshes
}

// If bigger, advance the output watermark
Expand Down Expand Up @@ -2328,6 +2385,7 @@ func (ss *stageState) updateWatermarks(em *ElementManager) set[string] {

// Update this stage's output watermark, and then propagate that to downstream stages
refreshes := set[string]{}
ss.consumersWithNewData = nil
ss.output = newOut
for _, outputCol := range ss.outputIDs {
consumers := em.consumers[outputCol]
Expand Down Expand Up @@ -2435,13 +2493,20 @@ func (ss *stageState) bundleReady(em *ElementManager, emNow mtime.Time) (mtime.T
previousInputW := ss.previousInput

_, isOrdinaryStage := ss.kind.(*ordinaryStageKind)
if isOrdinaryStage && len(ss.sides) == 0 {
_, isStatefulStage := ss.kind.(*statefulStageKind)
switch {
case isOrdinaryStage && len(ss.sides) == 0:
// For ordinary stage with no side inputs, we use whether there are pending elements to determine
// whether a bundle is ready or not.
if len(ss.pending) == 0 {
return mtime.MinTimestamp, false, ptimeEventsReady, injectedReady
}
} else if inputW == upstreamW && previousInputW == inputW {
case isStatefulStage && len(ss.sides) == 0 && em.sawResidual.Load() && ss.hasBuildableDataLocked(upstreamW):
// A stateful stage processes pending data at whatever the current
// watermark is, so data alone makes it ready once something is self
// checkpointing. Side input readiness comes from the watermark, so
// stages that read one keep waiting.
case inputW == upstreamW && previousInputW == inputW:
// Otherwise, use the progression of watermark to determine the bundle readiness.
slog.Debug("bundleReady: unchanged upstream watermark",
slog.String("stage", ss.ID),
Expand Down
Loading
Loading