Repository navigation
[Bug]: Unpersist cached Datasets after evaluation in Spark Structured Streaming runner - #40265
akashchamp wants to merge 2 commits into
Conversation
|
Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment |
|
Assigning reviewers: R: @chamikaramj 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). |
3393a1c to
8f2f937
Compare
|
Reminder, please take a look at this pr: @chamikaramj |
|
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment R: @chamikaramj added as fallback since no labels match configuration Available commands:
|
3351a92 to
b20bf04
Compare
|
Reminder, please take a look at this pr: @chamikaramj |
|
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment R: @kennknowles added as fallback since no labels match configuration Available commands:
|
kennknowles
left a comment
kennknowles
left a comment
There was a problem hiding this comment.
Seems generally good. We'll want to make sure everything is green. We can also run extra tests by touching the files for SparkRunner in ~/.github/trigger_files/
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #40265 +/- ##
============================================
+ Coverage 56.16% 58.73% +2.56%
- Complexity 2288 13920 +11632
============================================
Files 1124 2585 +1461
Lines 178336 271705 +93369
Branches 1489 11243 +9754
============================================
+ Hits 100170 159589 +59419
- Misses 75645 106112 +30467
- Partials 2521 6004 +3483
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:
|
b20bf04 to
154c9b1
Compare
…ing runner ParDoTranslatorBatch persists an intermediate, wide-row Dataset when a multi-output ParDo needs to be divided into its per-tag Datasets, to avoid re-evaluating the parent mapPartitions transformation once per tag. It never called unpersist() on that cached Dataset, which the code even acknowledged with a "FIXME What's the strategy to unpersist Datasets / RDDs?" comment. Every multi-output ParDo therefore leaked its cached Dataset into Spark's storage memory for the lifetime of the SparkSession, causing unbounded growth, disk spilling and OOMs on large batch pipelines. Add PipelineTranslator.TranslationState#cacheDataset(Dataset, StorageLevel), a helper translators use instead of calling Dataset#persist directly when caching a dataset outside the normal putDataset path. Track every dataset cached this way and thread that collection through to EvaluationContext, which now unpersists all of them in a finally block once leaf evaluation has finished, whether it completed normally, was stopped early, or failed with an exception. ParDoTranslatorBatch now goes through cacheDataset instead of calling persist() directly, removing the FIXME. Fixes apache#40243
runners/spark/4 keeps its own source-override copies of PipelineTranslatorStreaming and StreamingEvaluationContext because the Spark 4 streaming runner needs a different createEvaluationContext / EvaluationContext than the shared batch translator. The previous commit added a cachedDatasets parameter to PipelineTranslator.TranslationState#createEvaluationContext and to the EvaluationContext constructor, but didn't update these two Spark 4 overrides, which broke compilation of :runners:spark:4. Thread cachedDatasets through both override files the same way the base PipelineTranslator and EvaluationContext now do. Since StreamingEvaluationContext#evaluate() fully replaces the base evaluate() (it drives streaming queries instead of a per-leaf loop) rather than calling super.evaluate(), also unpersist the cached datasets once all queries have terminated, in its own finally block, so Datasets cached via TranslationState#cacheDataset (used by the shared, stateless-ParDo translator this streaming runner reuses) don't leak for streaming pipelines either. Expose EvaluationContext#unpersistCachedDatasets as protected so the subclass can call it. Also reflow two doc comments in EvaluationContext.java and PipelineTranslator.java to satisfy Spotless line-wrapping.
154c9b1 to
791fb5a
Compare
Fixes #40243
Problem
In the Spark Structured Streaming runner,
ParDoTranslatorBatchexplicitlypersist()s the intermediate, wide-rowDatasetit builds for a multi-outputParDo(so the per-mapPartitionstransformation backing every tagged outputisn't re-evaluated once per tag). The code even carried a
FIXMEacknowledgingthe gap:
// FIXME What's the strategy to unpersist Datasets / RDDs?No corresponding
unpersist()was ever called, so every multi-outputParDoin a pipeline leaves its cached
Datasetin Spark's storage memory for thelifetime of the
SparkSession. On a large batch pipeline with several suchtransforms this accumulates without bound and leads to storage-memory
exhaustion, disk spilling, and OOMs.
Fix
PipelineTranslator.TranslationState#cacheDataset(Dataset, StorageLevel),a small helper translators use instead of calling
Dataset#persistdirectlywhen they need to cache a dataset that isn't going through the normal
putDatasetpath. It persists the dataset and registers it for latercleanup.
PipelineTranslator.TranslatingVisitornow tracks every dataset cached thisway and threads that collection through to
EvaluationContext.EvaluationContext#evaluate()now unpersists all such datasets in afinallyblock once evaluation of the leaves has finished — whether itcompleted normally, was stopped early, or an evaluation failed with an
exception — so a cached dataset never outlives the pipeline run that created
it.
ParDoTranslatorBatchnow callscxt.cacheDataset(...)instead of callingpersist()directly on the wide-rowDataset, removing theFIXME.Testing
EvaluationContextTest: added two new tests asserting that datasetsregistered via the new cache-tracking collection are
unpersist()-ed onceevaluate()finishes, both on the happy path and when evaluating a leafthrows.
ParDoTest: updated the two existing tests that assert a multi-outputParDoproduces a cached RDD/Dataset (testPardoWithOutputTagsCachedRDD/...CachedDataset, renamed to...UnpersistsRDDAfterRun/...UnpersistsDatasetAfterRun) so that, after the pipeline run completes,they assert the opposite:
SESSION.hasCachedData()is nowfalse, i.e. thedataset was unpersisted rather than left cached.
ParDoTestregression tests(with their new "unpersisted after run" assertions) against the unfixed
production code — both fail with
AssertionError: Expected cached data to be unpersisted after run,confirming
SESSION.hasCachedData()is stilltrueafter the pipeline runcompletes, i.e. the
Datasetreally is left cached.EvaluationContextTestandParDoTestsuites via
./gradlew :runners:spark:3:test --tests ...EvaluationContextTest --tests ...batch.ParDoTest: 14 tests, 14 passed,0 failed. I did not run the
runners:spark:4variant locally (shares thesame
src/sources viarunners/spark/4); CI covers it.Notes
This only changes the Spark Structured Streaming runner's translation-time
bookkeeping; it does not change what gets computed or written, only when the
intermediate cached
Datasetused to fan out a multi-outputParDointo itsper-tag
Datasets is released.