chore: import upstream snapshot with attribution

python/ray/autoscaler/v2/sdk.py

@@ -0,0 +1,121 @@
+import time
+from collections import Counter
+from typing import List, NamedTuple
+
+from ray._raylet import GcsClient
+from ray.autoscaler.v2.schema import ClusterStatus, Stats
+from ray.autoscaler.v2.utils import ClusterStatusParser
+from ray.core.generated.autoscaler_pb2 import (
+    ClusterResourceState,
+    GetClusterResourceStateReply,
+    GetClusterStatusReply,
+)
+
+DEFAULT_RPC_TIMEOUT_S = 10
+
+
+class ResourceRequest(NamedTuple):
+    resources: dict
+    label_selector: dict
+
+
+def request_cluster_resources(
+    gcs_address: str,
+    to_request: List[dict],
+    timeout: int = DEFAULT_RPC_TIMEOUT_S,
+):
+    """Request resources from the autoscaler.
+
+    This will add a cluster resource constraint to GCS. GCS will asynchronously
+    pass the constraint to the autoscaler, and the autoscaler will try to provision the
+    requested minimal bundles in `to_request`.
+
+    If the cluster already has `to_request` resources, this will be an no-op.
+    Future requests submitted through this API will overwrite the previous requests.
+
+    Args:
+        gcs_address: The GCS address to query.
+        to_request: A list of resource requests to request the cluster to have.
+            Each resource request is a tuple of resources and a label_selector
+            to apply per-bundle. e.g.: [{"resources": {"CPU": 1, "GPU": 1}, "label_selector": {"accelerator-type": "A100"}}]
+        timeout: Timeout in seconds for the request to be timeout
+
+    """
+    assert len(gcs_address) > 0, "GCS address is not specified."
+
+    # Convert bundle dicts to ResourceRequest tuples.
+    normalized: List[ResourceRequest] = []
+    for r in to_request:
+        assert isinstance(
+            r, dict
+        ), f"Internal Error: Expected a dict, but got {type(r)}"
+        resources = r.get("resources", {})
+        selector = r.get("label_selector", {})
+        normalized.append(ResourceRequest(resources, selector))
+
+    to_request = normalized
+
+    # Aggregate bundle by shape
+    def keyfunc(r):
+        return (
+            frozenset(r.resources.items()),
+            frozenset(r.label_selector.items()),
+        )
+
+    grouped_requests = Counter(keyfunc(r) for r in to_request)
+
+    bundles: List[dict] = []
+    label_selectors: List[dict] = []
+    counts: List[int] = []
+
+    for (bundle, selector), count in grouped_requests.items():
+        bundles.append(dict(bundle))
+        label_selectors.append(dict(selector))
+        counts.append(count)
+
+    GcsClient(gcs_address).request_cluster_resource_constraint(
+        bundles, label_selectors, counts, timeout_s=timeout
+    )
+
+
+def get_cluster_status(
+    gcs_address: str, timeout: int = DEFAULT_RPC_TIMEOUT_S
+) -> ClusterStatus:
+    """
+    Get the cluster status from the autoscaler.
+
+    Args:
+        gcs_address: The GCS address to query.
+        timeout: Timeout in seconds for the request to be timeout
+
+    Returns:
+        A ClusterStatus object.
+    """
+    assert len(gcs_address) > 0, "GCS address is not specified."
+    req_time = time.time()
+    str_reply = GcsClient(gcs_address).get_cluster_status(timeout_s=timeout)
+    reply_time = time.time()
+    reply = GetClusterStatusReply()
+    reply.ParseFromString(str_reply)
+
+    # TODO(rickyx): To be more accurate, we could add a timestamp field from the reply.
+    return ClusterStatusParser.from_get_cluster_status_reply(
+        reply,
+        stats=Stats(gcs_request_time_s=reply_time - req_time, request_ts_s=req_time),
+    )
+
+
+def get_cluster_resource_state(gcs_client: GcsClient) -> ClusterResourceState:
+    """
+    Get the cluster resource state from GCS.
+    Args:
+        gcs_client: The GCS client to query.
+    Returns:
+        A ClusterResourceState object
+    Raises:
+        Exception: If the request times out or failed.
+    """
+    str_reply = gcs_client.get_cluster_resource_state()
+    reply = GetClusterResourceStateReply()
+    reply.ParseFromString(str_reply)
+    return reply.cluster_resource_state