chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,539 @@
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import logging
|
||||
import pprint
|
||||
import time
|
||||
import traceback
|
||||
from collections import defaultdict
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Callable, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
import ray
|
||||
import ray._common.test_utils as test_utils
|
||||
from ray._private.gcs_utils import GcsChannel
|
||||
from ray._raylet import GcsClient
|
||||
from ray.actor import ActorHandle
|
||||
from ray.dashboard.state_aggregator import (
|
||||
StateAPIManager,
|
||||
)
|
||||
from ray.util.state import list_tasks, list_workers
|
||||
from ray.util.state.common import (
|
||||
DEFAULT_LIMIT,
|
||||
DEFAULT_RPC_TIMEOUT,
|
||||
ListApiOptions,
|
||||
PredicateType,
|
||||
SupportedFilterType,
|
||||
)
|
||||
from ray.util.state.state_manager import StateDataSourceClient
|
||||
|
||||
import psutil
|
||||
|
||||
|
||||
@dataclass
|
||||
class StateAPIMetric:
|
||||
latency_sec: float
|
||||
result_size: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class StateAPICallSpec:
|
||||
api: Callable
|
||||
verify_cb: Callable
|
||||
kwargs: Dict = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class StateAPIStats:
|
||||
pending_calls: int = 0
|
||||
total_calls: int = 0
|
||||
calls: Dict = field(default_factory=lambda: defaultdict(list))
|
||||
|
||||
|
||||
GLOBAL_STATE_STATS = StateAPIStats()
|
||||
|
||||
STATE_LIST_LIMIT = int(1e6) # 1m
|
||||
STATE_LIST_TIMEOUT = 600 # 10min
|
||||
|
||||
|
||||
def invoke_state_api(
|
||||
verify_cb: Callable,
|
||||
state_api_fn: Callable,
|
||||
state_stats: StateAPIStats = GLOBAL_STATE_STATS,
|
||||
key_suffix: Optional[str] = None,
|
||||
print_result: Optional[bool] = False,
|
||||
err_msg: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Invoke a State API
|
||||
|
||||
Args:
|
||||
verify_cb: Callback that takes in the response from `state_api_fn` and
|
||||
returns a boolean, indicating the correctness of the results.
|
||||
state_api_fn: Function of the state API.
|
||||
state_stats: Stats container that tracks latency and call counts.
|
||||
key_suffix: Optional suffix appended to the stats key for this call.
|
||||
print_result: If True, pretty-print the API result.
|
||||
err_msg: Optional message included in the assertion error if the
|
||||
verify callback fails.
|
||||
**kwargs: Keyword arguments to be forwarded to the `state_api_fn`.
|
||||
|
||||
Returns:
|
||||
The response from ``state_api_fn``.
|
||||
"""
|
||||
if "timeout" not in kwargs:
|
||||
kwargs["timeout"] = STATE_LIST_TIMEOUT
|
||||
|
||||
# Suppress missing output warning
|
||||
kwargs["raise_on_missing_output"] = False
|
||||
|
||||
res = None
|
||||
try:
|
||||
state_stats.total_calls += 1
|
||||
state_stats.pending_calls += 1
|
||||
|
||||
t_start = time.perf_counter()
|
||||
res = state_api_fn(**kwargs)
|
||||
t_end = time.perf_counter()
|
||||
|
||||
if print_result:
|
||||
pprint.pprint(res)
|
||||
|
||||
metric = StateAPIMetric(t_end - t_start, len(res))
|
||||
if key_suffix:
|
||||
key = f"{state_api_fn.__name__}_{key_suffix}"
|
||||
else:
|
||||
key = state_api_fn.__name__
|
||||
state_stats.calls[key].append(metric)
|
||||
assert verify_cb(
|
||||
res
|
||||
), f"Calling State API failed. len(res)=({len(res)}): {err_msg}"
|
||||
except Exception as e:
|
||||
traceback.print_exc()
|
||||
assert (
|
||||
False
|
||||
), f"Calling {state_api_fn.__name__}({kwargs}) failed with {repr(e)}."
|
||||
finally:
|
||||
state_stats.pending_calls -= 1
|
||||
|
||||
return res
|
||||
|
||||
|
||||
def invoke_state_api_n(*args, **kwargs):
|
||||
def verify():
|
||||
NUM_API_CALL_SAMPLES = 10
|
||||
for _ in range(NUM_API_CALL_SAMPLES):
|
||||
invoke_state_api(*args, **kwargs)
|
||||
return True
|
||||
|
||||
test_utils.wait_for_condition(verify, retry_interval_ms=2000, timeout=30)
|
||||
|
||||
|
||||
def aggregate_perf_results(state_stats: StateAPIStats = GLOBAL_STATE_STATS):
|
||||
"""Aggregate stats of state API calls
|
||||
|
||||
Return:
|
||||
This returns a dict of below fields:
|
||||
- max_{api_key_name}_latency_sec:
|
||||
Max latency of call to {api_key_name}
|
||||
- {api_key_name}_result_size_with_max_latency:
|
||||
The size of the result (or the number of bytes for get_log API)
|
||||
for the max latency invocation
|
||||
- avg/p99/p95/p50_{api_key_name}_latency_sec:
|
||||
The percentile latency stats
|
||||
- avg_state_api_latency_sec:
|
||||
The average latency of all the state apis tracked
|
||||
"""
|
||||
# Prevent iteration when modifying error
|
||||
state_stats = deepcopy(state_stats)
|
||||
perf_result = {}
|
||||
for api_key_name, metrics in state_stats.calls.items():
|
||||
# Per api aggregation
|
||||
# Max latency
|
||||
latency_key = f"max_{api_key_name}_latency_sec"
|
||||
size_key = f"{api_key_name}_result_size_with_max_latency"
|
||||
metric = max(metrics, key=lambda metric: metric.latency_sec)
|
||||
|
||||
perf_result[latency_key] = metric.latency_sec
|
||||
perf_result[size_key] = metric.result_size
|
||||
|
||||
latency_list = np.array([metric.latency_sec for metric in metrics])
|
||||
# avg latency
|
||||
key = f"avg_{api_key_name}_latency_sec"
|
||||
perf_result[key] = np.average(latency_list)
|
||||
|
||||
# p99 latency
|
||||
key = f"p99_{api_key_name}_latency_sec"
|
||||
perf_result[key] = np.percentile(latency_list, 99)
|
||||
|
||||
# p95 latency
|
||||
key = f"p95_{api_key_name}_latency_sec"
|
||||
perf_result[key] = np.percentile(latency_list, 95)
|
||||
|
||||
# p50 latency
|
||||
key = f"p50_{api_key_name}_latency_sec"
|
||||
perf_result[key] = np.percentile(latency_list, 50)
|
||||
|
||||
all_state_api_latency = sum(
|
||||
metric.latency_sec
|
||||
for metric_samples in state_stats.calls.values()
|
||||
for metric in metric_samples
|
||||
)
|
||||
|
||||
perf_result["avg_state_api_latency_sec"] = (
|
||||
(all_state_api_latency / state_stats.total_calls)
|
||||
if state_stats.total_calls != 0
|
||||
else -1
|
||||
)
|
||||
|
||||
return perf_result
|
||||
|
||||
|
||||
@ray.remote(num_cpus=0)
|
||||
class StateAPIGeneratorActor:
|
||||
def __init__(
|
||||
self,
|
||||
apis: List[StateAPICallSpec],
|
||||
call_interval_s: float = 5.0,
|
||||
print_interval_s: float = 20.0,
|
||||
wait_after_stop: bool = True,
|
||||
print_result: bool = False,
|
||||
) -> None:
|
||||
"""An actor that periodically issues state API
|
||||
|
||||
Args:
|
||||
apis: List of StateAPICallSpec.
|
||||
call_interval_s: State apis in the `apis` will be issued
|
||||
every `call_interval_s` seconds.
|
||||
print_interval_s: How frequent state api stats will be dumped.
|
||||
wait_after_stop: When true, call to `ray.get(actor.stop.remote())`
|
||||
will wait for all pending state APIs to return.
|
||||
Setting it to `False` might miss some long-running state apis calls.
|
||||
print_result: True if result of each API call is printed. Default False.
|
||||
"""
|
||||
# Configs
|
||||
self._apis = apis
|
||||
self._call_interval_s = call_interval_s
|
||||
self._print_interval_s = print_interval_s
|
||||
self._wait_after_cancel = wait_after_stop
|
||||
self._logger = logging.getLogger(self.__class__.__name__)
|
||||
self._print_result = print_result
|
||||
|
||||
# States
|
||||
self._tasks = None
|
||||
self._fut_queue = None
|
||||
self._executor = None
|
||||
self._loop = None
|
||||
self._stopping = False
|
||||
self._stopped = False
|
||||
self._stats = StateAPIStats()
|
||||
|
||||
async def start(self):
|
||||
# Run the periodic api generator
|
||||
self._fut_queue = asyncio.Queue()
|
||||
self._executor = concurrent.futures.ThreadPoolExecutor()
|
||||
|
||||
self._tasks = [
|
||||
asyncio.ensure_future(awt)
|
||||
for awt in [
|
||||
self._run_generator(),
|
||||
self._run_result_waiter(),
|
||||
self._run_stats_reporter(),
|
||||
]
|
||||
]
|
||||
await asyncio.gather(*self._tasks)
|
||||
|
||||
def call(self, fn, verify_cb, **kwargs):
|
||||
def run_fn():
|
||||
try:
|
||||
self._logger.debug(f"calling {fn.__name__}({kwargs})")
|
||||
return invoke_state_api(
|
||||
verify_cb,
|
||||
fn,
|
||||
state_stats=self._stats,
|
||||
print_result=self._print_result,
|
||||
**kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
self._logger.warning(f"{fn.__name__}({kwargs}) failed with: {repr(e)}")
|
||||
return None
|
||||
|
||||
fut = asyncio.get_running_loop().run_in_executor(self._executor, run_fn)
|
||||
return fut
|
||||
|
||||
async def _run_stats_reporter(self):
|
||||
while not self._stopped:
|
||||
# Keep the reporter running until all pending apis finish and the bool
|
||||
# `self._stopped` is then True
|
||||
self._logger.info(pprint.pprint(aggregate_perf_results(self._stats)))
|
||||
try:
|
||||
await asyncio.sleep(self._print_interval_s)
|
||||
except asyncio.CancelledError:
|
||||
self._logger.info(
|
||||
"_run_stats_reporter cancelled, "
|
||||
f"waiting for all api {self._stats.pending_calls}calls to return..."
|
||||
)
|
||||
|
||||
async def _run_generator(self):
|
||||
try:
|
||||
while not self._stopping:
|
||||
# Run the state API in another thread
|
||||
for api_spec in self._apis:
|
||||
fut = self.call(api_spec.api, api_spec.verify_cb, **api_spec.kwargs)
|
||||
self._fut_queue.put_nowait(fut)
|
||||
|
||||
await asyncio.sleep(self._call_interval_s)
|
||||
except asyncio.CancelledError:
|
||||
# Stop running
|
||||
self._logger.info("_run_generator cancelled, now stopping...")
|
||||
return
|
||||
|
||||
async def _run_result_waiter(self):
|
||||
try:
|
||||
while not self._stopping:
|
||||
fut = await self._fut_queue.get()
|
||||
await fut
|
||||
except asyncio.CancelledError:
|
||||
self._logger.info(
|
||||
f"_run_result_waiter cancelled, cancelling {self._fut_queue.qsize()} "
|
||||
"pending futures..."
|
||||
)
|
||||
while not self._fut_queue.empty():
|
||||
fut = self._fut_queue.get_nowait()
|
||||
if self._wait_after_cancel:
|
||||
await fut
|
||||
else:
|
||||
# Ignore the queue futures if we are not
|
||||
# waiting on them after stop() called
|
||||
fut.cancel()
|
||||
return
|
||||
|
||||
def get_stats(self):
|
||||
# deep copy to prevent race between reporting and modifying stats
|
||||
return aggregate_perf_results(self._stats)
|
||||
|
||||
def ready(self):
|
||||
pass
|
||||
|
||||
def stop(self):
|
||||
self._stopping = True
|
||||
self._logger.debug(f"calling stop, canceling {len(self._tasks)} tasks")
|
||||
for task in self._tasks:
|
||||
task.cancel()
|
||||
|
||||
# This will block the stop() function until all futures are cancelled
|
||||
# if _wait_after_cancel=True. When _wait_after_cancel=False, it will still
|
||||
# wait for any in-progress futures.
|
||||
# See: https://docs.python.org/3.8/library/concurrent.futures.html
|
||||
self._executor.shutdown(wait=self._wait_after_cancel)
|
||||
self._stopped = True
|
||||
|
||||
|
||||
def periodic_invoke_state_apis_with_actor(*args, **kwargs) -> ActorHandle:
|
||||
current_node_ip = ray._private.worker.global_worker.node_ip_address
|
||||
# Schedule the actor on the current node.
|
||||
actor = StateAPIGeneratorActor.options(
|
||||
resources={f"node:{current_node_ip}": 0.001}
|
||||
).remote(*args, **kwargs)
|
||||
print("Waiting for state api actor to be ready...")
|
||||
ray.get(actor.ready.remote())
|
||||
print("State api actor is ready now.")
|
||||
actor.start.remote()
|
||||
return actor
|
||||
|
||||
|
||||
def get_state_api_manager(gcs_address: str) -> StateAPIManager:
|
||||
gcs_client = GcsClient(address=gcs_address)
|
||||
gcs_channel = GcsChannel(gcs_address=gcs_address, aio=True)
|
||||
gcs_channel.connect()
|
||||
state_api_data_source_client = StateDataSourceClient(
|
||||
gcs_channel.channel(), gcs_client
|
||||
)
|
||||
return StateAPIManager(
|
||||
state_api_data_source_client,
|
||||
thread_pool_executor=ThreadPoolExecutor(
|
||||
thread_name_prefix="state_api_test_utils"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def summarize_worker_startup_time():
|
||||
workers = list_workers(
|
||||
detail=True,
|
||||
filters=[("worker_type", "=", "WORKER")],
|
||||
limit=10000,
|
||||
raise_on_missing_output=False,
|
||||
)
|
||||
time_to_launch = []
|
||||
time_to_initialize = []
|
||||
for worker in workers:
|
||||
launch_time = worker.get("worker_launch_time_ms")
|
||||
launched_time = worker.get("worker_launched_time_ms")
|
||||
start_time = worker.get("start_time_ms")
|
||||
|
||||
if launched_time > 0:
|
||||
time_to_launch.append(launched_time - launch_time)
|
||||
if start_time:
|
||||
time_to_initialize.append(start_time - launched_time)
|
||||
time_to_launch.sort()
|
||||
time_to_initialize.sort()
|
||||
|
||||
def print_latencies(latencies):
|
||||
print(f"Avg: {round(sum(latencies) / len(latencies), 2)} ms")
|
||||
print(f"P25: {round(latencies[int(len(latencies) * 0.25)], 2)} ms")
|
||||
print(f"P50: {round(latencies[int(len(latencies) * 0.5)], 2)} ms")
|
||||
print(f"P95: {round(latencies[int(len(latencies) * 0.95)], 2)} ms")
|
||||
print(f"P99: {round(latencies[int(len(latencies) * 0.99)], 2)} ms")
|
||||
|
||||
print("Time to launch workers")
|
||||
print_latencies(time_to_launch)
|
||||
print("=======================")
|
||||
print("Time to initialize workers")
|
||||
print_latencies(time_to_initialize)
|
||||
|
||||
|
||||
def verify_failed_task(
|
||||
name: str, error_type: str, error_message: Union[str, List[str], None] = None
|
||||
) -> bool:
|
||||
"""
|
||||
Check if a task with 'name' has failed with the exact error type 'error_type'
|
||||
and 'error_message' in the error message.
|
||||
"""
|
||||
tasks = list_tasks(filters=[("name", "=", name)], detail=True)
|
||||
assert len(tasks) == 1, tasks
|
||||
t = tasks[0]
|
||||
assert t["state"] == "FAILED", t
|
||||
assert t["error_type"] == error_type, t
|
||||
if error_message is not None:
|
||||
if isinstance(error_message, str):
|
||||
error_message = [error_message]
|
||||
for msg in error_message:
|
||||
assert msg in t.get("error_message", None), t
|
||||
return True
|
||||
|
||||
|
||||
def wait_for_task_states(name_to_state: Dict[str, str], timeout: float = 30) -> None:
|
||||
"""
|
||||
Block until every task in ``name_to_state`` is observed in its expected
|
||||
state via the State API, or raise if the timeout expires.
|
||||
"""
|
||||
|
||||
def _check():
|
||||
for name, state in name_to_state.items():
|
||||
tasks = list_tasks(filters=[("name", "=", name), ("state", "=", state)])
|
||||
assert len(tasks) == 1, f"{name} not in {state}"
|
||||
return True
|
||||
|
||||
test_utils.wait_for_condition(_check, timeout=timeout)
|
||||
|
||||
|
||||
def _is_actor_task_running(actor_pid: int, task_name: str):
|
||||
"""
|
||||
Check whether the actor task `task_name` is running on the actor process
|
||||
with pid `actor_pid`.
|
||||
|
||||
Args:
|
||||
actor_pid: The pid of the actor process.
|
||||
task_name: The name of the actor task.
|
||||
|
||||
Returns:
|
||||
True if the actor task is running, False otherwise.
|
||||
|
||||
Limitation:
|
||||
If the actor task name is set using options.name and is a substring of
|
||||
the actor name, this function may return true even if the task is not
|
||||
running on the actor process. To resolve this issue, we can possibly
|
||||
pass in the actor name.
|
||||
"""
|
||||
if not psutil.pid_exists(actor_pid):
|
||||
return False
|
||||
|
||||
"""
|
||||
Why use both `psutil.Process.name()` and `psutil.Process.cmdline()`?
|
||||
|
||||
1. Core worker processes call `setproctitle` to set the process title before
|
||||
and after executing tasks. However, the definition of "title" is a bit
|
||||
complex.
|
||||
|
||||
[ref]: https://github.com/dvarrazzo/py-setproctitle
|
||||
|
||||
> The process title is usually visible in files such as /proc/PID/cmdline,
|
||||
/proc/PID/status, /proc/PID/comm, depending on the operating system and
|
||||
kernel version. This information is used by user-space tools such as ps
|
||||
and top.
|
||||
|
||||
Ideally, we would only need to check `psutil.Process.cmdline()`, but I decided
|
||||
to check both `psutil.Process.name()` and `psutil.Process.cmdline()` based on
|
||||
the definition of "title" stated above.
|
||||
|
||||
2. Additionally, the definition of `psutil.Process.name()` is not consistent
|
||||
with the definition of "title" in `setproctitle`. The length of `/proc/PID/comm` and
|
||||
the prefix of `/proc/PID/cmdline` affect the return value of
|
||||
`psutil.Process.name()`.
|
||||
|
||||
In addition, executing `setproctitle` in different threads within the same
|
||||
process may result in different outcomes.
|
||||
|
||||
To learn more details, please refer to the source code of `psutil`:
|
||||
|
||||
[ref]:
|
||||
https://github.com/giampaolo/psutil/blob/a17550784b0d3175da01cdb02cee1bc6b61637dc/psutil/__init__.py#L664-L693
|
||||
|
||||
3. `/proc/PID/comm` will be truncated to TASK_COMM_LEN (16) characters
|
||||
(including the terminating null byte).
|
||||
|
||||
[ref]:
|
||||
https://man7.org/linux/man-pages/man5/proc_pid_comm.5.html
|
||||
"""
|
||||
name = psutil.Process(actor_pid).name()
|
||||
if task_name in name and name.startswith("ray::"):
|
||||
return True
|
||||
|
||||
cmdline = psutil.Process(actor_pid).cmdline()
|
||||
# If `options.name` is set, the format is `ray::<task_name>`. If not,
|
||||
# the format is `ray::<actor_name>.<task_name>`.
|
||||
if cmdline and task_name in cmdline[0] and cmdline[0].startswith("ray::"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def verify_schema(state, result_dict: dict, detail: bool = False):
|
||||
"""
|
||||
Verify the schema of the result_dict is the same as the state.
|
||||
"""
|
||||
state_fields_columns = set()
|
||||
if detail:
|
||||
state_fields_columns = state.columns()
|
||||
else:
|
||||
state_fields_columns = state.base_columns()
|
||||
|
||||
for k in state_fields_columns:
|
||||
assert k in result_dict
|
||||
|
||||
for k in result_dict:
|
||||
assert k in state_fields_columns
|
||||
|
||||
# Make the field values can be converted without error as well
|
||||
state(**result_dict)
|
||||
|
||||
|
||||
def create_api_options(
|
||||
timeout: int = DEFAULT_RPC_TIMEOUT,
|
||||
limit: int = DEFAULT_LIMIT,
|
||||
filters: List[Tuple[str, PredicateType, SupportedFilterType]] = None,
|
||||
detail: bool = False,
|
||||
exclude_driver: bool = True,
|
||||
):
|
||||
if not filters:
|
||||
filters = []
|
||||
return ListApiOptions(
|
||||
limit=limit,
|
||||
timeout=timeout,
|
||||
filters=filters,
|
||||
server_timeout_multiplier=1.0,
|
||||
detail=detail,
|
||||
exclude_driver=exclude_driver,
|
||||
)
|
||||
Reference in New Issue
Block a user