Utils

class ResolvedPlacementGroup:
    """Wrapper around Ray PlacementGroup that resolves physical ordering of bundles and stores reordered bundle indices.

    Ray placement groups don't guarantee bundle ordering (bundles on the same node
    may not have consecutive indices). This wrapper probes the PG once on first access
    and caches the full (bundle_idx, node_id, gpu_id) mapping sorted by (node_id, gpu_id).

    All attributes are lazy and computed on first access.
    Use ``.pg`` to access the underlying Ray PlacementGroup for Ray APIs.

    Attributes:
        reordered_bundle_indices: Raw bundle indices sorted by (node_id, gpu_id).
        bundle_node_ids: Node ID for each reordered bundle index.
        bundle_gpu_ids: Physical GPU ID for each reordered bundle index.
        num_nodes: Number of distinct nodes in the placement group.
        num_gpus_per_node: Number of GPUs per node (assumes uniform distribution).
    """

    def __init__(self, pg: PlacementGroup):
        self.pg = pg
        self._bundle_placement = None

    def _get_bundle_placement(self):
        if self._bundle_placement is None:
            self._bundle_placement = _probe_bundle_placement(self.pg)
        return self._bundle_placement

    @functools.cached_property
    def reordered_bundle_indices(self):
        return [info[0] for info in self._get_bundle_placement()]

    @functools.cached_property
    def bundle_node_ids(self):
        """Node ID for each reordered bundle index."""
        return [info[1] for info in self._get_bundle_placement()]

    @functools.cached_property
    def bundle_gpu_ids(self):
        """Physical GPU ID for each reordered bundle index."""
        return [info[2] for info in self._get_bundle_placement()]

    @functools.cached_property
    def num_nodes(self):
        return len(set(self.bundle_node_ids))

    @functools.cached_property
    def num_gpus_per_node(self):
        return len(self._get_bundle_placement()) // self.num_nodes