wehub-resource-sync
69 lines
2.7 KiB
Python
69 lines
2.7 KiB
Python
"""Placement group utilities for Ray LLM Serve."""
|
|
|
|
from collections import defaultdict
|
|
from typing import Dict, List, Optional
|
|
|
|
from ray.serve._private.utils import get_head_node_id
|
|
from ray.util.placement_group import PlacementGroup, placement_group_table
|
|
|
|
|
|
def _sort_bundle_indices_by_node(
|
|
bundles_to_node_id: Dict[int, str],
|
|
driver_node_id: str,
|
|
) -> List[int]:
|
|
"""Sort bundle indices so that same-node bundles are adjacent, driver node first.
|
|
|
|
Args:
|
|
bundles_to_node_id: Mapping from bundle index to node ID.
|
|
driver_node_id: The node ID of the driver node.
|
|
|
|
Returns:
|
|
List of bundle indices sorted with driver node bundles first,
|
|
then remaining nodes in deterministic order.
|
|
"""
|
|
node_to_bundles: Dict[str, List[int]] = defaultdict(list)
|
|
# bundle_idx is already in ascending order: created sequentially during
|
|
# placement group creation and preserved by the GCS protobuf.
|
|
for bundle_idx, node_id in bundles_to_node_id.items():
|
|
node_to_bundles[node_id].append(bundle_idx)
|
|
|
|
result: List[int] = []
|
|
if driver_node_id in node_to_bundles:
|
|
result.extend(node_to_bundles.pop(driver_node_id))
|
|
for node_id in sorted(node_to_bundles.keys()):
|
|
result.extend(node_to_bundles[node_id])
|
|
|
|
return result
|
|
|
|
|
|
def get_bundle_indices_sorted_by_node(
|
|
pg: PlacementGroup,
|
|
driver_node_id: Optional[str] = None,
|
|
) -> List[int]:
|
|
"""Return bundle indices sorted such that same-node bundles are adjacent, driver node first.
|
|
|
|
When a placement group is provisioned, adjacent bundle indices don't
|
|
necessarily map to the same physical node. This utility reorders bundle
|
|
indices so that bundles on the same node are grouped together.
|
|
|
|
The driver node's bundles come first so that global rank 0 (which hosts the
|
|
distributed rendezvous store) is co-located with the driver.
|
|
|
|
Args:
|
|
pg: A ready placement group.
|
|
driver_node_id: Node ID whose bundles should be ordered first. Callers
|
|
should pass the node that advertises the distributed master address
|
|
so that global rank 0 is co-located with it. Defaults to the cluster
|
|
head node, which may not own any of the placement group's bundles
|
|
(e.g. a CPU-only head node in a GPU cluster), in which case rank 0
|
|
falls back to the lexicographically-first node.
|
|
|
|
Returns:
|
|
List of bundle indices sorted such that same-node bundles are adjacent, driver node first.
|
|
"""
|
|
table = placement_group_table(pg)
|
|
bundles_to_node_id = table["bundles_to_node_id"]
|
|
if driver_node_id is None:
|
|
driver_node_id = get_head_node_id()
|
|
return _sort_bundle_indices_by_node(bundles_to_node_id, driver_node_id)
|