ray-project--ray/python/ray/util/placement_group.py

import warnings
from typing import Dict, List, Optional, Union

import ray
from ray._common.utils import PLACEMENT_GROUP_BUNDLE_RESOURCE_NAME, hex_to_binary
from ray._private.auto_init_hook import auto_init_ray
from ray._private.client_mode_hook import client_mode_should_convert, client_mode_wrap
from ray._private.label_utils import validate_label_selector
from ray._private.utils import get_ray_doc_version
from ray._raylet import PlacementGroupID
from ray.util.annotations import DeveloperAPI, PublicAPI

VALID_PLACEMENT_GROUP_STRATEGIES = {
    "PACK",
    "SPREAD",
    "STRICT_PACK",
    "STRICT_SPREAD",
}

# Keep this in sync with RAY_NODE_ID_KEY in common.pxi. The module cannot import
# that value because it forms a circular dependency.
NODE_ID_LABEL_KEY = "ray.io/node-id"


@PublicAPI
class PlacementGroup:
    """A handle to a placement group."""

    @staticmethod
    def empty() -> "PlacementGroup":
        return PlacementGroup(PlacementGroupID.nil())

    def __init__(
        self,
        id: "ray._raylet.PlacementGroupID",
        bundle_cache: Optional[List[Dict]] = None,
    ):
        self.id = id
        self.bundle_cache = bundle_cache

    @property
    def is_empty(self):
        return self.id.is_nil()

    def ready(self) -> "ray._raylet.ObjectRef":
        """Returns an ObjectRef to check ready status.

        Returns:
            An ``ObjectRef`` that resolves once the placement group has been
            created and all bundles are scheduled.

        Example:
            .. testcode::

                import ray

                pg = ray.util.placement_group([{"CPU": 1}])
                ray.get(pg.ready())

                pg = ray.util.placement_group([{"CPU": 1}])
                ray.wait([pg.ready()])

        """
        if self.is_empty:
            return ray.put(self)

        return _call_placement_group_ready_async(self)

    def wait(self, timeout_seconds: Union[float, int] = 30) -> bool:
        """Wait for the placement group to be ready within the specified time.

        Args:
             timeout_seconds: Timeout in seconds.

        Returns:
             True if the placement group is created. False otherwise.
        """
        return _call_placement_group_ready(self.id, timeout_seconds)

    @property
    def bundle_specs(self) -> List[Dict]:
        """List[Dict]: Return bundles belonging to this placement group."""
        self._fill_bundle_cache_if_needed()
        return self.bundle_cache

    @property
    def bundle_count(self) -> int:
        self._fill_bundle_cache_if_needed()
        return len(self.bundle_cache)

    def _fill_bundle_cache_if_needed(self) -> None:
        if not self.bundle_cache:
            self.bundle_cache = _get_bundle_cache(self.id)

    def __eq__(self, other):
        if not isinstance(other, PlacementGroup):
            return False
        return self.id == other.id

    def __hash__(self):
        return hash(self.id)


@client_mode_wrap
def _call_placement_group_ready_async(pg: PlacementGroup) -> "ray._raylet.ObjectRef":
    worker = ray._private.worker.global_worker
    worker.check_connected()
    # Serialize pg so that ray.get() returns the PlacementGroup
    serialized = worker.get_serialization_context().serialize(pg)
    return worker.core_worker.async_wait_placement_group_ready(pg.id, serialized)


@client_mode_wrap
def _call_placement_group_ready(pg_id: PlacementGroupID, timeout_seconds: int) -> bool:
    worker = ray._private.worker.global_worker
    worker.check_connected()

    return worker.core_worker.wait_placement_group_ready(pg_id, timeout_seconds)


@client_mode_wrap
def _get_bundle_cache(pg_id: PlacementGroupID) -> List[Dict]:
    worker = ray._private.worker.global_worker
    worker.check_connected()

    return list(
        ray._private.state.state.placement_group_table(pg_id)["bundles"].values()
    )


@PublicAPI
@client_mode_wrap
def placement_group(
    bundles: List[Dict[str, float]],
    strategy: Optional[str] = None,
    name: str = "",
    lifetime: Optional[str] = None,
    _soft_target_node_id: Optional[str] = None,
    bundle_label_selector: List[Dict[str, str]] = None,
    topology_strategy: Optional[Dict[str, str]] = None,
) -> PlacementGroup:
    """Asynchronously creates a PlacementGroup.

    Args:
        bundles: A list of bundles which
            represent the resources requirements.
        strategy: The strategy to create the placement group.

         - "PACK": Packs Bundles into as few nodes as possible.
         - "SPREAD": Places Bundles across distinct nodes as even as possible.
         - "STRICT_PACK": Packs Bundles into one node. The group is
           not allowed to span multiple nodes.
         - "STRICT_SPREAD": Packs Bundles across distinct nodes.

        name: The name of the placement group.
        lifetime: Either `None`, which defaults to the placement group
            will fate share with its creator and will be deleted once its
            creator is dead, or "detached", which means the placement group
            will live as a global object independent of the creator.
        _soft_target_node_id: (Private, Experimental) Soft hint where bundles of
            this placement group should be placed.
            The target node is specified by it's hex ID.
            If the target node has no available resources or died,
            bundles can be placed elsewhere.
            This currently only works with STRICT_PACK pg.
        bundle_label_selector: A list of label selectors to apply to a
            placement group on a per-bundle level.
        topology_strategy: Topology-aware placement. A dict mapping each
            topology label key to a placement strategy (e.g.,
            ``{"ray.io/gpu-domain": "STRICT_PACK"}``).
            Mutually exclusive with `strategy`.

    Raises:
        ValueError: if bundle type is not a list.
        ValueError: if empty bundle or empty resource bundles are given.
        ValueError: if the wrong lifetime arguments are given.
        ValueError: if `topology_strategy` and `strategy` are both provided,
            or if invalid `topology_strategy` is passed.

    Returns:
        PlacementGroup: Placement group object.
    """
    worker = ray._private.worker.global_worker
    worker.check_connected()

    validate_placement_group(
        bundles=bundles,
        strategy=strategy,
        lifetime=lifetime,
        _soft_target_node_id=_soft_target_node_id,
        bundle_label_selector=bundle_label_selector,
        topology_strategy=topology_strategy,
    )

    if bundle_label_selector is None:
        bundle_label_selector = []

    node_level_strategy = _derive_node_level_strategy(strategy, topology_strategy)

    # Current implementation derives node level strategy from topology_strategy,
    # while we pass a topology strategy with node level strategy stripped.
    if topology_strategy is not None:
        stripped_topology_strategy = {
            k: v for k, v in topology_strategy.items() if k != NODE_ID_LABEL_KEY
        }
    else:
        stripped_topology_strategy = {}

    if lifetime == "detached":
        detached = True
    else:
        detached = False

    placement_group_id = worker.core_worker.create_placement_group(
        name,
        bundles,
        node_level_strategy,
        detached,
        _soft_target_node_id,
        bundle_label_selector,
        stripped_topology_strategy,
    )

    return PlacementGroup(
        placement_group_id,
        bundle_cache=[{k: float(v) for k, v in bundle.items()} for bundle in bundles],
    )


@PublicAPI
@client_mode_wrap
def remove_placement_group(placement_group: PlacementGroup) -> None:
    """Asynchronously remove placement group.

    Args:
        placement_group: The placement group to delete.
    """
    assert placement_group is not None
    worker = ray._private.worker.global_worker
    worker.check_connected()

    worker.core_worker.remove_placement_group(placement_group.id)


@PublicAPI
@client_mode_wrap
def get_placement_group(placement_group_name: str) -> PlacementGroup:
    """Get a placement group object with a global name.

    Args:
        placement_group_name: Global name of the placement group to look up.

    Returns:
        None if can't find a placement group with the given name.
        The placement group object otherwise.
    """
    if not placement_group_name:
        raise ValueError("Please supply a non-empty value to get_placement_group")
    worker = ray._private.worker.global_worker
    worker.check_connected()
    placement_group_info = ray._private.state.state.get_placement_group_by_name(
        placement_group_name, worker.namespace
    )
    if placement_group_info is None:
        raise ValueError(
            f"Failed to look up placement group with name: {placement_group_name}"
        )
    else:
        return PlacementGroup(
            PlacementGroupID(hex_to_binary(placement_group_info["placement_group_id"]))
        )


@DeveloperAPI
@client_mode_wrap
def placement_group_table(placement_group: PlacementGroup = None) -> dict:
    """Get the state of the placement group from GCS.

    Args:
        placement_group: placement group to see
            states.

    Returns:
        A dictionary describing the state of the given placement group, or
        the table of all placement groups if ``placement_group`` is None.
    """
    worker = ray._private.worker.global_worker
    worker.check_connected()
    placement_group_id = placement_group.id if (placement_group is not None) else None
    return ray._private.state.state.placement_group_table(placement_group_id)


@PublicAPI
def get_current_placement_group() -> Optional[PlacementGroup]:
    """Get the current placement group which a task or actor is using.

    It returns None if there's no current placement group for the worker.
    For example, if you call this method in your driver, it returns None
    (because drivers never belong to any placement group).

    Examples:
        .. testcode::

            import ray
            from ray.util.placement_group import get_current_placement_group
            from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

            @ray.remote
            def f():
                # This returns the placement group the task f belongs to.
                # It means this pg is identical to the pg created below.
                return get_current_placement_group()

            pg = ray.util.placement_group([{"CPU": 2}])
            assert ray.get(f.options(
                    scheduling_strategy=PlacementGroupSchedulingStrategy(
                        placement_group=pg)).remote()) == pg

            # Driver doesn't belong to any placement group,
            # so it returns None.
            assert get_current_placement_group() is None

    Returns:
        PlacementGroup: Placement group object.
            None if the current task or actor wasn't
            created with any placement group.
    """
    auto_init_ray()
    if client_mode_should_convert():
        # Client mode is only a driver.
        return None
    worker = ray._private.worker.global_worker
    worker.check_connected()
    pg_id = worker.placement_group_id
    if pg_id.is_nil():
        return None
    return PlacementGroup(pg_id)


def check_placement_group_index(
    placement_group: PlacementGroup, bundle_index: int
) -> None:
    assert placement_group is not None
    if placement_group.id.is_nil():
        if bundle_index != -1:
            raise ValueError(
                "If placement group is not set, "
                "the value of bundle index must be -1."
            )
    elif bundle_index >= placement_group.bundle_count or bundle_index < -1:
        raise ValueError(
            f"placement group bundle index {bundle_index} "
            f"is invalid. Valid placement group indexes: "
            f"0-{placement_group.bundle_count}"
        )


def _derive_node_level_strategy(
    strategy: Optional[str], topology_strategy: Optional[Dict[str, str]]
) -> str:
    """Assumes valid strategy and topology strategy, and derives the node level
    strategy from the corresponding fields accordingly.
    """
    if topology_strategy is not None:
        return topology_strategy.get(NODE_ID_LABEL_KEY, "PACK")
    return strategy if strategy is not None else "PACK"


def validate_placement_group(
    bundles: List[Dict[str, float]],
    strategy: Optional[str] = None,
    lifetime: Optional[str] = None,
    _soft_target_node_id: Optional[str] = None,
    bundle_label_selector: List[Dict[str, str]] = None,
    topology_strategy: Optional[Dict[str, str]] = None,
) -> bool:
    """Validates inputs for placement_group.

    Raises ValueError if inputs are invalid.
    """
    # Mutual exclusion: `strategy` and `topology_strategy` cannot both be passed.
    if topology_strategy is not None and strategy is not None:
        raise ValueError(
            "`strategy` and `topology_strategy` cannot both be specified. "
            "Pass node-level packing via "
            f"`topology_strategy['{NODE_ID_LABEL_KEY}']`."
        )

    if strategy is not None and strategy not in VALID_PLACEMENT_GROUP_STRATEGIES:
        raise ValueError(
            f"Invalid placement group strategy {strategy}. "
            f"Supported strategies are: {VALID_PLACEMENT_GROUP_STRATEGIES}."
        )

    if topology_strategy is not None:
        _validate_topology_strategy(topology_strategy)

    # Resolve the strategy for the rest of validation that depends on it.
    node_level_strategy = _derive_node_level_strategy(strategy, topology_strategy)

    if _soft_target_node_id and node_level_strategy != "STRICT_PACK":
        raise ValueError(
            "_soft_target_node_id currently only works "
            f"with STRICT_PACK but got {node_level_strategy}"
        )

    if _soft_target_node_id and ray.NodeID.from_hex(_soft_target_node_id).is_nil():
        raise ValueError(
            f"Invalid hex ID of _soft_target_node_id, got {_soft_target_node_id}"
        )

    _validate_bundles(bundles)

    if bundle_label_selector is not None:
        if len(bundles) != len(bundle_label_selector):
            raise ValueError(
                f"Invalid bundle label selector {bundle_label_selector}. "
                f"The length of `bundle_label_selector` should equal the length of `bundles`."
            )
        _validate_bundle_label_selector(bundle_label_selector)

    if lifetime not in [None, "detached"]:
        raise ValueError(
            "Placement group `lifetime` argument must be either `None` or "
            f"'detached'. Got {lifetime}."
        )


def _validate_topology_strategy(topology_strategy: Dict[str, str]) -> None:
    """Validates topology_strategy shape.

    Currently accepts a dict containing "ray.io/node-id" and at most one other
    topology label. The "ray.io/node-id" entry is equivalent to the `strategy=`
    parameter and accepts any value in VALID_PLACEMENT_GROUP_STRATEGIES. The
    other (topology) label is restricted to "STRICT_PACK" for now.
    """
    if not isinstance(topology_strategy, dict):
        raise ValueError(
            "`topology_strategy` must be a dict, "
            f"got {type(topology_strategy).__name__}."
        )

    if not (0 <= len(topology_strategy) <= 2):
        raise ValueError(
            "`topology_strategy` must contain 0, 1, or 2 entries: "
            f"`{NODE_ID_LABEL_KEY}` plus an optional topology label. "
            f"Got {len(topology_strategy)} entries."
        )

    topology_label_keys = []
    for key, value in topology_strategy.items():
        if not isinstance(key, str) or not key:
            raise ValueError(
                "`topology_strategy` keys must be non-empty strings, " f"got {key!r}."
            )
        if value not in VALID_PLACEMENT_GROUP_STRATEGIES:
            raise ValueError(
                f"Invalid topology strategy {value!r} for label {key!r}. "
                f"Supported strategies are: {VALID_PLACEMENT_GROUP_STRATEGIES}."
            )
        # The node-id entry is just a redirect to the existing `strategy=`
        # parameter and accepts any valid placement-group strategy. For other
        # topology labels (e.g. "ray.io/gpu-domain"), currently only supports
        # STRICT_PACK.
        if not key == NODE_ID_LABEL_KEY:
            topology_label_keys.append(key)
            if value != "STRICT_PACK":
                raise ValueError(
                    f"Topology strategy {value!r} for label {key!r} is not "
                    "supported yet; only 'STRICT_PACK' is supported for topology "
                    f"labels other than `{NODE_ID_LABEL_KEY}` for now."
                )

    if len(topology_label_keys) > 1:
        raise ValueError(
            "`topology_strategy` currently supports at most one topology label "
            f"other than `{NODE_ID_LABEL_KEY}`. Got {topology_label_keys}."
        )


def _validate_bundles(bundles: List[Dict[str, float]]):
    """Validates each bundle and raises a ValueError if any bundle is invalid."""

    if not isinstance(bundles, list):
        raise ValueError(
            "Placement group bundles must be a list, " f"got {type(bundles)}."
        )

    if len(bundles) == 0:
        raise ValueError(
            "Bundles must be a non-empty list of resource "
            'dictionaries. For example: `[{"CPU": 1.0}, {"GPU": 1.0}]`. '
            "Got empty list instead."
        )

    for bundle in bundles:
        if (
            not isinstance(bundle, dict)
            or not all(isinstance(k, str) for k in bundle.keys())
            or not all(isinstance(v, (int, float)) for v in bundle.values())
        ):
            raise ValueError(
                "Bundles must be a non-empty list of "
                "resource dictionaries. For example: "
                '`[{"CPU": 1.0}, {"GPU": 1.0}]`.'
            )

        if len(bundle) == 0 or all(
            resource_value == 0 for resource_value in bundle.values()
        ):
            raise ValueError(
                "Bundles cannot be an empty dictionary or "
                f"resources with only 0 values. Bundles: {bundles}"
            )

        if "object_store_memory" in bundle.keys():
            warnings.warn(
                "Setting 'object_store_memory' for"
                " bundles is deprecated since it doesn't actually"
                " reserve the required object store memory."
                f" Use object spilling that's enabled by default (https://docs.ray.io/en/{get_ray_doc_version()}/ray-core/objects/object-spilling.html) "  # noqa: E501
                "instead to bypass the object store memory size limitation.",
                DeprecationWarning,
                stacklevel=1,
            )


def _validate_bundle_label_selector(bundle_label_selector: List[Dict[str, str]]):
    """Validates each label selector and raises a ValueError if any label selector is invalid."""

    if not isinstance(bundle_label_selector, list):
        raise ValueError(
            "Placement group bundle_label_selector must be a list, "
            f"got {type(bundle_label_selector)}."
        )

    if len(bundle_label_selector) == 0:
        # No label selectors provided, no-op.
        return

    for label_selector in bundle_label_selector:
        if (
            not isinstance(label_selector, dict)
            or not all(isinstance(k, str) for k in label_selector.keys())
            or not all(isinstance(v, str) for v in label_selector.values())
        ):
            raise ValueError(
                "Bundle label selector must be a list of string dictionary"
                " label selectors. For example: "
                '`[{ray.io/market_type": "spot"}, {"ray.io/accelerator-type": "A100"}]`.'
            )
        # Call helper function to validate label selector key-value syntax.
        error_message = validate_label_selector(label_selector)
        if error_message:
            raise ValueError(
                f"Invalid label selector provided in bundle_label_selector list."
                f" Detailed error: '{error_message}'"
            )

    gpu_domain_accelerator = None
    for label_selector in bundle_label_selector:
        accel = label_selector.get("ray.io/accelerator-type")
        if accel in {"GB200", "GB300"}:
            gpu_domain_accelerator = accel
            break

    if gpu_domain_accelerator is not None:
        for label_selector in bundle_label_selector:
            if label_selector.get("ray.io/accelerator-type") != gpu_domain_accelerator:
                raise ValueError(
                    f"Invalid bundle label selector {label_selector}. "
                    "GPU-domain scheduling requires all bundles to have "
                    f"'ray.io/accelerator-type: {gpu_domain_accelerator}'"
                    " in their label selector."
                )


def _valid_resource_shape(resources, bundle_specs):
    """
    If the resource shape cannot fit into every
    bundle spec, return False
    """
    for bundle in bundle_specs:
        fit_in_bundle = True
        for resource, requested_val in resources.items():
            # Skip "bundle" resource as it is automatically added
            # to all nodes with bundles by the placement group.
            if resource == PLACEMENT_GROUP_BUNDLE_RESOURCE_NAME:
                continue
            if bundle.get(resource, 0) < requested_val:
                fit_in_bundle = False
                break
        if fit_in_bundle:
            # If resource request fits in any bundle, it is valid.
            return True
    return False


def _validate_resource_shape(
    placement_group, resources, placement_resources, task_or_actor_repr
):
    bundles = placement_group.bundle_specs
    resources_valid = _valid_resource_shape(resources, bundles)
    placement_resources_valid = _valid_resource_shape(placement_resources, bundles)

    if not resources_valid:
        raise ValueError(
            f"Cannot schedule {task_or_actor_repr} with "
            "the placement group because the resource request "
            f"{resources} cannot fit into any bundles for "
            f"the placement group, {bundles}."
        )
    if not placement_resources_valid:
        # Happens for the default actor case.
        # placement_resources is not an exposed concept to users,
        # so we should write more specialized error messages.
        raise ValueError(
            f"Cannot schedule {task_or_actor_repr} with "
            "the placement group because the actor requires "
            f"{placement_resources.get('CPU', 0)} CPU for "
            "creation, but it cannot "
            f"fit into any bundles for the placement group, "
            f"{bundles}. Consider "
            "creating a placement group with CPU resources."
        )


def _configure_placement_group_based_on_context(
    placement_group_capture_child_tasks: bool,
    bundle_index: int,
    resources: Dict,
    placement_resources: Dict,
    task_or_actor_repr: str,
    placement_group: Union[PlacementGroup, str, None] = "default",
) -> PlacementGroup:
    """Configure the placement group based on the given context.

    Based on the given context, this API returns the placement group instance
    for task/actor scheduling.

    Params:
        placement_group_capture_child_tasks: Whether or not the
            placement group needs to be captured from the global
            context.
        bundle_index: The bundle index for tasks/actor scheduling.
        resources: The scheduling resources.
        placement_resources: The scheduling placement resources for
            actors.
        task_or_actor_repr: The repr of task or actor
            function/class descriptor.
        placement_group: The placement group instance.
            - "default": Default placement group argument. Currently,
                the default behavior is to capture the parent task'
                placement group if placement_group_capture_child_tasks
                is set.
            - None: means placement group is explicitly not configured.
            - Placement group instance: In this case, do nothing.

    Returns:
        Placement group instance based on the given context.

    Raises:
        ValueError: If the bundle index is invalid for the placement group
            or the requested resources shape doesn't fit to any
            bundles.
    """
    # Validate inputs.
    assert placement_group_capture_child_tasks is not None
    assert resources is not None

    # Validate and get the PlacementGroup instance.
    # Placement group could be None, default, or placement group.
    # Default behavior is "do not capture child tasks".
    if placement_group != "default":
        if not placement_group:
            placement_group = PlacementGroup.empty()
    elif placement_group == "default":
        if placement_group_capture_child_tasks:
            placement_group = get_current_placement_group()
        else:
            placement_group = PlacementGroup.empty()

    if not placement_group:
        placement_group = PlacementGroup.empty()
    assert isinstance(placement_group, PlacementGroup)

    # Validate the index.
    check_placement_group_index(placement_group, bundle_index)

    # Validate the shape.
    if not placement_group.is_empty:
        _validate_resource_shape(
            placement_group, resources, placement_resources, task_or_actor_repr
        )
    return placement_group