wehub-resource-sync
120 lines
5.0 KiB
Python
120 lines
5.0 KiB
Python
# coding: utf-8
|
|
import os
|
|
import sys
|
|
import unittest
|
|
from queue import Queue
|
|
from unittest.mock import patch
|
|
|
|
import pytest # noqa
|
|
|
|
from ray._private.test_utils import load_test_config
|
|
from ray.autoscaler.tags import TAG_RAY_NODE_KIND
|
|
from ray.autoscaler.v2.instance_manager.config import AutoscalingConfig
|
|
from ray.autoscaler.v2.instance_manager.instance_storage import InstanceStorage
|
|
from ray.autoscaler.v2.instance_manager.ray_installer import RayInstaller
|
|
from ray.autoscaler.v2.instance_manager.storage import InMemoryStorage
|
|
from ray.autoscaler.v2.instance_manager.subscribers.threaded_ray_installer import (
|
|
ThreadedRayInstaller,
|
|
)
|
|
from ray.core.generated.instance_manager_pb2 import Instance, NodeKind
|
|
from ray.tests.autoscaler_test_utils import MockProcessRunner, MockProvider
|
|
|
|
|
|
class ThreadedRayInstallerTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.base_provider = MockProvider()
|
|
self.config = AutoscalingConfig(load_test_config("test_ray_complex.yaml"))
|
|
self.runner = MockProcessRunner()
|
|
self.ray_installer = RayInstaller(self.base_provider, self.config, self.runner)
|
|
self.instance_storage = InstanceStorage(
|
|
cluster_id="test_cluster_id",
|
|
storage=InMemoryStorage(),
|
|
)
|
|
self.error_queue = Queue()
|
|
self.threaded_ray_installer = ThreadedRayInstaller(
|
|
head_node_ip="127.0.0.1",
|
|
instance_storage=self.instance_storage,
|
|
ray_installer=self.ray_installer,
|
|
error_queue=self.error_queue,
|
|
)
|
|
|
|
def test_install_ray_on_new_node_version_mismatch(self):
|
|
self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
|
|
instance = Instance(
|
|
instance_id="0",
|
|
instance_type="worker_nodes1",
|
|
cloud_instance_id="0",
|
|
status=Instance.RAY_INSTALLING,
|
|
node_kind=NodeKind.WORKER,
|
|
)
|
|
success, verison = self.instance_storage.upsert_instance(instance)
|
|
assert success
|
|
self.runner.respond_to_call("json .Config.Env", ["[]" for i in range(1)])
|
|
|
|
self.threaded_ray_installer._install_ray_on_single_node(instance)
|
|
instances, _ = self.instance_storage.get_instances(
|
|
instance_ids={instance.instance_id}
|
|
)
|
|
assert instances[instance.instance_id].status == Instance.RAY_INSTALLING
|
|
assert instances[instance.instance_id].version == verison
|
|
|
|
@patch.object(RayInstaller, "install_ray")
|
|
def test_install_ray_on_new_node_install_failed(self, mock_method):
|
|
self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
|
|
instance = Instance(
|
|
instance_id="0",
|
|
instance_type="worker_nodes1",
|
|
cloud_instance_id="0",
|
|
status=Instance.RAY_INSTALLING,
|
|
node_kind=NodeKind.WORKER,
|
|
)
|
|
success, verison = self.instance_storage.upsert_instance(instance)
|
|
assert success
|
|
instance.version = verison
|
|
|
|
mock_method.side_effect = RuntimeError("Installation failed")
|
|
self.threaded_ray_installer._install_retry_interval = 0
|
|
self.threaded_ray_installer._max_install_attempts = 1
|
|
self.threaded_ray_installer._install_ray_on_single_node(instance)
|
|
|
|
instances, _ = self.instance_storage.get_instances(
|
|
instance_ids={instance.instance_id}
|
|
)
|
|
# Make sure the instance status is not updated by the ThreadedRayInstaller
|
|
# since it should be updated by the Reconciler.
|
|
assert instances[instance.instance_id].status == Instance.RAY_INSTALLING
|
|
# Make sure the error is added to the error queue.
|
|
error = self.error_queue.get()
|
|
assert error.im_instance_id == instance.instance_id
|
|
assert "Installation failed" in error.details
|
|
|
|
def test_install_ray_on_new_nodes(self):
|
|
self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
|
|
instance = Instance(
|
|
instance_id="0",
|
|
instance_type="worker_nodes1",
|
|
cloud_instance_id="0",
|
|
status=Instance.RAY_INSTALLING,
|
|
node_kind=NodeKind.WORKER,
|
|
)
|
|
success, verison = self.instance_storage.upsert_instance(instance)
|
|
assert success
|
|
instance.version = verison
|
|
self.runner.respond_to_call("json .Config.Env", ["[]" for i in range(1)])
|
|
|
|
self.threaded_ray_installer._install_ray_on_new_nodes(instance.instance_id)
|
|
self.threaded_ray_installer._ray_installation_executor.shutdown(wait=True)
|
|
instances, _ = self.instance_storage.get_instances(
|
|
instance_ids={instance.instance_id}
|
|
)
|
|
# Make sure the instance status is not updated by the ThreadedRayInstaller
|
|
# since it should be updated by the Reconciler.
|
|
assert instances[instance.instance_id].status == Instance.RAY_INSTALLING
|
|
|
|
|
|
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__]))
|