wehub-resource-sync
415 lines
15 KiB
Python
415 lines
15 KiB
Python
import os
|
|
import sys
|
|
import unittest
|
|
from collections import defaultdict
|
|
from unittest.mock import MagicMock
|
|
|
|
# coding: utf-8
|
|
import pytest
|
|
|
|
from ray.autoscaler.v2.instance_manager.instance_manager import InstanceManager
|
|
from ray.autoscaler.v2.instance_manager.instance_storage import InstanceStorage
|
|
from ray.autoscaler.v2.instance_manager.storage import InMemoryStorage, StoreStatus
|
|
from ray.autoscaler.v2.tests.util import MockSubscriber
|
|
from ray.core.generated.instance_manager_pb2 import (
|
|
GetInstanceManagerStateRequest,
|
|
Instance,
|
|
InstanceUpdateEvent,
|
|
NodeKind,
|
|
StatusCode,
|
|
UpdateInstanceManagerStateRequest,
|
|
)
|
|
|
|
|
|
class InstanceManagerTest(unittest.TestCase):
|
|
def test_instances_version_mismatch(self):
|
|
ins_storage = MagicMock()
|
|
subscriber = MockSubscriber()
|
|
im = InstanceManager(
|
|
ins_storage, instance_status_update_subscribers=[subscriber]
|
|
)
|
|
# Version mismatch on reading from the storage.
|
|
ins_storage.get_instances.return_value = ({}, 1)
|
|
|
|
update = InstanceUpdateEvent(
|
|
instance_id="id-1",
|
|
new_instance_status=Instance.QUEUED,
|
|
instance_type="type-1",
|
|
upsert=True,
|
|
)
|
|
reply = im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=0,
|
|
updates=[update],
|
|
)
|
|
)
|
|
assert reply.status.code == StatusCode.VERSION_MISMATCH
|
|
assert len(subscriber.events) == 0
|
|
|
|
# Version OK.
|
|
reply = im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=1,
|
|
updates=[update],
|
|
)
|
|
)
|
|
assert reply.status.code == StatusCode.OK
|
|
assert len(subscriber.events) == 1
|
|
assert subscriber.events[0].new_instance_status == Instance.QUEUED
|
|
|
|
# Version mismatch when writing to the storage (race happens)
|
|
ins_storage.batch_upsert_instances.return_value = StoreStatus(
|
|
False, 2 # No longer 1
|
|
)
|
|
subscriber.clear()
|
|
reply = im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=1,
|
|
updates=[update],
|
|
)
|
|
)
|
|
assert reply.status.code == StatusCode.VERSION_MISMATCH
|
|
assert len(subscriber.events) == 0
|
|
|
|
# Non-version mismatch error.
|
|
ins_storage.batch_upsert_instances.return_value = StoreStatus(
|
|
False, 1 # Still 1
|
|
)
|
|
reply = im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=1,
|
|
updates=[update],
|
|
)
|
|
)
|
|
assert reply.status.code == StatusCode.UNKNOWN_ERRORS
|
|
assert len(subscriber.events) == 0
|
|
|
|
def test_get_and_updates(self):
|
|
ins_storage = InstanceStorage(
|
|
"cluster-id",
|
|
InMemoryStorage(),
|
|
)
|
|
subscriber = MockSubscriber()
|
|
im = InstanceManager(
|
|
ins_storage, instance_status_update_subscribers=[subscriber]
|
|
)
|
|
|
|
# Empty storage.
|
|
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
|
|
assert reply.status.code == StatusCode.OK
|
|
assert list(reply.state.instances) == []
|
|
|
|
# Launch nodes.
|
|
reply = im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=0,
|
|
updates=[
|
|
InstanceUpdateEvent(
|
|
instance_type="type-1",
|
|
instance_id="id-1",
|
|
new_instance_status=Instance.QUEUED,
|
|
upsert=True,
|
|
),
|
|
InstanceUpdateEvent(
|
|
instance_type="type-2",
|
|
instance_id="id-2",
|
|
new_instance_status=Instance.QUEUED,
|
|
upsert=True,
|
|
),
|
|
InstanceUpdateEvent(
|
|
instance_type="type-2",
|
|
instance_id="id-3",
|
|
new_instance_status=Instance.QUEUED,
|
|
upsert=True,
|
|
),
|
|
],
|
|
)
|
|
)
|
|
assert reply.status.code == StatusCode.OK
|
|
assert len(subscriber.events) == 3
|
|
for e in subscriber.events:
|
|
assert e.new_instance_status == Instance.QUEUED
|
|
|
|
# Get launched nodes.
|
|
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
|
|
assert reply.status.code == StatusCode.OK
|
|
assert len(reply.state.instances) == 3
|
|
|
|
instance_ids = [ins.instance_id for ins in reply.state.instances]
|
|
|
|
types_count = defaultdict(int)
|
|
for ins in reply.state.instances:
|
|
types_count[ins.instance_type] += 1
|
|
assert ins.status == Instance.QUEUED
|
|
|
|
assert types_count["type-1"] == 1
|
|
assert types_count["type-2"] == 2
|
|
|
|
# Update node status.
|
|
subscriber.clear()
|
|
reply = im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=1,
|
|
updates=[
|
|
InstanceUpdateEvent(
|
|
instance_id=instance_ids[0],
|
|
new_instance_status=Instance.REQUESTED,
|
|
instance_type="type-1",
|
|
launch_request_id="l1",
|
|
),
|
|
InstanceUpdateEvent(
|
|
instance_id=instance_ids[1],
|
|
new_instance_status=Instance.REQUESTED,
|
|
launch_request_id="l1",
|
|
instance_type="type-1",
|
|
),
|
|
],
|
|
)
|
|
)
|
|
|
|
assert reply.status.code == StatusCode.OK
|
|
assert len(subscriber.events) == 2
|
|
for e in subscriber.events:
|
|
assert e.new_instance_status == Instance.REQUESTED
|
|
|
|
# Get updated nodes.
|
|
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
|
|
assert reply.status.code == StatusCode.OK
|
|
assert len(reply.state.instances) == 3
|
|
|
|
types_count = defaultdict(int)
|
|
for ins in reply.state.instances:
|
|
types_count[ins.instance_type] += 1
|
|
if ins.instance_id in [instance_ids[0], instance_ids[1]]:
|
|
assert ins.status == Instance.REQUESTED
|
|
else:
|
|
assert ins.status == Instance.QUEUED
|
|
|
|
# Invalid instances status update.
|
|
subscriber.clear()
|
|
with pytest.raises(AssertionError):
|
|
reply = im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=2,
|
|
updates=[
|
|
InstanceUpdateEvent(
|
|
instance_id=instance_ids[2],
|
|
# Not requested yet.
|
|
new_instance_status=Instance.RAY_RUNNING,
|
|
),
|
|
],
|
|
)
|
|
)
|
|
assert len(subscriber.events) == 0
|
|
|
|
# Invalid versions.
|
|
reply = im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=0, # Invalid version, outdated.
|
|
updates=[
|
|
InstanceUpdateEvent(
|
|
instance_id=instance_ids[2],
|
|
new_instance_status=Instance.REQUESTED,
|
|
instance_type="type-2",
|
|
),
|
|
],
|
|
)
|
|
)
|
|
assert reply.status.code == StatusCode.VERSION_MISMATCH
|
|
assert len(subscriber.events) == 0
|
|
|
|
def test_insert(self):
|
|
ins_storage = InstanceStorage(
|
|
"cluster-id",
|
|
InMemoryStorage(),
|
|
)
|
|
subscriber = MockSubscriber()
|
|
im = InstanceManager(
|
|
ins_storage, instance_status_update_subscribers=[subscriber]
|
|
)
|
|
|
|
im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=0,
|
|
updates=[
|
|
InstanceUpdateEvent(
|
|
instance_type="type-1",
|
|
instance_id="id-1",
|
|
new_instance_status=Instance.QUEUED,
|
|
upsert=True,
|
|
),
|
|
InstanceUpdateEvent(
|
|
instance_id="id-2",
|
|
new_instance_status=Instance.TERMINATING,
|
|
cloud_instance_id="cloud-id-2",
|
|
upsert=True,
|
|
),
|
|
InstanceUpdateEvent(
|
|
instance_id="id-3",
|
|
new_instance_status=Instance.ALLOCATED,
|
|
cloud_instance_id="cloud-id-3",
|
|
node_kind=NodeKind.WORKER,
|
|
instance_type="type-3",
|
|
upsert=True,
|
|
),
|
|
],
|
|
)
|
|
)
|
|
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
|
|
assert len(reply.state.instances) == 3
|
|
instance_by_ids = {ins.instance_id: ins for ins in reply.state.instances}
|
|
assert instance_by_ids["id-1"].status == Instance.QUEUED
|
|
assert instance_by_ids["id-1"].instance_type == "type-1"
|
|
assert instance_by_ids["id-2"].status == Instance.TERMINATING
|
|
assert instance_by_ids["id-3"].status == Instance.ALLOCATED
|
|
assert instance_by_ids["id-3"].cloud_instance_id == "cloud-id-3"
|
|
version = reply.state.version
|
|
|
|
# With non-upsert flags.
|
|
with pytest.raises(AssertionError):
|
|
reply = im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=version,
|
|
updates=[
|
|
InstanceUpdateEvent(
|
|
instance_type="type-1",
|
|
instance_id="id-999",
|
|
new_instance_status=Instance.QUEUED,
|
|
),
|
|
],
|
|
)
|
|
)
|
|
|
|
# With invalid statuses
|
|
all_statuses = set(Instance.InstanceStatus.values())
|
|
non_insertable_statuses = all_statuses - {
|
|
Instance.QUEUED,
|
|
Instance.TERMINATING,
|
|
Instance.ALLOCATED,
|
|
}
|
|
|
|
for status in non_insertable_statuses:
|
|
subscriber.clear()
|
|
with pytest.raises(AssertionError):
|
|
reply = im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=version,
|
|
updates=[
|
|
InstanceUpdateEvent(
|
|
instance_id="id-999",
|
|
new_instance_status=status,
|
|
),
|
|
],
|
|
)
|
|
)
|
|
assert len(subscriber.events) == 0
|
|
|
|
def test_apply_update(self):
|
|
ins_storage = InstanceStorage(
|
|
"cluster-id",
|
|
InMemoryStorage(),
|
|
)
|
|
subscriber = MockSubscriber()
|
|
im = InstanceManager(
|
|
ins_storage, instance_status_update_subscribers=[subscriber]
|
|
)
|
|
|
|
# Insert a new instance.
|
|
im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=0,
|
|
updates=[
|
|
InstanceUpdateEvent(
|
|
instance_type="type-1",
|
|
instance_id="id-1",
|
|
new_instance_status=Instance.QUEUED,
|
|
upsert=True,
|
|
),
|
|
],
|
|
)
|
|
)
|
|
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
|
|
assert len(reply.state.instances) == 1
|
|
assert reply.state.instances[0].status == Instance.QUEUED
|
|
assert reply.state.instances[0].instance_type == "type-1"
|
|
|
|
# Request
|
|
im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=1,
|
|
updates=[
|
|
InstanceUpdateEvent(
|
|
instance_id="id-1",
|
|
new_instance_status=Instance.REQUESTED,
|
|
launch_request_id="l1",
|
|
instance_type="type-1",
|
|
),
|
|
],
|
|
)
|
|
)
|
|
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
|
|
assert len(reply.state.instances) == 1
|
|
assert reply.state.instances[0].status == Instance.REQUESTED
|
|
assert reply.state.instances[0].launch_request_id == "l1"
|
|
|
|
# ALLOCATED
|
|
im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=2,
|
|
updates=[
|
|
InstanceUpdateEvent(
|
|
instance_id="id-1",
|
|
new_instance_status=Instance.ALLOCATED,
|
|
cloud_instance_id="cloud-id-1",
|
|
node_kind=NodeKind.WORKER,
|
|
instance_type="type-1",
|
|
),
|
|
],
|
|
)
|
|
)
|
|
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
|
|
assert len(reply.state.instances) == 1
|
|
assert reply.state.instances[0].status == Instance.ALLOCATED
|
|
assert reply.state.instances[0].cloud_instance_id == "cloud-id-1"
|
|
|
|
# RAY_RUNNING
|
|
im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=3,
|
|
updates=[
|
|
InstanceUpdateEvent(
|
|
instance_id="id-1",
|
|
new_instance_status=Instance.RAY_RUNNING,
|
|
ray_node_id="ray-node-1",
|
|
),
|
|
],
|
|
)
|
|
)
|
|
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
|
|
assert len(reply.state.instances) == 1
|
|
assert reply.state.instances[0].status == Instance.RAY_RUNNING
|
|
assert reply.state.instances[0].node_id == "ray-node-1"
|
|
|
|
# TERMINATED
|
|
im.update_instance_manager_state(
|
|
UpdateInstanceManagerStateRequest(
|
|
expected_version=4,
|
|
updates=[
|
|
InstanceUpdateEvent(
|
|
instance_id="id-1",
|
|
new_instance_status=Instance.TERMINATED,
|
|
),
|
|
],
|
|
)
|
|
)
|
|
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
|
|
assert len(reply.state.instances) == 1
|
|
assert reply.state.instances[0].status == Instance.TERMINATED
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if os.environ.get("PARALLEL_CI"):
|
|
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
|
|
else:
|
|
sys.exit(pytest.main(["-sv", __file__]))
|