Files
2026-07-13 13:17:40 +08:00

122 lines
3.4 KiB
Python

import os
import time
import tqdm
from many_nodes_tests.dashboard_test import DashboardTestAtScale
import ray
import ray._common.test_utils
import ray._private.test_utils as test_utils
from ray.util.placement_group import placement_group, remove_placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
is_smoke_test = True
if "SMOKE_TEST" in os.environ:
MAX_PLACEMENT_GROUPS = 20
else:
MAX_PLACEMENT_GROUPS = 1000
is_smoke_test = False
def test_many_placement_groups():
# @ray.remote(num_cpus=1, resources={"node": 0.02})
@ray.remote
class C1:
def ping(self):
return "pong"
# @ray.remote(num_cpus=1)
@ray.remote
class C2:
def ping(self):
return "pong"
# @ray.remote(resources={"node": 0.02})
@ray.remote
class C3:
def ping(self):
return "pong"
bundle1 = {"node": 0.02, "CPU": 1}
bundle2 = {"CPU": 1}
bundle3 = {"node": 0.02}
pgs = []
for _ in tqdm.trange(MAX_PLACEMENT_GROUPS, desc="Creating pgs"):
pg = placement_group(bundles=[bundle1, bundle2, bundle3])
pgs.append(pg)
for pg in tqdm.tqdm(pgs, desc="Waiting for pgs to be ready"):
ray.get(pg.ready())
actors = []
for pg in tqdm.tqdm(pgs, desc="Scheduling tasks"):
actors.append(
C1.options(
scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)
).remote()
)
actors.append(
C2.options(
scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)
).remote()
)
actors.append(
C3.options(
scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)
).remote()
)
not_ready = [actor.ping.remote() for actor in actors]
for _ in tqdm.trange(len(actors)):
ready, not_ready = ray.wait(not_ready)
assert ray.get(*ready) == "pong"
for pg in tqdm.tqdm(pgs, desc="Cleaning up pgs"):
remove_placement_group(pg)
def no_resource_leaks():
return test_utils.no_resource_leaks_excluding_node_resources()
addr = ray.init(address="auto")
ray._common.test_utils.wait_for_condition(no_resource_leaks)
monitor_actor = test_utils.monitor_memory_usage()
dashboard_test = DashboardTestAtScale(addr)
start_time = time.time()
test_many_placement_groups()
end_time = time.time()
ray.get(monitor_actor.stop_run.remote())
used_gb, usage = ray.get(monitor_actor.get_peak_memory_info.remote())
print(f"Peak memory usage: {round(used_gb, 2)}GB")
print(f"Peak memory usage per processes:\n {usage}")
del monitor_actor
ray._common.test_utils.wait_for_condition(no_resource_leaks)
rate = MAX_PLACEMENT_GROUPS / (end_time - start_time)
print(
f"Success! Started {MAX_PLACEMENT_GROUPS} pgs in "
f"{end_time - start_time}s. ({rate} pgs/s)"
)
results = {
"pgs_per_second": rate,
"num_pgs": MAX_PLACEMENT_GROUPS,
"time": end_time - start_time,
"_peak_memory": round(used_gb, 2),
"_peak_process_memory": usage,
}
if not is_smoke_test:
results["perf_metrics"] = [
{
"perf_metric_name": "pgs_per_second",
"perf_metric_value": rate,
"perf_metric_type": "THROUGHPUT",
}
]
dashboard_test.update_release_test_result(results)
test_utils.safe_write_to_results_json(results)