wehub-resource-sync
90 lines
3.6 KiB
Python
90 lines
3.6 KiB
Python
from typing import Optional
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from ray.serve._private.common import TargetCapacityDirection
|
|
from ray.serve._private.config import DeploymentConfig, ReplicaConfig
|
|
from ray.serve._private.deployment_info import DeploymentInfo
|
|
|
|
|
|
@patch("ray.serve._private.deployment_info.ray.get_runtime_context")
|
|
@pytest.mark.parametrize("target_capacity", [0, 25.809, 73, 100.0, None])
|
|
@pytest.mark.parametrize(
|
|
"target_capacity_direction",
|
|
[TargetCapacityDirection.UP, TargetCapacityDirection.DOWN, None],
|
|
)
|
|
def test_deployment_info_serialization(
|
|
mock_runtime_context,
|
|
target_capacity: Optional[float],
|
|
target_capacity_direction: Optional[TargetCapacityDirection],
|
|
):
|
|
"""Checks that deploymet infos can be serialized without losing data."""
|
|
# Mock out the runtime_context call, so Ray doesn't start.
|
|
class MockRuntimeContext:
|
|
def get_job_id(self) -> str:
|
|
return ""
|
|
|
|
mock_runtime_context.return_value = MockRuntimeContext()
|
|
|
|
def compare_replica_configs(config1: ReplicaConfig, config2: ReplicaConfig):
|
|
"""Check whether two replica configs contain the same data."""
|
|
|
|
assert config1.serialized_deployment_def == config2.serialized_deployment_def
|
|
assert config1.serialized_init_args == config2.serialized_init_args
|
|
assert config1.serialized_init_kwargs == config2.serialized_init_kwargs
|
|
assert config1.ray_actor_options == config2.ray_actor_options
|
|
assert config1.placement_group_bundles == config2.placement_group_bundles
|
|
assert config1.placement_group_strategy == config2.placement_group_strategy
|
|
assert config1.max_replicas_per_node == config2.max_replicas_per_node
|
|
|
|
def compare_deployment_info(info1: DeploymentInfo, info2: DeploymentInfo):
|
|
"""Checks that two deployment infos contain the same data.
|
|
|
|
Does not compare job_ids because those are mocked out to avoid starting
|
|
Ray
|
|
"""
|
|
|
|
assert info1.version == info2.version
|
|
assert info1.deployment_config == info2.deployment_config
|
|
compare_replica_configs(info1.replica_config, info2.replica_config)
|
|
assert info1.start_time_ms == info2.start_time_ms
|
|
assert info1.target_capacity == info2.target_capacity
|
|
assert info1.target_capacity_direction == info2.target_capacity_direction
|
|
|
|
deployment_info_args = dict(
|
|
version="123",
|
|
deployment_config=DeploymentConfig(num_replicas=1),
|
|
replica_config=ReplicaConfig.create(lambda x: x),
|
|
start_time_ms=0,
|
|
deployer_job_id="",
|
|
)
|
|
|
|
# Check that serialization works correctly when the DeploymentInfo is
|
|
# initialized with the target_capacity info directly.
|
|
info = DeploymentInfo(
|
|
target_capacity=target_capacity,
|
|
target_capacity_direction=target_capacity_direction,
|
|
**deployment_info_args
|
|
)
|
|
serialized_info = info.to_proto()
|
|
reconstructed_info = DeploymentInfo.from_proto(serialized_info)
|
|
compare_deployment_info(reconstructed_info, info)
|
|
|
|
# Check that serialization works correctly when the DeploymentInfo's
|
|
# target_capacity has been set using set_target_capacity().
|
|
info = DeploymentInfo(**deployment_info_args)
|
|
info.set_target_capacity(
|
|
new_target_capacity=target_capacity,
|
|
new_target_capacity_direction=target_capacity_direction,
|
|
)
|
|
serialized_info = info.to_proto()
|
|
reconstructed_info = DeploymentInfo.from_proto(serialized_info)
|
|
compare_deployment_info(reconstructed_info, info)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-s", __file__]))
|