Use a stub to store Spark StageInfo #1525
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: Ahmed Hussein (amahussein) <[email protected]> Fixes NVIDIA#1524 This commit uses a smaller class `StageInfoStub` to store Spark's StageInfo. This class is common between all the spark implementations but it has more fields to the constructor across different versions. Currently we only use a subset of the class fields. The remaining fields represent an overhead or redundant storage; especially when it comes to store the accumulables and taskMetrics for each stage. To evaluate the memory optimization, a new `Checkpoint` mechanism was added to allow gathering information at separate stages of the execution. The `checkpoint` design and implementation can be further improved and extended to build a performance profile to compare different tradeoffs.
amahussein
added
core_tools
parthosa
reviewed
Feb 5, 2025
There was a problem hiding this comment.
Thanks @amahussein. A minor typo. This framework for logging memory used looks great.
core/src/main/scala/org/apache/spark/rapids/tool/benchmarks/DevRuntimeCheckpoint.scala
Outdated
Show resolved
Hide resolved
Signed-off-by: Ahmed Hussein (amahussein) <[email protected]>
amahussein
commented
Feb 5, 2025
core/src/main/scala/org/apache/spark/rapids/tool/benchmarks/DevRuntimeCheckpoint.scala
Outdated
Show resolved
Hide resolved
parthosa
approved these changes
Feb 5, 2025
| private def initStageInfo(newStageInfo: StageInfo): StageInfo = { | ||
| newStageInfo | ||
| private def initStageInfo(newStageInfo: StageInfo): StageInfoStub = { | ||
| StageInfoStub.fromStageInfo(newStageInfo) |
There was a problem hiding this comment.
@amahussein Currently we are reassigning the StageInfo object and updating the StageModel class with the incoming StageInfo object.
Now that we are using a Stub and creating a new object, can we not use the existing Stub object in case of updates to StageModel and just update its variables. Currently we are doing a new Stub allocation in all the cases.
There was a problem hiding this comment.
mmm, yeah it is possible we do only update to some fields that get changed.
The idea that:
- This update should only happens a single time when the stage is completed. This implies that this is not very frequent event.
- Updating some fields could lead to bugs. When we extend this object in the future, the dev will have to make sure that they are handling the fields correctly (which one could be updated vs which one are not).
- allocating the new object in that case made the code look easier especially to maintain moving fwd.
amahussein
commented
Feb 5, 2025
| private def initStageInfo(newStageInfo: StageInfo): StageInfo = { | ||
| newStageInfo | ||
| private def initStageInfo(newStageInfo: StageInfo): StageInfoStub = { | ||
| StageInfoStub.fromStageInfo(newStageInfo) |
There was a problem hiding this comment.
mmm, yeah it is possible we do only update to some fields that get changed.
The idea that:
- This update should only happens a single time when the stage is completed. This implies that this is not very frequent event.
- Updating some fields could lead to bugs. When we extend this object in the future, the dev will have to make sure that they are handling the fields correctly (which one could be updated vs which one are not).
- allocating the new object in that case made the code look easier especially to maintain moving fwd.
cindyyuanjiang
approved these changes
Feb 6, 2025
| } else { // loads the noOp implementation by default | ||
| new NoOpRuntimeCheckpoint | ||
| } | ||
| } |
There was a problem hiding this comment.
nit: new line between 2 defs
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Signed-off-by: Ahmed Hussein (amahussein) [email protected]
Fixes #1524
This commit uses a smaller class
StageInfoStubto store Spark's StageInfo. This class is common between all the spark implementations but it has more fields to the constructor across different versions.Currently we only use a subset of the class fields. The remaining fields represent an overhead or redundant storage; especially when it comes to store the accumulables and taskMetrics for each stage.
To evaluate the memory optimization, a new
Checkpointmechanism was added to allow gathering information at separate stages of the execution.The
checkpointdesign and implementation can be further improved and extended to build a performance profile to compare different tradeoffs.This pull request introduces a new runtime checkpointing feature for performance benchmarking and debugging in the RAPIDS tools. The main changes include adding new classes for runtime checkpoints, updating build properties, and modifying existing classes to integrate the new checkpointing functionality.
Memory evaluation:
benchmarks.checkpointstodev. This is achieved by changing the pom.xml file or passing the property as an arhument to the mvn command.Code Changes
Main core changes
core/src/main/scala/org/apache/spark/sql/rapids/tool/util/stubs/StageInfoStub.scala: Added a new classStageInfoStubto provide a consistent interface forStageInfoacross different Spark versions.core/src/main/scala/org/apache/spark/sql/rapids/tool/store/StageModel.scala: UpdatedStageModelto useStageInfoStubfor compatibility with different Spark versions. [1] [2] [3]Build and Configuration Updates:
core/pom.xml: Added a new propertybenchmarks.checkpointsto manage the checkpointing feature.core/src/main/resources/configs/build.properties: Updated build properties to includebenchmarks.checkpoints. [1] [2]Runtime Checkpointing Feature:
core/src/main/scala/org/apache/spark/rapids/tool/benchmarks/DevRuntimeCheckpoint.scala: Added a new classDevRuntimeCheckpointto insert memory markers and print memory information during runtime for performance metrics.core/src/main/scala/org/apache/spark/rapids/tool/benchmarks/NoOpRuntimeCheckpoint.scala: Added a new classNoOpRuntimeCheckpointas a default no-operation implementation for checkpoints.core/src/main/scala/org/apache/spark/rapids/tool/benchmarks/RuntimeCheckpointTrait.scala: Introduced a new traitRuntimeCheckpointTraitdefining the API for inserting runtime checkpoints.core/src/main/scala/org/apache/spark/rapids/tool/benchmarks/RuntimeInjector.scala: Added a new objectRuntimeInjectorto manage and insert runtime checkpoints based on build properties.Integration of Checkpoints:
core/src/main/scala/org/apache/spark/sql/rapids/tool/AppBase.scala: IntegratedRuntimeInjectorto insert a memory marker after processing events. [1] [2]core/src/main/scala/org/apache/spark/sql/rapids/tool/util/RuntimeUtil.scala: Added a methodgetJVMHeapInfoto retrieve JVM heap information. [1] [2]