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(
+            update_expert_location_with_recovery(
-                expert_location_metadata,
+                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]

python/sglang/srt/model_executor/model_runner.py

@@ -25,7 +25,7 @@ import threading
 import time
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any, List, Optional, Union
+from typing import Any, Optional, Union
 
 import torch
 import torch.distributed as dist
@@ -96,7 +96,6 @@ from sglang.srt.eplb.expert_distribution import (
     set_global_expert_distribution_recorder,
 )
 from sglang.srt.eplb.expert_location import (
-    ExpertLocationMetadata,
     broadcast_global_expert_location_metadata,
     compute_initial_expert_location_metadata,
     format_expert_location_layout,
@@ -1649,52 +1648,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             set_global_lplb_solver(lid, solver)
         logger.info(f"Initialized LPLB solvers for {metadata.num_layers} layers")
 
-    def update_expert_location(
-        self,
-        new_expert_location_metadata: ExpertLocationMetadata,
-        update_layer_ids: List[int],
-    ):
-        p2p_missing_logical_experts = self.expert_location_updater.update(
-            self.model.routed_experts_weights_of_layer,
-            new_expert_location_metadata,
-            update_layer_ids=update_layer_ids,
-            nnodes=self.server_args.nnodes,
-            rank=self.tp_rank,
-        )
-
-        if len(p2p_missing_logical_experts) > 0:
-            # Load the missing expert weights from disk
-            if callable(getattr(self.model, "generate_weight_name_filter", None)):
-                # Filter and load only missing expert weights
-                weight_name_filter = self.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 (
-                self.expert_backup_client is not None
-                and self.expert_backup_client.use_backup
-            ):
-                # Load the missing weights from the DRAM backup
-                self.expert_backup_client.update_weights(weight_name_filter)
-            else:
-                # Load the missing weights from disk
-                self.weight_updater.update_weights_from_disk(
-                    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 self.server_args.ep_dispatch_algorithm == "lp":
-            self._init_lplb_solvers()
-
     def maybe_recover_ep_ranks(self):
         # TODO(perf): `active_ranks.all()` on a CUDA tensor triggers host-device
         # synchronization, and this function is on the forward-path.