chore: import upstream snapshot with attribution

python/ray/autoscaler/v2/schema.py

@@ -0,0 +1,621 @@
+import re
+import time
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Dict, List, Optional, Tuple
+
+from ray.autoscaler.v2.instance_manager.common import InstanceUtil
+from ray.core.generated.autoscaler_pb2 import NodeState, NodeStatus
+from ray.core.generated.instance_manager_pb2 import Instance
+
+# TODO(rickyx): once we have graceful shutdown, we could populate
+# the failure detail with the actual termination message. As of now,
+# we will use a more generic message to include cases such as:
+# (idle termination, node death, crash, preemption, etc)
+NODE_DEATH_CAUSE_RAYLET_DIED = "NodeTerminated"
+
+
+# e.g., cpu_4_ondemand.
+NodeType = str
+
+
+@dataclass
+class ResourceUsage:
+    # Resource name.
+    resource_name: str = ""
+    # Total resource.
+    total: float = 0.0
+    # Resource used.
+    used: float = 0.0
+
+
+@dataclass
+class NodeUsage:
+    # The node resource usage.
+    usage: List[ResourceUsage]
+    # How long the node has been idle.
+    idle_time_ms: int
+
+
+@dataclass
+class NodeInfo:
+    # The instance type name, e.g. p3.2xlarge
+    instance_type_name: str
+    # ray node type name.
+    ray_node_type_name: str
+    # Cloud instance id.
+    instance_id: str
+    # Ip address of the node when alive.
+    ip_address: str
+    # The status of the node. Optional for pending nodes.
+    node_status: Optional[str] = None
+    # ray node id in hex. None if still pending.
+    node_id: Optional[str] = None
+    # Resource usage breakdown if node is running.
+    resource_usage: Optional[NodeUsage] = None
+    # Failure detail if the node failed.
+    failure_detail: Optional[str] = None
+    # Descriptive details.
+    details: Optional[str] = None
+    # Activity on the node.
+    node_activity: Optional[List[str]] = None
+    # Ray node labels.
+    labels: Optional[Dict[str, str]] = None
+
+    def total_resources(self) -> Dict[str, float]:
+        if self.resource_usage is None:
+            return {}
+        return {r.resource_name: r.total for r in self.resource_usage.usage}
+
+    def available_resources(self) -> Dict[str, float]:
+        if self.resource_usage is None:
+            return {}
+        return {r.resource_name: r.total - r.used for r in self.resource_usage.usage}
+
+    def used_resources(self) -> Dict[str, float]:
+        if self.resource_usage is None:
+            return {}
+        return {r.resource_name: r.used for r in self.resource_usage.usage}
+
+
+@dataclass
+class LaunchRequest:
+    class Status(Enum):
+        FAILED = "FAILED"
+        PENDING = "PENDING"
+
+    # The instance type name, e.g. p3.2xlarge
+    instance_type_name: str
+    # ray node type name.
+    ray_node_type_name: str
+    # count.
+    count: int
+    # State: (e.g. PENDING, FAILED)
+    state: Status
+    # When the launch request was made in unix timestamp in secs.
+    request_ts_s: int
+    # When the launch request failed unix timestamp in secs if failed.
+    failed_ts_s: Optional[int] = None
+    # Request details, e.g. error reason if the launch request failed.
+    details: Optional[str] = None
+
+
+@dataclass
+class ResourceRequestByCount:
+    # Bundles in the demand.
+    bundle: Dict[str, float]
+    # Number of bundles with the same shape.
+    count: int
+
+    def __str__(self) -> str:
+        return f"[{self.count} {self.bundle}]"
+
+
+@dataclass
+class ResourceDemand:
+    # The bundles in the demand with shape and count info.
+    bundles_by_count: List[ResourceRequestByCount]
+
+
+@dataclass
+class PlacementGroupResourceDemand(ResourceDemand):
+    # Details string (parsed into below information)
+    details: str
+    # Placement group's id.
+    pg_id: Optional[str] = None
+    # Strategy, e.g. STRICT_SPREAD
+    strategy: Optional[str] = None
+    # Placement group's state, e.g. PENDING
+    state: Optional[str] = None
+
+    def __post_init__(self):
+        if not self.details:
+            return
+
+        # Details in the format of <pg_id>:<strategy>|<state>, parse
+        # it into the above fields.
+        pattern = r"^.*:.*\|.*$"
+        match = re.match(pattern, self.details)
+        if not match:
+            return
+
+        pg_id, details = self.details.split(":")
+        strategy, state = details.split("|")
+        self.pg_id = pg_id
+        self.strategy = strategy
+        self.state = state
+
+
+@dataclass
+class RayTaskActorDemand(ResourceDemand):
+    pass
+
+
+@dataclass
+class ClusterConstraintDemand(ResourceDemand):
+    pass
+
+
+@dataclass
+class ResourceDemandSummary:
+    # Placement group demand.
+    placement_group_demand: List[PlacementGroupResourceDemand] = field(
+        default_factory=list
+    )
+    # Ray task actor demand.
+    ray_task_actor_demand: List[RayTaskActorDemand] = field(default_factory=list)
+    # Cluster constraint demand.
+    cluster_constraint_demand: List[ClusterConstraintDemand] = field(
+        default_factory=list
+    )
+
+
+@dataclass
+class Stats:
+    # How long it took to get the GCS request.
+    gcs_request_time_s: float = 0.0
+    # How long it took to get all live instances from node provider.
+    none_terminated_node_request_time_s: Optional[float] = None
+    # How long for autoscaler to process the scaling decision.
+    autoscaler_iteration_time_s: Optional[float] = None
+    # The last seen autoscaler state version from Ray.
+    autoscaler_version: Optional[str] = None
+    # The last seen cluster state resource version.
+    cluster_resource_state_version: Optional[str] = None
+    # Request made time unix timestamp: when the data was pulled from GCS.
+    request_ts_s: Optional[int] = None
+
+
+@dataclass
+class ClusterStatus:
+    # Healthy nodes information (non-idle)
+    active_nodes: List[NodeInfo] = field(default_factory=list)
+    # Idle node information
+    idle_nodes: List[NodeInfo] = field(default_factory=list)
+    # Pending launches.
+    pending_launches: List[LaunchRequest] = field(default_factory=list)
+    # Failed launches.
+    failed_launches: List[LaunchRequest] = field(default_factory=list)
+    # Pending nodes.
+    pending_nodes: List[NodeInfo] = field(default_factory=list)
+    # Failures
+    failed_nodes: List[NodeInfo] = field(default_factory=list)
+    # Resource usage summary for entire cluster.
+    cluster_resource_usage: List[ResourceUsage] = field(default_factory=list)
+    # Demand summary.
+    resource_demands: ResourceDemandSummary = field(
+        default_factory=ResourceDemandSummary
+    )
+    # Query metics
+    stats: Stats = field(default_factory=Stats)
+
+    def total_resources(self) -> Dict[str, float]:
+        return {r.resource_name: r.total for r in self.cluster_resource_usage}
+
+    def available_resources(self) -> Dict[str, float]:
+        return {r.resource_name: r.total - r.used for r in self.cluster_resource_usage}
+
+    # TODO(rickyx): we don't show infeasible requests as of now.
+    # (They will just be pending forever as part of the demands)
+    # We should show them properly in the future.
+
+
+IPPRSpecsSchema = {
+    # JSON schema for IPPR (In-Place Pod Resize) specs provided via the
+    # Kubernetes annotation `ray.io/ippr` on a RayCluster CR.
+    #
+    # Structure:
+    # {
+    #   "groups": {
+    #     "<groupName>": {
+    #       "max-cpu":     string|number,  # K8s quantity (e.g. "2", "1500m")
+    #       "max-memory":  string|integer, # K8s quantity (e.g. "8Gi", 2147483648)
+    #       "resize-timeout": integer      # Seconds to wait for a pod resize to
+    #                                      # complete before considering it timed out
+    #     },
+    #     ...
+    #   }
+    # }
+    #
+    # Notes:
+    # - The set of valid <groupName> keys corresponds to the RayCluster
+    #   `workerGroupSpecs[].groupName` plus the implicit "headgroup".
+    # - The minimal CPU/memory values (min_*) are derived from the pod template
+    #   requests/limits and are not part of this schema.
+    "type": "object",
+    "properties": {
+        "groups": {
+            "type": "object",
+            "additionalProperties": {
+                "type": "object",
+                "properties": {
+                    "max-cpu": {"type": ["string", "number"]},
+                    "max-memory": {"type": ["string", "integer"]},
+                    "resize-timeout": {"type": "integer"},
+                },
+                "required": ["max-cpu", "max-memory", "resize-timeout"],
+            },
+        }
+    },
+}
+
+
+@dataclass
+class IPPRGroupSpec:
+    """Per-group IPPR limits and baseline resources.
+
+    This mirrors a single Ray group (worker group or head group). The minimal
+    resources are derived from the pod template's container resources (either
+    requests or limits if present), and the maximal resources and timeout are
+    provided by the IPPR spec annotation validated by ``IPPRSpecsSchema``.
+
+    Attributes:
+        min_cpu: Baseline CPU in cores derived from the pod template (float, e.g., 1.5).
+        max_cpu: Maximum CPU in cores allowed for in-place resize for this group.
+        min_memory: Baseline memory in bytes derived from the pod template.
+        max_memory: Maximum memory in bytes allowed for in-place resize for this group.
+        resize_timeout: Timeout in seconds for a single resize operation before
+            it is considered timed out.
+    """
+
+    min_cpu: float
+    max_cpu: float
+    min_memory: int
+    max_memory: int
+    resize_timeout: int
+
+
+@dataclass
+class IPPRSpecs:
+    """Typed, validated IPPR specs across Ray groups.
+
+    Attributes:
+        groups: Mapping from Ray group name (e.g., worker group ``groupName`` or
+            ``"headgroup"``) to its ``IPPRGroupSpec``.
+    """
+
+    groups: Dict[str, IPPRGroupSpec]
+
+
+@dataclass
+class IPPRStatus:
+    """Represents the current and target resources for a pod under IPPR.
+
+    This structure is the working state used by the autoscaler to decide if and
+    when to apply in-place pod resizes and when to synchronize resource changes
+    with the Raylet.
+
+    Attributes:
+        cloud_instance_id: Cloud instance identifier for the pod (K8s pod name).
+        spec: The group-level limits and baselines for this pod.
+        current_cpu: Current CPU allocation in cores from the pod status.
+        current_memory: Current memory allocation in bytes from the pod status.
+        desired_cpu: Target CPU allocation in cores.
+        desired_memory: Target memory allocation in bytes.
+        resizing_at: Unix timestamp (seconds) when a resize request was issued to
+            Kubernetes, or None if not pending/needed.
+        k8s_resize_status: Lower-cased status from pod conditions for IPPR, e.g.
+            "inprogress", "deferred", "infeasible", "error"; None indicates no
+            active resize and is treated as finished.
+        k8s_resize_message: Message from the pod condition describing the resize
+            state or failure, if any.
+        suggested_max_cpu: Current-iteration suggested CPU cap (cores), computed from
+            node remaining capacity plus existing gap when resize is deferred/infeasible.
+        suggested_max_memory: Current-iteration suggested memory cap (bytes), computed
+            from node remaining capacity plus existing gap when resize is deferred/infeasible.
+        last_failed_at: Unix timestamp (seconds) when this pod most recently hit a
+            terminal IPPR failure. Once set, the autoscaler stops sending further
+            IPPR requests for this pod.
+        last_failed_reason: Human-readable reason for the terminal IPPR failure.
+        raylet_id: Raylet node id (hex) running in this pod, used to sync Ray's
+            internal resource view when K8s successfully changed pod resources.
+    """
+
+    cloud_instance_id: str
+    spec: IPPRGroupSpec
+    current_cpu: float
+    current_memory: int
+    desired_cpu: float
+    desired_memory: int
+    resizing_at: Optional[int] = None
+    k8s_resize_status: Optional[str] = None
+    k8s_resize_message: Optional[str] = None
+    suggested_max_cpu: Optional[float] = None
+    suggested_max_memory: Optional[int] = None
+    last_failed_at: Optional[int] = None
+    last_failed_reason: Optional[str] = None
+    raylet_id: Optional[str] = None
+
+    def queue_resize_request(
+        self,
+        desired_cpu: Optional[float] = None,
+        desired_memory: Optional[int] = None,
+    ) -> bool:
+        """Queue the new desired resources and reset resize tracking state.
+
+        Queues the new desired CPU/memory if provided, associates the Raylet id,
+        and marks the resize state as "new" so the scheduler can identify the IPPR
+        action before the next iteration.
+
+        Args:
+            desired_cpu: Optional new desired CPU in cores.
+            desired_memory: Optional new desired memory in bytes.
+
+        Returns:
+            bool: True if the resize request is queued.
+        """
+        updated = False
+        if desired_cpu is not None and desired_cpu != self.desired_cpu:
+            self.desired_cpu = desired_cpu
+            updated = True
+        if desired_memory is not None and desired_memory != self.desired_memory:
+            self.desired_memory = desired_memory
+            updated = True
+        if updated:
+            self.resizing_at = None
+            self.k8s_resize_status = "new"
+            self.k8s_resize_message = None
+        return updated
+
+    def has_resize_request_to_send(self) -> bool:
+        """Whether this pod should be sent an IPPR request now.
+
+        Returns True if there is a Raylet id and the status is marked as "new".
+        """
+        return self.raylet_id is not None and self.k8s_resize_status == "new"
+
+    def is_in_progress(self) -> bool:
+        """Whether a resize is on going or about to be issued.
+
+        True if a resize was already issued to K8s (``resizing_at`` set), or if
+        the status is newly queued for resize (``has_resize_request_to_send``).
+        """
+        return self.resizing_at is not None or self.has_resize_request_to_send()
+
+    def is_k8s_resize_finished(self) -> bool:
+        """Whether the Kubernetes-side resize is considered finished.
+
+        We treat ``k8s_resize_status is None`` as finished. After K8s completes a
+        resize, the provider clears ``k8s_resize_status`` as None.
+        """
+        return self.k8s_resize_status is None
+
+    def is_raylet_synced(self) -> bool:
+        """Whether the Raylet's internal resources have been updated.
+
+        After the Raylet has been synchronized,
+        the provider clears ``resizing_at``.
+        """
+        return self.resizing_at is None
+
+    def need_sync_with_raylet(self) -> bool:
+        """Whether the Raylet's internal resources need to be updated.
+
+        Returns True when all of the following hold:
+        - We know the ``raylet_id`` for this pod.
+        - A resize was previously issued to Kubernetes (``resizing_at`` set).
+        - Kubernetes has finished applying the resize (``k8s_resize_status`` is None).
+        - The pod's current resources equal the desired resources.
+
+        In this case, the provider will call GCS's resize_raylet_resource_instances to update the raylet's local
+        resource instances and then clear ``resizing_at`` on the pod annotation.
+        """
+        return (
+            self.raylet_id is not None
+            and self.resizing_at is not None
+            and self.k8s_resize_status is None
+            and self.desired_cpu == self.current_cpu
+            and self.desired_memory == self.current_memory
+        )
+
+    def max_cpu(self) -> float:
+        """Effective maximum CPU cores allowed for this pod.
+
+        Preference order:
+        1) ``suggested_max_cpu`` (discovered limit)
+        2) ``spec.max_cpu`` (static limit)
+        """
+        if self.suggested_max_cpu is not None:
+            return self.suggested_max_cpu
+        return self.spec.max_cpu
+
+    def max_memory(self) -> int:
+        """Effective maximum memory bytes allowed for this pod.
+
+        Preference order:
+        1) ``suggested_max_memory`` (discovered limit)
+        2) ``spec.max_memory`` (static limit)
+        """
+        if self.suggested_max_memory is not None:
+            return self.suggested_max_memory
+        return self.spec.max_memory
+
+    def can_resize_up(self) -> bool:
+        """Whether the pod can still be scaled up within allowed limits.
+
+        Only returns True when no resize is in progress and the current
+        CPU/memory are below the effective max limits. Pods that have already
+        hit a terminal IPPR failure are permanently excluded from future IPPR.
+        """
+        return (
+            self.last_failed_at is None
+            and self.is_k8s_resize_finished()
+            and self.is_raylet_synced()
+            and (
+                self.current_cpu < self.max_cpu()
+                or self.current_memory < self.max_memory()
+            )
+        )
+
+    def is_timeout(self) -> bool:
+        """Whether an in-flight resize has exceeded the group's timeout.
+
+        Returns True when a resize was issued and now current time exceeds
+        ``resizing_at + spec.resize_timeout``.
+        """
+        return (
+            self.resizing_at is not None
+            and not self.need_sync_with_raylet()
+            and (self.resizing_at + self.spec.resize_timeout) < time.time()
+        )
+
+    def is_errored(self) -> bool:
+        """Whether the last resize attempt reported an error from Kubernetes."""
+        return self.k8s_resize_status == "error"
+
+    def record_failure(self, reason: str, failed_at: Optional[int] = None) -> None:
+        """Record the IPPR failure."""
+        self.last_failed_at = int(time.time()) if failed_at is None else failed_at
+        self.last_failed_reason = reason
+
+
+@dataclass
+class AutoscalerInstance:
+    """
+    AutoscalerInstance represents an instance that's managed by the autoscaler.
+    This includes two states:
+        1. the instance manager state: information of the underlying cloud instance.
+        2. the ray node state, e.g. resources, ray node status.
+
+    The two states are linked by the cloud instance id, which should be set
+    when the ray node is started.
+    """
+
+    # The cloud instance id. It could be None if the instance hasn't been assigned
+    # a cloud instance id, e.g. the instance is still in QUEUED or REQUESTED status.
+    cloud_instance_id: Optional[str] = None
+
+    # The ray node state status. It could be None when no ray node is running
+    # or has run on the cloud instance: for example, ray is still being installed
+    # or the instance manager hasn't had a cloud instance assigned (e.g. QUEUED,
+    # REQUESTED).
+    ray_node: Optional[NodeState] = None
+
+    # The instance manager instance state. It would be None when the ray_node is not
+    # None.
+    # It could be None iff:
+    #   1. There's a ray node, but the instance manager hasn't discovered the
+    #   cloud instance that's running this ray process yet. This could happen since
+    #   the instance manager only discovers instances periodically.
+    #
+    #   2. There was a ray node running on the cloud instance, which was already stopped
+    #   and removed from the instance manager state. But the ray state is still lagging
+    #   behind.
+    #
+    #   3. There is a ray node that's unmanaged by the instance manager.
+    #
+    im_instance: Optional[Instance] = None
+
+    # | cloud_instance_id | ray_node | im_instance |
+    # |-------------------|----------|-------------|
+    # | None              | None     | None        | Not possible.
+    # | None              | None     | not None    | OK. An instance hasn't had ray running on it yet. # noqa E501
+    # | None              | Not None | None        | OK. Possible if the ray node is not started by autoscaler. # noqa E501
+    # | None              | Not None | not None    | Not possible - no way to link im instance with ray node. # noqa E501
+    # | not None          | None     | None        | Not possible since cloud instance id is either part of im state or ray node. # noqa E501
+    # | not None          | None     | not None    | OK. e.g. An instance that's not running ray yet. # noqa E501
+    # | not None          | Not None | None        | OK. See scenario 1, 2, 3 above.
+    # | not None          | Not None | not None    | OK. An instance that's running ray.
+    def validate(self) -> Tuple[bool, str]:
+        """Validate the autoscaler instance state.
+
+        Returns:
+            A tuple of (valid, error_msg) where:
+            - valid is whether the state is valid
+            - error_msg is the error message for the validation results.
+        """
+
+        state_combinations = {
+            # (cloud_instance_id is None, ray_node is None, im_instance is None): (valid, error_msg) # noqa E501
+            (True, True, True): (False, "Not possible"),
+            (True, True, False): (True, ""),
+            (True, False, True): (
+                True,
+                "There's a ray node w/o cloud instance id, must be started not "
+                "by autoscaler",
+            ),
+            (True, False, False): (
+                False,
+                "Not possible - no way to link im instance with ray node",
+            ),
+            (False, True, True): (
+                False,
+                "Not possible since cloud instance id is either part of "
+                "im state or ray node",
+            ),
+            (False, True, False): (True, ""),
+            (False, False, True): (True, ""),
+            (False, False, False): (True, ""),
+        }
+
+        valid, error_msg = state_combinations[
+            (
+                self.cloud_instance_id is None,
+                self.ray_node is None,
+                self.im_instance is None,
+            )
+        ]
+        if not valid:
+            return valid, error_msg
+
+        if self.im_instance is not None and self.ray_node is None:
+            # We don't see a ray node, but tracking an im instance.
+            if self.cloud_instance_id is None:
+                if InstanceUtil.is_cloud_instance_allocated(self.im_instance.status):
+                    return (
+                        False,
+                        "instance should be in a status where cloud instance "
+                        "is not allocated.",
+                    )
+            else:
+                if not InstanceUtil.is_cloud_instance_allocated(
+                    self.im_instance.status
+                ):
+                    return (
+                        False,
+                        "instance should be in a status where cloud instance is "
+                        "allocated.",
+                    )
+
+        if self.ray_node is not None:
+            if self.cloud_instance_id != self.ray_node.instance_id:
+                return False, "cloud instance id doesn't match."
+
+        if self.im_instance is not None and self.cloud_instance_id is not None:
+            if self.cloud_instance_id != self.im_instance.cloud_instance_id:
+                return False, "cloud instance id doesn't match."
+
+        return True, ""
+
+    def is_ray_running(self) -> bool:
+        """Whether the ray node is running."""
+        return self.ray_node is not None and self.ray_node.status in [
+            NodeStatus.RUNNING,
+            NodeStatus.IDLE,
+        ]
+
+    def is_ray_stop(self) -> bool:
+        """Whether the ray node is stopped."""
+        return self.ray_node is None or self.ray_node.status in [
+            NodeStatus.DEAD,
+        ]