Skip to content

Commit

Permalink
add comment about the accumulator usage
Browse files Browse the repository at this point in the history
Signed-off-by: Zach Puller <zpuller@nvidia.com>
  • Loading branch information
zpuller committed Sep 24, 2024
1 parent ebb0803 commit 4e81016
Showing 1 changed file with 5 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -214,6 +214,11 @@ class GpuTaskMetrics extends Serializable {
def updateMaxGpuMemory(taskAttemptId: Long): Unit = {
val maxMem = RmmSpark.getAndResetGpuMaxMemoryAllocated(taskAttemptId)
if (maxMem > 0) {
// This metric tracks the max amount of memory that is allocated on the gpu during
// the lifespan of a task. However, this update function only gets called once on task
// completion, whereas the actual logic tracking of the max value during memory allocations
// lives in the JNI. Therefore we can stick the convention here of calling the add method
// instead of adding a dedicated max method to the accumulator.
maxDeviceMemoryBytes.add(maxMem)
}
}
Expand Down

0 comments on commit 4e81016

Please sign in to comment.