Add in a layer of indirection for task completion callbacks [databricks] #9009
Conversation
Signed-off-by: Robert (Bobby) Evans <[email protected]>
|
build |
sameerz
added
the
reliability
|
build |
sql-plugin/src/main/scala/com/nvidia/spark/rapids/ScalableTaskCompletion.scala
Outdated
Show resolved
Hide resolved
tests/src/test/scala/com/nvidia/spark/rapids/GpuSemaphoreSuite.scala
Outdated
Show resolved
Hide resolved
|
build |
sql-plugin/src/main/scala/com/nvidia/spark/rapids/ScalableTaskCompletion.scala
Outdated
Show resolved
Hide resolved
| private var wasCalled = false | ||
|
|
||
| override def removeCallback(): Unit = { | ||
| val topLevel = ScalableTaskCompletion.getTopLevel(tc.taskAttemptId()) |
There was a problem hiding this comment.
why not call remove on the hashmap blindly, rather than get and remove?
There was a problem hiding this comment.
This API is removing a callback from a task. We get the collection object that holds task callbacks and check to see if it (the collection object) is null. If it is not we then remove ourselves, the callback, from that collection. I guess each handle could hold onto a reference to the original collection associated with it instead of asking the singleton for it. Is that what you are asking?
| * When the current task completes, if there is a current task call the given function. | ||
| * @return a handle that can be used to remove the callback. | ||
| */ | ||
| def onTaskCompletionIfNotTest(f: => Unit): TaskCompletionCallbackHandle = { |
There was a problem hiding this comment.
I am a little unclear on when I should use onTaskCompletionIfNotTest vs onTaskCompletion and it feels like having multiple choices is an area where we could introduce bugs. Is there any way we could change the logic if we know we are in the scala tests, but the calling code is unaware and has 1 choice to register?
There was a problem hiding this comment.
onTaskCompletionIfNotTest inserts a callback if a TaskContext is available. If none is available, then it is not inserted. This was done by a lot of code for unit tests. The problem is that I don't want to encourage this behavior. The API itself is somewhat dangerous because we might or might not get a callback. I am fine with removing it and then forcing everywhere it is used to put in their own workaround.
|
build |
This fixes #8482
It touches a lot of files because I wanted to move as much of the old API over to the new one so that I could properly test it. I think there may be other places where removing a callback might be nice too, but I didn't take that route here.