ray-project--ray/python/ray/data/_internal/execution/resource_manager.py

import logging
import math
import time
from abc import ABC, abstractmethod
from collections import defaultdict
from typing import TYPE_CHECKING, Callable, Dict, Iterable, List, Optional

from ray._common.utils import env_bool, env_float
from ray.data._internal.execution import create_resource_allocator
from ray.data._internal.execution.block_ref_counter import BlockRefCounter
from ray.data._internal.execution.interfaces.execution_options import (
    ExecutionOptions,
    ExecutionResources,
)
from ray.data._internal.execution.interfaces.physical_operator import (
    PhysicalOperator,
    ReportsExtraResourceUsage,
)
from ray.data._internal.execution.operators.base_physical_operator import (
    AllToAllOperator,
)
from ray.data._internal.execution.operators.hash_shuffle import (
    HashShufflingOperatorBase,
)
from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
from ray.data._internal.execution.operators.shuffle_operators.shuffle_map_operator import (  # noqa: E501
    ShuffleMapOp,
)
from ray.data._internal.execution.operators.zip_operator import ZipOperator
from ray.data._internal.execution.util import memory_string
from ray.data.context import DataContext
from ray.util.debug import log_once

if TYPE_CHECKING:
    from ray.data._internal.execution.streaming_executor_state import OpState, Topology


logger = logging.getLogger(__name__)


LOG_DEBUG_TELEMETRY_FOR_RESOURCE_MANAGER_OVERRIDE: Optional[bool] = env_bool(
    "RAY_DATA_DEBUG_RESOURCE_MANAGER", None
)

# Only warn that the cluster can't run any task once the operator has been starved of
# its minimum resources for this long. This avoids spurious warnings while the cluster
# is still scaling up or waiting for a response from the autoscaling coordinator.
#
# I arbitrarily chose the default delay.
STARVATION_WARNING_DELAY_S = env_float("RAY_DATA_STARVATION_WARNING_DELAY_S", 60)


# Following list is a list of *blocking* materializing operators, that prevent
# operators downstream from them from starting execution until these operators
# finish executing.
_BLOCKING_MATERIALIZING_OPERATORS = (
    HashShufflingOperatorBase,
    AllToAllOperator,
    ShuffleMapOp,
    # TODO remove after zip made fully streaming
    ZipOperator,
)


def terminal_operator_from_topology(topology: "Topology") -> PhysicalOperator:
    """Return the executor sink: the unique op with no in-DAG downstream consumers.

    ``build_streaming_topology`` is rooted at the same node passed to
    ``StreamingExecutor``; that root is the only operator whose
    ``output_dependencies`` is empty.
    """
    if not topology:
        raise ValueError("topology must be non-empty")
    sinks = [op for op in topology if not op.output_dependencies]
    if len(sinks) == 1:
        return sinks[0]
    if not sinks:
        raise ValueError(
            "No terminal operator found in topology (expected exactly one "
            "operator with empty output_dependencies)"
        )
    raise ValueError(
        "Expected exactly one terminal operator in topology, found "
        f"{len(sinks)}: {sinks!r}"
    )


class ResourceManager:
    """A class that manages the resource usage of a streaming executor."""

    # The interval in seconds at which the global resource limits are refreshed.
    GLOBAL_LIMITS_UPDATE_INTERVAL_S = 1

    # The fraction of the object store capacity that will be used as the default object
    # store memory limit for the streaming executor,
    # when `OpResourceAllocator` is enabled.
    DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION = env_float(
        "RAY_DATA_OBJECT_STORE_MEMORY_LIMIT_FRACTION", 0.5
    )

    # The fraction of the object store capacity that will be used as the default object
    # store memory limit for the streaming executor,
    # when `OpResourceAllocator` is not enabled.
    DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION_NO_RESERVATION = 0.25

    def __init__(
        self,
        topology: "Topology",
        options: ExecutionOptions,
        get_total_resources: Callable[[], ExecutionResources],
        data_context: DataContext,
        block_ref_counter: BlockRefCounter,
    ):
        self._topology = topology
        self._options = options
        self._get_total_resources = get_total_resources
        self._global_limits = ExecutionResources.zero()
        self._global_limits_last_update_time = 0
        self._global_usage = ExecutionResources.zero()
        self._global_running_usage = ExecutionResources.zero()
        self._global_pending_usage = ExecutionResources.zero()
        self._op_usages: Dict[PhysicalOperator, ExecutionResources] = {}
        self._op_running_usages: Dict[PhysicalOperator, ExecutionResources] = {}
        self._op_pending_usages: Dict[PhysicalOperator, ExecutionResources] = {}
        # Object store memory usage of pending task outputs (blocks being
        # generated by running tasks but not yet yielded).
        self._mem_op_internal: Dict[PhysicalOperator, int] = defaultdict(int)
        # Object store memory usage of the operator's outputs, including:
        # internal output queue, external output buffer in OpState, and
        # downstream operators' input buffers (inqueue + pending task inputs).
        self._mem_op_outputs: Dict[PhysicalOperator, int] = defaultdict(int)

        # Bytes buffered by external consumers (iterators) consuming Batches
        # (including the prefetched blocks). For example,
        # - ds.iter_batches -> one iterator
        # - streaming_split -> multiple iterators
        self._external_consumer_bytes: int = 0
        self._has_external_consumer: bool = False

        # Executor sink (DAG root: unique op with no output_dependencies).
        # Iterator/streaming_split prefetch bytes are charged on this
        # operator's output usage.
        self._output_operator = terminal_operator_from_topology(topology)

        self._block_ref_counter = block_ref_counter

        self._op_resource_allocator: Optional[
            "OpResourceAllocator"
        ] = create_resource_allocator(self, data_context)

        self._object_store_memory_limit_fraction = (
            data_context.override_object_store_memory_limit_fraction
            if data_context.override_object_store_memory_limit_fraction is not None
            else (
                self.DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION
                if self.op_resource_allocator_enabled()
                else self.DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION_NO_RESERVATION
            )
        )

    @property
    def has_external_consumer(self) -> bool:
        """Return whether there is any external consumer."""
        return self._has_external_consumer

    def set_external_consumer_bytes(self, num_bytes: int) -> None:
        """Set the bytes buffered by external consumers."""
        assert (
            num_bytes >= 0
        ), f"external consumer bytes must be non-negative, got {num_bytes}"
        self._external_consumer_bytes = num_bytes
        self._has_external_consumer = True

    def get_external_consumer_bytes(self) -> int:
        """Get the bytes buffered by external consumers."""
        return self._external_consumer_bytes

    def _estimate_object_store_memory_usage(
        self, op: "PhysicalOperator", state: "OpState"
    ) -> int:
        # Don't count input refs towards dynamic memory usage, as they have been
        # pre-created already outside this execution.
        if isinstance(op, InputDataBuffer):
            if op is self._output_operator:
                self._mem_op_internal[op] = 0
                self._mem_op_outputs[op] = self._external_consumer_bytes
                return self._external_consumer_bytes
            return 0

        usage = op.estimate_object_store_usage(state)
        self._mem_op_internal[op] = usage.internal
        self._mem_op_outputs[op] = usage.outputs

        # Attribute iterator / streaming_split prefetch to the executor sink only.
        if op is self._output_operator:
            self._mem_op_outputs[op] += self._external_consumer_bytes

        return self._mem_op_outputs[op] + self._mem_op_internal[op]

    def update_usages(self):
        """Recalculate resource usages."""
        # TODO(hchen): This method will be called frequently during the execution loop.
        # And some computations are redundant. We should either remove redundant
        # computations or remove this method entirely and compute usages on demand.
        self._op_usages.clear()
        self._op_running_usages.clear()
        self._op_pending_usages.clear()

        # Iterate from last to first operator.
        for op, state in reversed(self._topology.items()):
            # Update `self._op_usages`, `self._op_running_usages`,
            # and `self._op_pending_usages`.
            op_usage = op.current_logical_usage()
            op_running_usage = op.running_logical_usage()
            op_pending_usage = op.pending_logical_usage()

            assert not op_usage.object_store_memory
            assert not op_running_usage.object_store_memory
            assert not op_pending_usage.object_store_memory

            used_object_store = self._estimate_object_store_memory_usage(op, state)

            op_usage = op_usage.copy(object_store_memory=used_object_store)
            op_running_usage = op_running_usage.copy(
                object_store_memory=used_object_store
            )

            if isinstance(op, ReportsExtraResourceUsage):
                op_usage.add(op.extra_resource_usage())

            self._op_usages[op] = op_usage
            self._op_running_usages[op] = op_running_usage
            self._op_pending_usages[op] = op_pending_usage

            # Update operator's object store usage, which is used by
            # DatasetStats and updated on the Ray Data dashboard.
            op._metrics.obj_store_mem_used = op_usage.object_store_memory

        # Roll the per-op usages up into the global totals in a single pass
        # each (one allocation per total instead of one per operator).
        self._global_usage = ExecutionResources.combine_sum(self._op_usages.values())
        self._global_running_usage = ExecutionResources.combine_sum(
            self._op_running_usages.values()
        )
        self._global_pending_usage = ExecutionResources.combine_sum(
            self._op_pending_usages.values()
        )

        if self._op_resource_allocator is not None:
            self._update_allocated_budgets()

    def _update_allocated_budgets(self):
        completed_ops_usage = self._get_completed_ops_usage()

        available_limits = (
            self.get_global_limits()
            .subtract(completed_ops_usage)
            .max(ExecutionResources.zero())
        )

        self._op_resource_allocator.update_budgets(limits=available_limits)

    def get_global_usage(self) -> ExecutionResources:
        """Return the global resource usage at the current time."""
        assert (
            self._global_usage.is_non_negative()
        ), f"Global usage should be non-negative, got {self._global_usage}"
        return self._global_usage

    def get_global_running_usage(self) -> ExecutionResources:
        """Return the global running resource usage at the current time."""
        return self._global_running_usage

    def get_global_pending_usage(self) -> ExecutionResources:
        """Return the global pending resource usage at the current time."""
        return self._global_pending_usage

    def get_global_limits(self) -> ExecutionResources:
        """Return the global resource limits at the current time.

        This method autodetects any unspecified execution resource limits based on the
        current cluster size, refreshing these values periodically to support cluster
        autoscaling.
        """
        if (
            time.time() - self._global_limits_last_update_time
            < self.GLOBAL_LIMITS_UPDATE_INTERVAL_S
        ):
            return self._global_limits

        self._global_limits_last_update_time = time.time()
        default_limits = self._options.resource_limits
        exclude = self._options.exclude_resources
        total_resources = self._get_total_resources()
        default_mem_fraction = self._object_store_memory_limit_fraction
        total_resources = total_resources.copy(
            object_store_memory=total_resources.object_store_memory
            * default_mem_fraction
        )
        # Clamp to non-negative because exclude_resources (e.g., training worker
        # CPUs) can exceed the total resources reported by the cluster autoscaler,
        # such as when Ray Train reserves more CPUs than are visible to Ray Data.
        self._global_limits = (
            default_limits.min(total_resources)
            .subtract(exclude)
            .max(ExecutionResources.zero())
        )
        return self._global_limits

    def get_op_usage(
        self, op: PhysicalOperator, include_ineligible_downstream: bool = False
    ) -> ExecutionResources:
        """Return the resource usage of the given operator at the current time."""
        own_usage = self._op_usages[op]

        if not include_ineligible_downstream:
            return own_usage

        return own_usage.add(self._get_downstream_ineligible_ops_usage(op))

    def _get_downstream_ineligible_ops_usage(
        self, op: PhysicalOperator
    ) -> ExecutionResources:
        return ExecutionResources.combine_sum(
            self.get_op_usage(op) for op in self._get_downstream_ineligible_ops(op)
        )

    def get_mem_op_internal(self, op: PhysicalOperator) -> int:
        """Return the memory usage of pending task outputs for the given operator."""
        return self._mem_op_internal[op]

    def get_mem_op_outputs(
        self,
        op: PhysicalOperator,
        include_ineligible_downstream: bool = False,
    ) -> int:
        """Return the memory usage of the outputs of the given operator."""
        # Outputs usage of the current operator.
        op_outputs_usage = self._mem_op_outputs[op]

        if not include_ineligible_downstream:
            return op_outputs_usage

        # Also account the downstream ineligible operators' memory usage.
        return (
            op_outputs_usage
            + self._get_downstream_ineligible_ops_usage(op).object_store_memory
        )

    def get_op_usage_str(self, op: PhysicalOperator, *, verbose: bool) -> str:
        """Return a human-readable string representation of the resource usage of
        the given operator."""
        # Handle case where operator is not in _op_running_usages dict
        if op not in self._op_running_usages:
            usage_str = "n/a"
        else:
            usage_str = f"{self._op_running_usages[op].cpu:.1f} CPU"
            if self._op_running_usages[op].memory:
                usage_str += f", {self._op_running_usages[op].memory_str()} memory"
            if self._op_running_usages[op].gpu:
                usage_str += f", {self._op_running_usages[op].gpu:.1f} GPU"
            usage_str += f", {self._op_running_usages[op].object_store_memory_str()} object store"

        # NOTE: Config can override requested verbosity level
        if LOG_DEBUG_TELEMETRY_FOR_RESOURCE_MANAGER_OVERRIDE is not None:
            verbose = LOG_DEBUG_TELEMETRY_FOR_RESOURCE_MANAGER_OVERRIDE

        if verbose:
            usage_str += (
                f" (in={memory_string(self.get_mem_op_internal(op))},"
                f"out={memory_string(self.get_mem_op_outputs(op))}"
            )
            # External-consumer bytes (iterator / streaming_split prefetch) are
            # only attached to the output operator. Surface them in its line so
            # users can see how much of `out` is held by the downstream iterator
            # vs. the operator's own output queues.
            if op is self._output_operator and self._has_external_consumer:
                usage_str += (
                    f",external_consumer="
                    f"{memory_string(self._external_consumer_bytes)}"
                )
            usage_str += ")"
            if self._op_resource_allocator is not None:
                allocation = self._op_resource_allocator.get_allocation(op)
                if allocation:
                    usage_str += f", alloc=(cpu={allocation.cpu:.1f}"
                    usage_str += f",mem={allocation.memory_str()}"
                    usage_str += f",gpu={allocation.gpu:.1f}"
                    usage_str += f",obj_store={allocation.object_store_memory_str()})"

                budget = self._op_resource_allocator.get_budget(op)
                if budget:
                    usage_str += f", budget=(cpu={budget.cpu:.1f}"
                    usage_str += f",mem={budget.memory_str()}"
                    usage_str += f",gpu={budget.gpu:.1f}"
                    usage_str += f",obj_store={budget.object_store_memory_str()}"

                    # Remaining memory budget for producing new task outputs.
                    if isinstance(
                        self._op_resource_allocator, ReservationOpResourceAllocator
                    ):
                        reserved_for_output = memory_string(
                            self._op_resource_allocator._output_budgets.get(op, 0)
                        )
                        usage_str += f",out={reserved_for_output})"

        return usage_str

    def op_resource_allocator_enabled(self) -> bool:
        """Return whether OpResourceAllocator is enabled."""
        return self._op_resource_allocator is not None

    @property
    def op_resource_allocator(self) -> "OpResourceAllocator":
        """Return the OpResourceAllocator."""
        assert self._op_resource_allocator is not None
        return self._op_resource_allocator

    def get_budget(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
        """Return the budget for the given operator, or None if the operator
        has unlimited budget."""
        if self._op_resource_allocator is None:
            return None
        return self._op_resource_allocator.get_budget(op)

    def get_allocation(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
        """Return the allocation of the given operator, or None if the operator
        doesn't have a designated allocation."""
        if self._op_resource_allocator is None:
            return None
        return self._op_resource_allocator.get_allocation(op)

    def is_op_eligible(self, op: PhysicalOperator) -> bool:
        """Whether the op is eligible for memory reservation."""
        return (
            not op.throttling_disabled()
            # As long as the op has finished execution, even if there are still
            # non-taken outputs, we don't need to allocate resources for it.
            and not op.has_execution_finished()
        )

    def _get_downstream_ineligible_ops(
        self, op: PhysicalOperator
    ) -> Iterable[PhysicalOperator]:
        """Get the downstream ineligible operators of the given operator.

        E.g.,
          - "cur_map->downstream_map" will return an empty list.
          - "cur_map->limit1->limit2->downstream_map" will return [limit1, limit2].
        """
        for next_op in op.output_dependencies:
            if not self.is_op_eligible(next_op):
                yield next_op
                yield from self._get_downstream_ineligible_ops(next_op)

    def get_downstream_eligible_ops(
        self, op: PhysicalOperator
    ) -> Iterable[PhysicalOperator]:
        """Get the downstream eligible operators of the given operator, ignoring
        intermediate ineligible operators.

        E.g.,
          - "cur_map->downstream_map" will return [downstream_map].
          - "cur_map->limit1->limit2->downstream_map" will return [downstream_map].
        """
        for next_op in op.output_dependencies:
            if self.is_op_eligible(next_op):
                yield next_op
            else:
                yield from self.get_downstream_eligible_ops(next_op)

    def max_task_output_bytes_to_read(self, op: PhysicalOperator) -> Optional[int]:
        if self._op_resource_allocator is not None:
            return self._op_resource_allocator.max_task_output_bytes_to_read(op)
        return None

    def _get_completed_ops_usage(self) -> ExecutionResources:
        """
        Resource reservation is based on the number of eligible operators.
        However, there might be completed operators that still have blocks in their output queue, which we need to exclude them from the reservation.
        And we also need to exclude the downstream ineligible operators.

        E.g., for the following pipeline:
        ```
        map1 (completed, but still has blocks in its output queue) -> limit1 (ineligible, not completed) -> map2 (not completed) -> limit2 -> map3
        ```

        The reservation is based on the number of eligible operators (map2 and map3), but we need to exclude map1 and limit1 from the reservation.
        """

        last_completed_ops = []
        ops_to_exclude = []
        # Traverse operator tree collecting all operators that have already finished
        for op in self._topology:
            if not op.has_execution_finished():
                for dep in op.input_dependencies:
                    if dep.has_execution_finished():
                        last_completed_ops.append(dep)

        # In addition to completed operators,
        # filter out downstream ineligible operators since they are omitted from reservation calculations.
        for op in last_completed_ops:
            ops_to_exclude.extend(list(self._get_downstream_ineligible_ops(op)))
            ops_to_exclude.append(op)

        completed_ops = list(set(ops_to_exclude))
        completed_ops_usage = ExecutionResources.combine_sum(
            self.get_op_usage(op) for op in completed_ops
        )

        return completed_ops_usage

    def _is_blocking_materializing_op(self, op: PhysicalOperator) -> bool:
        """This method checks whether either

        1. Operator itself or
        2. One of operator's immediate *ineligible* downstream dependencies

        Are blocking, materializing operators.

        NOTE: That downstream ineligible operators are considered an "extension" of their
              first preceding eligible operator from resource allocation standpoint.
        """

        # Check if Op itself is a blocking, materializing operator
        if isinstance(op, _BLOCKING_MATERIALIZING_OPERATORS):
            return True

        # Check if any of its direct *ineligible* downstream dependencies are
        # blocking, materializing operators.
        #
        # NOTE: We only check ineligible downstream deps, since eligible downstream
        #       deps will have their own allocation that is adjusted appropriately
        return any(
            isinstance(op, _BLOCKING_MATERIALIZING_OPERATORS)
            for op in self._get_downstream_ineligible_ops(op)
        )


def _get_first_pending_materializing_op(topology: "Topology") -> int:
    for idx, op in enumerate(topology):
        if isinstance(op, _BLOCKING_MATERIALIZING_OPERATORS) and not op.has_completed():
            return idx

    return -1


class OpResourceAllocator(ABC):
    """An interface for dynamic operator resource allocation.

    This interface allows dynamically allocating available resources to each operator,
    limiting how many tasks each operator can submit, and how much data each operator
    can read from its running tasks.
    """

    def __init__(self, resource_manager: "ResourceManager"):
        self._resource_manager = resource_manager
        self._topology = resource_manager._topology

    @abstractmethod
    def update_budgets(
        self,
        *,
        limits: ExecutionResources,
    ):
        """Callback to update resource usages."""
        ...

    @abstractmethod
    def can_submit_new_task(self, op: PhysicalOperator) -> bool:
        """Return whether the given operator can submit a new task."""
        ...

    @abstractmethod
    def max_task_output_bytes_to_read(self, op: PhysicalOperator) -> Optional[int]:
        """Return the maximum bytes of pending task outputs can be read for
        the given operator. None means no limit."""
        ...

    @abstractmethod
    def get_budget(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
        """Returns the budget for the given operator or `None` if the operator
        has unlimited budget. Operator's budget is defined as:

            Budget = Allocation - Usage
        """
        ...

    @abstractmethod
    def get_output_budget(self, op: PhysicalOperator) -> Optional[int]:
        """Returns the budget for operator's outputs (in object store bytes) or
        `None` if there's no limit.
        """
        ...

    @abstractmethod
    def get_allocation(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
        """Returns allocation for the given operator or `None` if operator's
        allocation is unlimited."""
        ...

    def _get_eligible_ops(self) -> List[PhysicalOperator]:
        """Returns a list of operators eligible for allocation.

        Only operators upstream of the first non-completed *materializing* Operator,
        (like shuffle, etc) are able to receive inputs and start execution.

        Therefore only these operators are eligible for resource allocation.
        """
        first_pending_materializing_op_idx = _get_first_pending_materializing_op(
            self._topology
        )
        return [
            op
            for idx, op in enumerate(self._topology)
            if self._is_op_eligible(op)
            and (
                first_pending_materializing_op_idx == -1
                or idx <= first_pending_materializing_op_idx
            )
        ]

    @staticmethod
    def _is_op_eligible(op: PhysicalOperator) -> bool:
        """Whether the op is eligible for memory reservation."""
        return (
            not op.throttling_disabled()
            # As long as the op has finished execution, even if there are still
            # non-taken outputs, we don't need to allocate resources for it.
            and not op.has_execution_finished()
        )


class ReservationOpResourceAllocator(OpResourceAllocator):
    """An OpResourceAllocator implementation that reserves resources for each operator.

    This class reserves memory and CPU resources for eligible operators, and considers
    runtime resource usages to limit the resources that each operator can use.

    It works in the following way:
    1. An operator is eligible for resource reservation, if it has enabled throttling
       and hasn't completed. Ineligible operators are not throttled, but
       their usage will be accounted for their upstream eligible operators. E.g., for
       such a dataset "map1->limit->map2->streaming_split", we'll treat "map1->limit" as
       a group and "map2->streaming_split" as another group.
    2. For each eligible operator, we reserve `reservation_ratio * global_resources /
        num_eligible_ops` resources, half of which is reserved only for the operator
        outputs, excluding pending task outputs.
    3. Non-reserved resources are shared among all operators.
    4. In each scheduling iteration, each eligible operator will get "remaining of their
       own reserved resources" + "remaining of shared resources / num_eligible_ops"
       resources.

    The `reservation_ratio` is set to 50% by default. Users can tune this value to
    adjust how aggressive or conservative the resource allocation is. A higher value
    will make the resource allocation more even, but may lead to underutilization and
    worse performance. And vice versa.
    """

    def __init__(self, resource_manager: ResourceManager, reservation_ratio: float):
        super().__init__(resource_manager)

        self._reservation_ratio = reservation_ratio
        assert 0.0 <= self._reservation_ratio <= 1.0
        # Per-op reserved resources, excluding `_reserved_for_op_outputs`.
        self._op_reserved: Dict[PhysicalOperator, ExecutionResources] = {}
        # Memory reserved exclusively for the outputs of each operator.
        # "Op outputs" refer to blocks that have been taken out of an operator,
        # i.e., `RessourceManager._mem_op_outputs`.
        #
        # Note, if we don't reserve memory for op outputs, all the budget may be used by
        # the pending task outputs, and/or op's internal output buffers (the latter can
        # happen when `preserve_order=True`).
        # Then we'll have no budget to pull blocks from the op.
        self._reserved_for_op_outputs: Dict[PhysicalOperator, float] = {}
        # Total shared resources.
        self._total_shared = ExecutionResources.zero()
        # Resource budgets for each operator, excluding `_reserved_for_op_outputs`.
        self._op_budgets: Dict[PhysicalOperator, ExecutionResources] = {}
        # Remaining memory budget for generating new task outputs, per operator.
        self._output_budgets: Dict[PhysicalOperator, float] = {}
        # Whether each operator has reserved the minimum resources to run
        # at least one task.
        # This is used to avoid edge cases where the entire resource limits are not
        # enough to run one task of each op.
        # See `test_no_deadlock_on_small_cluster_resources` as an example.
        self._reserved_min_resources: Dict[PhysicalOperator, bool] = {}
        # `time.monotonic()` timestamp at which each operator most recently became
        # starved of its minimum resources, or None if it currently has them.
        self._op_starved_since: Dict[PhysicalOperator, Optional[float]] = {}

    def _update_reservation(self, limits: ExecutionResources):
        eligible_ops = self._get_eligible_ops()

        self._op_reserved.clear()
        self._reserved_for_op_outputs.clear()
        self._reserved_min_resources.clear()

        if len(eligible_ops) == 0:
            return

        remaining = limits.copy()

        # Reserve `reservation_ratio * global_limits / num_ops` resources for each
        # operator.
        default_reserved = limits.scale(self._reservation_ratio / (len(eligible_ops)))
        for index, op in enumerate(eligible_ops):
            # Reserve at least half of the default reserved resources for the outputs.
            # This makes sure that we will have enough budget to pull blocks from the
            # op.
            reserved_for_outputs = ExecutionResources(
                0, 0, max(default_reserved.object_store_memory / 2, 1)
            )

            reserved_for_tasks = default_reserved.subtract(reserved_for_outputs)

            min_resource_usage, max_resource_usage = op.min_max_resource_requirements()

            if min_resource_usage is not None:
                reserved_for_tasks = reserved_for_tasks.max(min_resource_usage)
            if max_resource_usage is not None:
                reserved_for_tasks = reserved_for_tasks.min(max_resource_usage)

            # Check if the remaining resources are enough for both reserved_for_tasks
            # and reserved_for_outputs. Note, we only consider CPU and GPU, but not
            # object_store_memory, because object_store_memory can be oversubscribed,
            # but CPU/GPU cannot.
            if reserved_for_tasks.add(reserved_for_outputs).satisfies_limit(
                remaining, ignore_object_store_memory=True
            ):
                self._reserved_min_resources[op] = True
                self._op_starved_since[op] = None
            else:
                self._reserved_min_resources[op] = False
                if self._op_starved_since.get(op) is None:
                    self._op_starved_since[op] = time.monotonic()

                # If the remaining resources are not enough to reserve the minimum
                # resources for this operator, we'll only reserve the minimum object
                # store memory, but not the CPU and GPU resources.
                # Because Ray Core doesn't allow CPU/GPU resources to be oversubscribed.
                # NOTE: we prioritize upstream operators for minimum resource reservation.
                # ops. It's fine that downstream ops don't get the minimum reservation,
                # because they can wait for upstream ops to finish and release resources.
                reserved_for_tasks = ExecutionResources(
                    0, 0, min_resource_usage.object_store_memory
                )

            # Log a warning if even the first operator cannot reserve the minimum
            # resources.
            if index == 0:
                self._warn_if_op_starved_too_long(op)

            self._op_reserved[op] = reserved_for_tasks
            self._reserved_for_op_outputs[op] = reserved_for_outputs.object_store_memory

            op_total_reserved = reserved_for_tasks.add(reserved_for_outputs)
            remaining = remaining.subtract(op_total_reserved)
            remaining = remaining.max(ExecutionResources.zero())

        self._total_shared = remaining

    def _warn_if_op_starved_too_long(self, op: PhysicalOperator) -> None:
        # The operator isn't starved. Return early.
        if self._op_starved_since.get(op) is None:
            return

        op_starved_duration = time.monotonic() - self._op_starved_since[op]
        if (
            op_starved_duration >= STARVATION_WARNING_DELAY_S
            # Add `id(self)` to the log_once key so that it will be logged once per
            # execution.
            and log_once(f"starvation_warning_{id(self)}")
        ):
            logger.warning(
                f"Cluster resources are not enough to run any task from {op}."
                " The job may hang forever unless the cluster scales up."
            )

    def can_submit_new_task(self, op: PhysicalOperator) -> bool:
        """Return whether the given operator can submit a new task based on budget."""
        budget = self.get_budget(op)

        if budget is None:
            return True

        return (
            op.incremental_resource_usage().satisfies_limit(budget)
            and
            # Avoid scheduling if there's no more Object Store budget (for
            # task outputs)
            budget.object_store_memory
            >= (op.metrics.obj_store_mem_max_pending_output_per_task or 0)
        )

    def get_budget(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
        return self._op_budgets.get(op)

    def get_output_budget(self, op: PhysicalOperator) -> Optional[int]:
        return self._output_budgets.get(op)

    def get_allocation(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
        budget = self.get_budget(op)
        if budget is None:
            return None
        return budget.add(self._resource_manager.get_op_usage(op))

    def _get_total_reserved(self, op: PhysicalOperator) -> ExecutionResources:
        """Get total reserved resources for an operator, including outputs reservation."""
        op_reserved = self._op_reserved[op]
        reserved_for_outputs = self._reserved_for_op_outputs[op]
        return op_reserved.copy(
            object_store_memory=op_reserved.object_store_memory + reserved_for_outputs
        )

    def max_task_output_bytes_to_read(self, op: PhysicalOperator) -> Optional[int]:
        if op not in self._op_budgets:
            return None

        res = self._op_budgets[op].object_store_memory
        # Add the remaining of `_reserved_for_op_outputs`.
        op_outputs_usage = self._resource_manager.get_mem_op_outputs(
            op, include_ineligible_downstream=True
        )

        res += max(self._reserved_for_op_outputs[op] - op_outputs_usage, 0)
        if math.isinf(res):
            self._output_budgets[op] = res
            return None

        res = int(res)
        assert res >= 0
        self._output_budgets[op] = res
        return res

    def update_budgets(
        self,
        *,
        limits: ExecutionResources,
    ):
        # Remaining resources to be distributed across operators
        remaining_shared = self._update_reservation(limits)

        self._op_budgets.clear()
        eligible_ops = self._get_eligible_ops()
        if len(eligible_ops) == 0:
            return

        # Remaining of shared resources.
        remaining_shared = self._total_shared
        for op in eligible_ops:
            # Calculate the memory usage of the operator.
            op_mem_usage = 0
            # Add the memory usage of the operator itself,
            # excluding `_reserved_for_op_outputs`.
            op_mem_usage += self._resource_manager.get_mem_op_internal(op)
            # Add the portion of op outputs usage that has
            # exceeded `_reserved_for_op_outputs`.
            op_outputs_usage = self._resource_manager.get_mem_op_outputs(
                op, include_ineligible_downstream=True
            )
            op_mem_usage += max(op_outputs_usage - self._reserved_for_op_outputs[op], 0)

            op_usage = self._resource_manager.get_op_usage(op).copy(
                object_store_memory=op_mem_usage
            )

            op_reserved = self._op_reserved[op]
            # How much of the reserved resources are remaining.
            op_reserved_remaining = op_reserved.subtract(op_usage).max(
                ExecutionResources.zero()
            )

            self._op_budgets[op] = op_reserved_remaining
            # How much of the reserved resources are exceeded.
            # If exceeded, we need to subtract from the remaining shared resources.
            op_reserved_exceeded = op_usage.subtract(op_reserved).max(
                ExecutionResources.zero()
            )
            remaining_shared = remaining_shared.subtract(op_reserved_exceeded)

        remaining_shared = remaining_shared.max(ExecutionResources.zero())

        # Allocate the remaining shared resources to each operator.
        for i, op in enumerate(reversed(eligible_ops)):
            # By default, divide the remaining shared resources equally.
            op_shared = remaining_shared.scale(1.0 / (len(eligible_ops) - i))
            # But if the op's budget is less than `min_scheduling_resources`,
            # it will be useless. So we'll let the downstream operator
            # borrow some resources from the upstream operator, if remaining_shared
            # is still enough.
            to_borrow = (
                op.min_scheduling_resources()
                .subtract(self._op_budgets[op].add(op_shared))
                .max(ExecutionResources.zero())
            )
            if not to_borrow.is_zero() and op_shared.add(to_borrow).satisfies_limit(
                remaining_shared
            ):
                op_shared = op_shared.add(to_borrow)

            # Cap op_shared so that total allocation doesn't exceed max_resource_usage.
            # Total allocation = max(total_reserved, op_usage) + op_shared
            # This ensures excess resources stay in remaining_shared for other operators.
            _, max_resource_usage = op.min_max_resource_requirements()
            if max_resource_usage != ExecutionResources.inf():
                total_reserved = self._get_total_reserved(op)
                op_usage = self._resource_manager.get_op_usage(op)
                current_allocation = total_reserved.max(op_usage)
                max_shared = max_resource_usage.subtract(current_allocation).max(
                    ExecutionResources.zero()
                )
                op_shared = op_shared.min(max_shared)

            remaining_shared = remaining_shared.subtract(op_shared)
            assert remaining_shared.is_non_negative(), (
                remaining_shared,
                op,
                op_shared,
                to_borrow,
            )

            self._op_budgets[op] = self._op_budgets[op].add(op_shared)

        # Give any remaining shared resources to the most downstream uncapped op.
        # This can happen when some ops have their shared allocation capped.
        if eligible_ops and not remaining_shared.is_zero():
            for op in reversed(eligible_ops):
                _, max_resource_usage = op.min_max_resource_requirements()
                if max_resource_usage == ExecutionResources.inf():
                    self._op_budgets[op] = self._op_budgets[op].add(remaining_shared)
                    break

        # A materializing operator like `AllToAllOperator` waits for all its input
        # operator's outputs before processing data. This often forces the input
        # operator to exceed its object store memory budget. To prevent deadlock, we
        # disable object store memory backpressure for the input operator.
        for op in eligible_ops:
            if self._resource_manager._is_blocking_materializing_op(op):
                self._op_budgets[op] = self._op_budgets[op].copy(
                    object_store_memory=float("inf")
                )