chore: import upstream snapshot with attribution

python/ray/autoscaler/_private/node_provider_availability_tracker.py

@@ -0,0 +1,165 @@
+import threading
+import time
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional, Tuple
+
+from ray.autoscaler._private.constants import (
+    AUTOSCALER_NODE_AVAILABILITY_MAX_STALENESS_S,
+)
+from ray.autoscaler.node_launch_exception import NodeLaunchException
+
+
+@dataclass
+class UnavailableNodeInformation:
+    category: str
+    description: str
+
+
+@dataclass
+class NodeAvailabilityRecord:
+    node_type: str
+    is_available: bool
+    last_checked_timestamp: float
+    unavailable_node_information: Optional[UnavailableNodeInformation]
+
+
+@dataclass
+class NodeAvailabilitySummary:
+    node_availabilities: Dict[
+        str, NodeAvailabilityRecord
+    ]  # Mapping from node type to node availability record.
+
+    @classmethod
+    def from_fields(cls, **fields) -> Optional["NodeAvailabilitySummary"]:
+        """Implement marshalling from nested fields. pydantic isn't a core dependency
+        so we're implementing this by hand instead."""
+        parsed = {}
+
+        node_availabilites_dict = fields.get("node_availabilities", {})
+
+        for node_type, node_availability_record_dict in node_availabilites_dict.items():
+            unavailable_information_dict = node_availability_record_dict.pop(
+                "unavailable_node_information", None
+            )
+            unavaiable_information = None
+            if unavailable_information_dict is not None:
+                unavaiable_information = UnavailableNodeInformation(
+                    **unavailable_information_dict
+                )
+
+            parsed[node_type] = NodeAvailabilityRecord(
+                unavailable_node_information=unavaiable_information,
+                **node_availability_record_dict,
+            )
+
+        return NodeAvailabilitySummary(node_availabilities=parsed)
+
+    def __eq__(self, other: "NodeAvailabilitySummary"):
+        return self.node_availabilities == other.node_availabilities
+
+    def __bool__(self) -> bool:
+        return bool(self.node_availabilities)
+
+
+class NodeProviderAvailabilityTracker:
+    """A thread safe, TTL cache of node provider availability. We don't use
+    cachetools.TTLCache because it always sets the expiration time relative to
+    insertion time, but in our case, we want entries to expire relative to when
+    the node creation was attempted (and entries aren't necessarily added in
+    order). We want the entries to expire because the information grows stale
+    over time.
+    """
+
+    def __init__(
+        self,
+        timer: Callable[[], float] = time.time,
+        ttl: float = AUTOSCALER_NODE_AVAILABILITY_MAX_STALENESS_S,
+    ):
+        """A cache that tracks the availability of nodes and throw away
+        entries which have grown too stale.
+
+        Args:
+          timer: A function that returns the current time in seconds.
+          ttl: The ttl from the insertion timestamp of an entry.
+        """
+        self.timer = timer
+        self.ttl = ttl
+        # Mapping from node type to (eviction_time, record)
+        self.store: Dict[str, Tuple[float, NodeAvailabilityRecord]] = {}
+        # A global lock to simplify thread safety handling.
+        self.lock = threading.RLock()
+
+    def _update_node_availability_requires_lock(
+        self,
+        node_type: str,
+        timestamp: int,
+        node_launch_exception: Optional[NodeLaunchException],
+    ) -> None:
+        if node_launch_exception is None:
+            record = NodeAvailabilityRecord(
+                node_type=node_type,
+                is_available=True,
+                last_checked_timestamp=timestamp,
+                unavailable_node_information=None,
+            )
+        else:
+            info = UnavailableNodeInformation(
+                category=node_launch_exception.category,
+                description=node_launch_exception.description,
+            )
+            record = NodeAvailabilityRecord(
+                node_type=node_type,
+                is_available=False,
+                last_checked_timestamp=timestamp,
+                unavailable_node_information=info,
+            )
+
+        expiration_time = timestamp + self.ttl
+
+        # TODO (Alex): In theory it would be nice to make this dictionary
+        # ordered by expiration time, unfortunately that's a bit difficult
+        # since `update_node_availability` can be called with out of order
+        # timestamps.
+        self.store[node_type] = (expiration_time, record)
+
+        self._remove_old_entries()
+
+    def update_node_availability(
+        self,
+        node_type: str,
+        timestamp: int,
+        node_launch_exception: Optional[NodeLaunchException],
+    ) -> None:
+        """
+        Update the availability and details of a single ndoe type.
+
+        Args:
+          node_type: The node type.
+          timestamp: The timestamp that this information is accurate as of.
+          node_launch_exception: Details about why the node launch failed. If
+            empty, the node type will be considered available."""
+        with self.lock:
+            self._update_node_availability_requires_lock(
+                node_type, timestamp, node_launch_exception
+            )
+
+    def summary(self) -> NodeAvailabilitySummary:
+        """
+        Returns a summary of node availabilities and their staleness.
+
+        Returns
+            A summary of node availabilities and their staleness.
+        """
+        with self.lock:
+            self._remove_old_entries()
+            return NodeAvailabilitySummary(
+                {node_type: record for node_type, (_, record) in self.store.items()}
+            )
+
+    def _remove_old_entries(self):
+        """Remove any expired entries from the cache."""
+        cur_time = self.timer()
+        with self.lock:
+            for key, (expiration_time, _) in list(self.store.items()):
+                if expiration_time < cur_time:
+                    del self.store[key]