chore: import upstream snapshot with attribution

python/ray/train/v2/_internal/execution/scaling_policy/__init__.py

@@ -0,0 +1,29 @@
+# isort: off
+from .scaling_policy import ScalingDecision, ScalingPolicy, NoopDecision, ResizeDecision
+from .scaling_policy import (
+    AUTOSCALING_REQUESTS_EXPIRE_TIME_S,
+    AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
+    AUTOSCALING_REQUESTS_INTERVAL_S,
+)
+from .elastic import ElasticScalingPolicy
+from .fixed import FixedScalingPolicy
+from .factory import create_scaling_policy
+
+# isort: on
+
+
+__all__ = [
+    "AUTOSCALING_REQUESTS_EXPIRE_TIME_S",
+    "AUTOSCALING_REQUESTS_GET_TIMEOUT_S",
+    "AUTOSCALING_REQUESTS_INTERVAL_S",
+    "ScalingPolicy",
+    "ElasticScalingPolicy",
+    "FixedScalingPolicy",
+    "ScalingDecision",
+    "NoopDecision",
+    "ResizeDecision",
+    "create_scaling_policy",
+]
+
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/train/v2/_internal/execution/scaling_policy/elastic.py

@@ -0,0 +1,291 @@
+import logging
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+import ray
+from ray.train.v2._internal.execution.scaling_policy import (
+    NoopDecision,
+    ResizeDecision,
+    ScalingDecision,
+    ScalingPolicy,
+)
+from ray.train.v2._internal.execution.scaling_policy.scaling_policy import (
+    AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
+)
+from ray.train.v2._internal.execution.worker_group import (
+    WorkerGroupPollStatus,
+    WorkerGroupState,
+)
+from ray.train.v2._internal.util import time_monotonic
+from ray.train.v2.api.config import ScalingConfig
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from ray.data._internal.cluster_autoscaler.default_autoscaling_coordinator import (
+        ResourceDict,
+    )
+
+
+class ElasticScalingPolicy(ScalingPolicy):
+
+    # Minimum interval in seconds between querying the AutoscalingCoordinator for allocated resources.
+    GET_ALLOCATED_RESOURCES_INTERVAL_S = 1
+    # Minimum interval in seconds between logging warnings about insufficient workers.
+    INSUFFICIENT_WORKERS_WARNING_INTERVAL_S = 30
+
+    def __init__(self, scaling_config: ScalingConfig):
+        super().__init__(scaling_config)
+
+        self._latest_monitor_time = float("-inf")
+        self._latest_insufficient_workers_warning_time = float("-inf")
+        self._latest_allocated_resources_query_time = float("-inf")
+        self._latest_allocated_resources: Optional[List["ResourceDict"]] = None
+
+    def _get_num_workers_for_resource_request(self) -> int:
+        return self.scaling_config.max_workers
+
+    def _count_possible_workers(
+        self, allocated_resources: List[Dict[str, float]]
+    ) -> int:
+        """Count the number of workers that can be started/restarted with the given
+        the list of node resources. The returned number is capped at the maximum
+        number of workers.
+
+        For GPUs, this divides raw allocated resources by per-worker requirements.
+        For TPUs, an additional check ensures workers align with physically intact
+        TPU slices (see ``_get_strict_tpu_worker_count``).
+
+        Args:
+            allocated_resources: The resources currently allocated by the AutoscalingCoordinator.
+
+        Returns:
+            The number of workers that can be started/restarted with the current resources.
+        """
+        # TODO: Fractional resources do not work well here.
+        single_worker_resources = self.scaling_config._resources_per_worker_not_none
+        total_num_workers = 0
+
+        # If workers require no resources, we can run as many as we want.
+        if sum(single_worker_resources.values()) == 0:
+            return self.scaling_config.max_workers
+
+        for resources in allocated_resources:
+            num_workers = min(
+                [
+                    resources.get(resource, 0.0) // single_worker_resources[resource]
+                    for resource in single_worker_resources
+                    if single_worker_resources[resource] > 0
+                ]
+            )
+            total_num_workers += num_workers
+
+        total_num_workers = min(int(total_num_workers), self.scaling_config.max_workers)
+
+        # Multi-host TPUs are scheduled atomically in interconnected slices defined by a topology.
+        if (
+            self.scaling_config.use_tpu
+            and self.scaling_config.topology
+            and self.scaling_config.accelerator_type
+        ):
+            total_num_workers = self._get_strict_tpu_worker_count(
+                total_num_workers=total_num_workers,
+            )
+
+        return total_num_workers
+
+    def _get_strict_tpu_worker_count(self, total_num_workers: int) -> int:
+        """Calculate the number of workers that can run on intact TPU slices.
+
+        The Autoscaler's allocated resources might overestimate the number of
+        schedulable TPU workers because it counts raw resources. TPUs require
+        atomic, interconnected slices. This function checks the cluster for
+        physically intact slices to prevent scaling onto fractional/broken
+        topologies.
+
+        The worker count is: min(resource_based_slices, intact_slices) *
+        workers_per_slice, where resource_based_slices =
+        total_num_workers // workers_per_slice.
+
+        Args:
+            total_num_workers: The initial estimate of workers based on raw
+                allocated resources.
+
+        Returns:
+            The number of workers aligned to fully intact TPU slices.
+        """
+        from ray.util.tpu import get_num_tpu_slices, get_tpu_worker_resources
+
+        single_worker_resources = self.scaling_config._resources_per_worker_not_none
+
+        try:
+            workers_per_slice, _ = get_tpu_worker_resources(
+                topology=self.scaling_config.topology,
+                accelerator_type=self.scaling_config.accelerator_type,
+                resources_per_unit=single_worker_resources,
+                num_slices=1,
+            )
+
+            if workers_per_slice == 0:
+                # A single worker requires more resources than exist in a
+                # full slice — impossible scheduling configuration for TPU.
+                return 0
+
+            num_slices_from_resources = total_num_workers // workers_per_slice
+
+            if num_slices_from_resources > 0:
+                try:
+                    num_intact_slices = get_num_tpu_slices(
+                        topology=self.scaling_config.topology,
+                        accelerator_type=self.scaling_config.accelerator_type,
+                    )
+                    num_slices_from_resources = min(
+                        num_slices_from_resources, num_intact_slices
+                    )
+                except Exception as e:
+                    logger.warning(
+                        f"Failed to check cluster state for intact TPU slices: {e}"
+                    )
+
+            return num_slices_from_resources * workers_per_slice
+
+        except Exception as e:
+            logger.warning(
+                f"Could not calculate TPU slice boundaries for elastic scaling: {e}. "
+                "Worker counts may not align with TPU topology."
+            )
+
+        return 0
+
+    def _get_resize_decision(self, num_workers: int) -> ResizeDecision:
+        return ResizeDecision(
+            num_workers=num_workers,
+            resources_per_worker=self.scaling_config._resources_per_worker_not_none,
+        )
+
+    def make_decision_for_non_running_worker_group(self) -> ScalingDecision:
+        self._maybe_send_resource_request()
+
+        allocated_resources = self._get_allocated_resources()
+        if allocated_resources is None:
+            return NoopDecision()
+
+        num_workers = self._count_possible_workers(allocated_resources)
+
+        if num_workers < self.scaling_config.min_workers:
+            now = time_monotonic()
+            # Only log this warning periodically to avoid spamming logs
+            if (
+                now - self._latest_insufficient_workers_warning_time
+                >= self.INSUFFICIENT_WORKERS_WARNING_INTERVAL_S
+            ):
+                logger.info(
+                    f"Detected ready resources for {num_workers} workers "
+                    "in the cluster. "
+                    "Deciding NOT to start/restart training due to the number of workers "
+                    "falling below the minimum "
+                    f"(min_workers={self.scaling_config.min_workers})."
+                )
+                self._latest_insufficient_workers_warning_time = now
+            return NoopDecision()
+
+        logger.info(
+            f"Detected ready resources for {num_workers} workers "
+            "in the cluster. "
+            "Deciding to start/restart training with this worker group size."
+        )
+        return self._get_resize_decision(num_workers)
+
+    def make_decision_for_running_worker_group(
+        self,
+        worker_group_state: WorkerGroupState,
+        worker_group_status: WorkerGroupPollStatus,
+    ) -> ScalingDecision:
+        self._maybe_send_resource_request()
+
+        # Ensure that we don't make resizing decisions too frequently.
+        # The latest restart time and the latest monitor time (whichever is later)
+        # determine the time of the next resize consideration.
+        latest_consideration_time = max(
+            worker_group_state.start_time, self._latest_monitor_time
+        )
+
+        now = time_monotonic()
+        time_since_latest_consideration = now - latest_consideration_time
+        if (
+            time_since_latest_consideration
+            < self.scaling_config.elastic_resize_monitor_interval_s
+        ):
+            logger.debug(
+                "Skipping resize decision due to the latest resizing consideration "
+                "happening too recently: "
+                "%.2f seconds < ScalingConfig(elastic_resize_monitor_interval_s=%.2f seconds).",
+                time_since_latest_consideration,
+                self.scaling_config.elastic_resize_monitor_interval_s,
+            )
+            return NoopDecision()
+
+        self._latest_monitor_time = now
+
+        allocated_resources = self._get_allocated_resources()
+        if allocated_resources is None:
+            return NoopDecision()
+
+        num_workers = self._count_possible_workers(allocated_resources)
+
+        if num_workers == worker_group_state.num_workers:
+            logger.info(
+                "Did not detect any changes in the cluster resources. "
+                "Training will continue with the same worker group size "
+                f"({num_workers})."
+            )
+            return NoopDecision()
+        elif num_workers < self.scaling_config.min_workers:
+            # This covers an edge case where allocated resources decrease to less
+            # than the minimum number of workers.
+            # This situation is rare, since cluster downsizing typically involves
+            # worker failures. However, this check is still useful to fully
+            # avoid entering an invalid state with fewer workers than the minimum.
+            return NoopDecision()
+
+        logger.info(
+            "Detected changes in the cluster resources. "
+            "Deciding to resize the worker group from "
+            f"{worker_group_state.num_workers} -> {num_workers} workers."
+        )
+        return self._get_resize_decision(num_workers)
+
+    # ---------------------------------------------------
+    # Methods for interacting with AutoscalingCoordinator
+    # ---------------------------------------------------
+
+    def _get_allocated_resources(self) -> Optional[List["ResourceDict"]]:
+        """Get allocated resources from AutoscalingCoordinator.
+        Return None if there is an error."""
+        now = time_monotonic()
+        time_since_last_call = now - self._latest_allocated_resources_query_time
+
+        if time_since_last_call < self.GET_ALLOCATED_RESOURCES_INTERVAL_S:
+            return self._latest_allocated_resources
+
+        allocated_resources = None
+        try:
+            allocated_resources = ray.get(
+                self._autoscaling_coordinator.get_allocated_resources.remote(
+                    self._requester_id
+                ),
+                timeout=AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
+            )
+        except Exception:
+            msg = (
+                f"Failed to get allocated resources for {self._requester_id}."
+                " Will not resize the worker group."
+                " If this only happens transiently during network partition or"
+                " CPU being overloaded, it's safe to ignore this error."
+                " If this error persists, file a GitHub issue."
+            )
+            logger.warning(msg, exc_info=True)
+        finally:
+            self._latest_allocated_resources_query_time = time_monotonic()
+            self._latest_allocated_resources = allocated_resources
+
+        return self._latest_allocated_resources

python/ray/train/v2/_internal/execution/scaling_policy/factory.py

@@ -0,0 +1,16 @@
+from ray.train.v2._internal.execution.scaling_policy import (
+    ElasticScalingPolicy,
+    FixedScalingPolicy,
+    ScalingPolicy,
+)
+from ray.train.v2.api.config import ScalingConfig
+
+
+def create_scaling_policy(scaling_config: ScalingConfig) -> ScalingPolicy:
+    """Create a scaling policy from the given scaling config.
+
+    Defaults to the `FixedScalingPolicy` implementation.
+    """
+    if scaling_config.elasticity_enabled:
+        return ElasticScalingPolicy(scaling_config=scaling_config)
+    return FixedScalingPolicy(scaling_config=scaling_config)

python/ray/train/v2/_internal/execution/scaling_policy/fixed.py

@@ -0,0 +1,30 @@
+from ray.train.v2._internal.execution.scaling_policy import (
+    NoopDecision,
+    ResizeDecision,
+    ScalingDecision,
+    ScalingPolicy,
+)
+from ray.train.v2._internal.execution.worker_group import (
+    WorkerGroupPollStatus,
+    WorkerGroupState,
+)
+
+
+class FixedScalingPolicy(ScalingPolicy):
+    def _get_num_workers_for_resource_request(self) -> int:
+        return self.scaling_config.num_workers
+
+    def make_decision_for_non_running_worker_group(self) -> ScalingDecision:
+        self._maybe_send_resource_request()
+        return ResizeDecision(
+            num_workers=self.scaling_config.num_workers,
+            resources_per_worker=self.scaling_config._resources_per_worker_not_none,
+        )
+
+    def make_decision_for_running_worker_group(
+        self,
+        worker_group_state: WorkerGroupState,
+        worker_group_status: WorkerGroupPollStatus,
+    ) -> ScalingDecision:
+        self._maybe_send_resource_request()
+        return NoopDecision()

python/ray/train/v2/_internal/execution/scaling_policy/scaling_policy.py

@@ -0,0 +1,183 @@
+import abc
+import logging
+import uuid
+from dataclasses import dataclass
+from functools import cached_property
+from typing import Dict
+
+import ray
+from ray.train.v2._internal.execution.callback import ControllerCallback
+from ray.train.v2._internal.execution.context import TrainRunContext
+from ray.train.v2._internal.execution.worker_group import (
+    WorkerGroupPollStatus,
+    WorkerGroupState,
+)
+from ray.train.v2._internal.util import time_monotonic
+from ray.train.v2.api.config import ScalingConfig
+
+logger = logging.getLogger(__name__)
+
+# The time in seconds after which an autoscaling request will expire.
+AUTOSCALING_REQUESTS_EXPIRE_TIME_S = 180
+# Timeout in seconds for getting the result of a call to the AutoscalingCoordinator.
+AUTOSCALING_REQUESTS_GET_TIMEOUT_S = 5
+# Interval in seconds between resource requests to the AutoscalingCoordinator.
+AUTOSCALING_REQUESTS_INTERVAL_S = 20
+
+
+@dataclass
+class ScalingDecision:
+    pass
+
+
+@dataclass
+class NoopDecision(ScalingDecision):
+    pass
+
+
+@dataclass
+class ResizeDecision(ScalingDecision):
+    num_workers: int
+    resources_per_worker: Dict[str, float]
+
+
+class ScalingPolicy(abc.ABC, ControllerCallback):
+    """A policy that determines when and how to scale a worker group.
+
+    This can be used to implement elasticity and fault tolerance.
+
+    Recovery decisions are made when workers are in an inactive or unhealthy state.
+    Upscale decisions are optional and are made when workers are healthy.
+
+    Note: When adding new scaling policies, revisit the shared defaults- particularly if:
+    - AutoscalingCoordinator integration is not needed or a different interface
+      becomes available
+    - Timeout/expiry constants need to diverge between policies
+    - _get_num_workers_for_resource_request() needs variable worker counts
+    - Controller lifecycle behavior diverges
+    """
+
+    # TODO: Restructure these APIs to consider different TrainControllerStates
+    # instead of just running and non-running worker groups.
+
+    def __init__(self, scaling_config: ScalingConfig):
+        self.scaling_config = scaling_config
+        self._requester_id = "train-" + uuid.uuid4().hex
+        self._latest_autoscaling_request_time = float("-inf")
+
+    @abc.abstractmethod
+    def make_decision_for_non_running_worker_group(self) -> ScalingDecision:
+        """Makes a scaling decision when the worker group is initializing
+        or recovering from an error."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def make_decision_for_running_worker_group(
+        self,
+        worker_group_state: WorkerGroupState,
+        worker_group_status: WorkerGroupPollStatus,
+    ) -> ScalingDecision:
+        """Makes a scaling decision when monitoring healthy, running workers."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def _get_num_workers_for_resource_request(self) -> int:
+        """Return the number of workers to request resources for."""
+        raise NotImplementedError
+
+    # ---------------------------------------------------
+    # Methods for interacting with AutoscalingCoordinator
+    # ---------------------------------------------------
+
+    @cached_property
+    def _autoscaling_coordinator(self):
+        from ray.data._internal.cluster_autoscaler.default_autoscaling_coordinator import (
+            get_or_create_autoscaling_coordinator,
+        )
+
+        return get_or_create_autoscaling_coordinator()
+
+    def _maybe_send_resource_request(self):
+        """Send a resource request to AutoscalingCoordinator,
+        if AUTOSCALING_REQUESTS_INTERVAL_S has passed since the last send."""
+        now = time_monotonic()
+        if (
+            now - self._latest_autoscaling_request_time
+            < AUTOSCALING_REQUESTS_INTERVAL_S
+        ):
+            return
+        self._send_resource_request()
+
+    def _send_resource_request(self):
+        """Register training resources with the AutoscalingCoordinator."""
+        resources_per_worker = self.scaling_config._resources_per_worker_not_none
+        num_workers = self._get_num_workers_for_resource_request()
+        label_selectors = self.scaling_config._label_selector_per_worker(num_workers)
+        try:
+            from ray.data._internal.cluster_autoscaler.default_autoscaling_coordinator import (
+                ResourceRequestPriority,
+            )
+
+            ray.get(
+                self._autoscaling_coordinator.request_resources.remote(
+                    requester_id=self._requester_id,
+                    resources=[resources_per_worker] * num_workers,
+                    label_selectors=label_selectors,
+                    expire_after_s=AUTOSCALING_REQUESTS_EXPIRE_TIME_S,
+                    priority=ResourceRequestPriority.HIGH,
+                ),
+                timeout=AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
+            )
+            self._latest_autoscaling_request_time = time_monotonic()
+        except Exception:
+            msg = (
+                f"Failed to send resource request for {self._requester_id}."
+                " If this only happens transiently during network partition or"
+                " CPU being overloaded, it's safe to ignore this error."
+                " If this error persists, file a GitHub issue."
+            )
+            logger.warning(msg, exc_info=True)
+
+    def _cancel_resource_request(self):
+        """Cancel the resource request to AutoscalingCoordinator."""
+        try:
+            ray.get(
+                self._autoscaling_coordinator.cancel_request.remote(
+                    requester_id=self._requester_id,
+                ),
+                timeout=AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
+            )
+        except Exception:
+            msg = (
+                f"Failed to cancel resource request for {self._requester_id}."
+                " The request will still expire after the timeout of"
+                f" {AUTOSCALING_REQUESTS_EXPIRE_TIME_S} seconds."
+            )
+            logger.warning(msg, exc_info=True)
+
+    # --------------------------
+    # ControllerCallback
+    # --------------------------
+
+    def after_controller_start(self, train_run_context: TrainRunContext):
+        """Register training resources with the AutoscalingCoordinator."""
+        self._requester_id = f"train-{train_run_context.run_id}"
+        resources_per_worker = self.scaling_config._resources_per_worker_not_none
+        num_workers = self._get_num_workers_for_resource_request()
+        label_selectors = self.scaling_config._label_selector_per_worker(num_workers)
+        if label_selectors:
+            logger.info(
+                f"Requesting resources: {resources_per_worker} * {num_workers} "
+                f"with label_selectors={label_selectors}"
+            )
+        else:
+            logger.info(f"Requesting resources: {resources_per_worker} * {num_workers}")
+        self._send_resource_request()
+
+    async def before_controller_shutdown(self):
+        """Cancel the resource request when the controller shuts down."""
+        self._cancel_resource_request()
+
+    def before_controller_abort(self):
+        """Cancel the resource request when the controller is aborted."""
+        self._cancel_resource_request()