You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by ni...@apache.org on 2026/08/20 15:16:40 UTC
(celeborn) branch main updated: [CELEBORN-2060] Refactor master slot allocation with strategy abstraction
This is an automated email from the ASF dual-hosted git repository.
SteNicholas 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 7b050089b [CELEBORN-2060] Refactor master slot allocation with strategy abstraction
7b050089b is described below
commit 7b050089b837055a3813b26f05b1fce88aa32e0b
Author: Kalvin2077 <wk...@outlook.com>
AuthorDate: Thu Aug 20 23:15:54 2026 +0800
[CELEBORN-2060] Refactor master slot allocation with strategy abstraction
### What changes were proposed in this pull request?
This PR refactors master-side slot allocation so that strategy-specific slot-budget computation is
separated from the shared allocation workflow. The existing round-robin and load-aware algorithms
are moved behind `SlotsAssignStrategy`, while `SlotsAllocator` retains the common progressive
allocation and best-effort fallback behavior.
The change is intentionally structural. Existing names and allocation behavior are preserved where
possible so that the old implementation can be compared directly with the refactored code.
The refactor also exposed several potential correctness issues that are intentionally left for
separate follow-up changes:
- With replication enabled, remote-only storage and rack awareness disabled, primary and replica
locations may be assigned to the same worker.
- During budget-free rack-aware fallback, a replica may reuse the primary worker's `StorageInfo`,
which can produce an incorrect mount-point hint and slot accounting when workers have different
mount points.
- During budget-free rack-aware fallback with local storage enabled, replica selection does not
verify that the selected worker has a local disk, so allocation may be deferred to a worker that
cannot satisfy the requested local placement.
### Why are the changes needed?
`SlotsAllocator` previously mixed strategy-specific budget calculation with worker selection,
storage selection, rack awareness, interruption awareness, and fallback handling. Separating the
strategy boundary makes the allocation stages easier to follow and allows strategies to evolve
without duplicating the common allocation workflow.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [ ] Yes
### How was this patch tested?
- `build/sbt "celeborn-master/testOnly org.apache.celeborn.service.deploy.master.slotsalloc.*"`
Closes #3781 from Kalvin2077/refactor/slots-alloc-cleanup.
Authored-by: Kalvin2077 <wk...@outlook.com>
Signed-off-by: Nicholas Jiang <pr...@163.com>
---
.../org/apache/celeborn/common/CelebornConf.scala | 29 +-
docs/configuration/master.md | 12 +-
docs/developers/sbt.md | 6 +-
docs/developers/slotsallocation.md | 35 +
docs/migration.md | 7 +
.../service/deploy/master/SlotsAllocator.java | 854 ---------------------
.../slotsalloc/LoadAwareSlotsAssignStrategy.java | 262 +++++++
.../LoadAwareSlotsAssignStrategyProvider.java | 39 +
.../slotsalloc/RoundRobinSlotsAssignStrategy.java | 51 ++
.../RoundRobinSlotsAssignStrategyProvider.java | 34 +
.../deploy/master/slotsalloc/SlotsAllocator.java | 595 ++++++++++++++
.../master/slotsalloc/SlotsAssignStrategy.java | 41 +
.../slotsalloc/SlotsAssignStrategyManager.java | 176 +++++
.../slotsalloc/SlotsAssignStrategyProvider.java | 46 ++
.../deploy/master/slotsalloc/UsableDiskInfo.java | 35 +
...y.master.slotsalloc.SlotsAssignStrategyProvider | 19 +
.../celeborn/service/deploy/master/Master.scala | 43 +-
.../{ => slotsalloc}/BuildStorageInfoSuiteJ.java | 262 +++----
.../LoadAwareSlotsAssignStrategySuiteJ.java | 134 ++++
.../SlotsAllocatorJmhBenchmark.java | 11 +-
.../SlotsAllocatorRackAwareSuiteJ.java | 48 +-
.../{ => slotsalloc}/SlotsAllocatorSuiteJ.java | 107 ++-
.../SlotsAssignStrategyManagerSuiteJ.java | 382 +++++++++
.../test/NullSlotsAssignStrategyProvider.java | 38 +
.../ThirdPartySlotsAssignStrategyProvider.java | 88 +++
.../test/TransientSlotsAssignStrategyProvider.java | 46 ++
...y.master.slotsalloc.SlotsAssignStrategyProvider | 20 +
27 files changed, 2288 insertions(+), 1132 deletions(-)
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 6dce12fec..2a928c611 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -656,6 +656,10 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
// //////////////////////////////////////////////////////
// Master //
// //////////////////////////////////////////////////////
+ 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")
def masterSlotAssignPolicy: SlotsAssignPolicy =
SlotsAssignPolicy.valueOf(get(MASTER_SLOT_ASSIGN_POLICY))
@@ -3098,13 +3102,12 @@ object CelebornConf extends Logging {
.withAlternative("celeborn.slots.assign.policy")
.categories("master")
.version("0.3.0")
- .doc("Policy for master to assign slots, Celeborn supports two types of policy: roundrobin and loadaware. " +
+ .doc("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`")
+ .dynamic
.stringConf
.transform(_.toUpperCase(Locale.ROOT))
- .checkValues(Set(
- SlotsAssignPolicy.ROUNDROBIN.name,
- SlotsAssignPolicy.LOADAWARE.name))
.createWithDefault(SlotsAssignPolicy.ROUNDROBIN.name)
val MASTER_SLOT_ASSIGN_INTERRUPTION_AWARE: ConfigEntry[Boolean] =
@@ -3132,8 +3135,10 @@ object CelebornConf extends Logging {
.categories("master")
.doc("This configuration is a guidance for load-aware slot allocation algorithm. " +
"This value is control how many disk groups will be created.")
+ .dynamic
.version("0.3.0")
.intConf
+ .checkValue(_ > 0, "Value must be positive")
.createWithDefault(5)
val MASTER_SLOT_ASSIGN_LOADAWARE_DISKGROUP_GRADIENT: ConfigEntry[Double] =
@@ -3142,8 +3147,12 @@ object CelebornConf extends Logging {
.categories("master")
.doc("This value means how many more workload will be placed into a faster disk group " +
"than a slower group.")
+ .dynamic
.version("0.3.0")
.doubleConf
+ .checkValue(
+ value => java.lang.Double.isFinite(value) && value >= 0.0,
+ "Value must be finite and non-negative")
.createWithDefault(0.1)
val MASTER_SLOT_ASSIGN_LOADAWARE_FLUSHTIME_WEIGHT: ConfigEntry[Double] =
@@ -3152,8 +3161,12 @@ object CelebornConf extends Logging {
.categories("master")
.doc(
"Weight of average flush time when calculating ordering in load-aware assignment strategy")
+ .dynamic
.version("0.3.0")
.doubleConf
+ .checkValue(
+ value => java.lang.Double.isFinite(value) && value >= 0.0,
+ "Value must be finite and non-negative")
.createWithDefault(0)
val MASTER_SLOT_ASSIGN_LOADAWARE_FETCHTIME_WEIGHT: ConfigEntry[Double] =
@@ -3162,8 +3175,12 @@ object CelebornConf extends Logging {
.categories("master")
.doc(
"Weight of average fetch time when calculating ordering in load-aware assignment strategy")
+ .dynamic
.version("0.3.0")
.doubleConf
+ .checkValue(
+ value => java.lang.Double.isFinite(value) && value >= 0.0,
+ "Value must be finite and non-negative")
.createWithDefault(1)
val MASTER_SLOT_ASSIGN_LOADAWARE_ACTIVE_SLOTS_WEIGHT: ConfigEntry[Double] =
@@ -3171,8 +3188,12 @@ object CelebornConf extends Logging {
.categories("master")
.doc(
"Weight of active slots when calculating ordering in load-aware assignment strategy")
+ .dynamic
.version("0.7.0")
.doubleConf
+ .checkValue(
+ value => java.lang.Double.isFinite(value) && value >= 0.0,
+ "Value must be finite and non-negative")
.createWithDefault(0)
val MASTER_SLOT_ASSIGN_EXTRA_SLOTS: ConfigEntry[Int] =
diff --git a/docs/configuration/master.md b/docs/configuration/master.md
index 1f1c900c9..a01bca3f6 100644
--- a/docs/configuration/master.md
+++ b/docs/configuration/master.md
@@ -77,14 +77,14 @@ license: |
| celeborn.master.slot.assign.extraSlots | 2 | false | Extra slots number when master assign slots. Provided enough workers are available. | 0.3.0 | celeborn.slots.assign.extraSlots |
| celeborn.master.slot.assign.interruptionAware | false | false | If this is set to true, Celeborn master will prioritize partition placement on workers that are not in scope for maintenance soon. | 0.7.0 | |
| celeborn.master.slot.assign.interruptionAware.threshold | 50 | false | This controls what percentage of hosts would be selected for slot selection in the first iteration of creating partitions. Default is 50%. | 0.7.0 | |
-| celeborn.master.slot.assign.loadAware.activeSlotsWeight | 0.0 | false | Weight of active slots when calculating ordering in load-aware assignment strategy | 0.7.0 | |
-| celeborn.master.slot.assign.loadAware.diskGroupGradient | 0.1 | false | This value means how many more workload will be placed into a faster disk group than a slower group. | 0.3.0 | celeborn.slots.assign.loadAware.diskGroupGradient |
-| celeborn.master.slot.assign.loadAware.fetchTimeWeight | 1.0 | false | Weight of average fetch time when calculating ordering in load-aware assignment strategy | 0.3.0 | celeborn.slots.assign.loadAware.fetchTimeWeight |
-| celeborn.master.slot.assign.loadAware.flushTimeWeight | 0.0 | false | Weight of average flush time when calculating ordering in load-aware assignment strategy | 0.3.0 | celeborn.slots.assign.loadAware.flushTimeWeight |
-| celeborn.master.slot.assign.loadAware.numDiskGroups | 5 | false | This configuration is a guidance for load-aware slot allocation algorithm. This value is control how many disk groups will be created. | 0.3.0 | celeborn.slots.assign.loadAware.numDiskGroups |
+| celeborn.master.slot.assign.loadAware.activeSlotsWeight | 0.0 | true | Weight of active slots when calculating ordering in load-aware assignment strategy | 0.7.0 | |
+| celeborn.master.slot.assign.loadAware.diskGroupGradient | 0.1 | true | This value means how many more workload will be placed into a faster disk group than a slower group. | 0.3.0 | celeborn.slots.assign.loadAware.diskGroupGradient |
+| celeborn.master.slot.assign.loadAware.fetchTimeWeight | 1.0 | true | Weight of average fetch time when calculating ordering in load-aware assignment strategy | 0.3.0 | celeborn.slots.assign.loadAware.fetchTimeWeight |
+| celeborn.master.slot.assign.loadAware.flushTimeWeight | 0.0 | true | Weight of average flush time when calculating ordering in load-aware assignment strategy | 0.3.0 | celeborn.slots.assign.loadAware.flushTimeWeight |
+| celeborn.master.slot.assign.loadAware.numDiskGroups | 5 | true | This configuration is a guidance for load-aware slot allocation algorithm. This value is control how many disk groups will be created. | 0.3.0 | celeborn.slots.assign.loadAware.numDiskGroups |
| 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 | false | Policy for master to assign slots, Celeborn supports two types of policy: roundrobin and loadaware. Loadaware policy will be ignored when `HDFS` is enabled in `celeborn.storage.availableTypes` | 0.3.0 | celeborn.slots.assign.policy |
+| 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.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/developers/sbt.md b/docs/developers/sbt.md
index 63506e6f2..f10d2a948 100644
--- a/docs/developers/sbt.md
+++ b/docs/developers/sbt.md
@@ -187,17 +187,17 @@ $ ./build/sbt
You can run a single test suite using the `testOnly` command. For example, to run the `SlotsAllocatorSuiteJ`:
```
-> testOnly org.apache.celeborn.service.deploy.master.SlotsAllocatorSuiteJ
+> testOnly org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAllocatorSuiteJ
```
The `testOnly` command accepts wildcards; e.g., you can also run the `SlotsAllocatorSuiteJ` with:
```
> testOnly *SlotsAllocatorSuiteJ
```
-Or you could run all the tests in the `master` package:
+Or you could run all the tests in the `slotsalloc` package:
```
-> testOnly org.apache.celeborn.service.deploy.master.*
+> testOnly org.apache.celeborn.service.deploy.master.slotsalloc.*
```
If you’d like to run just a single Java test in the `SlotsAllocatorSuiteJ`, e.g., a test that with the name `testAllocateSlotsForSinglePartitionId`, you run the following command in the sbt console:
diff --git a/docs/developers/slotsallocation.md b/docs/developers/slotsallocation.md
index f1dcb040e..6057d63f2 100644
--- a/docs/developers/slotsallocation.md
+++ b/docs/developers/slotsallocation.md
@@ -76,6 +76,41 @@ Roundrobin slots allocation will distribute all slots into all registered worker
all workers as an array and place 1 slots in a worker until all slots are allocated.
If a worker has multiple disks, the chosen disk index is `(monotone increasing disk index +1) % disk count`.
+## Custom strategies
+
+Custom slot assignment strategies are loaded through Java's service provider interface (SPI).
+Implement `SlotsAssignStrategy` for the allocation algorithm and `SlotsAssignStrategyProvider`
+to give it a configuration name and create it from `CelebornConf`. The provider must have a public
+no-argument constructor and must be listed in
+`META-INF/services/org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAssignStrategyProvider`
+inside a JAR on the master class path. Set `celeborn.master.slot.assign.policy` to the provider name;
+provider names are matched case-insensitively.
+
+This SPI exposes Celeborn Master's slot-allocation model and is therefore version-coupled to the
+Master rather than being a cross-version plugin API. Compile the provider against the same
+`celeborn-master_<scala.binary.version>` artifact and Celeborn version used by the target cluster,
+and rebuild the provider when either version changes.
+
+Custom providers read arbitrary namespaced keys directly from the `CelebornConf` supplied to
+`create`. For example, a provider can read
+`celeborn.master.slot.assign.acme.maxSlotsPerDisk` with `conf.getLong(key, defaultValue)`. The same
+key can be set either in Celeborn's static configuration or at `SYSTEM` level in the dynamic
+configuration service. Dynamic values take precedence over static values.
+
+After a system-level dynamic configuration update, the Master creates a new `CelebornConf`
+snapshot and calls the selected provider's `create` method again. Once creation succeeds, the new
+strategy atomically replaces the old one. If configuration validation fails, the Master logs the
+failure and keeps using the last valid strategy. The provider itself can also be changed at runtime
+by updating `celeborn.master.slot.assign.policy` at `SYSTEM` level.
+
+The manager does not retry an unchanged system-level configuration snapshot. If provider creation
+fails, subsequent refreshes with the same snapshot are ignored while the last valid strategy
+remains active. Changing the dynamic configuration triggers another creation attempt.
+
+Providers should validate and capture their configuration in `create`, then return a strategy with
+fixed configuration. `computeSlotBudgets` should not read configuration or perform configuration
+I/O.
+
## Celeborn Worker's Behavior
1. When reserve slots Celeborn worker will decide a slot be placed in local disks or HDFS when reserve slots.
2. If a partition is evicted from memory, the partition might be placed in HDFS.
diff --git a/docs/migration.md b/docs/migration.md
index 75d693193..b015ef3d8 100644
--- a/docs/migration.md
+++ b/docs/migration.md
@@ -21,6 +21,13 @@ license: |
# Migration Guide
+# Upgrading from 0.7 to 1.0
+
+- Since 1.0.0, `celeborn.master.slot.assign.policy` supports custom strategy names registered through
+ the `SlotsAssignStrategyProvider` SPI and can be updated through system-level dynamic
+ configuration. Existing `ROUNDROBIN` and `LOADAWARE` values remain supported and require no
+ migration. See [Slots allocation](./developers/slotsallocation.md#custom-strategies) for details.
+
# Upgrading from 0.6 to 0.7
- Since 0.7.0, Celeborn removed `ReleaseSlots`.
diff --git a/master/src/main/java/org/apache/celeborn/service/deploy/master/SlotsAllocator.java b/master/src/main/java/org/apache/celeborn/service/deploy/master/SlotsAllocator.java
deleted file mode 100644
index 481b0aa10..000000000
--- a/master/src/main/java/org/apache/celeborn/service/deploy/master/SlotsAllocator.java
+++ /dev/null
@@ -1,854 +0,0 @@
-/*
- * 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.service.deploy.master;
-
-import static org.apache.celeborn.common.protocol.StorageInfo.Type.*;
-
-import java.util.*;
-import java.util.function.IntUnaryOperator;
-import java.util.stream.Collectors;
-
-import scala.Tuple2;
-import scala.Tuple3;
-
-import org.apache.commons.lang3.StringUtils;
-import org.roaringbitmap.RoaringBitmap;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import org.apache.celeborn.common.meta.DiskInfo;
-import org.apache.celeborn.common.meta.DiskStatus;
-import org.apache.celeborn.common.meta.WorkerInfo;
-import org.apache.celeborn.common.protocol.PartitionLocation;
-import org.apache.celeborn.common.protocol.StorageInfo;
-
-public class SlotsAllocator {
- static class UsableDiskInfo {
- DiskInfo diskInfo;
- long usableSlots;
-
- /** @param diskInfo will be used as source for usableSlots. */
- UsableDiskInfo(DiskInfo diskInfo) {
- this.diskInfo = diskInfo;
- this.usableSlots = diskInfo.getAvailableSlots();
- }
-
- UsableDiskInfo(DiskInfo diskInfo, long usableSlots) {
- this.diskInfo = diskInfo;
- this.usableSlots = usableSlots;
- }
- }
-
- private static final Logger logger = LoggerFactory.getLogger(SlotsAllocator.class);
- private static final Random rand = new Random();
- private static boolean initialized = false;
- private static double[] taskAllocationRatio = null;
-
- public static Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>>
- offerSlotsRoundRobin(
- List<WorkerInfo> workers,
- List<Integer> partitionIds,
- boolean shouldReplicate,
- boolean shouldRackAware,
- int availableStorageTypes,
- boolean interruptionAware,
- int interruptionAwareThreshold) {
- if (partitionIds.isEmpty()) {
- return new HashMap<>();
- }
- if (workers.size() < 2 && shouldReplicate) {
- return new HashMap<>();
- }
- Map<WorkerInfo, List<UsableDiskInfo>> slotsRestrictions = new HashMap<>();
- for (WorkerInfo worker : workers) {
- List<UsableDiskInfo> usableDisks =
- slotsRestrictions.computeIfAbsent(worker, v -> new ArrayList<>());
- for (DiskInfo diskInfo : worker.diskInfos().values()) {
- if (DiskStatus.HEALTHY.equals(diskInfo.status())
- && StorageInfo.isAvailable(diskInfo.storageType(), availableStorageTypes)) {
- usableDisks.add(new UsableDiskInfo(diskInfo));
- }
- }
- }
- return locateSlots(
- partitionIds,
- workers,
- slotsRestrictions,
- shouldReplicate,
- shouldRackAware,
- availableStorageTypes,
- interruptionAware,
- interruptionAwareThreshold);
- }
-
- /**
- * It assumes that all disks whose available space is greater than the minimum space are divided
- * into multiple groups. A faster group will allocate more allocations than a slower group by
- * diskGroupGradient.
- */
- public static Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>>
- offerSlotsLoadAware(
- List<WorkerInfo> workers,
- List<Integer> partitionIds,
- boolean shouldReplicate,
- boolean shouldRackAware,
- int diskGroupCount,
- double diskGroupGradient,
- double flushTimeWeight,
- double fetchTimeWeight,
- double activeSlotsWeight,
- int availableStorageTypes,
- boolean interruptionAware,
- int interruptionAwareThreshold) {
- if (partitionIds.isEmpty()) {
- return new HashMap<>();
- }
- if (workers.size() < 2 && shouldReplicate) {
- return new HashMap<>();
- }
- if (StorageInfo.HDFSOnly(availableStorageTypes)
- || StorageInfo.S3Only(availableStorageTypes)
- || StorageInfo.OSSOnly(availableStorageTypes)) {
- return offerSlotsRoundRobin(
- workers,
- partitionIds,
- shouldReplicate,
- shouldRackAware,
- availableStorageTypes,
- interruptionAware,
- interruptionAwareThreshold);
- }
-
- List<DiskInfo> usableDisks = new ArrayList<>();
- Map<DiskInfo, WorkerInfo> diskToWorkerMap = new HashMap<>();
-
- workers.forEach(
- i ->
- i.diskInfos()
- .forEach(
- (key, diskInfo) -> {
- diskToWorkerMap.put(diskInfo, i);
- if (diskInfo.actualUsableSpace() > 0
- && diskInfo.status().equals(DiskStatus.HEALTHY)
- && !diskInfo.storageType().isDFS()) {
- usableDisks.add(diskInfo);
- }
- }));
-
- boolean noUsableDisks =
- usableDisks.isEmpty()
- || (shouldReplicate
- && (usableDisks.size() == 1
- || usableDisks.stream().map(diskToWorkerMap::get).distinct().count() <= 1));
- boolean noAvailableSlots =
- usableDisks.stream().mapToLong(DiskInfo::getAvailableSlots).sum() <= 0;
-
- if (noUsableDisks || noAvailableSlots) {
- logger.warn(
- "offer slots for {} fallback to roundrobin because there is no {}",
- StringUtils.join(partitionIds, ','),
- noUsableDisks ? "usable disks" : "available slots");
- return offerSlotsRoundRobin(
- workers,
- partitionIds,
- shouldReplicate,
- shouldRackAware,
- availableStorageTypes,
- interruptionAware,
- interruptionAwareThreshold);
- }
-
- if (!initialized) {
- initLoadAwareAlgorithm(diskGroupCount, diskGroupGradient);
- }
-
- Map<WorkerInfo, List<UsableDiskInfo>> slotsRestrictions =
- getSlotsRestrictionsByLoadAwareAlgorithm(
- placeDisksToGroups(
- usableDisks, diskGroupCount, flushTimeWeight, fetchTimeWeight, activeSlotsWeight),
- diskToWorkerMap,
- shouldReplicate ? partitionIds.size() * 2 : partitionIds.size());
- return locateSlots(
- partitionIds,
- workers,
- slotsRestrictions,
- shouldReplicate,
- shouldRackAware,
- availableStorageTypes,
- interruptionAware,
- interruptionAwareThreshold);
- }
-
- static StorageInfo buildStorageInfo(
- List<WorkerInfo> workers,
- int workerIndex,
- Map<WorkerInfo, List<UsableDiskInfo>> restrictions,
- Map<WorkerInfo, Integer> workerDiskIndex,
- int availableStorageTypes) {
- WorkerInfo selectedWorker = workers.get(workerIndex);
- StorageInfo storageInfo;
- if (restrictions != null) {
- List<UsableDiskInfo> usableDiskInfos = restrictions.get(selectedWorker);
- int diskIndex =
- workerDiskIndex.computeIfAbsent(
- selectedWorker, v -> rand.nextInt(usableDiskInfos.size()));
- while (usableDiskInfos.get(diskIndex).usableSlots <= 0) {
- diskIndex = (diskIndex + 1) % usableDiskInfos.size();
- }
- usableDiskInfos.get(diskIndex).usableSlots--;
- DiskInfo selectedDiskInfo = usableDiskInfos.get(diskIndex).diskInfo;
- if (selectedDiskInfo.storageType().isDFS()) {
- storageInfo = new StorageInfo("", selectedDiskInfo.storageType(), availableStorageTypes);
- } else {
- storageInfo =
- new StorageInfo(
- selectedDiskInfo.mountPoint(),
- selectedDiskInfo.storageType(),
- availableStorageTypes);
- workerDiskIndex.put(selectedWorker, (diskIndex + 1) % usableDiskInfos.size());
- }
- } else {
- if (StorageInfo.localDiskAvailable(availableStorageTypes)) {
- DiskInfo[] diskInfos =
- selectedWorker.diskInfos().values().stream()
- .filter(p -> !p.storageType().isDFS())
- .collect(Collectors.toList())
- .toArray(new DiskInfo[0]);
- int diskIndex =
- workerDiskIndex.computeIfAbsent(selectedWorker, v -> rand.nextInt(diskInfos.length));
- storageInfo =
- new StorageInfo(
- diskInfos[diskIndex].mountPoint(),
- diskInfos[diskIndex].storageType(),
- availableStorageTypes);
- workerDiskIndex.put(selectedWorker, (diskIndex + 1) % diskInfos.length);
- } else if (StorageInfo.S3Available(availableStorageTypes)) {
- storageInfo = new StorageInfo("", S3, availableStorageTypes);
- } else if (StorageInfo.OSSAvailable(availableStorageTypes)) {
- storageInfo = new StorageInfo("", OSS, availableStorageTypes);
- } else if (StorageInfo.HDFSAvailable(availableStorageTypes)) {
- storageInfo = new StorageInfo("", HDFS, availableStorageTypes);
- } else if (StorageInfo.memoryAvailable(availableStorageTypes)) {
- storageInfo = new StorageInfo("", StorageInfo.Type.MEMORY, availableStorageTypes);
- } else {
- throw new IllegalStateException("no storage type available");
- }
- }
- return storageInfo;
- }
-
- /**
- * If interruptionAware = true, select workers based on 2 main criteria: <br>
- * 1. Workers that have no nextInterruptionNotice are the first priority and are included in the
- * 1st pass for slot selection. <br>
- * 2. Workers that have a later interruption notice are a little less deprioritized, and are
- * included in the 2nd pass for slot selection. This is determined by nextInterruptionNotice above
- * a certain percentage threshold.<br>
- * All other workers are considered least priority, and are only included for slot selection in
- * the worst case. <br>
- */
- static Tuple3<List<WorkerInfo>, List<WorkerInfo>, List<WorkerInfo>>
- prioritizeWorkersBasedOnInterruptionNotice(
- List<WorkerInfo> workers,
- boolean shouldReplicate,
- boolean shouldRackAware,
- double percentileThreshold) {
- Map<Boolean, List<WorkerInfo>> partitioned =
- workers.stream().collect(Collectors.partitioningBy(WorkerInfo::hasInterruptionNotice));
- List<WorkerInfo> workersWithInterruptions = partitioned.get(true);
- List<WorkerInfo> workersWithoutInterruptions = partitioned.get(false);
- // Timestamps towards the boundary of `percentileThreshold` might be the same. Given this
- // is a stable sort, it makes sense to randomize these hosts so that the same hosts are not
- // consistently selected.
- Collections.shuffle(workersWithInterruptions);
- workersWithInterruptions.sort(
- Comparator.comparingLong(WorkerInfo::nextInterruptionNotice).reversed());
- int requiredNodes =
- (int) Math.floor((percentileThreshold * workersWithInterruptions.size()) / 100.0);
-
- List<WorkerInfo> workersWithLateInterruptions =
- new ArrayList<>(workersWithInterruptions.subList(0, requiredNodes));
- List<WorkerInfo> workersWithEarlyInterruptions =
- new ArrayList<>(
- workersWithInterruptions.subList(requiredNodes, workersWithInterruptions.size()));
- if (shouldReplicate && shouldRackAware) {
- return Tuple3.apply(
- generateRackAwareWorkers(workersWithoutInterruptions),
- Collections.unmodifiableList(workersWithLateInterruptions),
- generateRackAwareWorkers(workersWithEarlyInterruptions));
- }
- return Tuple3.apply(
- Collections.unmodifiableList(workersWithoutInterruptions),
- Collections.unmodifiableList(workersWithLateInterruptions),
- Collections.unmodifiableList(workersWithEarlyInterruptions));
- }
-
- /**
- * Progressive locate slots for all partitions <br>
- * 1. try to allocate for all partitions under restrictions, on workers with no interruption
- * notice if interruptionAware = true. <br>
- * 2. try to allocate for all partitions, and attempt the replica selection to be
- * interruptionAware if interruptionAware = true <br>
- * 3. allocate remain partitions to all workers <br>
- * 4. allocate remain partitions to all workers again without considering rack aware <br>
- */
- private static Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>>
- locateSlots(
- List<Integer> partitionIds,
- List<WorkerInfo> workersList,
- Map<WorkerInfo, List<UsableDiskInfo>> slotRestrictions,
- boolean shouldReplicate,
- boolean shouldRackAware,
- int availableStorageTypes,
- boolean interruptionAware,
- int interruptionAwareThreshold) {
-
- List<WorkerInfo> workersFromSlotRestrictions = new ArrayList<>(slotRestrictions.keySet());
- List<WorkerInfo> workers = workersList;
- if (shouldReplicate && shouldRackAware) {
- workersFromSlotRestrictions = generateRackAwareWorkers(workersFromSlotRestrictions);
- workers = generateRackAwareWorkers(workers);
- }
-
- Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots =
- new HashMap<>();
- List<WorkerInfo> workersWithoutInterruptions;
- List<WorkerInfo> workersWithLateInterruptions;
- List<WorkerInfo> workersWithEarlyInterruptions;
- if (interruptionAware) {
- Tuple3<List<WorkerInfo>, List<WorkerInfo>, List<WorkerInfo>>
- workersBasedOnInterruptionNotice =
- prioritizeWorkersBasedOnInterruptionNotice(
- workersFromSlotRestrictions,
- shouldReplicate,
- shouldRackAware,
- interruptionAwareThreshold);
- workersWithoutInterruptions = workersBasedOnInterruptionNotice._1();
- workersWithLateInterruptions = workersBasedOnInterruptionNotice._2();
- workersWithEarlyInterruptions = workersBasedOnInterruptionNotice._3();
- } else {
- workersWithoutInterruptions = workersFromSlotRestrictions;
- workersWithLateInterruptions = null;
- workersWithEarlyInterruptions = null;
- }
- // In the first pass, we try to place all partitions (primary and replica) from
- // `workersWithoutInterruptions`.
- List<Integer> remain =
- roundRobin(
- slots,
- partitionIds,
- workersWithoutInterruptions,
- workersWithoutInterruptions,
- slotRestrictions,
- shouldReplicate,
- shouldRackAware,
- availableStorageTypes,
- true);
- logger.debug(
- "Remaining number of partitionIds after 1st pass slot selection: {}", remain.size());
- // Do an extra pass for partition placement if interruptionAware = true, to see if we can
- // assign the remaining partitions with slot restriction still set in place. The goal during
- // this pass
- // is to see if we can place primary from `workersWithoutInterruptions +
- // workersWithLateInterruptions`, while replica can be in
- // `workersWithEarlyInterruptions`.
- // This is to avoid the degenerate case in which both primary and replica may end up in
- // `workersWithEarlyInterruptions`.
- if (interruptionAware && !remain.isEmpty()) {
- List<WorkerInfo> primaryWorkerCandidates = new ArrayList<>(workersWithoutInterruptions);
- primaryWorkerCandidates.addAll(workersWithLateInterruptions);
- if (shouldReplicate && shouldRackAware) {
- primaryWorkerCandidates = generateRackAwareWorkers(primaryWorkerCandidates);
- }
- remain =
- roundRobin(
- slots,
- remain,
- primaryWorkerCandidates,
- workersWithEarlyInterruptions,
- null,
- shouldReplicate,
- shouldRackAware,
- availableStorageTypes,
- false);
- logger.debug(
- "Remaining number of partitionIds after 2nd pass slot selection: {}", remain.size());
- }
- // If partitions are remaining from this point on, and interruptionAware = true, then
- // this becomes the degenerate case where both primary and replica are likely chosen on
- // workers with interruptions that are sooner.
- if (!remain.isEmpty()) {
- remain =
- roundRobin(
- slots,
- remain,
- workers,
- workers,
- null,
- shouldReplicate,
- shouldRackAware,
- availableStorageTypes,
- true);
- logger.debug(
- "Remaining number of partitionIds after 3rd pass slot selection: {}", remain.size());
- }
- if (!remain.isEmpty()) {
- roundRobin(
- slots,
- remain,
- workers,
- workers,
- null,
- shouldReplicate,
- false,
- availableStorageTypes,
- true);
- }
- return slots;
- }
-
- /**
- * The rack distribution of the input workers list is essentially random, and in degenerate cases
- * the rack aware slot selection ends up skipping sub list's of hosts (in same rack) - which
- * results in uneven distribution of replica selection. For example given worker list: [h1r1,
- * h2r1, h3r1, h4r2, h5r2, h6r2] if primary is h3r1 and replica index is pointing to h1r1, it will
- * skip both h2r1 and h3r1 in order to pick h4r2; and for the next slot, primary will be h4r2 and
- * will skip all the r2 hosts in order to pick h1r1. This ends up being suboptimal where some
- * hosts are picked a lot more than others (due to the worker and worker/rack distribution). In
- * order to mitigate this, we reorder the worker list by redistributing the workers based on rack
- * to increase the rack diversity between adjoining workers, so that we minimize skipping over
- * consecutive hosts.
- */
- static List<WorkerInfo> generateRackAwareWorkers(List<WorkerInfo> workers) {
-
- List<Map.Entry<String, LinkedList<WorkerInfo>>> sortedRackToHosts;
- {
- Map<String, LinkedList<WorkerInfo>> map = new HashMap<>();
- for (WorkerInfo worker : workers) {
- map.computeIfAbsent(worker.networkLocation(), key -> new LinkedList<>()).add(worker);
- }
- sortedRackToHosts = new ArrayList<>(map.entrySet());
- // reverse sort by number of hosts per rack
- sortedRackToHosts.sort(
- (o1, o2) -> Integer.compare(o2.getValue().size(), o1.getValue().size()));
- }
-
- ArrayList<WorkerInfo> result = new ArrayList<>(workers.size());
- int count = 0;
- final int numWorkers = workers.size();
- while (count < numWorkers) {
- Iterator<Map.Entry<String, LinkedList<WorkerInfo>>> iter = sortedRackToHosts.iterator();
- while (iter.hasNext()) {
- LinkedList<WorkerInfo> workerList = iter.next().getValue();
- result.add(workerList.removeFirst());
- count++;
- if (workerList.isEmpty()) {
- iter.remove();
- }
- }
- }
-
- return Collections.unmodifiableList(result);
- }
-
- /**
- * Assigns slots in a roundrobin fashion given lists of primary and replica worker candidates and
- * other restrictions.
- *
- * @param slots the slots that have been assigned for each partitionId
- * @param partitionIds the partitionIds that require slot selection still
- * @param primaryWorkers list of worker candidates that can be used for primary workers.
- * @param replicaWorkers list of worker candidates that can be used for replica workers.
- * @param slotsRestrictions restrictions for each available slot based on worker characteristics
- * @param shouldReplicate if replication is enabled within the cluster
- * @param shouldRackAware if rack-aware replication is enabled within the cluster.
- * @param availableStorageTypes available storage types coming from the offer slots request.
- * @param skipLocationsOnSameWorkerCheck if the worker candidates list for primaries and replicas
- * is the same. This is to prevent index mismatch while assigning slots across both lists.
- * @return the partitionIds that were not able to be assigned slots in this iteration with the
- * current primary and replica worker candidates and slot restrictions.
- */
- private static List<Integer> roundRobin(
- Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots,
- List<Integer> partitionIds,
- List<WorkerInfo> primaryWorkers,
- List<WorkerInfo> replicaWorkers,
- Map<WorkerInfo, List<UsableDiskInfo>> slotsRestrictions,
- boolean shouldReplicate,
- boolean shouldRackAware,
- int availableStorageTypes,
- boolean skipLocationsOnSameWorkerCheck) {
- if (primaryWorkers.isEmpty() || (shouldReplicate && replicaWorkers.isEmpty())) {
- return partitionIds;
- }
- // workerInfo -> (diskIndexForPrimaryAndReplica)
- Map<WorkerInfo, Integer> workerDiskIndex = new HashMap<>();
- List<Integer> partitionIdList = new LinkedList<>(partitionIds);
-
- final int primaryWorkersSize = primaryWorkers.size();
- final int replicaWorkersSize = replicaWorkers.size();
- final IntUnaryOperator primaryWorkersIncrementIndex = v -> (v + 1) % primaryWorkersSize;
- int primaryIndex = rand.nextInt(primaryWorkersSize);
- final IntUnaryOperator replicaWorkersIncrementIndex;
- int replicaIndex;
- if (shouldReplicate) {
- replicaWorkersIncrementIndex = v -> (v + 1) % replicaWorkersSize;
- replicaIndex = rand.nextInt(replicaWorkersSize);
- } else {
- replicaWorkersIncrementIndex = null;
- replicaIndex = -1;
- }
-
- ListIterator<Integer> iter = partitionIdList.listIterator(partitionIdList.size());
- // Iterate from the end to preserve O(1) removal of processed partitions.
- // This is important when we have a high number of concurrent apps that have a
- // high number of partitions.
- outer:
- while (iter.hasPrevious()) {
- int nextPrimaryInd = primaryIndex;
-
- int partitionId = iter.previous();
- StorageInfo storageInfo;
- if (slotsRestrictions != null && !slotsRestrictions.isEmpty()) {
- // this means that we'll select a mount point
- while (!haveUsableSlots(slotsRestrictions, primaryWorkers, nextPrimaryInd)) {
- nextPrimaryInd = primaryWorkersIncrementIndex.applyAsInt(nextPrimaryInd);
- if (nextPrimaryInd == primaryIndex) {
- break outer;
- }
- }
- storageInfo =
- buildStorageInfo(
- primaryWorkers,
- nextPrimaryInd,
- slotsRestrictions,
- workerDiskIndex,
- availableStorageTypes);
- } else {
- if (StorageInfo.localDiskAvailable(availableStorageTypes)) {
- while (!primaryWorkers.get(nextPrimaryInd).haveDisk()) {
- nextPrimaryInd = primaryWorkersIncrementIndex.applyAsInt(nextPrimaryInd);
- if (nextPrimaryInd == primaryIndex) {
- break outer;
- }
- }
- }
- storageInfo =
- buildStorageInfo(
- primaryWorkers, nextPrimaryInd, null, workerDiskIndex, availableStorageTypes);
- }
- PartitionLocation primaryPartition =
- createLocation(partitionId, primaryWorkers.get(nextPrimaryInd), null, storageInfo, true);
-
- if (shouldReplicate) {
- int nextReplicaInd = replicaIndex;
- if (slotsRestrictions != null) {
- while ((nextReplicaInd == nextPrimaryInd && skipLocationsOnSameWorkerCheck)
- || !haveUsableSlots(slotsRestrictions, replicaWorkers, nextReplicaInd)
- || !satisfyRackAware(
- shouldRackAware,
- primaryWorkers,
- nextPrimaryInd,
- replicaWorkers,
- nextReplicaInd)) {
- nextReplicaInd = replicaWorkersIncrementIndex.applyAsInt(nextReplicaInd);
- if (nextReplicaInd == replicaIndex) {
- break outer;
- }
- }
- storageInfo =
- buildStorageInfo(
- replicaWorkers,
- nextReplicaInd,
- slotsRestrictions,
- workerDiskIndex,
- availableStorageTypes);
- } else if (shouldRackAware) {
- while ((nextReplicaInd == nextPrimaryInd && skipLocationsOnSameWorkerCheck)
- || !satisfyRackAware(
- true, primaryWorkers, nextPrimaryInd, replicaWorkers, nextReplicaInd)) {
- nextReplicaInd = replicaWorkersIncrementIndex.applyAsInt(nextReplicaInd);
- if (nextReplicaInd == replicaIndex) {
- break outer;
- }
- }
- } else {
- if (StorageInfo.localDiskAvailable(availableStorageTypes)) {
- while ((nextReplicaInd == nextPrimaryInd && skipLocationsOnSameWorkerCheck)
- || !replicaWorkers.get(nextReplicaInd).haveDisk()) {
- nextReplicaInd = replicaWorkersIncrementIndex.applyAsInt(nextReplicaInd);
- if (nextReplicaInd == replicaIndex) {
- break outer;
- }
- }
- }
- storageInfo =
- buildStorageInfo(
- replicaWorkers, nextReplicaInd, null, workerDiskIndex, availableStorageTypes);
- }
- PartitionLocation replicaPartition =
- createLocation(
- partitionId,
- replicaWorkers.get(nextReplicaInd),
- primaryPartition,
- storageInfo,
- false);
- primaryPartition.setPeer(replicaPartition);
- Tuple2<List<PartitionLocation>, List<PartitionLocation>> locations =
- slots.computeIfAbsent(
- replicaWorkers.get(nextReplicaInd),
- v -> new Tuple2<>(new ArrayList<>(), new ArrayList<>()));
- locations._2.add(replicaPartition);
- replicaIndex = replicaWorkersIncrementIndex.applyAsInt(nextReplicaInd);
- }
-
- Tuple2<List<PartitionLocation>, List<PartitionLocation>> locations =
- slots.computeIfAbsent(
- primaryWorkers.get(nextPrimaryInd),
- v -> new Tuple2<>(new ArrayList<>(), new ArrayList<>()));
- locations._1.add(primaryPartition);
- primaryIndex = primaryWorkersIncrementIndex.applyAsInt(nextPrimaryInd);
- iter.remove();
- }
- return partitionIdList;
- }
-
- private static boolean haveUsableSlots(
- Map<WorkerInfo, List<UsableDiskInfo>> restrictions, List<WorkerInfo> workers, int index) {
- return restrictions.get(workers.get(index)).stream().mapToLong(i -> i.usableSlots).sum() > 0;
- }
-
- private static boolean satisfyRackAware(
- boolean shouldRackAware,
- List<WorkerInfo> primaryWorkers,
- int primaryIndex,
- List<WorkerInfo> replicaWorkers,
- int nextReplicaInd) {
- return !shouldRackAware
- || !Objects.equals(
- primaryWorkers.get(primaryIndex).networkLocation(),
- replicaWorkers.get(nextReplicaInd).networkLocation());
- }
-
- private static void initLoadAwareAlgorithm(int diskGroups, double diskGroupGradient) {
- taskAllocationRatio = new double[diskGroups];
- double totalAllocations = 0;
-
- for (int i = 0; i < diskGroups; i++) {
- totalAllocations += Math.pow(1 + diskGroupGradient, diskGroups - 1 - i);
- }
- for (int i = 0; i < diskGroups; i++) {
- taskAllocationRatio[i] =
- Math.pow(1 + diskGroupGradient, diskGroups - 1 - i) / totalAllocations;
- }
- logger.info(
- "load-aware offer slots algorithm init with taskAllocationRatio {}",
- StringUtils.join(taskAllocationRatio, ','));
- initialized = true;
- }
-
- private static List<List<DiskInfo>> placeDisksToGroups(
- List<DiskInfo> usableDisks,
- int diskGroupCount,
- double flushTimeWeight,
- double fetchTimeWeight,
- double activeSlotsWeight) {
- List<List<DiskInfo>> diskGroups = new ArrayList<>();
- usableDisks.sort(
- (o1, o2) -> {
- double delta =
- (o1.avgFlushTime() * flushTimeWeight + o1.avgFetchTime() * fetchTimeWeight)
- + o1.activeSlots() * activeSlotsWeight
- - (o2.avgFlushTime() * flushTimeWeight
- + o2.avgFetchTime() * fetchTimeWeight
- + o2.activeSlots() * activeSlotsWeight);
- return delta < 0 ? -1 : (delta > 0 ? 1 : 0);
- });
- int diskCount = usableDisks.size();
- int startIndex = 0;
- int groupSizeSize = (int) Math.ceil(usableDisks.size() / (double) diskGroupCount);
- for (int i = 0; i < diskGroupCount; i++) {
- List<DiskInfo> diskList = new ArrayList<>();
- if (startIndex >= usableDisks.size()) {
- continue;
- }
- if (startIndex + groupSizeSize <= diskCount) {
- diskList.addAll(usableDisks.subList(startIndex, startIndex + groupSizeSize));
- startIndex += groupSizeSize;
- } else {
- diskList.addAll(usableDisks.subList(startIndex, usableDisks.size()));
- startIndex = usableDisks.size();
- }
- diskGroups.add(diskList);
- }
- return diskGroups;
- }
-
- /**
- * This method implement the load aware slots allocation algorithm. See details at
- * /docs/developers/slotsallocation.md
- */
- private static Map<WorkerInfo, List<UsableDiskInfo>> getSlotsRestrictionsByLoadAwareAlgorithm(
- List<List<DiskInfo>> groups, Map<DiskInfo, WorkerInfo> diskWorkerMap, int partitionCnt) {
- int groupSize = groups.size();
- long[] groupAllocations = new long[groupSize];
- Map<WorkerInfo, List<UsableDiskInfo>> restrictions = new HashMap<>();
- long[] groupAvailableSlots = new long[groupSize];
- for (int i = 0; i < groupSize; i++) {
- for (DiskInfo disk : groups.get(i)) {
- groupAvailableSlots[i] += disk.getAvailableSlots();
- }
- }
- double[] currentAllocation = new double[groupSize];
- double currentAllocationSum = 0;
- for (int i = 0; i < groupSize; i++) {
- if (!groups.get(i).isEmpty()) {
- currentAllocationSum += taskAllocationRatio[i];
- }
- }
- for (int i = 0; i < groupSize; i++) {
- if (!groups.get(i).isEmpty()) {
- currentAllocation[i] = taskAllocationRatio[i] / currentAllocationSum;
- }
- }
- long toNextGroup = 0;
- long left = partitionCnt;
- for (int i = 0; i < groupSize; i++) {
- if (left <= 0) {
- break;
- }
- long estimateAllocation = (int) Math.ceil(partitionCnt * currentAllocation[i]);
- if (estimateAllocation > left) {
- estimateAllocation = left;
- }
- if (estimateAllocation + toNextGroup > groupAvailableSlots[i]) {
- groupAllocations[i] = groupAvailableSlots[i];
- toNextGroup = estimateAllocation - groupAvailableSlots[i] + toNextGroup;
- } else {
- groupAllocations[i] = estimateAllocation + toNextGroup;
- }
- left -= groupAllocations[i];
- }
-
- long groupLeft = 0;
- for (int i = 0; i < groups.size(); i++) {
- int disksInsideGroup = groups.get(i).size();
- long groupRequired = groupAllocations[i] + groupLeft;
- for (DiskInfo disk : groups.get(i)) {
- if (groupRequired <= 0) {
- break;
- }
- List<UsableDiskInfo> diskAllocation =
- restrictions.computeIfAbsent(diskWorkerMap.get(disk), v -> new ArrayList<>());
- long allocated =
- (int) Math.ceil((groupAllocations[i] + groupLeft) / (double) disksInsideGroup);
- if (allocated > disk.getAvailableSlots()) {
- allocated = disk.getAvailableSlots();
- }
- if (allocated > groupRequired) {
- allocated = groupRequired;
- }
- diskAllocation.add(new UsableDiskInfo(disk, Math.toIntExact(allocated)));
- groupRequired -= allocated;
- }
- groupLeft = groupRequired;
- }
-
- if (logger.isDebugEnabled()) {
- StringBuilder sb = new StringBuilder();
- for (int i = 0; i < groups.size(); i++) {
- sb.append("| group ").append(i).append(" ");
- for (DiskInfo diskInfo : groups.get(i)) {
- WorkerInfo workerInfo = diskWorkerMap.get(diskInfo);
- String workerHost = workerInfo.host();
- long allocation = 0;
- if (restrictions.get(workerInfo) != null) {
- for (UsableDiskInfo usableInfo : restrictions.get(workerInfo)) {
- if (usableInfo.diskInfo.equals(diskInfo)) {
- allocation = usableInfo.usableSlots;
- }
- }
- }
- sb.append(workerHost)
- .append("-")
- .append(diskInfo.mountPoint())
- .append(" flushtime:")
- .append(diskInfo.avgFlushTime())
- .append(" fetchtime:")
- .append(diskInfo.avgFetchTime())
- .append(" allocation: ")
- .append(allocation)
- .append(" ");
- }
- sb.append(" | ");
- }
- logger.debug(
- "total {} allocate with group {} with allocations {}",
- partitionCnt,
- StringUtils.join(groupAllocations, ','),
- sb);
- }
- return restrictions;
- }
-
- private static PartitionLocation createLocation(
- int partitionIndex,
- WorkerInfo workerInfo,
- PartitionLocation peer,
- StorageInfo storageInfo,
- boolean isPrimary) {
- return new PartitionLocation(
- partitionIndex,
- 0,
- workerInfo.host(),
- workerInfo.rpcPort(),
- workerInfo.pushPort(),
- workerInfo.fetchPort(),
- workerInfo.replicatePort(),
- isPrimary ? PartitionLocation.Mode.PRIMARY : PartitionLocation.Mode.REPLICA,
- peer,
- storageInfo,
- new RoaringBitmap());
- }
-
- public static Map<WorkerInfo, Map<String, Integer>> slotsToDiskAllocations(
- Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots) {
- Iterator<WorkerInfo> workers = slots.keySet().iterator();
- Map<WorkerInfo, Map<String, Integer>> workerToSlots = new HashMap<>();
- while (workers.hasNext()) {
- WorkerInfo worker = workers.next();
- Map<String, Integer> slotsPerDisk =
- workerToSlots.computeIfAbsent(worker, v -> new HashMap<>());
- List<PartitionLocation> jointLocations = new ArrayList<>();
- jointLocations.addAll(slots.get(worker)._1);
- jointLocations.addAll(slots.get(worker)._2);
- for (PartitionLocation location : jointLocations) {
- String mountPoint = location.getStorageInfo().getMountPoint();
- // skip non local disks slots
- if (!mountPoint.isEmpty()) {
- if (slotsPerDisk.containsKey(mountPoint)) {
- slotsPerDisk.put(mountPoint, slotsPerDisk.get(mountPoint) + 1);
- } else {
- slotsPerDisk.put(mountPoint, 1);
- }
- }
- }
- }
- return workerToSlots;
- }
-}
diff --git a/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/LoadAwareSlotsAssignStrategy.java b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/LoadAwareSlotsAssignStrategy.java
new file mode 100644
index 000000000..617aa820f
--- /dev/null
+++ b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/LoadAwareSlotsAssignStrategy.java
@@ -0,0 +1,262 @@
+/*
+ * 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.service.deploy.master.slotsalloc;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.commons.lang3.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.meta.DiskInfo;
+import org.apache.celeborn.common.meta.DiskStatus;
+import org.apache.celeborn.common.meta.WorkerInfo;
+import org.apache.celeborn.common.protocol.StorageInfo;
+
+public class LoadAwareSlotsAssignStrategy implements SlotsAssignStrategy {
+
+ private static final Logger logger = LoggerFactory.getLogger(LoadAwareSlotsAssignStrategy.class);
+
+ private final int diskGroupCount;
+ private final double flushTimeWeight;
+ private final double fetchTimeWeight;
+ private final double activeSlotsWeight;
+ private final double[] taskAllocationRatio;
+ private final RoundRobinSlotsAssignStrategy fallback = new RoundRobinSlotsAssignStrategy();
+
+ public LoadAwareSlotsAssignStrategy(
+ int diskGroupCount,
+ double diskGroupGradient,
+ double flushTimeWeight,
+ double fetchTimeWeight,
+ double activeSlotsWeight) {
+ this.diskGroupCount = diskGroupCount;
+ this.flushTimeWeight = flushTimeWeight;
+ this.fetchTimeWeight = fetchTimeWeight;
+ this.activeSlotsWeight = activeSlotsWeight;
+ this.taskAllocationRatio = computeTaskAllocationRatio(diskGroupCount, diskGroupGradient);
+ }
+
+ @Override
+ public Map<WorkerInfo, List<UsableDiskInfo>> computeSlotBudgets(
+ List<WorkerInfo> workers,
+ List<Integer> partitionIds,
+ boolean shouldReplicate,
+ int availableStorageTypes) {
+ if (StorageInfo.HDFSOnly(availableStorageTypes)
+ || StorageInfo.S3Only(availableStorageTypes)
+ || StorageInfo.OSSOnly(availableStorageTypes)) {
+ return fallback.computeSlotBudgets(
+ workers, partitionIds, shouldReplicate, availableStorageTypes);
+ }
+
+ List<DiskInfo> usableDisks = new ArrayList<>();
+ Map<DiskInfo, WorkerInfo> diskToWorkerMap = new HashMap<>();
+
+ workers.forEach(
+ i ->
+ i.diskInfos()
+ .forEach(
+ (key, diskInfo) -> {
+ diskToWorkerMap.put(diskInfo, i);
+ if (diskInfo.actualUsableSpace() > 0
+ && diskInfo.status().equals(DiskStatus.HEALTHY)
+ && !diskInfo.storageType().isDFS()) {
+ usableDisks.add(diskInfo);
+ }
+ }));
+
+ boolean noUsableDisks =
+ usableDisks.isEmpty()
+ || (shouldReplicate
+ && (usableDisks.size() == 1
+ || usableDisks.stream().map(diskToWorkerMap::get).distinct().count() <= 1));
+ boolean noAvailableSlots =
+ usableDisks.stream().mapToLong(DiskInfo::getAvailableSlots).sum() <= 0;
+
+ if (noUsableDisks || noAvailableSlots) {
+ logger.warn(
+ "offer slots for {} fallback to roundrobin because there is no {}",
+ StringUtils.join(partitionIds, ','),
+ noUsableDisks ? "usable disks" : "available slots");
+ return fallback.computeSlotBudgets(
+ workers, partitionIds, shouldReplicate, availableStorageTypes);
+ }
+
+ List<List<DiskInfo>> groups =
+ placeDisksToGroups(
+ usableDisks, diskGroupCount, flushTimeWeight, fetchTimeWeight, activeSlotsWeight);
+ int partitionCnt = shouldReplicate ? partitionIds.size() * 2 : partitionIds.size();
+ int groupSize = groups.size();
+ long[] groupAllocations = new long[groupSize];
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets = new HashMap<>();
+ long[] groupAvailableSlots = new long[groupSize];
+ for (int i = 0; i < groupSize; i++) {
+ for (DiskInfo disk : groups.get(i)) {
+ groupAvailableSlots[i] += disk.getAvailableSlots();
+ }
+ }
+ double[] currentAllocation = new double[groupSize];
+ double currentAllocationSum = 0;
+ for (int i = 0; i < groupSize; i++) {
+ if (!groups.get(i).isEmpty()) {
+ currentAllocationSum += taskAllocationRatio[i];
+ }
+ }
+ for (int i = 0; i < groupSize; i++) {
+ if (!groups.get(i).isEmpty()) {
+ currentAllocation[i] = taskAllocationRatio[i] / currentAllocationSum;
+ }
+ }
+ long toNextGroup = 0;
+ long left = partitionCnt;
+ for (int i = 0; i < groupSize; i++) {
+ if (left <= 0) {
+ break;
+ }
+ long estimateAllocation = (int) Math.ceil(partitionCnt * currentAllocation[i]);
+ if (estimateAllocation > left) {
+ estimateAllocation = left;
+ }
+ if (estimateAllocation + toNextGroup > groupAvailableSlots[i]) {
+ groupAllocations[i] = groupAvailableSlots[i];
+ toNextGroup = estimateAllocation - groupAvailableSlots[i] + toNextGroup;
+ } else {
+ groupAllocations[i] = estimateAllocation + toNextGroup;
+ }
+ left -= groupAllocations[i];
+ }
+
+ long groupLeft = 0;
+ for (int i = 0; i < groups.size(); i++) {
+ int disksInsideGroup = groups.get(i).size();
+ long groupRequired = groupAllocations[i] + groupLeft;
+ for (DiskInfo disk : groups.get(i)) {
+ if (groupRequired <= 0) {
+ break;
+ }
+ List<UsableDiskInfo> diskAllocation =
+ budgets.computeIfAbsent(diskToWorkerMap.get(disk), v -> new ArrayList<>());
+ long allocated =
+ (int) Math.ceil((groupAllocations[i] + groupLeft) / (double) disksInsideGroup);
+ if (allocated > disk.getAvailableSlots()) {
+ allocated = disk.getAvailableSlots();
+ }
+ if (allocated > groupRequired) {
+ allocated = groupRequired;
+ }
+ diskAllocation.add(new UsableDiskInfo(disk, Math.toIntExact(allocated)));
+ groupRequired -= allocated;
+ }
+ groupLeft = groupRequired;
+ }
+
+ if (logger.isDebugEnabled()) {
+ StringBuilder sb = new StringBuilder();
+ for (int i = 0; i < groups.size(); i++) {
+ sb.append("| group ").append(i).append(" ");
+ for (DiskInfo diskInfo : groups.get(i)) {
+ WorkerInfo workerInfo = diskToWorkerMap.get(diskInfo);
+ String workerHost = workerInfo.host();
+ long allocation = 0;
+ if (budgets.get(workerInfo) != null) {
+ for (UsableDiskInfo usableInfo : budgets.get(workerInfo)) {
+ if (usableInfo.diskInfo.equals(diskInfo)) {
+ allocation = usableInfo.usableSlots;
+ }
+ }
+ }
+ sb.append(workerHost)
+ .append("-")
+ .append(diskInfo.mountPoint())
+ .append(" flushtime:")
+ .append(diskInfo.avgFlushTime())
+ .append(" fetchtime:")
+ .append(diskInfo.avgFetchTime())
+ .append(" allocation: ")
+ .append(allocation)
+ .append(" ");
+ }
+ sb.append(" | ");
+ }
+ logger.debug(
+ "total {} allocate with group {} with allocations {}",
+ partitionCnt,
+ StringUtils.join(groupAllocations, ','),
+ sb);
+ }
+ return budgets;
+ }
+
+ private static double[] computeTaskAllocationRatio(int diskGroups, double diskGroupGradient) {
+ double[] taskAllocationRatio = new double[diskGroups];
+ double totalAllocations = 0;
+
+ for (int i = 0; i < diskGroups; i++) {
+ totalAllocations += Math.pow(1 + diskGroupGradient, diskGroups - 1 - i);
+ }
+ for (int i = 0; i < diskGroups; i++) {
+ taskAllocationRatio[i] =
+ Math.pow(1 + diskGroupGradient, diskGroups - 1 - i) / totalAllocations;
+ }
+ logger.info(
+ "load-aware offer slots algorithm init with taskAllocationRatio {}",
+ StringUtils.join(taskAllocationRatio, ','));
+ return taskAllocationRatio;
+ }
+
+ private static List<List<DiskInfo>> placeDisksToGroups(
+ List<DiskInfo> usableDisks,
+ int diskGroupCount,
+ double flushTimeWeight,
+ double fetchTimeWeight,
+ double activeSlotsWeight) {
+ List<List<DiskInfo>> diskGroups = new ArrayList<>();
+ usableDisks.sort(
+ (o1, o2) -> {
+ double delta =
+ (o1.avgFlushTime() * flushTimeWeight + o1.avgFetchTime() * fetchTimeWeight)
+ + o1.activeSlots() * activeSlotsWeight
+ - (o2.avgFlushTime() * flushTimeWeight
+ + o2.avgFetchTime() * fetchTimeWeight
+ + o2.activeSlots() * activeSlotsWeight);
+ return delta < 0 ? -1 : (delta > 0 ? 1 : 0);
+ });
+ int diskCount = usableDisks.size();
+ int startIndex = 0;
+ int groupSizeSize = (int) Math.ceil(usableDisks.size() / (double) diskGroupCount);
+ for (int i = 0; i < diskGroupCount; i++) {
+ List<DiskInfo> diskList = new ArrayList<>();
+ if (startIndex >= usableDisks.size()) {
+ continue;
+ }
+ if (startIndex + groupSizeSize <= diskCount) {
+ diskList.addAll(usableDisks.subList(startIndex, startIndex + groupSizeSize));
+ startIndex += groupSizeSize;
+ } else {
+ diskList.addAll(usableDisks.subList(startIndex, usableDisks.size()));
+ startIndex = usableDisks.size();
+ }
+ diskGroups.add(diskList);
+ }
+ return diskGroups;
+ }
+}
diff --git a/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/LoadAwareSlotsAssignStrategyProvider.java b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/LoadAwareSlotsAssignStrategyProvider.java
new file mode 100644
index 000000000..d83dfed2d
--- /dev/null
+++ b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/LoadAwareSlotsAssignStrategyProvider.java
@@ -0,0 +1,39 @@
+/*
+ * 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.service.deploy.master.slotsalloc;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.protocol.SlotsAssignPolicy;
+
+public final class LoadAwareSlotsAssignStrategyProvider implements SlotsAssignStrategyProvider {
+
+ @Override
+ public String getName() {
+ return SlotsAssignPolicy.LOADAWARE.name();
+ }
+
+ @Override
+ public SlotsAssignStrategy create(CelebornConf conf) {
+ return new LoadAwareSlotsAssignStrategy(
+ conf.masterSlotAssignLoadAwareDiskGroupNum(),
+ conf.masterSlotAssignLoadAwareDiskGroupGradient(),
+ conf.masterSlotAssignLoadAwareFlushTimeWeight(),
+ conf.masterSlotAssignLoadAwareFetchTimeWeight(),
+ conf.masterSlotAssignLoadAwareActiveSlotsWeight());
+ }
+}
diff --git a/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/RoundRobinSlotsAssignStrategy.java b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/RoundRobinSlotsAssignStrategy.java
new file mode 100644
index 000000000..c048dd3bb
--- /dev/null
+++ b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/RoundRobinSlotsAssignStrategy.java
@@ -0,0 +1,51 @@
+/*
+ * 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.service.deploy.master.slotsalloc;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.celeborn.common.meta.DiskInfo;
+import org.apache.celeborn.common.meta.DiskStatus;
+import org.apache.celeborn.common.meta.WorkerInfo;
+import org.apache.celeborn.common.protocol.StorageInfo;
+
+public class RoundRobinSlotsAssignStrategy implements SlotsAssignStrategy {
+
+ @Override
+ public Map<WorkerInfo, List<UsableDiskInfo>> computeSlotBudgets(
+ List<WorkerInfo> workers,
+ List<Integer> partitionIds,
+ boolean shouldReplicate,
+ int availableStorageTypes) {
+ Map<WorkerInfo, List<UsableDiskInfo>> slotBudgets = new HashMap<>();
+ for (WorkerInfo worker : workers) {
+ List<UsableDiskInfo> usableDisks =
+ slotBudgets.computeIfAbsent(worker, v -> new ArrayList<>());
+ for (DiskInfo diskInfo : worker.diskInfos().values()) {
+ if (DiskStatus.HEALTHY.equals(diskInfo.status())
+ && StorageInfo.isAvailable(diskInfo.storageType(), availableStorageTypes)) {
+ usableDisks.add(new UsableDiskInfo(diskInfo));
+ }
+ }
+ }
+ return slotBudgets;
+ }
+}
diff --git a/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/RoundRobinSlotsAssignStrategyProvider.java b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/RoundRobinSlotsAssignStrategyProvider.java
new file mode 100644
index 000000000..4d2c76f94
--- /dev/null
+++ b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/RoundRobinSlotsAssignStrategyProvider.java
@@ -0,0 +1,34 @@
+/*
+ * 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.service.deploy.master.slotsalloc;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.protocol.SlotsAssignPolicy;
+
+public final class RoundRobinSlotsAssignStrategyProvider implements SlotsAssignStrategyProvider {
+
+ @Override
+ public String getName() {
+ return SlotsAssignPolicy.ROUNDROBIN.name();
+ }
+
+ @Override
+ public SlotsAssignStrategy create(CelebornConf conf) {
+ return new RoundRobinSlotsAssignStrategy();
+ }
+}
diff --git a/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocator.java b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocator.java
new file mode 100644
index 000000000..4de3ed5c2
--- /dev/null
+++ b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocator.java
@@ -0,0 +1,595 @@
+/*
+ * 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.service.deploy.master.slotsalloc;
+
+import static org.apache.celeborn.common.protocol.StorageInfo.Type.*;
+
+import java.util.*;
+import java.util.function.IntPredicate;
+import java.util.stream.Collectors;
+
+import scala.Tuple2;
+import scala.Tuple3;
+
+import org.roaringbitmap.RoaringBitmap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.meta.DiskInfo;
+import org.apache.celeborn.common.meta.WorkerInfo;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.protocol.StorageInfo;
+
+public class SlotsAllocator {
+ private static final Logger logger = LoggerFactory.getLogger(SlotsAllocator.class);
+ private static final Random rand = new Random();
+
+ public static Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>>
+ offerSlots(
+ List<WorkerInfo> workers,
+ List<Integer> partitionIds,
+ boolean shouldReplicate,
+ boolean shouldRackAware,
+ int availableStorageTypes,
+ boolean interruptionAware,
+ int interruptionAwareThreshold,
+ SlotsAssignStrategy strategy) {
+ if (partitionIds.isEmpty()) {
+ return new HashMap<>();
+ }
+ if (workers.size() < 2 && shouldReplicate) {
+ return new HashMap<>();
+ }
+ Map<WorkerInfo, List<UsableDiskInfo>> slotBudgets =
+ strategy.computeSlotBudgets(workers, partitionIds, shouldReplicate, availableStorageTypes);
+ return locateSlots(
+ partitionIds,
+ workers,
+ slotBudgets,
+ shouldReplicate,
+ shouldRackAware,
+ availableStorageTypes,
+ interruptionAware,
+ interruptionAwareThreshold);
+ }
+
+ static StorageInfo buildStorageInfo(
+ WorkerInfo selectedWorker,
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets,
+ Map<WorkerInfo, Integer> workerDiskIndex,
+ int availableStorageTypes) {
+ StorageInfo storageInfo;
+ if (budgets != null) {
+ List<UsableDiskInfo> usableDiskInfos = budgets.get(selectedWorker);
+ int diskIndex =
+ workerDiskIndex.computeIfAbsent(
+ selectedWorker, v -> rand.nextInt(usableDiskInfos.size()));
+ while (usableDiskInfos.get(diskIndex).usableSlots <= 0) {
+ diskIndex = (diskIndex + 1) % usableDiskInfos.size();
+ }
+ usableDiskInfos.get(diskIndex).usableSlots--;
+ DiskInfo selectedDiskInfo = usableDiskInfos.get(diskIndex).diskInfo;
+ if (selectedDiskInfo.storageType().isDFS()) {
+ storageInfo = new StorageInfo("", selectedDiskInfo.storageType(), availableStorageTypes);
+ } else {
+ storageInfo =
+ new StorageInfo(
+ selectedDiskInfo.mountPoint(),
+ selectedDiskInfo.storageType(),
+ availableStorageTypes);
+ workerDiskIndex.put(selectedWorker, (diskIndex + 1) % usableDiskInfos.size());
+ }
+ } else {
+ if (StorageInfo.localDiskAvailable(availableStorageTypes)) {
+ DiskInfo[] diskInfos =
+ selectedWorker.diskInfos().values().stream()
+ .filter(p -> !p.storageType().isDFS())
+ .collect(Collectors.toList())
+ .toArray(new DiskInfo[0]);
+ int diskIndex =
+ workerDiskIndex.computeIfAbsent(selectedWorker, v -> rand.nextInt(diskInfos.length));
+ storageInfo =
+ new StorageInfo(
+ diskInfos[diskIndex].mountPoint(),
+ diskInfos[diskIndex].storageType(),
+ availableStorageTypes);
+ workerDiskIndex.put(selectedWorker, (diskIndex + 1) % diskInfos.length);
+ } else if (StorageInfo.S3Available(availableStorageTypes)) {
+ storageInfo = new StorageInfo("", S3, availableStorageTypes);
+ } else if (StorageInfo.OSSAvailable(availableStorageTypes)) {
+ storageInfo = new StorageInfo("", OSS, availableStorageTypes);
+ } else if (StorageInfo.HDFSAvailable(availableStorageTypes)) {
+ storageInfo = new StorageInfo("", HDFS, availableStorageTypes);
+ } else if (StorageInfo.memoryAvailable(availableStorageTypes)) {
+ storageInfo = new StorageInfo("", StorageInfo.Type.MEMORY, availableStorageTypes);
+ } else {
+ throw new IllegalStateException("no storage type available");
+ }
+ }
+ return storageInfo;
+ }
+
+ /**
+ * If interruptionAware = true, select workers based on 2 main criteria: <br>
+ * 1. Workers that have no nextInterruptionNotice are the first priority and are included in the
+ * 1st pass for slot selection. <br>
+ * 2. Workers that have a later interruption notice are a little less deprioritized, and are
+ * included in the 2nd pass for slot selection. This is determined by nextInterruptionNotice above
+ * a certain percentage threshold.<br>
+ * All other workers are considered least priority, and are only included for slot selection in
+ * the worst case. <br>
+ */
+ static Tuple3<List<WorkerInfo>, List<WorkerInfo>, List<WorkerInfo>>
+ prioritizeWorkersBasedOnInterruptionNotice(
+ List<WorkerInfo> workers,
+ boolean shouldReplicate,
+ boolean shouldRackAware,
+ double percentileThreshold) {
+ Map<Boolean, List<WorkerInfo>> partitioned =
+ workers.stream().collect(Collectors.partitioningBy(WorkerInfo::hasInterruptionNotice));
+ List<WorkerInfo> workersWithInterruptions = partitioned.get(true);
+ List<WorkerInfo> workersWithoutInterruptions = partitioned.get(false);
+ // Timestamps towards the boundary of `percentileThreshold` might be the same. Given this
+ // is a stable sort, it makes sense to randomize these hosts so that the same hosts are not
+ // consistently selected.
+ Collections.shuffle(workersWithInterruptions);
+ workersWithInterruptions.sort(
+ Comparator.comparingLong(WorkerInfo::nextInterruptionNotice).reversed());
+ int requiredNodes =
+ (int) Math.floor((percentileThreshold * workersWithInterruptions.size()) / 100.0);
+
+ List<WorkerInfo> workersWithLateInterruptions =
+ new ArrayList<>(workersWithInterruptions.subList(0, requiredNodes));
+ List<WorkerInfo> workersWithEarlyInterruptions =
+ new ArrayList<>(
+ workersWithInterruptions.subList(requiredNodes, workersWithInterruptions.size()));
+ if (shouldReplicate && shouldRackAware) {
+ return Tuple3.apply(
+ generateRackAwareWorkers(workersWithoutInterruptions),
+ Collections.unmodifiableList(workersWithLateInterruptions),
+ generateRackAwareWorkers(workersWithEarlyInterruptions));
+ }
+ return Tuple3.apply(
+ Collections.unmodifiableList(workersWithoutInterruptions),
+ Collections.unmodifiableList(workersWithLateInterruptions),
+ Collections.unmodifiableList(workersWithEarlyInterruptions));
+ }
+
+ /**
+ * Allocates slots by progressively relaxing soft constraints: strategy budgets, interruption
+ * preference, and finally rack awareness.
+ */
+ private static Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>>
+ locateSlots(
+ List<Integer> partitionIds,
+ List<WorkerInfo> workersList,
+ Map<WorkerInfo, List<UsableDiskInfo>> slotBudgets,
+ boolean shouldReplicate,
+ boolean shouldRackAware,
+ int availableStorageTypes,
+ boolean interruptionAware,
+ int interruptionAwareThreshold) {
+
+ List<WorkerInfo> budgetWorkers = new ArrayList<>(slotBudgets.keySet());
+ List<WorkerInfo> allWorkers = workersList;
+ if (shouldReplicate && shouldRackAware) {
+ budgetWorkers = generateRackAwareWorkers(budgetWorkers);
+ allWorkers = generateRackAwareWorkers(allWorkers);
+ }
+
+ Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots =
+ new HashMap<>();
+ List<WorkerInfo> preferredWorkers;
+ List<WorkerInfo> workersWithLateInterruptions = Collections.emptyList();
+ List<WorkerInfo> workersWithEarlyInterruptions = Collections.emptyList();
+ if (interruptionAware) {
+ Tuple3<List<WorkerInfo>, List<WorkerInfo>, List<WorkerInfo>>
+ workersBasedOnInterruptionNotice =
+ prioritizeWorkersBasedOnInterruptionNotice(
+ budgetWorkers, shouldReplicate, shouldRackAware, interruptionAwareThreshold);
+ preferredWorkers = workersBasedOnInterruptionNotice._1();
+ workersWithLateInterruptions = workersBasedOnInterruptionNotice._2();
+ workersWithEarlyInterruptions = workersBasedOnInterruptionNotice._3();
+ } else {
+ preferredWorkers = budgetWorkers;
+ }
+
+ // First honor the budgets computed by the selected strategy.
+ List<Integer> remainingPartitionIds =
+ tryAllocateWithinBudgets(
+ slots,
+ partitionIds,
+ preferredWorkers,
+ slotBudgets,
+ shouldReplicate,
+ shouldRackAware,
+ availableStorageTypes);
+ logger.debug(
+ "Remaining number of partitionIds after budget-constrained allocation: {}",
+ remainingPartitionIds.size());
+
+ // Relax strategy budgets while retaining the interruption-aware worker preference.
+ if (interruptionAware && !remainingPartitionIds.isEmpty()) {
+ remainingPartitionIds =
+ tryAllocateInterruptionAware(
+ slots,
+ remainingPartitionIds,
+ preferredWorkers,
+ workersWithLateInterruptions,
+ workersWithEarlyInterruptions,
+ shouldReplicate,
+ shouldRackAware,
+ availableStorageTypes);
+ logger.debug(
+ "Remaining number of partitionIds after interruption-aware fallback: {}",
+ remainingPartitionIds.size());
+ }
+
+ // Budgets and interruption preference are soft constraints. Use every worker if necessary.
+ if (!remainingPartitionIds.isEmpty()) {
+ remainingPartitionIds =
+ tryAllocateBestEffort(
+ slots,
+ remainingPartitionIds,
+ allWorkers,
+ shouldReplicate,
+ shouldRackAware,
+ availableStorageTypes);
+ logger.debug(
+ "Remaining number of partitionIds after best-effort allocation: {}",
+ remainingPartitionIds.size());
+ }
+
+ // Rack awareness is also a soft constraint and is relaxed only as the final fallback.
+ if (shouldReplicate && shouldRackAware && !remainingPartitionIds.isEmpty()) {
+ remainingPartitionIds =
+ tryAllocateBestEffort(
+ slots,
+ remainingPartitionIds,
+ allWorkers,
+ shouldReplicate,
+ false,
+ availableStorageTypes);
+ logger.debug(
+ "Remaining number of partitionIds after disabling rack-aware placement: {}",
+ remainingPartitionIds.size());
+ }
+ return slots;
+ }
+
+ private static List<Integer> tryAllocateWithinBudgets(
+ Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots,
+ List<Integer> partitionIds,
+ List<WorkerInfo> workers,
+ Map<WorkerInfo, List<UsableDiskInfo>> slotBudgets,
+ boolean shouldReplicate,
+ boolean shouldRackAware,
+ int availableStorageTypes) {
+ return tryAllocateSlots(
+ slots,
+ partitionIds,
+ workers,
+ workers,
+ slotBudgets,
+ shouldReplicate,
+ shouldRackAware,
+ availableStorageTypes,
+ true);
+ }
+
+ private static List<Integer> tryAllocateInterruptionAware(
+ Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots,
+ List<Integer> partitionIds,
+ List<WorkerInfo> preferredWorkers,
+ List<WorkerInfo> workersWithLateInterruptions,
+ List<WorkerInfo> workersWithEarlyInterruptions,
+ boolean shouldReplicate,
+ boolean shouldRackAware,
+ int availableStorageTypes) {
+ List<WorkerInfo> primaryWorkerCandidates = new ArrayList<>(preferredWorkers);
+ primaryWorkerCandidates.addAll(workersWithLateInterruptions);
+ if (shouldReplicate && shouldRackAware) {
+ primaryWorkerCandidates = generateRackAwareWorkers(primaryWorkerCandidates);
+ }
+ return tryAllocateSlots(
+ slots,
+ partitionIds,
+ primaryWorkerCandidates,
+ workersWithEarlyInterruptions,
+ null,
+ shouldReplicate,
+ shouldRackAware,
+ availableStorageTypes,
+ false);
+ }
+
+ private static List<Integer> tryAllocateBestEffort(
+ Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots,
+ List<Integer> partitionIds,
+ List<WorkerInfo> workers,
+ boolean shouldReplicate,
+ boolean shouldRackAware,
+ int availableStorageTypes) {
+ return tryAllocateSlots(
+ slots,
+ partitionIds,
+ workers,
+ workers,
+ null,
+ shouldReplicate,
+ shouldRackAware,
+ availableStorageTypes,
+ true);
+ }
+
+ /**
+ * The rack distribution of the input workers list is essentially random, and in degenerate cases
+ * the rack aware slot selection ends up skipping sub list's of hosts (in same rack) - which
+ * results in uneven distribution of replica selection. For example given worker list: [h1r1,
+ * h2r1, h3r1, h4r2, h5r2, h6r2] if primary is h3r1 and replica index is pointing to h1r1, it will
+ * skip both h2r1 and h3r1 in order to pick h4r2; and for the next slot, primary will be h4r2 and
+ * will skip all the r2 hosts in order to pick h1r1. This ends up being suboptimal where some
+ * hosts are picked a lot more than others (due to the worker and worker/rack distribution). In
+ * order to mitigate this, we reorder the worker list by redistributing the workers based on rack
+ * to increase the rack diversity between adjoining workers, so that we minimize skipping over
+ * consecutive hosts.
+ */
+ static List<WorkerInfo> generateRackAwareWorkers(List<WorkerInfo> workers) {
+
+ List<Map.Entry<String, LinkedList<WorkerInfo>>> sortedRackToHosts;
+ {
+ Map<String, LinkedList<WorkerInfo>> map = new HashMap<>();
+ for (WorkerInfo worker : workers) {
+ map.computeIfAbsent(worker.networkLocation(), key -> new LinkedList<>()).add(worker);
+ }
+ sortedRackToHosts = new ArrayList<>(map.entrySet());
+ // reverse sort by number of hosts per rack
+ sortedRackToHosts.sort(
+ (o1, o2) -> Integer.compare(o2.getValue().size(), o1.getValue().size()));
+ }
+
+ ArrayList<WorkerInfo> result = new ArrayList<>(workers.size());
+ int count = 0;
+ final int numWorkers = workers.size();
+ while (count < numWorkers) {
+ Iterator<Map.Entry<String, LinkedList<WorkerInfo>>> iter = sortedRackToHosts.iterator();
+ while (iter.hasNext()) {
+ LinkedList<WorkerInfo> workerList = iter.next().getValue();
+ result.add(workerList.removeFirst());
+ count++;
+ if (workerList.isEmpty()) {
+ iter.remove();
+ }
+ }
+ }
+
+ return Collections.unmodifiableList(result);
+ }
+
+ /**
+ * Attempts to allocate the given partitions by scanning primary and replica worker candidates in
+ * round-robin order.
+ *
+ * @param slots the slots that have been assigned for each partitionId
+ * @param partitionIds the partitionIds that require slot selection still
+ * @param primaryWorkers list of worker candidates that can be used for primary workers.
+ * @param replicaWorkers list of worker candidates that can be used for replica workers.
+ * @param slotBudgets budgets for each available slot based on worker characteristics
+ * @param shouldReplicate if replication is enabled within the cluster
+ * @param shouldRackAware if rack-aware replication is enabled within the cluster.
+ * @param availableStorageTypes available storage types coming from the offer slots request.
+ * @param sameWorkerCandidates whether primary and replica indexes refer to the same worker
+ * candidate list, in which case equal indexes identify the same worker
+ * @return the partitionIds that were not able to be assigned slots in this iteration with the
+ * current primary and replica worker candidates and slot budgets.
+ */
+ private static List<Integer> tryAllocateSlots(
+ Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots,
+ List<Integer> partitionIds,
+ List<WorkerInfo> primaryWorkers,
+ List<WorkerInfo> replicaWorkers,
+ Map<WorkerInfo, List<UsableDiskInfo>> slotBudgets,
+ boolean shouldReplicate,
+ boolean shouldRackAware,
+ int availableStorageTypes,
+ boolean sameWorkerCandidates) {
+ if (primaryWorkers.isEmpty() || (shouldReplicate && replicaWorkers.isEmpty())) {
+ return partitionIds;
+ }
+ // Tracks the next disk to try for each worker during this allocation pass.
+ Map<WorkerInfo, Integer> workerDiskIndex = new HashMap<>();
+ List<Integer> remainingPartitionIds = new LinkedList<>(partitionIds);
+
+ int primaryIndex = rand.nextInt(primaryWorkers.size());
+ int replicaIndex = shouldReplicate ? rand.nextInt(replicaWorkers.size()) : -1;
+
+ // Preserve the original allocation order while keeping removal from the linked list O(1).
+ ListIterator<Integer> partitionIterator =
+ remainingPartitionIds.listIterator(remainingPartitionIds.size());
+ while (partitionIterator.hasPrevious()) {
+ int partitionId = partitionIterator.previous();
+ int selectedPrimaryIndex =
+ findNextWorkerIndex(
+ primaryWorkers.size(),
+ primaryIndex,
+ index -> canAssign(slotBudgets, primaryWorkers.get(index), availableStorageTypes));
+ if (selectedPrimaryIndex < 0) {
+ break;
+ }
+ WorkerInfo primaryWorker = primaryWorkers.get(selectedPrimaryIndex);
+
+ StorageInfo primaryStorageInfo =
+ buildStorageInfo(primaryWorker, slotBudgets, workerDiskIndex, availableStorageTypes);
+ PartitionLocation primaryPartition =
+ createLocation(
+ partitionId, primaryWorker, null, primaryStorageInfo, PartitionLocation.Mode.PRIMARY);
+
+ if (shouldReplicate) {
+ int selectedReplicaIndex;
+ if (slotBudgets != null) {
+ selectedReplicaIndex =
+ findNextWorkerIndex(
+ replicaWorkers.size(),
+ replicaIndex,
+ index ->
+ !(sameWorkerCandidates && index == selectedPrimaryIndex)
+ && haveUsableSlots(slotBudgets, replicaWorkers.get(index))
+ && satisfyRackAware(
+ shouldRackAware, primaryWorker, replicaWorkers.get(index)));
+ } else if (shouldRackAware) {
+ selectedReplicaIndex =
+ findNextWorkerIndex(
+ replicaWorkers.size(),
+ replicaIndex,
+ index ->
+ !(sameWorkerCandidates && index == selectedPrimaryIndex)
+ && satisfyRackAware(true, primaryWorker, replicaWorkers.get(index)));
+ } else if (StorageInfo.localDiskAvailable(availableStorageTypes)) {
+ selectedReplicaIndex =
+ findNextWorkerIndex(
+ replicaWorkers.size(),
+ replicaIndex,
+ index ->
+ !(sameWorkerCandidates && index == selectedPrimaryIndex)
+ && replicaWorkers.get(index).haveDisk());
+ } else {
+ selectedReplicaIndex = replicaIndex;
+ }
+ if (selectedReplicaIndex < 0) {
+ break;
+ }
+ WorkerInfo replicaWorker = replicaWorkers.get(selectedReplicaIndex);
+
+ StorageInfo replicaStorageInfo =
+ slotBudgets == null && shouldRackAware
+ ? primaryStorageInfo
+ : buildStorageInfo(
+ replicaWorker, slotBudgets, workerDiskIndex, availableStorageTypes);
+ PartitionLocation replicaPartition =
+ createLocation(
+ partitionId,
+ replicaWorker,
+ primaryPartition,
+ replicaStorageInfo,
+ PartitionLocation.Mode.REPLICA);
+ primaryPartition.setPeer(replicaPartition);
+ addLocation(slots, replicaWorker, replicaPartition);
+ replicaIndex = (selectedReplicaIndex + 1) % replicaWorkers.size();
+ }
+
+ addLocation(slots, primaryWorker, primaryPartition);
+ primaryIndex = (selectedPrimaryIndex + 1) % primaryWorkers.size();
+ partitionIterator.remove();
+ }
+ return remainingPartitionIds;
+ }
+
+ private static int findNextWorkerIndex(int workerCount, int startIndex, IntPredicate isEligible) {
+ int workerIndex = startIndex;
+ do {
+ if (isEligible.test(workerIndex)) {
+ return workerIndex;
+ }
+ workerIndex = (workerIndex + 1) % workerCount;
+ } while (workerIndex != startIndex);
+ return -1;
+ }
+
+ private static boolean canAssign(
+ Map<WorkerInfo, List<UsableDiskInfo>> slotBudgets,
+ WorkerInfo worker,
+ int availableStorageTypes) {
+ if (slotBudgets != null) {
+ return haveUsableSlots(slotBudgets, worker);
+ }
+ return !StorageInfo.localDiskAvailable(availableStorageTypes) || worker.haveDisk();
+ }
+
+ private static boolean haveUsableSlots(
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets, WorkerInfo worker) {
+ List<UsableDiskInfo> usableDiskInfos = budgets.get(worker);
+ return usableDiskInfos != null
+ && usableDiskInfos.stream().anyMatch(disk -> disk.usableSlots > 0);
+ }
+
+ private static boolean satisfyRackAware(
+ boolean shouldRackAware, WorkerInfo primaryWorker, WorkerInfo replicaWorker) {
+ return !shouldRackAware
+ || !Objects.equals(primaryWorker.networkLocation(), replicaWorker.networkLocation());
+ }
+
+ private static PartitionLocation createLocation(
+ int partitionIndex,
+ WorkerInfo workerInfo,
+ PartitionLocation peer,
+ StorageInfo storageInfo,
+ PartitionLocation.Mode mode) {
+ return new PartitionLocation(
+ partitionIndex,
+ 0,
+ workerInfo.host(),
+ workerInfo.rpcPort(),
+ workerInfo.pushPort(),
+ workerInfo.fetchPort(),
+ workerInfo.replicatePort(),
+ mode,
+ peer,
+ storageInfo,
+ new RoaringBitmap());
+ }
+
+ private static void addLocation(
+ Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots,
+ WorkerInfo worker,
+ PartitionLocation location) {
+ Tuple2<List<PartitionLocation>, List<PartitionLocation>> locations =
+ slots.computeIfAbsent(
+ worker, ignored -> new Tuple2<>(new ArrayList<>(), new ArrayList<>()));
+ if (location.getMode() == PartitionLocation.Mode.PRIMARY) {
+ locations._1.add(location);
+ } else {
+ locations._2.add(location);
+ }
+ }
+
+ public static Map<WorkerInfo, Map<String, Integer>> slotsToDiskAllocations(
+ Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots) {
+ Iterator<WorkerInfo> workers = slots.keySet().iterator();
+ Map<WorkerInfo, Map<String, Integer>> workerToSlots = new HashMap<>();
+ while (workers.hasNext()) {
+ WorkerInfo worker = workers.next();
+ Map<String, Integer> slotsPerDisk =
+ workerToSlots.computeIfAbsent(worker, v -> new HashMap<>());
+ List<PartitionLocation> jointLocations = new ArrayList<>();
+ jointLocations.addAll(slots.get(worker)._1);
+ jointLocations.addAll(slots.get(worker)._2);
+ for (PartitionLocation location : jointLocations) {
+ String mountPoint = location.getStorageInfo().getMountPoint();
+ // skip non local disks slots
+ if (!mountPoint.isEmpty()) {
+ if (slotsPerDisk.containsKey(mountPoint)) {
+ slotsPerDisk.put(mountPoint, slotsPerDisk.get(mountPoint) + 1);
+ } else {
+ slotsPerDisk.put(mountPoint, 1);
+ }
+ }
+ }
+ }
+ return workerToSlots;
+ }
+}
diff --git a/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAssignStrategy.java b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAssignStrategy.java
new file mode 100644
index 000000000..52b8982ae
--- /dev/null
+++ b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAssignStrategy.java
@@ -0,0 +1,41 @@
+/*
+ * 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.service.deploy.master.slotsalloc;
+
+import java.util.List;
+import java.util.Map;
+
+import org.apache.celeborn.common.meta.WorkerInfo;
+
+/**
+ * Computes per-disk slot budgets for {@link SlotsAllocator}. The allocator owns worker placement,
+ * replication, rack awareness, interruption awareness, and best-effort placement after the budgets
+ * are exhausted.
+ */
+public interface SlotsAssignStrategy {
+
+ /**
+ * Returns a fresh, mutable budget map for one allocation request. The allocator consumes the
+ * returned budgets while placing partitions.
+ */
+ Map<WorkerInfo, List<UsableDiskInfo>> computeSlotBudgets(
+ List<WorkerInfo> workers,
+ List<Integer> partitionIds,
+ boolean shouldReplicate,
+ int availableStorageTypes);
+}
diff --git a/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAssignStrategyManager.java b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAssignStrategyManager.java
new file mode 100644
index 000000000..98c002c90
--- /dev/null
+++ b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAssignStrategyManager.java
@@ -0,0 +1,176 @@
+/*
+ * 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.service.deploy.master.slotsalloc;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Objects;
+import java.util.ServiceLoader;
+import java.util.TreeSet;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.server.common.service.config.ConfigService;
+import org.apache.celeborn.server.common.service.config.SystemConfig;
+
+/**
+ * Discovers slot assignment strategy providers and atomically switches the active strategy after
+ * dynamic configuration updates.
+ */
+public final class SlotsAssignStrategyManager {
+
+ private static final Logger LOG = LoggerFactory.getLogger(SlotsAssignStrategyManager.class);
+
+ private final CelebornConf staticConf;
+ private final ConfigService configService;
+ private final Map<String, SlotsAssignStrategyProvider> providersByName;
+
+ private volatile ConfiguredStrategy configuredStrategy;
+ private Map<String, String> lastAttemptedDynamicConfigs;
+
+ public SlotsAssignStrategyManager(CelebornConf staticConf, ConfigService configService) {
+ this.staticConf = staticConf.clone();
+ this.configService = configService;
+ this.providersByName =
+ Collections.unmodifiableMap(
+ loadProviders(ServiceLoader.load(SlotsAssignStrategyProvider.class)));
+
+ Map<String, String> initialDynamicConfigs = currentDynamicConfigs();
+ this.lastAttemptedDynamicConfigs = initialDynamicConfigs;
+ try {
+ this.configuredStrategy = createStrategy(initialDynamicConfigs);
+ } catch (RuntimeException e) {
+ if (initialDynamicConfigs.isEmpty()) {
+ throw e;
+ }
+ LOG.error(
+ "Failed to initialize slots assignment strategy from dynamic configuration; "
+ + "falling back to static configuration",
+ e);
+ this.configuredStrategy = createStrategy(Collections.emptyMap());
+ }
+
+ if (configService != null) {
+ configService.registerListenerOnConfigUpdate(this::reload);
+ // Close the gap between taking the initial snapshot and registering the listener.
+ reload();
+ }
+ }
+
+ public SlotsAssignStrategy getStrategy() {
+ return configuredStrategy.strategy;
+ }
+
+ static Map<String, SlotsAssignStrategyProvider> loadProviders(
+ Iterable<SlotsAssignStrategyProvider> providers) {
+ Map<String, SlotsAssignStrategyProvider> providersByName = new LinkedHashMap<>();
+ for (SlotsAssignStrategyProvider provider : providers) {
+ String providerName = provider.getName();
+ if (providerName == null || providerName.isEmpty()) {
+ throw new IllegalStateException(
+ "Slots assignment strategy provider "
+ + provider.getClass().getName()
+ + " has an empty name");
+ }
+
+ String normalizedName = providerName.toUpperCase(Locale.ROOT);
+ SlotsAssignStrategyProvider previous = providersByName.put(normalizedName, provider);
+ if (previous != null) {
+ throw new IllegalStateException(
+ "Multiple slots assignment strategy providers are registered for '"
+ + providerName
+ + "': "
+ + previous.getClass().getName()
+ + " and "
+ + provider.getClass().getName());
+ }
+ }
+ return providersByName;
+ }
+
+ private synchronized void reload() {
+ Map<String, String> latestDynamicConfigs = currentDynamicConfigs();
+ if (latestDynamicConfigs.equals(lastAttemptedDynamicConfigs)) {
+ return;
+ }
+
+ // Record the snapshot before creating the strategy so an unchanged invalid configuration is
+ // reported only once instead of being retried on every polling cycle.
+ lastAttemptedDynamicConfigs = latestDynamicConfigs;
+ try {
+ ConfiguredStrategy updatedStrategy = createStrategy(latestDynamicConfigs);
+ configuredStrategy = updatedStrategy;
+ LOG.info(
+ "Reloaded slots assignment strategy provider {} after dynamic configuration update",
+ updatedStrategy.providerName);
+ } catch (RuntimeException e) {
+ LOG.error(
+ "Failed to reload slots assignment strategy; keeping provider {}",
+ configuredStrategy.providerName,
+ e);
+ }
+ }
+
+ private ConfiguredStrategy createStrategy(Map<String, String> dynamicConfigs) {
+ CelebornConf effectiveConf = staticConf.clone();
+ dynamicConfigs.forEach(effectiveConf::set);
+
+ String providerName = effectiveConf.masterSlotAssignPolicyName();
+ SlotsAssignStrategyProvider provider =
+ providersByName.get(providerName.toUpperCase(Locale.ROOT));
+ if (provider == null) {
+ throw new IllegalArgumentException(
+ "No slots assignment strategy provider is registered for '"
+ + providerName
+ + "'. Available providers: "
+ + new TreeSet<>(providersByName.keySet()));
+ }
+ SlotsAssignStrategy strategy =
+ Objects.requireNonNull(
+ provider.create(effectiveConf),
+ "Slots assignment strategy provider '" + providerName + "' returned null");
+ return new ConfiguredStrategy(providerName, strategy);
+ }
+
+ private Map<String, String> currentDynamicConfigs() {
+ if (configService == null) {
+ return Collections.emptyMap();
+ }
+
+ SystemConfig systemConfig = configService.getSystemConfigFromCache();
+ if (systemConfig == null || systemConfig.getConfigs() == null) {
+ return Collections.emptyMap();
+ }
+ return Collections.unmodifiableMap(new HashMap<>(systemConfig.getConfigs()));
+ }
+
+ private static final class ConfiguredStrategy {
+ private final String providerName;
+ private final SlotsAssignStrategy strategy;
+
+ private ConfiguredStrategy(String providerName, SlotsAssignStrategy strategy) {
+ this.providerName = providerName;
+ this.strategy = strategy;
+ }
+ }
+}
diff --git a/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAssignStrategyProvider.java b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAssignStrategyProvider.java
new file mode 100644
index 000000000..c0ef82d06
--- /dev/null
+++ b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAssignStrategyProvider.java
@@ -0,0 +1,46 @@
+/*
+ * 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.service.deploy.master.slotsalloc;
+
+import org.apache.celeborn.common.CelebornConf;
+
+/**
+ * Service provider interface for named slot assignment strategies.
+ *
+ * <p>Implementations must have a public no-argument constructor and be registered in {@code
+ * META-INF/services/org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAssignStrategyProvider}.
+ * Provider names are matched case-insensitively against {@code celeborn.master.slot.assign.policy}.
+ * This is a master extension point: providers must be compiled against the same Celeborn release
+ * and Scala binary version as the master where they are installed.
+ */
+public interface SlotsAssignStrategyProvider {
+
+ /** Returns the name used to select this provider from configuration. */
+ String getName();
+
+ /**
+ * Creates a configured strategy instance for a master.
+ *
+ * <p>A custom provider may read its own namespaced keys directly from the supplied Celeborn
+ * configuration. When system-level dynamic configuration is enabled, it overrides the static
+ * configuration in this snapshot and the master calls {@code create} after an update. The master
+ * does not retry an unchanged system-level configuration snapshot. Configuration should therefore
+ * be validated and captured when the strategy is created rather than read during slot allocation.
+ */
+ SlotsAssignStrategy create(CelebornConf conf);
+}
diff --git a/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/UsableDiskInfo.java b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/UsableDiskInfo.java
new file mode 100644
index 000000000..e242e580e
--- /dev/null
+++ b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/UsableDiskInfo.java
@@ -0,0 +1,35 @@
+/*
+ * 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.service.deploy.master.slotsalloc;
+
+import org.apache.celeborn.common.meta.DiskInfo;
+
+public class UsableDiskInfo {
+ DiskInfo diskInfo;
+ long usableSlots;
+
+ public UsableDiskInfo(DiskInfo diskInfo) {
+ this.diskInfo = diskInfo;
+ this.usableSlots = diskInfo.getAvailableSlots();
+ }
+
+ public UsableDiskInfo(DiskInfo diskInfo, long usableSlots) {
+ this.diskInfo = diskInfo;
+ this.usableSlots = usableSlots;
+ }
+}
diff --git a/master/src/main/resources/META-INF/services/org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAssignStrategyProvider b/master/src/main/resources/META-INF/services/org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAssignStrategyProvider
new file mode 100644
index 000000000..5030c56e8
--- /dev/null
+++ b/master/src/main/resources/META-INF/services/org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAssignStrategyProvider
@@ -0,0 +1,19 @@
+#
+# 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.
+#
+
+org.apache.celeborn.service.deploy.master.slotsalloc.RoundRobinSlotsAssignStrategyProvider
+org.apache.celeborn.service.deploy.master.slotsalloc.LoadAwareSlotsAssignStrategyProvider
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 75ceb951a..ca1b67bee 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
@@ -57,6 +57,7 @@ import org.apache.celeborn.service.deploy.master.audit.ShuffleAuditLogger
import org.apache.celeborn.service.deploy.master.clustermeta.SingleMasterMetaManager
import org.apache.celeborn.service.deploy.master.clustermeta.ha.{HAHelper, HAMasterMetaManager, MetaHandler}
import org.apache.celeborn.service.deploy.master.quota.QuotaManager
+import org.apache.celeborn.service.deploy.master.slotsalloc.{SlotsAllocator, SlotsAssignStrategyManager}
import org.apache.celeborn.service.deploy.master.tags.TagsManager
private[celeborn] class Master(
@@ -207,12 +208,8 @@ private[celeborn] class Master(
private val slotsAssignMaxWorkers = conf.masterSlotAssignMaxWorkers
private val slotsAssignMinWorkers = conf.masterSlotAssignMinWorkers
private val slotsAssignExtraSlots = conf.masterSlotAssignExtraSlots
- private val slotsAssignLoadAwareDiskGroupNum = conf.masterSlotAssignLoadAwareDiskGroupNum
- private val slotsAssignLoadAwareDiskGroupGradient =
- conf.masterSlotAssignLoadAwareDiskGroupGradient
- private val loadAwareFlushTimeWeight = conf.masterSlotAssignLoadAwareFlushTimeWeight
- private val loadAwareFetchTimeWeight = conf.masterSlotAssignLoadAwareFetchTimeWeight
- private val loadAwareActiveSlotsWeight = conf.masterSlotAssignLoadAwareActiveSlotsWeight
+ private val slotsAssignStrategyManager =
+ new SlotsAssignStrategyManager(conf, configService)
private val estimatedPartitionSizeUpdaterInitialDelay =
conf.estimatedPartitionSizeUpdaterInitialDelay
@@ -233,7 +230,6 @@ private[celeborn] class Master(
estimatedPartitionSizeUpdaterInitialDelay,
estimatedPartitionSizeForEstimationUpdateInterval,
TimeUnit.MILLISECONDS)
- private val slotsAssignPolicy = conf.masterSlotAssignPolicy
private val slotsAssignInterruptionAware = conf.masterSlotAssignInterruptionAware
private val slotsAssignInterruptionAwareThreshold =
conf.masterSlotsAssignInterruptionAwareThreshold
@@ -977,30 +973,15 @@ private[celeborn] class Master(
val slots =
masterSource.sample(MasterSource.OFFER_SLOTS_TIME, s"offerSlots-${Random.nextInt()}") {
statusSystem.workersMap.synchronized {
- if (slotsAssignPolicy == SlotsAssignPolicy.LOADAWARE) {
- SlotsAllocator.offerSlotsLoadAware(
- selectedWorkers,
- requestSlots.partitionIdList,
- requestSlots.shouldReplicate,
- requestSlots.shouldRackAware,
- slotsAssignLoadAwareDiskGroupNum,
- slotsAssignLoadAwareDiskGroupGradient,
- loadAwareFlushTimeWeight,
- loadAwareFetchTimeWeight,
- loadAwareActiveSlotsWeight,
- requestSlots.availableStorageTypes,
- slotsAssignInterruptionAware,
- slotsAssignInterruptionAwareThreshold)
- } else {
- SlotsAllocator.offerSlotsRoundRobin(
- selectedWorkers,
- requestSlots.partitionIdList,
- requestSlots.shouldReplicate,
- requestSlots.shouldRackAware,
- requestSlots.availableStorageTypes,
- slotsAssignInterruptionAware,
- slotsAssignInterruptionAwareThreshold)
- }
+ SlotsAllocator.offerSlots(
+ selectedWorkers,
+ requestSlots.partitionIdList,
+ requestSlots.shouldReplicate,
+ requestSlots.shouldRackAware,
+ requestSlots.availableStorageTypes,
+ slotsAssignInterruptionAware,
+ slotsAssignInterruptionAwareThreshold,
+ slotsAssignStrategyManager.getStrategy)
}
}
diff --git a/master/src/test/java/org/apache/celeborn/service/deploy/master/BuildStorageInfoSuiteJ.java b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/BuildStorageInfoSuiteJ.java
similarity index 55%
rename from master/src/test/java/org/apache/celeborn/service/deploy/master/BuildStorageInfoSuiteJ.java
rename to master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/BuildStorageInfoSuiteJ.java
index cefd81e1b..f3c0d1f03 100644
--- a/master/src/test/java/org/apache/celeborn/service/deploy/master/BuildStorageInfoSuiteJ.java
+++ b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/BuildStorageInfoSuiteJ.java
@@ -15,7 +15,7 @@
* limitations under the License.
*/
-package org.apache.celeborn.service.deploy.master;
+package org.apache.celeborn.service.deploy.master.slotsalloc;
import static org.junit.Assert.*;
@@ -33,10 +33,10 @@ import org.apache.celeborn.common.protocol.StorageInfo;
* <p>The method has two main paths:
*
* <ol>
- * <li><b>With restrictions</b> – a restrictions map is provided, and a disk is selected from the
- * worker's {@link SlotsAllocator.UsableDiskInfo} list.
- * <li><b>Without restrictions</b> – the restrictions map is {@code null}, and storage type is
- * derived from the {@code availableStorageTypes} bitmask.
+ * <li><b>With budgets</b> – a budget map is provided, and a disk is selected from the worker's
+ * {@link UsableDiskInfo} list.
+ * <li><b>Without budgets</b> – the budget map is {@code null}, and storage type is derived from
+ * the {@code availableStorageTypes} bitmask.
* </ol>
*/
public class BuildStorageInfoSuiteJ {
@@ -53,38 +53,32 @@ public class BuildStorageInfoSuiteJ {
return new WorkerInfo(host, 9001, 9002, 9003, 9004, 9005, disks, null);
}
- private Map<WorkerInfo, List<SlotsAllocator.UsableDiskInfo>> restrictionsFor(
- WorkerInfo worker, List<SlotsAllocator.UsableDiskInfo> diskList) {
- Map<WorkerInfo, List<SlotsAllocator.UsableDiskInfo>> restrictions = new HashMap<>();
- restrictions.put(worker, diskList);
- return restrictions;
+ private Map<WorkerInfo, List<UsableDiskInfo>> budgetsFor(
+ WorkerInfo worker, List<UsableDiskInfo> diskList) {
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets = new HashMap<>();
+ budgets.put(worker, diskList);
+ return budgets;
}
// ---------------------------------------------------------------------------
- // Tests: restrictions != null
+ // Tests: budgets != null
// ---------------------------------------------------------------------------
- /**
- * An HDD disk in the restrictions list should produce a StorageInfo with the disk mount point.
- */
+ /** An HDD disk in the budget list should produce a StorageInfo with the disk mount point. */
@Test
- public void testWithRestrictions_HDDDisk() {
+ public void testWithBudgets_HDDDisk() {
DiskInfo disk = makeDiskInfo("/mnt/hdd1", StorageInfo.Type.HDD);
WorkerInfo worker = makeWorker("host1", Collections.singletonMap("/mnt/hdd1", disk));
- SlotsAllocator.UsableDiskInfo usable = new SlotsAllocator.UsableDiskInfo(disk, 10);
- Map<WorkerInfo, List<SlotsAllocator.UsableDiskInfo>> restrictions =
- restrictionsFor(worker, Collections.singletonList(usable));
+ UsableDiskInfo usable = new UsableDiskInfo(disk, 10);
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets =
+ budgetsFor(worker, Collections.singletonList(usable));
Map<WorkerInfo, Integer> workerDiskIndex = new HashMap<>();
workerDiskIndex.put(worker, 0);
StorageInfo result =
SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker),
- 0,
- restrictions,
- workerDiskIndex,
- StorageInfo.ALL_TYPES_AVAILABLE_MASK);
+ worker, budgets, workerDiskIndex, StorageInfo.ALL_TYPES_AVAILABLE_MASK);
assertEquals(StorageInfo.Type.HDD, result.getType());
assertEquals("/mnt/hdd1", result.getMountPoint());
@@ -92,27 +86,21 @@ public class BuildStorageInfoSuiteJ {
assertEquals(9, usable.usableSlots); // consumed one slot
}
- /**
- * An SSD disk in the restrictions list should produce a StorageInfo with the disk mount point.
- */
+ /** An SSD disk in the budget list should produce a StorageInfo with the disk mount point. */
@Test
- public void testWithRestrictions_SSDDisk() {
+ public void testWithBudgets_SSDDisk() {
DiskInfo disk = makeDiskInfo("/mnt/ssd1", StorageInfo.Type.SSD);
WorkerInfo worker = makeWorker("host1", Collections.singletonMap("/mnt/ssd1", disk));
- SlotsAllocator.UsableDiskInfo usable = new SlotsAllocator.UsableDiskInfo(disk, 5);
- Map<WorkerInfo, List<SlotsAllocator.UsableDiskInfo>> restrictions =
- restrictionsFor(worker, Collections.singletonList(usable));
+ UsableDiskInfo usable = new UsableDiskInfo(disk, 5);
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets =
+ budgetsFor(worker, Collections.singletonList(usable));
Map<WorkerInfo, Integer> workerDiskIndex = new HashMap<>();
workerDiskIndex.put(worker, 0);
StorageInfo result =
SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker),
- 0,
- restrictions,
- workerDiskIndex,
- StorageInfo.LOCAL_DISK_MASK);
+ worker, budgets, workerDiskIndex, StorageInfo.LOCAL_DISK_MASK);
assertEquals(StorageInfo.Type.SSD, result.getType());
assertEquals("/mnt/ssd1", result.getMountPoint());
@@ -121,27 +109,22 @@ public class BuildStorageInfoSuiteJ {
}
/**
- * An HDFS disk in the restrictions list should produce a StorageInfo with an empty mount point
- * and HDFS type, regardless of the actual mount-point string stored in DiskInfo.
+ * An HDFS disk in the budget list should produce a StorageInfo with an empty mount point and HDFS
+ * type, regardless of the actual mount-point string stored in DiskInfo.
*/
@Test
- public void testWithRestrictions_HDFSDisk_emptyMountPoint() {
+ public void testWithBudgets_HDFSDisk_emptyMountPoint() {
DiskInfo disk = makeDiskInfo("HDFS", StorageInfo.Type.HDFS);
WorkerInfo worker = makeWorker("host1", Collections.singletonMap("HDFS", disk));
- SlotsAllocator.UsableDiskInfo usable = new SlotsAllocator.UsableDiskInfo(disk, 50);
- Map<WorkerInfo, List<SlotsAllocator.UsableDiskInfo>> restrictions =
- restrictionsFor(worker, Collections.singletonList(usable));
+ UsableDiskInfo usable = new UsableDiskInfo(disk, 50);
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets =
+ budgetsFor(worker, Collections.singletonList(usable));
Map<WorkerInfo, Integer> workerDiskIndex = new HashMap<>();
workerDiskIndex.put(worker, 0);
StorageInfo result =
- SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker),
- 0,
- restrictions,
- workerDiskIndex,
- StorageInfo.HDFS_MASK);
+ SlotsAllocator.buildStorageInfo(worker, budgets, workerDiskIndex, StorageInfo.HDFS_MASK);
assertEquals(StorageInfo.Type.HDFS, result.getType());
assertEquals("", result.getMountPoint());
@@ -150,27 +133,22 @@ public class BuildStorageInfoSuiteJ {
}
/**
- * An S3 disk in the restrictions list should produce a StorageInfo with an empty mount point and
- * S3 type.
+ * An S3 disk in the budget list should produce a StorageInfo with an empty mount point and S3
+ * type.
*/
@Test
- public void testWithRestrictions_S3Disk_emptyMountPoint() {
+ public void testWithBudgets_S3Disk_emptyMountPoint() {
DiskInfo disk = makeDiskInfo("S3", StorageInfo.Type.S3);
WorkerInfo worker = makeWorker("host1", Collections.singletonMap("S3", disk));
- SlotsAllocator.UsableDiskInfo usable = new SlotsAllocator.UsableDiskInfo(disk, 20);
- Map<WorkerInfo, List<SlotsAllocator.UsableDiskInfo>> restrictions =
- restrictionsFor(worker, Collections.singletonList(usable));
+ UsableDiskInfo usable = new UsableDiskInfo(disk, 20);
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets =
+ budgetsFor(worker, Collections.singletonList(usable));
Map<WorkerInfo, Integer> workerDiskIndex = new HashMap<>();
workerDiskIndex.put(worker, 0);
StorageInfo result =
- SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker),
- 0,
- restrictions,
- workerDiskIndex,
- StorageInfo.S3_MASK);
+ SlotsAllocator.buildStorageInfo(worker, budgets, workerDiskIndex, StorageInfo.S3_MASK);
assertEquals(StorageInfo.Type.S3, result.getType());
assertEquals("", result.getMountPoint());
@@ -179,27 +157,22 @@ public class BuildStorageInfoSuiteJ {
}
/**
- * An OSS disk in the restrictions list should produce a StorageInfo with an empty mount point and
- * OSS type.
+ * An OSS disk in the budget list should produce a StorageInfo with an empty mount point and OSS
+ * type.
*/
@Test
- public void testWithRestrictions_OSSDisk_emptyMountPoint() {
+ public void testWithBudgets_OSSDisk_emptyMountPoint() {
DiskInfo disk = makeDiskInfo("OSS", StorageInfo.Type.OSS);
WorkerInfo worker = makeWorker("host1", Collections.singletonMap("OSS", disk));
- SlotsAllocator.UsableDiskInfo usable = new SlotsAllocator.UsableDiskInfo(disk, 30);
- Map<WorkerInfo, List<SlotsAllocator.UsableDiskInfo>> restrictions =
- restrictionsFor(worker, Collections.singletonList(usable));
+ UsableDiskInfo usable = new UsableDiskInfo(disk, 30);
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets =
+ budgetsFor(worker, Collections.singletonList(usable));
Map<WorkerInfo, Integer> workerDiskIndex = new HashMap<>();
workerDiskIndex.put(worker, 0);
StorageInfo result =
- SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker),
- 0,
- restrictions,
- workerDiskIndex,
- StorageInfo.OSS_MASK);
+ SlotsAllocator.buildStorageInfo(worker, budgets, workerDiskIndex, StorageInfo.OSS_MASK);
assertEquals(StorageInfo.Type.OSS, result.getType());
assertEquals("", result.getMountPoint());
@@ -208,33 +181,28 @@ public class BuildStorageInfoSuiteJ {
}
/**
- * When the first disk in the restrictions list has zero usable slots, the method must skip it and
+ * When the first disk in the budget list has zero usable slots, the method must skip it and
* select the next disk with available capacity.
*/
@Test
- public void testWithRestrictions_skipExhaustedDisk() {
+ public void testWithBudgets_skipExhaustedDisk() {
DiskInfo disk1 = makeDiskInfo("/mnt/disk1", StorageInfo.Type.HDD);
DiskInfo disk2 = makeDiskInfo("/mnt/disk2", StorageInfo.Type.HDD);
WorkerInfo worker = makeWorker("host1", new HashMap<>());
- SlotsAllocator.UsableDiskInfo exhausted = new SlotsAllocator.UsableDiskInfo(disk1, 0);
- SlotsAllocator.UsableDiskInfo active = new SlotsAllocator.UsableDiskInfo(disk2, 8);
- List<SlotsAllocator.UsableDiskInfo> diskList = new ArrayList<>();
+ UsableDiskInfo exhausted = new UsableDiskInfo(disk1, 0);
+ UsableDiskInfo active = new UsableDiskInfo(disk2, 8);
+ List<UsableDiskInfo> diskList = new ArrayList<>();
diskList.add(exhausted);
diskList.add(active);
- Map<WorkerInfo, List<SlotsAllocator.UsableDiskInfo>> restrictions =
- restrictionsFor(worker, diskList);
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets = budgetsFor(worker, diskList);
Map<WorkerInfo, Integer> workerDiskIndex = new HashMap<>();
workerDiskIndex.put(worker, 0); // start at the exhausted disk
StorageInfo result =
SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker),
- 0,
- restrictions,
- workerDiskIndex,
- StorageInfo.ALL_TYPES_AVAILABLE_MASK);
+ worker, budgets, workerDiskIndex, StorageInfo.ALL_TYPES_AVAILABLE_MASK);
assertEquals(StorageInfo.Type.HDD, result.getType());
assertEquals("/mnt/disk2", result.getMountPoint()); // disk2 was chosen
@@ -247,26 +215,21 @@ public class BuildStorageInfoSuiteJ {
* advance to the next disk (round-robin).
*/
@Test
- public void testWithRestrictions_localDiskAdvancesDiskIndex() {
+ public void testWithBudgets_localDiskAdvancesDiskIndex() {
DiskInfo disk1 = makeDiskInfo("/mnt/disk1", StorageInfo.Type.HDD);
DiskInfo disk2 = makeDiskInfo("/mnt/disk2", StorageInfo.Type.HDD);
WorkerInfo worker = makeWorker("host1", new HashMap<>());
- List<SlotsAllocator.UsableDiskInfo> diskList = new ArrayList<>();
- diskList.add(new SlotsAllocator.UsableDiskInfo(disk1, 10));
- diskList.add(new SlotsAllocator.UsableDiskInfo(disk2, 10));
+ List<UsableDiskInfo> diskList = new ArrayList<>();
+ diskList.add(new UsableDiskInfo(disk1, 10));
+ diskList.add(new UsableDiskInfo(disk2, 10));
- Map<WorkerInfo, List<SlotsAllocator.UsableDiskInfo>> restrictions =
- restrictionsFor(worker, diskList);
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets = budgetsFor(worker, diskList);
Map<WorkerInfo, Integer> workerDiskIndex = new HashMap<>();
workerDiskIndex.put(worker, 0);
SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker),
- 0,
- restrictions,
- workerDiskIndex,
- StorageInfo.ALL_TYPES_AVAILABLE_MASK);
+ worker, budgets, workerDiskIndex, StorageInfo.ALL_TYPES_AVAILABLE_MASK);
assertEquals(Integer.valueOf(1), workerDiskIndex.get(worker));
}
@@ -276,48 +239,43 @@ public class BuildStorageInfoSuiteJ {
* advanced after slot assignment, because there is only a single logical endpoint.
*/
@Test
- public void testWithRestrictions_HDFSDiskDoesNotAdvanceDiskIndex() {
+ public void testWithBudgets_HDFSDiskDoesNotAdvanceDiskIndex() {
DiskInfo disk = makeDiskInfo("HDFS", StorageInfo.Type.HDFS);
WorkerInfo worker = makeWorker("host1", new HashMap<>());
- Map<WorkerInfo, List<SlotsAllocator.UsableDiskInfo>> restrictions =
- restrictionsFor(
- worker, Collections.singletonList(new SlotsAllocator.UsableDiskInfo(disk, 10)));
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets =
+ budgetsFor(worker, Collections.singletonList(new UsableDiskInfo(disk, 10)));
Map<WorkerInfo, Integer> workerDiskIndex = new HashMap<>();
workerDiskIndex.put(worker, 0);
- SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker), 0, restrictions, workerDiskIndex, StorageInfo.HDFS_MASK);
+ SlotsAllocator.buildStorageInfo(worker, budgets, workerDiskIndex, StorageInfo.HDFS_MASK);
assertEquals(Integer.valueOf(0), workerDiskIndex.get(worker));
}
/**
- * Consecutive calls with two local disks in the restrictions list must cycle through the disks in
+ * Consecutive calls with two local disks in the budget list must cycle through the disks in
* order, demonstrating round-robin selection.
*/
@Test
- public void testWithRestrictions_roundRobinAcrossLocalDisks() {
+ public void testWithBudgets_roundRobinAcrossLocalDisks() {
DiskInfo disk1 = makeDiskInfo("/mnt/disk1", StorageInfo.Type.HDD);
DiskInfo disk2 = makeDiskInfo("/mnt/disk2", StorageInfo.Type.HDD);
WorkerInfo worker = makeWorker("host1", new HashMap<>());
- SlotsAllocator.UsableDiskInfo usable1 = new SlotsAllocator.UsableDiskInfo(disk1, 10);
- SlotsAllocator.UsableDiskInfo usable2 = new SlotsAllocator.UsableDiskInfo(disk2, 10);
- List<SlotsAllocator.UsableDiskInfo> diskList = new ArrayList<>();
+ UsableDiskInfo usable1 = new UsableDiskInfo(disk1, 10);
+ UsableDiskInfo usable2 = new UsableDiskInfo(disk2, 10);
+ List<UsableDiskInfo> diskList = new ArrayList<>();
diskList.add(usable1);
diskList.add(usable2);
- Map<WorkerInfo, List<SlotsAllocator.UsableDiskInfo>> restrictions =
- restrictionsFor(worker, diskList);
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets = budgetsFor(worker, diskList);
Map<WorkerInfo, Integer> workerDiskIndex = new HashMap<>();
workerDiskIndex.put(worker, 0);
- List<WorkerInfo> workers = Collections.singletonList(worker);
-
// First call: disk at index 0
StorageInfo result1 =
SlotsAllocator.buildStorageInfo(
- workers, 0, restrictions, workerDiskIndex, StorageInfo.ALL_TYPES_AVAILABLE_MASK);
+ worker, budgets, workerDiskIndex, StorageInfo.ALL_TYPES_AVAILABLE_MASK);
assertEquals("/mnt/disk1", result1.getMountPoint());
assertEquals(9, usable1.usableSlots);
assertEquals(Integer.valueOf(1), workerDiskIndex.get(worker));
@@ -325,22 +283,22 @@ public class BuildStorageInfoSuiteJ {
// Second call: disk at index 1
StorageInfo result2 =
SlotsAllocator.buildStorageInfo(
- workers, 0, restrictions, workerDiskIndex, StorageInfo.ALL_TYPES_AVAILABLE_MASK);
+ worker, budgets, workerDiskIndex, StorageInfo.ALL_TYPES_AVAILABLE_MASK);
assertEquals("/mnt/disk2", result2.getMountPoint());
assertEquals(9, usable2.usableSlots);
assertEquals(Integer.valueOf(0), workerDiskIndex.get(worker)); // wrapped around
}
// ---------------------------------------------------------------------------
- // Tests: restrictions == null
+ // Tests: budgets == null
// ---------------------------------------------------------------------------
/**
- * When restrictions are {@code null} and all storage types are available (mask = 0), the method
- * must pick a local disk from the worker's disk map.
+ * When budgets are {@code null} and all storage types are available (mask = 0), the method must
+ * pick a local disk from the worker's disk map.
*/
@Test
- public void testWithoutRestrictions_allTypesAvailable_picksLocalDisk() {
+ public void testWithoutBudgets_allTypesAvailable_picksLocalDisk() {
DiskInfo disk = makeDiskInfo("/mnt/hdd1", StorageInfo.Type.HDD);
WorkerInfo worker = makeWorker("host1", Collections.singletonMap("/mnt/hdd1", disk));
Map<WorkerInfo, Integer> workerDiskIndex = new HashMap<>();
@@ -348,11 +306,7 @@ public class BuildStorageInfoSuiteJ {
StorageInfo result =
SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker),
- 0,
- null,
- workerDiskIndex,
- StorageInfo.ALL_TYPES_AVAILABLE_MASK);
+ worker, null, workerDiskIndex, StorageInfo.ALL_TYPES_AVAILABLE_MASK);
assertEquals(StorageInfo.Type.HDD, result.getType());
assertEquals("/mnt/hdd1", result.getMountPoint());
@@ -360,23 +314,18 @@ public class BuildStorageInfoSuiteJ {
}
/**
- * When restrictions are {@code null} and only LOCAL_DISK_MASK is set, the method must pick an SSD
- * disk from the worker and record its mount point.
+ * When budgets are {@code null} and only LOCAL_DISK_MASK is set, the method must pick an SSD disk
+ * from the worker and record its mount point.
*/
@Test
- public void testWithoutRestrictions_localDiskMask() {
+ public void testWithoutBudgets_localDiskMask() {
DiskInfo disk = makeDiskInfo("/mnt/ssd1", StorageInfo.Type.SSD);
WorkerInfo worker = makeWorker("host1", Collections.singletonMap("/mnt/ssd1", disk));
Map<WorkerInfo, Integer> workerDiskIndex = new HashMap<>();
workerDiskIndex.put(worker, 0);
StorageInfo result =
- SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker),
- 0,
- null,
- workerDiskIndex,
- StorageInfo.LOCAL_DISK_MASK);
+ SlotsAllocator.buildStorageInfo(worker, null, workerDiskIndex, StorageInfo.LOCAL_DISK_MASK);
assertEquals(StorageInfo.Type.SSD, result.getType());
assertEquals("/mnt/ssd1", result.getMountPoint());
@@ -384,16 +333,15 @@ public class BuildStorageInfoSuiteJ {
}
/**
- * When restrictions are {@code null} and only S3_MASK is set, the method must return a
- * StorageInfo with empty mount point and S3 type without touching any worker disks.
+ * When budgets are {@code null} and only S3_MASK is set, the method must return a StorageInfo
+ * with empty mount point and S3 type without touching any worker disks.
*/
@Test
- public void testWithoutRestrictions_S3Only() {
+ public void testWithoutBudgets_S3Only() {
WorkerInfo worker = makeWorker("host1", new HashMap<>());
StorageInfo result =
- SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker), 0, null, new HashMap<>(), StorageInfo.S3_MASK);
+ SlotsAllocator.buildStorageInfo(worker, null, new HashMap<>(), StorageInfo.S3_MASK);
assertEquals(StorageInfo.Type.S3, result.getType());
assertEquals("", result.getMountPoint());
@@ -401,16 +349,15 @@ public class BuildStorageInfoSuiteJ {
}
/**
- * When restrictions are {@code null} and only OSS_MASK is set, the method must return a
- * StorageInfo with empty mount point and OSS type.
+ * When budgets are {@code null} and only OSS_MASK is set, the method must return a StorageInfo
+ * with empty mount point and OSS type.
*/
@Test
- public void testWithoutRestrictions_OSSOnly() {
+ public void testWithoutBudgets_OSSOnly() {
WorkerInfo worker = makeWorker("host1", new HashMap<>());
StorageInfo result =
- SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker), 0, null, new HashMap<>(), StorageInfo.OSS_MASK);
+ SlotsAllocator.buildStorageInfo(worker, null, new HashMap<>(), StorageInfo.OSS_MASK);
assertEquals(StorageInfo.Type.OSS, result.getType());
assertEquals("", result.getMountPoint());
@@ -418,16 +365,15 @@ public class BuildStorageInfoSuiteJ {
}
/**
- * When restrictions are {@code null} and only HDFS_MASK is set, the method must return a
- * StorageInfo with empty mount point and HDFS type.
+ * When budgets are {@code null} and only HDFS_MASK is set, the method must return a StorageInfo
+ * with empty mount point and HDFS type.
*/
@Test
- public void testWithoutRestrictions_HDFSOnly() {
+ public void testWithoutBudgets_HDFSOnly() {
WorkerInfo worker = makeWorker("host1", new HashMap<>());
StorageInfo result =
- SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker), 0, null, new HashMap<>(), StorageInfo.HDFS_MASK);
+ SlotsAllocator.buildStorageInfo(worker, null, new HashMap<>(), StorageInfo.HDFS_MASK);
assertEquals(StorageInfo.Type.HDFS, result.getType());
assertEquals("", result.getMountPoint());
@@ -435,16 +381,15 @@ public class BuildStorageInfoSuiteJ {
}
/**
- * When restrictions are {@code null} and only MEMORY_MASK is set, the method must return a
- * StorageInfo with empty mount point and MEMORY type.
+ * When budgets are {@code null} and only MEMORY_MASK is set, the method must return a StorageInfo
+ * with empty mount point and MEMORY type.
*/
@Test
- public void testWithoutRestrictions_memoryOnly() {
+ public void testWithoutBudgets_memoryOnly() {
WorkerInfo worker = makeWorker("host1", new HashMap<>());
StorageInfo result =
- SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker), 0, null, new HashMap<>(), StorageInfo.MEMORY_MASK);
+ SlotsAllocator.buildStorageInfo(worker, null, new HashMap<>(), StorageInfo.MEMORY_MASK);
assertEquals(StorageInfo.Type.MEMORY, result.getType());
assertEquals("", result.getMountPoint());
@@ -452,24 +397,23 @@ public class BuildStorageInfoSuiteJ {
}
/**
- * When restrictions are {@code null} and the bitmask does not correspond to any known storage
- * type, the method must throw {@link IllegalStateException}.
+ * When budgets are {@code null} and the bitmask does not correspond to any known storage type,
+ * the method must throw {@link IllegalStateException}.
*/
@Test(expected = IllegalStateException.class)
- public void testWithoutRestrictions_noValidStorageType_throwsIllegalState() {
+ public void testWithoutBudgets_noValidStorageType_throwsIllegalState() {
WorkerInfo worker = makeWorker("host1", new HashMap<>());
// 0b100000 = 32 has none of the bits used by known storage types
int unknownMask = 0b100000;
- SlotsAllocator.buildStorageInfo(
- Collections.singletonList(worker), 0, null, new HashMap<>(), unknownMask);
+ SlotsAllocator.buildStorageInfo(worker, null, new HashMap<>(), unknownMask);
}
/**
- * When restrictions are {@code null} and a local disk is selected, the disk index in {@code
+ * When budgets are {@code null} and a local disk is selected, the disk index in {@code
* workerDiskIndex} must advance so that the next call picks the following disk (round-robin).
*/
@Test
- public void testWithoutRestrictions_localDiskAdvancesDiskIndex() {
+ public void testWithoutBudgets_localDiskAdvancesDiskIndex() {
DiskInfo disk1 = makeDiskInfo("/mnt/disk1", StorageInfo.Type.HDD);
DiskInfo disk2 = makeDiskInfo("/mnt/disk2", StorageInfo.Type.HDD);
@@ -481,19 +425,17 @@ public class BuildStorageInfoSuiteJ {
Map<WorkerInfo, Integer> workerDiskIndex = new HashMap<>();
workerDiskIndex.put(worker, 0); // force first disk
- List<WorkerInfo> workers = Collections.singletonList(worker);
-
// First call selects disk1 and advances the index to 1
StorageInfo result1 =
SlotsAllocator.buildStorageInfo(
- workers, 0, null, workerDiskIndex, StorageInfo.ALL_TYPES_AVAILABLE_MASK);
+ worker, null, workerDiskIndex, StorageInfo.ALL_TYPES_AVAILABLE_MASK);
assertEquals("/mnt/disk1", result1.getMountPoint());
assertEquals(Integer.valueOf(1), workerDiskIndex.get(worker));
// Second call selects disk2 and wraps the index back to 0
StorageInfo result2 =
SlotsAllocator.buildStorageInfo(
- workers, 0, null, workerDiskIndex, StorageInfo.ALL_TYPES_AVAILABLE_MASK);
+ worker, null, workerDiskIndex, StorageInfo.ALL_TYPES_AVAILABLE_MASK);
assertEquals("/mnt/disk2", result2.getMountPoint());
assertEquals(Integer.valueOf(0), workerDiskIndex.get(worker));
}
diff --git a/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/LoadAwareSlotsAssignStrategySuiteJ.java b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/LoadAwareSlotsAssignStrategySuiteJ.java
new file mode 100644
index 000000000..718bde558
--- /dev/null
+++ b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/LoadAwareSlotsAssignStrategySuiteJ.java
@@ -0,0 +1,134 @@
+/*
+ * 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.service.deploy.master.slotsalloc;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import org.junit.Test;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.meta.DiskInfo;
+import org.apache.celeborn.common.meta.WorkerInfo;
+import org.apache.celeborn.common.protocol.StorageInfo;
+
+public class LoadAwareSlotsAssignStrategySuiteJ {
+
+ @Test
+ public void testProviderCreatesIndependentStrategiesFromConfiguration() {
+ List<WorkerInfo> workers = createWorkers(3);
+ List<Integer> partitionIds = createPartitionIds(70);
+
+ SlotsAssignStrategy evenStrategy = createStrategy(0);
+ Map<WorkerInfo, List<UsableDiskInfo>> evenBudgets =
+ evenStrategy.computeSlotBudgets(workers, partitionIds, false, StorageInfo.LOCAL_DISK_MASK);
+ assertEquals(24, allocatedSlots(evenBudgets, workers.get(0)));
+ assertEquals(24, allocatedSlots(evenBudgets, workers.get(1)));
+ assertEquals(22, allocatedSlots(evenBudgets, workers.get(2)));
+
+ SlotsAssignStrategy weightedStrategy = createStrategy(1);
+ Map<WorkerInfo, List<UsableDiskInfo>> weightedBudgets =
+ weightedStrategy.computeSlotBudgets(
+ workers, partitionIds, false, StorageInfo.LOCAL_DISK_MASK);
+ assertEquals(40, allocatedSlots(weightedBudgets, workers.get(0)));
+ assertEquals(20, allocatedSlots(weightedBudgets, workers.get(1)));
+ assertEquals(10, allocatedSlots(weightedBudgets, workers.get(2)));
+ }
+
+ @Test
+ public void testProviderRejectsInvalidConfiguration() {
+ assertInvalidConfiguration(
+ CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_DISKGROUP_NUM().key(), "0");
+ assertInvalidConfiguration(
+ CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_DISKGROUP_NUM().key(), "-1");
+
+ assertInvalidDoubleConfiguration(
+ CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_DISKGROUP_GRADIENT().key());
+ assertInvalidDoubleConfiguration(
+ CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_FLUSHTIME_WEIGHT().key());
+ assertInvalidDoubleConfiguration(
+ CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_FETCHTIME_WEIGHT().key());
+ assertInvalidDoubleConfiguration(
+ CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_ACTIVE_SLOTS_WEIGHT().key());
+ }
+
+ private static SlotsAssignStrategy createStrategy(double diskGroupGradient) {
+ CelebornConf conf = validLoadAwareConf();
+ conf.set(
+ CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_DISKGROUP_GRADIENT().key(),
+ String.valueOf(diskGroupGradient));
+ return new LoadAwareSlotsAssignStrategyProvider().create(conf);
+ }
+
+ private static CelebornConf validLoadAwareConf() {
+ CelebornConf conf = new CelebornConf();
+ conf.set(CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_DISKGROUP_NUM().key(), "3");
+ conf.set(CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_FLUSHTIME_WEIGHT().key(), "1");
+ conf.set(CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_FETCHTIME_WEIGHT().key(), "10");
+ conf.set(CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_ACTIVE_SLOTS_WEIGHT().key(), "100");
+ return conf;
+ }
+
+ private static void assertInvalidDoubleConfiguration(String key) {
+ assertInvalidConfiguration(key, "-1");
+ assertInvalidConfiguration(key, "NaN");
+ assertInvalidConfiguration(key, "Infinity");
+ }
+
+ private static void assertInvalidConfiguration(String key, String value) {
+ CelebornConf conf = validLoadAwareConf().set(key, value);
+
+ IllegalArgumentException error =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> new LoadAwareSlotsAssignStrategyProvider().create(conf));
+
+ assertTrue(error.getMessage().contains(key));
+ }
+
+ private static List<WorkerInfo> createWorkers(int workerCount) {
+ List<WorkerInfo> workers = new ArrayList<>();
+ for (int i = 0; i < workerCount; i++) {
+ String mountPoint = "/mnt/disk" + i;
+ DiskInfo disk =
+ new DiskInfo(mountPoint, 1024, i == 0 ? 1 : 0, i == 1 ? 1 : 0, i == 2 ? 1 : 0);
+ disk.availableSlots_$eq(100);
+ Map<String, DiskInfo> disks = new HashMap<>();
+ disks.put(mountPoint, disk);
+ workers.add(new WorkerInfo("host" + i, i, i, i, i, i, disks, null));
+ }
+ return workers;
+ }
+
+ private static List<Integer> createPartitionIds(int partitionCount) {
+ return IntStream.range(0, partitionCount).boxed().collect(Collectors.toList());
+ }
+
+ private static long allocatedSlots(
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets, WorkerInfo worker) {
+ return budgets.get(worker).stream().mapToLong(info -> info.usableSlots).sum();
+ }
+}
diff --git a/master/src/test/java/org/apache/celeborn/service/deploy/master/SlotsAllocatorJmhBenchmark.java b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocatorJmhBenchmark.java
similarity index 89%
rename from master/src/test/java/org/apache/celeborn/service/deploy/master/SlotsAllocatorJmhBenchmark.java
rename to master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocatorJmhBenchmark.java
index e972a399c..79ff9aa38 100644
--- a/master/src/test/java/org/apache/celeborn/service/deploy/master/SlotsAllocatorJmhBenchmark.java
+++ b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocatorJmhBenchmark.java
@@ -15,9 +15,9 @@
* limitations under the License.
*/
-package org.apache.celeborn.service.deploy.master;
+package org.apache.celeborn.service.deploy.master.slotsalloc;
-import static org.apache.celeborn.service.deploy.master.SlotsAllocatorSuiteJ.*;
+import static org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAllocatorSuiteJ.*;
import static org.openjdk.jmh.annotations.Mode.*;
import java.util.Collections;
@@ -44,6 +44,8 @@ public class SlotsAllocatorJmhBenchmark {
private static final boolean HAS_DISKS = true;
private static final int NUM_NETWORK_LOCATIONS = 20;
private static final int NUM_PARTITIONS = 100000;
+ private static final SlotsAssignStrategy SLOTS_ASSIGN_STRATEGY =
+ new RoundRobinSlotsAssignStrategy();
@State(Scope.Thread)
public static class BenchmarkState {
@@ -78,14 +80,15 @@ public class SlotsAllocatorJmhBenchmark {
public void benchmarkSlotSelection(Blackhole blackhole, BenchmarkState state) {
blackhole.consume(
- SlotsAllocator.offerSlotsRoundRobin(
+ SlotsAllocator.offerSlots(
state.workers,
state.partitionIds,
true,
true,
StorageInfo.ALL_TYPES_AVAILABLE_MASK,
false,
- 0));
+ 0,
+ SLOTS_ASSIGN_STRATEGY));
}
public static void main(String[] args) throws Exception {
diff --git a/master/src/test/java/org/apache/celeborn/service/deploy/master/SlotsAllocatorRackAwareSuiteJ.java b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocatorRackAwareSuiteJ.java
similarity index 94%
rename from master/src/test/java/org/apache/celeborn/service/deploy/master/SlotsAllocatorRackAwareSuiteJ.java
rename to master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocatorRackAwareSuiteJ.java
index 0e7e34ced..ef13ef162 100644
--- a/master/src/test/java/org/apache/celeborn/service/deploy/master/SlotsAllocatorRackAwareSuiteJ.java
+++ b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocatorRackAwareSuiteJ.java
@@ -15,7 +15,7 @@
* limitations under the License.
*/
-package org.apache.celeborn.service.deploy.master;
+package org.apache.celeborn.service.deploy.master.slotsalloc;
import static org.apache.hadoop.fs.CommonConfigurationKeysPublic.NET_TOPOLOGY_NODE_SWITCH_MAPPING_IMPL_KEY;
import static org.apache.hadoop.fs.CommonConfigurationKeysPublic.NET_TOPOLOGY_TABLE_MAPPING_FILE_KEY;
@@ -51,6 +51,8 @@ public class SlotsAllocatorRackAwareSuiteJ {
private static final int NUM_ATTEMPTS =
Integer.getInteger("SlotsAllocatorRackAwareSuiteJ.NUM_ATTEMPTS", 100);
+ private static final SlotsAssignStrategy SLOTS_ASSIGN_STRATEGY =
+ new RoundRobinSlotsAssignStrategy();
@Test
public void offerSlotsRoundRobinWithRackAware() throws IOException {
@@ -80,8 +82,15 @@ public class SlotsAllocatorRackAwareSuiteJ {
List<WorkerInfo> workers = prepareWorkers(resolver);
Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots =
- SlotsAllocator.offerSlotsRoundRobin(
- workers, partitionIds, true, true, StorageInfo.ALL_TYPES_AVAILABLE_MASK, false, 0);
+ SlotsAllocator.offerSlots(
+ workers,
+ partitionIds,
+ true,
+ true,
+ StorageInfo.ALL_TYPES_AVAILABLE_MASK,
+ false,
+ 0,
+ SLOTS_ASSIGN_STRATEGY);
Consumer<PartitionLocation> assertCustomer =
new Consumer<PartitionLocation>() {
@@ -120,8 +129,15 @@ public class SlotsAllocatorRackAwareSuiteJ {
List<WorkerInfo> workers = prepareWorkers(resolver);
Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots =
- SlotsAllocator.offerSlotsRoundRobin(
- workers, partitionIds, true, true, StorageInfo.ALL_TYPES_AVAILABLE_MASK, false, 0);
+ SlotsAllocator.offerSlots(
+ workers,
+ partitionIds,
+ true,
+ true,
+ StorageInfo.ALL_TYPES_AVAILABLE_MASK,
+ false,
+ 0,
+ SLOTS_ASSIGN_STRATEGY);
Consumer<PartitionLocation> assertConsumer =
new Consumer<PartitionLocation>() {
@@ -207,8 +223,15 @@ public class SlotsAllocatorRackAwareSuiteJ {
IntStream.range(0, numPartitions).boxed().collect(Collectors.toList());
Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots =
- SlotsAllocator.offerSlotsRoundRobin(
- workers, partitionIds, true, true, StorageInfo.ALL_TYPES_AVAILABLE_MASK, false, 0);
+ SlotsAllocator.offerSlots(
+ workers,
+ partitionIds,
+ true,
+ true,
+ StorageInfo.ALL_TYPES_AVAILABLE_MASK,
+ false,
+ 0,
+ SLOTS_ASSIGN_STRATEGY);
Map<String, Long> numReplicaPerHost =
slots.entrySet().stream()
@@ -245,8 +268,15 @@ public class SlotsAllocatorRackAwareSuiteJ {
List<WorkerInfo> workers = test.generateWorkers();
Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots =
- SlotsAllocator.offerSlotsRoundRobin(
- workers, partitionIds, true, true, StorageInfo.ALL_TYPES_AVAILABLE_MASK, false, 0);
+ SlotsAllocator.offerSlots(
+ workers,
+ partitionIds,
+ true,
+ true,
+ StorageInfo.ALL_TYPES_AVAILABLE_MASK,
+ false,
+ 0,
+ SLOTS_ASSIGN_STRATEGY);
Map<String, Long> numReplicaPerHost =
slots.entrySet().stream()
diff --git a/master/src/test/java/org/apache/celeborn/service/deploy/master/SlotsAllocatorSuiteJ.java b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocatorSuiteJ.java
similarity index 93%
rename from master/src/test/java/org/apache/celeborn/service/deploy/master/SlotsAllocatorSuiteJ.java
rename to master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocatorSuiteJ.java
index f7c70d26e..ecbd04331 100644
--- a/master/src/test/java/org/apache/celeborn/service/deploy/master/SlotsAllocatorSuiteJ.java
+++ b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocatorSuiteJ.java
@@ -15,7 +15,7 @@
* limitations under the License.
*/
-package org.apache.celeborn.service.deploy.master;
+package org.apache.celeborn.service.deploy.master.slotsalloc;
import static org.junit.Assert.*;
@@ -161,19 +161,15 @@ public class SlotsAllocatorSuiteJ {
lightlyReservedDisk.activeSlots_$eq(0);
final Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots =
- SlotsAllocator.offerSlotsLoadAware(
+ SlotsAllocator.offerSlots(
workers,
Collections.singletonList(0),
false,
false,
- 2,
- 1,
- 0,
- 0,
- 1,
StorageInfo.ALL_TYPES_AVAILABLE_MASK,
false,
- 0);
+ 0,
+ new LoadAwareSlotsAssignStrategy(2, 1, 0, 0, 1));
assertTrue(slots.containsKey(workers.get(1)));
assertFalse(slots.containsKey(workers.get(0)));
@@ -199,33 +195,25 @@ public class SlotsAllocatorSuiteJ {
CelebornConf conf = new CelebornConf();
conf.set(CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_DISKGROUP_NUM().key(), "2");
conf.set(CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_DISKGROUP_GRADIENT().key(), "1");
- Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots;
- if (roundrobin) {
- slots =
- SlotsAllocator.offerSlotsRoundRobin(
- workers,
- partitionIds,
- shouldReplicate,
- false,
- StorageInfo.ALL_TYPES_AVAILABLE_MASK,
- interruptionAware,
- interruptionAwareThreshold);
- } else {
- slots =
- SlotsAllocator.offerSlotsLoadAware(
- workers,
- partitionIds,
- shouldReplicate,
- false,
- conf.masterSlotAssignLoadAwareDiskGroupNum(),
- conf.masterSlotAssignLoadAwareDiskGroupGradient(),
- conf.masterSlotAssignLoadAwareFlushTimeWeight(),
- conf.masterSlotAssignLoadAwareFetchTimeWeight(),
- conf.masterSlotAssignLoadAwareActiveSlotsWeight(),
- StorageInfo.ALL_TYPES_AVAILABLE_MASK,
- interruptionAware,
- interruptionAwareThreshold);
- }
+ SlotsAssignStrategy strategy =
+ roundrobin
+ ? new RoundRobinSlotsAssignStrategy()
+ : new LoadAwareSlotsAssignStrategy(
+ conf.masterSlotAssignLoadAwareDiskGroupNum(),
+ conf.masterSlotAssignLoadAwareDiskGroupGradient(),
+ conf.masterSlotAssignLoadAwareFlushTimeWeight(),
+ conf.masterSlotAssignLoadAwareFetchTimeWeight(),
+ conf.masterSlotAssignLoadAwareActiveSlotsWeight());
+ Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots =
+ SlotsAllocator.offerSlots(
+ workers,
+ partitionIds,
+ shouldReplicate,
+ false,
+ StorageInfo.ALL_TYPES_AVAILABLE_MASK,
+ interruptionAware,
+ interruptionAwareThreshold,
+ strategy);
if (expectSuccess) {
if (shouldReplicate) {
slots.forEach(
@@ -318,27 +306,22 @@ public class SlotsAllocatorSuiteJ {
conf.set("celeborn.active.storage.levels", "HDFS");
availableStorageTypes = StorageInfo.HDFS_MASK;
}
- Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots;
- if (roundRobin) {
- slots =
- SlotsAllocator.offerSlotsRoundRobin(
- workers, partitionIds, shouldReplicate, false, availableStorageTypes, false, 0);
- } else {
- slots =
- SlotsAllocator.offerSlotsLoadAware(
- workers,
- partitionIds,
- shouldReplicate,
- false,
- 3,
- 0.1,
- 0,
- 1,
- 0,
- StorageInfo.LOCAL_DISK_MASK | availableStorageTypes,
- false,
- 0);
- }
+ SlotsAssignStrategy strategy =
+ roundRobin
+ ? new RoundRobinSlotsAssignStrategy()
+ : new LoadAwareSlotsAssignStrategy(3, 0.1, 0, 1, 0);
+ int requestedStorageTypes =
+ roundRobin ? availableStorageTypes : StorageInfo.LOCAL_DISK_MASK | availableStorageTypes;
+ Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots =
+ SlotsAllocator.offerSlots(
+ workers,
+ partitionIds,
+ shouldReplicate,
+ false,
+ requestedStorageTypes,
+ false,
+ 0,
+ strategy);
int allocatedPartitionCount = 0;
for (Map.Entry<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>>
workerToPartitions : slots.entrySet()) {
@@ -568,14 +551,15 @@ public class SlotsAllocatorSuiteJ {
IntStream.range(0, 150).boxed().collect(Collectors.toList());
Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>>
slotsFromBestCasePartitionIds =
- SlotsAllocator.offerSlotsRoundRobin(
+ SlotsAllocator.offerSlots(
workers,
bestCasePartitionIds,
shouldReplicate,
shouldRackAware,
StorageInfo.ALL_TYPES_AVAILABLE_MASK,
true,
- (int) interruptionAwarePercentileThreshold);
+ (int) interruptionAwarePercentileThreshold,
+ new RoundRobinSlotsAssignStrategy());
slotsFromBestCasePartitionIds
.values()
.forEach(
@@ -623,14 +607,15 @@ public class SlotsAllocatorSuiteJ {
}
Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>>
slotsFromSpillOverCasePartitionIds =
- SlotsAllocator.offerSlotsRoundRobin(
+ SlotsAllocator.offerSlots(
workers,
spillOverCasePartitionIds,
shouldReplicate,
shouldRackAware,
StorageInfo.ALL_TYPES_AVAILABLE_MASK,
true,
- (int) interruptionAwarePercentileThreshold);
+ (int) interruptionAwarePercentileThreshold,
+ new RoundRobinSlotsAssignStrategy());
slotsFromSpillOverCasePartitionIds
.values()
.forEach(
@@ -663,7 +648,7 @@ public class SlotsAllocatorSuiteJ {
}
}
});
- // With the slot restrictions in place for LoadAware, we expect to spill replicas into
+ // With the load-aware slot budgets in place, we expect to spill replicas into
// workersWithEarlyInterruptionsHosts.
// But primaries should be in workersWithoutInterruptions + workersWithLateInterruptions.
Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>>
diff --git a/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAssignStrategyManagerSuiteJ.java b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAssignStrategyManagerSuiteJ.java
new file mode 100644
index 000000000..238545c11
--- /dev/null
+++ b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAssignStrategyManagerSuiteJ.java
@@ -0,0 +1,382 @@
+/*
+ * 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.service.deploy.master.slotsalloc;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.meta.DiskInfo;
+import org.apache.celeborn.common.meta.WorkerInfo;
+import org.apache.celeborn.common.protocol.StorageInfo;
+import org.apache.celeborn.server.common.service.config.ConfigService;
+import org.apache.celeborn.server.common.service.config.SystemConfig;
+import org.apache.celeborn.service.deploy.master.slotsalloc.test.NullSlotsAssignStrategyProvider;
+import org.apache.celeborn.service.deploy.master.slotsalloc.test.ThirdPartySlotsAssignStrategyProvider;
+import org.apache.celeborn.service.deploy.master.slotsalloc.test.TransientSlotsAssignStrategyProvider;
+
+public class SlotsAssignStrategyManagerSuiteJ {
+
+ @Test
+ public void testRoundRobinProviderIsDiscovered() {
+ SlotsAssignStrategyManager manager = new SlotsAssignStrategyManager(conf("roundrobin"), null);
+
+ assertTrue(manager.getStrategy() instanceof RoundRobinSlotsAssignStrategy);
+ }
+
+ @Test
+ public void testLoadAwareProviderIsDiscovered() {
+ SlotsAssignStrategyManager manager = new SlotsAssignStrategyManager(conf("loadaware"), null);
+
+ assertTrue(manager.getStrategy() instanceof LoadAwareSlotsAssignStrategy);
+ }
+
+ @Test
+ public void testThirdPartyProviderLoadsCustomConfigAndCreatesStrategy() {
+ SlotsAssignStrategyManager manager = new SlotsAssignStrategyManager(thirdPartyConf("7"), null);
+
+ assertEquals(7L, slotBudget(manager.getStrategy()));
+ }
+
+ @Test
+ public void testThirdPartyProviderValidatesCustomConfig() {
+ SlotsAssignStrategyProvider provider = new ThirdPartySlotsAssignStrategyProvider();
+
+ IllegalArgumentException error =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> provider.create(conf(ThirdPartySlotsAssignStrategyProvider.NAME)));
+
+ assertTrue(
+ error.getMessage().contains(ThirdPartySlotsAssignStrategyProvider.MAX_SLOTS_PER_DISK_KEY));
+ }
+
+ @Test
+ public void testUnknownProviderFailsWithAvailableNames() {
+ IllegalArgumentException error =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> new SlotsAssignStrategyManager(conf("missing"), null));
+
+ assertTrue(error.getMessage().contains("MISSING"));
+ assertTrue(error.getMessage().contains("ROUNDROBIN"));
+ assertTrue(error.getMessage().contains("LOADAWARE"));
+ assertTrue(error.getMessage().contains(ThirdPartySlotsAssignStrategyProvider.NAME));
+ }
+
+ @Test
+ public void testDuplicateProviderNamesFailFast() {
+ SlotsAssignStrategyProvider first = new FixedProvider("duplicate");
+ SlotsAssignStrategyProvider second = new FixedProvider("DUPLICATE");
+
+ IllegalStateException error =
+ assertThrows(
+ IllegalStateException.class,
+ () -> SlotsAssignStrategyManager.loadProviders(Arrays.asList(first, second)));
+
+ assertTrue(error.getMessage().contains("Multiple slots assignment strategy providers"));
+ }
+
+ @Test
+ public void testNullStrategyProviderFailsFast() {
+ NullPointerException error =
+ assertThrows(
+ NullPointerException.class,
+ () -> new SlotsAssignStrategyManager(conf(NullSlotsAssignStrategyProvider.NAME), null));
+
+ assertTrue(error.getMessage().contains(NullSlotsAssignStrategyProvider.NAME));
+ }
+
+ @Test
+ public void testThirdPartyStrategyReloadsCustomConfigFromCelebornConf() {
+ CelebornConf staticConf = thirdPartyConf("7");
+ SystemConfig dynamicConfig = new SystemConfig(staticConf);
+ dynamicConfig.setConfigs(Collections.emptyMap());
+ ConfigService configService = mock(ConfigService.class);
+ when(configService.getSystemConfigFromCache()).thenReturn(dynamicConfig);
+
+ SlotsAssignStrategyManager manager = new SlotsAssignStrategyManager(staticConf, configService);
+ ArgumentCaptor<Runnable> listener = ArgumentCaptor.forClass(Runnable.class);
+ verify(configService).registerListenerOnConfigUpdate(listener.capture());
+
+ assertEquals(7L, slotBudget(manager.getStrategy()));
+
+ dynamicConfig.setConfigs(
+ Collections.singletonMap(
+ ThirdPartySlotsAssignStrategyProvider.MAX_SLOTS_PER_DISK_KEY, "3"));
+ listener.getValue().run();
+
+ assertEquals(3L, slotBudget(manager.getStrategy()));
+ }
+
+ @Test
+ public void testDynamicConfigSwitchesProvider() {
+ CelebornConf staticConf = conf("roundrobin");
+ SystemConfig dynamicConfig = new SystemConfig(staticConf);
+ dynamicConfig.setConfigs(Collections.emptyMap());
+ ConfigService configService = mock(ConfigService.class);
+ when(configService.getSystemConfigFromCache()).thenReturn(dynamicConfig);
+
+ SlotsAssignStrategyManager manager = new SlotsAssignStrategyManager(staticConf, configService);
+ ArgumentCaptor<Runnable> listener = ArgumentCaptor.forClass(Runnable.class);
+ verify(configService).registerListenerOnConfigUpdate(listener.capture());
+ assertTrue(manager.getStrategy() instanceof RoundRobinSlotsAssignStrategy);
+
+ Map<String, String> updatedConfigs = new HashMap<>();
+ updatedConfigs.put(
+ CelebornConf.MASTER_SLOT_ASSIGN_POLICY().key(), ThirdPartySlotsAssignStrategyProvider.NAME);
+ updatedConfigs.put(ThirdPartySlotsAssignStrategyProvider.MAX_SLOTS_PER_DISK_KEY, "3");
+ dynamicConfig.setConfigs(updatedConfigs);
+ listener.getValue().run();
+
+ assertEquals(3L, slotBudget(manager.getStrategy()));
+ }
+
+ @Test
+ public void testInvalidInitialDynamicProviderFallsBackToStaticStrategy() {
+ CelebornConf staticConf = conf("roundrobin");
+ SystemConfig dynamicConfig = new SystemConfig(staticConf);
+ dynamicConfig.setConfigs(
+ Collections.singletonMap(CelebornConf.MASTER_SLOT_ASSIGN_POLICY().key(), "missing"));
+ ConfigService configService = mock(ConfigService.class);
+ when(configService.getSystemConfigFromCache()).thenReturn(dynamicConfig);
+
+ SlotsAssignStrategyManager manager = new SlotsAssignStrategyManager(staticConf, configService);
+ ArgumentCaptor<Runnable> listener = ArgumentCaptor.forClass(Runnable.class);
+ verify(configService).registerListenerOnConfigUpdate(listener.capture());
+ assertTrue(manager.getStrategy() instanceof RoundRobinSlotsAssignStrategy);
+
+ dynamicConfig.setConfigs(
+ Collections.singletonMap(CelebornConf.MASTER_SLOT_ASSIGN_POLICY().key(), "loadaware"));
+ listener.getValue().run();
+
+ assertTrue(manager.getStrategy() instanceof LoadAwareSlotsAssignStrategy);
+ }
+
+ @Test
+ public void testInvalidInitialDynamicParameterFallsBackToStaticStrategy() {
+ CelebornConf staticConf = conf("roundrobin");
+ SystemConfig dynamicConfig = new SystemConfig(staticConf);
+ Map<String, String> invalidConfigs = new HashMap<>();
+ invalidConfigs.put(CelebornConf.MASTER_SLOT_ASSIGN_POLICY().key(), "loadaware");
+ invalidConfigs.put(CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_DISKGROUP_NUM().key(), "0");
+ dynamicConfig.setConfigs(invalidConfigs);
+ ConfigService configService = mock(ConfigService.class);
+ when(configService.getSystemConfigFromCache()).thenReturn(dynamicConfig);
+
+ SlotsAssignStrategyManager manager = new SlotsAssignStrategyManager(staticConf, configService);
+ ArgumentCaptor<Runnable> listener = ArgumentCaptor.forClass(Runnable.class);
+ verify(configService).registerListenerOnConfigUpdate(listener.capture());
+ assertTrue(manager.getStrategy() instanceof RoundRobinSlotsAssignStrategy);
+
+ Map<String, String> validConfigs = new HashMap<>();
+ validConfigs.put(CelebornConf.MASTER_SLOT_ASSIGN_POLICY().key(), "loadaware");
+ validConfigs.put(CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_DISKGROUP_NUM().key(), "3");
+ dynamicConfig.setConfigs(validConfigs);
+ listener.getValue().run();
+
+ assertTrue(manager.getStrategy() instanceof LoadAwareSlotsAssignStrategy);
+ }
+
+ @Test
+ public void testUnknownDynamicProviderKeepsPreviousStrategy() {
+ CelebornConf staticConf = conf("roundrobin");
+ SystemConfig dynamicConfig = new SystemConfig(staticConf);
+ dynamicConfig.setConfigs(Collections.emptyMap());
+ ConfigService configService = mock(ConfigService.class);
+ when(configService.getSystemConfigFromCache()).thenReturn(dynamicConfig);
+
+ SlotsAssignStrategyManager manager = new SlotsAssignStrategyManager(staticConf, configService);
+ ArgumentCaptor<Runnable> listener = ArgumentCaptor.forClass(Runnable.class);
+ verify(configService).registerListenerOnConfigUpdate(listener.capture());
+ SlotsAssignStrategy previousStrategy = manager.getStrategy();
+
+ dynamicConfig.setConfigs(
+ Collections.singletonMap(CelebornConf.MASTER_SLOT_ASSIGN_POLICY().key(), "missing"));
+ listener.getValue().run();
+
+ assertSame(previousStrategy, manager.getStrategy());
+ }
+
+ @Test
+ public void testNullDynamicStrategyKeepsPreviousStrategy() {
+ CelebornConf staticConf = conf("roundrobin");
+ SystemConfig dynamicConfig = new SystemConfig(staticConf);
+ dynamicConfig.setConfigs(Collections.emptyMap());
+ ConfigService configService = mock(ConfigService.class);
+ when(configService.getSystemConfigFromCache()).thenReturn(dynamicConfig);
+
+ SlotsAssignStrategyManager manager = new SlotsAssignStrategyManager(staticConf, configService);
+ ArgumentCaptor<Runnable> listener = ArgumentCaptor.forClass(Runnable.class);
+ verify(configService).registerListenerOnConfigUpdate(listener.capture());
+ SlotsAssignStrategy previousStrategy = manager.getStrategy();
+
+ dynamicConfig.setConfigs(
+ Collections.singletonMap(
+ CelebornConf.MASTER_SLOT_ASSIGN_POLICY().key(), NullSlotsAssignStrategyProvider.NAME));
+ listener.getValue().run();
+
+ assertSame(previousStrategy, manager.getStrategy());
+ }
+
+ @Test
+ public void testFailedDynamicStrategyIsRetriedOnlyAfterConfigChange() {
+ CelebornConf staticConf = conf("roundrobin");
+ SystemConfig dynamicConfig = new SystemConfig(staticConf);
+ dynamicConfig.setConfigs(Collections.emptyMap());
+ ConfigService configService = mock(ConfigService.class);
+ when(configService.getSystemConfigFromCache()).thenReturn(dynamicConfig);
+
+ SlotsAssignStrategyManager manager = new SlotsAssignStrategyManager(staticConf, configService);
+ ArgumentCaptor<Runnable> listener = ArgumentCaptor.forClass(Runnable.class);
+ verify(configService).registerListenerOnConfigUpdate(listener.capture());
+ SlotsAssignStrategy previousStrategy = manager.getStrategy();
+
+ dynamicConfig.setConfigs(
+ Collections.singletonMap(
+ CelebornConf.MASTER_SLOT_ASSIGN_POLICY().key(),
+ TransientSlotsAssignStrategyProvider.NAME));
+ listener.getValue().run();
+ assertSame(previousStrategy, manager.getStrategy());
+
+ listener.getValue().run();
+ assertSame(previousStrategy, manager.getStrategy());
+
+ Map<String, String> changedConfigs = new HashMap<>();
+ changedConfigs.put(
+ CelebornConf.MASTER_SLOT_ASSIGN_POLICY().key(), TransientSlotsAssignStrategyProvider.NAME);
+ changedConfigs.put(TransientSlotsAssignStrategyProvider.REVISION_KEY, "1");
+ dynamicConfig.setConfigs(changedConfigs);
+ listener.getValue().run();
+ assertNotSame(previousStrategy, manager.getStrategy());
+ assertTrue(manager.getStrategy() instanceof RoundRobinSlotsAssignStrategy);
+ }
+
+ @Test
+ public void testInvalidDynamicConfigKeepsPreviousStrategy() {
+ CelebornConf staticConf = thirdPartyConf("7");
+ SystemConfig dynamicConfig = new SystemConfig(staticConf);
+ dynamicConfig.setConfigs(
+ Collections.singletonMap(
+ ThirdPartySlotsAssignStrategyProvider.MAX_SLOTS_PER_DISK_KEY, "3"));
+ ConfigService configService = mock(ConfigService.class);
+ when(configService.getSystemConfigFromCache()).thenReturn(dynamicConfig);
+
+ SlotsAssignStrategyManager manager = new SlotsAssignStrategyManager(staticConf, configService);
+ ArgumentCaptor<Runnable> listener = ArgumentCaptor.forClass(Runnable.class);
+ verify(configService).registerListenerOnConfigUpdate(listener.capture());
+ SlotsAssignStrategy previousStrategy = manager.getStrategy();
+
+ dynamicConfig.setConfigs(
+ Collections.singletonMap(
+ ThirdPartySlotsAssignStrategyProvider.MAX_SLOTS_PER_DISK_KEY, "0"));
+ listener.getValue().run();
+
+ assertSame(previousStrategy, manager.getStrategy());
+ assertEquals(3L, slotBudget(manager.getStrategy()));
+ }
+
+ @Test
+ public void testInvalidDynamicLoadAwareConfigKeepsPreviousStrategy() {
+ CelebornConf staticConf = conf("loadaware");
+ SystemConfig dynamicConfig = new SystemConfig(staticConf);
+ dynamicConfig.setConfigs(Collections.emptyMap());
+ ConfigService configService = mock(ConfigService.class);
+ when(configService.getSystemConfigFromCache()).thenReturn(dynamicConfig);
+
+ SlotsAssignStrategyManager manager = new SlotsAssignStrategyManager(staticConf, configService);
+ ArgumentCaptor<Runnable> listener = ArgumentCaptor.forClass(Runnable.class);
+ verify(configService).registerListenerOnConfigUpdate(listener.capture());
+ SlotsAssignStrategy previousStrategy = manager.getStrategy();
+
+ dynamicConfig.setConfigs(
+ Collections.singletonMap(
+ CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_DISKGROUP_NUM().key(), "0"));
+ listener.getValue().run();
+ assertSame(previousStrategy, manager.getStrategy());
+
+ dynamicConfig.setConfigs(
+ Collections.singletonMap(
+ CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_DISKGROUP_NUM().key(), "3"));
+ listener.getValue().run();
+ assertNotSame(previousStrategy, manager.getStrategy());
+ assertTrue(manager.getStrategy() instanceof LoadAwareSlotsAssignStrategy);
+ }
+
+ private static CelebornConf thirdPartyConf(String maxSlotsPerDisk) {
+ return conf(ThirdPartySlotsAssignStrategyProvider.NAME)
+ .set(ThirdPartySlotsAssignStrategyProvider.MAX_SLOTS_PER_DISK_KEY, maxSlotsPerDisk);
+ }
+
+ private static CelebornConf conf(String policyName) {
+ return new CelebornConf().set(CelebornConf.MASTER_SLOT_ASSIGN_POLICY().key(), policyName);
+ }
+
+ private static long slotBudget(SlotsAssignStrategy strategy) {
+ WorkerInfo worker = workerWithAvailableSlots(20);
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets =
+ strategy.computeSlotBudgets(
+ Collections.singletonList(worker),
+ Collections.singletonList(0),
+ false,
+ StorageInfo.LOCAL_DISK_MASK);
+ return budgets.get(worker).get(0).usableSlots;
+ }
+
+ private static WorkerInfo workerWithAvailableSlots(int availableSlots) {
+ String mountPoint = "/mnt/third-party";
+ DiskInfo disk = new DiskInfo(mountPoint, 1024, 0, 0, 0);
+ disk.availableSlots_$eq(availableSlots);
+ Map<String, DiskInfo> disks = new HashMap<>();
+ disks.put(mountPoint, disk);
+ return new WorkerInfo("third-party-worker", 1, 2, 3, 4, 5, disks, null);
+ }
+
+ private static final class FixedProvider implements SlotsAssignStrategyProvider {
+ private final String name;
+
+ private FixedProvider(String name) {
+ this.name = name;
+ }
+
+ @Override
+ public String getName() {
+ return name;
+ }
+
+ @Override
+ public SlotsAssignStrategy create(CelebornConf conf) {
+ return new RoundRobinSlotsAssignStrategy();
+ }
+ }
+}
diff --git a/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/test/NullSlotsAssignStrategyProvider.java b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/test/NullSlotsAssignStrategyProvider.java
new file mode 100644
index 000000000..ddac0ea91
--- /dev/null
+++ b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/test/NullSlotsAssignStrategyProvider.java
@@ -0,0 +1,38 @@
+/*
+ * 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.service.deploy.master.slotsalloc.test;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAssignStrategy;
+import org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAssignStrategyProvider;
+
+/** Test provider that violates the SPI contract by returning a null strategy. */
+public final class NullSlotsAssignStrategyProvider implements SlotsAssignStrategyProvider {
+
+ public static final String NAME = "TEST_NULL";
+
+ @Override
+ public String getName() {
+ return NAME;
+ }
+
+ @Override
+ public SlotsAssignStrategy create(CelebornConf conf) {
+ return null;
+ }
+}
diff --git a/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/test/ThirdPartySlotsAssignStrategyProvider.java b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/test/ThirdPartySlotsAssignStrategyProvider.java
new file mode 100644
index 000000000..1e38694da
--- /dev/null
+++ b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/test/ThirdPartySlotsAssignStrategyProvider.java
@@ -0,0 +1,88 @@
+/*
+ * 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.service.deploy.master.slotsalloc.test;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.meta.DiskInfo;
+import org.apache.celeborn.common.meta.DiskStatus;
+import org.apache.celeborn.common.meta.WorkerInfo;
+import org.apache.celeborn.common.protocol.StorageInfo;
+import org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAssignStrategy;
+import org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAssignStrategyProvider;
+import org.apache.celeborn.service.deploy.master.slotsalloc.UsableDiskInfo;
+
+/** Simulates a third-party provider packaged outside Celeborn's slots-allocation package. */
+public final class ThirdPartySlotsAssignStrategyProvider implements SlotsAssignStrategyProvider {
+
+ public static final String NAME = "THIRD_PARTY_LIMITED";
+ public static final String MAX_SLOTS_PER_DISK_KEY =
+ "celeborn.master.slot.assign.thirdParty.maxSlotsPerDisk";
+
+ @Override
+ public String getName() {
+ return NAME;
+ }
+
+ @Override
+ public SlotsAssignStrategy create(CelebornConf conf) {
+ if (!conf.contains(MAX_SLOTS_PER_DISK_KEY)) {
+ throw new IllegalArgumentException(
+ "Missing third-party configuration: " + MAX_SLOTS_PER_DISK_KEY);
+ }
+
+ long maxSlotsPerDisk = conf.getLong(MAX_SLOTS_PER_DISK_KEY, -1L);
+ if (maxSlotsPerDisk <= 0) {
+ throw new IllegalArgumentException(MAX_SLOTS_PER_DISK_KEY + " must be positive");
+ }
+ return new LimitedSlotsAssignStrategy(maxSlotsPerDisk);
+ }
+
+ private static final class LimitedSlotsAssignStrategy implements SlotsAssignStrategy {
+ private final long maxSlotsPerDisk;
+
+ private LimitedSlotsAssignStrategy(long maxSlotsPerDisk) {
+ this.maxSlotsPerDisk = maxSlotsPerDisk;
+ }
+
+ @Override
+ public Map<WorkerInfo, List<UsableDiskInfo>> computeSlotBudgets(
+ List<WorkerInfo> workers,
+ List<Integer> partitionIds,
+ boolean shouldReplicate,
+ int availableStorageTypes) {
+ Map<WorkerInfo, List<UsableDiskInfo>> budgets = new HashMap<>();
+ for (WorkerInfo worker : workers) {
+ for (DiskInfo disk : worker.diskInfos().values()) {
+ if (DiskStatus.HEALTHY.equals(disk.status())
+ && StorageInfo.isAvailable(disk.storageType(), availableStorageTypes)
+ && disk.getAvailableSlots() > 0) {
+ budgets
+ .computeIfAbsent(worker, ignored -> new ArrayList<>())
+ .add(new UsableDiskInfo(disk, Math.min(maxSlotsPerDisk, disk.getAvailableSlots())));
+ }
+ }
+ }
+ return budgets;
+ }
+ }
+}
diff --git a/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/test/TransientSlotsAssignStrategyProvider.java b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/test/TransientSlotsAssignStrategyProvider.java
new file mode 100644
index 000000000..1480ecd34
--- /dev/null
+++ b/master/src/test/java/org/apache/celeborn/service/deploy/master/slotsalloc/test/TransientSlotsAssignStrategyProvider.java
@@ -0,0 +1,46 @@
+/*
+ * 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.service.deploy.master.slotsalloc.test;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.service.deploy.master.slotsalloc.RoundRobinSlotsAssignStrategy;
+import org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAssignStrategy;
+import org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAssignStrategyProvider;
+
+/** Test provider that fails its first creation attempt and succeeds after configuration changes. */
+public final class TransientSlotsAssignStrategyProvider implements SlotsAssignStrategyProvider {
+
+ public static final String NAME = "TEST_TRANSIENT";
+ public static final String REVISION_KEY = "celeborn.master.slot.assign.test.revision";
+
+ private boolean firstCreation = true;
+
+ @Override
+ public String getName() {
+ return NAME;
+ }
+
+ @Override
+ public SlotsAssignStrategy create(CelebornConf conf) {
+ if (firstCreation) {
+ firstCreation = false;
+ throw new IllegalStateException("Transient strategy creation failure");
+ }
+ return new RoundRobinSlotsAssignStrategy();
+ }
+}
diff --git a/master/src/test/resources/META-INF/services/org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAssignStrategyProvider b/master/src/test/resources/META-INF/services/org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAssignStrategyProvider
new file mode 100644
index 000000000..361af1bf0
--- /dev/null
+++ b/master/src/test/resources/META-INF/services/org.apache.celeborn.service.deploy.master.slotsalloc.SlotsAssignStrategyProvider
@@ -0,0 +1,20 @@
+#
+# 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.
+#
+
+org.apache.celeborn.service.deploy.master.slotsalloc.test.ThirdPartySlotsAssignStrategyProvider
+org.apache.celeborn.service.deploy.master.slotsalloc.test.NullSlotsAssignStrategyProvider
+org.apache.celeborn.service.deploy.master.slotsalloc.test.TransientSlotsAssignStrategyProvider