Extract expert location updating into EPLBManager (#31149)

python/sglang/srt/eplb/eplb_manager.py

@@ -3,6 +3,7 @@ import time
 from typing import TYPE_CHECKING, List
 
 import torch.cuda
+from torch import nn
 
 from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
@@ -12,6 +13,8 @@ from sglang.srt.eplb.expert_location import (
     format_expert_location_layout_diff,
     get_global_expert_location_metadata,
 )
+from sglang.srt.eplb.expert_location_updater import ExpertLocationUpdater
+from sglang.srt.runtime_context import get_server_args
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.model_runner import ModelRunner
@@ -94,9 +97,17 @@ class EPLBManager:
         for chunk_layer_ids in update_layer_ids_chunks:
             if len(update_layer_ids_chunks) > 1:
                 yield
-            self._model_runner.update_expert_location(
-                expert_location_metadata,
+            update_expert_location_with_recovery(
+                expert_location_updater=self._model_runner.expert_location_updater,
+                model=self._model_runner.model,
+                new_expert_location_metadata=expert_location_metadata,
                 update_layer_ids=chunk_layer_ids,
+                nnodes=self._model_runner.server_args.nnodes,
+                tp_rank=self._model_runner.tp_rank,
+                expert_backup_client=self._model_runner.expert_backup_client,
+                update_weights_from_disk_callable=self._model_runner.weight_updater.update_weights_from_disk,
+                ep_dispatch_algorithm=self._model_runner.server_args.ep_dispatch_algorithm,
+                init_lplb_solvers_callable=self._model_runner._init_lplb_solvers,
             )
 
         self._log_rebalance_layout_after_update(update_layer_ids=all_update_layer_ids)
@@ -181,6 +192,58 @@ class EPLBManager:
         )
 
 
+def update_expert_location_with_recovery(
+    *,
+    expert_location_updater: ExpertLocationUpdater,
+    model: nn.Module,
+    new_expert_location_metadata: ExpertLocationMetadata,
+    update_layer_ids: List[int],
+    nnodes: int,
+    tp_rank: int,
+    expert_backup_client,
+    update_weights_from_disk_callable,
+    ep_dispatch_algorithm: str,
+    init_lplb_solvers_callable,
+):
+    p2p_missing_logical_experts = expert_location_updater.update(
+        model.routed_experts_weights_of_layer,
+        new_expert_location_metadata,
+        update_layer_ids=update_layer_ids,
+        nnodes=nnodes,
+        rank=tp_rank,
+    )
+
+    if len(p2p_missing_logical_experts) > 0:
+        # Load the missing expert weights from disk
+        if callable(getattr(model, "generate_weight_name_filter", None)):
+            # Filter and load only missing expert weights
+            weight_name_filter = model.generate_weight_name_filter(
+                p2p_missing_logical_experts
+            )
+        else:
+            # Do a full reload from disk/DRAM
+            logger.info(
+                "[Elastic EP] Model does not implement generate_weight_name_filter. "
+                "Performing full weight reload."
+            )
+            weight_name_filter = None
+
+        if expert_backup_client is not None and expert_backup_client.use_backup:
+            # Load the missing weights from the DRAM backup
+            expert_backup_client.update_weights(weight_name_filter)
+        else:
+            # Load the missing weights from disk
+            update_weights_from_disk_callable(
+                get_server_args().model_path,
+                get_server_args().load_format,
+                weight_name_filter=weight_name_filter,
+            )
+
+    # Re-init LPLB solvers after expert location update
+    if ep_dispatch_algorithm == "lp":
+        init_lplb_solvers_callable()
+
+
 def _chunk_list(items: List, chunk_size):
     for start_index in range(0, len(items), chunk_size):
         yield items[start_index : start_index + chunk_size]