Kalvin2077 commented on code in PR #3775:
URL: https://github.com/apache/celeborn/pull/3775#discussion_r3726735918
##########
master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala:
##########
@@ -1062,34 +1068,116 @@ private[celeborn] class Master(
s"extraSlots=$offerSlotsExtraSize"))
if (authEnabled) {
- pushApplicationMetaToWorkers(requestSlots, slots)
+ pushApplicationMetaToWorkers(requestSlots.applicationId, slots.keySet())
}
context.reply(RequestSlotsResponse(
StatusCode.SUCCESS,
slots.asInstanceOf[WorkerResource],
requestSlots.packed))
}
+ private def selectWorkersForRequest(
+ requestWorkers: PbRequestWorkers,
+ availableWorkers: util.List[WorkerInfo]): util.List[WorkerInfo] = {
+ val maxWorkers =
+ if (requestWorkers.getMaxWorkers <= 0) splitSlotAssignMaxWorkers
+ else Math.min(splitSlotAssignMaxWorkers, requestWorkers.getMaxWorkers)
+ val storageType = StorageInfo.typesMap.get(requestWorkers.getStorageType)
+ if (storageType == null) {
+ return Collections.emptyList()
+ }
+ val eligibleWorkers = new util.ArrayList[WorkerInfo]()
+ availableWorkers.asScala
+ .filter { worker =>
+ (storageType != StorageInfo.Type.HDD && storageType !=
StorageInfo.Type.SSD) ||
+ worker.haveDisk
+ }
+ .foreach(eligibleWorkers.add)
+ if (eligibleWorkers.isEmpty) {
+ return Collections.emptyList()
+ }
+
+ val minWorkers = if (requestWorkers.getShouldReplicate) 2 else 1
+ val selectedWorkerCount =
+ Math.min(Math.max(minWorkers, maxWorkers), eligibleWorkers.size)
+ val startIndex = Random.nextInt(eligibleWorkers.size)
Review Comment:
Thx. I've used `ThreadLocalRandom`.
##########
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:
Thx. I've already changed the type of `storageTypes` to bitmask(int32).
##########
master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala:
##########
@@ -1062,34 +1068,112 @@ private[celeborn] class Master(
s"extraSlots=$offerSlotsExtraSize"))
if (authEnabled) {
- pushApplicationMetaToWorkers(requestSlots, slots)
+ pushApplicationMetaToWorkers(requestSlots.applicationId, slots.keySet())
}
context.reply(RequestSlotsResponse(
StatusCode.SUCCESS,
slots.asInstanceOf[WorkerResource],
requestSlots.packed))
}
+ private def selectWorkersForRequest(
+ requestWorkers: PbRequestWorkers,
+ availableWorkers: util.List[WorkerInfo]): util.List[WorkerInfo] = {
+ val maxWorkers =
+ if (requestWorkers.getMaxWorkers <= 0) splitSlotAssignMaxWorkers
+ else Math.min(splitSlotAssignMaxWorkers, requestWorkers.getMaxWorkers)
+ val eligibleWorkers = new util.ArrayList[WorkerInfo]()
+ availableWorkers.asScala
+ .filter { worker =>
+
!StorageInfo.localDiskAvailable(requestWorkers.getAvailableStorageTypes) ||
+ worker.haveDisk
+ }
+ .foreach(eligibleWorkers.add)
+ if (eligibleWorkers.isEmpty) {
+ return Collections.emptyList()
+ }
+
+ val minWorkers = if (requestWorkers.getShouldReplicate) 2 else 1
+ val selectedWorkerCount =
+ Math.min(Math.max(minWorkers, maxWorkers), eligibleWorkers.size)
+ val startIndex = ThreadLocalRandom.current().nextInt(eligibleWorkers.size)
Review Comment:
This is intentional. RequestWorkers is a best-effort discovery RPC, so
SUCCESS means the available candidate set was returned, not that replicated
slot allocation can proceed. ChangePartitionManager.handleRequestPartitions
independently returns SLOT_NOT_AVAILABLE when replication is enabled and the
merged candidate set has fewer than two workers. Keeping the partial result
also allows it to be merged with existing snapshot candidates.
##########
master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala:
##########
@@ -1062,34 +1068,112 @@ private[celeborn] class Master(
s"extraSlots=$offerSlotsExtraSize"))
if (authEnabled) {
- pushApplicationMetaToWorkers(requestSlots, slots)
+ pushApplicationMetaToWorkers(requestSlots.applicationId, slots.keySet())
}
context.reply(RequestSlotsResponse(
StatusCode.SUCCESS,
slots.asInstanceOf[WorkerResource],
requestSlots.packed))
}
+ private def selectWorkersForRequest(
Review Comment:
> This is a similar code as handleRequestSlots. We could have this used in
handleRequestSlots as well?
Thx.
The similar code has been extracted.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
