This is an automated email from the ASF dual-hosted git repository.
zaynt4606 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 6b6efe2c2e [CELEBORN-2398] Allow ChangePartitionManager to refresh
candidate workers without allocating slots
6b6efe2c2e is described below
commit 6b6efe2c2e7a0efef26b12419c7f2f50152767f0
Author: Kalvin2077 <[email protected]>
AuthorDate: Tue Sep 15 10:46:52 2026 +0800
[CELEBORN-2398] Allow ChangePartitionManager to refresh candidate workers
without allocating slots
### What changes were proposed in this pull request?
This PR introduces a Protobuf-based
`RequestWorkers`/`RequestWorkersResponse` RPC for read-only worker discovery.
On the Master side, the new handler selects currently available workers
after applying client exclusions and tag filters. It limits the response using
the smaller of `celeborn.client.slot.assign.maxWorkers` and
`celeborn.master.splitSlot.assign.maxWorkers`, while preserving the minimum
worker count required for replication. When authentication is enabled,
application metadata is pushed to the selected workers.
On the client side, LifecycleManager periodically requests workers from the
Master, creates endpoints for newly selected workers, records connection
failures, and maintains an endpoint-ready worker pool. ChangePartitionManager
uses that pool for change-partition requests and falls back to the shuffle's
existing worker snapshots when no refreshed candidates are available.
The PR also makes the following configuration changes:
- removes `celeborn.client.shuffle.dynamicResourceFactor`;
- adds `celeborn.client.shuffle.dynamicResource.updateTime`, defaulting to
`30s`;
- adds `celeborn.master.splitSlot.assign.maxWorkers`, defaulting to `500`;
- documents the breaking client configuration change in the migration guide.
### Why are the changes needed?
The existing factor-based logic refreshes candidates only after enough
workers from the shuffle's original allocation become unavailable. It therefore
cannot use newly added workers during normal workers scale-out while the
original workers remain healthy.
Requesting slots merely to discover workers also mixes worker discovery
with slot allocation and mutates Master shuffle state. A dedicated read-only
RPC allows ChangePartitionManager to use the current cluster membership without
those side effects, while rate limiting and worker-count caps prevent excessive
RPC connections and oversized responses.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [x] Yes
Users of `celeborn.client.shuffle.dynamicResourceFactor` must migrate to
`celeborn.client.shuffle.dynamicResource.updateTime`.
### How was this patch tested?
- Added unit tests for Worker selection, storage eligibility, dynamic
Worker merging, deduplication, and excluded-Worker filtering.
- Built with JDK 8 and deployed to a four-Worker Celeborn cluster.
- Ran an E2E shuffle that started with one Worker, added three Workers, and
paused the original Worker. The client detected the failure, revived partitions
on new Workers, and successfully validated 8 million rows with fallback
disabled.
- Restored all Workers and reran the standard Spark/YARN smoke test
successfully.
Closes #3775 from Kalvin2077/feat/dynamic-resource.
Authored-by: Kalvin2077 <[email protected]>
Signed-off-by: zhengtao <[email protected]>
AI-Contributed/Feature: 0/471
AI-Contributed/UT: 0/431
---
.../celeborn/client/ChangePartitionManager.scala | 98 +++--------
.../apache/celeborn/client/LifecycleManager.scala | 146 ++++++++++++++-
.../celeborn/client/WorkerStatusTracker.scala | 15 ++
.../client/ChangePartitionManagerSuite.scala | 90 ++++++++++
.../celeborn/client/WorkerStatusTrackerSuite.scala | 25 +++
common/src/main/proto/TransportMessages.proto | 17 ++
.../org/apache/celeborn/common/CelebornConf.scala | 43 +++--
.../common/protocol/message/ControlMessages.scala | 12 ++
.../apache/celeborn/common/CelebornConfSuite.scala | 25 +++
.../apache/celeborn/common/util/UtilsSuite.scala | 33 +++-
docs/configuration/client.md | 4 +-
docs/configuration/master.md | 1 +
docs/migration.md | 5 +
.../celeborn/service/deploy/master/Master.scala | 130 +++++++++++---
.../service/deploy/master/MasterSuite.scala | 196 ++++++++++++++++++++-
.../ChangePartitionManagerUpdateWorkersSuite.scala | 60 +++++--
.../celeborn/tests/spark/RetryReviveTest.scala | 2 +-
17 files changed, 760 insertions(+), 142 deletions(-)
diff --git
a/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
b/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
index 71096a952f..c0af4e9113 100644
---
a/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
+++
b/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
@@ -23,7 +23,6 @@ import java.util.concurrent.{ConcurrentHashMap,
ScheduledExecutorService, Schedu
import scala.collection.JavaConverters._
-import org.apache.celeborn.client.LifecycleManager.ShuffleFailedWorkers
import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.internal.Logging
import org.apache.celeborn.common.meta.{ShufflePartitionLocationInfo,
WorkerInfo}
@@ -77,7 +76,6 @@ class ChangePartitionManager(
private val testRetryRevive = conf.testRetryRevive
private val dynamicResourceEnabled = conf.clientShuffleDynamicResourceEnabled
- private val dynamicResourceUnavailableFactor =
conf.clientShuffleDynamicResourceFactor
def start(): Unit = {
batchHandleChangePartition = batchHandleChangePartitionSchedulerThread.map
{
@@ -284,73 +282,7 @@ class ChangePartitionManager(
}
}
- val candidates = new util.HashSet[WorkerInfo]()
- val newlyRequestedLocations = new WorkerResource()
-
- val snapshotCandidates =
- lifecycleManager
- .workerSnapshots(shuffleId)
- .asScala
- .values
- .map(_.workerInfo)
- .filter(lifecycleManager.workerStatusTracker.workerAvailable)
- .toSet
- .asJava
- candidates.addAll(snapshotCandidates)
-
- if (dynamicResourceEnabled) {
- val shuffleAllocatedWorkers =
lifecycleManager.workerSnapshots(shuffleId).size()
- val unavailableWorkerRatio = 1 - (snapshotCandidates.size * 1.0 /
shuffleAllocatedWorkers)
- if (candidates.size < 1 || (pushReplicateEnabled && candidates.size < 2)
- || (unavailableWorkerRatio >= dynamicResourceUnavailableFactor)) {
-
- // get new available workers for the request partition ids
- val partitionIds = new util.ArrayList[Integer](
-
changePartitions.map(_.partitionId).map(Integer.valueOf).toList.asJava)
- // The partition id value is not important here because we're just
trying to get the workers to use
- val requestSlotsRes =
- lifecycleManager.requestMasterRequestSlotsWithRetry(shuffleId,
partitionIds)
-
- requestSlotsRes.status match {
- case StatusCode.REQUEST_FAILED =>
- logInfo(s"ChangePartition requestSlots RPC request failed for
$shuffleId!")
- case StatusCode.SLOT_NOT_AVAILABLE =>
- logInfo(s"ChangePartition requestSlots for $shuffleId failed, have
no available slots.")
- case StatusCode.SUCCESS =>
- logDebug(
- s"ChangePartition requestSlots request for workers Success!
shuffleId: $shuffleId availableWorkers Info:
${requestSlotsRes.workerResource.keySet()}")
- case StatusCode.WORKER_EXCLUDED =>
- logInfo(s"ChangePartition requestSlots request for workers for
$shuffleId failed due to all workers be excluded!")
- case _ => // won't happen
- throw new UnsupportedOperationException()
- }
-
- if (requestSlotsRes.status.equals(StatusCode.SUCCESS)) {
- requestSlotsRes.workerResource.keySet().asScala.foreach {
workerInfo: WorkerInfo =>
- newlyRequestedLocations.computeIfAbsent(workerInfo,
lifecycleManager.newLocationFunc)
- }
-
- // SetupEndpoint for new Workers
- val workersRequireEndpoints = new util.HashSet[WorkerInfo](
- requestSlotsRes.workerResource.keySet()
- .asScala
- .filter(lifecycleManager.workerStatusTracker.workerAvailable)
- .asJava)
-
- val connectFailedWorkers = new ShuffleFailedWorkers()
- lifecycleManager.setupEndpoints(
- workersRequireEndpoints,
- shuffleId,
- connectFailedWorkers)
-
workersRequireEndpoints.removeAll(connectFailedWorkers.asScala.keys.toList.asJava)
- candidates.addAll(workersRequireEndpoints)
-
- // Update worker status
-
lifecycleManager.workerStatusTracker.recordWorkerFailure(connectFailedWorkers)
-
lifecycleManager.workerStatusTracker.removeFromExcludedWorkers(candidates)
- }
- }
- }
+ val candidates = collectCandidateWorkers(shuffleId)
if (candidates.size < 1 || (pushReplicateEnabled && candidates.size < 2)) {
logError("[Update partition] failed for not enough candidates for
revive.")
@@ -374,10 +306,7 @@ class ChangePartitionManager(
return
}
- // newlyRequestedLocations is empty if dynamicResourceEnabled is false
- newlyRequestedLocations.putAll(newlyAllocatedLocations)
-
- val newPrimaryLocations = newlyRequestedLocations.asScala.flatMap {
+ val newPrimaryLocations = newlyAllocatedLocations.asScala.flatMap {
case (workInfo, (primaryLocations, replicaLocations)) =>
// Add all re-allocated slots to worker snapshots.
val partitionLocationInfo =
lifecycleManager.workerSnapshots(shuffleId).computeIfAbsent(
@@ -410,6 +339,29 @@ class ChangePartitionManager(
replySuccess(newPrimaryLocations.toArray)
}
+ private[client] def collectCandidateWorkers(shuffleId: Int):
util.HashSet[WorkerInfo] = {
+ if (dynamicResourceEnabled) {
+ lifecycleManager.refreshEndpointReadyWorkersFromMaster(shuffleId)
+ }
+
+ val snapshotCandidates =
+ lifecycleManager
+ .workerSnapshots(shuffleId)
+ .asScala
+ .values
+ .map(_.workerInfo)
+ .filter(lifecycleManager.workerStatusTracker.workerAvailable)
+ .toSet
+ val candidates = new util.HashSet[WorkerInfo](snapshotCandidates.asJava)
+ if (dynamicResourceEnabled) {
+ candidates.addAll(
+ lifecycleManager.workerStatusTracker.endpointReadyWorkers
+ .filter(lifecycleManager.workerStatusTracker.workerAvailable)
+ .asJava)
+ }
+ candidates
+ }
+
private def reallocateChangePartitionRequestSlotsFromCandidates(
changePartitionRequests: List[ChangePartitionRequest],
candidates: List[WorkerInfo]): WorkerResource = {
diff --git
a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
index f9508cfe6c..88b778395f 100644
--- a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
@@ -124,6 +124,11 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
private val excludedWorkersFilter =
conf.registerShuffleFilterExcludedWorkerEnabled
+ private val dynamicResourceUpdateTime =
conf.clientShuffleDynamicResourceUpdateTime
+ private val endpointReadyWorkersRefreshLock = new Object
+ private var endpointReadyWorkersRefreshInProgress = false
+ private var lastEndpointReadyWorkersRefreshAttemptTime = 0L
+
private val registerShuffleResponseRpcCache: Cache[Int, ByteBuffer] =
CacheBuilder.newBuilder()
.concurrencyLevel(rpcCacheConcurrencyLevel)
.expireAfterAccess(rpcCacheExpireTime, TimeUnit.MILLISECONDS)
@@ -859,6 +864,7 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
partitionLocationInfo.addReplicaPartitions(replicaLocations)
allocatedWorkers.put(workerInfo.toUniqueId, partitionLocationInfo)
}
+
workerStatusTracker.addEndpointReadyWorkers(candidatesWorkers.asScala.toSet)
shuffleAllocatedWorkers.put(shuffleId, allocatedWorkers)
registeredShuffle.add(shuffleId)
commitManager.registerShuffle(
@@ -1850,12 +1856,7 @@ class LifecycleManager(val appUniqueId: String, val
conf: CelebornConf) extends
def requestMasterRequestSlotsWithRetry(
shuffleId: Int,
ids: util.ArrayList[Integer]): RequestSlotsResponse = {
- val excludedWorkerSet =
- if (excludedWorkersFilter) {
- workerStatusTracker.excludedWorkers.asScala.keys.toSet
- } else {
- Set.empty[WorkerInfo]
- }
+ val excludedWorkerSet = currentExcludedWorkerSet
// UserResourceConsumption and DiskInfo are eliminated from WorkerInfo
// during serialization of RequestSlots
val req =
@@ -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)
+ .setAvailableStorageTypes(availableStorageTypes)
+ .addAllExcludedWorkerSet(excludedWorkerSet.map(
+ PbSerDeUtils.toPbWorkerInfo(_, true, true)).asJava)
+ .build()
+ val res = requestMasterRequestWorkers(req)
+ if (StatusCode.fromValue(res.getStatus) == StatusCode.REQUEST_FAILED) {
+ requestMasterRequestWorkers(req)
+ } else {
+ res
+ }
+ }
+
+ private def currentExcludedWorkerSet: Set[WorkerInfo] = {
+ if (excludedWorkersFilter) {
+ workerStatusTracker.excludedWorkers.asScala.keys.toSet
+ } else {
+ Set.empty
+ }
+ }
+
+ private def requestMasterRequestWorkers(
+ message: PbRequestWorkers): PbRequestWorkersResponse = {
+ try {
+ masterClient.askSync[PbRequestWorkersResponse](
+ message,
+ classOf[PbRequestWorkersResponse])
+ } catch {
+ case e: Exception =>
+ logError("AskSync request workers failed.", e)
+ PbRequestWorkersResponse.newBuilder()
+ .setStatus(StatusCode.REQUEST_FAILED.getValue)
+ .build()
+ }
+ }
+
private def requestMasterRequestSlots(message: RequestSlots):
RequestSlotsResponse = {
val shuffleKey = Utils.makeShuffleKey(message.applicationId,
message.shuffleId)
try {
diff --git
a/client/src/main/scala/org/apache/celeborn/client/WorkerStatusTracker.scala
b/client/src/main/scala/org/apache/celeborn/client/WorkerStatusTracker.scala
index 2c1243acb1..67e56a08e1 100644
--- a/client/src/main/scala/org/apache/celeborn/client/WorkerStatusTracker.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/WorkerStatusTracker.scala
@@ -41,6 +41,21 @@ class WorkerStatusTracker(
val excludedWorkers = new ShuffleFailedWorkers()
val shuttingWorkers: JSet[WorkerInfo] =
ConcurrentHashMap.newKeySet[WorkerInfo]()
+ // These workers already have client-local RPC endpoints and can reserve
slots directly.
+ private val endpointReadyWorkerSet: JSet[WorkerInfo] =
+ ConcurrentHashMap.newKeySet[WorkerInfo]()
+
+ private[client] def endpointReadyWorkers: Set[WorkerInfo] =
+ endpointReadyWorkerSet.asScala.toSet
+
+ private[client] def addEndpointReadyWorkers(workers: Set[WorkerInfo]): Unit
= {
+ workers.filter(workerAvailable).foreach(endpointReadyWorkerSet.add)
+ }
+
+ private[client] def removeEndpointReadyWorkers(workers: Set[WorkerInfo]):
Unit = {
+ workers.foreach(endpointReadyWorkerSet.remove)
+ }
+
def registerWorkerStatusListener(workerStatusListener:
WorkerStatusListener): Unit = {
workerStatusListeners.add(workerStatusListener)
}
diff --git
a/client/src/test/scala/org/apache/celeborn/client/ChangePartitionManagerSuite.scala
b/client/src/test/scala/org/apache/celeborn/client/ChangePartitionManagerSuite.scala
new file mode 100644
index 0000000000..f57d1499ed
--- /dev/null
+++
b/client/src/test/scala/org/apache/celeborn/client/ChangePartitionManagerSuite.scala
@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.client
+
+import java.util
+
+import scala.collection.JavaConverters._
+
+import org.mockito.Mockito.{doAnswer, mock, verify, when}
+import org.mockito.invocation.InvocationOnMock
+import org.mockito.stubbing.Answer
+
+import org.apache.celeborn.CelebornFunSuite
+import org.apache.celeborn.common.CelebornConf
+import org.apache.celeborn.common.meta.{ShufflePartitionLocationInfo,
WorkerInfo}
+import org.apache.celeborn.common.protocol.message.StatusCode
+
+class ChangePartitionManagerSuite extends CelebornFunSuite {
+ test("collectCandidateWorkers combines snapshot and endpoint-ready workers")
{
+ val conf = dynamicResourceConf()
+ val lifecycleManager = mock(classOf[LifecycleManager])
+ val statusTracker = new WorkerStatusTracker(conf, lifecycleManager)
+ when(lifecycleManager.workerStatusTracker).thenReturn(statusTracker)
+
+ val snapshotWorker = worker("snapshot")
+ val duplicateSnapshotWorker = worker("snapshot")
+ val dynamicWorker = worker("dynamic")
+
when(lifecycleManager.workerSnapshots(1)).thenReturn(snapshot(snapshotWorker))
+ statusTracker.addEndpointReadyWorkers(Set(duplicateSnapshotWorker,
dynamicWorker))
+
+ val manager = new ChangePartitionManager(conf, lifecycleManager)
+ val candidates = manager.collectCandidateWorkers(1)
+
+ verify(lifecycleManager).refreshEndpointReadyWorkersFromMaster(1)
+ assert(candidates.asScala.toSet === Set(snapshotWorker, dynamicWorker))
+ assert(candidates.asScala.find(_ == snapshotWorker).get eq snapshotWorker)
+ }
+
+ test("collectCandidateWorkers filters workers excluded during refresh") {
+ val conf = dynamicResourceConf()
+ val lifecycleManager = mock(classOf[LifecycleManager])
+ val statusTracker = new WorkerStatusTracker(conf, lifecycleManager)
+ when(lifecycleManager.workerStatusTracker).thenReturn(statusTracker)
+
+ val failedWorker = worker("failed")
+
when(lifecycleManager.workerSnapshots(1)).thenReturn(snapshot(failedWorker))
+ statusTracker.addEndpointReadyWorkers(Set(failedWorker))
+ doAnswer(new Answer[AnyRef] {
+ override def answer(invocation: InvocationOnMock): AnyRef = {
+ statusTracker.excludedWorkers.put(
+ failedWorker,
+ (StatusCode.WORKER_UNRESPONSIVE, System.currentTimeMillis()))
+ null
+ }
+ }).when(lifecycleManager).refreshEndpointReadyWorkersFromMaster(1)
+
+ val manager = new ChangePartitionManager(conf, lifecycleManager)
+
+ assert(manager.collectCandidateWorkers(1).isEmpty)
+ }
+
+ private def dynamicResourceConf(): CelebornConf = {
+ new
CelebornConf().set(CelebornConf.CLIENT_SHUFFLE_DYNAMIC_RESOURCE_ENABLED, true)
+ }
+
+ private def worker(host: String): WorkerInfo = {
+ new WorkerInfo(host, 1001, 1002, 1003, 1004)
+ }
+
+ private def snapshot(worker: WorkerInfo): util.Map[String,
ShufflePartitionLocationInfo] = {
+ val workers = new util.HashMap[String, ShufflePartitionLocationInfo]()
+ workers.put(worker.toUniqueId, new ShufflePartitionLocationInfo(worker))
+ workers
+ }
+}
diff --git
a/client/src/test/scala/org/apache/celeborn/client/WorkerStatusTrackerSuite.scala
b/client/src/test/scala/org/apache/celeborn/client/WorkerStatusTrackerSuite.scala
index 90a1d53bbd..43cbefd45e 100644
---
a/client/src/test/scala/org/apache/celeborn/client/WorkerStatusTrackerSuite.scala
+++
b/client/src/test/scala/org/apache/celeborn/client/WorkerStatusTrackerSuite.scala
@@ -34,6 +34,31 @@ import
org.apache.celeborn.common.protocol.message.ControlMessages.HeartbeatFrom
import org.apache.celeborn.common.protocol.message.StatusCode
class WorkerStatusTrackerSuite extends CelebornFunSuite {
+ test("maintain endpoint-ready worker pool") {
+ val statusTracker = new WorkerStatusTracker(new CelebornConf(), null)
+ val worker1 = mock("host1")
+ val worker2 = mock("host2")
+ val worker3 = mock("host3")
+
+ statusTracker.addEndpointReadyWorkers(Set(worker1, worker2))
+ assert(statusTracker.endpointReadyWorkers == Set(worker1, worker2))
+
+ val replacementWorker2 = mock("host2")
+ statusTracker.addEndpointReadyWorkers(Set(replacementWorker2, worker3))
+ assert(statusTracker.endpointReadyWorkers == Set(worker1, worker2,
worker3))
+ // Keep the existing WorkerInfo instance because it owns the client-local
RPC endpoint.
+ assert(statusTracker.endpointReadyWorkers.find(_ == worker2).get eq
worker2)
+
+ statusTracker.removeEndpointReadyWorkers(Set(worker1))
+ assert(statusTracker.endpointReadyWorkers == Set(worker2, worker3))
+
+ statusTracker.excludedWorkers.put(
+ worker1,
+ (StatusCode.WORKER_EXCLUDED, System.currentTimeMillis()))
+ statusTracker.addEndpointReadyWorkers(Set(worker1))
+ assert(statusTracker.endpointReadyWorkers == Set(worker2, worker3))
+ }
+
test("handleHeartbeatResponse without availableWorkers") {
val celebornConf = new CelebornConf()
celebornConf.set(CLIENT_EXCLUDED_WORKER_EXPIRE_TIMEOUT, 2000L)
diff --git a/common/src/main/proto/TransportMessages.proto
b/common/src/main/proto/TransportMessages.proto
index a813a9e501..e57055391e 100644
--- a/common/src/main/proto/TransportMessages.proto
+++ b/common/src/main/proto/TransportMessages.proto
@@ -117,6 +117,8 @@ enum MessageType {
READ_REDUCER_PARTITION_END = 94;
READ_REDUCER_PARTITION_END_RESPONSE = 95;
REGISTER_APPLICATION_INFO = 96;
+ REQUEST_WORKERS = 97;
+ REQUEST_WORKERS_RESPONSE = 98;
}
enum StreamType {
@@ -325,6 +327,16 @@ message PbRequestSlots {
string tagsExpr = 14;
}
+message PbRequestWorkers {
+ string applicationId = 1;
+ PbUserIdentifier userIdentifier = 2;
+ int32 maxWorkers = 3;
+ string tagsExpr = 4;
+ repeated PbWorkerInfo excludedWorkerSet = 5;
+ bool shouldReplicate = 6;
+ int32 availableStorageTypes = 7;
+}
+
message PbSlotInfo {
map<string, int32> slot = 1;
}
@@ -335,6 +347,11 @@ message PbRequestSlotsResponse {
map<string, PbPackedWorkerResource> packedWorkerResource = 3;
}
+message PbRequestWorkersResponse {
+ int32 status = 1;
+ repeated PbWorkerInfo workers = 2;
+}
+
message PbRevivePartitionInfo {
int32 partitionId = 1;
int32 epoch = 2;
diff --git
a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index 2a928c611b..1e34ac34f9 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -659,7 +659,7 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable
with Logging with Se
def masterSlotAssignPolicyName: String = get(MASTER_SLOT_ASSIGN_POLICY)
/** Returns the configured built-in policy. Use `masterSlotAssignPolicyName`
for SPI providers. */
- @deprecated("Use masterSlotAssignPolicyName for SPI provider selection",
"0.7.0")
+ @deprecated("Use masterSlotAssignPolicyName for SPI provider selection",
"1.0.0")
def masterSlotAssignPolicy: SlotsAssignPolicy =
SlotsAssignPolicy.valueOf(get(MASTER_SLOT_ASSIGN_POLICY))
@@ -688,6 +688,7 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable
with Logging with Se
def masterSlotAssignExtraSlots: Int = get(MASTER_SLOT_ASSIGN_EXTRA_SLOTS)
def masterSlotAssignMaxWorkers: Int = get(MASTER_SLOT_ASSIGN_MAX_WORKERS)
def masterSlotAssignMinWorkers: Int = get(MASTER_SLOT_ASSIGN_MIN_WORKERS)
+ def masterSplitSlotAssignMaxWorkers: Int =
get(MASTER_SPLIT_SLOT_ASSIGN_MAX_WORKERS)
def initialEstimatedPartitionSize: Long =
get(ESTIMATED_PARTITION_SIZE_INITIAL_SIZE)
def estimatedPartitionSizeUpdaterInitialDelay: Long =
get(ESTIMATED_PARTITION_SIZE_UPDATE_INITIAL_DELAY)
@@ -961,7 +962,8 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable
with Logging with Se
def clientCommitFilesIgnoreExcludedWorkers: Boolean =
get(CLIENT_COMMIT_IGNORE_EXCLUDED_WORKERS)
def clientShuffleDynamicResourceEnabled: Boolean =
get(CLIENT_SHUFFLE_DYNAMIC_RESOURCE_ENABLED)
- def clientShuffleDynamicResourceFactor: Double =
get(CLIENT_SHUFFLE_DYNAMIC_RESOURCE_FACTOR)
+ def clientShuffleDynamicResourceUpdateTime: Long =
+ get(CLIENT_SHUFFLE_DYNAMIC_RESOURCE_UPDATE_TIME)
def appHeartbeatTimeoutMs: Long = get(APPLICATION_HEARTBEAT_TIMEOUT)
def dfsExpireDirsTimeoutMS: Long = get(DFS_EXPIRE_DIRS_TIMEOUT)
def appHeartbeatIntervalMs: Long = get(APPLICATION_HEARTBEAT_INTERVAL)
@@ -3222,6 +3224,18 @@ object CelebornConf extends Logging {
.intConf
.createWithDefault(100)
+ val MASTER_SPLIT_SLOT_ASSIGN_MAX_WORKERS: ConfigEntry[Int] =
+ buildConf("celeborn.master.splitSlot.assign.maxWorkers")
+ .categories("master")
+ .version("1.0.0")
+ .doc("Maximum workers returned by each dynamic candidate refresh. The
request limit is the " +
+ "smaller positive value of this setting and
`celeborn.client.slot.assign.maxWorkers`. " +
+ "For replicated shuffle, an effective limit of one is raised to two.
Workers already " +
+ "present in a shuffle snapshot are not counted against this limit.")
+ .intConf
+ .checkValue(_ > 0, "Must be positive.")
+ .createWithDefault(500)
+
val ESTIMATED_PARTITION_SIZE_INITIAL_SIZE: ConfigEntry[Long] =
buildConf("celeborn.master.estimatedPartitionSize.initialSize")
.withAlternative("celeborn.shuffle.initialEstimatedPartitionSize")
@@ -5603,21 +5617,24 @@ object CelebornConf extends Logging {
buildConf("celeborn.client.shuffle.dynamicResourceEnabled")
.categories("client")
.version("0.6.0")
- .doc("When enabled, the ChangePartitionManager will obtain candidate
workers from the availableWorkers pool " +
- "during heartbeats when worker resource change.")
+ .doc("When enabled, ChangePartitionManager refreshes endpoint-ready
worker candidates from " +
+ "the Master on demand while handling change-partition requests, and
combines them with " +
+ "workers already present in the shuffle snapshot.")
.booleanConf
.createWithDefault(false)
- val CLIENT_SHUFFLE_DYNAMIC_RESOURCE_FACTOR: ConfigEntry[Double] =
- buildConf("celeborn.client.shuffle.dynamicResourceFactor")
+ val CLIENT_SHUFFLE_DYNAMIC_RESOURCE_UPDATE_TIME: ConfigEntry[Long] =
+ buildConf("celeborn.client.shuffle.dynamicResource.updateTime")
.categories("client")
- .version("0.6.0")
- .doc("The ChangePartitionManager will check whether (unavailable workers
/ shuffle allocated workers) " +
- "is more than the factor before obtaining candidate workers from the
requestSlots RPC response " +
- s"when `${CLIENT_SHUFFLE_DYNAMIC_RESOURCE_ENABLED.key}` set true")
- .doubleConf
- .checkValue(v => v >= 0.0 && v <= 1.0, "Should be in [0.0, 1.0].")
- .createWithDefault(0.5)
+ .version("1.0.0")
+ .doc(
+ "Minimum interval after a worker-candidate refresh attempt completes
before " +
+ s"ChangePartitionManager may try again when
`${CLIENT_SHUFFLE_DYNAMIC_RESOURCE_ENABLED.key}` " +
+ "is true. Set to 0 to allow each change-partition handling cycle to
refresh when no " +
+ "refresh is already in progress.")
+ .timeConf(TimeUnit.MILLISECONDS)
+ .checkValue(_ >= 0, "Must be non-negative.")
+ .createWithDefaultString("30s")
val CLIENT_PUSH_STAGE_END_TIMEOUT: ConfigEntry[Long] =
buildConf("celeborn.client.push.stageEnd.timeout")
diff --git
a/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
b/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
index f1e34aa54f..b305539b65 100644
---
a/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
+++
b/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
@@ -679,6 +679,9 @@ object ControlMessages extends Logging {
.build().toByteArray
new TransportMessage(MessageType.REQUEST_SLOTS, payload)
+ case pb: PbRequestWorkers =>
+ new TransportMessage(MessageType.REQUEST_WORKERS, pb.toByteArray)
+
case RequestSlotsResponse(status, workerResource, packed) =>
val builder = PbRequestSlotsResponse.newBuilder()
.setStatus(status.getValue)
@@ -693,6 +696,9 @@ object ControlMessages extends Logging {
val payload = builder.build().toByteArray
new TransportMessage(MessageType.REQUEST_SLOTS_RESPONSE, payload)
+ case pb: PbRequestWorkersResponse =>
+ new TransportMessage(MessageType.REQUEST_WORKERS_RESPONSE,
pb.toByteArray)
+
case Revive(shuffleId, mapIds, reviveRequests, serdeVersion) =>
val builder = PbRevive.newBuilder()
.setShuffleId(shuffleId)
@@ -1177,6 +1183,12 @@ object ControlMessages extends Logging {
StatusCode.fromValue(pbRequestSlotsResponse.getStatus),
workerResource)
+ case REQUEST_WORKERS_VALUE =>
+ PbRequestWorkers.parseFrom(message.getPayload)
+
+ case REQUEST_WORKERS_RESPONSE_VALUE =>
+ PbRequestWorkersResponse.parseFrom(message.getPayload)
+
case CHANGE_LOCATION_VALUE =>
val pbRevive = PbRevive.parseFrom(message.getPayload)
val shuffleId = pbRevive.getShuffleId
diff --git
a/common/src/test/scala/org/apache/celeborn/common/CelebornConfSuite.scala
b/common/src/test/scala/org/apache/celeborn/common/CelebornConfSuite.scala
index 18d71e23eb..392d5edcda 100644
--- a/common/src/test/scala/org/apache/celeborn/common/CelebornConfSuite.scala
+++ b/common/src/test/scala/org/apache/celeborn/common/CelebornConfSuite.scala
@@ -24,6 +24,31 @@ import org.apache.celeborn.common.protocol.StorageInfo
class CelebornConfSuite extends CelebornFunSuite {
+ test("master split slot assign max workers must be positive") {
+ Seq("0", "-1").foreach { value =>
+ val error = intercept[IllegalArgumentException] {
+ new CelebornConf()
+ .set(MASTER_SPLIT_SLOT_ASSIGN_MAX_WORKERS.key, value)
+ .masterSplitSlotAssignMaxWorkers
+ }
+ assert(error.getMessage.contains("Must be positive."))
+ }
+ }
+
+ test("dynamic resource update time must be non-negative") {
+ assert(
+ new CelebornConf()
+ .set(CLIENT_SHUFFLE_DYNAMIC_RESOURCE_UPDATE_TIME.key, "0")
+ .clientShuffleDynamicResourceUpdateTime == 0L)
+
+ val error = intercept[IllegalArgumentException] {
+ new CelebornConf()
+ .set(CLIENT_SHUFFLE_DYNAMIC_RESOURCE_UPDATE_TIME.key, "-1ms")
+ .clientShuffleDynamicResourceUpdateTime
+ }
+ assert(error.getMessage.contains("Must be non-negative."))
+ }
+
test("JVMQuake thresholds should preserve configured time units") {
val conf = new CelebornConf()
.set(WORKER_JVM_QUAKE_DUMP_THRESHOLD.key, "30s")
diff --git
a/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
b/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
index c98f3bd498..a3e03f7c9b 100644
--- a/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
+++ b/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
@@ -29,9 +29,10 @@ import org.apache.celeborn.CelebornFunSuite
import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.client.{MasterEndpointResolver,
StaticMasterEndpointResolver}
import org.apache.celeborn.common.exception.CelebornException
-import org.apache.celeborn.common.identity.DefaultIdentityProvider
+import org.apache.celeborn.common.identity.{DefaultIdentityProvider,
UserIdentifier}
+import org.apache.celeborn.common.meta.WorkerInfo
import org.apache.celeborn.common.network.protocol.SerdeVersion
-import org.apache.celeborn.common.protocol.{PartitionLocation,
PbReviseLostShuffles, PbReviseLostShufflesResponse, TransportModuleConstants}
+import org.apache.celeborn.common.protocol.{PartitionLocation,
PbRequestWorkers, PbRequestWorkersResponse, PbReviseLostShuffles,
PbReviseLostShufflesResponse, StorageInfo, TransportModuleConstants}
import
org.apache.celeborn.common.protocol.message.ControlMessages.{GetReducerFileGroupResponse,
MapperEnd, ReviseLostShuffles, ReviseLostShufflesResponse}
import org.apache.celeborn.common.protocol.message.StatusCode
@@ -197,6 +198,34 @@ class UtilsSuite extends CelebornFunSuite {
mapperEnd.bytesWrittenPerPartition.array should contain
theSameElementsInOrderAs mapperEndTrans.bytesWrittenPerPartition
}
+ test("PbRequestWorkers messages convert with TransportMessage") {
+ val excludedWorker = new WorkerInfo("host1", 1001, 1002, 1003, 1004)
+ val request = PbRequestWorkers.newBuilder()
+ .setApplicationId("app-1")
+ .setUserIdentifier(
+ PbSerDeUtils.toPbUserIdentifier(new UserIdentifier("tenant", "user")))
+ .setMaxWorkers(10)
+ .setTagsExpr("tag-a,tag-b")
+ .setShouldReplicate(true)
+ .setAvailableStorageTypes(StorageInfo.LOCAL_DISK_MASK)
+ .addExcludedWorkerSet(PbSerDeUtils.toPbWorkerInfo(excludedWorker, true,
true))
+ .build()
+ val convertedRequest =
+
Utils.fromTransportMessage(Utils.toTransportMessage(request)).asInstanceOf[PbRequestWorkers]
+ assert(convertedRequest == request)
+
+ val response = PbRequestWorkersResponse.newBuilder()
+ .setStatus(StatusCode.SUCCESS.getValue)
+ .addWorkers(PbSerDeUtils.toPbWorkerInfo(excludedWorker, true,
true).toBuilder
+ .setNetworkLocation("/rack-1")
+ .build())
+ .build()
+ val convertedResponse =
+ Utils.fromTransportMessage(Utils.toTransportMessage(response))
+ .asInstanceOf[PbRequestWorkersResponse]
+ assert(convertedResponse == response)
+ }
+
test("ReviseLostShuffles class convert with pb") {
val req = ReviseLostShuffles("app-1", util.Arrays.asList[Integer](1, 2,
3), "req-1")
val reqTrans = Utils.fromTransportMessage(Utils.toTransportMessage(req))
diff --git a/docs/configuration/client.md b/docs/configuration/client.md
index be5422c236..dbb89561e4 100644
--- a/docs/configuration/client.md
+++ b/docs/configuration/client.md
@@ -111,8 +111,8 @@ license: |
| celeborn.client.shuffle.compression.codec | LZ4 | false | The codec used to
compress shuffle data. By default, Celeborn provides three codecs: `lz4`,
`zstd`, `none`. `none` means that shuffle compression is disabled. Since Flink
version 1.16, zstd is supported for Flink shuffle client. | 0.3.0 |
celeborn.shuffle.compression.codec,remote-shuffle.job.compression.codec |
| celeborn.client.shuffle.compression.zstd.level | 1 | false | Compression
level for Zstd compression codec, its value should be an integer between -5 and
22. Increasing the compression level will result in better compression at the
expense of more CPU and memory. | 0.3.0 |
celeborn.shuffle.compression.zstd.level |
| celeborn.client.shuffle.decompression.lz4.xxhash.instance |
<undefined> | false | Decompression XXHash instance for Lz4. Available
options: JNI, JAVASAFE, JAVAUNSAFE. | 0.3.2 | |
-| celeborn.client.shuffle.dynamicResourceEnabled | false | false | When
enabled, the ChangePartitionManager will obtain candidate workers from the
availableWorkers pool during heartbeats when worker resource change. | 0.6.0 |
|
-| celeborn.client.shuffle.dynamicResourceFactor | 0.5 | false | The
ChangePartitionManager will check whether (unavailable workers / shuffle
allocated workers) is more than the factor before obtaining candidate workers
from the requestSlots RPC response when
`celeborn.client.shuffle.dynamicResourceEnabled` set true | 0.6.0 | |
+| celeborn.client.shuffle.dynamicResource.updateTime | 30s | false | Minimum
interval after a worker-candidate refresh attempt completes before
ChangePartitionManager may try again when
`celeborn.client.shuffle.dynamicResourceEnabled` is true. Set to 0 to allow
each change-partition handling cycle to refresh when no refresh is already in
progress. | 1.0.0 | |
+| celeborn.client.shuffle.dynamicResourceEnabled | false | false | When
enabled, ChangePartitionManager refreshes endpoint-ready worker candidates from
the Master on demand while handling change-partition requests, and combines
them with workers already present in the shuffle snapshot. | 0.6.0 | |
| celeborn.client.shuffle.expired.checkInterval | 60s | false | Interval for
client to check expired shuffles. | 0.3.0 |
celeborn.shuffle.expired.checkInterval |
| celeborn.client.shuffle.integrityCheck.enabled | false | false | When
`true`, enables end-to-end integrity checks for Spark and Flink workloads. |
0.6.1 | |
| celeborn.client.shuffle.manager.port | 0 | false | Port used by the
LifecycleManager on the Driver. | 0.3.0 | celeborn.shuffle.manager.port |
diff --git a/docs/configuration/master.md b/docs/configuration/master.md
index a01bca3f69..86a4158668 100644
--- a/docs/configuration/master.md
+++ b/docs/configuration/master.md
@@ -85,6 +85,7 @@ license: |
| celeborn.master.slot.assign.maxWorkers | 10000 | false | Max workers that
slots of one shuffle can be allocated on. Will choose the smaller positive one
from Master side and Client side, see `celeborn.client.slot.assign.maxWorkers`.
| 0.3.1 | |
| celeborn.master.slot.assign.minWorkers | 100 | false | Min workers that
slots of one shuffle should be allocated on. Provided enough workers are
available. | 0.6.0 | |
| celeborn.master.slot.assign.policy | ROUNDROBIN | true | Policy for master
to assign slots. Built-in policies are roundrobin and loadaware. Additional
policies can be registered through the SlotsAssignStrategyProvider SPI.
Loadaware policy will be ignored when `HDFS` is enabled in
`celeborn.storage.availableTypes` | 0.3.0 | celeborn.slots.assign.policy |
+| celeborn.master.splitSlot.assign.maxWorkers | 500 | false | Maximum workers
returned by each dynamic candidate refresh. The request limit is the smaller
positive value of this setting and `celeborn.client.slot.assign.maxWorkers`.
For replicated shuffle, an effective limit of one is raised to two. Workers
already present in a shuffle snapshot are not counted against this limit. |
1.0.0 | |
| celeborn.master.userResourceConsumption.metrics.enabled | false | false |
Whether to enable resource consumption metrics. | 0.6.0 | |
| celeborn.master.userResourceConsumption.update.interval | 30s | false | Time
length for a window about compute user resource consumption. | 0.3.0 | |
| celeborn.master.workerUnavailableInfo.expireTimeout | 1800s | false | Worker
unavailable info would be cleared when the retention period is expired. Set -1
to disable the expiration. | 0.3.1 | |
diff --git a/docs/migration.md b/docs/migration.md
index b015ef3d86..7c6f654b7a 100644
--- a/docs/migration.md
+++ b/docs/migration.md
@@ -28,6 +28,11 @@ license: |
configuration. Existing `ROUNDROBIN` and `LOADAWARE` values remain supported
and require no
migration. See [Slots
allocation](./developers/slotsallocation.md#custom-strategies) for details.
+- Since 1.0.0, Celeborn removed
`celeborn.client.shuffle.dynamicResourceFactor`, which used the
+ unavailable-worker ratio to trigger a refresh. Dynamic worker refresh is now
performed on demand
+ while handling change-partition requests and throttled by
+ `celeborn.client.shuffle.dynamicResource.updateTime` (30s by default).
+
# Upgrading from 0.6 to 0.7
- Since 0.7.0, Celeborn removed `ReleaseSlots`.
diff --git
a/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
b/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
index ca1b67bee8..81a5027384 100644
---
a/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
+++
b/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
@@ -22,7 +22,7 @@ import java.net.BindException
import java.util
import java.util.{Map => JMap}
import java.util.Collections
-import java.util.concurrent.{ExecutorService, ScheduledFuture, TimeUnit}
+import java.util.concurrent.{ExecutorService, ScheduledFuture,
ThreadLocalRandom, TimeUnit}
import java.util.concurrent.atomic.AtomicBoolean
import java.util.function.ToLongFunction
@@ -206,6 +206,7 @@ private[celeborn] class Master(
private val tagsManager = new TagsManager(Option(configService))
private val slotsAssignMaxWorkers = conf.masterSlotAssignMaxWorkers
+ private val splitSlotAssignMaxWorkers = conf.masterSplitSlotAssignMaxWorkers
private val slotsAssignMinWorkers = conf.masterSlotAssignMinWorkers
private val slotsAssignExtraSlots = conf.masterSlotAssignExtraSlots
private val slotsAssignStrategyManager =
@@ -528,6 +529,11 @@ private[celeborn] class Master(
checkAuth(context, applicationId)
executeWithLeaderChecker(context, handleRequestSlots(context,
requestSlots))
+ case requestWorkers: PbRequestWorkers =>
+ logTrace(s"Received RequestWorkers request $requestWorkers.")
+ checkAuth(context, requestWorkers.getApplicationId)
+ executeWithLeaderChecker(context, handleRequestWorkers(context,
requestWorkers))
+
case pb: PbBatchUnregisterShuffles =>
val applicationId = pb.getAppId
val shuffleIds = pb.getShuffleIdsList.asScala.toList
@@ -953,22 +959,11 @@ private[celeborn] class Master(
return
}
- val numWorkers = Math.min(
- Math.max(
- if (requestSlots.shouldReplicate) 2 else 1,
- if (requestSlots.maxWorkers <= 0) slotsAssignMaxWorkers
- else Math.min(slotsAssignMaxWorkers, requestSlots.maxWorkers)),
- numAvailableWorkers)
- val startIndex = Random.nextInt(numAvailableWorkers)
- val selectedWorkers = new util.ArrayList[WorkerInfo](numWorkers)
- selectedWorkers.addAll(availableWorkers.subList(
- startIndex,
- Math.min(numAvailableWorkers, startIndex + numWorkers)))
- if (startIndex + numWorkers > numAvailableWorkers) {
- selectedWorkers.addAll(availableWorkers.subList(
- 0,
- startIndex + numWorkers - numAvailableWorkers))
- }
+ val selectedWorkers = selectWorkers(
+ availableWorkers,
+ requestSlots.maxWorkers,
+ slotsAssignMaxWorkers,
+ requestSlots.shouldReplicate)
// offer slots
val slots =
masterSource.sample(MasterSource.OFFER_SLOTS_TIME,
s"offerSlots-${Random.nextInt()}") {
@@ -1043,7 +1038,7 @@ private[celeborn] class Master(
s"extraSlots=$offerSlotsExtraSize"))
if (authEnabled) {
- pushApplicationMetaToWorkers(requestSlots, slots)
+ pushApplicationMetaToWorkers(requestSlots.applicationId, slots.keySet())
}
context.reply(RequestSlotsResponse(
StatusCode.SUCCESS,
@@ -1051,26 +1046,111 @@ private[celeborn] class Master(
requestSlots.packed))
}
+ private def selectWorkers(
+ candidates: util.List[WorkerInfo],
+ numSelectMax: Int,
+ numAssignMax: Int,
+ shouldReplicate: Boolean): util.List[WorkerInfo] = {
+ val numCandidates = candidates.size()
+ if (numCandidates == 0) {
+ return Collections.emptyList()
+ }
+
+ val numWorkers = Math.min(
+ Math.max(
+ if (shouldReplicate) 2 else 1,
+ if (numSelectMax <= 0) numAssignMax
+ else Math.min(numAssignMax, numSelectMax)),
+ numCandidates)
+ val startIndex = ThreadLocalRandom.current().nextInt(numCandidates)
+ val selectedWorkers = new util.ArrayList[WorkerInfo](numWorkers)
+ selectedWorkers.addAll(candidates.subList(
+ startIndex,
+ Math.min(numCandidates, startIndex + numWorkers)))
+ if (startIndex + numWorkers > numCandidates) {
+ selectedWorkers.addAll(candidates.subList(
+ 0,
+ startIndex + numWorkers - numCandidates))
+ }
+ selectedWorkers
+ }
+
+ def handleRequestWorkers(context: RpcCallContext, requestWorkers:
PbRequestWorkers): Unit = {
+
+ val excludedWorkerSet =
+ requestWorkers.getExcludedWorkerSetList.asScala
+ .map(PbSerDeUtils.fromPbWorkerInfo)
+ .toSet
+ var availableWorkers = workersAvailable(excludedWorkerSet)
+ if (conf.tagsEnabled) {
+ availableWorkers = tagsManager.getTaggedWorkers(
+ PbSerDeUtils.fromPbUserIdentifier(requestWorkers.getUserIdentifier),
+ requestWorkers.getTagsExpr,
+ availableWorkers)
+ }
+ if (availableWorkers.isEmpty) {
+ logWarning(
+ s"Offer workers for ${requestWorkers.getApplicationId} failed due to
no available workers.")
+ context.reply(PbRequestWorkersResponse.newBuilder()
+ .setStatus(StatusCode.WORKER_EXCLUDED.getValue)
+ .build())
+ return
+ }
+
+ val candidates = new util.ArrayList[WorkerInfo]()
+ availableWorkers.asScala
+ .filter { worker =>
+
!StorageInfo.localDiskAvailable(requestWorkers.getAvailableStorageTypes) ||
+ worker.haveDisk
+ }
+ .foreach(candidates.add)
+ val selectedWorkers = selectWorkers(
+ candidates,
+ requestWorkers.getMaxWorkers,
+ splitSlotAssignMaxWorkers,
+ requestWorkers.getShouldReplicate)
+ if (selectedWorkers.isEmpty) {
+ logWarning(
+ s"Offer workers for ${requestWorkers.getApplicationId} failed due to
no eligible workers.")
+ context.reply(PbRequestWorkersResponse.newBuilder()
+ .setStatus(StatusCode.SLOT_NOT_AVAILABLE.getValue)
+ .build())
+ return
+ }
+
+ if (authEnabled) {
+ pushApplicationMetaToWorkers(requestWorkers.getApplicationId,
selectedWorkers)
+ }
+ context.reply(PbRequestWorkersResponse.newBuilder()
+ .setStatus(StatusCode.SUCCESS.getValue)
+ .addAllWorkers(
+ selectedWorkers.asScala.map { worker =>
+ PbSerDeUtils.toPbWorkerInfo(worker, true, true).toBuilder
+ .setNetworkLocation(worker.networkLocation)
+ .build()
+ }.asJava)
+ .build())
+ }
+
def pushApplicationMetaToWorkers(
- requestSlots: RequestSlots,
- slots: util.Map[WorkerInfo, (util.List[PartitionLocation],
util.List[PartitionLocation])])
- : Unit = {
+ applicationId: String,
+ workers: util.Collection[WorkerInfo]): Unit = {
// Pass application registration information to the workers
val pbApplicationMeta = PbApplicationMeta.newBuilder()
- .setAppId(requestSlots.applicationId)
- .setSecret(secretRegistry.getSecretKey(requestSlots.applicationId))
+ .setAppId(applicationId)
+ .setSecret(secretRegistry.getSecretKey(applicationId))
.build()
val transportMessage =
new TransportMessage(MessageType.APPLICATION_META,
pbApplicationMeta.toByteArray)
val workerSet = workersAssignedToApp.computeIfAbsent(
- requestSlots.applicationId,
+ applicationId,
new util.function.Function[String, util.Set[WorkerInfo]] {
override def apply(key: String): util.Set[WorkerInfo] =
util.Collections.newSetFromMap(JavaUtils.newConcurrentHashMap[
WorkerInfo,
java.lang.Boolean]())
})
- slots.keySet().asScala.foreach { worker =>
+ workers.asScala.foreach { worker =>
// The app meta info is send to a Worker only if it wasn't previously
sent.
if (workerSet.add(worker)) {
sendApplicationMetaExecutor.submit(new Runnable {
diff --git
a/master/src/test/scala/org/apache/celeborn/service/deploy/master/MasterSuite.scala
b/master/src/test/scala/org/apache/celeborn/service/deploy/master/MasterSuite.scala
index 48ad9fe8d9..0067a69fa5 100644
---
a/master/src/test/scala/org/apache/celeborn/service/deploy/master/MasterSuite.scala
+++
b/master/src/test/scala/org/apache/celeborn/service/deploy/master/MasterSuite.scala
@@ -20,6 +20,8 @@ package org.apache.celeborn.service.deploy.master
import java.nio.file.Files
import java.util
+import scala.collection.JavaConverters._
+
import org.mockito.ArgumentCaptor
import org.mockito.Mockito.{mock, verify, when}
import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}
@@ -27,13 +29,16 @@ import org.scalatest.funsuite.AnyFunSuite
import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.identity.UserIdentifier
+import org.apache.celeborn.common.meta.{DiskInfo, WorkerInfo}
import org.apache.celeborn.common.network.client.{RpcResponseCallback,
TransportClient}
-import org.apache.celeborn.common.protocol.{PbApplicationMetaRequest,
PbCheckForWorkerTimeout, PbRegisterWorker}
+import org.apache.celeborn.common.protocol.{PbApplicationMetaRequest,
PbCheckForWorkerTimeout, PbRegisterWorker, PbRequestWorkers,
PbRequestWorkersResponse}
+import org.apache.celeborn.common.protocol.StorageInfo
import
org.apache.celeborn.common.protocol.message.ControlMessages.{RequestSlots,
RequestSlotsResponse, ReviseLostShuffles}
import org.apache.celeborn.common.protocol.message.StatusCode
+import org.apache.celeborn.common.quota.ResourceConsumption
import org.apache.celeborn.common.rpc.{RpcAddress, RpcCallContext}
import org.apache.celeborn.common.rpc.netty.{NettyRpcEnv,
RemoteNettyRpcCallContext}
-import org.apache.celeborn.common.util.{CelebornExitKind, ThreadUtils}
+import org.apache.celeborn.common.util.{CelebornExitKind, PbSerDeUtils,
ThreadUtils}
class MasterSuite extends AnyFunSuite
with BeforeAndAfterAll
@@ -52,6 +57,16 @@ class MasterSuite extends AnyFunSuite
client)
}
+ private def requestWorkers(
+ master: Master,
+ request: PbRequestWorkers): PbRequestWorkersResponse = {
+ val context = mock(classOf[RpcCallContext])
+ val captor = ArgumentCaptor.forClass(classOf[Any])
+ master.handleRequestWorkers(context, request)
+ verify(context).reply(captor.capture())
+ captor.getValue.asInstanceOf[PbRequestWorkersResponse]
+ }
+
def getTmpDir(): String = {
val tmpDir = Files.createTempDirectory(null).toFile
tmpDir.deleteOnExit()
@@ -212,6 +227,183 @@ class MasterSuite extends AnyFunSuite
master.rpcEnv.shutdown()
}
+ test("handleRequestWorkers returns one failure when no workers are
available") {
+ val conf = new CelebornConf()
+ conf.set(CelebornConf.HA_ENABLED.key, "false")
+ conf.set(CelebornConf.MASTER_HTTP_PORT.key, selectRandomPort().toString)
+ val masterArgs = new MasterArguments(
+ Array("-h", "localhost", "-p", selectRandomPort().toString),
+ conf)
+ val master = new Master(conf, masterArgs)
+ val request = PbRequestWorkers.newBuilder()
+ .setApplicationId("app1")
+ .setUserIdentifier(
+ PbSerDeUtils.toPbUserIdentifier(new UserIdentifier("tenant", "user")))
+ .setMaxWorkers(10)
+ .build()
+ val response = requestWorkers(master, request)
+ assert(StatusCode.fromValue(response.getStatus) ===
StatusCode.WORKER_EXCLUDED)
+ assert(response.getWorkersCount === 0)
+ master.rpcEnv.shutdown()
+ }
+
+ test("handleRequestWorkers applies limits without allocating slots") {
+ val conf = new CelebornConf()
+ conf.set(CelebornConf.HA_ENABLED.key, "false")
+ conf.set(CelebornConf.MASTER_HTTP_PORT.key, selectRandomPort().toString)
+ conf.set(CelebornConf.MASTER_SPLIT_SLOT_ASSIGN_MAX_WORKERS.key, "2")
+ val masterArgs = new MasterArguments(
+ Array("-h", "localhost", "-p", selectRandomPort().toString),
+ conf)
+ val master = new Master(conf, masterArgs)
+ val workers = (1 to 4).map { index =>
+ val worker =
+ new WorkerInfo(s"host$index", 1000 + index, 2000 + index, 3000 +
index, 4000 + index)
+ worker.networkLocation = s"/rack-$index"
+ worker
+ }
+ master.statusSystem.availableWorkers.addAll(workers.asJava)
+ val registeredShufflesBefore =
master.statusSystem.registeredAppAndShuffles.size()
+ val request = PbRequestWorkers.newBuilder()
+ .setApplicationId("app1")
+ .setUserIdentifier(
+ PbSerDeUtils.toPbUserIdentifier(new UserIdentifier("tenant", "user")))
+ .setMaxWorkers(3)
+ .setAvailableStorageTypes(StorageInfo.MEMORY_MASK)
+ .build()
+ val response = requestWorkers(master, request)
+ val responseWorkers =
+ response.getWorkersList.asScala.map(PbSerDeUtils.fromPbWorkerInfo).toSet
+ assert(StatusCode.fromValue(response.getStatus) === StatusCode.SUCCESS)
+ assert(responseWorkers.size === 2)
+ assert(responseWorkers.subsetOf(workers.toSet))
+ val networkLocationsByHost = workers.map(worker => worker.host ->
worker.networkLocation).toMap
+ response.getWorkersList.asScala.foreach { worker =>
+ assert(worker.getNetworkLocation ===
networkLocationsByHost(worker.getHost))
+ }
+ assert(master.statusSystem.registeredAppAndShuffles.size() ===
registeredShufflesBefore)
+
+ val defaultLimitResponse = requestWorkers(master,
request.toBuilder.setMaxWorkers(0).build())
+ assert(StatusCode.fromValue(defaultLimitResponse.getStatus) ===
StatusCode.SUCCESS)
+ assert(defaultLimitResponse.getWorkersCount === 2)
+
+ val singleWorkerResponse = requestWorkers(master,
request.toBuilder.setMaxWorkers(1).build())
+ assert(StatusCode.fromValue(singleWorkerResponse.getStatus) ===
StatusCode.SUCCESS)
+ assert(singleWorkerResponse.getWorkersCount === 1)
+
+ val replicatedResponse = requestWorkers(
+ master,
+ request.toBuilder
+ .setMaxWorkers(1)
+ .setShouldReplicate(true)
+ .build())
+ assert(StatusCode.fromValue(replicatedResponse.getStatus) ===
StatusCode.SUCCESS)
+ assert(replicatedResponse.getWorkersCount === 2)
+
+ val excludedWorkers = workers.take(3)
+ val excludedResponse = requestWorkers(
+ master,
+ request.toBuilder
+ .addAllExcludedWorkerSet(
+ excludedWorkers.map(PbSerDeUtils.toPbWorkerInfo(_, true,
true)).asJava)
+ .build())
+ assert(StatusCode.fromValue(excludedResponse.getStatus) ===
StatusCode.SUCCESS)
+ assert(
+
excludedResponse.getWorkersList.asScala.map(PbSerDeUtils.fromPbWorkerInfo).toSet
===
+ workers.drop(3).toSet)
+
+ val partialReplicaResponse = requestWorkers(
+ master,
+ request.toBuilder
+ .setShouldReplicate(true)
+ .addAllExcludedWorkerSet(
+ excludedWorkers.map(PbSerDeUtils.toPbWorkerInfo(_, true,
true)).asJava)
+ .build())
+ assert(
+ StatusCode.fromValue(partialReplicaResponse.getStatus) ===
+ StatusCode.SUCCESS)
+ assert(partialReplicaResponse.getWorkersCount === 1)
+
+ val taggedResponse =
+ requestWorkers(master,
request.toBuilder.setTagsExpr("missing-tag").build())
+ assert(StatusCode.fromValue(taggedResponse.getStatus) ===
StatusCode.WORKER_EXCLUDED)
+ assert(taggedResponse.getWorkersCount === 0)
+
+ master.rpcEnv.shutdown()
+ }
+
+ test("handleRequestWorkers applies storage eligibility") {
+ val conf = new CelebornConf()
+ conf.set(CelebornConf.HA_ENABLED.key, "false")
+ conf.set(CelebornConf.MASTER_HTTP_PORT.key, selectRandomPort().toString)
+ val masterArgs = new MasterArguments(
+ Array("-h", "localhost", "-p", selectRandomPort().toString),
+ conf)
+ val master = new Master(conf, masterArgs)
+
+ def workerWithDisk(host: String, port: Int): WorkerInfo = {
+ val disk = new DiskInfo("/disk", 1024L, 0L, 0L, 0L, StorageInfo.Type.HDD)
+ val disks = new util.HashMap[String, DiskInfo]()
+ disks.put(disk.mountPoint, disk)
+ new WorkerInfo(
+ host,
+ port,
+ port + 1,
+ port + 2,
+ port + 3,
+ port + 4,
+ disks,
+ new util.HashMap[UserIdentifier, ResourceConsumption]())
+ }
+
+ val diskWorker = workerWithDisk("disk", 1000)
+ val disklessWorker = new WorkerInfo("diskless", 2000, 2001, 2002, 2003)
+ val workers = Seq(diskWorker, disklessWorker)
+ master.statusSystem.availableWorkers.addAll(workers.asJava)
+
+ val request = PbRequestWorkers.newBuilder()
+ .setApplicationId("app1")
+ .setUserIdentifier(
+ PbSerDeUtils.toPbUserIdentifier(new UserIdentifier("tenant", "user")))
+ .setMaxWorkers(10)
+ .build()
+
+ val localResponse = requestWorkers(
+ master,
+ request.toBuilder
+ .setAvailableStorageTypes(StorageInfo.LOCAL_DISK_MASK)
+ .build())
+ assert(StatusCode.fromValue(localResponse.getStatus) ===
StatusCode.SUCCESS)
+ assert(localResponse.getWorkersList.asScala.map(_.getHost).toSet ===
Set(diskWorker.host))
+
+ val noLocalDiskResponse = requestWorkers(
+ master,
+ request.toBuilder
+ .setAvailableStorageTypes(StorageInfo.LOCAL_DISK_MASK)
+ .addExcludedWorkerSet(PbSerDeUtils.toPbWorkerInfo(diskWorker, true,
true))
+ .build())
+ assert(StatusCode.fromValue(noLocalDiskResponse.getStatus) ===
StatusCode.SLOT_NOT_AVAILABLE)
+ assert(noLocalDiskResponse.getWorkersCount === 0)
+
+ val remoteResponse = requestWorkers(
+ master,
+ request.toBuilder
+ .setAvailableStorageTypes(StorageInfo.HDFS_MASK)
+ .build())
+ assert(StatusCode.fromValue(remoteResponse.getStatus) ===
StatusCode.SUCCESS)
+ assert(remoteResponse.getWorkersList.asScala.map(_.getHost).toSet ===
workers.map(_.host).toSet)
+
+ val mixedResponse = requestWorkers(
+ master,
+ request.toBuilder
+ .setAvailableStorageTypes(StorageInfo.MEMORY_MASK |
StorageInfo.LOCAL_DISK_MASK)
+ .build())
+ assert(StatusCode.fromValue(mixedResponse.getStatus) ===
StatusCode.SUCCESS)
+ assert(mixedResponse.getWorkersList.asScala.map(_.getHost).toSet ===
Set(diskWorker.host))
+
+ master.rpcEnv.shutdown()
+ }
+
test("PbApplicationMetaRequest rejects a caller requesting another
application's secret") {
val conf = new CelebornConf()
val randomMasterPort = selectRandomPort()
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/client/ChangePartitionManagerUpdateWorkersSuite.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/client/ChangePartitionManagerUpdateWorkersSuite.scala
index 6e4b081230..3ecdf8a282 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/client/ChangePartitionManagerUpdateWorkersSuite.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/client/ChangePartitionManagerUpdateWorkersSuite.scala
@@ -22,6 +22,10 @@ import java.util.Collections
import scala.collection.JavaConverters.{collectionAsScalaIterableConverter,
mapAsScalaMapConverter}
+import org.scalatest.concurrent.Eventually.eventually
+import org.scalatest.concurrent.Futures.{interval, timeout}
+import org.scalatest.time.SpanSugar.convertIntToGrainOfTime
+
import org.apache.celeborn.client.{ChangePartitionManager,
ChangePartitionRequest, LifecycleManager, WithShuffleClientSuite}
import org.apache.celeborn.client.LifecycleManager.ShuffleFailedWorkers
import org.apache.celeborn.common.CelebornConf
@@ -57,7 +61,7 @@ class ChangePartitionManagerUpdateWorkersSuite extends
WithShuffleClientSuite
conf.set(CelebornConf.CLIENT_PUSH_MAX_REVIVE_TIMES.key, "3")
.set(CelebornConf.CLIENT_BATCH_HANDLE_CHANGE_PARTITION_ENABLED.key,
"false")
.set(CelebornConf.CLIENT_SHUFFLE_DYNAMIC_RESOURCE_ENABLED.key, "true")
- .set(CelebornConf.CLIENT_SHUFFLE_DYNAMIC_RESOURCE_FACTOR.key, "0.0")
+ .set(CelebornConf.CLIENT_SHUFFLE_DYNAMIC_RESOURCE_UPDATE_TIME.key, "0")
val lifecycleManager: LifecycleManager = new LifecycleManager(APP, conf)
val changePartitionManager: ChangePartitionManager =
@@ -100,7 +104,8 @@ class ChangePartitionManagerUpdateWorkersSuite extends
WithShuffleClientSuite
setUpWorkers(workerConfForAdding, 2)
assert(workerInfos.size == 3)
- 0 until 10 foreach { partitionId: Int =>
+ var partitionId = 0
+ eventually(timeout(10.seconds), interval(100.milliseconds)) {
val req = ChangePartitionRequest(
null,
shuffleId,
@@ -115,9 +120,9 @@ class ChangePartitionManagerUpdateWorkersSuite extends
WithShuffleClientSuite
shuffleId,
Array(req),
lifecycleManager.commitManager.isSegmentGranularityVisible(shuffleId))
+ partitionId += 1
+ assert(lifecycleManager.workerSnapshots(shuffleId).size() > 1)
}
- Thread.sleep(5000)
- assert(lifecycleManager.workerSnapshots(shuffleId).size() > 1)
lifecycleManager.stop()
}
@@ -131,7 +136,7 @@ class ChangePartitionManagerUpdateWorkersSuite extends
WithShuffleClientSuite
conf.set(CelebornConf.CLIENT_PUSH_MAX_REVIVE_TIMES.key, "3")
.set(CelebornConf.CLIENT_BATCH_HANDLE_CHANGE_PARTITION_ENABLED.key,
"false")
.set(CelebornConf.CLIENT_SHUFFLE_DYNAMIC_RESOURCE_ENABLED.key, "true")
- .set(CelebornConf.CLIENT_SHUFFLE_DYNAMIC_RESOURCE_FACTOR.key, "0.5")
+ .set(CelebornConf.CLIENT_SHUFFLE_DYNAMIC_RESOURCE_UPDATE_TIME.key, "0")
val lifecycleManager: LifecycleManager = new LifecycleManager(APP, conf)
val changePartitionManager: ChangePartitionManager =
@@ -202,7 +207,8 @@ class ChangePartitionManagerUpdateWorkersSuite extends
WithShuffleClientSuite
setUpWorkers(workerConfForAdding, 1)
assert(workerInfos.size == 3)
- 0 until 10 foreach { partitionId: Int =>
+ var partitionId = 0
+ eventually(timeout(10.seconds), interval(100.milliseconds)) {
val req = ChangePartitionRequest(
null,
shuffleId,
@@ -217,27 +223,27 @@ class ChangePartitionManagerUpdateWorkersSuite extends
WithShuffleClientSuite
shuffleId,
Array(req),
lifecycleManager.commitManager.isSegmentGranularityVisible(shuffleId))
+ partitionId += 1
+
+ val snapshotCandidates =
+ lifecycleManager
+ .workerSnapshots(shuffleId)
+ .asScala
+ .values
+ .map(_.workerInfo)
+ .filter(lifecycleManager.workerStatusTracker.workerAvailable)
+ assert(snapshotCandidates.size == 2)
}
-
- val snapshotCandidates =
- lifecycleManager
- .workerSnapshots(shuffleId)
- .asScala
- .values
- .map(_.workerInfo)
- .filter(lifecycleManager.workerStatusTracker.workerAvailable)
-
- assert(snapshotCandidates.size == 2)
lifecycleManager.stop()
}
- test("test changePartition with available workers and factor") {
+ test("test changePartition honors worker pool update time") {
val shuffleId = nextShuffleId
val conf = celebornConf.clone
conf.set(CelebornConf.CLIENT_PUSH_MAX_REVIVE_TIMES.key, "3")
.set(CelebornConf.CLIENT_BATCH_HANDLE_CHANGE_PARTITION_ENABLED.key,
"false")
.set(CelebornConf.CLIENT_SHUFFLE_DYNAMIC_RESOURCE_ENABLED.key, "true")
- .set(CelebornConf.CLIENT_SHUFFLE_DYNAMIC_RESOURCE_FACTOR.key, "1.0")
+ .set(CelebornConf.CLIENT_SHUFFLE_DYNAMIC_RESOURCE_UPDATE_TIME.key,
"2000s")
val lifecycleManager: LifecycleManager = new LifecycleManager(APP, conf)
val changePartitionManager: ChangePartitionManager =
@@ -278,10 +284,28 @@ class ChangePartitionManagerUpdateWorkersSuite extends
WithShuffleClientSuite
}
assert(lifecycleManager.workerSnapshots(shuffleId).size() == workerNum)
+ // The first change refreshes the pool and starts the update-time window.
+ val initialRequest = ChangePartitionRequest(
+ null,
+ shuffleId,
+ 0,
+ -1,
+ null,
+ None)
+ changePartitionManager.changePartitionRequests.computeIfAbsent(
+ shuffleId,
+ changePartitionManager.rpcContextRegisterFunc)
+ changePartitionManager.handleRequestPartitions(
+ shuffleId,
+ Array(initialRequest),
+ lifecycleManager.commitManager.isSegmentGranularityVisible(shuffleId))
+ assert(lifecycleManager.workerSnapshots(shuffleId).size() == workerNum)
+
// total workerNum is 1 + 2 = 3 now
setUpWorkers(workerConfForAdding, 2)
assert(workerInfos.size == 3)
+ // The refresh interval has not elapsed, so the newly registered workers
are not candidates.
0 until 10 foreach { partitionId: Int =>
val req = ChangePartitionRequest(
null,
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryReviveTest.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryReviveTest.scala
index ffb8e7721d..31d37cf244 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryReviveTest.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryReviveTest.scala
@@ -68,7 +68,7 @@ class RetryReviveTest extends AnyFunSuite
.set(s"spark.${CelebornConf.TEST_CLIENT_RETRY_REVIVE.key}", "true")
.set(s"spark.${CelebornConf.CLIENT_PUSH_MAX_REVIVE_TIMES.key}", "3")
.set(s"spark.${CelebornConf.CLIENT_SHUFFLE_DYNAMIC_RESOURCE_ENABLED.key}",
"true")
-
.set(s"spark.${CelebornConf.CLIENT_SHUFFLE_DYNAMIC_RESOURCE_FACTOR.key}", "0")
+
.set(s"spark.${CelebornConf.CLIENT_SHUFFLE_DYNAMIC_RESOURCE_UPDATE_TIME.key}",
"0")
.set(s"spark.${CelebornConf.MASTER_SLOT_ASSIGN_EXTRA_SLOTS.key}", "0")
.setAppName("celeborn-demo").setMaster("local[2]")
val ss = SparkSession.builder()
