Repository navigation
Conversation
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
📋 Review Summary
This PR enables reporting of bytes processed for Splittable DoFns (SplittableProcessFn) in the Dataflow streaming worker so that autoscaling can accurately estimate backlog clearance time alongside SDF backlog reporting. The implementation cleanly propagates the output coder to ParDoOperation, attaches the byte counter in ComputationWorkExecutorFactory, and aggregates the processed bytes in StreamingModeExecutionContext with thorough unit test coverage.
🔍 General Feedback
- The changes are well-structured, adhere to existing worker conventions, and include comprehensive unit tests across
IntrinsicMapTaskExecutorFactoryTest,MapTaskExecutorTest,ComputationWorkExecutorFactoryTest, andStreamingModeExecutionContextTest. - Calling
getAndReset()on the sharedCounterwhen multiple SDFParDoOperations exist in the sameMapTaskworks safely (since the first operation drains the shared counter and subsequent ones return0), and the added comment inStreamingModeExecutionContexthelpful documents the existing interaction withStreamingWorkerStatusReporter.
💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #40639 +/- ##
=============================================
+ Coverage 56.16% 76.19% +20.02%
- Complexity 2288 2290 +2
=============================================
Files 1124 260 -864
Lines 178336 14411 -163925
Branches 1489 1489
=============================================
- Hits 100170 10980 -89190
+ Misses 75645 2736 -72909
+ Partials 2521 695 -1826
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Assigning reviewers: R: @kennknowles added as fallback since no labels match configuration Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
| readNode.getParallelInstruction().getOriginalName(), | ||
| readNode.getParallelInstruction().getSystemName(), | ||
| readNode.getParallelInstruction().getName()); | ||
| operation.receivers[0].addOutputCounter( |
There was a problem hiding this comment.
why track just the first receiver? If a operator is multioutput are we missing tracking some of it?
There was a problem hiding this comment.
I'm not entirely sure. I didn't actually change this code, just added some extra null checks. I guess as long as the SplittableFn is not outputting the same element across multiple outputs we don't risk double-counting. I think that might more closely match what the portable worker does.
| new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>(readCoder), | ||
| mapTaskExecutor.getOutputCounters(), | ||
| nameContext) | ||
| .setSamplingPeriod(100) |
There was a problem hiding this comment.
do we want to sample this often? should it be time-based?
There was a problem hiding this comment.
I'm happy to change it if we can agree on a better time. This is just what we currently do for custom source reporting already.
| // An SDF's SplittableProcessFn behaves like a source, so we need to report its bytes read even | ||
| // though it's not a read operator for autoscaling to work properly. | ||
| Coder<?> bytesReadCoder = null; | ||
| if ("SplittableProcessFn".equals(cloudUserFn.getClassName())) { |
There was a problem hiding this comment.
this seems brittle. I'm not very familiar with this code but perhaps Kenn will have a cleaner way of doing this.
There was a problem hiding this comment.
My first thought was to add something more generically to ParDo, but I feel that this is more of a Dataflow implementation detail that we even need to do this, so I'd rather not add something to the generic API to specify which ParDo's should be reporting their bytes to autoscaling. Maybe there's some way to add something to the CloudObject in SplittableProcessFnFactory.java?
There was a problem hiding this comment.
Here were some alternatives provided by AI, though it's first recommendation is what I didn't really want to do:
- Add a capability method to ParDoFn (Recommended)
Add default boolean isSplittable() { return false; } (or shouldReportBytesRead()) on ParDoFn (and delegate in ForwardingParDoFn).
Determine it via the actual DoFn class rather than any string: SplittableProcessFnFactory.ProcessFnExtractor wraps the user DoFn in SplittableParDoViaKeyedWorkItems.ProcessFn, which is already accessible at construction time via helpers.doFnInstanceManager.peek().getDoFn() instanceof ProcessFn inside SimpleParDoFnHelpers (used by both SimpleParDoFn and StreamingKeyedWorkItemSideInputParDoFn).
Alternatively, SplittableProcessFnFactory.createDefault() (or UserParDoFnFactory.streamingKeyedWorkItem, which is already true exclusively for SplittableProcessFnFactory) can set/wrap that flag on the returned ParDoFn.
IntrinsicMapTaskExecutorFactory then simply checks if (fn.isSplittable()), keeping all CloudObject dispatch encapsulated inside DefaultParDoFnFactory.
- Have SplittableProcessFnFactory wrap with a dedicated ForwardingParDoFn subclass (or marker interface)
If you don't want to touch the public ParDoFn interface, SplittableProcessFnFactory can wrap the ParDoFn produced by UserParDoFnFactory in a package-private SplittableParDoFn extends ForwardingParDoFn (or marker interface), allowing fn instanceof SplittableParDoFn in IntrinsicMapTaskExecutorFactory.
- Share a constant if keeping the CloudObject check
If you prefer keeping the check on cloudUserFn.getClassName(), extracting a shared constant (e.g., SplittableProcessFnFactory.OBJECT_TYPE_NAME = "SplittableProcessFn") used by DefaultParDoFnFactory, IntrinsicMapTaskExecutorFactory, and tests avoids duplicating the string literal (similar to ComputationWorkExecutorFactory.isCustomSource), though Option 1 or 2 is cleaner since it delegates the classification to ParDoFnFactory.
There was a problem hiding this comment.
Is there a way to get the DoFnSignature and examine it's processElement.isSplittable()?
That seems to be done elsewhere to determine if something is a splittable dofn.
There was a problem hiding this comment.
I think I came up with a solution. It involves getting the (already deserialized) DoFnInfo from ParDoFn and then checking its signature as you suggested.
|
R: @scwhittle |
|
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment |
| return (ReadOperation) readOperation; | ||
| } | ||
|
|
||
| public List<ParDoOperation> getReadReportingParDoOperations() { |
There was a problem hiding this comment.
would it be better to have a more generic method to get all the pardo operations or all operations and do the specific filtering to the bytesread ones done at the caller? seems like that would keep down on method proliferation
| // An SDF's SplittableProcessFn behaves like a source, so we need to report its bytes read even | ||
| // though it's not a read operator for autoscaling to work properly. | ||
| Coder<?> bytesReadCoder = null; | ||
| if ("SplittableProcessFn".equals(cloudUserFn.getClassName())) { |
There was a problem hiding this comment.
Is there a way to get the DoFnSignature and examine it's processElement.isSplittable()?
That seems to be done elsewhere to determine if something is a splittable dofn.
| // windows. | ||
| helpers.registerStateCleanup( | ||
| (WindowingStrategy<?, W>) getDoFnInfo().getWindowingStrategy(), | ||
| (WindowingStrategy<?, W>) helpers.fnInfo.getWindowingStrategy(), |
There was a problem hiding this comment.
Note I changed these to call helpers.fnInfo directly since getDoFnInfo() now potentially throws Exception (from helpers.doFnInstanceManager.peek()). There were already instances of calling helpers.fnInfo in this method so this seemed cleaner than adding more exception throwing.
There was a problem hiding this comment.
maybe we should rename helpers fnInfo to activeFnInfo or something indicating it is the fninfo on the active thread?
See comment below
| void abort() throws Exception; | ||
|
|
||
| default @Nullable DoFnInfo<?, ?> getDoFnInfo() throws Exception { | ||
| return null; |
There was a problem hiding this comment.
perhaps it would be better to remove the default so that new classes (and perhaps other existing ones) return an appropriate dofninfo?
| @VisibleForTesting | ||
| @Nullable DoFnInfo<?, ?> getDoFnInfo() { | ||
| return helpers.fnInfo; | ||
| /** Returns the {@link DoFnInfo} used by this {@link StreamingKeyedWorkItemSideInputParDoFn}. */ |
There was a problem hiding this comment.
hmm, it appears these DoFnInfo are stateful in that they contain a DoFn which goes through the dofn lifecycle. previously we would return only the active dofninfo (or null if not processing). Now we always return one but it might not be an active one or one that is then activated in another thread.
We don't want callers of getDoFnInfo to call methods that expect to just be used from the active thread (like dofn lifecycle etc). It seems like it would be cleanest if all the immutable stuff within DoFnInfo was a different interface which was returned so that getDoFnInfo couldn't be misused.
| // windows. | ||
| helpers.registerStateCleanup( | ||
| (WindowingStrategy<?, W>) getDoFnInfo().getWindowingStrategy(), | ||
| (WindowingStrategy<?, W>) helpers.fnInfo.getWindowingStrategy(), |
There was a problem hiding this comment.
maybe we should rename helpers fnInfo to activeFnInfo or something indicating it is the fninfo on the active thread?
See comment below
SDF's SplittableProcessFn's report their backlog for autoscaling, so we need to also report their bytes processed for autoscaling to correctly compute time to clear backlog.
SplittableProcessFn's are tagged when constructing MapTaskExecutors and their output coders are attached to the ParDoOperation so the counter values can be updated while processing.
Also adds tests for the new code.