[6/N] (Elastic EP) Recover failed ranks (#15771)

python/sglang/srt/eplb/eplb_manager.py

@@ -41,6 +41,9 @@ class EPLBManager:
     def on_forward_pass_end(self):
         next(self._main_generator)
 
+    def reset_generator(self):
+        self._main_generator = self._entrypoint()
+
     # can be more complex if needed
     def _entrypoint(self):
         while True:

python/sglang/srt/eplb/expert_location.py

@@ -318,6 +318,52 @@ def set_global_expert_location_metadata(value):
     _global_expert_location_metadata = value
 
 
+def broadcast_global_expert_location_metadata(
+    src_rank: int = 0, group: Optional[torch.distributed.ProcessGroup] = None
+):
+    """Broadcast the global ExpertLocationMetadata from src_rank to all ranks.
+
+    This is used in Elastic EP rank recovery to ensure that all ranks (including
+    newly recovered ones) share exactly the same expert location metadata.
+
+    Note: The caller must ensure src_rank is a healthy rank. In recovery scenarios,
+    this function is called after try_recover_ranks succeeds, at which point all
+    ranks (including src_rank=0) have recovered and are ready.
+    """
+    metadata = get_global_expert_location_metadata()
+    assert metadata is not None
+
+    # Ensure device tensors are contiguous before broadcasting in-place
+    metadata.physical_to_logical_map = metadata.physical_to_logical_map.contiguous()
+    metadata.logical_to_all_physical_map = (
+        metadata.logical_to_all_physical_map.contiguous()
+    )
+    metadata.logical_to_all_physical_map_num_valid = (
+        metadata.logical_to_all_physical_map_num_valid.contiguous()
+    )
+    if metadata.logical_to_rank_dispatch_physical_map is not None:
+        metadata.logical_to_rank_dispatch_physical_map = (
+            metadata.logical_to_rank_dispatch_physical_map.contiguous()
+        )
+
+    device_tensors = [
+        metadata.physical_to_logical_map,
+        metadata.logical_to_all_physical_map,
+        metadata.logical_to_all_physical_map_num_valid,
+    ]
+    if metadata.logical_to_rank_dispatch_physical_map is not None:
+        device_tensors.append(metadata.logical_to_rank_dispatch_physical_map)
+
+    for tensor in device_tensors:
+        torch.distributed.broadcast(tensor, src=src_rank, group=group)
+
+    # After broadcasting device tensors, refresh corresponding CPU copies
+    metadata.physical_to_logical_map_cpu = metadata.physical_to_logical_map.cpu()
+    metadata.logical_to_all_physical_map_cpu = (
+        metadata.logical_to_all_physical_map.cpu()
+    )
+
+
 def _compute_logical_to_all_physical_map(
     server_args: ServerArgs,
     physical_to_logical_map: torch.Tensor,