Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Support profiling for specific stages on a limited number of tasks #11708

Open
wants to merge 8 commits into
base: branch-24.12
Choose a base branch
from
Open

Support profiling for specific stages on a limited number of tasks #11708

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
e4a5dd6
wip
thirtiseven Sep 30, 2024
b920cef
Merge branch 'branch-24.12' into profiler_task_limit
thirtiseven Oct 14, 2024
392da62
wip
thirtiseven Oct 25, 2024
549beb4
Merge branch 'branch-24.12' into profiler_task_limit
thirtiseven Oct 30, 2024
d523766
Merge branch 'branch-24.12' into profiler_task_limit
thirtiseven Nov 6, 2024
2ef86af
use task end listener
thirtiseven Nov 7, 2024
c8b90a7
clean up
thirtiseven Nov 7, 2024
dfb6c7e
Merge branch 'branch-24.12' into profiler_task_limit
thirtiseven Nov 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: Would be nice to locate this next to the related config, PROFILE_STAGES

.doc("Limit the number of tasks to profile per stage. A value <= 0 will profile all tasks.")
.internal()
.integerConf
.createWithDefault(0)

// ENABLE/DISABLE PROCESSING

val SQL_ENABLED = conf("spark.rapids.sql.enabled")
Expand Down Expand Up @@ -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)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: Would be nice to locate this next to the related config, profileStages.


lazy val isSqlEnabled: Boolean = get(SQL_ENABLED)

lazy val isSqlExecuteOnGPU: Boolean = get(SQL_MODE).equals("executeongpu")
Expand Down
27 changes: 24 additions & 3 deletions sql-plugin/src/main/scala/com/nvidia/spark/rapids/profiler.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
private var taskLimit = 0
private var stageTaskLimit = 0


def init(pluginCtx: PluginContext, conf: RapidsConf): Unit = {
require(writer.isEmpty, "Already initialized")
timeRanges = conf.profileTimeRangesSeconds.map(parseTimeRanges)
taskLimit = conf.profileTaskLimitPerStage
jobRanges = new RangeConfMatcher(conf, RapidsConf.PROFILE_JOBS)
stageRanges = new RangeConfMatcher(conf, RapidsConf.PROFILE_STAGES)
driverPollMillis = conf.profileDriverPollMillis
Expand Down Expand Up @@ -119,9 +122,27 @@ object ProfilerOnExecutor extends Logging {
val stageId = taskCtx.stageId
if (stageRanges.contains(stageId)) {
synchronized {
activeStages.add(taskCtx.stageId)
enable()
startPollingDriver()
if (taskLimit <= 0) {
// Unlimited tasks per stage
activeStages.add(taskCtx.stageId)
enable()
startPollingDriver()
} else {
// Limited tasks per stage
if (stageTaskCount.getOrElse(stageId, 0) < taskLimit) {
activeStages.add(taskCtx.stageId)
enable()
startPollingDriver()
}
taskCtx.addTaskCompletionListener[Unit] { _ =>
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
Comment on lines +138 to +140
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not multithread safe. Spark guarantees the listener callbacks of a task are called serially, but makes no such guarantee across tasks.

Comment on lines +138 to +140
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

} else {
disable()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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()
                  }
                }

}
}
}
}
}
}
Expand Down
Loading