[CELEBORN-1602] do hard split for push merged data RPC with disk full

### What changes were proposed in this pull request? return hard split for push merged data RPC with disk full ### Why are the changes needed? prevent worker disk from writing to 100% capacity. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Cluster tested Closes apache#2747 from ErikFang/handle-merged-data-disk-full. Authored-by: Erik.fang <[email protected]> Signed-off-by: mingji <[email protected]>
SteNicholas · Sep 20, 2024 · 4381dd3 · 4381dd3
1 parent b11d111
commit 4381dd3
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
@@ -533,6 +533,15 @@ class PushDataHandler(val workerSource: WorkerSource) extends BaseMessageHandler
       callbackWithTimer.onFailure(new CelebornIOException(cause))
       return
     }
+
+    if (fileWriters.exists(checkDiskFull(_) == true)) {
+      val (mapId, attemptId) = getMapAttempt(body)
+      logWarning(
+        s"return hard split for disk full with shuffle $shuffleKey map $mapId attempt $attemptId")
+      callbackWithTimer.onSuccess(ByteBuffer.wrap(Array[Byte](StatusCode.HARD_SPLIT.getValue)))
+      return
+    }
+
     fileWriters.foreach(_.incrementPendingWrites())
 
     val closedFileWriter = fileWriters.find(_.isClosed)