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 71a7d0afa [CELEBORN-2257] Add reporting of remote disks during
registration
71a7d0afa is described below
commit 71a7d0afa21b807f360a9c5e3d21e45258a6c441
Author: Filip Darmanovic <[email protected]>
AuthorDate: Thu May 14 09:56:37 2026 +0800
[CELEBORN-2257] Add reporting of remote disks during registration
### What changes were proposed in this pull request?
1. Disks reported to the master on registration now include remote disks
(HDFS, S3, OSS)
2. Refactored method names to clarify difference between local and remote
disks.
3. Embedded disk type information into the enum.
4. Refactored unnecessarily complicated code in the slot assignment and
worker registration path.
### Why are the changes needed?
1. Before the first heartbeat, the master won't be able to assign slots
from the remote disks on the worker.
2. All other changes are in preparation for better support of remote disks.
### Does this PR resolve a correctness bug?
Not a correctness bug
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
**Important**: I want help from the community on how to write tests for
this.
Closes #3597 from Dzeri96/CELEBORN-2257.
Authored-by: Filip Darmanovic <[email protected]>
Signed-off-by: SteNicholas <[email protected]>
---
.../celeborn/common/protocol/StorageInfo.java | 31 +++++++++---
.../apache/celeborn/common/meta/WorkerInfo.scala | 16 +++---
.../service/deploy/master/SlotsAllocator.java | 59 +++++++---------------
.../tests/spark/CelebornHashCheckDiskSuite.scala | 4 +-
.../service/deploy/worker/Controller.scala | 2 +-
.../celeborn/service/deploy/worker/Worker.scala | 17 +++----
.../deploy/worker/storage/StorageManager.scala | 32 +++++++-----
.../service/deploy/MiniClusterFeature.scala | 2 +-
.../service/deploy/worker/WorkerSuite.scala | 56 ++++++++++++++++++--
.../worker/storage/StorageManagerSuite.scala | 2 +-
10 files changed, 131 insertions(+), 90 deletions(-)
diff --git
a/common/src/main/java/org/apache/celeborn/common/protocol/StorageInfo.java
b/common/src/main/java/org/apache/celeborn/common/protocol/StorageInfo.java
index 1ab97309e..c204907c2 100644
--- a/common/src/main/java/org/apache/celeborn/common/protocol/StorageInfo.java
+++ b/common/src/main/java/org/apache/celeborn/common/protocol/StorageInfo.java
@@ -22,22 +22,34 @@ import java.util.*;
public class StorageInfo implements Serializable {
public enum Type {
- MEMORY(0),
- HDD(1),
- SSD(2),
- HDFS(3),
- OSS(4),
- S3(5);
+ MEMORY(0, false, MEMORY_MASK),
+ HDD(1, false, LOCAL_DISK_MASK),
+ SSD(2, false, LOCAL_DISK_MASK),
+ HDFS(3, true, HDFS_MASK),
+ OSS(4, true, OSS_MASK),
+ S3(5, true, S3_MASK);
private final int value;
+ private final boolean isDFS;
+ private final int mask;
- Type(int value) {
+ Type(int value, boolean isDFS, int mask) {
this.value = value;
+ this.isDFS = isDFS;
+ this.mask = mask;
}
public int getValue() {
return value;
}
+
+ public boolean isDFS() {
+ return isDFS;
+ }
+
+ public int getMask() {
+ return mask;
+ }
}
public static final Map<Integer, Type> typesMap = new HashMap<>();
@@ -232,6 +244,11 @@ public class StorageInfo implements Serializable {
return S3Available(availableStorageTypes);
}
+ public static boolean isAvailable(Type type, int availableStorageTypes) {
+ return availableStorageTypes == ALL_TYPES_AVAILABLE_MASK
+ || (availableStorageTypes & type.getMask()) > 0;
+ }
+
@Override
public boolean equals(Object o) {
if (this == o) return true;
diff --git
a/common/src/main/scala/org/apache/celeborn/common/meta/WorkerInfo.scala
b/common/src/main/scala/org/apache/celeborn/common/meta/WorkerInfo.scala
index 10a37cf07..0304bd423 100644
--- a/common/src/main/scala/org/apache/celeborn/common/meta/WorkerInfo.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/meta/WorkerInfo.scala
@@ -214,6 +214,10 @@ class WorkerInfo(
for (newDisk <- newDiskInfos.values().asScala) {
val mountPoint: String = newDisk.mountPoint
val curDisk = diskInfos.get(mountPoint)
+ if (estimatedPartitionSize.nonEmpty && !newDisk.storageType.isDFS) {
+ newDisk.maxSlots = newDisk.totalSpace / estimatedPartitionSize.get
+ newDisk.availableSlots = newDisk.actualUsableSpace /
estimatedPartitionSize.get
+ }
if (curDisk != null) {
curDisk.actualUsableSpace = newDisk.actualUsableSpace
curDisk.totalSpace = newDisk.totalSpace
@@ -221,18 +225,10 @@ class WorkerInfo(
curDisk.activeSlots = newDisk.activeSlots
curDisk.avgFlushTime = newDisk.avgFlushTime
curDisk.avgFetchTime = newDisk.avgFetchTime
- if (estimatedPartitionSize.nonEmpty && curDisk.storageType !=
StorageInfo.Type.HDFS
- && curDisk.storageType != StorageInfo.Type.S3 &&
curDisk.storageType != StorageInfo.Type.OSS) {
- curDisk.maxSlots = curDisk.totalSpace / estimatedPartitionSize.get
- curDisk.availableSlots = curDisk.actualUsableSpace /
estimatedPartitionSize.get
- }
+ curDisk.maxSlots = newDisk.maxSlots
+ curDisk.availableSlots = newDisk.availableSlots
curDisk.setStatus(newDisk.status)
} else {
- if (estimatedPartitionSize.nonEmpty && newDisk.storageType !=
StorageInfo.Type.HDFS
- && newDisk.storageType != StorageInfo.Type.S3 &&
newDisk.storageType != StorageInfo.Type.OSS) {
- newDisk.maxSlots = newDisk.totalSpace / estimatedPartitionSize.get
- newDisk.availableSlots = newDisk.actualUsableSpace /
estimatedPartitionSize.get
- }
diskInfos.put(mountPoint, newDisk)
}
}
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
index 5580cd341..07b817427 100644
---
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
@@ -17,6 +17,8 @@
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;
@@ -40,6 +42,12 @@ public class SlotsAllocator {
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;
@@ -70,31 +78,10 @@ public class SlotsAllocator {
for (WorkerInfo worker : workers) {
List<UsableDiskInfo> usableDisks =
slotsRestrictions.computeIfAbsent(worker, v -> new ArrayList<>());
- for (Map.Entry<String, DiskInfo> diskInfoEntry :
worker.diskInfos().entrySet()) {
- if (diskInfoEntry.getValue().status().equals(DiskStatus.HEALTHY)) {
- if (StorageInfo.localDiskAvailable(availableStorageTypes)
- && diskInfoEntry.getValue().storageType() !=
StorageInfo.Type.HDFS
- && diskInfoEntry.getValue().storageType() != StorageInfo.Type.S3
- && diskInfoEntry.getValue().storageType() !=
StorageInfo.Type.OSS) {
- usableDisks.add(
- new UsableDiskInfo(
- diskInfoEntry.getValue(),
diskInfoEntry.getValue().getAvailableSlots()));
- } else if (StorageInfo.HDFSAvailable(availableStorageTypes)
- && diskInfoEntry.getValue().storageType() ==
StorageInfo.Type.HDFS) {
- usableDisks.add(
- new UsableDiskInfo(
- diskInfoEntry.getValue(),
diskInfoEntry.getValue().getAvailableSlots()));
- } else if (StorageInfo.S3Available(availableStorageTypes)
- && diskInfoEntry.getValue().storageType() ==
StorageInfo.Type.S3) {
- usableDisks.add(
- new UsableDiskInfo(
- diskInfoEntry.getValue(),
diskInfoEntry.getValue().getAvailableSlots()));
- } else if (StorageInfo.OSSAvailable(availableStorageTypes)
- && diskInfoEntry.getValue().storageType() ==
StorageInfo.Type.OSS) {
- usableDisks.add(
- new UsableDiskInfo(
- diskInfoEntry.getValue(),
diskInfoEntry.getValue().availableSlots()));
- }
+ for (DiskInfo diskInfo : worker.diskInfos().values()) {
+ if (DiskStatus.HEALTHY.equals(diskInfo.status())
+ && StorageInfo.isAvailable(diskInfo.storageType(),
availableStorageTypes)) {
+ usableDisks.add(new UsableDiskInfo(diskInfo));
}
}
}
@@ -157,9 +144,7 @@ public class SlotsAllocator {
diskToWorkerMap.put(diskInfo, i);
if (diskInfo.actualUsableSpace() > 0
&& diskInfo.status().equals(DiskStatus.HEALTHY)
- && diskInfo.storageType() != StorageInfo.Type.HDFS
- && diskInfo.storageType() != StorageInfo.Type.S3
- && diskInfo.storageType() != StorageInfo.Type.OSS) {
+ && !diskInfo.storageType().isDFS()) {
usableDisks.add(diskInfo);
}
}));
@@ -225,12 +210,8 @@ public class SlotsAllocator {
}
usableDiskInfos.get(diskIndex).usableSlots--;
DiskInfo selectedDiskInfo = usableDiskInfos.get(diskIndex).diskInfo;
- if (selectedDiskInfo.storageType() == StorageInfo.Type.HDFS) {
- storageInfo = new StorageInfo("", StorageInfo.Type.HDFS,
availableStorageTypes);
- } else if (selectedDiskInfo.storageType() == StorageInfo.Type.S3) {
- storageInfo = new StorageInfo("", StorageInfo.Type.S3,
availableStorageTypes);
- } else if (selectedDiskInfo.storageType() == StorageInfo.Type.OSS) {
- storageInfo = new StorageInfo("", StorageInfo.Type.OSS,
availableStorageTypes);
+ if (selectedDiskInfo.storageType().isDFS()) {
+ storageInfo = new StorageInfo("", selectedDiskInfo.storageType(),
availableStorageTypes);
} else {
storageInfo =
new StorageInfo(
@@ -243,9 +224,7 @@ public class SlotsAllocator {
if (StorageInfo.localDiskAvailable(availableStorageTypes)) {
DiskInfo[] diskInfos =
selectedWorker.diskInfos().values().stream()
- .filter(p -> p.storageType() != StorageInfo.Type.HDFS)
- .filter(p -> p.storageType() != StorageInfo.Type.S3)
- .filter(p -> p.storageType() != StorageInfo.Type.OSS)
+ .filter(p -> !p.storageType().isDFS())
.collect(Collectors.toList())
.toArray(new DiskInfo[0]);
int diskIndex =
@@ -257,11 +236,11 @@ public class SlotsAllocator {
availableStorageTypes);
workerDiskIndex.put(selectedWorker, (diskIndex + 1) %
diskInfos.length);
} else if (StorageInfo.S3Available(availableStorageTypes)) {
- storageInfo = new StorageInfo("", StorageInfo.Type.S3,
availableStorageTypes);
+ storageInfo = new StorageInfo("", S3, availableStorageTypes);
} else if (StorageInfo.OSSAvailable(availableStorageTypes)) {
- storageInfo = new StorageInfo("", StorageInfo.Type.OSS,
availableStorageTypes);
+ storageInfo = new StorageInfo("", OSS, availableStorageTypes);
} else if (StorageInfo.HDFSAvailable(availableStorageTypes)) {
- storageInfo = new StorageInfo("", StorageInfo.Type.HDFS,
availableStorageTypes);
+ storageInfo = new StorageInfo("", HDFS, availableStorageTypes);
} else if (StorageInfo.memoryAvailable(availableStorageTypes)) {
storageInfo = new StorageInfo("", StorageInfo.Type.MEMORY,
availableStorageTypes);
} else {
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornHashCheckDiskSuite.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornHashCheckDiskSuite.scala
index 7ac2ac48c..4f22982bb 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornHashCheckDiskSuite.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornHashCheckDiskSuite.scala
@@ -75,7 +75,7 @@ class CelebornHashCheckDiskSuite extends SparkTestBase {
// shuffle key not expired, diskInfo.actualUsableSpace <= 0, no space
workers.foreach { worker =>
worker.storageManager.updateDiskInfos()
- worker.storageManager.disksSnapshot().foreach { diskInfo =>
+ worker.storageManager.localDisksSnapshot().foreach { diskInfo =>
assert(diskInfo.actualUsableSpace <= 0)
}
}
@@ -89,7 +89,7 @@ class CelebornHashCheckDiskSuite extends SparkTestBase {
assert(t.size() === 0)
}
// after shuffle key expired, diskInfo.actualUsableSpace will equal
capacity=1000
- worker.storageManager.disksSnapshot().foreach { diskInfo =>
+ worker.storageManager.localDisksSnapshot().foreach { diskInfo =>
assert(diskInfo.actualUsableSpace === 1000)
}
}
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Controller.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Controller.scala
index ee959e4d6..565acb441 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Controller.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Controller.scala
@@ -190,7 +190,7 @@ private[deploy] class Controller(
return
}
- if (storageManager.healthyWorkingDirs().size <= 0 &&
remoteStorageDirs.isEmpty) {
+ if (storageManager.healthyLocalWorkingDirs().size <= 0 &&
remoteStorageDirs.isEmpty) {
val msg = "Local storage has no available dirs!"
logError(s"[handleReserveSlots] $msg")
context.reply(ReserveSlotsResponse(StatusCode.NO_AVAILABLE_WORKING_DIR,
msg))
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
index da2cab1c3..1c5bc2018 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
@@ -284,11 +284,10 @@ private[celeborn] class Worker(
storageManager.updateDiskInfos()
storageManager.startDeviceMonitor()
- // WorkerInfo's diskInfos is a reference to storageManager.diskInfos
- val diskInfos = JavaUtils.newConcurrentHashMap[String, DiskInfo]()
- storageManager.disksSnapshot().foreach { diskInfo =>
- diskInfos.put(diskInfo.mountPoint, diskInfo)
- }
+ private val diskInfos = storageManager
+ .allDisksSnapshot()
+ .map { diskInfo => diskInfo.mountPoint -> diskInfo }
+ .toMap.asJava
val workerInfo =
new WorkerInfo(
@@ -515,10 +514,10 @@ private[celeborn] class Worker(
activeShuffleKeys.addAll(partitionLocationInfo.shuffleKeySet)
activeShuffleKeys.addAll(storageManager.shuffleKeySet())
storageManager.updateDiskInfos()
- val diskInfos =
- workerInfo.updateThenGetDiskInfos(storageManager.disksSnapshot().map {
disk =>
- disk.mountPoint -> disk
- }.toMap.asJava).values().asScala.toSeq ++
storageManager.remoteDiskInfos.getOrElse(Set.empty)
+ val currentDiskMap = storageManager.allDisksSnapshot().map { disk =>
+ disk.mountPoint -> disk
+ }.toMap.asJava
+ val diskInfos =
workerInfo.updateThenGetDiskInfos(currentDiskMap).asScala.values.toSeq
workerStatusManager.checkIfNeedTransitionStatus()
val response = masterClient.askSync[HeartbeatFromWorkerResponse](
HeartbeatFromWorker(
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
index 25014470a..cde13f29a 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
@@ -103,19 +103,23 @@ final private[worker] class StorageManager(conf:
CelebornConf, workerSource: Abs
if (diskInfoSet.nonEmpty) Some(diskInfoSet) else None
}
- def disksSnapshot(): List[DiskInfo] = {
+ def localDisksSnapshot(): List[DiskInfo] = {
diskInfos.synchronized {
val disks = new util.ArrayList[DiskInfo](diskInfos.values())
disks.asScala.toList
}
}
- def healthyWorkingDirs(): List[File] =
- disksSnapshot().filter(_.status == DiskStatus.HEALTHY).flatMap(_.dirs)
+ def allDisksSnapshot(): List[DiskInfo] = {
+ localDisksSnapshot() ++ remoteDiskInfos.getOrElse(Nil)
+ }
+
+ def healthyLocalWorkingDirs(): List[File] =
+ localDisksSnapshot().filter(_.status == DiskStatus.HEALTHY).flatMap(_.dirs)
private val diskOperators: ConcurrentHashMap[String, ThreadPoolExecutor] = {
val cleaners = JavaUtils.newConcurrentHashMap[String, ThreadPoolExecutor]()
- disksSnapshot().foreach {
+ localDisksSnapshot().foreach {
diskInfo =>
cleaners.put(
diskInfo.mountPoint,
@@ -127,7 +131,7 @@ final private[worker] class StorageManager(conf:
CelebornConf, workerSource: Abs
}
val tmpDiskInfos = JavaUtils.newConcurrentHashMap[String, DiskInfo]()
- disksSnapshot().foreach { diskInfo =>
+ localDisksSnapshot().foreach { diskInfo =>
tmpDiskInfos.put(diskInfo.mountPoint, diskInfo)
}
private val deviceMonitor =
@@ -142,7 +146,7 @@ final private[worker] class StorageManager(conf:
CelebornConf, workerSource: Abs
_totalLocalFlusherThread: Int) = {
val flushers = JavaUtils.newConcurrentHashMap[String, LocalFlusher]()
var totalThread = 0
- disksSnapshot().foreach { diskInfo =>
+ localDisksSnapshot().foreach { diskInfo =>
if (!flushers.containsKey(diskInfo.mountPoint)) {
val flusher = new LocalFlusher(
workerSource,
@@ -269,7 +273,7 @@ final private[worker] class StorageManager(conf:
CelebornConf, workerSource: Abs
private val counter = new AtomicInteger()
private val counterOperator = new IntUnaryOperator() {
override def applyAsInt(operand: Int): Int = {
- val dirs = healthyWorkingDirs()
+ val dirs = healthyLocalWorkingDirs()
if (dirs.nonEmpty) {
(operand + 1) % dirs.length
} else 0
@@ -489,7 +493,7 @@ final private[worker] class StorageManager(conf:
CelebornConf, workerSource: Abs
userIdentifier: UserIdentifier,
partitionSplitEnabled: Boolean,
isSegmentGranularityVisible: Boolean): PartitionDataWriter = {
- if (healthyWorkingDirs().isEmpty && remoteStorageDirs.isEmpty) {
+ if (healthyLocalWorkingDirs().isEmpty && remoteStorageDirs.isEmpty) {
throw new IOException("No available working dirs!")
}
val partitionDataWriterContext = new PartitionDataWriterContext(
@@ -687,7 +691,7 @@ final private[worker] class StorageManager(conf:
CelebornConf, workerSource: Abs
}
}
val (appId, shuffleId) = Utils.splitShuffleKey(shuffleKey)
- disksSnapshot().filter(diskInfo =>
+ localDisksSnapshot().filter(diskInfo =>
diskInfo.status == DiskStatus.HEALTHY
|| diskInfo.status == DiskStatus.HIGH_DISK_USAGE).foreach {
diskInfo =>
diskInfo.dirs.foreach { dir =>
@@ -751,7 +755,7 @@ final private[worker] class StorageManager(conf:
CelebornConf, workerSource: Abs
TimeUnit.MINUTES)
private def cleanupExpiredAppDirs(expireDuration: Long): Unit = {
- val diskInfoAndAppDirs = disksSnapshot()
+ val diskInfoAndAppDirs = localDisksSnapshot()
.filter(diskInfo =>
diskInfo.status == DiskStatus.HEALTHY
|| diskInfo.status == DiskStatus.HIGH_DISK_USAGE)
@@ -801,7 +805,7 @@ final private[worker] class StorageManager(conf:
CelebornConf, workerSource: Abs
val appIds = shuffleKeySet().asScala.map(key =>
Utils.splitShuffleKey(key)._1)
while (retryTimes < conf.workerCheckFileCleanMaxRetries) {
val localCleaned =
- !disksSnapshot().filter(_.status != DiskStatus.IO_HANG).exists {
diskInfo =>
+ !localDisksSnapshot().filter(_.status != DiskStatus.IO_HANG).exists {
diskInfo =>
diskInfo.dirs.exists {
case workingDir if workingDir.exists() =>
// Don't check appDirs that store information in the fileInfos
@@ -946,7 +950,7 @@ final private[worker] class StorageManager(conf:
CelebornConf, workerSource: Abs
}
def updateDiskInfos(): Unit = this.synchronized {
- disksSnapshot()
+ localDisksSnapshot()
.filter(diskInfo =>
diskInfo.status != DiskStatus.IO_HANG && diskInfo.status !=
DiskStatus.READ_OR_WRITE_FAILURE)
.foreach {
@@ -985,7 +989,7 @@ final private[worker] class StorageManager(conf:
CelebornConf, workerSource: Abs
diskInfo.updateFlushTime()
diskInfo.updateFetchTime()
}
- logInfo(s"Updated diskInfos:\n${disksSnapshot().mkString("\n")}")
+ logInfo(s"Updated diskInfos:\n${localDisksSnapshot().mkString("\n")}")
}
def getFileSystemReportedSpace(mountPoint: String): (Long, Long) = {
@@ -1148,7 +1152,7 @@ final private[worker] class StorageManager(conf:
CelebornConf, workerSource: Abs
logInfo(s"Disk(${diskInfo.mountPoint}) unavailable for
$suggestedMountPoint, return all healthy" +
s" working dirs.")
}
- healthyWorkingDirs()
+ healthyLocalWorkingDirs()
}
if (dirs.isEmpty && hdfsFlusher.isEmpty && s3Flusher.isEmpty &&
ossFlusher.isEmpty) {
throw new IOException(s"No available disks! suggested mountPoint
$suggestedMountPoint")
diff --git
a/worker/src/test/scala/org/apache/celeborn/service/deploy/MiniClusterFeature.scala
b/worker/src/test/scala/org/apache/celeborn/service/deploy/MiniClusterFeature.scala
index 95d69fc12..2fab21e71 100644
---
a/worker/src/test/scala/org/apache/celeborn/service/deploy/MiniClusterFeature.scala
+++
b/worker/src/test/scala/org/apache/celeborn/service/deploy/MiniClusterFeature.scala
@@ -213,12 +213,12 @@ trait MiniClusterFeature extends Logging {
val workers = new Array[Worker](workerNum)
val flagUpdateLock = new ReentrantLock()
val threads = (1 to workerNum).map { i =>
+ val worker = createWorker(workerConf)
val workerThread = new RunnerWrap({
var workerStartRetry = 0
var workerStarted = false
while (!workerStarted) {
try {
- val worker = createWorker(workerConf)
flagUpdateLock.lock()
workers(i - 1) = worker
flagUpdateLock.unlock()
diff --git
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerSuite.scala
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerSuite.scala
index 2e13ef1d6..26a1cb1b6 100644
---
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerSuite.scala
+++
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerSuite.scala
@@ -21,27 +21,30 @@ import java.io.File
import java.nio.file.{Files, Paths}
import java.util
import java.util.{HashSet => JHashSet}
-import java.util.concurrent.ConcurrentHashMap
import scala.collection.JavaConverters._
-import scala.collection.mutable.ArrayBuffer
import org.junit.Assert
+import org.mockito.{ArgumentCaptor, ArgumentMatchers, MockedConstruction,
Mockito}
+import org.mockito.MockedConstruction.MockInitializer
+import org.mockito.Mockito.mockConstruction
import org.mockito.MockitoSugar._
-import org.scalatest.{shortstacks, BeforeAndAfterEach}
+import org.scalatest.BeforeAndAfterEach
import org.scalatest.funsuite.AnyFunSuite
import org.apache.celeborn.common.CelebornConf
+import org.apache.celeborn.common.client.MasterClient
import org.apache.celeborn.common.identity.UserIdentifier
-import org.apache.celeborn.common.protocol.{PartitionLocation,
PartitionSplitMode, PartitionType}
+import org.apache.celeborn.common.protocol._
import
org.apache.celeborn.common.protocol.message.ControlMessages.CommitFilesResponse
import org.apache.celeborn.common.protocol.message.StatusCode
import org.apache.celeborn.common.quota.ResourceConsumption
import org.apache.celeborn.common.rpc.RpcCallContext
import org.apache.celeborn.common.util.{CelebornExitKind, JavaUtils,
ThreadUtils}
+import org.apache.celeborn.service.deploy.MiniClusterFeature
import org.apache.celeborn.service.deploy.worker.storage.PartitionDataWriter
-class WorkerSuite extends AnyFunSuite with BeforeAndAfterEach {
+class WorkerSuite extends AnyFunSuite with BeforeAndAfterEach with
MiniClusterFeature {
private var worker: Worker = _
private val conf = new CelebornConf()
private val workerArgs = new WorkerArguments(Array(), conf)
@@ -303,4 +306,47 @@ class WorkerSuite extends AnyFunSuite with
BeforeAndAfterEach {
assert(shuffleCommitTime.get(shuffleKey).get(epoch2) == null)
assert(epochCommitMap.get(epoch2).response.status == StatusCode.SUCCESS)
}
+
+ test("CELEBORN-2257: Properly reports remote disks on worker registration") {
+ val mockInitializer = {
+ // Old syntax needed for scala 2.11
+ new MockInitializer[MasterClient] {
+ override def prepare(instance: MasterClient, context:
MockedConstruction.Context): Unit = {
+ doReturn(PbRegisterWorkerResponse
+ .newBuilder()
+ .setSuccess(true)
+ .build())
+ .when(instance)
+ .askSync(
+ ArgumentMatchers.any(classOf[PbRegisterWorker]),
+ ArgumentMatchers.eq(classOf[PbRegisterWorkerResponse]))
+ }
+ }
+ }
+ val mockedMasterClient = mockConstruction(classOf[MasterClient],
mockInitializer)
+ val argCaptor = ArgumentCaptor.forClass(classOf[PbRegisterWorker])
+ val workerConf: Map[String, String] = Map(
+ CelebornConf.ACTIVE_STORAGE_TYPES.key -> "HDFS",
+ CelebornConf.HDFS_DIR.key -> "file:///")
+ setupMiniClusterWithRandomPorts(workerNum = 1, workerConf = workerConf);
+
+ try {
+ val createdMocks = mockedMasterClient.constructed();
+ assert(createdMocks.size() == 1)
+ verify(createdMocks.get(0), timeout(5000).atLeast(1))
+ .askSync(argCaptor.capture(),
ArgumentMatchers.eq(classOf[PbRegisterWorkerResponse]))
+ val registrationMessage = argCaptor.getValue;
+
+ Assert.assertEquals(2, registrationMessage.getDisksCount)
+ val maybeS3DiskInfo = registrationMessage
+ .getDisksList.asScala
+ .find(diskInfo => diskInfo.getStorageType ==
StorageInfo.Type.HDFS.getValue)
+ assert(maybeS3DiskInfo.nonEmpty)
+ } catch {
+ case e: Throwable => throw e;
+ } finally {
+ shutdownMiniCluster()
+ mockedMasterClient.close()
+ }
+ }
}
diff --git
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
index d393df083..fb1d3e2d2 100644
---
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
+++
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
@@ -54,7 +54,7 @@ class StorageManagerSuite extends CelebornFunSuite with
MockitoHelper {
diskInfo.setUsableSpace(-1L)
var diskSetSpace = (0L, 0L)
- doReturn(List(diskInfo)).when(spyStorageManager).disksSnapshot()
+ doReturn(List(diskInfo)).when(spyStorageManager).localDisksSnapshot()
doAnswer(diskSetSpace).when(spyStorageManager).getFileSystemReportedSpace(any)
// disk usable 80g, total 80g, worker config 8EB
