wehub-resource-sync
122 lines
3.8 KiB
Python
122 lines
3.8 KiB
Python
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
|