fix

Author: sibianl

src/backend/server/rpc_connection_handler.py

@@ -63,6 +63,7 @@ def node_update(self, message):
                 current_requests=node.current_requests,
                 layer_latency_ms=node.layer_latency_ms,
                 new_rtt_to_nodes=node.rtt_to_nodes,
+                is_active=node.is_active,
             )
             return {}
         except Exception as e:

src/scheduling/scheduler.py

@@ -88,7 +88,7 @@ def __init__(
         self._pending_joins: "queue.Queue[Node]" = queue.Queue()
         self._pending_leaves: "queue.Queue[str]" = queue.Queue()
         self._pending_node_updates: (
-            "queue.Queue[Tuple[str, Optional[int], Optional[float], Optional[Dict[str, float]]]]"
+            "queue.Queue[Tuple[str, Optional[int], Optional[float], Optional[Dict[str, float]], Optional[bool]]]"
         ) = queue.Queue()
 
         # Concurrency controls
@@ -192,6 +192,7 @@ def update_node_info(
         current_requests: Optional[int] = None,
         layer_latency_ms: Optional[float] = None,
         new_rtt_to_nodes: Optional[Dict[str, float]] = None,
+        is_active: Optional[bool] = None,
     ) -> None:
         """Update the info of a node."""
         if current_requests is not None:
@@ -200,6 +201,8 @@ def update_node_info(
             node.set_layer_latency_ms(layer_latency_ms)
         if new_rtt_to_nodes is not None:
             node.rtt_to_nodes.update(new_rtt_to_nodes)
+        if is_active is not None:
+            node.is_active = is_active
         node.last_heartbeat = time.time()
         # logger.debug(
         #     "Node updated: %s (requests=%s, latency_ms=%s, rtt_updates=%s)",
@@ -228,19 +231,13 @@ def enqueue_node_update(
         current_requests: Optional[int] = None,
         layer_latency_ms: Optional[float] = None,
         new_rtt_to_nodes: Optional[Dict[str, float]] = None,
+        is_active: Optional[bool] = None,
     ) -> None:
         """Enqueue a node update event."""
         self._pending_node_updates.put(
-            (node_id, current_requests, layer_latency_ms, new_rtt_to_nodes)
+            (node_id, current_requests, layer_latency_ms, new_rtt_to_nodes, is_active)
         )
         self._wake_event.set()
-        # logger.debug(
-        #     "Enqueued node update: %s (requests=%s, latency_ms=%s, rtt_updates=%s)",
-        #     node_id,
-        #     current_requests,
-        #     layer_latency_ms,
-        #     0 if new_rtt_to_nodes is None else len(new_rtt_to_nodes),
-        # )
 
     def checking_node_heartbeat(self) -> None:
         """Check the heartbeat of all nodes."""
@@ -457,14 +454,15 @@ def _process_node_updates(self) -> None:
         """Apply pending node stats updates from the queue."""
         while True:
             try:
-                node_id, cur, lat, rtts = self._pending_node_updates.get_nowait()
+                node_id, cur, lat, rtts, is_active = self._pending_node_updates.get_nowait()
             except queue.Empty:
                 break
             self.update_node_info(
                 self.node_id_to_node[node_id],
                 current_requests=cur,
                 layer_latency_ms=lat,
                 new_rtt_to_nodes=rtts,
+                is_active=is_active,
             )
 
     def _process_joins(self) -> None: