wehub-resource-sync
180 lines
6.5 KiB
Python
180 lines
6.5 KiB
Python
import csv
|
|
from collections import namedtuple
|
|
from typing import Dict, Optional, Tuple
|
|
|
|
from ray_release.bazel import bazel_runfile
|
|
from ray_release.logger import logger
|
|
from ray_release.template import load_test_cluster_compute
|
|
from ray_release.test import Test
|
|
|
|
# Keep 10% for the buffer.
|
|
limit = int(15784 * 0.9)
|
|
|
|
|
|
Condition = namedtuple(
|
|
"Condition", ["min_gpu", "max_gpu", "min_cpu", "max_cpu", "group", "limit"]
|
|
)
|
|
|
|
aws_gpu_cpu_to_concurrency_groups = [
|
|
Condition(min_gpu=9, max_gpu=-1, min_cpu=0, max_cpu=-1, group="large-gpu", limit=4),
|
|
Condition(
|
|
min_gpu=1, max_gpu=9, min_cpu=0, max_cpu=-128, group="small-gpu", limit=8
|
|
),
|
|
Condition(
|
|
min_gpu=0, max_gpu=0, min_cpu=1025, max_cpu=-1, group="enormous", limit=1
|
|
),
|
|
Condition(min_gpu=0, max_gpu=0, min_cpu=513, max_cpu=1024, group="large", limit=8),
|
|
Condition(min_gpu=0, max_gpu=0, min_cpu=129, max_cpu=512, group="medium", limit=6),
|
|
Condition(min_gpu=0, max_gpu=0, min_cpu=9, max_cpu=32, group="tiny", limit=32),
|
|
Condition(min_gpu=0, max_gpu=0, min_cpu=0, max_cpu=8, group="minuscule", limit=128),
|
|
# Make sure "small" is the last in the list, because it is the fallback.
|
|
Condition(min_gpu=0, max_gpu=0, min_cpu=0, max_cpu=128, group="small", limit=16),
|
|
]
|
|
|
|
gce_gpu_cpu_to_concurrent_groups = [
|
|
Condition(min_gpu=8, max_gpu=-1, min_cpu=0, max_cpu=-1, group="gpu-gce", limit=4),
|
|
Condition(min_gpu=4, max_gpu=-1, min_cpu=0, max_cpu=-1, group="gpu-gce", limit=8),
|
|
Condition(min_gpu=2, max_gpu=-1, min_cpu=0, max_cpu=-1, group="gpu-gce", limit=16),
|
|
Condition(min_gpu=1, max_gpu=-1, min_cpu=0, max_cpu=-1, group="gpu-gce", limit=32),
|
|
Condition(
|
|
min_gpu=0, max_gpu=0, min_cpu=1025, max_cpu=-1, group="enormous-gce", limit=1
|
|
),
|
|
Condition(
|
|
min_gpu=0, max_gpu=0, min_cpu=513, max_cpu=1024, group="large-gce", limit=8
|
|
),
|
|
Condition(
|
|
min_gpu=0, max_gpu=0, min_cpu=129, max_cpu=512, group="medium-gce", limit=6
|
|
),
|
|
Condition(min_gpu=0, max_gpu=0, min_cpu=9, max_cpu=32, group="tiny-gce", limit=32),
|
|
Condition(
|
|
min_gpu=0, max_gpu=0, min_cpu=0, max_cpu=8, group="minuscule-gce", limit=128
|
|
),
|
|
Condition(
|
|
min_gpu=0, max_gpu=0, min_cpu=0, max_cpu=128, group="small-gce", limit=16
|
|
),
|
|
]
|
|
|
|
|
|
# Obtained from https://cloud.google.com/compute/docs/accelerator-optimized-machines
|
|
gcp_gpu_instances = {
|
|
"a2-highgpu-1g": (12, 1),
|
|
"a2-highgpu-2g": (24, 2),
|
|
"a2-highgpu-4g": (48, 4),
|
|
"a2-highgpu-8g": (96, 8),
|
|
"a2-megagpu-16g": (96, 16),
|
|
"n1-standard-16-nvidia-tesla-t4-1": (16, 1),
|
|
"n1-standard-64-nvidia-tesla-t4-4": (64, 4),
|
|
"n1-standard-32-nvidia-tesla-t4-2": (32, 2),
|
|
"n1-highmem-64-nvidia-tesla-v100-8": (64, 8),
|
|
"n1-highmem-96-nvidia-tesla-v100-8": (96, 8),
|
|
}
|
|
|
|
|
|
def load_instance_types(path: Optional[str] = None) -> Dict[str, Tuple[int, int]]:
|
|
if not path:
|
|
path = bazel_runfile(
|
|
"release/ray_release/buildkite/aws_instance_types.csv",
|
|
)
|
|
|
|
instance_to_resources = {}
|
|
with open(path, "rt") as fp:
|
|
reader = csv.DictReader(fp)
|
|
for row in reader:
|
|
instance_to_resources[row["instance"]] = (
|
|
int(row["cpus"]),
|
|
int(row["gpus"]),
|
|
)
|
|
|
|
return instance_to_resources
|
|
|
|
|
|
def parse_instance_resources(instance: str) -> Tuple[int, int]:
|
|
"""Parse (GCP) instance strings to resources"""
|
|
# Assumes that GPU instances have already been parsed
|
|
num_cpus = int(instance.split("-")[-1])
|
|
num_gpus = 0
|
|
return num_cpus, num_gpus
|
|
|
|
|
|
def parse_condition(cond: int, limit: float = float("inf")) -> float:
|
|
return cond if cond > -1 else limit
|
|
|
|
|
|
def get_concurrency_group(test: Test) -> Tuple[str, int]:
|
|
if test.get("env", None) == "gce":
|
|
concurrent_group = gce_gpu_cpu_to_concurrent_groups
|
|
else:
|
|
concurrent_group = aws_gpu_cpu_to_concurrency_groups
|
|
default_concurrent = concurrent_group[-1]
|
|
try:
|
|
test_cpus, test_gpus = get_test_resources(test)
|
|
except Exception as e:
|
|
logger.warning(f"Couldn't get test resources for test {test['name']}: {e}")
|
|
return default_concurrent.group, default_concurrent.limit
|
|
|
|
for condition in concurrent_group:
|
|
min_gpu = parse_condition(condition.min_gpu, float("-inf"))
|
|
max_gpu = parse_condition(condition.max_gpu, float("inf"))
|
|
min_cpu = parse_condition(condition.min_cpu, float("-inf"))
|
|
max_cpu = parse_condition(condition.max_cpu, float("inf"))
|
|
|
|
if min_cpu <= test_cpus <= max_cpu and min_gpu <= test_gpus <= max_gpu:
|
|
return condition.group, condition.limit
|
|
|
|
# Return default
|
|
logger.warning(
|
|
f"Could not find concurrency group for test {test['name']} "
|
|
f"based on used resources."
|
|
)
|
|
return default_concurrent.group, default_concurrent.limit
|
|
|
|
|
|
def get_test_resources(test: Test) -> Tuple[int, int]:
|
|
cluster_compute = load_test_cluster_compute(test)
|
|
return get_test_resources_from_cluster_compute(
|
|
cluster_compute, is_new_schema=test.uses_anyscale_sdk_2026()
|
|
)
|
|
|
|
|
|
def get_test_resources_from_cluster_compute(
|
|
cluster_compute: Dict, is_new_schema: bool = False
|
|
) -> Tuple[int, int]:
|
|
instances = []
|
|
|
|
if is_new_schema:
|
|
# New schema: head_node, worker_nodes, min_nodes/max_nodes
|
|
head_node = cluster_compute.get("head_node", {})
|
|
if head_node.get("instance_type"):
|
|
instances.append((head_node["instance_type"], 1))
|
|
|
|
for w in cluster_compute.get("worker_nodes", []):
|
|
if w.get("instance_type"):
|
|
instances.append(
|
|
(w["instance_type"], w.get("max_nodes", w.get("min_nodes", 1)))
|
|
)
|
|
else:
|
|
# Legacy schema: head_node_type, worker_node_types, min_workers/max_workers
|
|
instances.append((cluster_compute["head_node_type"]["instance_type"], 1))
|
|
|
|
instances.extend(
|
|
(w["instance_type"], w.get("max_workers", w.get("min_workers", 1)))
|
|
for w in cluster_compute["worker_node_types"]
|
|
)
|
|
|
|
aws_instance_types = load_instance_types()
|
|
total_cpus = 0
|
|
total_gpus = 0
|
|
|
|
for instance, count in instances:
|
|
if instance in aws_instance_types:
|
|
instance_cpus, instance_gpus = aws_instance_types[instance]
|
|
elif instance in gcp_gpu_instances:
|
|
instance_cpus, instance_gpus = gcp_gpu_instances[instance]
|
|
else:
|
|
instance_cpus, instance_gpus = parse_instance_resources(instance)
|
|
|
|
total_cpus += instance_cpus * count
|
|
total_gpus += instance_gpus * count
|
|
|
|
return total_cpus, total_gpus
|