wehub-resource-sync
97 lines
4.0 KiB
Python
97 lines
4.0 KiB
Python
import logging
|
|
import subprocess
|
|
|
|
from ray.autoscaler._private.updater import (
|
|
STATUS_UP_TO_DATE,
|
|
TAG_RAY_NODE_STATUS,
|
|
NodeUpdater,
|
|
)
|
|
from ray.autoscaler._private.util import with_envs, with_head_node_ip
|
|
from ray.autoscaler.node_provider import NodeProvider as NodeProviderV1
|
|
from ray.autoscaler.v2.instance_manager.config import AutoscalingConfig
|
|
from ray.core.generated.instance_manager_pb2 import Instance
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RayInstaller(object):
|
|
"""
|
|
RayInstaller is responsible for installing ray on the target instance.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
provider: NodeProviderV1,
|
|
config: AutoscalingConfig,
|
|
process_runner=subprocess,
|
|
) -> None:
|
|
self._provider = provider
|
|
self._config = config
|
|
self._process_runner = process_runner
|
|
|
|
def install_ray(self, instance: Instance, head_node_ip: str) -> None:
|
|
"""
|
|
Install ray on the target instance synchronously.
|
|
TODO:(rickyx): This runs in another thread, and errors are silently
|
|
ignored. We should propagate the error to the main thread.
|
|
"""
|
|
setup_commands = self._config.get_worker_setup_commands(instance.instance_type)
|
|
ray_start_commands = self._config.get_worker_start_ray_commands()
|
|
docker_config = self._config.get_docker_config(instance.instance_type)
|
|
|
|
logger.info(
|
|
f"Creating new (spawn_updater) updater thread for node"
|
|
f" {instance.cloud_instance_id}."
|
|
)
|
|
provider_instance_type_name = self._config.get_provider_instance_type(
|
|
instance.instance_type
|
|
)
|
|
updater = NodeUpdater(
|
|
node_id=instance.cloud_instance_id,
|
|
provider_config=self._config.get_config("provider"),
|
|
provider=self._provider,
|
|
auth_config=self._config.get_config("auth"),
|
|
cluster_name=self._config.get_config("cluster_name"),
|
|
file_mounts=self._config.get_config("file_mounts"),
|
|
initialization_commands=with_head_node_ip(
|
|
self._config.get_initialization_commands(instance.instance_type),
|
|
head_node_ip,
|
|
),
|
|
setup_commands=with_head_node_ip(setup_commands, head_node_ip),
|
|
# This will prepend envs to the begin of the ray start commands, e.g.
|
|
# `RAY_HEAD_IP=<head_node_ip> \
|
|
# RAY_CLOUD_INSTANCE_ID=<instance_id> \
|
|
# ray start --head ...`
|
|
# See src/ray/common/constants.h for ENV name definitions.
|
|
ray_start_commands=with_envs(
|
|
ray_start_commands,
|
|
{
|
|
"RAY_HEAD_IP": head_node_ip,
|
|
"RAY_CLOUD_INSTANCE_ID": instance.cloud_instance_id,
|
|
"RAY_NODE_TYPE_NAME": instance.instance_type,
|
|
"RAY_CLOUD_INSTANCE_TYPE_NAME": provider_instance_type_name,
|
|
},
|
|
),
|
|
runtime_hash=self._config.runtime_hash,
|
|
file_mounts_contents_hash=self._config.file_mounts_contents_hash,
|
|
is_head_node=False,
|
|
cluster_synced_files=self._config.get_config("cluster_synced_files"),
|
|
rsync_options={
|
|
"rsync_exclude": self._config.get_config("rsync_exclude"),
|
|
"rsync_filter": self._config.get_config("rsync_filter"),
|
|
},
|
|
use_internal_ip=True,
|
|
docker_config=docker_config,
|
|
node_resources=self._config.get_node_resources(instance.instance_type),
|
|
node_labels=self._config.get_node_labels(instance.instance_type),
|
|
process_runner=self._process_runner,
|
|
)
|
|
updater.run()
|
|
# check if the updater was successful by checking the node tags
|
|
# since the updater could hide exceptions and just set the status tag
|
|
tags = self._provider.node_tags(instance.cloud_instance_id)
|
|
if tags.get(TAG_RAY_NODE_STATUS) != STATUS_UP_TO_DATE:
|
|
raise Exception(
|
|
f"Ray installation failed with unexpected status: {tags.get(TAG_RAY_NODE_STATUS)}"
|
|
)
|