chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,530 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user