Support profiling for specific stages on a limited number of tasks #11708
Conversation
Signed-off-by: Haoyang Li <[email protected]>
Signed-off-by: Haoyang Li <[email protected]>
Signed-off-by: Haoyang Li <[email protected]>
Signed-off-by: Haoyang Li <[email protected]>
|
build |
| @@ -792,6 +792,12 @@ val GPU_COREDUMP_PIPE_PATTERN = conf("spark.rapids.gpu.coreDump.pipePattern") | |||
| .bytesConf(ByteUnit.BYTE) | |||
| .createWithDefault(8 * 1024 * 1024) | |||
|
|
|||
| val PROFILE_TASK_LIMIT_PER_STAGE = conf("spark.rapids.profile.taskLimitPerStage") | |||
There was a problem hiding this comment.
Nit: Would be nice to locate this next to the related config, PROFILE_STAGES
| @@ -2603,6 +2609,8 @@ class RapidsConf(conf: Map[String, String]) extends Logging { | |||
|
|
|||
| lazy val profileWriteBufferSize: Long = get(PROFILE_WRITE_BUFFER_SIZE) | |||
|
|
|||
| lazy val profileTaskLimitPerStage: Int = get(PROFILE_TASK_LIMIT_PER_STAGE) | |||
There was a problem hiding this comment.
Nit: Would be nice to locate this next to the related config, profileStages.
| enable() | ||
| startPollingDriver() | ||
| } | ||
| taskCtx.addTaskCompletionListener[Unit] { _ => |
There was a problem hiding this comment.
Seems like we only need to add completions on tasks in the above conditional while we are still under the task limit. Once we're beyond the limit, it's just overhead.
| val currentCount = stageTaskCount.getOrElse(stageId, 0) | ||
| if (currentCount < taskLimit) { | ||
| stageTaskCount(stageId) = currentCount + 1 |
There was a problem hiding this comment.
This is not multithread safe. Spark guarantees the listener callbacks of a task are called serially, but makes no such guarantee across tasks.
| if (currentCount < taskLimit) { | ||
| stageTaskCount(stageId) = currentCount + 1 | ||
| } else { | ||
| disable() |
There was a problem hiding this comment.
Just because a stage has completed a task does not mean we want to disable profiling. Tasks for other stages/jobs that have been flagged for profiling could be active, and we want to keep profiling enabled in that case. See the logic in updateActiveFromDriver. Instead, this needs to do something like the following. Instead of copying this code, it should be refactored from the version in updateActiveFromDriver and reused.
synchronized {
if (activeJobs.isEmpty && activeStages.isEmpty) {
disable()
stopPollingDriver()
}
}| @@ -50,10 +50,13 @@ object ProfilerOnExecutor extends Logging { | |||
| private var isProfileActive = false | |||
| private var currentContextMethod: Method = null | |||
| private var getContextMethod: Method = null | |||
| private val stageTaskCount = mutable.HashMap[Int, Int]() | |||
| private var taskLimit = 0 | |||
There was a problem hiding this comment.
This name is a bit too generic. Is it a task limit across all tasks being profiled? Turns out no, but from the name one cannot tell.
| private var taskLimit = 0 | |
| private var stageTaskLimit = 0 |
| val currentCount = stageTaskCount.getOrElse(stageId, 0) | ||
| if (currentCount < taskLimit) { | ||
| stageTaskCount(stageId) = currentCount + 1 |
There was a problem hiding this comment.
stageId is leaked in the map and never removed. Can solve that by storing a 0 in the map for the stage when the first task for a stage to be profiled is seen. Task completion callback would try to lookup stage in the map. If it gets nothing, there's nothing for it to do. If it gets a task count, it increments it. If we're still under the limit then it stores the updated count, otherwise it removes the count for that stage from the map.
Close #11666
In a customer use case we found that the profiling files were sometimes very large even when we enabled stage/time range limit config on a job.
This pr add support for profiling on only a limited number of tasks for specific stages.
before this pr:
after this pr, task limit = 3:
This nsys-rep file size dropped from 10.4 Mb to 451 Kb.