Fix invalid parameter issue when force to close endpoint #25

Signed-off-by: Brian Sheng <[email protected]>
jeynmann · Mar 17, 2023 · ab39e5c · ab39e5c
1 parent 493d6b0
commit ab39e5c
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/src/main/scala/org/apache/spark/shuffle/ucx/UcxWorkerWrapper.scala b/src/main/scala/org/apache/spark/shuffle/ucx/UcxWorkerWrapper.scala
@@ -188,7 +188,14 @@ case class UcxWorkerWrapper(worker: UcpWorker, transport: UcxShuffleTransport, i
   def connectByWorkerAddress(executorId: transport.ExecutorId, workerAddress: ByteBuffer): Unit = {
     logDebug(s"Worker $this connecting back to $executorId by worker address")
     val ep = worker.newEndpoint(new UcpEndpointParams().setName(s"Server connection to $executorId")
-      .setUcpAddress(workerAddress))
+      .setUcpAddress(workerAddress)
+      .setPeerErrorHandlingMode()
+      .setErrorHandler(new UcpEndpointErrorHandler() {
+        override def onError(ep: UcpEndpoint, status: Int, errorMsg: String): Unit = {
+          logError(s"Endpoint to $executorId got an error: $errorMsg")
+          connections.remove(executorId)
+        }
+      }))
     connections.put(executorId, ep)
   }