import asyncio import fnmatch import io import json import logging import os import pathlib import random import socket import subprocess import sys import tempfile import time import timeit import traceback import uuid from collections.abc import Hashable from contextlib import contextmanager, redirect_stderr, redirect_stdout from datetime import datetime from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type from urllib.parse import quote, urlparse import requests import yaml import ray import ray._private.memory_monitor as memory_monitor import ray._private.services import ray._private.services as services import ray._private.utils import ray.dashboard.consts as dashboard_consts from ray._common.network_utils import build_address, parse_address from ray._common.test_utils import ( MetricSamplePattern, PrometheusTimeseries, fetch_prometheus_metric_timeseries, fetch_prometheus_timeseries, wait_for_condition, ) from ray._common.tls_utils import generate_self_signed_tls_certs from ray._common.utils import get_or_create_event_loop from ray._private import ( ray_constants, ) from ray._private.internal_api import memory_summary from ray._private.services import ProcessInfo from ray._private.worker import RayContext from ray._raylet import Config, GcsClient, GcsClientOptions, GlobalStateAccessor from ray.core.generated import ( gcs_pb2, gcs_service_pb2, node_manager_pb2, ) from ray.util.queue import Empty, Queue, _QueueActor from ray.util.state import get_actor, list_actors import psutil # We must import psutil after ray because we bundle it with ray. logger = logging.getLogger(__name__) EXE_SUFFIX = ".exe" if sys.platform == "win32" else "" RAY_PATH = os.path.abspath(os.path.dirname(os.path.dirname(__file__))) REDIS_EXECUTABLE = os.path.join( RAY_PATH, "core/src/ray/thirdparty/redis/src/redis-server" + EXE_SUFFIX ) def make_global_state_accessor(ray_context): gcs_options = GcsClientOptions.create( ray_context.address_info["gcs_address"], None, allow_cluster_id_nil=True, fetch_cluster_id_if_nil=False, ) global_state_accessor = GlobalStateAccessor(gcs_options) global_state_accessor.connect() return global_state_accessor def external_redis_test_enabled(): return os.environ.get("TEST_EXTERNAL_REDIS") == "1" def rocksdb_gcs_test_enabled(): """True when the test suite should run against the RocksDB GCS backend (REP-64). Set by the buildkite ":ray: core: rocksdb tests" job. """ return os.environ.get("TEST_GCS_ROCKSDB") == "1" def persistent_gcs_test_enabled(): """True when the GCS backend under test persists state across restart (external Redis or RocksDB). Use this — not external_redis_test_enabled — to branch test assertions on "is GCS state durable across restart?". """ return external_redis_test_enabled() or rocksdb_gcs_test_enabled() def redis_replicas(): return int(os.environ.get("TEST_EXTERNAL_REDIS_REPLICAS", "1")) def redis_sentinel_replicas(): return int(os.environ.get("TEST_EXTERNAL_REDIS_SENTINEL_REPLICAS", "2")) def get_redis_cli(port, enable_tls): try: # If there is no redis libs installed, skip the check. # This could happen In minimal test, where we don't have # redis. import redis except Exception: return True params = {} if enable_tls: from ray._raylet import Config params = {"ssl": True, "ssl_cert_reqs": "required"} if Config.REDIS_CA_CERT(): params["ssl_ca_certs"] = Config.REDIS_CA_CERT() if Config.REDIS_CLIENT_CERT(): params["ssl_certfile"] = Config.REDIS_CLIENT_CERT() if Config.REDIS_CLIENT_KEY(): params["ssl_keyfile"] = Config.REDIS_CLIENT_KEY() return redis.Redis("localhost", str(port), **params) def start_redis_sentinel_instance( session_dir_path: str, port: int, redis_master_port: int, password: Optional[str] = None, enable_tls: bool = False, db_dir=None, free_port=0, ): config_file = os.path.join( session_dir_path, "redis-sentinel-" + uuid.uuid4().hex + ".conf" ) config_lines = [] # Port for this Sentinel instance if enable_tls: config_lines.append(f"port {free_port}") else: config_lines.append(f"port {port}") # Monitor the Redis master config_lines.append(f"sentinel monitor redis-test 127.0.0.1 {redis_master_port} 1") config_lines.append( "sentinel down-after-milliseconds redis-test 1000" ) # failover after 1 second config_lines.append("sentinel failover-timeout redis-test 5000") # config_lines.append("sentinel parallel-syncs redis-test 1") if password: config_lines.append(f"sentinel auth-pass redis-test {password}") if enable_tls: config_lines.append(f"tls-port {port}") if Config.REDIS_CA_CERT(): config_lines.append(f"tls-ca-cert-file {Config.REDIS_CA_CERT()}") # Check and add TLS client certificate file if Config.REDIS_CLIENT_CERT(): config_lines.append(f"tls-cert-file {Config.REDIS_CLIENT_CERT()}") # Check and add TLS client key file if Config.REDIS_CLIENT_KEY(): config_lines.append(f"tls-key-file {Config.REDIS_CLIENT_KEY()}") config_lines.append("tls-auth-clients no") config_lines.append("sentinel tls-auth-clients redis-test no") if db_dir: config_lines.append(f"dir {db_dir}") with open(config_file, "w") as f: f.write("\n".join(config_lines)) command = [REDIS_EXECUTABLE, config_file, "--sentinel"] process_info = ray._private.services.start_ray_process( command, ray_constants.PROCESS_TYPE_REDIS_SERVER, fate_share=False, ) return process_info def start_redis_instance( session_dir_path: str, port: int, redis_max_clients: Optional[int] = None, num_retries: int = 20, stdout_file: Optional[str] = None, stderr_file: Optional[str] = None, password: Optional[str] = None, fate_share: Optional[bool] = None, port_denylist: Optional[List[int]] = None, listen_to_localhost_only: bool = False, enable_tls: bool = False, replica_of: Optional[int] = None, leader_id: Optional[bytes] = None, db_dir: Optional[str] = None, free_port: int = 0, ): """Start a single Redis server. Notes: We will initially try to start the Redis instance at the given port, and then try at most `num_retries - 1` times to start the Redis instance at successive random ports. Args: session_dir_path: Path to the session directory of this Ray cluster. port: Try to start a Redis server at this port. redis_max_clients: If this is provided, Ray will attempt to configure Redis with this maxclients number. num_retries: The number of times to attempt to start Redis at successive ports. stdout_file: A file handle opened for writing to redirect stdout to. If no redirection should happen, then this should be None. stderr_file: A file handle opened for writing to redirect stderr to. If no redirection should happen, then this should be None. password: Prevents external clients without the password from connecting to Redis if provided. fate_share: If True, the Redis process is bound to the parent's job on Windows so it terminates with the parent. port_denylist: A set of denylist ports that shouldn't be used when allocating a new port. listen_to_localhost_only: Redis server only listens to localhost (127.0.0.1) if it's true, otherwise it listens to all network interfaces. enable_tls: Enable the TLS/SSL in Redis or not replica_of: When set, configure this server as a replica of the given primary Redis port. leader_id: Cluster node id of the leader to replicate when running with multiple replicas. db_dir: Directory passed to ``--dir`` so Redis persists data here. free_port: Plaintext port used alongside ``--tls-port`` when TLS is enabled. Returns: A tuple of the port used by Redis and ProcessInfo for the process that was started. If a port is passed in, then the returned port value is the same. Raises: Exception: An exception is raised if Redis could not be started. """ assert os.path.isfile(REDIS_EXECUTABLE) # Construct the command to start the Redis server. command = [REDIS_EXECUTABLE] if password: if " " in password: raise ValueError("Spaces not permitted in redis password.") command += ["--requirepass", password] if redis_replicas() > 1: command += ["--cluster-enabled", "yes", "--cluster-config-file", f"node-{port}"] if enable_tls: command += [ "--tls-port", str(port), "--loglevel", "warning", "--port", str(free_port), ] else: command += ["--port", str(port), "--loglevel", "warning"] if listen_to_localhost_only: command += ["--bind", "127.0.0.1"] pidfile = os.path.join(session_dir_path, "redis-" + uuid.uuid4().hex + ".pid") command += ["--pidfile", pidfile] if enable_tls: if Config.REDIS_CA_CERT(): command += ["--tls-ca-cert-file", Config.REDIS_CA_CERT()] if Config.REDIS_CLIENT_CERT(): command += ["--tls-cert-file", Config.REDIS_CLIENT_CERT()] if Config.REDIS_CLIENT_KEY(): command += ["--tls-key-file", Config.REDIS_CLIENT_KEY()] if replica_of is not None: command += ["--tls-replication", "yes"] command += ["--tls-auth-clients", "no", "--tls-cluster", "yes"] if sys.platform != "win32": command += ["--save", "", "--appendonly", "no"] if db_dir is not None: command += ["--dir", str(db_dir)] process_info = ray._private.services.start_ray_process( command, ray_constants.PROCESS_TYPE_REDIS_SERVER, stdout_file=stdout_file, stderr_file=stderr_file, fate_share=fate_share, ) node_id = None if redis_replicas() > 1: # Setup redis cluster import redis while True: try: redis_cli = get_redis_cli(port, enable_tls) if replica_of is None: slots = [str(i) for i in range(16384)] redis_cli.cluster("addslots", *slots) else: logger.info(redis_cli.cluster("meet", "127.0.0.1", str(replica_of))) logger.info(redis_cli.cluster("replicate", leader_id)) node_id = redis_cli.cluster("myid") break except ( redis.exceptions.ConnectionError, redis.exceptions.ResponseError, ) as e: from time import sleep logger.info( f"Waiting for redis to be up. Check failed with error: {e}. " "Will retry in 0.1s" ) if process_info.process.poll() is not None: raise Exception( f"Redis process exited unexpectedly: {process_info}. " f"Exit code: {process_info.process.returncode}" ) sleep(0.1) logger.info( f"Redis started with node_id {node_id} and pid {process_info.process.pid}" ) return node_id, process_info def _pid_alive(pid: int): """Check if the process with this PID is alive or not. Args: pid: The pid to check. Returns: This returns false if the process is dead. Otherwise, it returns true. """ alive = True try: proc = psutil.Process(pid) if proc.status() == psutil.STATUS_ZOMBIE: alive = False except psutil.NoSuchProcess: alive = False return alive def _check_call_windows(main, argv, capture_stdout=False, capture_stderr=False): # We use this function instead of calling the "ray" command to work around # some deadlocks that occur when piping ray's output on Windows stream = io.TextIOWrapper(io.BytesIO(), encoding=sys.stdout.encoding) old_argv = sys.argv[:] try: sys.argv = argv[:] try: with redirect_stderr(stream if capture_stderr else sys.stderr): with redirect_stdout(stream if capture_stdout else sys.stdout): main() finally: stream.flush() except SystemExit as ex: if ex.code: output = stream.buffer.getvalue() raise subprocess.CalledProcessError(ex.code, argv, output) except Exception as ex: output = stream.buffer.getvalue() raise subprocess.CalledProcessError(1, argv, output, ex.args[0]) finally: sys.argv = old_argv if capture_stdout: sys.stdout.buffer.write(stream.buffer.getvalue()) elif capture_stderr: sys.stderr.buffer.write(stream.buffer.getvalue()) return stream.buffer.getvalue() def check_call_subprocess(argv, capture_stdout=False, capture_stderr=False): # We use this function instead of calling the "ray" command to work around # some deadlocks that occur when piping ray's output on Windows from ray.scripts.scripts import main as ray_main if sys.platform == "win32": result = _check_call_windows( ray_main, argv, capture_stdout=capture_stdout, capture_stderr=capture_stderr ) else: stdout_redir = None stderr_redir = None if capture_stdout: stdout_redir = subprocess.PIPE if capture_stderr and capture_stdout: stderr_redir = subprocess.STDOUT elif capture_stderr: stderr_redir = subprocess.PIPE proc = subprocess.Popen(argv, stdout=stdout_redir, stderr=stderr_redir) (stdout, stderr) = proc.communicate() if proc.returncode: raise subprocess.CalledProcessError(proc.returncode, argv, stdout, stderr) result = b"".join([s for s in [stdout, stderr] if s is not None]) return result def check_call_ray(args, capture_stdout=False, capture_stderr=False): check_call_subprocess(["ray"] + args, capture_stdout, capture_stderr) def get_dashboard_agent_address(gcs_client: GcsClient, node_id: str): result = gcs_client.internal_kv_get( f"{dashboard_consts.DASHBOARD_AGENT_ADDR_NODE_ID_PREFIX}{node_id}".encode(), namespace=ray_constants.KV_NAMESPACE_DASHBOARD, timeout=dashboard_consts.GCS_RPC_TIMEOUT_SECONDS, ) if result: # Returns [ip, http_port, grpc_port] ip, _, grpc_port = json.loads(result) return f"{ip}:{grpc_port}" return None def wait_for_dashboard_agent_available(cluster): gcs_client = GcsClient(address=cluster.address) wait_for_condition( lambda: get_dashboard_agent_address(gcs_client, cluster.head_node.node_id) is not None ) def wait_for_aggregator_agent(address: str, node_id: str, timeout: float = 10) -> None: """Wait for the aggregator agent to be ready by checking socket connectivity.""" gcs_client = GcsClient(address=address) # Wait for the agent to publish its address wait_for_condition( lambda: get_dashboard_agent_address(gcs_client, node_id) is not None ) # Get the agent address and test socket connectivity agent_address = get_dashboard_agent_address(gcs_client, node_id) parsed = urlparse(f"grpc://{agent_address}") def _can_connect() -> bool: try: with socket.create_connection((parsed.hostname, parsed.port), timeout=1): return True except OSError: return False wait_for_condition(_can_connect, timeout=timeout) def wait_for_aggregator_agent_if_enabled( address: str, node_id: str, timeout: float = 10 ) -> None: """Wait for aggregator agent only if aggregator mode is enabled. Checks RAY_enable_core_worker_ray_event_to_aggregator env var. """ if os.environ.get("RAY_enable_core_worker_ray_event_to_aggregator") == "1": wait_for_aggregator_agent(address, node_id, timeout) def wait_for_pid_to_exit(pid: int, timeout: float = 20): start_time = time.time() while time.time() - start_time < timeout: if not _pid_alive(pid): return time.sleep(0.1) raise TimeoutError(f"Timed out while waiting for process {pid} to exit.") def wait_for_children_of_pid(pid, num_children=1, timeout=20): p = psutil.Process(pid) start_time = time.time() alive = [] while time.time() - start_time < timeout: alive = p.children(recursive=False) num_alive = len(alive) if num_alive >= num_children: return time.sleep(0.1) raise TimeoutError( f"Timed out while waiting for process {pid} children to start " f"({num_alive}/{num_children} started: {alive})." ) def wait_for_children_of_pid_to_exit(pid, timeout=20): children = psutil.Process(pid).children() if len(children) == 0: return _, alive = psutil.wait_procs(children, timeout=timeout) if len(alive) > 0: raise TimeoutError( "Timed out while waiting for process children to exit." " Children still alive: {}.".format([p.name() for p in alive]) ) def kill_process_by_name(name, SIGKILL=False): for p in psutil.process_iter(attrs=["name"]): if p.info["name"] == name + ray._private.services.EXE_SUFFIX: if SIGKILL: p.kill() else: p.terminate() def kill_processes(process_infos: List[ProcessInfo]): """ Forcefully kills the list of given processes. Ignores processes that are already dead. Args: process_infos: The list of ProcessInfo representing the processes to kill. Raises: TimeoutError: If the process did not exit within 5 seconds. """ for process_info in process_infos: try: process_info.process.kill() process_info.process.wait(timeout=5) except ProcessLookupError: # Process already dead pass except subprocess.TimeoutExpired as exception: raise TimeoutError( f"Process {process_info.process.pid} did not exit within 5 seconds " "after SIGKILL" ) from exception def run_string_as_driver_stdout_stderr( driver_script: str, env: Dict = None, encode: str = "utf-8" ) -> Tuple[str, 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: Text encoding used to send the script to the subprocess and decode its stdout/stderr. Returns: The script's stdout and stderr. """ proc = subprocess.Popen( [sys.executable, "-"], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env, ) with proc: outputs_bytes = proc.communicate(driver_script.encode(encoding=encode)) out_str, err_str = [ ray._common.utils.decode(output, encode_type=encode) for output in outputs_bytes ] if proc.returncode: print(out_str) print(err_str) raise subprocess.CalledProcessError( proc.returncode, proc.args, out_str, err_str ) return out_str, err_str def run_string_as_driver_nonblocking(driver_script: str, env: Dict = None): """Start a driver as a separate process and return immediately. Args: driver_script: A string to run as a Python script. env: The environment variables for the driver. Returns: A handle to the driver process. """ script = "; ".join( [ "import sys", "script = sys.stdin.read()", "sys.stdin.close()", "del sys", 'exec("del script\\n" + script)', ] ) proc = subprocess.Popen( [sys.executable, "-c", script], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env, ) proc.stdin.write(driver_script.encode("ascii")) proc.stdin.close() return proc def convert_actor_state(state): if not state: return None return gcs_pb2.ActorTableData.ActorState.DESCRIPTOR.values_by_number[state].name def wait_for_num_actors(num_actors, state=None, timeout=10): state = convert_actor_state(state) start_time = time.time() while time.time() - start_time < timeout: if ( len( list_actors( filters=[("state", "=", state)] if state else None, limit=num_actors, ) ) >= num_actors ): return time.sleep(0.1) raise TimeoutError("Timed out while waiting for global state.") def kill_actor_and_wait_for_failure(actor, timeout=10, retry_interval_ms=100): actor_id = actor._actor_id.hex() current_num_restarts = get_actor(id=actor_id).num_restarts ray.kill(actor) start = time.time() while time.time() - start <= timeout: actor_state = get_actor(id=actor_id) if ( actor_state.state == "DEAD" or actor_state.num_restarts > current_num_restarts ): return time.sleep(retry_interval_ms / 1000.0) raise RuntimeError("It took too much time to kill an actor: {}".format(actor_id)) def wait_for_assertion( assertion_predictor: Callable, timeout: int = 10, retry_interval_ms: int = 100, raise_exceptions: bool = False, **kwargs: Any, ): """Wait until an assertion is met or time out with an exception. Args: assertion_predictor: A function that predicts the assertion. timeout: Maximum timeout in seconds. retry_interval_ms: Retry interval in milliseconds. raise_exceptions: If true, exceptions that occur while executing assertion_predictor won't be caught and instead will be raised. **kwargs: Arguments to pass to the condition_predictor. Raises: RuntimeError: If the assertion is not met before the timeout expires. """ def _assertion_to_condition(): try: assertion_predictor(**kwargs) return True except AssertionError: return False try: wait_for_condition( _assertion_to_condition, timeout=timeout, retry_interval_ms=retry_interval_ms, raise_exceptions=raise_exceptions, **kwargs, ) except RuntimeError: assertion_predictor(**kwargs) # Should fail assert def get_metric_check_condition( metrics_to_check: List[MetricSamplePattern], timeseries: PrometheusTimeseries, export_addr: Optional[str] = None, ) -> Callable[[], bool]: """A condition to check if a prometheus metrics reach a certain value. This is a blocking check that can be passed into a `wait_for_condition` style function. Args: metrics_to_check: A list of MetricSamplePattern. The fields that aren't `None` will be matched. timeseries: A PrometheusTimeseries object to store the metrics. export_addr: Optional address to export metrics to. Returns: A function that returns True if all the metrics are emitted. """ node_info = ray.nodes()[0] metrics_export_port = node_info["MetricsExportPort"] addr = node_info["NodeManagerAddress"] prom_addr = export_addr or build_address(addr, metrics_export_port) def f(): for metric_pattern in metrics_to_check: metric_samples = fetch_prometheus_timeseries( [prom_addr], timeseries ).metric_samples.values() for metric_sample in metric_samples: if metric_pattern.matches(metric_sample): break else: logger.info( f"Didn't find {metric_pattern} in all samples: {metric_samples}", ) return False return True return f def wait_until_succeeded_without_exception( func: Callable, exceptions: Tuple[Type[BaseException], ...], *args, timeout_ms: int = 1000, retry_interval_ms: int = 100, raise_last_ex: bool = False, ): """A helper function that waits until a given function completes without exceptions. Args: func: A function to run. exceptions: Exceptions that are supposed to occur. *args: arguments to pass for a given func timeout_ms: Maximum timeout in milliseconds. retry_interval_ms: Retry interval in milliseconds. raise_last_ex: Raise the last exception when timeout. Returns: Whether ``func`` succeeded within the timeout. """ if isinstance(type(exceptions), tuple): raise Exception("exceptions arguments should be given as a tuple") time_elapsed = 0 start = time.time() last_ex = None while time_elapsed <= timeout_ms: try: func(*args) return True except exceptions as ex: last_ex = ex time_elapsed = (time.time() - start) * 1000 time.sleep(retry_interval_ms / 1000.0) if raise_last_ex: ex_stack = ( traceback.format_exception(type(last_ex), last_ex, last_ex.__traceback__) if last_ex else [] ) ex_stack = "".join(ex_stack) raise Exception(f"Timed out while testing, {ex_stack}") return False def recursive_fnmatch(dirpath, pattern): """Looks at a file directory subtree for a filename pattern. Similar to glob.glob(..., recursive=True) but also supports 2.7 """ matches = [] for root, dirnames, filenames in os.walk(dirpath): for filename in fnmatch.filter(filenames, pattern): matches.append(os.path.join(root, filename)) return matches def generate_system_config_map(**kwargs): ray_kwargs = { "_system_config": kwargs, } return ray_kwargs def same_elements(elems_a, elems_b): """Checks if two iterables (such as lists) contain the same elements. Elements do not have to be hashable (this allows us to compare sets of dicts for example). This comparison is not necessarily efficient. """ a = list(elems_a) b = list(elems_b) for x in a: if x not in b: return False for x in b: if x not in a: return False return True @ray.remote def _put(obj): return obj def put_object(obj, use_ray_put): if use_ray_put: return ray.put(obj) else: return _put.remote(obj) def wait_until_server_available(address, timeout_ms=5000, retry_interval_ms=100): ip, port_str = parse_address(address) port = int(port_str) time_elapsed = 0 start = time.time() while time_elapsed <= timeout_ms: s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) s.settimeout(1) try: s.connect((ip, port)) except Exception: time_elapsed = (time.time() - start) * 1000 time.sleep(retry_interval_ms / 1000.0) s.close() continue s.close() return True return False def get_other_nodes(cluster, exclude_head=False): """Get all nodes except the one that we're connected to.""" return [ node for node in cluster.list_all_nodes() if node._raylet_socket_name != ray._private.worker._global_node._raylet_socket_name and (exclude_head is False or node.head is False) ] def get_non_head_nodes(cluster): """Get all non-head nodes.""" return list(filter(lambda x: x.head is False, cluster.list_all_nodes())) def init_error_pubsub(): """Initialize error info pub/sub""" s = ray._raylet.GcsErrorSubscriber( address=ray._private.worker.global_worker.gcs_client.address ) s.subscribe() return s def get_error_message(subscriber, num=1e6, error_type=None, timeout=20): """Gets errors from GCS subscriber. Returns maximum `num` error strings within `timeout`. Only returns errors of `error_type` if specified. """ deadline = time.time() + timeout msgs = [] while time.time() < deadline and len(msgs) < num: _, error_data = subscriber.poll(timeout=deadline - time.time()) if not error_data: # Timed out before any data is received. break if error_type is None or error_type == error_data["type"]: msgs.append(error_data) else: time.sleep(0.01) return msgs def init_log_pubsub(): """Initialize log pub/sub""" s = ray._raylet.GcsLogSubscriber( address=ray._private.worker.global_worker.gcs_client.address ) s.subscribe() return s def get_log_data( subscriber, num: int = 1e6, timeout: float = 20, job_id: Optional[str] = None, matcher=None, ) -> List[dict]: deadline = time.time() + timeout msgs = [] while time.time() < deadline and len(msgs) < num: logs_data = subscriber.poll(timeout=deadline - time.time()) if not logs_data: # Timed out before any data is received. break if job_id and job_id != logs_data["job"]: continue if matcher and all(not matcher(line) for line in logs_data["lines"]): continue msgs.append(logs_data) return msgs def get_log_message( subscriber, num: int = 1e6, timeout: float = 20, job_id: Optional[str] = None, matcher=None, ) -> List[List[str]]: """Gets log lines through GCS subscriber. Returns maximum `num` of log messages, within `timeout`. If `job_id` or `match` is specified, only returns log lines from `job_id` or when `matcher` is true. """ msgs = get_log_data(subscriber, num, timeout, job_id, matcher) return [msg["lines"] for msg in msgs] def get_log_sources( subscriber, num: int = 1e6, timeout: float = 20, job_id: Optional[str] = None, matcher=None, ): """Get the source of all log messages""" msgs = get_log_data(subscriber, num, timeout, job_id, matcher) return {msg["pid"] for msg in msgs} def get_log_batch( subscriber, num: int, timeout: float = 20, job_id: Optional[str] = None, matcher=None, ) -> List[str]: """Gets log batches through GCS subscriber. Returns maximum `num` batches of logs. Each batch is a dict that includes metadata such as `pid`, `job_id`, and `lines` of log messages. If `job_id` or `match` is specified, only returns log batches from `job_id` or when `matcher` is true. """ deadline = time.time() + timeout batches = [] while time.time() < deadline and len(batches) < num: logs_data = subscriber.poll(timeout=deadline - time.time()) if not logs_data: # Timed out before any data is received. break if job_id and job_id != logs_data["job"]: continue if matcher and not matcher(logs_data): continue batches.append(logs_data) return batches def format_web_url(url): """Format web url.""" url = url.replace("localhost", "http://127.0.0.1") if not url.startswith("http://"): return "http://" + url return url def client_test_enabled() -> bool: return ray._private.client_mode_hook.is_client_mode_enabled def object_memory_usage() -> bool: """Returns the number of bytes used in the object store.""" total = ray.cluster_resources().get("object_store_memory", 0) avail = ray.available_resources().get("object_store_memory", 0) return total - avail def raw_metric_timeseries( info: RayContext, result: PrometheusTimeseries ) -> Dict[str, List[Any]]: """Return prometheus timeseries from a RayContext""" metrics_page = "localhost:{}".format(info.address_info["metrics_export_port"]) return fetch_prometheus_metric_timeseries([metrics_page], result) def get_system_metric_for_component( system_metric: str, component: str, prometheus_server_address: str, max_attempts: int = 3, backoff_base_s: float = 1.0, request_timeout_s: float = 30.0, ) -> List[float]: """Get the system metric for a given component from a Prometheus server address. Please note: - This function requires the availability of the Prometheus server. Therefore, it requires the server address. - It assumes the system metric has a `Component` label and `pid` label. `pid` is the process id, so it can be used to uniquely identify the process. Retries up to ``max_attempts`` times on connection errors, timeouts, and HTTP 5xx responses, with exponential backoff (``backoff_base_s`` * 2^(attempt-1) seconds). HTTP 4xx responses are not retried — they indicate a bad query. """ session_name = os.path.basename( ray._private.worker._global_node.get_session_dir_path() ) query = f"sum({system_metric}{{Component='{component}',SessionName='{session_name}'}}) by (pid)" url = f"{prometheus_server_address}/api/v1/query?query={quote(query)}" for attempt in range(1, max_attempts + 1): backoff_s = backoff_base_s * (2 ** (attempt - 1)) try: resp = requests.get(url, timeout=request_timeout_s) resp.raise_for_status() except requests.exceptions.RequestException as e: # HTTPError from raise_for_status() sets e.response; other # RequestException subclasses (ConnectionError, Timeout, ...) leave it None. err_resp = e.response status_code = err_resp.status_code if err_resp is not None else None body_truncated = (err_resp.text or "")[:500] if err_resp is not None else "" # 4xx indicates a malformed query — don't retry. is_retryable = status_code is None or status_code >= 500 if is_retryable and attempt < max_attempts: logger.warning( "Prometheus query failed (attempt %d/%d), retrying in %.1fs: " "error=%r, url=%s, query=%s", attempt, max_attempts, backoff_s, e, url, query, ) time.sleep(backoff_s) continue if status_code is not None: raise RuntimeError( f"Failed to query Prometheus after {attempt} attempts: " f"last_status={status_code}, url={url}, query={query}, " f"response_body={body_truncated!r}" ) from e raise RuntimeError( f"Failed to query Prometheus after {attempt} attempts: " f"last_error={e!r}, url={url}, query={query}" ) from e if attempt > 1: logger.info( "Prometheus query succeeded on attempt %d/%d", attempt, max_attempts ) result = resp.json() return [float(item["value"][1]) for item in result["data"]["result"]] def get_test_config_path(config_file_name): """Resolve the test config path from the config file dir""" here = os.path.realpath(__file__) path = pathlib.Path(here) grandparent = path.parent.parent return os.path.join(grandparent, "tests/test_cli_patterns", config_file_name) def load_test_config(config_file_name): """Loads a config yaml from tests/test_cli_patterns.""" config_path = get_test_config_path(config_file_name) config = yaml.safe_load(open(config_path).read()) return config def set_setup_func(): import ray._private.runtime_env as runtime_env runtime_env.VAR = "hello world" class BatchQueue(Queue): def __init__(self, maxsize: int = 0, actor_options: Optional[Dict] = None) -> None: actor_options = actor_options or {} self.maxsize = maxsize self.actor = ( ray.remote(_BatchQueueActor).options(**actor_options).remote(self.maxsize) ) def get_batch( self, batch_size: int = None, total_timeout: Optional[float] = None, first_timeout: Optional[float] = None, ) -> List[Any]: """Gets batch of items from the queue and returns them in a list in order. Args: batch_size: Max number of items to return. ``None`` means drain everything currently in the queue (subject to the timeouts). total_timeout: Total time, in seconds, to wait for the entire batch. first_timeout: Time, in seconds, to wait for the first item before raising ``Empty``. Returns: List of items pulled off the queue, in arrival order. Raises: Empty: if the queue does not contain the desired number of items """ return ray.get( self.actor.get_batch.remote(batch_size, total_timeout, first_timeout) ) class _BatchQueueActor(_QueueActor): async def get_batch(self, batch_size=None, total_timeout=None, first_timeout=None): start = timeit.default_timer() try: first = await asyncio.wait_for(self.queue.get(), first_timeout) batch = [first] if total_timeout: end = timeit.default_timer() total_timeout = max(total_timeout - (end - start), 0) except asyncio.TimeoutError: raise Empty if batch_size is None: if total_timeout is None: total_timeout = 0 while True: try: start = timeit.default_timer() batch.append( await asyncio.wait_for(self.queue.get(), total_timeout) ) if total_timeout: end = timeit.default_timer() total_timeout = max(total_timeout - (end - start), 0) except asyncio.TimeoutError: break else: for _ in range(batch_size - 1): try: start = timeit.default_timer() batch.append( await asyncio.wait_for(self.queue.get(), total_timeout) ) if total_timeout: end = timeit.default_timer() total_timeout = max(total_timeout - (end - start), 0) except asyncio.TimeoutError: break return batch def is_placement_group_removed(pg): table = ray.util.placement_group_table(pg) if "state" not in table: return False return table["state"] == "REMOVED" def placement_group_assert_no_leak(pgs_created): for pg in pgs_created: ray.util.remove_placement_group(pg) def wait_for_pg_removed(): for pg_entry in ray.util.placement_group_table().values(): if pg_entry["state"] != "REMOVED": return False return True wait_for_condition(wait_for_pg_removed) cluster_resources = ray.cluster_resources() cluster_resources.pop("memory") cluster_resources.pop("object_store_memory") def wait_for_resource_recovered(): for resource, val in ray.available_resources().items(): if resource in cluster_resources and cluster_resources[resource] != val: return False if "_group_" in resource: return False return True wait_for_condition(wait_for_resource_recovered) def monitor_memory_usage( print_interval_s: int = 30, record_interval_s: int = 5, warning_threshold: float = 0.9, ): """Run the memory monitor actor that prints the memory usage. The monitor will run on the same node as this function is called. Args: print_interval_s: How often, in seconds, memory usage information is logged. record_interval_s: How often, in seconds, the monitor samples and records memory usage between log lines. warning_threshold: The threshold where the memory usage warning is printed. Returns: The memory monitor actor. """ assert ray.is_initialized(), "The API is only available when Ray is initialized." @ray.remote(num_cpus=0) class MemoryMonitorActor: def __init__( self, print_interval_s: float = 20, record_interval_s: float = 5, warning_threshold: float = 0.9, n: int = 10, ): """The actor that monitor the memory usage of the cluster. Params: print_interval_s: The interval where memory usage is printed. record_interval_s: The interval where memory usage is recorded. warning_threshold: The threshold where memory warning is printed n: When memory usage is printed, top n entries are printed. """ # -- Interval the monitor prints the memory usage information. -- self.print_interval_s = print_interval_s # -- Interval the monitor records the memory usage information. -- self.record_interval_s = record_interval_s # -- Whether or not the monitor is running. -- self.is_running = False # -- The used_gb/total_gb threshold where warning message omits. -- self.warning_threshold = warning_threshold # -- The monitor that calculates the memory usage of the node. -- self.monitor = memory_monitor.MemoryMonitor() # -- The top n memory usage of processes are printed. -- self.n = n # -- The peak memory usage in GB during lifetime of monitor. -- self.peak_memory_usage = 0 # -- The top n memory usage of processes # during peak memory usage. -- self.peak_top_n_memory_usage = "" # -- The last time memory usage was printed -- self._last_print_time = 0 # -- logger. -- logging.basicConfig(level=logging.INFO) def ready(self): pass async def run(self): """Run the monitor.""" self.is_running = True while self.is_running: now = time.time() used_gb, total_gb = self.monitor.get_memory_usage() top_n_memory_usage = memory_monitor.get_top_n_memory_usage(n=self.n) if used_gb > self.peak_memory_usage: self.peak_memory_usage = used_gb self.peak_top_n_memory_usage = top_n_memory_usage if used_gb > total_gb * self.warning_threshold: logging.warning( "The memory usage is high: " f"{used_gb / total_gb * 100}%" ) if now - self._last_print_time > self.print_interval_s: logging.info(f"Memory usage: {used_gb} / {total_gb}") logging.info(f"Top {self.n} process memory usage:") logging.info(top_n_memory_usage) self._last_print_time = now await asyncio.sleep(self.record_interval_s) async def stop_run(self): """Stop running the monitor. Returns: True if the monitor is stopped. False otherwise. """ was_running = self.is_running self.is_running = False return was_running async def get_peak_memory_info(self): """Return the tuple of the peak memory usage and the top n process information during the peak memory usage. """ return self.peak_memory_usage, self.peak_top_n_memory_usage current_node_ip = ray._private.worker.global_worker.node_ip_address # Schedule the actor on the current node. memory_monitor_actor = MemoryMonitorActor.options( resources={f"node:{current_node_ip}": 0.001} ).remote( print_interval_s=print_interval_s, record_interval_s=record_interval_s, warning_threshold=warning_threshold, ) print("Waiting for memory monitor actor to be ready...") ray.get(memory_monitor_actor.ready.remote()) print("Memory monitor actor is ready now.") memory_monitor_actor.run.remote() return memory_monitor_actor def setup_tls(): """Sets up required environment variables for tls""" import pytest if sys.platform == "darwin": pytest.skip("Cryptography doesn't install in Mac build pipeline") cert, key = generate_self_signed_tls_certs() temp_dir = tempfile.mkdtemp("ray-test-certs") cert_filepath = os.path.join(temp_dir, "server.crt") key_filepath = os.path.join(temp_dir, "server.key") with open(cert_filepath, "w") as fh: fh.write(cert) with open(key_filepath, "w") as fh: fh.write(key) os.environ["RAY_USE_TLS"] = "1" os.environ["RAY_TLS_SERVER_CERT"] = cert_filepath os.environ["RAY_TLS_SERVER_KEY"] = key_filepath os.environ["RAY_TLS_CA_CERT"] = cert_filepath return key_filepath, cert_filepath, temp_dir def teardown_tls(key_filepath, cert_filepath, temp_dir): os.remove(key_filepath) os.remove(cert_filepath) os.removedirs(temp_dir) del os.environ["RAY_USE_TLS"] del os.environ["RAY_TLS_SERVER_CERT"] del os.environ["RAY_TLS_SERVER_KEY"] del os.environ["RAY_TLS_CA_CERT"] class ResourceKillerActor: """Abstract base class used to implement resource killers for chaos testing.""" def __init__( self, head_node_id, kill_interval_s: float = 60, kill_delay_s: float = 0, max_to_kill: Optional[int] = 2, batch_size_to_kill: int = 1, kill_filter_fn: Optional[Callable] = None, ): self.kill_interval_s = kill_interval_s self.kill_delay_s = kill_delay_s self.is_running = False self.head_node_id = head_node_id # Set to track the killed nodes. self.killed = set() self.done = get_or_create_event_loop().create_future() self.max_to_kill = max_to_kill self.batch_size_to_kill = batch_size_to_kill self.kill_filter_fn = kill_filter_fn self.kill_immediately_after_found = False # -- logger. -- logging.basicConfig(level=logging.INFO) def ready(self): pass async def run(self): self.is_running = True time.sleep(self.kill_delay_s) while self.is_running: to_kills = await self._find_resources_to_kill() if not self.is_running: break if self.kill_immediately_after_found: sleep_interval = 0 else: sleep_interval = random.random() * self.kill_interval_s time.sleep(sleep_interval) results = await asyncio.gather( *[self._kill_resource(*to_kill) for to_kill in to_kills], return_exceptions=True, ) for to_kill, result in zip(to_kills, results): if isinstance(result, Exception): logger.error( f"Failed to kill resource {to_kill}, may retry later. Error: {result}" ) elif result is True: logger.info(f"Successfully killed resource: {to_kill}") self.killed.add(to_kill) if self.max_to_kill is not None and len(self.killed) >= self.max_to_kill: break await asyncio.sleep(self.kill_interval_s - sleep_interval) self.done.set_result(True) await self.stop_run() async def _find_resources_to_kill(self) -> List[Hashable]: """Implemented by subclasses to discover resources to kill. Should return a list of "resources" to kill, which will be passed into _kill_resource. """ raise NotImplementedError async def _kill_resource(self, *args: Hashable) -> bool: """Implemented by subclasses to kill resources. The method should return False or raise an exception if killing the resource failed, in which case it may be retried. """ raise NotImplementedError async def stop_run(self): was_running = self.is_running self.is_running = False return was_running async def get_killed_nodes(self) -> Set[Hashable]: """Get the set of nodes that were killed.""" await self.done return self.killed.copy() class NodeKillerBase(ResourceKillerActor): async def _find_resources_to_kill(self) -> List[Tuple[str, str, str]]: def _resource_from_node_info(n: Dict) -> Tuple[str, str, str]: return (n["NodeID"], n["NodeManagerAddress"], n["NodeManagerPort"]) nodes_to_kill = [] while not nodes_to_kill and self.is_running: worker_nodes = [ node for node in ray.nodes() if node["Alive"] and (node["NodeID"] != self.head_node_id) and (_resource_from_node_info(node) not in self.killed) ] if self.kill_filter_fn: candidates = list(filter(self.kill_filter_fn(), worker_nodes)) else: candidates = worker_nodes # Ensure at least one worker node remains alive. if len(worker_nodes) < self.batch_size_to_kill + 1: # Give the cluster some time to start. await asyncio.sleep(1) continue # Collect nodes to kill, limited by batch size. for candidate in candidates[: self.batch_size_to_kill]: nodes_to_kill.append(_resource_from_node_info(candidate)) return nodes_to_kill @ray.remote(num_cpus=0) class RayletKiller(NodeKillerBase): async def _kill_resource( self, node_id: str, node_to_kill_ip: str, node_to_kill_port: int ): if node_to_kill_port is None: return False self._kill_raylet(node_to_kill_ip, node_to_kill_port, graceful=False) return True def _kill_raylet(self, ip, port, graceful=False): import grpc from grpc._channel import _InactiveRpcError from ray.core.generated import node_manager_pb2_grpc raylet_address = build_address(ip, port) channel = grpc.insecure_channel(raylet_address) stub = node_manager_pb2_grpc.NodeManagerServiceStub(channel) try: stub.ShutdownRaylet( node_manager_pb2.ShutdownRayletRequest(graceful=graceful) ) except _InactiveRpcError: assert not graceful @ray.remote(num_cpus=0) class EC2InstanceTerminator(NodeKillerBase): async def _kill_resource( self, node_id: str, node_to_kill_ip: str, node_to_kill_port: int ): if node_to_kill_ip is None: return False _terminate_ec2_instance(node_to_kill_ip) return True @ray.remote(num_cpus=0) class EC2InstanceTerminatorWithGracePeriod(NodeKillerBase): def __init__(self, *args, grace_period_s: int = 30, **kwargs): super().__init__(*args, **kwargs) self._grace_period_s = grace_period_s async def _kill_resource( self, node_id: str, node_to_kill_ip: str, node_to_kill_port: int ): self._drain_node(node_id) await asyncio.sleep(self._grace_period_s) # Anyscale extends the drain deadline if you shut down the instance # directly. To work around this, we force-stop Ray on the node. # Anyscale should then terminate it shortly after without updating # the drain deadline. _execute_command_on_node("ray stop --force", node_to_kill_ip) return True def _drain_node(self, node_id: str) -> None: # We need to lazily import this object. Otherwise, Ray can't serialize the # class. from ray.core.generated import autoscaler_pb2 assert ray.NodeID.from_hex(node_id) != ray.NodeID.nil() logging.info(f"Draining node {node_id=}") address = services.canonicalize_bootstrap_address_or_die(addr="auto") gcs_client = ray._raylet.GcsClient(address=address) deadline_timestamp_ms = (time.time_ns() // 1e6) + (self._grace_period_s * 1e3) try: is_accepted, _ = gcs_client.drain_node( node_id, autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"), "", deadline_timestamp_ms, ) except ray.exceptions.RayError as e: logger.error(f"Failed to drain node {node_id=}") raise e assert is_accepted, "Drain node request was rejected" @ray.remote(num_cpus=0) class WorkerKillerActor(ResourceKillerActor): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # Kill worker immediately so that the task does # not finish successfully on its own. self.kill_immediately_after_found = True from ray.util.state.api import StateApiClient from ray.util.state.common import ListApiOptions self.client = StateApiClient() self.task_options = ListApiOptions( filters=[ ("state", "=", "RUNNING"), ("name", "!=", "WorkerKillActor.run"), ] ) async def _find_resources_to_kill(self) -> List[Tuple[str, int, str]]: from ray.util.state.common import StateResource process_to_kill_task_id = None process_to_kill_pid = None process_to_kill_node_id = None while process_to_kill_pid is None and self.is_running: tasks = self.client.list( StateResource.TASKS, options=self.task_options, raise_on_missing_output=False, ) if self.kill_filter_fn is not None: tasks = list(filter(self.kill_filter_fn(), tasks)) for task in tasks: if task.worker_id is not None and task.node_id is not None: process_to_kill_task_id = task.task_id process_to_kill_pid = task.worker_pid process_to_kill_node_id = task.node_id break # Give the cluster some time to start. await asyncio.sleep(0.1) return [(process_to_kill_task_id, process_to_kill_pid, process_to_kill_node_id)] async def _kill_resource( self, process_to_kill_task_id: str, process_to_kill_pid: int, process_to_kill_node_id: str, ): if process_to_kill_pid is None: return False @ray.remote def kill_process(pid: int): proc = psutil.Process(pid) proc.kill() label_selector = {ray._raylet.RAY_NODE_ID_KEY: process_to_kill_node_id} await kill_process.options(label_selector=label_selector).remote( process_to_kill_pid ) return True def get_and_run_resource_killer( resource_killer_cls, kill_interval_s, namespace=None, lifetime=None, no_start=False, max_to_kill=2, batch_size_to_kill=1, kill_delay_s=0, kill_filter_fn=None, ): assert ray.is_initialized(), "The API is only available when Ray is initialized." head_node_id = ray.get_runtime_context().get_node_id() # Schedule the actor on the current node. resource_killer = resource_killer_cls.options( label_selector={ray._raylet.RAY_NODE_ID_KEY: head_node_id}, namespace=namespace, name="ResourceKiller", lifetime=lifetime, ).remote( head_node_id, kill_interval_s=kill_interval_s, kill_delay_s=kill_delay_s, max_to_kill=max_to_kill, batch_size_to_kill=batch_size_to_kill, kill_filter_fn=kill_filter_fn, ) print("Waiting for ResourceKiller to be ready...") ray.get(resource_killer.ready.remote()) print("ResourceKiller is ready now.") if not no_start: resource_killer.run.remote() return resource_killer def get_actor_node_id(actor_handle: "ray.actor.ActorHandle") -> str: return ray.get( actor_handle.__ray_call__.remote( lambda self: ray.get_runtime_context().get_node_id() ) ) @contextmanager def chdir(d: str): old_dir = os.getcwd() os.chdir(d) try: yield finally: os.chdir(old_dir) def test_get_directory_size_bytes(): with tempfile.TemporaryDirectory() as tmp_dir, chdir(tmp_dir): assert ray._private.utils.get_directory_size_bytes(tmp_dir) == 0 with open("test_file", "wb") as f: f.write(os.urandom(100)) assert ray._private.utils.get_directory_size_bytes(tmp_dir) == 100 with open("test_file_2", "wb") as f: f.write(os.urandom(50)) assert ray._private.utils.get_directory_size_bytes(tmp_dir) == 150 os.mkdir("subdir") with open("subdir/subdir_file", "wb") as f: f.write(os.urandom(2)) assert ray._private.utils.get_directory_size_bytes(tmp_dir) == 152 def check_local_files_gced(cluster): for node in cluster.list_all_nodes(): for subdir in ["conda", "pip", "working_dir_files", "py_modules_files"]: all_files = os.listdir( os.path.join(node.get_runtime_env_dir_path(), subdir) ) # Check that there are no files remaining except for .lock files # and generated requirements.txt files. # Note: On Windows the top folder is not deleted as it is in use. # TODO(architkulkarni): these files should get cleaned up too! items = list(filter(lambda f: not f.endswith((".lock", ".txt")), all_files)) if len(items) > 0: print(f"runtime_env files not GC'd from subdir '{subdir}': {items}") return False return True def generate_runtime_env_dict(field, spec_format, tmp_path, pip_list=None): if pip_list is None: pip_list = ["pip-install-test==0.5"] if field == "conda": conda_dict = {"dependencies": ["pip", {"pip": pip_list}]} if spec_format == "file": conda_file = tmp_path / f"environment-{hash(str(pip_list))}.yml" conda_file.write_text(yaml.dump(conda_dict)) conda = str(conda_file) elif spec_format == "python_object": conda = conda_dict runtime_env = {"conda": conda} elif field == "pip": if spec_format == "file": pip_file = tmp_path / f"requirements-{hash(str(pip_list))}.txt" pip_file.write_text("\n".join(pip_list)) pip = str(pip_file) elif spec_format == "python_object": pip = pip_list runtime_env = {"pip": pip} return runtime_env def check_spilled_mb(address, spilled=None, restored=None, fallback=None): def ok(): s = memory_summary(address=address["address"], stats_only=True) print(s) if restored: if "Restored {} MiB".format(restored) not in s: return False else: if "Restored" in s: return False if spilled: if not isinstance(spilled, list): spilled_lst = [spilled] else: spilled_lst = spilled found = False for n in spilled_lst: if "Spilled {} MiB".format(n) in s: found = True if not found: return False else: if "Spilled" in s: return False if fallback: if "Plasma filesystem mmap usage: {} MiB".format(fallback) not in s: return False else: if "Plasma filesystem mmap usage:" in s: return False return True wait_for_condition(ok, timeout=3, retry_interval_ms=1000) def no_resource_leaks_excluding_node_resources(): cluster_resources = ray.cluster_resources() available_resources = ray.available_resources() for r in ray.cluster_resources(): if "node" in r: del cluster_resources[r] del available_resources[r] return cluster_resources == available_resources def job_hook(**kwargs): """Function called by reflection by test_cli_integration.""" cmd = " ".join(kwargs["entrypoint"]) print(f"hook intercepted: {cmd}") sys.exit(0) def wandb_setup_api_key_hook(): """ Example external hook to set up W&B API key in WandbIntegrationTest.testWandbLoggerConfig """ return "abcd" # Get node stats from node manager. def get_node_stats(raylet, num_retry=5, timeout=2): import grpc from ray._private.grpc_utils import init_grpc_channel from ray.core.generated import node_manager_pb2_grpc raylet_address = build_address( raylet["NodeManagerAddress"], raylet["NodeManagerPort"] ) channel = init_grpc_channel(raylet_address) stub = node_manager_pb2_grpc.NodeManagerServiceStub(channel) for _ in range(num_retry): try: reply = stub.GetNodeStats( node_manager_pb2.GetNodeStatsRequest(), timeout=timeout ) break except grpc.RpcError: continue assert reply is not None return reply # Gets resource usage assuming gcs is local. def get_resource_usage(gcs_address, timeout=10): from ray._private.grpc_utils import init_grpc_channel from ray.core.generated import gcs_service_pb2_grpc if not gcs_address: gcs_address = ray.worker._global_node.gcs_address gcs_channel = init_grpc_channel( gcs_address, ray_constants.GLOBAL_GRPC_OPTIONS, asynchronous=False ) gcs_node_resources_stub = gcs_service_pb2_grpc.NodeResourceInfoGcsServiceStub( gcs_channel ) request = gcs_service_pb2.GetAllResourceUsageRequest() response = gcs_node_resources_stub.GetAllResourceUsage(request, timeout=timeout) resources_batch_data = response.resource_usage_data return resources_batch_data # Gets the load metrics report assuming gcs is local. def get_load_metrics_report(webui_url): webui_url = format_web_url(webui_url) response = requests.get(f"{webui_url}/api/cluster_status") response.raise_for_status() return response.json()["data"]["clusterStatus"]["loadMetricsReport"] # Send a RPC to the raylet to have it self-destruct its process. def kill_raylet(raylet, graceful=False): import grpc from grpc._channel import _InactiveRpcError from ray.core.generated import node_manager_pb2_grpc raylet_address = build_address( raylet["NodeManagerAddress"], raylet["NodeManagerPort"] ) channel = grpc.insecure_channel(raylet_address) stub = node_manager_pb2_grpc.NodeManagerServiceStub(channel) try: stub.ShutdownRaylet(node_manager_pb2.ShutdownRayletRequest(graceful=graceful)) except _InactiveRpcError: assert not graceful def get_gcs_memory_used(): import psutil m = { proc.info["name"]: proc.info["memory_info"].rss for proc in psutil.process_iter(["status", "name", "memory_info"]) if ( proc.info["status"] not in (psutil.STATUS_ZOMBIE, psutil.STATUS_DEAD) and proc.info["name"] in ("gcs_server", "redis-server") ) } assert "gcs_server" in m return sum(m.values()) def safe_write_to_results_json( result: dict, default_file_name: str = "/tmp/release_test_output.json", env_var: Optional[str] = "TEST_OUTPUT_JSON", ): """ Safe (atomic) write to file to guard against malforming the json if the job gets interrupted in the middle of writing. """ test_output_json = os.environ.get(env_var, default_file_name) test_output_json_tmp = f"{test_output_json}.tmp.{str(uuid.uuid4())}" with open(test_output_json_tmp, "wt") as f: json.dump(result, f) f.flush() os.replace(test_output_json_tmp, test_output_json) logger.info(f"Wrote results to {test_output_json}") logger.info(json.dumps(result)) def get_current_unused_port(): """ Returns a port number that is not currently in use. This is useful for testing when we need to bind to a port but don't care which one. Returns: A port number that is not currently in use. (Note that this port might become used by the time you try to bind to it.) """ sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) # Bind the socket to a local address with a random port number sock.bind(("localhost", 0)) port = sock.getsockname()[1] sock.close() return port # Global counter to test different return values # for external_ray_cluster_activity_hook1. ray_cluster_activity_hook_counter = 0 ray_cluster_activity_hook_5_counter = 0 def external_ray_cluster_activity_hook1(): """ Example external hook for test_component_activities_hook. Returns valid response and increments counter in `reason` field on each call. """ global ray_cluster_activity_hook_counter ray_cluster_activity_hook_counter += 1 from pydantic import BaseModel, Extra class TestRayActivityResponse(BaseModel, extra=Extra.allow): """ Redefinition of dashboard.modules.api.api_head.RayActivityResponse used in test_component_activities_hook to mimic typical usage of redefining or extending response type. """ is_active: str reason: Optional[str] = None timestamp: float return { "test_component1": TestRayActivityResponse( is_active="ACTIVE", reason=f"Counter: {ray_cluster_activity_hook_counter}", timestamp=datetime.now().timestamp(), ) } def external_ray_cluster_activity_hook2(): """ Example external hook for test_component_activities_hook. Returns invalid output because the value of `test_component2` should be of type RayActivityResponse. """ return {"test_component2": "bad_output"} def external_ray_cluster_activity_hook3(): """ Example external hook for test_component_activities_hook. Returns invalid output because return type is not Dict[str, RayActivityResponse] """ return "bad_output" def external_ray_cluster_activity_hook4(): """ Example external hook for test_component_activities_hook. Errors during execution. """ raise Exception("Error in external cluster activity hook") def external_ray_cluster_activity_hook5(): """ Example external hook for test_component_activities_hook. Returns valid response and increments counter in `reason` field on each call. """ global ray_cluster_activity_hook_5_counter ray_cluster_activity_hook_5_counter += 1 return { "test_component5": { "is_active": "ACTIVE", "reason": f"Counter: {ray_cluster_activity_hook_5_counter}", "timestamp": datetime.now().timestamp(), } } # TODO(rickyx): We could remove this once we unify the autoscaler v1 and v2 # code path for ray status def reset_autoscaler_v2_enabled_cache(): import ray.autoscaler.v2.utils as u u.cached_is_autoscaler_v2 = None def _terminate_ec2_instance(node_ip: str) -> None: logging.info(f"Terminating instance {node_ip}") # This command uses IMDSv2 to get the host instance id and region. # After that it terminates itself using aws cli. command = ( 'instanceId=$(curl -H "X-aws-ec2-metadata-token: $TOKEN" http://169.254.169.254/latest/meta-data/instance-id/);' # noqa: E501 'region=$(curl -H "X-aws-ec2-metadata-token: $TOKEN" http://169.254.169.254/latest/meta-data/placement/region);' # noqa: E501 "aws ec2 terminate-instances --region $region --instance-ids $instanceId" # noqa: E501 ) _execute_command_on_node(command, node_ip) def _execute_command_on_node(command: str, node_ip: str): logging.debug(f"Executing command on node {node_ip}: {command}") multi_line_command = ( 'TOKEN=$(curl -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 21600");' # noqa: E501 f"{command}" ) # This is a feature on Anyscale platform that enables # easy ssh access to worker nodes. ssh_command = f"ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -p 2222 ray@{node_ip} '{multi_line_command}'" # noqa: E501 # Strip library path overrides so that the system ssh binary # doesn't pick up a conflicting OpenSSL from the Ray/conda env. env = os.environ.copy() env.pop("LD_LIBRARY_PATH", None) env.pop("DYLD_LIBRARY_PATH", None) try: subprocess.run( ssh_command, shell=True, capture_output=True, text=True, check=True, env=env, ) except subprocess.CalledProcessError as e: logger.error( f"Command failed on node {node_ip}: {command}, " f"exit code: {e.returncode}, stderr: {e.stderr}" ) raise RPC_FAILURE_MAP = { "request": { "req_failure_prob": 100, "resp_failure_prob": 0, "in_flight_failure_prob": 0, }, "response": { "req_failure_prob": 0, "resp_failure_prob": 100, "in_flight_failure_prob": 0, }, "in_flight": { "req_failure_prob": 0, "resp_failure_prob": 0, "in_flight_failure_prob": 100, }, } RPC_FAILURE_TYPES = list(RPC_FAILURE_MAP.keys())