diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py index e313ca73..6c621e72 100644 --- a/src/exo/worker/plan.py +++ b/src/exo/worker/plan.py @@ -235,9 +235,8 @@ def _ready_to_warmup( assert device_rank < world_size assert device_rank >= 0 - # TODO: Ensure these align with MLX distributeds expectations. - # Rank < n-1 - accepting_ranks_ready = device_rank < world_size - 1 and all( + # Rank != 0 + accepting_ranks_ready = device_rank > 0 and all( isinstance( all_runners.get(global_runner_id, None), (RunnerLoaded, RunnerWarmingUp), @@ -245,8 +244,8 @@ def _ready_to_warmup( for global_runner_id in shard_assignments.runner_to_shard ) - # Rank = n-1 - connecting_rank_ready = device_rank == world_size - 1 and all( + # Rank = 0 + connecting_rank_ready = device_rank == 0 and all( isinstance(all_runners.get(global_runner_id, None), RunnerWarmingUp) for global_runner_id in shard_assignments.runner_to_shard if global_runner_id != runner_id