wehub-resource-sync
531 lines
16 KiB
Python
531 lines
16 KiB
Python
"""Test utilities for Ray.
|
|
|
|
This module contains test utility classes that are distributed with the Ray package
|
|
and can be used by external libraries and tests. These utilities must remain in
|
|
_common/ (not in tests/) to be accessible in the Ray package distribution.
|
|
"""
|
|
|
|
import asyncio
|
|
import inspect
|
|
import logging
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
import traceback
|
|
import uuid
|
|
from collections import defaultdict
|
|
from collections.abc import Awaitable
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from typing import Any, Callable, Dict, Iterator, List, Optional, Set
|
|
|
|
import ray
|
|
import ray._common.usage.usage_lib as ray_usage_lib
|
|
import ray._private.utils
|
|
from ray._common.network_utils import build_address
|
|
from ray._common.utils import decode
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
try:
|
|
from prometheus_client.core import Metric
|
|
from prometheus_client.parser import Sample, text_string_to_metric_families
|
|
except (ImportError, ModuleNotFoundError):
|
|
Metric = None
|
|
Sample = None
|
|
|
|
def text_string_to_metric_families(*args, **kwargs):
|
|
raise ModuleNotFoundError("`prometheus_client` not found")
|
|
|
|
|
|
@ray.remote(num_cpus=0)
|
|
class SignalActor:
|
|
"""A Ray actor for coordinating test execution through signals.
|
|
|
|
Useful for testing async coordination, waiting for specific states,
|
|
and synchronizing multiple actors or tasks in tests.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.ready_event = asyncio.Event()
|
|
self.num_waiters = 0
|
|
|
|
def send(self, clear: bool = False):
|
|
self.ready_event.set()
|
|
if clear:
|
|
self.ready_event.clear()
|
|
|
|
async def wait(self, should_wait: bool = True):
|
|
if should_wait:
|
|
self.num_waiters += 1
|
|
await self.ready_event.wait()
|
|
self.num_waiters -= 1
|
|
|
|
async def cur_num_waiters(self) -> int:
|
|
return self.num_waiters
|
|
|
|
|
|
@ray.remote(num_cpus=0)
|
|
class Semaphore:
|
|
"""A Ray actor implementing a semaphore for test coordination.
|
|
|
|
Useful for testing resource limiting, concurrency control,
|
|
and coordination between multiple actors or tasks.
|
|
"""
|
|
|
|
def __init__(self, value: int = 1):
|
|
self._sema = asyncio.Semaphore(value=value)
|
|
|
|
async def acquire(self):
|
|
await self._sema.acquire()
|
|
|
|
async def release(self):
|
|
self._sema.release()
|
|
|
|
async def locked(self) -> bool:
|
|
return self._sema.locked()
|
|
|
|
|
|
__all__ = ["SignalActor", "Semaphore"]
|
|
|
|
|
|
def wait_for_condition(
|
|
condition_predictor: Callable[..., bool],
|
|
timeout: float = 10,
|
|
retry_interval_ms: float = 100,
|
|
raise_exceptions: bool = False,
|
|
**kwargs: Any,
|
|
):
|
|
"""Wait until a condition is met or time out with an exception.
|
|
|
|
Args:
|
|
condition_predictor: A function that predicts the condition.
|
|
timeout: Maximum timeout in seconds.
|
|
retry_interval_ms: Retry interval in milliseconds.
|
|
raise_exceptions: If true, exceptions that occur while executing
|
|
condition_predictor won't be caught and instead will be raised.
|
|
**kwargs: Arguments to pass to the condition_predictor.
|
|
|
|
Returns:
|
|
None: Returns when the condition is met.
|
|
|
|
Raises:
|
|
RuntimeError: If the condition is not met before the timeout expires.
|
|
"""
|
|
start = time.monotonic()
|
|
last_ex = None
|
|
while time.monotonic() - start <= timeout:
|
|
try:
|
|
if condition_predictor(**kwargs):
|
|
return
|
|
except Exception:
|
|
if raise_exceptions:
|
|
raise
|
|
last_ex = ray._private.utils.format_error_message(traceback.format_exc())
|
|
time.sleep(retry_interval_ms / 1000.0)
|
|
message = "The condition wasn't met before the timeout expired."
|
|
if last_ex is not None:
|
|
message += f" Last exception: {last_ex}"
|
|
raise RuntimeError(message)
|
|
|
|
|
|
async def async_wait_for_condition(
|
|
condition_predictor: Callable[..., Awaitable[bool]],
|
|
timeout: float = 10,
|
|
retry_interval_ms: float = 100,
|
|
**kwargs: Any,
|
|
):
|
|
"""Wait until a condition is met or time out with an exception.
|
|
|
|
Args:
|
|
condition_predictor: A function that predicts the condition.
|
|
timeout: Maximum timeout in seconds.
|
|
retry_interval_ms: Retry interval in milliseconds.
|
|
**kwargs: Arguments to pass to the condition_predictor.
|
|
|
|
Returns:
|
|
None: Returns when the condition is met.
|
|
|
|
Raises:
|
|
RuntimeError: If the condition is not met before the timeout expires.
|
|
"""
|
|
start = time.monotonic()
|
|
last_ex = None
|
|
while time.monotonic() - start <= timeout:
|
|
try:
|
|
if inspect.iscoroutinefunction(condition_predictor):
|
|
if await condition_predictor(**kwargs):
|
|
return
|
|
else:
|
|
if condition_predictor(**kwargs):
|
|
return
|
|
except Exception as ex:
|
|
last_ex = ex
|
|
await asyncio.sleep(retry_interval_ms / 1000.0)
|
|
message = "The condition wasn't met before the timeout expired."
|
|
if last_ex is not None:
|
|
message += f" Last exception: {last_ex}"
|
|
raise RuntimeError(message)
|
|
|
|
|
|
@contextmanager
|
|
def simulate_s3_bucket(
|
|
port: int = 5002,
|
|
region: str = "us-west-2",
|
|
) -> Iterator[str]:
|
|
"""Context manager that simulates an S3 bucket and yields the URI.
|
|
|
|
Args:
|
|
port: The port of the localhost endpoint where S3 is being served.
|
|
region: The S3 region.
|
|
|
|
Yields:
|
|
str: URI for the simulated S3 bucket.
|
|
"""
|
|
from moto.server import ThreadedMotoServer
|
|
|
|
old_env = os.environ
|
|
os.environ["AWS_ACCESS_KEY_ID"] = "testing"
|
|
os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
|
|
os.environ["AWS_SECURITY_TOKEN"] = "testing"
|
|
os.environ["AWS_SESSION_TOKEN"] = "testing"
|
|
|
|
s3_server = f"http://{build_address('localhost', port)}"
|
|
server = ThreadedMotoServer(port=port)
|
|
server.start()
|
|
url = f"s3://{uuid.uuid4().hex}?region={region}&endpoint_override={s3_server}"
|
|
yield url
|
|
server.stop()
|
|
os.environ = old_env
|
|
|
|
|
|
class TelemetryCallsite(Enum):
|
|
DRIVER = "driver"
|
|
ACTOR = "actor"
|
|
TASK = "task"
|
|
|
|
|
|
def _get_library_usages() -> Set[str]:
|
|
return set(
|
|
ray_usage_lib.get_library_usages_to_report(
|
|
ray.experimental.internal_kv.internal_kv_get_gcs_client()
|
|
)
|
|
)
|
|
|
|
|
|
def _get_extra_usage_tags() -> Dict[str, str]:
|
|
return ray_usage_lib.get_extra_usage_tags_to_report(
|
|
ray.experimental.internal_kv.internal_kv_get_gcs_client()
|
|
)
|
|
|
|
|
|
def check_library_usage_telemetry(
|
|
use_lib_fn: Callable[[], None],
|
|
*,
|
|
callsite: TelemetryCallsite,
|
|
expected_library_usages: List[Set[str]],
|
|
expected_extra_usage_tags: Optional[Dict[str, str]] = None,
|
|
):
|
|
"""Helper for writing tests to validate library usage telemetry.
|
|
|
|
`use_lib_fn` is a callable that will be called from the provided callsite.
|
|
After calling it, the telemetry data to export will be validated against
|
|
expected_library_usages and expected_extra_usage_tags.
|
|
"""
|
|
assert len(_get_library_usages()) == 0, _get_library_usages()
|
|
|
|
if callsite == TelemetryCallsite.DRIVER:
|
|
use_lib_fn()
|
|
elif callsite == TelemetryCallsite.ACTOR:
|
|
|
|
@ray.remote
|
|
class A:
|
|
def __init__(self):
|
|
use_lib_fn()
|
|
|
|
a = A.remote()
|
|
ray.get(a.__ray_ready__.remote())
|
|
elif callsite == TelemetryCallsite.TASK:
|
|
|
|
@ray.remote
|
|
def f():
|
|
use_lib_fn()
|
|
|
|
ray.get(f.remote())
|
|
else:
|
|
assert False, f"Unrecognized callsite: {callsite}"
|
|
|
|
library_usages = _get_library_usages()
|
|
extra_usage_tags = _get_extra_usage_tags()
|
|
|
|
assert library_usages in expected_library_usages, library_usages
|
|
if expected_extra_usage_tags:
|
|
assert all(
|
|
[extra_usage_tags[k] == v for k, v in expected_extra_usage_tags.items()]
|
|
), extra_usage_tags
|
|
|
|
|
|
class FakeTimer:
|
|
def __init__(self, start_time: Optional[float] = None):
|
|
self._lock = threading.Lock()
|
|
self.reset(start_time=start_time)
|
|
|
|
def reset(self, start_time: Optional[float] = None):
|
|
with self._lock:
|
|
if start_time is None:
|
|
start_time = time.time()
|
|
self._curr = start_time
|
|
|
|
def time(self) -> float:
|
|
return self._curr
|
|
|
|
def advance(self, by: float):
|
|
with self._lock:
|
|
self._curr += by
|
|
|
|
def realistic_sleep(self, amt: float):
|
|
with self._lock:
|
|
self._curr += amt + 0.001
|
|
|
|
|
|
def is_named_tuple(cls):
|
|
"""Return True if cls is a namedtuple and False otherwise."""
|
|
b = cls.__bases__
|
|
if len(b) != 1 or b[0] is not tuple:
|
|
return False
|
|
f = getattr(cls, "_fields", None)
|
|
if not isinstance(f, tuple):
|
|
return False
|
|
return all(type(n) is str for n in f)
|
|
|
|
|
|
def assert_tensors_equivalent(obj1, obj2):
|
|
"""
|
|
Recursively compare objects with special handling for torch.Tensor.
|
|
|
|
Tensors are considered equivalent if:
|
|
- Same dtype and shape
|
|
- Same device type (e.g., both 'cpu' or both 'cuda'), index ignored
|
|
- Values are equal (or close for floats)
|
|
"""
|
|
import torch
|
|
|
|
if isinstance(obj1, torch.Tensor) and isinstance(obj2, torch.Tensor):
|
|
# 1. dtype
|
|
assert obj1.dtype == obj2.dtype, f"dtype mismatch: {obj1.dtype} vs {obj2.dtype}"
|
|
# 2. shape
|
|
assert obj1.shape == obj2.shape, f"shape mismatch: {obj1.shape} vs {obj2.shape}"
|
|
# 3. device type must match (cpu/cpu or cuda/cuda), ignore index
|
|
assert (
|
|
obj1.device.type == obj2.device.type
|
|
), f"Device type mismatch: {obj1.device} vs {obj2.device}"
|
|
|
|
# 4. Compare values safely on CPU
|
|
t1_cpu = obj1.cpu()
|
|
t2_cpu = obj2.cpu()
|
|
if obj1.dtype.is_floating_point or obj1.dtype.is_complex:
|
|
assert torch.allclose(
|
|
t1_cpu, t2_cpu, atol=1e-6, rtol=1e-5
|
|
), "Floating-point tensors not close"
|
|
else:
|
|
assert torch.equal(t1_cpu, t2_cpu), "Integer/bool tensors not equal"
|
|
return
|
|
|
|
# Type must match
|
|
if type(obj1) is not type(obj2):
|
|
raise AssertionError(f"Type mismatch: {type(obj1)} vs {type(obj2)}")
|
|
|
|
# Handle namedtuples
|
|
if is_named_tuple(type(obj1)):
|
|
assert len(obj1) == len(obj2)
|
|
for a, b in zip(obj1, obj2):
|
|
assert_tensors_equivalent(a, b)
|
|
elif isinstance(obj1, dict):
|
|
assert obj1.keys() == obj2.keys()
|
|
for k in obj1:
|
|
assert_tensors_equivalent(obj1[k], obj2[k])
|
|
elif isinstance(obj1, (list, tuple)):
|
|
assert len(obj1) == len(obj2)
|
|
for a, b in zip(obj1, obj2):
|
|
assert_tensors_equivalent(a, b)
|
|
elif hasattr(obj1, "__dict__") and hasattr(obj2, "__dict__"):
|
|
# Compare user-defined objects by their public attributes
|
|
keys1 = {
|
|
k
|
|
for k in obj1.__dict__.keys()
|
|
if not k.startswith("_ray_") and k != "_pytype_"
|
|
}
|
|
keys2 = {
|
|
k
|
|
for k in obj2.__dict__.keys()
|
|
if not k.startswith("_ray_") and k != "_pytype_"
|
|
}
|
|
assert keys1 == keys2, f"Object attribute keys differ: {keys1} vs {keys2}"
|
|
for k in keys1:
|
|
assert_tensors_equivalent(obj1.__dict__[k], obj2.__dict__[k])
|
|
else:
|
|
# Fallback for primitives: int, float, str, bool, etc.
|
|
assert obj1 == obj2, f"Non-tensor values differ: {obj1} vs {obj2}"
|
|
|
|
|
|
def run_string_as_driver(
|
|
driver_script: str, env: Dict = None, encode: str = "utf-8"
|
|
) -> str:
|
|
"""Run a driver as a separate process.
|
|
|
|
Args:
|
|
driver_script: A string to run as a Python script.
|
|
env: The environment variables for the driver.
|
|
encode: The encoding to use for the driver script.
|
|
|
|
Returns:
|
|
The script's output.
|
|
"""
|
|
proc = subprocess.Popen(
|
|
[sys.executable, "-"],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
env=env,
|
|
)
|
|
with proc:
|
|
output = proc.communicate(driver_script.encode(encoding=encode))[0]
|
|
if proc.returncode:
|
|
print(decode(output, encode_type=encode))
|
|
logger.error(proc.stderr)
|
|
raise subprocess.CalledProcessError(
|
|
proc.returncode, proc.args, output, proc.stderr
|
|
)
|
|
out = decode(output, encode_type=encode)
|
|
return out
|
|
|
|
|
|
@dataclass
|
|
class MetricSamplePattern:
|
|
name: Optional[str] = None
|
|
value: Optional[str] = None
|
|
partial_label_match: Optional[Dict[str, str]] = None
|
|
|
|
def matches(self, sample: "Sample"):
|
|
if self.name is not None:
|
|
if self.name != sample.name:
|
|
return False
|
|
|
|
if self.value is not None:
|
|
if self.value != sample.value:
|
|
return False
|
|
|
|
if self.partial_label_match is not None:
|
|
for label, value in self.partial_label_match.items():
|
|
if sample.labels.get(label) != value:
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
@dataclass
|
|
class PrometheusTimeseries:
|
|
"""A collection of timeseries from multiple addresses. Each timeseries is a
|
|
collection of samples with the same metric name and labels. Concretely:
|
|
- components_dict: a dictionary of addresses to the Component labels
|
|
- metric_descriptors: a dictionary of metric names to the Metric object
|
|
- metric_samples: the latest value of each label
|
|
"""
|
|
|
|
components_dict: Dict[str, Set[str]] = field(default_factory=dict)
|
|
metric_descriptors: Dict[str, "Metric"] = field(default_factory=dict)
|
|
metric_samples: Dict[frozenset, "Sample"] = field(default_factory=dict)
|
|
|
|
def flush(self):
|
|
self.components_dict.clear()
|
|
self.metric_descriptors.clear()
|
|
self.metric_samples.clear()
|
|
|
|
|
|
def fetch_raw_prometheus(prom_addresses, timeout=None):
|
|
# Local import so minimal dependency tests can run without requests
|
|
import requests
|
|
|
|
for address in prom_addresses:
|
|
try:
|
|
kwargs = {} if timeout is None else {"timeout": timeout}
|
|
response = requests.get(f"http://{address}/metrics", **kwargs)
|
|
yield address, response.text
|
|
except requests.exceptions.ConnectionError:
|
|
continue
|
|
except requests.exceptions.Timeout:
|
|
continue
|
|
|
|
|
|
def fetch_prometheus(prom_addresses, timeout=None):
|
|
components_dict = {}
|
|
metric_descriptors = {}
|
|
metric_samples = []
|
|
|
|
for address in prom_addresses:
|
|
if address not in components_dict:
|
|
components_dict[address] = set()
|
|
|
|
for address, response in fetch_raw_prometheus(prom_addresses, timeout=timeout):
|
|
for metric in text_string_to_metric_families(response):
|
|
for sample in metric.samples:
|
|
metric_descriptors[sample.name] = metric
|
|
metric_samples.append(sample)
|
|
if "Component" in sample.labels:
|
|
components_dict[address].add(sample.labels["Component"])
|
|
return components_dict, metric_descriptors, metric_samples
|
|
|
|
|
|
def fetch_prometheus_timeseries(
|
|
prom_addreses: List[str],
|
|
result: PrometheusTimeseries,
|
|
timeout=None,
|
|
) -> PrometheusTimeseries:
|
|
components_dict, metric_descriptors, metric_samples = fetch_prometheus(
|
|
prom_addreses, timeout=timeout
|
|
)
|
|
for address, components in components_dict.items():
|
|
if address not in result.components_dict:
|
|
result.components_dict[address] = set()
|
|
result.components_dict[address].update(components)
|
|
result.metric_descriptors.update(metric_descriptors)
|
|
for sample in metric_samples:
|
|
# udpate sample to the latest value
|
|
result.metric_samples[
|
|
frozenset(list(sample.labels.items()) + [("_metric_name_", sample.name)])
|
|
] = sample
|
|
return result
|
|
|
|
|
|
def fetch_prometheus_metrics(prom_addresses: List[str]) -> Dict[str, List[Any]]:
|
|
"""Return prometheus metrics from the given addresses.
|
|
|
|
Args:
|
|
prom_addresses: List of metrics_agent addresses to collect metrics from.
|
|
|
|
Returns:
|
|
Dict mapping from metric name to list of samples for the metric.
|
|
"""
|
|
_, _, samples = fetch_prometheus(prom_addresses)
|
|
samples_by_name = defaultdict(list)
|
|
for sample in samples:
|
|
samples_by_name[sample.name].append(sample)
|
|
return samples_by_name
|
|
|
|
|
|
def fetch_prometheus_metric_timeseries(
|
|
prom_addresses: List[str],
|
|
result: PrometheusTimeseries,
|
|
timeout=None,
|
|
) -> Dict[str, List[Any]]:
|
|
samples = fetch_prometheus_timeseries(
|
|
prom_addresses, result, timeout=timeout
|
|
).metric_samples.values()
|
|
samples_by_name = defaultdict(list)
|
|
for sample in samples:
|
|
samples_by_name[sample.name].append(sample)
|
|
return samples_by_name
|