diff --git a/python/sglang/srt/ray/engine.py b/python/sglang/srt/ray/engine.py index 9c1fe75cb..94c26436e 100644 --- a/python/sglang/srt/ray/engine.py +++ b/python/sglang/srt/ray/engine.py @@ -41,27 +41,36 @@ class RaySchedulerInitResult(SchedulerInitResult): scheduler_actors: list = dataclasses.field(default_factory=list) -def _get_rank0_node_ip(placement_group) -> str: - """Get the IP address of the node where rank 0 will run. - - Uses a probe task to discover the IP of the placement group's first bundle node. - This is needed because rank 0 starts the TCPStore server for torch.distributed, - so dist_init_addr must be the IP of the node where rank 0 runs, not the driver node. +def _find_engine_bundle(placement_group, nnodes: int) -> tuple[int, str]: + """Find which placement group bundle is on the same node as the Engine. + Rank0 scheduler must be co-located with the Engine. Returns (bundle_index, engine_ip). """ + engine_ip = ray.util.get_node_ip_address() @ray.remote(num_cpus=0, num_gpus=0) def get_node_ip(): return ray.util.get_node_ip_address() - return ray.get( - get_node_ip.options( - scheduling_strategy=PlacementGroupSchedulingStrategy( - placement_group=placement_group, - placement_group_bundle_index=0, - ), - ).remote() + bundle_ips = ray.get( + [ + get_node_ip.options( + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=placement_group, + placement_group_bundle_index=i, + ), + ).remote() + for i in range(nnodes) + ] ) + try: + return bundle_ips.index(engine_ip), engine_ip + except ValueError: + raise RuntimeError( + f"Engine node {engine_ip} not found in any placement group bundle {bundle_ips}. " + f"Rank-0 scheduler must be co-located with the Engine." + ) + class RayEngine(Engine): """Engine using Ray actors for scheduler processes.""" @@ -105,13 +114,20 @@ class RayEngine(Engine): f"Use {gpus_per_node} GPUs/node, world_size={world_size}" ) - rank0_node_ip = _get_rank0_node_ip(pg) + # co-located with the Engine and rank0 scheduler at the same node + engine_bundle, engine_ip = _find_engine_bundle(pg, nnodes) + bundle_for_node = [engine_bundle] + [ + i for i in range(nnodes) if i != engine_bundle + ] + + rank0_node_ip = engine_ip dist_init_addr = f"{rank0_node_ip}:{server_args.port + ZMQ_TCP_PORT_DELTA}" logger.info(f"dist_init_addr: {dist_init_addr}") scheduler_actors = [] for node_idx in range(nnodes): + bundle_idx = bundle_for_node[node_idx] pp_range, tp_range, pp_per_node, tp_per_node = _calculate_rank_ranges( nnodes, server_args.pp_size, server_args.tp_size, node_rank=node_idx ) @@ -131,7 +147,7 @@ class RayEngine(Engine): name=f"sglang_scheduler_rank0node={rank0_node_ip}_pp{pp_rank}_tp{tp_rank}", scheduling_strategy=PlacementGroupSchedulingStrategy( placement_group=pg, - placement_group_bundle_index=node_idx, + placement_group_bundle_index=bundle_idx, ), ).remote( server_args=server_args, diff --git a/python/sglang/srt/ray/scheduler_actor.py b/python/sglang/srt/ray/scheduler_actor.py index e68b78e94..68d9fc654 100644 --- a/python/sglang/srt/ray/scheduler_actor.py +++ b/python/sglang/srt/ray/scheduler_actor.py @@ -105,6 +105,10 @@ class SchedulerActor: def run_event_loop(self) -> None: """Run the scheduler's event loop. Blocks until shutdown.""" try: + import torch + + # Need to set the GPU id for the event loop for nccl to work + torch.cuda.set_device(self.scheduler.gpu_id) self.scheduler.run_event_loop() except Exception as e: logger.error(f"Scheduler PP{self._pp_rank} TP{self._tp_rank} crashed: {e}")