Re: [PR] [CELEBORN-2398] Allow ChangePartitionManager to refresh candidate workers without allocating slots [celeborn]

["zaynt4606 (via GitHub)" <gi...@apache.org>]

zaynt4606 commented on code in PR #3775:
URL: https://github.com/apache/celeborn/pull/3775#discussion_r3726513920


##########
client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala:
##########
@@ -1880,6 +1881,139 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
     }
   }
 
+  private def syncEndpointReadyWorkers(
+      shuffleId: Int,
+      workersFromMaster: Set[WorkerInfo]): Unit = {
+    val currentEndpointReadyWorkers = workerStatusTracker.endpointReadyWorkers
+    val workersToRemove = currentEndpointReadyWorkers.diff(workersFromMaster)
+    val workersToConnect = workersFromMaster.diff(currentEndpointReadyWorkers)
+    val connectFailedWorkers = new ShuffleFailedWorkers()
+    setupEndpoints(workersToConnect.asJava, shuffleId, connectFailedWorkers)
+    workerStatusTracker.recordWorkerFailure(connectFailedWorkers)
+
+    val connectedWorkers = workersToConnect.diff(connectFailedWorkers.asScala.keySet)
+    workerStatusTracker.addEndpointReadyWorkers(connectedWorkers)
+    workerStatusTracker.removeEndpointReadyWorkers(workersToRemove)
+  }
+
+  private[client] def refreshEndpointReadyWorkersFromMaster(shuffleId: Int): Unit = {
+    val shouldRefresh = endpointReadyWorkersRefreshLock.synchronized {
+      var waitedForRefresh = false
+      val waitDeadline = System.currentTimeMillis() + rpcAskTimeoutMs
+      var remainingWaitTime = rpcAskTimeoutMs
+      while (endpointReadyWorkersRefreshInProgress && remainingWaitTime > 0) {
+        // Reuse the in-flight result instead of letting this revive observe an incomplete pool.
+        waitedForRefresh = true
+        try {
+          endpointReadyWorkersRefreshLock.wait(remainingWaitTime)
+        } catch {
+          case _: InterruptedException =>
+            Thread.currentThread().interrupt()
+            return
+        }
+        remainingWaitTime = waitDeadline - System.currentTimeMillis()
+      }
+      if (endpointReadyWorkersRefreshInProgress) {
+        logWarning(
+          s"Timed out after ${rpcAskTimeoutMs}ms waiting for the in-flight endpoint-ready " +
+            "workers refresh; continue using the current worker pool.")
+      }
+
+      val currentTime = System.currentTimeMillis()
+      val refreshIntervalElapsed = lastEndpointReadyWorkersRefreshAttemptTime == 0L ||
+        currentTime - lastEndpointReadyWorkersRefreshAttemptTime >= dynamicResourceUpdateTime
+      if (!waitedForRefresh && refreshIntervalElapsed) {
+        endpointReadyWorkersRefreshInProgress = true
+        true
+      } else {
+        false
+      }
+    }
+    if (!shouldRefresh) {
+      return
+    }
+
+    try {
+      val requestWorkersRes = requestMasterRequestWorkersWithRetry()
+      StatusCode.fromValue(requestWorkersRes.getStatus) match {
+        case StatusCode.REQUEST_FAILED =>
+          logInfo("ChangePartition requestWorkers RPC request failed.")
+        case StatusCode.SUCCESS =>
+          val availableWorkers =
+            requestWorkersRes.getWorkersList.asScala.map { pbWorkerInfo =>
+              val workerInfo = PbSerDeUtils.fromPbWorkerInfo(pbWorkerInfo)
+              if (pbWorkerInfo.getNetworkLocation.nonEmpty) {
+                workerInfo.networkLocation = pbWorkerInfo.getNetworkLocation
+              }
+              workerInfo
+            }.toSet
+          syncEndpointReadyWorkers(shuffleId, availableWorkers)
+          logDebug(
+            s"ChangePartition requestWorkers succeeded with workers " +
+              s"$availableWorkers.")
+        case StatusCode.WORKER_EXCLUDED =>
+          syncEndpointReadyWorkers(shuffleId, Set.empty)
+          logInfo(s"Offer workers for appId $appUniqueId shuffleId $shuffleId failed.")
+        case StatusCode.SLOT_NOT_AVAILABLE =>
+          syncEndpointReadyWorkers(shuffleId, Set.empty)
+          logInfo(
+            s"No eligible workers are available for appId $appUniqueId shuffleId $shuffleId.")
+        case status =>
+          logWarning(
+            s"ChangePartition requestWorkers failed with status $status.")
+      }
+    } finally {
+      endpointReadyWorkersRefreshLock.synchronized {
+        lastEndpointReadyWorkersRefreshAttemptTime = System.currentTimeMillis()
+        endpointReadyWorkersRefreshInProgress = false
+        endpointReadyWorkersRefreshLock.notifyAll()
+      }
+    }
+  }
+
+  private def requestMasterRequestWorkersWithRetry(): PbRequestWorkersResponse = {
+    val excludedWorkerSet = currentExcludedWorkerSet
+    val req = PbRequestWorkers.newBuilder()
+      .setApplicationId(appUniqueId)
+      .setUserIdentifier(PbSerDeUtils.toPbUserIdentifier(userIdentifier))
+      .setMaxWorkers(slotsAssignMaxWorkers)
+      .setTagsExpr(clientTagsExpr)
+      .setShouldReplicate(pushReplicateEnabled)
+      .setStorageType(storageTypes.head.getValue)

Review Comment:
   Minor: `storageTypes` is likely a `Set`, so `.head` returns an arbitrary element. When multiple storage types are configured (e.g. `HDD,MEMORY`), the type sent to Master depends on Set iteration order. If `.head` returns a non-disk type, Master skips the `haveDisk` filter and may return disk-less workers. The impact is low — `requestSlots` does the real filtering — but with many application splits the extra unnecessary endpoints could add up. Consider picking a disk type preferentially, e.g. `storageTypes.find(t => t == HDD || t == SSD).getOrElse(...)`.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org