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.12% +19.95%
- Complexity 2288 2293 +5
=============================================
Files 1124 259 -865
Lines 178336 14440 -163896
Branches 1489 1492 +3
=============================================
- Hits 100170 10992 -89178
+ Misses 75645 2752 -72893
+ Partials 2521 696 -1825
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?
There was a problem hiding this comment.
There are 11 other production and 3 test classes that implement ParDoFn but don't actually wrap a DoFn:
- AssignWindowsParDoFn
- CreateIsmShardKeyAndSortKeyParDoFn
- GroupAlsoByWindowsParDoFn
- PairWithConstantKeyParDoFn
- BatchSideInputPGBKParDoFn
- StreamingSideInputPGBKParDoFn
- ReifyTimestampAndWindowsParDoFn
- SimplePartialGroupByKeyParDoFn
- StreamingPCollectionViewWriterParDoFn
- ToIsmRecordForMultimapParDoFn
- ValuesParDoFn
- NoopParDoFn in IntrinsicMapTaskExecutorTest.java and MapTaskExecutorTest.java
- TestParDoFn in ParDoOperationTest.java
There is one more class, BatchModeUngroupingParDoFn, that wraps ParDoFn, so we should probably have it's getDoFnInfo delegate to the underlying ParDoFn.
For all the others though, I think we'd just have to return null. I'll go ahead and remove the default method since I think it's nicer to at least force new classes to think about what to return.
| @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.
There was a problem hiding this comment.
Yeah, I think I see what you're saying. One option would be to have separate methods for all of the things on the underlying DoFn that are immutable, e.g. WindowingStrategy. But that seems like a lot of extra boilerplate. Also, for the DoFnSignature stuff, I need an instance of the DoFn object itself.
So maybe we go with a combination of what you've suggested: change helpers.fnInfo to helpers.activeFnInfo to make it more clear that this is a dynamic object. And then have getDoFnInfo and getActiveDoFnInfo with better comments. getActiveDoFnInfo doesn't need to be part of the generic ParDoFn interface.
| // 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
… and renaming helpers.doFnInfo to activeDoFnInfo.
github-actions
Bot
added
size/xl
|
The code wasn't handling multi-output SDF's properly since it assumed we would be using the same coder for each output. So I fixed that and added some new tests. |
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.