wehub-resource-sync
1724 lines
54 KiB
Python
1724 lines
54 KiB
Python
import asyncio
|
||
import datetime
|
||
import glob
|
||
import os
|
||
import random
|
||
import socket
|
||
import threading
|
||
import time
|
||
from contextlib import asynccontextmanager
|
||
from copy import copy, deepcopy
|
||
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
|
||
from unittest.mock import Mock
|
||
|
||
import grpc
|
||
import httpx
|
||
import requests
|
||
from starlette.requests import Request
|
||
|
||
import ray
|
||
from ray import serve
|
||
from ray._common.network_utils import build_address
|
||
from ray._common.test_utils import (
|
||
PrometheusTimeseries,
|
||
fetch_prometheus_metric_timeseries,
|
||
wait_for_condition,
|
||
)
|
||
from ray._common.utils import TimerBase
|
||
from ray.actor import ActorHandle
|
||
from ray.serve._private.client import ServeControllerClient
|
||
from ray.serve._private.common import (
|
||
CreatePlacementGroupRequest,
|
||
DeploymentID,
|
||
DeploymentStatus,
|
||
ReplicaID,
|
||
RequestProtocol,
|
||
RunningReplicaInfo,
|
||
)
|
||
from ray.serve._private.constants import (
|
||
RAY_SERVE_ENABLE_HA_PROXY,
|
||
RAY_SERVE_HAPROXY_SOCKET_PATH,
|
||
SERVE_DEFAULT_APP_NAME,
|
||
SERVE_NAMESPACE,
|
||
)
|
||
from ray.serve._private.deployment_info import DeploymentInfo
|
||
from ray.serve._private.deployment_scheduler import ReplicaSchedulingRequest
|
||
from ray.serve._private.deployment_state import (
|
||
ALL_REPLICA_STATES,
|
||
DeploymentVersion,
|
||
ReplicaStartupStatus,
|
||
ReplicaState,
|
||
)
|
||
from ray.serve._private.haproxy import HAProxyApi
|
||
from ray.serve._private.proxy import DRAINING_MESSAGE
|
||
from ray.serve._private.replica_result import ReplicaResult
|
||
from ray.serve._private.request_router import (
|
||
PendingRequest,
|
||
RunningReplica,
|
||
)
|
||
from ray.serve._private.usage import ServeUsageTag
|
||
from ray.serve.context import _get_global_client
|
||
from ray.serve.gang import GangContext
|
||
from ray.serve.generated import serve_pb2, serve_pb2_grpc
|
||
from ray.serve.schema import ApplicationStatus, ReplicaRank, TargetGroup
|
||
from ray.util.state import list_actors
|
||
|
||
TELEMETRY_ROUTE_PREFIX = "/telemetry"
|
||
STORAGE_ACTOR_NAME = "storage"
|
||
PROMETHEUS_METRICS_TIMEOUT_S = 5
|
||
|
||
|
||
def skip_if_haproxy(reason: str):
|
||
"""Skip a test when the HAProxy ingress is enabled.
|
||
|
||
The HAProxy ingress runs as a separate premerge step with
|
||
RAY_SERVE_ENABLE_HA_PROXY=1. Some tests exercise behavior HAProxy does not
|
||
support yet (e.g. gRPC ingress) or that probes the native Serve proxy
|
||
directly. Mark those with this decorator instead of maintaining a separate
|
||
test allowlist. The test still runs in the non-HAProxy steps.
|
||
"""
|
||
import pytest
|
||
|
||
return pytest.mark.skipif(
|
||
RAY_SERVE_ENABLE_HA_PROXY, reason=f"HAProxy ingress: {reason}"
|
||
)
|
||
|
||
|
||
# Global variable that is fetched during controller recovery that
|
||
# marks (simulates) which replicas have died since controller first
|
||
# recovered a list of live replica names.
|
||
# NOTE(zcin): This is necessary because the replica's `recover()` method
|
||
# is called in the controller's init function, instead of in the control
|
||
# loop, so we can't "mark" a replica dead through a method. This global
|
||
# state is cleared after each test that uses the fixtures in this file.
|
||
dead_replicas_context = set()
|
||
# Replicas registered in this set will report `was_initialized() == False` to
|
||
# the controller during recovery, simulating the case where a previous
|
||
# controller crashed before the actor finished its initial setup.
|
||
uninitialized_replicas_context = set()
|
||
replica_rank_context: Dict[str, ReplicaRank] = {}
|
||
|
||
|
||
class MockTimer(TimerBase):
|
||
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):
|
||
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
|
||
|
||
|
||
class MockAsyncTimer:
|
||
def __init__(self, start_time: Optional[float] = 0):
|
||
self.reset(start_time=start_time)
|
||
self._num_sleepers = 0
|
||
|
||
def reset(self, start_time: 0):
|
||
self._curr = start_time
|
||
|
||
def time(self) -> float:
|
||
return self._curr
|
||
|
||
async def sleep(self, amt: float):
|
||
self._num_sleepers += 1
|
||
end = self._curr + amt
|
||
|
||
# Give up the event loop
|
||
while self._curr < end:
|
||
await asyncio.sleep(0)
|
||
|
||
self._num_sleepers -= 1
|
||
|
||
def advance(self, amt: float):
|
||
self._curr += amt
|
||
|
||
def num_sleepers(self):
|
||
return self._num_sleepers
|
||
|
||
|
||
class MockKVStore:
|
||
def __init__(self):
|
||
self.store = dict()
|
||
|
||
def put(self, key: str, val: Any) -> bool:
|
||
if not isinstance(key, str):
|
||
raise TypeError("key must be a string, got: {}.".format(type(key)))
|
||
self.store[key] = val
|
||
return True
|
||
|
||
def get(self, key: str) -> Any:
|
||
if not isinstance(key, str):
|
||
raise TypeError("key must be a string, got: {}.".format(type(key)))
|
||
return self.store.get(key, None)
|
||
|
||
def delete(self, key: str) -> bool:
|
||
if not isinstance(key, str):
|
||
raise TypeError("key must be a string, got: {}.".format(type(key)))
|
||
|
||
if key in self.store:
|
||
del self.store[key]
|
||
return True
|
||
|
||
return False
|
||
|
||
|
||
class MockClusterNodeInfoCache:
|
||
def __init__(self):
|
||
self.alive_node_ids = set()
|
||
self.total_resources_per_node = dict()
|
||
self.available_resources_per_node = dict()
|
||
self.draining_nodes = dict()
|
||
self.node_labels = dict()
|
||
|
||
def get_alive_node_ids(self):
|
||
return self.alive_node_ids
|
||
|
||
def get_draining_nodes(self):
|
||
return self.draining_nodes
|
||
|
||
def get_active_node_ids(self):
|
||
return self.alive_node_ids - set(self.draining_nodes)
|
||
|
||
def get_node_az(self, node_id):
|
||
return None
|
||
|
||
def get_available_resources_per_node(self):
|
||
return self.available_resources_per_node
|
||
|
||
def get_total_resources_per_node(self):
|
||
return self.total_resources_per_node
|
||
|
||
def add_node(self, node_id: str, resources: Dict = None, labels: Dict = None):
|
||
self.alive_node_ids.add(node_id)
|
||
self.total_resources_per_node[node_id] = deepcopy(resources) or {}
|
||
self.available_resources_per_node[node_id] = deepcopy(resources) or {}
|
||
self.node_labels[node_id] = labels or {}
|
||
|
||
def set_available_resources_per_node(self, node_id: str, resources: Dict):
|
||
self.available_resources_per_node[node_id] = deepcopy(resources)
|
||
|
||
def get_node_labels(self, node_id: str):
|
||
return self.node_labels.get(node_id, {})
|
||
|
||
|
||
# Default value used by FakeRunningReplica when the test doesn't override.
|
||
FAKE_REPLICA_DEFAULT_MAX_ONGOING_REQUESTS = 10
|
||
# Every FakeRunningReplica is stamped with this deployment id. Router
|
||
# tests never care about app/name specifically; cross-test consistency
|
||
# matters more than differentiation.
|
||
FAKE_REPLICA_DEPLOYMENT_ID = DeploymentID(name="TEST_DEPLOYMENT")
|
||
|
||
|
||
class FakeRunningReplica(RunningReplica):
|
||
"""In-memory ``RunningReplica`` stand-in for request-router unit tests.
|
||
|
||
Shared between ``test_pow_2_request_router.py`` and
|
||
``test_consistent_hash_router.py``. Supports enough knobs for pow-2's
|
||
locality / cancellation / re-probe tests; consistent-hash only uses
|
||
the queue-length response hooks.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
replica_unique_id: str,
|
||
*,
|
||
node_id: str = "",
|
||
availability_zone: Optional[str] = None,
|
||
reset_after_response: bool = False,
|
||
model_ids: Optional[Set[str]] = None,
|
||
sleep_time_s: float = 0.0,
|
||
max_ongoing_requests: int = FAKE_REPLICA_DEFAULT_MAX_ONGOING_REQUESTS,
|
||
):
|
||
self._replica_id = ReplicaID(
|
||
unique_id=replica_unique_id,
|
||
deployment_id=FAKE_REPLICA_DEPLOYMENT_ID,
|
||
)
|
||
self._node_id = node_id
|
||
self._availability_zone = availability_zone
|
||
self._queue_len = 0
|
||
self._max_ongoing_requests = max_ongoing_requests
|
||
self._has_queue_len_response = asyncio.Event()
|
||
self._reset_after_response = reset_after_response
|
||
self._model_ids = model_ids or set()
|
||
self._sleep_time_s = sleep_time_s
|
||
self._exception: Optional[Exception] = None
|
||
|
||
self.get_queue_len_was_cancelled = False
|
||
self.queue_len_deadline_history = list()
|
||
self.num_get_queue_len_calls = 0
|
||
|
||
@property
|
||
def replica_id(self) -> ReplicaID:
|
||
return self._replica_id
|
||
|
||
@property
|
||
def node_id(self) -> str:
|
||
return self._node_id
|
||
|
||
@property
|
||
def availability_zone(self) -> Optional[str]:
|
||
return self._availability_zone
|
||
|
||
@property
|
||
def multiplexed_model_ids(self) -> Set[str]:
|
||
return self._model_ids
|
||
|
||
def update_replica_info(self, replica_info: RunningReplicaInfo) -> None:
|
||
self._model_ids = set(replica_info.multiplexed_model_ids)
|
||
|
||
@property
|
||
def max_ongoing_requests(self) -> int:
|
||
return self._max_ongoing_requests
|
||
|
||
def set_queue_len_response(
|
||
self,
|
||
queue_len: int,
|
||
exception: Optional[Exception] = None,
|
||
):
|
||
self._queue_len = queue_len
|
||
self._exception = exception
|
||
self._has_queue_len_response.set()
|
||
|
||
def push_proxy_handle(self, handle: ActorHandle):
|
||
pass
|
||
|
||
async def get_queue_len(self, *, deadline_s: float) -> int:
|
||
self.num_get_queue_len_calls += 1
|
||
self.queue_len_deadline_history.append(deadline_s)
|
||
try:
|
||
while not self._has_queue_len_response.is_set():
|
||
await self._has_queue_len_response.wait()
|
||
|
||
if self._sleep_time_s > 0:
|
||
await asyncio.sleep(self._sleep_time_s)
|
||
|
||
if self._reset_after_response:
|
||
self._has_queue_len_response.clear()
|
||
|
||
if self._exception is not None:
|
||
raise self._exception
|
||
|
||
return self._queue_len
|
||
except asyncio.CancelledError:
|
||
self.get_queue_len_was_cancelled = True
|
||
raise
|
||
|
||
def try_send_request(
|
||
self, pr: PendingRequest, with_rejection: bool
|
||
) -> ReplicaResult:
|
||
raise NotImplementedError()
|
||
|
||
def send_request_with_rejection(self, pr: PendingRequest) -> ReplicaResult:
|
||
raise NotImplementedError()
|
||
|
||
|
||
class FakeRemoteFunction:
|
||
def remote(self):
|
||
pass
|
||
|
||
|
||
class MockActorHandle:
|
||
def __init__(self, **kwargs):
|
||
self._options = kwargs
|
||
self._actor_id = "fake_id"
|
||
self.initialize_and_get_metadata_called = False
|
||
self.is_allocated_called = False
|
||
|
||
@property
|
||
def initialize_and_get_metadata(self):
|
||
self.initialize_and_get_metadata_called = True
|
||
# return a mock object so that we can call `remote()` on it.
|
||
return FakeRemoteFunction()
|
||
|
||
@property
|
||
def is_allocated(self):
|
||
self.is_allocated_called = True
|
||
return FakeRemoteFunction()
|
||
|
||
|
||
class MockActorClass:
|
||
def __init__(self):
|
||
self._init_args = ()
|
||
self._options = dict()
|
||
|
||
def options(self, **kwargs):
|
||
res = copy(self)
|
||
|
||
for k, v in kwargs.items():
|
||
res._options[k] = v
|
||
|
||
return res
|
||
|
||
def remote(self, *args) -> MockActorHandle:
|
||
return MockActorHandle(init_args=args, **self._options)
|
||
|
||
|
||
class MockPlacementGroup:
|
||
def __init__(self, request: CreatePlacementGroupRequest):
|
||
self._bundles = request.bundles
|
||
self._strategy = request.strategy
|
||
self._soft_target_node_id = request.target_node_id
|
||
self._name = request.name
|
||
self._lifetime = "detached"
|
||
|
||
|
||
class MockDeploymentHandle:
|
||
def __init__(self, deployment_name: str, app_name: str = SERVE_DEFAULT_APP_NAME):
|
||
self._deployment_name = deployment_name
|
||
self._app_name = app_name
|
||
self._protocol = RequestProtocol.UNDEFINED
|
||
self._running_replicas_populated = False
|
||
self._initialized = False
|
||
|
||
def is_initialized(self):
|
||
return self._initialized
|
||
|
||
def _init(self):
|
||
if self._initialized:
|
||
raise RuntimeError("already initialized")
|
||
|
||
self._initialized = True
|
||
|
||
def options(self, *args, **kwargs):
|
||
return self
|
||
|
||
def __eq__(self, dep: Tuple[str]):
|
||
other_deployment_name, other_app_name = dep
|
||
return (
|
||
self._deployment_name == other_deployment_name
|
||
and self._app_name == other_app_name
|
||
)
|
||
|
||
def _set_request_protocol(self, protocol: RequestProtocol):
|
||
self._protocol = protocol
|
||
|
||
def _get_or_create_router(self):
|
||
pass
|
||
|
||
def running_replicas_populated(self) -> bool:
|
||
return self._running_replicas_populated
|
||
|
||
def set_running_replicas_populated(self, val: bool):
|
||
self._running_replicas_populated = val
|
||
|
||
|
||
class MockDeploymentActorWrapper:
|
||
"""Mock for DeploymentActorWrapper with per-instance setters."""
|
||
|
||
def __init__(
|
||
self,
|
||
deployment_id: DeploymentID,
|
||
config,
|
||
code_version: str,
|
||
recovered_handle=None,
|
||
):
|
||
self._deployment_id = deployment_id
|
||
self._config = config
|
||
self._code_version = code_version
|
||
self._handle = recovered_handle
|
||
self._ready = False # Recovered starts pending until set_ready()
|
||
self._start_error_msg: Optional[str] = None
|
||
self._check_ready_error_msg: Optional[str] = None
|
||
self.killed = False
|
||
self.pending_killed = False
|
||
self._start_fail_count = 0
|
||
self._start_fail_msg = None
|
||
self._health_ok = True
|
||
self.reset_health_state_after_running_count = 0
|
||
|
||
@property
|
||
def actor_logical_name(self) -> str:
|
||
return self._config.name
|
||
|
||
@property
|
||
def code_version(self) -> str:
|
||
return self._code_version
|
||
|
||
def set_ready(self):
|
||
self._ready = True
|
||
|
||
def set_start_error(self, error_msg: str):
|
||
"""Make start() fail with this error (persists until cleared)."""
|
||
self._start_error_msg = error_msg
|
||
|
||
def set_check_ready_error(self, error_msg: str):
|
||
self._check_ready_error_msg = error_msg
|
||
|
||
def set_failed_to_start(self, error_msg: str = "constructor failed"):
|
||
"""Simulate start failure (like replica's set_failed_to_start). check_ready will return this error."""
|
||
self._check_ready_error_msg = error_msg
|
||
|
||
def start(self, deployment_runtime_env=None):
|
||
if self._start_error_msg is not None:
|
||
return False, self._start_error_msg
|
||
# Match real behavior: created actor starts as pending until ready.
|
||
return True, None
|
||
|
||
def check_ready(self):
|
||
if self._check_ready_error_msg is not None:
|
||
return False, self._check_ready_error_msg
|
||
if self._ready:
|
||
return True, None
|
||
return False, None
|
||
|
||
def set_health_ok(self, ok: bool) -> None:
|
||
"""If False, ``check_health`` returns False (failed health reconciliation)."""
|
||
self._health_ok = ok
|
||
|
||
def reset_health_state_after_running(self) -> None:
|
||
"""Match DeploymentActorWrapper; increments counter for unit tests."""
|
||
self.reset_health_state_after_running_count += 1
|
||
|
||
def check_health(self) -> bool:
|
||
"""Match DeploymentActorWrapper; controlled via ``set_health_ok``."""
|
||
return self._health_ok
|
||
|
||
def kill(self) -> None:
|
||
self.killed = True
|
||
|
||
|
||
class MockReplicaActorWrapper:
|
||
def __init__(
|
||
self,
|
||
replica_id: ReplicaID,
|
||
version: DeploymentVersion,
|
||
):
|
||
self._replica_id = replica_id
|
||
self._actor_name = replica_id.to_full_id_str()
|
||
# Will be set when `start()` is called.
|
||
self.started = False
|
||
# Will be set when `recover()` is called.
|
||
self.recovering = False
|
||
# Will be set when `start()` is called.
|
||
self.version = version
|
||
# Initial state for a replica is PENDING_ALLOCATION.
|
||
self.status = ReplicaStartupStatus.PENDING_ALLOCATION
|
||
# Will be set when `graceful_stop()` is called.
|
||
self.stopped = False
|
||
# Expected to be set in the test.
|
||
self.done_stopping = False
|
||
# Will be set when `force_stop()` is called.
|
||
self.force_stopped_counter = 0
|
||
# Will be set when `check_health()` is called.
|
||
self.health_check_called = False
|
||
# Returned by the health check.
|
||
self.healthy = True
|
||
self._is_cross_language = False
|
||
self._actor_handle = MockActorHandle()
|
||
self._node_id = None
|
||
self._node_ip = None
|
||
self._node_instance_id = None
|
||
self._node_id_is_set = False
|
||
self._log_file_path = None
|
||
self._actor_id = None
|
||
self._internal_grpc_port = None
|
||
self._http_port = None
|
||
self._pg_bundles = None
|
||
self._initialization_latency_s = -1
|
||
self._docs_path = None
|
||
self._rank = replica_rank_context.get(replica_id.unique_id, None)
|
||
self._assign_rank_callback = None
|
||
self._ingress = False
|
||
self._gang_context = None
|
||
self._gang_pg_index = None
|
||
self._unrecoverable = False
|
||
|
||
@property
|
||
def is_cross_language(self) -> bool:
|
||
return self._is_cross_language
|
||
|
||
@property
|
||
def replica_id(self) -> ReplicaID:
|
||
return self._replica_id
|
||
|
||
@property
|
||
def deployment_name(self) -> str:
|
||
return self._replica_id.deployment_id.name
|
||
|
||
@property
|
||
def actor_handle(self) -> MockActorHandle:
|
||
return self._actor_handle
|
||
|
||
@property
|
||
def max_ongoing_requests(self) -> int:
|
||
return self.version.deployment_config.max_ongoing_requests
|
||
|
||
@property
|
||
def graceful_shutdown_timeout_s(self) -> float:
|
||
return self.version.deployment_config.graceful_shutdown_timeout_s
|
||
|
||
@property
|
||
def health_check_period_s(self) -> float:
|
||
return self.version.deployment_config.health_check_period_s
|
||
|
||
@property
|
||
def health_check_timeout_s(self) -> float:
|
||
return self.version.deployment_config.health_check_timeout_s
|
||
|
||
@property
|
||
def pid(self) -> Optional[int]:
|
||
return None
|
||
|
||
@property
|
||
def actor_id(self) -> Optional[str]:
|
||
return self._actor_id
|
||
|
||
@property
|
||
def worker_id(self) -> Optional[str]:
|
||
return None
|
||
|
||
@property
|
||
def node_id(self) -> Optional[str]:
|
||
if self._node_id_is_set:
|
||
return self._node_id
|
||
if self.status == ReplicaStartupStatus.SUCCEEDED or self.started:
|
||
return "node-id"
|
||
return None
|
||
|
||
@property
|
||
def availability_zone(self) -> Optional[str]:
|
||
return None
|
||
|
||
@property
|
||
def node_ip(self) -> Optional[str]:
|
||
return None
|
||
|
||
@property
|
||
def node_instance_id(self) -> Optional[str]:
|
||
return None
|
||
|
||
@property
|
||
def log_file_path(self) -> Optional[str]:
|
||
return self._log_file_path
|
||
|
||
@property
|
||
def grpc_port(self) -> Optional[int]:
|
||
return None
|
||
|
||
@property
|
||
def placement_group_bundles(self) -> Optional[List[Dict[str, float]]]:
|
||
return None
|
||
|
||
@property
|
||
def initialization_latency_s(self) -> float:
|
||
return self._initialization_latency_s
|
||
|
||
@property
|
||
def reconfigure_start_time(self) -> Optional[float]:
|
||
return None
|
||
|
||
@property
|
||
def last_health_check_latency_ms(self) -> Optional[float]:
|
||
return None
|
||
|
||
@property
|
||
def last_health_check_failed(self) -> bool:
|
||
return False
|
||
|
||
def set_docs_path(self, docs_path: str):
|
||
self._docs_path = docs_path
|
||
|
||
@property
|
||
def docs_path(self) -> Optional[str]:
|
||
return self._docs_path
|
||
|
||
def set_status(self, status: ReplicaStartupStatus):
|
||
self.status = status
|
||
|
||
def set_ready(self, version: DeploymentVersion = None):
|
||
self.status = ReplicaStartupStatus.SUCCEEDED
|
||
# Mirror the real actor: a started replica has allocated a log file.
|
||
self._log_file_path = "serve/replica.log"
|
||
if version:
|
||
self.version_to_be_fetched_from_actor = version
|
||
else:
|
||
self.version_to_be_fetched_from_actor = self.version
|
||
|
||
def set_failed_to_start(self):
|
||
self.status = ReplicaStartupStatus.FAILED
|
||
|
||
def set_done_stopping(self):
|
||
self.done_stopping = True
|
||
|
||
def set_unhealthy(self):
|
||
self.healthy = False
|
||
|
||
def set_starting_version(self, version: DeploymentVersion):
|
||
"""Mocked deployment_worker return version from reconfigure()"""
|
||
self.starting_version = version
|
||
|
||
def set_node_id(self, node_id: str):
|
||
self._node_id = node_id
|
||
self._node_id_is_set = True
|
||
|
||
def set_actor_id(self, actor_id: str):
|
||
self._actor_id = actor_id
|
||
|
||
def start(
|
||
self,
|
||
deployment_info: DeploymentInfo,
|
||
assign_rank_callback: Callable[[ReplicaID], ReplicaRank],
|
||
gang_placement_group=None,
|
||
gang_pg_index=None,
|
||
gang_context=None,
|
||
):
|
||
self.started = True
|
||
self._gang_context = gang_context
|
||
self._gang_pg_index = gang_pg_index
|
||
self._assign_rank_callback = assign_rank_callback
|
||
self._rank = assign_rank_callback(self._replica_id.unique_id, node_id=-1)
|
||
replica_rank_context[self._replica_id.unique_id] = self._rank
|
||
|
||
def _on_scheduled_stub(*args, **kwargs):
|
||
pass
|
||
|
||
return ReplicaSchedulingRequest(
|
||
replica_id=self._replica_id,
|
||
actor_def=Mock(),
|
||
actor_resources={},
|
||
actor_options={"name": "placeholder"},
|
||
actor_init_args=(),
|
||
placement_group_bundles=(
|
||
deployment_info.replica_config.placement_group_bundles
|
||
),
|
||
on_scheduled=_on_scheduled_stub,
|
||
gang_placement_group=gang_placement_group,
|
||
gang_pg_index=gang_pg_index,
|
||
)
|
||
|
||
@property
|
||
def rank(self) -> Optional[ReplicaRank]:
|
||
return self._rank
|
||
|
||
def reconfigure(
|
||
self,
|
||
version: DeploymentVersion,
|
||
rank: ReplicaRank = None,
|
||
):
|
||
self.started = True
|
||
updating = self.version.requires_actor_reconfigure(version)
|
||
self.version = version
|
||
self._rank = rank
|
||
replica_rank_context[self._replica_id.unique_id] = rank
|
||
return updating
|
||
|
||
def recover(self, ingress: bool = False):
|
||
if self.replica_id in dead_replicas_context:
|
||
return False
|
||
|
||
self._ingress = ingress
|
||
self.recovering = True
|
||
self.started = False
|
||
self._rank = replica_rank_context.get(self._replica_id.unique_id, None)
|
||
# Mirror production: the `was_initialized` probe is fired here and
|
||
# observed asynchronously in `check_ready()`. Tests register
|
||
# uninitialized replicas via `uninitialized_replicas_context`.
|
||
self._unrecoverable = self.replica_id in uninitialized_replicas_context
|
||
return True
|
||
|
||
def check_ready(self) -> ReplicaStartupStatus:
|
||
# If the controller's async `was_initialized` probe came back False,
|
||
# report a failed-but-unrecoverable startup so the reconciler drops
|
||
# and replaces the replica without recording a deploy failure.
|
||
if self.recovering and self._unrecoverable:
|
||
return (
|
||
ReplicaStartupStatus.FAILED,
|
||
f"{self._replica_id} was found alive but never finished "
|
||
"its initial setup.",
|
||
)
|
||
|
||
ready = self.status
|
||
self.status = ReplicaStartupStatus.PENDING_INITIALIZATION
|
||
if ready == ReplicaStartupStatus.SUCCEEDED and self.recovering:
|
||
self.recovering = False
|
||
self.started = True
|
||
self.version = self.version_to_be_fetched_from_actor
|
||
return ready, None
|
||
|
||
def resource_requirements(self) -> Tuple[str, str]:
|
||
assert self.started
|
||
return str({"REQUIRED_RESOURCE": 1.0}), str({"AVAILABLE_RESOURCE": 1.0})
|
||
|
||
@property
|
||
def actor_resources(self) -> Dict[str, float]:
|
||
return {"CPU": 0.1}
|
||
|
||
@property
|
||
def available_resources(self) -> Dict[str, float]:
|
||
# Only used to print a warning.
|
||
return {}
|
||
|
||
def graceful_stop(self) -> None:
|
||
# `started` is only set after a successful `check_ready` transition
|
||
# to RUNNING. A replica force-stopped while still RECOVERING (e.g.,
|
||
# because the `was_initialized` probe failed) is a legitimate stop.
|
||
assert self.started or self.recovering
|
||
self.stopped = True
|
||
return self.graceful_shutdown_timeout_s
|
||
|
||
def check_stopped(self) -> bool:
|
||
return self.done_stopping
|
||
|
||
def force_stop(self, log_shutdown_message: bool = False):
|
||
self.force_stopped_counter += 1
|
||
|
||
def check_health(self):
|
||
self.health_check_called = True
|
||
return self.healthy
|
||
|
||
def get_routing_stats(self) -> Dict[str, Any]:
|
||
return {}
|
||
|
||
def get_outbound_deployments(self) -> Optional[List[DeploymentID]]:
|
||
return getattr(self, "_outbound_deployments", None)
|
||
|
||
@property
|
||
def route_patterns(self) -> Optional[List[str]]:
|
||
return None
|
||
|
||
@property
|
||
def gang_context(self) -> Optional[GangContext]:
|
||
return self._gang_context
|
||
|
||
@property
|
||
def unrecoverable(self) -> bool:
|
||
return self._unrecoverable
|
||
|
||
|
||
@serve.deployment
|
||
class GetPID:
|
||
def __call__(self):
|
||
return os.getpid()
|
||
|
||
|
||
get_pid_entrypoint = GetPID.bind()
|
||
|
||
|
||
def check_ray_stopped():
|
||
try:
|
||
requests.get("http://localhost:8265/api/ray/version")
|
||
return False
|
||
except Exception:
|
||
return True
|
||
|
||
|
||
def check_ray_started():
|
||
return requests.get("http://localhost:8265/api/ray/version").status_code == 200
|
||
|
||
|
||
def check_deployment_status(
|
||
name: str, expected_status: DeploymentStatus, app_name=SERVE_DEFAULT_APP_NAME
|
||
) -> bool:
|
||
app_status = serve.status().applications[app_name]
|
||
assert app_status.deployments[name].status == expected_status
|
||
return True
|
||
|
||
|
||
def get_num_alive_replicas(
|
||
deployment_name: str, app_name: str = SERVE_DEFAULT_APP_NAME
|
||
) -> int:
|
||
"""Get the replicas currently running for the given deployment."""
|
||
|
||
dep_id = DeploymentID(name=deployment_name, app_name=app_name)
|
||
actors = list_actors(
|
||
filters=[
|
||
("class_name", "=", dep_id.to_replica_actor_class_name()),
|
||
("state", "=", "ALIVE"),
|
||
]
|
||
)
|
||
return len(actors)
|
||
|
||
|
||
def expected_proxy_actors(num_proxy_nodes: int = 1) -> Dict[str, int]:
|
||
"""Proxy actors expected ALIVE by class name for the given number of proxy nodes.
|
||
|
||
Natively each proxy node runs one ProxyActor. Under HAProxy each proxy node runs an
|
||
HAProxyManager and the head node also runs a single fallback ProxyActor. Set-based
|
||
callers can take the keys, count-based callers use the counts directly.
|
||
"""
|
||
if RAY_SERVE_ENABLE_HA_PROXY:
|
||
return {"HAProxyManager": num_proxy_nodes, "ProxyActor": 1}
|
||
return {"ProxyActor": num_proxy_nodes}
|
||
|
||
|
||
def wait_for_haproxy_routing_to_replica(timeout: int = 30):
|
||
"""Block until HAProxy marks a replica's primary backend server UP.
|
||
|
||
serve.run returns once replicas are RUNNING, but until a replica's
|
||
direct-ingress server passes HAProxy's health check, HAProxy serves requests
|
||
from the head-node backup fallback proxy. That path adds proxy and router
|
||
spans, so trace-topology assertions must wait for direct routing first. No-op
|
||
without HAProxy, which has no fallback to race.
|
||
"""
|
||
if not RAY_SERVE_ENABLE_HA_PROXY:
|
||
return
|
||
|
||
socket_glob = os.path.join(
|
||
os.path.dirname(RAY_SERVE_HAPROXY_SOCKET_PATH), "*", "admin.sock"
|
||
)
|
||
|
||
def replica_backend_up():
|
||
for sock_path in glob.glob(socket_glob):
|
||
try:
|
||
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client:
|
||
client.settimeout(5.0)
|
||
client.connect(sock_path)
|
||
client.sendall(b"show stat\n")
|
||
data = b""
|
||
while chunk := client.recv(65536):
|
||
data += chunk
|
||
except OSError:
|
||
continue # stale socket from a prior run
|
||
stats = HAProxyApi._parse_haproxy_csv_stats(data.decode(errors="replace"))
|
||
if any(
|
||
name.startswith("SERVE_REPLICA") and server.is_up
|
||
for servers in stats.values()
|
||
for name, server in servers.items()
|
||
):
|
||
return True
|
||
return False
|
||
|
||
wait_for_condition(replica_backend_up, timeout=timeout)
|
||
|
||
|
||
def alive_actor_counts() -> Dict[str, int]:
|
||
"""Count of ALIVE actors by class name in the current Ray session."""
|
||
counts: Dict[str, int] = {}
|
||
for actor in list_actors(filters=[("STATE", "=", "ALIVE")]):
|
||
counts[actor["class_name"]] = counts.get(actor["class_name"], 0) + 1
|
||
return counts
|
||
|
||
|
||
def check_num_replicas_gte(
|
||
name: str, target: int, app_name: str = SERVE_DEFAULT_APP_NAME
|
||
) -> int:
|
||
"""Check if num replicas is >= target."""
|
||
|
||
assert get_num_alive_replicas(name, app_name) >= target
|
||
return True
|
||
|
||
|
||
def check_num_replicas_eq(
|
||
name: str,
|
||
target: int,
|
||
app_name: str = SERVE_DEFAULT_APP_NAME,
|
||
use_controller: bool = False,
|
||
) -> bool:
|
||
"""Check if num replicas is == target."""
|
||
|
||
if use_controller:
|
||
dep = serve.status().applications[app_name].deployments[name]
|
||
num_running_replicas = dep.replica_states.get(ReplicaState.RUNNING, 0)
|
||
assert num_running_replicas == target
|
||
else:
|
||
assert get_num_alive_replicas(name, app_name) == target
|
||
|
||
return True
|
||
|
||
|
||
def check_num_replicas_lte(
|
||
name: str, target: int, app_name: str = SERVE_DEFAULT_APP_NAME
|
||
) -> int:
|
||
"""Check if num replicas is <= target."""
|
||
|
||
assert get_num_alive_replicas(name, app_name) <= target
|
||
return True
|
||
|
||
|
||
def check_apps_running(apps: List):
|
||
status = serve.status()
|
||
|
||
for app_name in apps:
|
||
assert status.applications[app_name].status == ApplicationStatus.RUNNING
|
||
|
||
return True
|
||
|
||
|
||
def check_replica_counts(
|
||
controller: ActorHandle,
|
||
deployment_id: DeploymentID,
|
||
total: Optional[int] = None,
|
||
by_state: Optional[List[Tuple[ReplicaState, int, Callable]]] = None,
|
||
):
|
||
"""Uses _dump_replica_states_for_testing to check replica counts.
|
||
|
||
Args:
|
||
controller: A handle to the Serve controller.
|
||
deployment_id: The deployment to check replica counts for.
|
||
total: The total number of expected replicas for the deployment.
|
||
by_state: A list of tuples of the form
|
||
(replica state, number of replicas, filter function).
|
||
Used for more fine grained checks.
|
||
|
||
Returns:
|
||
True when all assertions pass (raises ``AssertionError`` otherwise).
|
||
"""
|
||
replicas = ray.get(
|
||
controller._dump_replica_states_for_testing.remote(deployment_id)
|
||
)
|
||
|
||
if total is not None:
|
||
replica_counts = {
|
||
state: len(replicas.get([state]))
|
||
for state in ALL_REPLICA_STATES
|
||
if replicas.get([state])
|
||
}
|
||
assert replicas.count() == total, replica_counts
|
||
|
||
if by_state is not None:
|
||
for state, count, check in by_state:
|
||
assert isinstance(state, ReplicaState)
|
||
assert isinstance(count, int) and count >= 0
|
||
if check:
|
||
filtered = {r for r in replicas.get(states=[state]) if check(r)}
|
||
curr_count = len(filtered)
|
||
else:
|
||
curr_count = replicas.count(states=[state])
|
||
msg = f"Expected {count} for state {state} but got {curr_count}."
|
||
assert curr_count == count, msg
|
||
|
||
return True
|
||
|
||
|
||
@ray.remote(name=STORAGE_ACTOR_NAME, namespace=SERVE_NAMESPACE, num_cpus=0)
|
||
class TelemetryStorage:
|
||
def __init__(self):
|
||
self.reports_received = 0
|
||
self.current_report = dict()
|
||
|
||
def store_report(self, report: Dict) -> None:
|
||
self.reports_received += 1
|
||
self.current_report = report
|
||
|
||
def get_report(self) -> Dict:
|
||
return self.current_report
|
||
|
||
def get_reports_received(self) -> int:
|
||
return self.reports_received
|
||
|
||
|
||
@serve.deployment(ray_actor_options={"num_cpus": 0})
|
||
class TelemetryReceiver:
|
||
def __init__(self):
|
||
self.storage = ray.get_actor(name=STORAGE_ACTOR_NAME, namespace=SERVE_NAMESPACE)
|
||
|
||
async def __call__(self, request: Request) -> bool:
|
||
report = await request.json()
|
||
ray.get(self.storage.store_report.remote(report))
|
||
return True
|
||
|
||
|
||
receiver_app = TelemetryReceiver.bind()
|
||
|
||
|
||
def start_telemetry_app():
|
||
"""Start a telemetry Serve app.
|
||
|
||
Ray should be initialized before calling this method.
|
||
|
||
NOTE: If you're running the TelemetryReceiver Serve app to check telemetry,
|
||
remember that the receiver itself is counted in the telemetry. E.g. if you
|
||
deploy a Serve app other than the receiver, the number of apps in the
|
||
cluster is 2- not 1– since the receiver is also running.
|
||
|
||
Returns a handle to a TelemetryStorage actor. You can use this actor
|
||
to access the latest telemetry reports.
|
||
"""
|
||
|
||
storage = TelemetryStorage.remote()
|
||
serve.run(receiver_app, name="telemetry", route_prefix=TELEMETRY_ROUTE_PREFIX)
|
||
return storage
|
||
|
||
|
||
def check_telemetry(
|
||
tag: ServeUsageTag, expected: Any, storage_actor_name: str = STORAGE_ACTOR_NAME
|
||
):
|
||
storage_handle = ray.get_actor(storage_actor_name, namespace=SERVE_NAMESPACE)
|
||
report = ray.get(storage_handle.get_report.remote())
|
||
print(report["extra_usage_tags"])
|
||
assert tag.get_value_from_report(report) == expected
|
||
return True
|
||
|
||
|
||
def ping_grpc_list_applications(channel, app_names, test_draining=False):
|
||
import pytest
|
||
|
||
stub = serve_pb2_grpc.RayServeAPIServiceStub(channel)
|
||
request = serve_pb2.ListApplicationsRequest()
|
||
if test_draining:
|
||
with pytest.raises(grpc.RpcError) as exception_info:
|
||
_, _ = stub.ListApplications.with_call(request=request)
|
||
rpc_error = exception_info.value
|
||
assert rpc_error.code() == grpc.StatusCode.UNAVAILABLE
|
||
assert rpc_error.details() == DRAINING_MESSAGE
|
||
else:
|
||
response, call = stub.ListApplications.with_call(request=request)
|
||
assert call.code() == grpc.StatusCode.OK
|
||
assert response.application_names == app_names
|
||
return True
|
||
|
||
|
||
def ping_grpc_healthz(channel, test_draining=False):
|
||
import pytest
|
||
|
||
stub = serve_pb2_grpc.RayServeAPIServiceStub(channel)
|
||
request = serve_pb2.HealthzRequest()
|
||
if test_draining:
|
||
with pytest.raises(grpc.RpcError) as exception_info:
|
||
_, _ = stub.Healthz.with_call(request=request)
|
||
rpc_error = exception_info.value
|
||
assert rpc_error.code() == grpc.StatusCode.UNAVAILABLE
|
||
assert rpc_error.details() == DRAINING_MESSAGE
|
||
else:
|
||
response, call = stub.Healthz.with_call(request=request)
|
||
assert call.code() == grpc.StatusCode.OK
|
||
if not RAY_SERVE_ENABLE_HA_PROXY:
|
||
assert response.message == "success"
|
||
return True
|
||
|
||
|
||
def ping_grpc_call_method(channel, app_name, test_not_found=False):
|
||
import pytest
|
||
|
||
stub = serve_pb2_grpc.UserDefinedServiceStub(channel)
|
||
request = serve_pb2.UserDefinedMessage(name="foo", num=30, foo="bar")
|
||
metadata = (("application", app_name),)
|
||
if test_not_found:
|
||
with pytest.raises(grpc.RpcError) as exception_info:
|
||
_, _ = stub.__call__.with_call(request=request, metadata=metadata)
|
||
rpc_error = exception_info.value
|
||
assert rpc_error.code() == grpc.StatusCode.NOT_FOUND, rpc_error.code()
|
||
assert f"Application '{app_name}' not found." in rpc_error.details()
|
||
else:
|
||
response, call = stub.__call__.with_call(request=request, metadata=metadata)
|
||
assert call.code() == grpc.StatusCode.OK, call.code()
|
||
assert response.greeting == "Hello foo from bar", response.greeting
|
||
|
||
|
||
def ping_grpc_another_method(channel, app_name):
|
||
stub = serve_pb2_grpc.UserDefinedServiceStub(channel)
|
||
request = serve_pb2.UserDefinedMessage(name="foo", num=30, foo="bar")
|
||
metadata = (("application", app_name),)
|
||
response = stub.Method1(request=request, metadata=metadata)
|
||
assert response.greeting == "Hello foo from method1"
|
||
|
||
|
||
def ping_grpc_model_multiplexing(channel, app_name):
|
||
stub = serve_pb2_grpc.UserDefinedServiceStub(channel)
|
||
request = serve_pb2.UserDefinedMessage(name="foo", num=30, foo="bar")
|
||
multiplexed_model_id = "999"
|
||
metadata = (
|
||
("application", app_name),
|
||
("multiplexed_model_id", multiplexed_model_id),
|
||
)
|
||
response = stub.Method2(request=request, metadata=metadata)
|
||
assert (
|
||
response.greeting
|
||
== f"Method2 called model, loading model: {multiplexed_model_id}"
|
||
)
|
||
|
||
|
||
def ping_grpc_streaming(channel, app_name):
|
||
stub = serve_pb2_grpc.UserDefinedServiceStub(channel)
|
||
request = serve_pb2.UserDefinedMessage(name="foo", num=30, foo="bar")
|
||
metadata = (("application", app_name),)
|
||
responses = stub.Streaming(request=request, metadata=metadata)
|
||
for idx, response in enumerate(responses):
|
||
assert response.greeting == f"{idx}: Hello foo from bar"
|
||
|
||
|
||
def ping_fruit_stand(channel, app_name):
|
||
stub = serve_pb2_grpc.FruitServiceStub(channel)
|
||
request = serve_pb2.FruitAmounts(orange=4, apple=8)
|
||
metadata = (("application", app_name),)
|
||
response = stub.FruitStand(request=request, metadata=metadata)
|
||
assert response.costs == 32
|
||
|
||
|
||
@asynccontextmanager
|
||
async def send_signal_on_cancellation(signal_actor: ActorHandle):
|
||
cancelled = False
|
||
try:
|
||
yield
|
||
await asyncio.sleep(100)
|
||
except asyncio.CancelledError:
|
||
cancelled = True
|
||
# Clear the context var to avoid Ray recursively cancelling this method call.
|
||
ray._raylet.async_task_id.set(None)
|
||
await signal_actor.send.remote()
|
||
|
||
if not cancelled:
|
||
raise RuntimeError(
|
||
"CancelledError wasn't raised during `send_signal_on_cancellation` block"
|
||
)
|
||
|
||
|
||
class FakeGrpcContext:
|
||
def __init__(self):
|
||
self._auth_context = {"key": "value"}
|
||
self._invocation_metadata = [("key", "value")]
|
||
self._peer = "peer"
|
||
self._peer_identities = b"peer_identities"
|
||
self._peer_identity_key = "peer_identity_key"
|
||
self._code = None
|
||
self._details = None
|
||
self._trailing_metadata = []
|
||
self._invocation_metadata = []
|
||
|
||
def auth_context(self):
|
||
return self._auth_context
|
||
|
||
def code(self):
|
||
return self._code
|
||
|
||
def details(self):
|
||
return self._details
|
||
|
||
def peer(self):
|
||
return self._peer
|
||
|
||
def peer_identities(self):
|
||
return self._peer_identities
|
||
|
||
def peer_identity_key(self):
|
||
return self._peer_identity_key
|
||
|
||
def trailing_metadata(self):
|
||
return self._trailing_metadata
|
||
|
||
def set_code(self, code):
|
||
self._code = code
|
||
|
||
def set_details(self, details):
|
||
self._details = details
|
||
|
||
def set_trailing_metadata(self, trailing_metadata):
|
||
self._trailing_metadata = trailing_metadata
|
||
|
||
def invocation_metadata(self):
|
||
return self._invocation_metadata
|
||
|
||
|
||
class FakeGauge:
|
||
def __init__(self, name: str = None, tag_keys: Tuple[str] = None):
|
||
self.name = name
|
||
self.values = dict()
|
||
|
||
self.tags = tag_keys or ()
|
||
self.default_tags = dict()
|
||
|
||
def set_default_tags(self, tags: Dict[str, str]):
|
||
for key, tag in tags.items():
|
||
assert key in self.tags
|
||
self.default_tags[key] = tag
|
||
|
||
def set(self, value: Union[int, float], tags: Dict[str, str] = None):
|
||
merged_tags = self.default_tags.copy()
|
||
merged_tags.update(tags or {})
|
||
assert set(merged_tags.keys()) == set(self.tags)
|
||
|
||
d = self.values
|
||
for tag in self.tags[:-1]:
|
||
tag_value = merged_tags[tag]
|
||
if tag_value not in d:
|
||
d[tag_value] = dict()
|
||
d = d[tag_value]
|
||
|
||
d[merged_tags[self.tags[-1]]] = value
|
||
|
||
def get_value(self, tags: Dict[str, str]):
|
||
value = self.values
|
||
for tag in self.tags:
|
||
tag_value = tags[tag]
|
||
value = value.get(tag_value)
|
||
if value is None:
|
||
return
|
||
|
||
return value
|
||
|
||
|
||
class FakeCounter:
|
||
def __init__(self, name: str = None, tag_keys: Tuple[str] = None):
|
||
self.name = name
|
||
self.counts = dict()
|
||
|
||
self.tags = tag_keys or ()
|
||
self.default_tags = dict()
|
||
|
||
def set_default_tags(self, tags: Dict[str, str]):
|
||
for key, tag in tags.items():
|
||
assert key in self.tags
|
||
self.default_tags[key] = tag
|
||
|
||
def inc(self, value: Union[int, float] = 1.0, tags: Dict[str, str] = None):
|
||
merged_tags = self.default_tags.copy()
|
||
merged_tags.update(tags or {})
|
||
assert set(merged_tags.keys()) == set(self.tags)
|
||
|
||
d = self.counts
|
||
for tag in self.tags[:-1]:
|
||
tag_value = merged_tags[tag]
|
||
if tag_value not in d:
|
||
d[tag_value] = dict()
|
||
d = d[tag_value]
|
||
|
||
key = merged_tags[self.tags[-1]]
|
||
d[key] = d.get(key, 0) + value
|
||
|
||
def get_count(self, tags: Dict[str, str]) -> int:
|
||
value = self.counts
|
||
for tag in self.tags:
|
||
tag_value = tags[tag]
|
||
value = value.get(tag_value)
|
||
if value is None:
|
||
return
|
||
|
||
return value
|
||
|
||
def get_tags(self):
|
||
return self.tags
|
||
|
||
|
||
@ray.remote
|
||
def get_node_id():
|
||
return ray.get_runtime_context().get_node_id()
|
||
|
||
|
||
def check_num_alive_nodes(target: int):
|
||
alive_nodes = [node for node in ray.nodes() if node["Alive"]]
|
||
assert len(alive_nodes) == target
|
||
return True
|
||
|
||
|
||
def get_deployment_details(
|
||
deployment_name: str,
|
||
app_name: str = SERVE_DEFAULT_APP_NAME,
|
||
_client: ServeControllerClient = None,
|
||
):
|
||
client = _client or _get_global_client()
|
||
details = client.get_serve_details()
|
||
return details["applications"][app_name]["deployments"][deployment_name]
|
||
|
||
|
||
@ray.remote(num_cpus=0)
|
||
class Barrier:
|
||
"""Block until n participants have called wait()."""
|
||
|
||
def __init__(self, n: int):
|
||
self.n = n
|
||
self.count = 0
|
||
self.event = asyncio.Event()
|
||
|
||
async def wait(self):
|
||
self.count += 1
|
||
if self.count >= self.n:
|
||
self.event.set()
|
||
else:
|
||
await self.event.wait()
|
||
|
||
def get_count(self) -> int:
|
||
return self.count
|
||
|
||
|
||
@ray.remote(num_cpus=0)
|
||
class Accumulator:
|
||
"""Collect items from multiple actors/replicas."""
|
||
|
||
def __init__(self):
|
||
self.items = []
|
||
|
||
def add(self, item):
|
||
self.items.append(item)
|
||
|
||
def get(self) -> list:
|
||
return list(self.items)
|
||
|
||
def count(self) -> int:
|
||
return len(self.items)
|
||
|
||
|
||
@ray.remote(num_cpus=0)
|
||
class FailedReplicaStore:
|
||
"""Controls replica constructor failure behavior for constructor failure tests.
|
||
|
||
Behavior is determined by ``fail_first``:
|
||
- ``fail_first=True`` (transient): first caller fails, rest succeed.
|
||
- ``fail_first=False`` (partial): first caller succeeds, rest fail.
|
||
|
||
All decisions are made in a single atomic actor call to avoid races
|
||
between concurrent replicas.
|
||
"""
|
||
|
||
def __init__(self, fail_first: bool = False):
|
||
self._fail_first = fail_first
|
||
self._first_caller = False
|
||
self._fail_count = 0
|
||
|
||
def get_fail_count(self) -> int:
|
||
"""Returns the number replica startup failures"""
|
||
return self._fail_count
|
||
|
||
def should_fail(self) -> bool:
|
||
"""Returns whether this replica should raise in its constructor."""
|
||
if not self._first_caller:
|
||
self._first_caller = True
|
||
result = self._fail_first
|
||
else:
|
||
result = not self._fail_first
|
||
if result:
|
||
self._fail_count += 1
|
||
return result
|
||
|
||
|
||
@ray.remote(num_cpus=0)
|
||
class FailedGangReplicaStore:
|
||
"""
|
||
Controls replica constructor failure behavior for gang scheduling tests.
|
||
- The first gang seen is allowed to run.
|
||
- The next gang has first replica flagged to fail.
|
||
- Retry gangs after the first failed gang have first replica flagged.
|
||
"""
|
||
|
||
def __init__(self):
|
||
self._good_gang_chosen = False
|
||
self._successful_gang_id = None
|
||
self._failed_gang_ids = set()
|
||
|
||
def mark_first_failing_gang(self, gang_id: str) -> bool:
|
||
"""Picks the good gang on the first call and flags the first replica of the next gang to fail."""
|
||
if not self._good_gang_chosen:
|
||
self._good_gang_chosen = True
|
||
self._successful_gang_id = gang_id
|
||
return False
|
||
if gang_id == self._successful_gang_id:
|
||
return False
|
||
if len(self._failed_gang_ids) == 0:
|
||
self._failed_gang_ids.add(gang_id)
|
||
return True
|
||
return False
|
||
|
||
def mark_retry_failing_gang(self, gang_id: str) -> bool:
|
||
"""Flags the first replica of each new retry gang.
|
||
This is called after ``mark_first_failing_gang``."""
|
||
if gang_id == self._successful_gang_id:
|
||
return False
|
||
if gang_id not in self._failed_gang_ids:
|
||
self._failed_gang_ids.add(gang_id)
|
||
return True
|
||
return False
|
||
|
||
def get_failed_gang_count(self) -> int:
|
||
"""Returns the number of distinct gangs that failed."""
|
||
return len(self._failed_gang_ids)
|
||
|
||
|
||
@ray.remote(num_cpus=0)
|
||
class SharedFlag:
|
||
"""Non-blocking boolean flag shared across actors."""
|
||
|
||
def __init__(self):
|
||
self.value = False
|
||
|
||
def set(self):
|
||
self.value = True
|
||
|
||
def clear(self):
|
||
self.value = False
|
||
|
||
def is_set(self) -> bool:
|
||
return self.value
|
||
|
||
|
||
@ray.remote(num_cpus=0)
|
||
class SharedCounter:
|
||
"""Simple shared counter for cross-actor coordination."""
|
||
|
||
def __init__(self):
|
||
self.count = 0
|
||
|
||
def inc(self) -> int:
|
||
self.count += 1
|
||
return self.count
|
||
|
||
def get(self) -> int:
|
||
return self.count
|
||
|
||
|
||
@ray.remote(num_cpus=0)
|
||
class Counter:
|
||
def __init__(self, target: int):
|
||
self.count = 0
|
||
self.target = target
|
||
self.ready_event = asyncio.Event()
|
||
|
||
def inc(self):
|
||
self.count += 1
|
||
if self.count == self.target:
|
||
self.ready_event.set()
|
||
|
||
async def wait(self):
|
||
await self.ready_event.wait()
|
||
|
||
|
||
def tlog(s: str, level: str = "INFO"):
|
||
"""Convenient logging method for testing."""
|
||
|
||
now = datetime.datetime.now().strftime("%H:%M:%S.%f")[:-3]
|
||
print(f"[{level}] {now} {s}")
|
||
|
||
|
||
def check_target_groups_ready(
|
||
client: ServeControllerClient,
|
||
app_name: str,
|
||
protocol: Union[str, RequestProtocol] = RequestProtocol.HTTP,
|
||
):
|
||
"""Wait for target groups to be ready for the given app and protocol.
|
||
|
||
Target groups are ready when there are at least one target for the given protocol. And it's
|
||
possible that target groups are not ready immediately. An example is when the controller
|
||
is recovering from a crash.
|
||
"""
|
||
target_groups = ray.get(client._controller.get_target_groups.remote(app_name))
|
||
target_groups = [
|
||
target_group
|
||
for target_group in target_groups
|
||
if target_group.protocol == protocol
|
||
]
|
||
all_targets = [
|
||
target for target_group in target_groups for target in target_group.targets
|
||
]
|
||
return len(all_targets) > 0
|
||
|
||
|
||
def get_application_urls(
|
||
protocol: Union[str, RequestProtocol] = RequestProtocol.HTTP,
|
||
app_name: str = SERVE_DEFAULT_APP_NAME,
|
||
use_localhost: bool = True,
|
||
is_websocket: bool = False,
|
||
exclude_route_prefix: bool = False,
|
||
from_proxy_manager: bool = False,
|
||
) -> List[str]:
|
||
"""Get the URL of the application.
|
||
|
||
Args:
|
||
protocol: The protocol to use for the application.
|
||
app_name: The name of the application.
|
||
use_localhost: Whether to use localhost instead of the IP address.
|
||
Set to True if Serve deployments are not exposed publicly or
|
||
for low latency benchmarking.
|
||
is_websocket: Whether the url should be served as a websocket.
|
||
exclude_route_prefix: The route prefix to exclude from the application.
|
||
from_proxy_manager: Whether the caller is a proxy manager.
|
||
Returns:
|
||
The URLs of the application.
|
||
"""
|
||
client = _get_global_client()
|
||
serve_details = client.get_serve_details()
|
||
assert (
|
||
app_name in serve_details["applications"]
|
||
), f"App {app_name} not found in serve details. Use this method only when the app is known to be running."
|
||
route_prefix = serve_details["applications"][app_name]["route_prefix"]
|
||
# route_prefix is set to None when route_prefix value is specifically set to None
|
||
# in the config used to deploy the app.
|
||
if exclude_route_prefix or route_prefix is None:
|
||
route_prefix = ""
|
||
if isinstance(protocol, str):
|
||
protocol = RequestProtocol(protocol)
|
||
target_groups: List[TargetGroup] = ray.get(
|
||
client._controller.get_target_groups.remote(app_name, from_proxy_manager)
|
||
)
|
||
target_groups = [
|
||
target_group
|
||
for target_group in target_groups
|
||
if target_group.protocol == protocol
|
||
]
|
||
if len(target_groups) == 0:
|
||
raise ValueError(
|
||
f"No target group found for app {app_name} with protocol {protocol} and route prefix {route_prefix}"
|
||
)
|
||
urls = []
|
||
for target_group in target_groups:
|
||
for target in target_group.targets:
|
||
ip = "localhost" if use_localhost else target.ip
|
||
if protocol == RequestProtocol.HTTP:
|
||
scheme = "ws" if is_websocket else "http"
|
||
url = f"{scheme}://{build_address(ip, target.port)}{route_prefix}"
|
||
elif protocol == RequestProtocol.GRPC:
|
||
if is_websocket:
|
||
raise ValueError(
|
||
"is_websocket=True is not supported with gRPC protocol."
|
||
)
|
||
url = build_address(ip, target.port)
|
||
else:
|
||
raise ValueError(f"Unsupported protocol: {protocol}")
|
||
url = url.rstrip("/")
|
||
urls.append(url)
|
||
return urls
|
||
|
||
|
||
def get_application_url(
|
||
protocol: Union[str, RequestProtocol] = RequestProtocol.HTTP,
|
||
app_name: str = SERVE_DEFAULT_APP_NAME,
|
||
use_localhost: bool = True,
|
||
is_websocket: bool = False,
|
||
exclude_route_prefix: bool = False,
|
||
from_proxy_manager: bool = False,
|
||
) -> str:
|
||
"""Get the URL of the application.
|
||
|
||
Args:
|
||
protocol: The protocol to use for the application.
|
||
app_name: The name of the application.
|
||
use_localhost: Whether to use localhost instead of the IP address.
|
||
Set to True if Serve deployments are not exposed publicly or
|
||
for low latency benchmarking.
|
||
is_websocket: Whether the url should be served as a websocket.
|
||
exclude_route_prefix: The route prefix to exclude from the application.
|
||
from_proxy_manager: Whether the caller is a proxy manager.
|
||
Returns:
|
||
The URL of the application. If there are multiple URLs, a random one is returned.
|
||
"""
|
||
return random.choice(
|
||
get_application_urls(
|
||
protocol,
|
||
app_name,
|
||
use_localhost,
|
||
is_websocket,
|
||
exclude_route_prefix,
|
||
from_proxy_manager,
|
||
)
|
||
)
|
||
|
||
|
||
def check_running(app_name: str = SERVE_DEFAULT_APP_NAME):
|
||
assert serve.status().applications[app_name].status == ApplicationStatus.RUNNING
|
||
return True
|
||
|
||
|
||
def request_with_retries(timeout=30, app_name=SERVE_DEFAULT_APP_NAME):
|
||
result_holder = {"resp": None}
|
||
|
||
def _attempt() -> bool:
|
||
try:
|
||
url = get_application_url("HTTP", app_name=app_name)
|
||
result_holder["resp"] = httpx.get(url, timeout=timeout)
|
||
return True
|
||
except (httpx.RequestError, IndexError):
|
||
return False
|
||
|
||
try:
|
||
wait_for_condition(_attempt, timeout=timeout)
|
||
return result_holder["resp"]
|
||
except RuntimeError as e:
|
||
# Preserve previous API by raising TimeoutError on expiry
|
||
raise TimeoutError from e
|
||
|
||
|
||
# Metrics test utilities
|
||
TEST_METRICS_EXPORT_PORT = 9999
|
||
|
||
|
||
def get_metric_float(
|
||
metric: str,
|
||
expected_tags: Optional[Dict[str, str]],
|
||
timeseries: Optional[PrometheusTimeseries] = None,
|
||
timeout: float = 20,
|
||
) -> float:
|
||
"""Gets the float value of metric.
|
||
|
||
If tags is specified, searched for metric with matching tags.
|
||
|
||
Returns -1 if the metric isn't available.
|
||
"""
|
||
if timeseries is None:
|
||
timeseries = PrometheusTimeseries()
|
||
samples = fetch_prometheus_metric_timeseries(
|
||
[f"localhost:{TEST_METRICS_EXPORT_PORT}"],
|
||
timeseries,
|
||
timeout=timeout,
|
||
).get(metric, [])
|
||
for sample in samples:
|
||
if expected_tags.items() <= sample.labels.items():
|
||
return sample.value
|
||
return -1
|
||
|
||
|
||
def check_metric_float_eq(
|
||
metric: str,
|
||
expected: float,
|
||
expected_tags: Optional[Dict[str, str]],
|
||
timeseries: Optional[PrometheusTimeseries] = None,
|
||
timeout: float = 20,
|
||
) -> bool:
|
||
"""Check if a metric's float value equals the expected value."""
|
||
metric_value = get_metric_float(metric, expected_tags, timeseries, timeout=timeout)
|
||
assert float(metric_value) == expected
|
||
return True
|
||
|
||
|
||
def get_metric_dictionaries(
|
||
name: str,
|
||
timeout: float = 20,
|
||
timeseries: Optional[PrometheusTimeseries] = None,
|
||
wait: bool = True,
|
||
) -> List[Dict]:
|
||
"""Get metric samples as list of label dicts.
|
||
|
||
Args:
|
||
name: Metric name to fetch.
|
||
timeout: Timeout for each fetch attempt.
|
||
timeseries: Optional shared timeseries to populate.
|
||
wait: If True (default), wait for metric to appear before returning.
|
||
If False, fetch once and return immediately (possibly empty).
|
||
|
||
Returns:
|
||
List of metric samples as label dicts.
|
||
"""
|
||
if timeseries is None:
|
||
timeseries = PrometheusTimeseries()
|
||
|
||
def metric_available() -> bool:
|
||
prom_timeseries = fetch_prometheus_metric_timeseries(
|
||
[f"localhost:{TEST_METRICS_EXPORT_PORT}"],
|
||
timeseries,
|
||
# pass timeout to fetch_prometheus_metric_timeseries
|
||
# so the test doesn't hang on requests.get
|
||
timeout=timeout,
|
||
)
|
||
assert (
|
||
name in prom_timeseries
|
||
), f"Metric {name} not found. Available metrics: {list(prom_timeseries.keys())}"
|
||
return True
|
||
|
||
if wait:
|
||
wait_for_condition(
|
||
metric_available, retry_interval_ms=1000, timeout=timeout * 4
|
||
)
|
||
else:
|
||
# Fetch once without asserting - allows outer wait_for_condition to retry
|
||
fetch_prometheus_metric_timeseries(
|
||
[f"localhost:{TEST_METRICS_EXPORT_PORT}"],
|
||
timeseries,
|
||
timeout=PROMETHEUS_METRICS_TIMEOUT_S,
|
||
)
|
||
|
||
metric_dicts = []
|
||
for sample in timeseries.metric_samples.values():
|
||
if sample.name == name:
|
||
metric_dicts.append(sample.labels)
|
||
|
||
return metric_dicts
|