[CELEBORN-2210] When a flushBuffer consolidation OOM exception occurs…

…, support setting the Buffer for fileInfo.

### What changes were proposed in this pull request?

When a flushBuffer consolidation OOM exception occurs, support setting the Buffer for fileInfo.

### Why are the changes needed?

When a flushBuffer consolidation OOM exception occurs, the current logic does not allow setting the Buffer for fileInfo.

### Does this PR resolve a correctness bug?

NO

### Does this PR introduce _any_ user-facing change?

NO

### How was this patch tested?

CI

Closes #3547 from xy2953396112/CELEBORN-2210.

Authored-by: xxx <[email protected]>
Signed-off-by: Shuang <[email protected]>

Author: xy2953396112

worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriter.scala

@@ -341,8 +341,18 @@ class MemoryTierWriter(
   }
 
   override def closeStreams(): Unit = {
-    flushBuffer.consolidate()
-    fileInfo.setBuffer(flushBuffer)
+    try {
+      flushBuffer.consolidate()
+    } catch {
+      case oom: OutOfMemoryError =>
+        logError(
+          s"MemoryTierWriter shuffleKey:${partitionDataWriterContext.getShuffleKey}, " +
+            s"partitionId:${partitionDataWriterContext.getPartitionLocation.getFileName} " +
+            s"failed to consolidate flush buffer due to OutOfMemoryError.",
+          oom)
+    } finally {
+      fileInfo.setBuffer(flushBuffer)
+    }
   }
 
   override def takeBufferInternal(): CompositeByteBuf = {