(celeborn) branch main updated: [CELEBORN-1726] Update WorkerInfo when transition worker state

[rexxiong]

This is an automated email from the ASF dual-hosted git repository.

rexxiong 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 bb96a5f31 [CELEBORN-1726] Update WorkerInfo when transition worker 
state
bb96a5f31 is described below

commit bb96a5f31f5655e4329501cf7fbb8ffe6d196e95
Author: Weijie Guo <res...@163.com>
AuthorDate: Wed Nov 20 22:26:10 2024 +0800

    [CELEBORN-1726] Update WorkerInfo when transition worker state
    
    ### What changes were proposed in this pull request?
    Update WorkerInfo when transition worker state
    
    ### Why are the changes needed?
    When we send a `getWorkerInfo` request to the `Worker` node after this 
worker has became `Decommission` state , it does not return the correct state.
    
    note: If we send this request to `Master` node instead of `Worker`, the 
return value is correct. We should update the worker state also for worker node 
itself.
    
    ### Does this PR introduce _any_ user-facing change?
    Yes
    
    ### How was this patch tested?
    Mannually
    
    Closes #2930 from reswqa/worker_info.
    
    Authored-by: Weijie Guo <res...@163.com>
    Signed-off-by: Shuang <lvshuang....@alibaba-inc.com>
---
 .../service/deploy/worker/WorkerStatusManager.scala  |  1 +
 .../deploy/worker/WorkerStatusManagerSuite.scala     | 20 ++++++++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git 
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManager.scala
 
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManager.scala
index 4291c5acd..1957cbfce 100644
--- 
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManager.scala
+++ 
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManager.scala
@@ -129,6 +129,7 @@ private[celeborn] class WorkerStatusManager(conf: 
CelebornConf) extends Logging
     if (allowStates != null && allowStates.contains(state)) {
       logInfo(s"Worker transition status from ${currentWorkerStatus.getState} 
to $state.")
       currentWorkerStatus = new WorkerStatus(state.getNumber, 
System.currentTimeMillis())
+      worker.workerInfo.setWorkerStatus(currentWorkerStatus)
     } else {
       logWarning(
         s"Worker transition status from ${currentWorkerStatus.getState} to 
$state is not allowed.")
diff --git 
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManagerSuite.scala
 
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManagerSuite.scala
index b1e7d1fc6..c787be688 100644
--- 
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManagerSuite.scala
+++ 
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManagerSuite.scala
@@ -19,13 +19,17 @@ package org.apache.celeborn.service.deploy.worker
 
 import java.util.concurrent.atomic.AtomicBoolean
 
-import com.google.common.collect.Sets
+import com.google.common.collect.{Maps, Sets}
 import org.junit.Assert
 import org.mockito.MockitoSugar._
 import org.scalatest.funsuite.AnyFunSuite
 
 import org.apache.celeborn.common.CelebornConf
+import org.apache.celeborn.common.identity.UserIdentifier
+import org.apache.celeborn.common.meta.WorkerInfo
 import org.apache.celeborn.common.protocol.{PbWorkerStatus, WorkerEventType}
+import org.apache.celeborn.common.quota.ResourceConsumption
+import org.apache.celeborn.common.util.JavaUtils
 import org.apache.celeborn.service.deploy.worker.storage.StorageManager
 
 class WorkerStatusManagerSuite extends AnyFunSuite {
@@ -36,15 +40,27 @@ class WorkerStatusManagerSuite extends AnyFunSuite {
     worker = mock[Worker]
     val storageManager = mock[StorageManager]
     val shuffleKeys = Sets.newHashSet("test")
+    val workerInfo = new WorkerInfo(
+      "host",
+      0,
+      0,
+      0,
+      0,
+      0,
+      Maps.newHashMap(),
+      JavaUtils.newConcurrentHashMap[UserIdentifier, ResourceConsumption])
     when(storageManager.shuffleKeySet()).thenReturn(shuffleKeys)
     when(worker.storageManager).thenReturn(storageManager)
     when(worker.shutdown).thenReturn(new AtomicBoolean())
-
+    when(worker.workerInfo).thenReturn(workerInfo)
     val statusManager = new WorkerStatusManager(conf)
     statusManager.init(worker)
 
     statusManager.doTransition(WorkerEventType.DecommissionThenIdle)
     Assert.assertEquals(statusManager.getWorkerState(), 
PbWorkerStatus.State.InDecommissionThenIdle)
+    Assert.assertEquals(
+      worker.workerInfo.getWorkerStatus().getStateValue,
+      PbWorkerStatus.State.InDecommissionThenIdle.getNumber)
 
     // Rerun state Transition
     statusManager.doTransition(WorkerEventType.DecommissionThenIdle)