wehub-resource-sync
70 lines
1.8 KiB
Python
70 lines
1.8 KiB
Python
import sys
|
|
from collections import defaultdict
|
|
from typing import Dict
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import (
|
|
PrometheusTimeseries,
|
|
wait_for_condition,
|
|
)
|
|
from ray._private.test_utils import raw_metric_timeseries
|
|
from ray._private.worker import RayContext
|
|
from ray.util.placement_group import remove_placement_group
|
|
|
|
_SYSTEM_CONFIG = {
|
|
"metrics_report_interval_ms": 200,
|
|
}
|
|
|
|
|
|
def groups_by_state(info: RayContext, timeseries: PrometheusTimeseries) -> Dict:
|
|
res = raw_metric_timeseries(info, timeseries)
|
|
info = defaultdict(int)
|
|
if "ray_placement_groups" in res:
|
|
for sample in res["ray_placement_groups"]:
|
|
info[sample.labels["State"]] += sample.value
|
|
for k, v in info.copy().items():
|
|
if v == 0:
|
|
del info[k]
|
|
print(f"Groups by state: {info}")
|
|
return info
|
|
|
|
|
|
def test_basic_states(shutdown_only):
|
|
info = ray.init(num_cpus=3, _system_config=_SYSTEM_CONFIG)
|
|
|
|
pg1 = ray.util.placement_group(bundles=[{"CPU": 1}])
|
|
pg2 = ray.util.placement_group(bundles=[{"CPU": 1}])
|
|
pg3 = ray.util.placement_group(bundles=[{"CPU": 4}])
|
|
ray.get([pg1.ready(), pg2.ready()])
|
|
|
|
expected = {
|
|
"CREATED": 2,
|
|
"PENDING": 1,
|
|
}
|
|
timeseries = PrometheusTimeseries()
|
|
wait_for_condition(
|
|
lambda: groups_by_state(info, timeseries) == expected,
|
|
timeout=20,
|
|
retry_interval_ms=500,
|
|
)
|
|
|
|
remove_placement_group(pg1)
|
|
remove_placement_group(pg2)
|
|
remove_placement_group(pg3)
|
|
|
|
expected = {
|
|
"REMOVED": 3,
|
|
}
|
|
timeseries = PrometheusTimeseries()
|
|
wait_for_condition(
|
|
lambda: groups_by_state(info, timeseries) == expected,
|
|
timeout=20,
|
|
retry_interval_ms=500,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|