wehub-resource-sync
132 lines
3.7 KiB
Python
132 lines
3.7 KiB
Python
import ray
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray.autoscaler.v2.sdk import get_cluster_status
|
|
import time
|
|
from logger import logger
|
|
from typing import Dict
|
|
|
|
ray.init("auto")
|
|
|
|
# Sync with the compute config.
|
|
HEAD_NODE_CPU = 0
|
|
WORKER_NODE_CPU = 4
|
|
IDLE_TERMINATION_S = 60 * 5 # 5 min
|
|
DEFAULT_RETRY_INTERVAL_MS = 15 * 1000 # 15 sec
|
|
|
|
|
|
def check_cluster(target_num_nodes: int, target_resources: Dict[str, float]):
|
|
gcs_address = ray.get_runtime_context().gcs_address
|
|
cluster_status = get_cluster_status(gcs_address)
|
|
|
|
assert (
|
|
len(cluster_status.active_nodes) + len(cluster_status.idle_nodes)
|
|
) == target_num_nodes
|
|
|
|
for k, v in target_resources.items():
|
|
assert cluster_status.total_resources().get(k, 0) == v
|
|
|
|
return True
|
|
|
|
|
|
ctx = {
|
|
"num_cpus": 0,
|
|
"num_nodes": 1,
|
|
}
|
|
logger.info(f"Starting cluster with {ctx['num_nodes']} nodes, {ctx['num_cpus']} cpus")
|
|
check_cluster(
|
|
target_num_nodes=ctx["num_nodes"], target_resources={"CPU": ctx["num_cpus"]}
|
|
)
|
|
|
|
|
|
# Request for cluster resources
|
|
def test_request_cluster_resources(ctx: dict):
|
|
from ray.autoscaler._private.commands import request_resources
|
|
|
|
request_resources(num_cpus=8)
|
|
|
|
ctx["num_cpus"] += 8
|
|
ctx["num_nodes"] += 8 // WORKER_NODE_CPU
|
|
|
|
# Assert on number of worker nodes.
|
|
logger.info(
|
|
f"Requesting cluster constraints: {ctx['num_nodes']} nodes, "
|
|
f"{ctx['num_cpus']} cpus"
|
|
)
|
|
wait_for_condition(
|
|
check_cluster,
|
|
timeout=60 * 5, # 5min
|
|
retry_interval_ms=DEFAULT_RETRY_INTERVAL_MS,
|
|
target_num_nodes=ctx["num_nodes"],
|
|
target_resources={"CPU": ctx["num_cpus"]},
|
|
)
|
|
|
|
# Reset the cluster constraints.
|
|
request_resources(num_cpus=0)
|
|
|
|
ctx["num_cpus"] -= 8
|
|
ctx["num_nodes"] -= 8 // WORKER_NODE_CPU
|
|
logger.info(
|
|
f"Waiting for cluster go idle after constraint cleared: {ctx['num_nodes']} "
|
|
f"nodes, {ctx['num_cpus']} cpus"
|
|
)
|
|
wait_for_condition(
|
|
check_cluster,
|
|
timeout=60 + IDLE_TERMINATION_S, # 1min + idle timeout
|
|
retry_interval_ms=DEFAULT_RETRY_INTERVAL_MS,
|
|
target_num_nodes=ctx["num_nodes"],
|
|
target_resources={"CPU": ctx["num_cpus"]},
|
|
)
|
|
|
|
|
|
# Run actors/tasks that exceed the cluster resources should upscale the cluster
|
|
def test_run_tasks_concurrent(ctx: dict):
|
|
num_tasks = 2
|
|
num_actors = 2
|
|
|
|
@ray.remote(num_cpus=WORKER_NODE_CPU)
|
|
def f():
|
|
while True:
|
|
time.sleep(1)
|
|
|
|
@ray.remote(num_cpus=WORKER_NODE_CPU)
|
|
class Actor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
tasks = [f.remote() for _ in range(num_tasks)]
|
|
actors = [Actor.remote() for _ in range(num_actors)]
|
|
|
|
ctx["num_cpus"] += (num_tasks + num_actors) * WORKER_NODE_CPU
|
|
ctx["num_nodes"] += num_tasks + num_actors
|
|
|
|
logger.info(f"Waiting for {ctx['num_nodes']} nodes, {ctx['num_cpus']} cpus")
|
|
wait_for_condition(
|
|
check_cluster,
|
|
timeout=60 * 5, # 5min
|
|
retry_interval_ms=DEFAULT_RETRY_INTERVAL_MS,
|
|
target_num_nodes=ctx["num_nodes"],
|
|
target_resources={"CPU": ctx["num_cpus"]},
|
|
)
|
|
|
|
[ray.cancel(task) for task in tasks]
|
|
[ray.kill(actor) for actor in actors]
|
|
|
|
ctx["num_cpus"] -= (num_actors + num_tasks) * WORKER_NODE_CPU
|
|
ctx["num_nodes"] -= num_actors + num_tasks
|
|
|
|
logger.info(
|
|
f"Waiting for cluster to scale down to {ctx['num_nodes']} nodes, "
|
|
f"{ctx['num_cpus']} cpus"
|
|
)
|
|
wait_for_condition(
|
|
check_cluster,
|
|
timeout=60 + IDLE_TERMINATION_S,
|
|
retry_interval_ms=DEFAULT_RETRY_INTERVAL_MS,
|
|
target_num_nodes=ctx["num_nodes"],
|
|
target_resources={"CPU": ctx["num_cpus"]},
|
|
)
|
|
|
|
|
|
test_request_cluster_resources(ctx)
|
|
test_run_tasks_concurrent(ctx)
|