wehub-resource-sync
95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
from typing import Tuple
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._raylet import NodeID
|
|
from ray.serve._private.autoscaling_state import AutoscalingStateManager
|
|
from ray.serve._private.deployment_state import DeploymentStateManager
|
|
from ray.serve._private.test_utils import (
|
|
MockClusterNodeInfoCache,
|
|
MockDeploymentActorWrapper,
|
|
MockKVStore,
|
|
MockReplicaActorWrapper,
|
|
MockTimer,
|
|
dead_replicas_context,
|
|
replica_rank_context,
|
|
uninitialized_replicas_context,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def disallow_ray_init(monkeypatch):
|
|
def raise_on_init():
|
|
raise RuntimeError("Unit tests should not depend on Ray being initialized.")
|
|
|
|
monkeypatch.setattr(ray, "init", raise_on_init)
|
|
# Unit tests don't run on a Ray cluster, so stub the runtime context
|
|
# that callers consult for worker/actor ids.
|
|
monkeypatch.setattr(ray, "get_runtime_context", Mock())
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_deployment_state_manager(
|
|
request,
|
|
) -> Tuple[DeploymentStateManager, MockTimer, Mock, Mock]:
|
|
"""Fully mocked deployment state manager.
|
|
|
|
i.e kv store and gcs client is mocked so we don't need to initialize
|
|
ray. Also, since this is used for some recovery tests, this yields a
|
|
method for creating a new mocked deployment state manager.
|
|
"""
|
|
|
|
timer = MockTimer()
|
|
with patch(
|
|
"ray.serve._private.deployment_state.ActorReplicaWrapper",
|
|
new=MockReplicaActorWrapper,
|
|
), patch(
|
|
"ray.serve._private.deployment_state.DeploymentActorWrapper",
|
|
new=MockDeploymentActorWrapper,
|
|
), patch(
|
|
"time.time", new=timer.time
|
|
), patch(
|
|
"ray.serve._private.long_poll.LongPollHost"
|
|
) as mock_long_poll, patch(
|
|
"ray.get_runtime_context"
|
|
):
|
|
kv_store = MockKVStore()
|
|
cluster_node_info_cache = MockClusterNodeInfoCache()
|
|
cluster_node_info_cache.add_node(NodeID.from_random().hex())
|
|
autoscaling_state_manager = AutoscalingStateManager()
|
|
|
|
def create_deployment_state_manager(
|
|
actor_names=None,
|
|
placement_group_names=None,
|
|
create_placement_group_fn_override=None,
|
|
):
|
|
if actor_names is None:
|
|
actor_names = []
|
|
|
|
if placement_group_names is None:
|
|
placement_group_names = []
|
|
|
|
return DeploymentStateManager(
|
|
kv_store,
|
|
mock_long_poll,
|
|
actor_names,
|
|
placement_group_names,
|
|
cluster_node_info_cache,
|
|
autoscaling_state_manager,
|
|
head_node_id_override="fake-head-node-id",
|
|
create_placement_group_fn_override=create_placement_group_fn_override,
|
|
)
|
|
|
|
yield (
|
|
create_deployment_state_manager,
|
|
timer,
|
|
cluster_node_info_cache,
|
|
autoscaling_state_manager,
|
|
)
|
|
|
|
dead_replicas_context.clear()
|
|
replica_rank_context.clear()
|
|
uninitialized_replicas_context.clear()
|