Skip to content

Commit

Permalink
Add log for GPU semaphore (#32)
Browse files Browse the repository at this point in the history
add log from semaphore

Signed-off-by: Firestarman <[email protected]>
  • Loading branch information
firestarman authored Mar 11, 2024
1 parent 429b354 commit a63b0e2
Showing 1 changed file with 5 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,7 @@ object GpuSemaphore {
* this is considered to be okay as there are other mechanisms in place, and it should be rather
* rare.
*/
private final class SemaphoreTaskInfo() extends Logging {
private final class SemaphoreTaskInfo(taskAttemptId: Long) extends Logging {
/**
* This holds threads that are not on the GPU yet. Most of the time they are
* blocked waiting for the semaphore to let them on, but it may hold one
Expand Down Expand Up @@ -251,6 +251,7 @@ private final class SemaphoreTaskInfo() extends Logging {
// We cannot be in a synchronized block and wait on the semaphore
// so we have to release it and grab it again afterwards.
semaphore.acquire(numPermits)
logInfo(s"==> Task $taskAttemptId acquiring GPU with $numPermits permits")
synchronized {
// We now own the semaphore so we need to wake up all of the other tasks that are
// waiting.
Expand Down Expand Up @@ -304,6 +305,7 @@ private final class SemaphoreTaskInfo() extends Logging {
activeThreads.remove(t)
if (hasSemaphore) {
semaphore.release(numPermits)
logInfo(s"==> Task $taskAttemptId releasing GPU with $numPermits permits")
hasSemaphore = false
}
// It should be impossible for the current thread to be blocked when releasing the semaphore
Expand All @@ -327,7 +329,7 @@ private final class GpuSemaphore() extends Logging {
val taskAttemptId = context.taskAttemptId()
val taskInfo = tasks.computeIfAbsent(taskAttemptId, _ => {
onTaskCompletion(context, completeTask)
new SemaphoreTaskInfo()
new SemaphoreTaskInfo(taskAttemptId)
})
if (taskInfo.tryAcquire(semaphore)) {
GpuDeviceManager.initializeFromTask()
Expand All @@ -351,7 +353,7 @@ private final class GpuSemaphore() extends Logging {
val taskAttemptId = context.taskAttemptId()
val taskInfo = tasks.computeIfAbsent(taskAttemptId, _ => {
onTaskCompletion(context, completeTask)
new SemaphoreTaskInfo()
new SemaphoreTaskInfo(taskAttemptId)
})
taskInfo.blockUntilReady(semaphore)
GpuDeviceManager.initializeFromTask()
Expand Down

0 comments on commit a63b0e2

Please sign in to comment.