chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,91 @@
|
||||
import os
|
||||
import subprocess
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
from ray.autoscaler._private.docker import with_docker_exec
|
||||
from ray.autoscaler.command_runner import CommandRunnerInterface
|
||||
|
||||
|
||||
class FakeDockerCommandRunner(CommandRunnerInterface):
|
||||
"""Command runner for the fke docker multinode cluster.
|
||||
|
||||
This command runner uses ``docker exec`` and ``docker cp`` to
|
||||
run commands and copy files, respectively.
|
||||
|
||||
The regular ``DockerCommandRunner`` is made for use in SSH settings
|
||||
where Docker runs on a remote hose. In contrast, this command runner
|
||||
does not wrap the docker commands in ssh calls.
|
||||
"""
|
||||
|
||||
def __init__(self, docker_config, **common_args):
|
||||
self.container_name = docker_config["container_name"]
|
||||
self.docker_config = docker_config
|
||||
self.home_dir = None
|
||||
self.initialized = False
|
||||
# Optionally use 'podman' instead of 'docker'
|
||||
use_podman = docker_config.get("use_podman", False)
|
||||
self.docker_cmd = "podman" if use_podman else "docker"
|
||||
|
||||
def _run_shell(self, cmd: str, timeout: int = 120) -> str:
|
||||
return subprocess.check_output(
|
||||
cmd, shell=True, timeout=timeout, encoding="utf-8"
|
||||
)
|
||||
|
||||
def run(
|
||||
self,
|
||||
cmd: str = None,
|
||||
timeout: int = 120,
|
||||
exit_on_fail: bool = False,
|
||||
port_forward: List[Tuple[int, int]] = None,
|
||||
with_output: bool = False,
|
||||
environment_variables: Dict[str, object] = None,
|
||||
run_env: str = "auto",
|
||||
ssh_options_override_ssh_key: str = "",
|
||||
shutdown_after_run: bool = False,
|
||||
) -> str:
|
||||
prefix = with_docker_exec(
|
||||
[cmd],
|
||||
container_name=self.container_name,
|
||||
with_interactive=False,
|
||||
docker_cmd=self.docker_cmd,
|
||||
)[0]
|
||||
return self._run_shell(prefix)
|
||||
|
||||
def run_init(
|
||||
self, *, as_head: bool, file_mounts: Dict[str, str], sync_run_yet: bool
|
||||
):
|
||||
pass
|
||||
|
||||
def remote_shell_command_str(self):
|
||||
return "{} exec -it {} bash".format(self.docker_cmd, self.container_name)
|
||||
|
||||
def run_rsync_down(self, source, target, options=None):
|
||||
docker_dir = os.path.dirname(self._docker_expand_user(source))
|
||||
|
||||
self._run_shell(f"docker cp {self.container_name}:{docker_dir} {target}")
|
||||
|
||||
def run_rsync_up(self, source, target, options=None):
|
||||
docker_dir = os.path.dirname(self._docker_expand_user(target))
|
||||
self.run(cmd=f"mkdir -p {docker_dir}")
|
||||
|
||||
self._run_shell(f"docker cp {source} {self.container_name}:{docker_dir}")
|
||||
|
||||
def _docker_expand_user(self, string, any_char=False):
|
||||
user_pos = string.find("~")
|
||||
if user_pos > -1:
|
||||
if self.home_dir is None:
|
||||
self.home_dir = self._run_shell(
|
||||
with_docker_exec(
|
||||
["printenv HOME"],
|
||||
container_name=self.container_name,
|
||||
docker_cmd=self.docker_cmd,
|
||||
)
|
||||
).strip()
|
||||
|
||||
if any_char:
|
||||
return string.replace("~/", self.home_dir + "/")
|
||||
|
||||
elif not any_char and user_pos == 0:
|
||||
return string.replace("~", self.home_dir, 1)
|
||||
|
||||
return string
|
||||
@@ -0,0 +1,246 @@
|
||||
"""Fake multinode docker monitoring script.
|
||||
|
||||
This script is the "docker compose server" for the fake_multinode
|
||||
provider using Docker compose. It should be started before running
|
||||
`RAY_FAKE_CLUSTER=1 ray up <cluster_config>`.
|
||||
|
||||
This script reads the volume directory from a supplied fake multinode
|
||||
docker cluster config file.
|
||||
It then waits until a docker-compose.yaml file is created in the same
|
||||
directory, which is done by the `ray up` command.
|
||||
|
||||
It then watches for changes in the docker-compose.yaml file and runs
|
||||
`docker compose up` whenever changes are detected. This will start docker
|
||||
containers as requested by the autoscaler.
|
||||
|
||||
Generally, the docker-compose.yaml will be mounted in the head node of the
|
||||
cluster, which will then continue to change it according to the autoscaler
|
||||
requirements.
|
||||
|
||||
Additionally, this script monitors the docker container status using
|
||||
`docker status` and writes it into a `status.json`. This information is
|
||||
again used by the autoscaler to determine if any nodes have died.
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
def _read_yaml(path: str):
|
||||
with open(path, "rt") as f:
|
||||
return yaml.safe_load(f)
|
||||
|
||||
|
||||
def _update_docker_compose(
|
||||
docker_compose_path: str, project_name: str, status: Optional[Dict[str, Any]]
|
||||
) -> bool:
|
||||
docker_compose_config = _read_yaml(docker_compose_path)
|
||||
|
||||
if not docker_compose_config:
|
||||
print("Docker compose currently empty")
|
||||
return False
|
||||
|
||||
cmd = ["up", "-d"]
|
||||
if status and len(status) > 0:
|
||||
cmd += ["--no-recreate"]
|
||||
|
||||
shutdown = False
|
||||
if not docker_compose_config["services"]:
|
||||
# If no more nodes, run `down` instead of `up`
|
||||
print("Shutting down nodes")
|
||||
cmd = ["down"]
|
||||
shutdown = True
|
||||
try:
|
||||
subprocess.check_call(
|
||||
["docker", "compose", "-f", docker_compose_path, "-p", project_name]
|
||||
+ cmd
|
||||
+ [
|
||||
"--remove-orphans",
|
||||
]
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Ran into error when updating docker compose: {e}")
|
||||
# Ignore error
|
||||
|
||||
return shutdown
|
||||
|
||||
|
||||
def _get_ip(
|
||||
project_name: str,
|
||||
container_name: str,
|
||||
override_network: Optional[str] = None,
|
||||
retry_times: int = 3,
|
||||
) -> Optional[str]:
|
||||
network = override_network or f"{project_name}_ray_local"
|
||||
|
||||
cmd = [
|
||||
"docker",
|
||||
"inspect",
|
||||
"-f",
|
||||
'"{{ .NetworkSettings.Networks' f".{network}.IPAddress" ' }}"',
|
||||
f"{container_name}",
|
||||
]
|
||||
for i in range(retry_times):
|
||||
try:
|
||||
ip_address = subprocess.check_output(cmd, encoding="utf-8")
|
||||
except Exception:
|
||||
time.sleep(1)
|
||||
else:
|
||||
return ip_address.strip().strip('"').strip('\\"')
|
||||
return None
|
||||
|
||||
|
||||
def _update_docker_status(
|
||||
docker_compose_path: str, project_name: str, docker_status_path: str
|
||||
):
|
||||
data_str = ""
|
||||
try:
|
||||
data_str = (
|
||||
subprocess.check_output(
|
||||
[
|
||||
"docker",
|
||||
"compose",
|
||||
"-f",
|
||||
docker_compose_path,
|
||||
"-p",
|
||||
project_name,
|
||||
"ps",
|
||||
"--format",
|
||||
"json",
|
||||
]
|
||||
)
|
||||
.decode("utf-8")
|
||||
.strip()
|
||||
.split("\n")
|
||||
)
|
||||
data: List[Dict[str, str]] = []
|
||||
for line in data_str:
|
||||
line = line.strip()
|
||||
if line:
|
||||
data.append(json.loads(line))
|
||||
except Exception as e:
|
||||
print(f"Ran into error when fetching status: {e}")
|
||||
print(f"docker compose ps output: {data_str}")
|
||||
return None
|
||||
|
||||
status = {}
|
||||
for container in data:
|
||||
node_id = container["Service"]
|
||||
container_name = container["Name"]
|
||||
if container["State"] == "running":
|
||||
ip = _get_ip(project_name, container_name)
|
||||
else:
|
||||
ip = ""
|
||||
container["IP"] = ip
|
||||
status[node_id] = container
|
||||
|
||||
with open(docker_status_path, "wt") as f:
|
||||
json.dump(status, f)
|
||||
|
||||
return status
|
||||
|
||||
|
||||
def monitor_docker(
|
||||
docker_compose_path: str,
|
||||
status_path: str,
|
||||
project_name: str,
|
||||
update_interval: float = 1.0,
|
||||
):
|
||||
while not os.path.exists(docker_compose_path):
|
||||
# Wait until cluster is created
|
||||
time.sleep(0.5)
|
||||
|
||||
print("Docker compose config detected, starting status monitoring")
|
||||
|
||||
# Make sure this is always writeable from inside the containers
|
||||
os.chmod(docker_compose_path, 0o777)
|
||||
|
||||
docker_config = {"force_update": True}
|
||||
|
||||
# Force update
|
||||
next_update = time.monotonic() - 1.0
|
||||
shutdown = False
|
||||
status = None
|
||||
|
||||
# Loop:
|
||||
# If the config changed, update cluster.
|
||||
# Every `update_interval` seconds, update docker status.
|
||||
while not shutdown:
|
||||
new_docker_config = _read_yaml(docker_compose_path)
|
||||
if new_docker_config != docker_config:
|
||||
# Update cluster
|
||||
shutdown = _update_docker_compose(docker_compose_path, project_name, status)
|
||||
|
||||
# Force status update
|
||||
next_update = time.monotonic() - 1.0
|
||||
|
||||
if time.monotonic() > next_update:
|
||||
# Update docker status
|
||||
status = _update_docker_status(
|
||||
docker_compose_path, project_name, status_path
|
||||
)
|
||||
next_update = time.monotonic() + update_interval
|
||||
|
||||
docker_config = new_docker_config
|
||||
time.sleep(0.1)
|
||||
|
||||
print("Cluster shut down, terminating monitoring script.")
|
||||
|
||||
|
||||
def start_monitor(config_file: str):
|
||||
cluster_config = _read_yaml(config_file)
|
||||
|
||||
provider_config = cluster_config["provider"]
|
||||
assert provider_config["type"] == "fake_multinode_docker", (
|
||||
f"The docker monitor only works with providers of type "
|
||||
f"`fake_multinode_docker`, got `{provider_config['type']}`"
|
||||
)
|
||||
|
||||
project_name = provider_config["project_name"]
|
||||
|
||||
volume_dir = provider_config["shared_volume_dir"]
|
||||
os.makedirs(volume_dir, mode=0o755, exist_ok=True)
|
||||
|
||||
# Create bootstrap config
|
||||
bootstrap_config_path = os.path.join(volume_dir, "bootstrap_config.yaml")
|
||||
shutil.copy(config_file, bootstrap_config_path)
|
||||
|
||||
# These two files usually don't exist, yet
|
||||
docker_compose_config_path = os.path.join(volume_dir, "docker-compose.yaml")
|
||||
|
||||
docker_status_path = os.path.join(volume_dir, "status.json")
|
||||
|
||||
if os.path.exists(docker_compose_config_path):
|
||||
# We wait until this file exists, so remove it if it exists
|
||||
# from a previous run.
|
||||
os.remove(docker_compose_config_path)
|
||||
|
||||
if os.path.exists(docker_status_path):
|
||||
os.remove(docker_status_path)
|
||||
# Create empty file so it can be mounted
|
||||
with open(docker_status_path, "wt") as f:
|
||||
f.write("{}")
|
||||
|
||||
print(
|
||||
f"Starting monitor process. Please start Ray cluster with:\n"
|
||||
f" RAY_FAKE_CLUSTER=1 ray up {config_file}"
|
||||
)
|
||||
monitor_docker(docker_compose_config_path, docker_status_path, project_name)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"config_file",
|
||||
help="Path to cluster config file containing a fake docker "
|
||||
"cluster configuration.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
start_monitor(args.config_file)
|
||||
@@ -0,0 +1,59 @@
|
||||
# Example command to start a cluster with this config:
|
||||
#
|
||||
# RAY_FAKE_CLUSTER=1 ray start --autoscaling-config=example.yaml --head --block
|
||||
#
|
||||
# Alternatively, you can programmatically create a fake autoscaling cluster
|
||||
# using ray.cluster_utils.AutoscalingCluster.
|
||||
cluster_name: fake_multinode
|
||||
max_workers: 8
|
||||
provider:
|
||||
type: fake_multinode
|
||||
# This must be true since the nodes share the same ip!
|
||||
use_node_id_as_ip: True
|
||||
disable_node_updaters: True
|
||||
disable_launch_config_check: True
|
||||
available_node_types:
|
||||
ray.head.default:
|
||||
# You must set this manually to your "head" node resources!! The head
|
||||
# node is launched via `ray start` and hence the autoscaler cannot
|
||||
# configure its resources. The resources specified for its node type
|
||||
# must line up with what Ray detects/is configured with on start.
|
||||
resources:
|
||||
CPU: 8 # <-- set this to num CPUs used/detected in `ray start`
|
||||
GPU: 0 # <-- set this to num GPUs used/detected in `ray start`
|
||||
node_config: {}
|
||||
max_workers: 0
|
||||
ray.worker.cpu:
|
||||
resources:
|
||||
CPU: 1
|
||||
object_store_memory: 1000000000
|
||||
node_config: {}
|
||||
min_workers: 0
|
||||
max_workers: 4
|
||||
ray.worker.gpu:
|
||||
resources:
|
||||
CPU: 4
|
||||
GPU: 1
|
||||
object_store_memory: 1000000000
|
||||
node_config: {}
|
||||
min_workers: 0
|
||||
max_workers: 2
|
||||
head_node_type: ray.head.default
|
||||
upscaling_speed: 1.0
|
||||
idle_timeout_minutes: 0.1
|
||||
#
|
||||
# !!! Configurations below are not supported in fake cluster mode !!!
|
||||
#
|
||||
auth: {}
|
||||
docker: {}
|
||||
initialization_commands: []
|
||||
setup_commands: []
|
||||
head_setup_commands: []
|
||||
worker_setup_commands: []
|
||||
head_start_ray_commands: []
|
||||
worker_start_ray_commands: []
|
||||
file_mounts: {}
|
||||
cluster_synced_files: []
|
||||
file_mounts_sync_continuously: false
|
||||
rsync_exclude: []
|
||||
rsync_filter: []
|
||||
@@ -0,0 +1,85 @@
|
||||
# This is an example config file to start a local
|
||||
# multi-node cluster using Docker compose.
|
||||
# It requires the ``docker compose`` plugin to be installed:
|
||||
# https://docs.docker.com/compose/cli-command/#installing-compose-v2
|
||||
# The resulting cluster will consist of docker containers
|
||||
# scheduled via docker compose. These containers behave just like
|
||||
# regular Ray nodes, have their own IPs, and can SSH into each other.
|
||||
# They are mostly used to test multi-node setups and autoscaling on
|
||||
# a single node.
|
||||
|
||||
# Example command to start a cluster with this config:
|
||||
#
|
||||
# python docker_monitor.py example_docker.yaml &
|
||||
# RAY_FAKE_DOCKER=1 ray up -y example_docker.yaml
|
||||
cluster_name: fake_docker
|
||||
max_workers: 8
|
||||
provider:
|
||||
type: fake_multinode_docker
|
||||
disable_launch_config_check: True
|
||||
disable_node_updaters: True
|
||||
# Docker-compose config
|
||||
project_name: fake_docker
|
||||
image: rayproject/ray:nightly
|
||||
shared_volume_dir: /tmp/fake_docker
|
||||
|
||||
# For now, this has to be set here separately again:
|
||||
head_resources:
|
||||
CPU: 4
|
||||
GPU: 0
|
||||
|
||||
auth:
|
||||
ssh_user: ubuntu
|
||||
|
||||
available_node_types:
|
||||
ray.head.default:
|
||||
# You must set this manually to your "head" node resources!! The head
|
||||
# node is launched via `ray start` and hence the autoscaler cannot
|
||||
# configure its resources. The resources specified for its node type
|
||||
# must line up with what Ray detects/is configured with on start.
|
||||
resources:
|
||||
CPU: 4
|
||||
GPU: 0
|
||||
node_config: {}
|
||||
max_workers: 0
|
||||
ray.worker.cpu:
|
||||
resources:
|
||||
CPU: 2
|
||||
object_store_memory: 1000000000
|
||||
node_config: {}
|
||||
min_workers: 1
|
||||
max_workers: 4
|
||||
ray.worker.gpu:
|
||||
resources:
|
||||
CPU: 4
|
||||
GPU: 1
|
||||
object_store_memory: 1000000000
|
||||
node_config: {}
|
||||
min_workers: 1
|
||||
max_workers: 2
|
||||
head_node_type: ray.head.default
|
||||
upscaling_speed: 1.0
|
||||
idle_timeout_minutes: 0.1
|
||||
|
||||
# The start commands currently don't work - docker doesn't seem to like docker exec
|
||||
# and Ray only works when including it in the docker-compose command
|
||||
head_start_ray_commands: []
|
||||
worker_start_ray_commands: []
|
||||
|
||||
# The docker config is currently not propagated to the node provider config.
|
||||
# Thus, docker-specific configuration is expected to go into the provider part
|
||||
# as demonstrated above.
|
||||
docker: {}
|
||||
|
||||
#
|
||||
# !!! Configurations below are not supported in fake cluster mode !!!
|
||||
#
|
||||
initialization_commands: []
|
||||
setup_commands: []
|
||||
head_setup_commands: []
|
||||
worker_setup_commands: []
|
||||
file_mounts: {}
|
||||
cluster_synced_files: []
|
||||
file_mounts_sync_continuously: false
|
||||
rsync_exclude: []
|
||||
rsync_filter: []
|
||||
@@ -0,0 +1,730 @@
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from threading import RLock
|
||||
from types import ModuleType
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import yaml
|
||||
|
||||
import ray
|
||||
import ray._private.ray_constants as ray_constants
|
||||
from ray._common.network_utils import build_address
|
||||
from ray.autoscaler._private.fake_multi_node.command_runner import (
|
||||
FakeDockerCommandRunner,
|
||||
)
|
||||
from ray.autoscaler.command_runner import CommandRunnerInterface
|
||||
from ray.autoscaler.node_provider import NodeProvider
|
||||
from ray.autoscaler.tags import (
|
||||
NODE_KIND_HEAD,
|
||||
NODE_KIND_WORKER,
|
||||
STATUS_UP_TO_DATE,
|
||||
TAG_RAY_NODE_KIND,
|
||||
TAG_RAY_NODE_NAME,
|
||||
TAG_RAY_NODE_STATUS,
|
||||
TAG_RAY_USER_NODE_TYPE,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# We generate the node ids deterministically in the fake node provider, so that
|
||||
# we can associate launched nodes with their resource reports. IDs increment
|
||||
# starting with fffff*00000 for the head node, fffff*00001, etc. for workers.
|
||||
FAKE_HEAD_NODE_ID = "fffffffffffffffffffffffffffffffffffffffffffffffffff00000"
|
||||
FAKE_HEAD_NODE_TYPE = "ray.head.default"
|
||||
|
||||
FAKE_DOCKER_DEFAULT_GCS_PORT = 16379
|
||||
FAKE_DOCKER_DEFAULT_OBJECT_MANAGER_PORT = 18076
|
||||
FAKE_DOCKER_DEFAULT_CLIENT_PORT = 10002
|
||||
|
||||
DOCKER_COMPOSE_SKELETON = {
|
||||
"services": {},
|
||||
"networks": {"ray_local": {}},
|
||||
}
|
||||
|
||||
DOCKER_NODE_SKELETON = {
|
||||
"networks": ["ray_local"],
|
||||
"mem_limit": "3000m",
|
||||
"mem_reservation": "3000m",
|
||||
"shm_size": "1200m",
|
||||
"volumes": [],
|
||||
}
|
||||
|
||||
DOCKER_HEAD_CMD = (
|
||||
'bash -c "'
|
||||
"sudo mkdir -p {volume_dir} && "
|
||||
"sudo chmod 777 {volume_dir} && "
|
||||
"touch {volume_dir}/.in_docker && "
|
||||
"sudo chown -R ray:users /cluster/node && "
|
||||
"sudo chmod -R 777 /cluster/node && "
|
||||
"sudo chown -R ray:users /cluster/shared && "
|
||||
"sudo chmod -R 777 /cluster/shared && "
|
||||
"sudo chmod 700 ~/.ssh && "
|
||||
"sudo chmod 600 ~/.ssh/authorized_keys && "
|
||||
"sudo chmod 600 ~/ray_bootstrap_key.pem && "
|
||||
"sudo chown ray:users "
|
||||
"~/.ssh ~/.ssh/authorized_keys ~/ray_bootstrap_key.pem && "
|
||||
"{ensure_ssh} && "
|
||||
"sleep 1 && "
|
||||
"RAY_FAKE_CLUSTER=1 ray start --head "
|
||||
"--autoscaling-config=~/ray_bootstrap_config.yaml "
|
||||
"--object-manager-port=8076 "
|
||||
"--num-cpus {num_cpus} "
|
||||
"--num-gpus {num_gpus} "
|
||||
# "--resources='{resources}' "
|
||||
'--block"'
|
||||
)
|
||||
|
||||
DOCKER_WORKER_CMD = (
|
||||
'bash -c "'
|
||||
"sudo mkdir -p {volume_dir} && "
|
||||
"sudo chmod 777 {volume_dir} && "
|
||||
"touch {volume_dir}/.in_docker && "
|
||||
"sudo chown -R ray:users /cluster/node && "
|
||||
"sudo chmod -R 777 /cluster/node && "
|
||||
"sudo chmod 700 ~/.ssh && "
|
||||
"sudo chmod 600 ~/.ssh/authorized_keys && "
|
||||
"sudo chown ray:users ~/.ssh ~/.ssh/authorized_keys && "
|
||||
"{ensure_ssh} && "
|
||||
"sleep 1 && "
|
||||
f"ray start --address={FAKE_HEAD_NODE_ID}:6379 "
|
||||
"--object-manager-port=8076 "
|
||||
"--num-cpus {num_cpus} "
|
||||
"--num-gpus {num_gpus} "
|
||||
# "--resources='{resources}' "
|
||||
'--block"'
|
||||
)
|
||||
|
||||
|
||||
def host_dir(container_dir: str):
|
||||
"""Replace local dir with potentially different host dir.
|
||||
|
||||
E.g. in docker-in-docker environments, the host dir might be
|
||||
different to the mounted directory in the container.
|
||||
|
||||
This method will do a simple global replace to adjust the paths.
|
||||
"""
|
||||
ray_tempdir = os.environ.get("RAY_TEMPDIR", None)
|
||||
ray_hostdir = os.environ.get("RAY_HOSTDIR", None)
|
||||
|
||||
if not ray_tempdir or not ray_hostdir:
|
||||
return container_dir
|
||||
|
||||
return container_dir.replace(ray_tempdir, ray_hostdir)
|
||||
|
||||
|
||||
def create_node_spec(
|
||||
head: bool,
|
||||
docker_image: str,
|
||||
mounted_cluster_dir: str,
|
||||
mounted_node_dir: str,
|
||||
num_cpus: int = 2,
|
||||
num_gpus: int = 0,
|
||||
resources: Optional[Dict] = None,
|
||||
env_vars: Optional[Dict] = None,
|
||||
host_gcs_port: int = 16379,
|
||||
host_object_manager_port: int = 18076,
|
||||
host_client_port: int = 10002,
|
||||
volume_dir: Optional[str] = None,
|
||||
node_state_path: Optional[str] = None,
|
||||
docker_status_path: Optional[str] = None,
|
||||
docker_compose_path: Optional[str] = None,
|
||||
bootstrap_config_path: Optional[str] = None,
|
||||
private_key_path: Optional[str] = None,
|
||||
public_key_path: Optional[str] = None,
|
||||
):
|
||||
node_spec = copy.deepcopy(DOCKER_NODE_SKELETON)
|
||||
node_spec["image"] = docker_image
|
||||
|
||||
bootstrap_cfg_path_on_container = "/home/ray/ray_bootstrap_config.yaml"
|
||||
bootstrap_key_path_on_container = "/home/ray/ray_bootstrap_key.pem"
|
||||
|
||||
resources = resources or {}
|
||||
|
||||
ensure_ssh = (
|
||||
(
|
||||
"((sudo apt update && sudo apt install -y openssh-server && "
|
||||
"sudo service ssh start) || true)"
|
||||
)
|
||||
if not bool(int(os.environ.get("RAY_HAS_SSH", "0") or "0"))
|
||||
else "sudo service ssh start"
|
||||
)
|
||||
|
||||
cmd_kwargs = dict(
|
||||
ensure_ssh=ensure_ssh,
|
||||
num_cpus=num_cpus,
|
||||
num_gpus=num_gpus,
|
||||
resources=json.dumps(resources, indent=None),
|
||||
volume_dir=volume_dir,
|
||||
autoscaling_config=bootstrap_cfg_path_on_container,
|
||||
)
|
||||
|
||||
env_vars = env_vars or {}
|
||||
|
||||
# Set to "auto" to mount current autoscaler directory to nodes for dev
|
||||
fake_cluster_dev_dir = os.environ.get("FAKE_CLUSTER_DEV", "")
|
||||
if fake_cluster_dev_dir:
|
||||
if fake_cluster_dev_dir == "auto":
|
||||
local_ray_dir = os.path.dirname(ray.__file__)
|
||||
else:
|
||||
local_ray_dir = fake_cluster_dev_dir
|
||||
os.environ["FAKE_CLUSTER_DEV"] = local_ray_dir
|
||||
|
||||
mj = sys.version_info.major
|
||||
mi = sys.version_info.minor
|
||||
|
||||
fake_modules_str = os.environ.get("FAKE_CLUSTER_DEV_MODULES", "autoscaler")
|
||||
fake_modules = fake_modules_str.split(",")
|
||||
|
||||
docker_ray_dir = f"/home/ray/anaconda3/lib/python{mj}.{mi}/site-packages/ray"
|
||||
|
||||
node_spec["volumes"] += [
|
||||
f"{local_ray_dir}/{module}:{docker_ray_dir}/{module}:ro"
|
||||
for module in fake_modules
|
||||
]
|
||||
env_vars["FAKE_CLUSTER_DEV"] = local_ray_dir
|
||||
env_vars["FAKE_CLUSTER_DEV_MODULES"] = fake_modules_str
|
||||
os.environ["FAKE_CLUSTER_DEV_MODULES"] = fake_modules_str
|
||||
|
||||
if head:
|
||||
node_spec["command"] = DOCKER_HEAD_CMD.format(**cmd_kwargs)
|
||||
# Expose ports so we can connect to the cluster from outside
|
||||
node_spec["ports"] = [
|
||||
f"{host_gcs_port}:{ray_constants.DEFAULT_PORT}",
|
||||
f"{host_object_manager_port}:8076",
|
||||
f"{host_client_port}:10001",
|
||||
]
|
||||
# Mount status and config files for the head node
|
||||
node_spec["volumes"] += [
|
||||
f"{host_dir(node_state_path)}:{node_state_path}",
|
||||
f"{host_dir(docker_status_path)}:{docker_status_path}",
|
||||
f"{host_dir(docker_compose_path)}:{docker_compose_path}",
|
||||
f"{host_dir(bootstrap_config_path)}:" f"{bootstrap_cfg_path_on_container}",
|
||||
f"{host_dir(private_key_path)}:{bootstrap_key_path_on_container}",
|
||||
]
|
||||
|
||||
# Create file if it does not exist on local filesystem
|
||||
for filename in [node_state_path, docker_status_path, bootstrap_config_path]:
|
||||
if not os.path.exists(filename):
|
||||
with open(filename, "wt") as f:
|
||||
f.write("{}")
|
||||
else:
|
||||
node_spec["command"] = DOCKER_WORKER_CMD.format(**cmd_kwargs)
|
||||
node_spec["depends_on"] = [FAKE_HEAD_NODE_ID]
|
||||
|
||||
# Mount shared directories and ssh access keys
|
||||
node_spec["volumes"] += [
|
||||
f"{host_dir(mounted_cluster_dir)}:/cluster/shared",
|
||||
f"{host_dir(mounted_node_dir)}:/cluster/node",
|
||||
f"{host_dir(public_key_path)}:/home/ray/.ssh/authorized_keys",
|
||||
]
|
||||
|
||||
# Pass these environment variables (to the head node)
|
||||
# These variables are propagated by the `docker compose` command.
|
||||
env_vars.setdefault("RAY_HAS_SSH", os.environ.get("RAY_HAS_SSH", ""))
|
||||
env_vars.setdefault("RAY_TEMPDIR", os.environ.get("RAY_TEMPDIR", ""))
|
||||
env_vars.setdefault("RAY_HOSTDIR", os.environ.get("RAY_HOSTDIR", ""))
|
||||
|
||||
node_spec["environment"] = [f"{k}={v}" for k, v in env_vars.items()]
|
||||
|
||||
return node_spec
|
||||
|
||||
|
||||
class FakeMultiNodeProvider(NodeProvider):
|
||||
"""A node provider that implements multi-node on a single machine.
|
||||
|
||||
This is used for laptop mode testing of autoscaling functionality."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider_config: Dict[str, Any],
|
||||
cluster_name: str,
|
||||
):
|
||||
"""Initialize the fake multi-node provider.
|
||||
|
||||
Args:
|
||||
provider_config: Configuration for the provider.
|
||||
cluster_name: Name of the cluster.
|
||||
"""
|
||||
|
||||
NodeProvider.__init__(self, provider_config, cluster_name)
|
||||
self.lock = RLock()
|
||||
if "RAY_FAKE_CLUSTER" not in os.environ:
|
||||
raise RuntimeError(
|
||||
"FakeMultiNodeProvider requires ray to be started with "
|
||||
"RAY_FAKE_CLUSTER=1 ray start ..."
|
||||
)
|
||||
# GCS address to use for the cluster
|
||||
self._gcs_address = provider_config.get("gcs_address", None)
|
||||
# Head node id
|
||||
self._head_node_id = provider_config.get("head_node_id", FAKE_HEAD_NODE_ID)
|
||||
# Whether to launch multiple nodes at once, or one by one regardless of
|
||||
# the count (default)
|
||||
self._launch_multiple = provider_config.get("launch_multiple", False)
|
||||
|
||||
# These are injected errors for testing purposes. If not None,
|
||||
# these will be raised on `create_node_with_resources_and_labels`` and
|
||||
# `terminate_node``, respectively.
|
||||
self._creation_error = None
|
||||
self._termination_errors = None
|
||||
|
||||
self._nodes = {
|
||||
self._head_node_id: {
|
||||
"tags": {
|
||||
TAG_RAY_NODE_KIND: NODE_KIND_HEAD,
|
||||
TAG_RAY_USER_NODE_TYPE: FAKE_HEAD_NODE_TYPE,
|
||||
TAG_RAY_NODE_NAME: self._head_node_id,
|
||||
TAG_RAY_NODE_STATUS: STATUS_UP_TO_DATE,
|
||||
}
|
||||
},
|
||||
}
|
||||
self._next_node_id = 0
|
||||
|
||||
def _next_hex_node_id(self):
|
||||
self._next_node_id += 1
|
||||
base = "fffffffffffffffffffffffffffffffffffffffffffffffffff"
|
||||
return base + str(self._next_node_id).zfill(5)
|
||||
|
||||
def non_terminated_nodes(self, tag_filters):
|
||||
with self.lock:
|
||||
nodes = []
|
||||
for node_id in self._nodes:
|
||||
tags = self.node_tags(node_id)
|
||||
ok = True
|
||||
for k, v in tag_filters.items():
|
||||
if tags.get(k) != v:
|
||||
ok = False
|
||||
if ok:
|
||||
nodes.append(node_id)
|
||||
|
||||
return nodes
|
||||
|
||||
def is_running(self, node_id):
|
||||
with self.lock:
|
||||
return node_id in self._nodes
|
||||
|
||||
def is_terminated(self, node_id):
|
||||
with self.lock:
|
||||
return node_id not in self._nodes
|
||||
|
||||
def node_tags(self, node_id):
|
||||
with self.lock:
|
||||
return self._nodes[node_id]["tags"]
|
||||
|
||||
def _get_ip(self, node_id: str) -> Optional[str]:
|
||||
return node_id
|
||||
|
||||
def external_ip(self, node_id):
|
||||
return self._get_ip(node_id)
|
||||
|
||||
def internal_ip(self, node_id):
|
||||
return self._get_ip(node_id)
|
||||
|
||||
def set_node_tags(self, node_id, tags):
|
||||
raise AssertionError("Readonly node provider cannot be updated")
|
||||
|
||||
def create_node_with_resources_and_labels(
|
||||
self, node_config, tags, count, resources, labels
|
||||
):
|
||||
if self._creation_error:
|
||||
raise self._creation_error
|
||||
|
||||
if self._launch_multiple:
|
||||
for _ in range(count):
|
||||
self._create_node_with_resources_and_labels(
|
||||
node_config, tags, count, resources, labels
|
||||
)
|
||||
else:
|
||||
self._create_node_with_resources_and_labels(
|
||||
node_config, tags, count, resources, labels
|
||||
)
|
||||
|
||||
def _create_node_with_resources_and_labels(
|
||||
self, node_config, tags, count, resources, labels
|
||||
):
|
||||
# This function calls `pop`. To avoid side effects, we make a
|
||||
# copy of `resources`.
|
||||
resources = copy.deepcopy(resources)
|
||||
with self.lock:
|
||||
node_type = tags[TAG_RAY_USER_NODE_TYPE]
|
||||
next_id = self._next_hex_node_id()
|
||||
ray_params = ray._private.parameter.RayParams(
|
||||
min_worker_port=0,
|
||||
max_worker_port=0,
|
||||
dashboard_port=None,
|
||||
dashboard_agent_listen_port=0,
|
||||
num_cpus=resources.pop("CPU", 0),
|
||||
num_gpus=resources.pop("GPU", 0),
|
||||
object_store_memory=resources.pop("object_store_memory", None),
|
||||
resources=resources,
|
||||
labels=labels,
|
||||
redis_address=build_address(
|
||||
ray._private.services.get_node_ip_address(), 6379
|
||||
)
|
||||
if not self._gcs_address
|
||||
else self._gcs_address,
|
||||
gcs_address=build_address(
|
||||
ray._private.services.get_node_ip_address(), 6379
|
||||
)
|
||||
if not self._gcs_address
|
||||
else self._gcs_address,
|
||||
env_vars={
|
||||
"RAY_OVERRIDE_NODE_ID_FOR_TESTING": next_id,
|
||||
"RAY_CLOUD_INSTANCE_ID": next_id,
|
||||
"RAY_NODE_TYPE_NAME": node_type,
|
||||
ray_constants.RESOURCES_ENVIRONMENT_VARIABLE: json.dumps(resources),
|
||||
ray_constants.LABELS_ENVIRONMENT_VARIABLE: json.dumps(labels),
|
||||
},
|
||||
)
|
||||
node = ray._private.node.Node(
|
||||
ray_params, head=False, shutdown_at_exit=False, spawn_reaper=False
|
||||
)
|
||||
all_tags = {
|
||||
TAG_RAY_NODE_KIND: NODE_KIND_WORKER,
|
||||
TAG_RAY_USER_NODE_TYPE: node_type,
|
||||
TAG_RAY_NODE_NAME: next_id,
|
||||
TAG_RAY_NODE_STATUS: STATUS_UP_TO_DATE,
|
||||
}
|
||||
all_tags.update(tags)
|
||||
self._nodes[next_id] = {
|
||||
"tags": all_tags,
|
||||
"node": node,
|
||||
}
|
||||
|
||||
def terminate_node(self, node_id):
|
||||
with self.lock:
|
||||
if self._termination_errors:
|
||||
raise self._termination_errors
|
||||
|
||||
try:
|
||||
node = self._nodes.pop(node_id)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
self._terminate_node(node)
|
||||
|
||||
def _terminate_node(self, node):
|
||||
node["node"].kill_all_processes(check_alive=False, allow_graceful=True)
|
||||
|
||||
@staticmethod
|
||||
def bootstrap_config(cluster_config):
|
||||
return cluster_config
|
||||
|
||||
############################
|
||||
# Test only methods
|
||||
############################
|
||||
def _test_set_creation_error(self, e: Exception):
|
||||
"""Set an error that will be raised on
|
||||
create_node_with_resources_and_labels."""
|
||||
self._creation_error = e
|
||||
|
||||
def _test_add_termination_errors(self, e: Exception):
|
||||
"""Set an error that will be raised on terminate_node."""
|
||||
self._termination_errors = e
|
||||
|
||||
|
||||
class FakeMultiNodeDockerProvider(FakeMultiNodeProvider):
|
||||
"""A node provider that implements multi-node on a single machine.
|
||||
|
||||
This is used for laptop mode testing of multi node functionality
|
||||
where each node has their own FS and IP."""
|
||||
|
||||
def __init__(self, provider_config, cluster_name):
|
||||
super(FakeMultiNodeDockerProvider, self).__init__(provider_config, cluster_name)
|
||||
|
||||
fake_head = copy.deepcopy(self._nodes)
|
||||
|
||||
self._project_name = self.provider_config["project_name"]
|
||||
self._docker_image = self.provider_config["image"]
|
||||
|
||||
self._host_gcs_port = self.provider_config.get(
|
||||
"host_gcs_port", FAKE_DOCKER_DEFAULT_GCS_PORT
|
||||
)
|
||||
self._host_object_manager_port = self.provider_config.get(
|
||||
"host_object_manager_port", FAKE_DOCKER_DEFAULT_OBJECT_MANAGER_PORT
|
||||
)
|
||||
self._host_client_port = self.provider_config.get(
|
||||
"host_client_port", FAKE_DOCKER_DEFAULT_CLIENT_PORT
|
||||
)
|
||||
|
||||
self._head_resources = self.provider_config["head_resources"]
|
||||
|
||||
# subdirs:
|
||||
# - ./shared (shared filesystem)
|
||||
# - ./nodes/<node_id> (node-specific mounted filesystem)
|
||||
self._volume_dir = self.provider_config["shared_volume_dir"]
|
||||
self._mounted_cluster_dir = os.path.join(self._volume_dir, "shared")
|
||||
|
||||
if not self.in_docker_container:
|
||||
# Only needed on host
|
||||
os.makedirs(self._mounted_cluster_dir, mode=0o755, exist_ok=True)
|
||||
|
||||
self._boostrap_config_path = os.path.join(
|
||||
self._volume_dir, "bootstrap_config.yaml"
|
||||
)
|
||||
|
||||
self._private_key_path = os.path.join(self._volume_dir, "bootstrap_key.pem")
|
||||
self._public_key_path = os.path.join(self._volume_dir, "bootstrap_key.pem.pub")
|
||||
|
||||
if not self.in_docker_container:
|
||||
# Create private key
|
||||
if not os.path.exists(self._private_key_path):
|
||||
subprocess.check_call(
|
||||
f'ssh-keygen -b 2048 -t rsa -q -N "" '
|
||||
f"-f {self._private_key_path}",
|
||||
shell=True,
|
||||
)
|
||||
|
||||
# Create public key
|
||||
if not os.path.exists(self._public_key_path):
|
||||
subprocess.check_call(
|
||||
f"ssh-keygen -y "
|
||||
f"-f {self._private_key_path} "
|
||||
f"> {self._public_key_path}",
|
||||
shell=True,
|
||||
)
|
||||
|
||||
self._docker_compose_config_path = os.path.join(
|
||||
self._volume_dir, "docker-compose.yaml"
|
||||
)
|
||||
self._docker_compose_config = None
|
||||
|
||||
self._node_state_path = os.path.join(self._volume_dir, "nodes.json")
|
||||
self._docker_status_path = os.path.join(self._volume_dir, "status.json")
|
||||
|
||||
self._load_node_state()
|
||||
if FAKE_HEAD_NODE_ID not in self._nodes:
|
||||
# Reset
|
||||
self._nodes = copy.deepcopy(fake_head)
|
||||
|
||||
self._nodes[FAKE_HEAD_NODE_ID][
|
||||
"node_spec"
|
||||
] = self._create_node_spec_with_resources(
|
||||
head=True, node_id=FAKE_HEAD_NODE_ID, resources=self._head_resources
|
||||
)
|
||||
self._possibly_terminated_nodes = dict()
|
||||
|
||||
self._cleanup_interval = provider_config.get("cleanup_interval", 9.5)
|
||||
|
||||
self._docker_status = {}
|
||||
|
||||
self._update_docker_compose_config()
|
||||
self._update_docker_status()
|
||||
self._save_node_state()
|
||||
|
||||
@property
|
||||
def in_docker_container(self):
|
||||
return os.path.exists(os.path.join(self._volume_dir, ".in_docker"))
|
||||
|
||||
def _create_node_spec_with_resources(
|
||||
self, head: bool, node_id: str, resources: Dict[str, Any]
|
||||
):
|
||||
resources = resources.copy()
|
||||
|
||||
# Create shared directory
|
||||
node_dir = os.path.join(self._volume_dir, "nodes", node_id)
|
||||
os.makedirs(node_dir, mode=0o777, exist_ok=True)
|
||||
|
||||
resource_str = json.dumps(resources, indent=None)
|
||||
|
||||
return create_node_spec(
|
||||
head=head,
|
||||
docker_image=self._docker_image,
|
||||
mounted_cluster_dir=self._mounted_cluster_dir,
|
||||
mounted_node_dir=node_dir,
|
||||
num_cpus=resources.pop("CPU", 0),
|
||||
num_gpus=resources.pop("GPU", 0),
|
||||
host_gcs_port=self._host_gcs_port,
|
||||
host_object_manager_port=self._host_object_manager_port,
|
||||
host_client_port=self._host_client_port,
|
||||
resources=resources,
|
||||
env_vars={
|
||||
"RAY_OVERRIDE_NODE_ID_FOR_TESTING": node_id,
|
||||
ray_constants.RESOURCES_ENVIRONMENT_VARIABLE: resource_str,
|
||||
**self.provider_config.get("env_vars", {}),
|
||||
},
|
||||
volume_dir=self._volume_dir,
|
||||
node_state_path=self._node_state_path,
|
||||
docker_status_path=self._docker_status_path,
|
||||
docker_compose_path=self._docker_compose_config_path,
|
||||
bootstrap_config_path=self._boostrap_config_path,
|
||||
public_key_path=self._public_key_path,
|
||||
private_key_path=self._private_key_path,
|
||||
)
|
||||
|
||||
def _load_node_state(self) -> bool:
|
||||
if not os.path.exists(self._node_state_path):
|
||||
return False
|
||||
try:
|
||||
with open(self._node_state_path, "rt") as f:
|
||||
nodes = json.load(f)
|
||||
except Exception:
|
||||
return False
|
||||
if not nodes:
|
||||
return False
|
||||
self._nodes = nodes
|
||||
return True
|
||||
|
||||
def _save_node_state(self):
|
||||
with open(self._node_state_path, "wt") as f:
|
||||
json.dump(self._nodes, f)
|
||||
|
||||
# Make sure this is always writeable from inside the containers
|
||||
if not self.in_docker_container:
|
||||
# Only chmod from the outer container
|
||||
os.chmod(self._node_state_path, 0o777)
|
||||
|
||||
def _update_docker_compose_config(self):
|
||||
config = copy.deepcopy(DOCKER_COMPOSE_SKELETON)
|
||||
config["services"] = {}
|
||||
for node_id, node in self._nodes.items():
|
||||
config["services"][node_id] = node["node_spec"]
|
||||
|
||||
with open(self._docker_compose_config_path, "wt") as f:
|
||||
yaml.safe_dump(config, f)
|
||||
|
||||
def _update_docker_status(self):
|
||||
if not os.path.exists(self._docker_status_path):
|
||||
return
|
||||
with open(self._docker_status_path, "rt") as f:
|
||||
self._docker_status = json.load(f)
|
||||
|
||||
def _update_nodes(self):
|
||||
for node_id in list(self._nodes):
|
||||
if not self._is_docker_running(node_id):
|
||||
self._possibly_terminated_nodes.setdefault(node_id, time.monotonic())
|
||||
else:
|
||||
self._possibly_terminated_nodes.pop(node_id, None)
|
||||
self._cleanup_nodes()
|
||||
|
||||
def _cleanup_nodes(self):
|
||||
for node_id, timestamp in list(self._possibly_terminated_nodes.items()):
|
||||
if time.monotonic() > timestamp + self._cleanup_interval:
|
||||
if not self._is_docker_running(node_id):
|
||||
self._nodes.pop(node_id, None)
|
||||
self._possibly_terminated_nodes.pop(node_id, None)
|
||||
self._save_node_state()
|
||||
|
||||
def _container_name(self, node_id):
|
||||
node_status = self._docker_status.get(node_id, {})
|
||||
timeout = time.monotonic() + 60
|
||||
while not node_status:
|
||||
if time.monotonic() > timeout:
|
||||
raise RuntimeError(f"Container for {node_id} never became available.")
|
||||
time.sleep(1)
|
||||
self._update_docker_status()
|
||||
node_status = self._docker_status.get(node_id, {})
|
||||
|
||||
return node_status["Name"]
|
||||
|
||||
def _is_docker_running(self, node_id):
|
||||
self._update_docker_status()
|
||||
|
||||
return self._docker_status.get(node_id, {}).get("State", None) == "running"
|
||||
|
||||
def non_terminated_nodes(self, tag_filters):
|
||||
self._update_nodes()
|
||||
return super(FakeMultiNodeDockerProvider, self).non_terminated_nodes(
|
||||
tag_filters
|
||||
)
|
||||
|
||||
def is_running(self, node_id):
|
||||
with self.lock:
|
||||
self._update_nodes()
|
||||
|
||||
return node_id in self._nodes and self._is_docker_running(node_id)
|
||||
|
||||
def is_terminated(self, node_id):
|
||||
with self.lock:
|
||||
self._update_nodes()
|
||||
|
||||
return node_id not in self._nodes and not self._is_docker_running(node_id)
|
||||
|
||||
def get_command_runner(
|
||||
self,
|
||||
log_prefix: str,
|
||||
node_id: str,
|
||||
auth_config: Dict[str, Any],
|
||||
cluster_name: str,
|
||||
process_runner: ModuleType,
|
||||
use_internal_ip: bool,
|
||||
docker_config: Optional[Dict[str, Any]] = None,
|
||||
) -> CommandRunnerInterface:
|
||||
if self.in_docker_container:
|
||||
return super(FakeMultiNodeProvider, self).get_command_runner(
|
||||
log_prefix,
|
||||
node_id,
|
||||
auth_config,
|
||||
cluster_name,
|
||||
process_runner,
|
||||
use_internal_ip,
|
||||
)
|
||||
|
||||
# Else, host command runner:
|
||||
common_args = {
|
||||
"log_prefix": log_prefix,
|
||||
"node_id": node_id,
|
||||
"provider": self,
|
||||
"auth_config": auth_config,
|
||||
"cluster_name": cluster_name,
|
||||
"process_runner": process_runner,
|
||||
"use_internal_ip": use_internal_ip,
|
||||
}
|
||||
|
||||
docker_config["container_name"] = self._container_name(node_id)
|
||||
docker_config["image"] = self._docker_image
|
||||
|
||||
return FakeDockerCommandRunner(docker_config, **common_args)
|
||||
|
||||
def _get_ip(self, node_id: str) -> Optional[str]:
|
||||
for i in range(3):
|
||||
self._update_docker_status()
|
||||
ip = self._docker_status.get(node_id, {}).get("IP", None)
|
||||
if ip:
|
||||
return ip
|
||||
time.sleep(3)
|
||||
return None
|
||||
|
||||
def set_node_tags(self, node_id, tags):
|
||||
assert node_id in self._nodes
|
||||
self._nodes[node_id]["tags"].update(tags)
|
||||
|
||||
def create_node_with_resources_and_labels(
|
||||
self, node_config, tags, count, resources, labels
|
||||
):
|
||||
with self.lock:
|
||||
is_head = tags[TAG_RAY_NODE_KIND] == NODE_KIND_HEAD
|
||||
|
||||
if is_head:
|
||||
next_id = FAKE_HEAD_NODE_ID
|
||||
else:
|
||||
next_id = self._next_hex_node_id()
|
||||
|
||||
self._nodes[next_id] = {
|
||||
"tags": tags,
|
||||
"node_spec": self._create_node_spec_with_resources(
|
||||
head=is_head, node_id=next_id, resources=resources
|
||||
),
|
||||
}
|
||||
self._update_docker_compose_config()
|
||||
self._save_node_state()
|
||||
|
||||
def create_node(
|
||||
self, node_config: Dict[str, Any], tags: Dict[str, str], count: int
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
resources = self._head_resources
|
||||
return self.create_node_with_resources_and_labels(
|
||||
node_config, tags, count, resources, {}
|
||||
)
|
||||
|
||||
def _terminate_node(self, node):
|
||||
self._update_docker_compose_config()
|
||||
self._save_node_state()
|
||||
|
||||
@staticmethod
|
||||
def bootstrap_config(cluster_config):
|
||||
return cluster_config
|
||||
@@ -0,0 +1,399 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import yaml
|
||||
|
||||
import ray
|
||||
from ray._common.network_utils import build_address
|
||||
from ray._private.dict import deep_update
|
||||
from ray.autoscaler._private.fake_multi_node.node_provider import (
|
||||
FAKE_DOCKER_DEFAULT_CLIENT_PORT,
|
||||
FAKE_DOCKER_DEFAULT_GCS_PORT,
|
||||
)
|
||||
from ray.util.queue import Empty, Queue
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_DOCKER_IMAGE = "rayproject/ray:nightly-py{major}{minor}-cpu"
|
||||
|
||||
|
||||
class ResourcesNotReadyError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class DockerCluster:
|
||||
"""Docker cluster wrapper.
|
||||
|
||||
Creates a directory for starting a fake multinode docker cluster.
|
||||
|
||||
Includes APIs to update the cluster config as needed in tests,
|
||||
and to start and connect to the cluster.
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[Dict[str, Any]] = None):
|
||||
self._base_config_file = os.path.join(
|
||||
os.path.dirname(__file__), "example_docker.yaml"
|
||||
)
|
||||
self._tempdir = None
|
||||
self._config_file = None
|
||||
self._nodes_file = None
|
||||
self._nodes = {}
|
||||
self._status_file = None
|
||||
self._status = {}
|
||||
self._partial_config = config
|
||||
self._cluster_config = None
|
||||
self._docker_image = None
|
||||
|
||||
self._monitor_script = os.path.join(
|
||||
os.path.dirname(__file__), "docker_monitor.py"
|
||||
)
|
||||
self._monitor_process = None
|
||||
|
||||
self._execution_thread = None
|
||||
self._execution_event = threading.Event()
|
||||
self._execution_queue = None
|
||||
|
||||
@property
|
||||
def config_file(self):
|
||||
return self._config_file
|
||||
|
||||
@property
|
||||
def cluster_config(self):
|
||||
return self._cluster_config
|
||||
|
||||
@property
|
||||
def cluster_dir(self):
|
||||
return self._tempdir
|
||||
|
||||
@property
|
||||
def gcs_port(self):
|
||||
return self._cluster_config.get("provider", {}).get(
|
||||
"host_gcs_port", FAKE_DOCKER_DEFAULT_GCS_PORT
|
||||
)
|
||||
|
||||
@property
|
||||
def client_port(self):
|
||||
return self._cluster_config.get("provider", {}).get(
|
||||
"host_client_port", FAKE_DOCKER_DEFAULT_CLIENT_PORT
|
||||
)
|
||||
|
||||
def connect(self, client: bool = True, timeout: int = 120, **init_kwargs):
|
||||
"""Connect to the docker-compose Ray cluster.
|
||||
|
||||
Assumes the cluster is at RAY_TESTHOST (defaults to
|
||||
``127.0.0.1``).
|
||||
|
||||
Args:
|
||||
client: If True, uses Ray client to connect to the
|
||||
cluster. If False, uses GCS to connect to the cluster.
|
||||
timeout: Connection timeout in seconds.
|
||||
**init_kwargs: kwargs to pass to ``ray.init()``.
|
||||
|
||||
"""
|
||||
host = os.environ.get("RAY_TESTHOST", "127.0.0.1")
|
||||
|
||||
if client:
|
||||
port = self.client_port
|
||||
address = f"ray://{build_address(host, port)}"
|
||||
else:
|
||||
port = self.gcs_port
|
||||
address = build_address(host, port)
|
||||
|
||||
timeout_at = time.monotonic() + timeout
|
||||
while time.monotonic() < timeout_at:
|
||||
try:
|
||||
ray.init(address, **init_kwargs)
|
||||
self.wait_for_resources({"CPU": 1})
|
||||
except ResourcesNotReadyError:
|
||||
time.sleep(1)
|
||||
continue
|
||||
else:
|
||||
break
|
||||
|
||||
try:
|
||||
ray.cluster_resources()
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Timed out connecting to Ray: {e}")
|
||||
|
||||
def remote_execution_api(self) -> "RemoteAPI":
|
||||
"""Create an object to control cluster state from within the cluster."""
|
||||
self._execution_queue = Queue(actor_options={"num_cpus": 0})
|
||||
stop_event = self._execution_event
|
||||
|
||||
def entrypoint():
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
cmd, kwargs = self._execution_queue.get(timeout=1)
|
||||
except Empty:
|
||||
continue
|
||||
|
||||
if cmd == "kill_node":
|
||||
self.kill_node(**kwargs)
|
||||
|
||||
self._execution_thread = threading.Thread(target=entrypoint)
|
||||
self._execution_thread.start()
|
||||
|
||||
return RemoteAPI(self._execution_queue)
|
||||
|
||||
@staticmethod
|
||||
def wait_for_resources(resources: Dict[str, float], timeout: int = 60):
|
||||
"""Wait until Ray cluster resources are available
|
||||
|
||||
Args:
|
||||
resources: Minimum resources needed before
|
||||
this function returns.
|
||||
timeout: Timeout in seconds.
|
||||
|
||||
"""
|
||||
timeout = time.monotonic() + timeout
|
||||
|
||||
available = ray.cluster_resources()
|
||||
while any(available.get(k, 0.0) < v for k, v in resources.items()):
|
||||
if time.monotonic() > timeout:
|
||||
raise ResourcesNotReadyError(
|
||||
f"Timed out waiting for resources: {resources}"
|
||||
)
|
||||
time.sleep(1)
|
||||
available = ray.cluster_resources()
|
||||
|
||||
def update_config(self, config: Optional[Dict[str, Any]] = None):
|
||||
"""Update autoscaling config.
|
||||
|
||||
Does a deep update of the base config with a new configuration.
|
||||
This can change autoscaling behavior.
|
||||
|
||||
Args:
|
||||
config: Partial config to update current
|
||||
config with.
|
||||
|
||||
"""
|
||||
assert self._tempdir, "Call setup() first"
|
||||
|
||||
config = config or {}
|
||||
|
||||
if config:
|
||||
self._partial_config = config
|
||||
|
||||
if not config.get("provider", {}).get("image"):
|
||||
# No image specified, trying to parse from buildkite
|
||||
docker_image = os.environ.get("RAY_DOCKER_IMAGE", None)
|
||||
|
||||
if not docker_image:
|
||||
# If still no docker image, use one according to Python version
|
||||
mj = sys.version_info.major
|
||||
mi = sys.version_info.minor
|
||||
|
||||
docker_image = DEFAULT_DOCKER_IMAGE.format(major=mj, minor=mi)
|
||||
|
||||
self._docker_image = docker_image
|
||||
|
||||
with open(self._base_config_file, "rt") as f:
|
||||
cluster_config = yaml.safe_load(f)
|
||||
|
||||
if self._partial_config:
|
||||
deep_update(cluster_config, self._partial_config, new_keys_allowed=True)
|
||||
|
||||
if self._docker_image:
|
||||
cluster_config["provider"]["image"] = self._docker_image
|
||||
|
||||
cluster_config["provider"]["shared_volume_dir"] = self._tempdir
|
||||
|
||||
self._cluster_config = cluster_config
|
||||
|
||||
with open(self._config_file, "wt") as f:
|
||||
yaml.safe_dump(self._cluster_config, f)
|
||||
|
||||
logging.info(f"Updated cluster config to: {self._cluster_config}")
|
||||
|
||||
def maybe_pull_image(self):
|
||||
if self._docker_image:
|
||||
try:
|
||||
images_str = subprocess.check_output(
|
||||
f"docker image inspect {self._docker_image}", shell=True
|
||||
)
|
||||
images = json.loads(images_str)
|
||||
except Exception as e:
|
||||
logger.error(f"Error inspecting image {self._docker_image}: {e}")
|
||||
return
|
||||
|
||||
if not images:
|
||||
try:
|
||||
subprocess.check_call(
|
||||
f"docker pull {self._docker_image}", shell=True
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error pulling image {self._docker_image}: {e}")
|
||||
|
||||
def setup(self):
|
||||
"""Setup docker compose cluster environment.
|
||||
|
||||
Creates the temporary directory, writes the initial config file,
|
||||
and pulls the docker image, if required.
|
||||
"""
|
||||
self._tempdir = tempfile.mkdtemp(dir=os.environ.get("RAY_TEMPDIR", None))
|
||||
os.chmod(self._tempdir, 0o777)
|
||||
self._config_file = os.path.join(self._tempdir, "cluster.yaml")
|
||||
self._nodes_file = os.path.join(self._tempdir, "nodes.json")
|
||||
self._status_file = os.path.join(self._tempdir, "status.json")
|
||||
self.update_config()
|
||||
self.maybe_pull_image()
|
||||
|
||||
def teardown(self, keep_dir: bool = False):
|
||||
"""Tear down docker compose cluster environment.
|
||||
|
||||
Args:
|
||||
keep_dir: If True, cluster directory
|
||||
will not be removed after termination.
|
||||
"""
|
||||
if not keep_dir:
|
||||
shutil.rmtree(self._tempdir)
|
||||
self._tempdir = None
|
||||
self._config_file = None
|
||||
|
||||
def _start_monitor(self):
|
||||
self._monitor_process = subprocess.Popen(
|
||||
[sys.executable, self._monitor_script, self.config_file]
|
||||
)
|
||||
time.sleep(2)
|
||||
|
||||
def _stop_monitor(self):
|
||||
if self._monitor_process:
|
||||
self._monitor_process.wait(timeout=30)
|
||||
if self._monitor_process.poll() is None:
|
||||
self._monitor_process.terminate()
|
||||
self._monitor_process = None
|
||||
|
||||
def start(self):
|
||||
"""Start docker compose cluster.
|
||||
|
||||
Starts the monitor process and runs ``ray up``.
|
||||
"""
|
||||
self._start_monitor()
|
||||
|
||||
subprocess.check_call(
|
||||
f"RAY_FAKE_CLUSTER=1 ray up -y {self.config_file}", shell=True
|
||||
)
|
||||
|
||||
def stop(self):
|
||||
"""Stop docker compose cluster.
|
||||
|
||||
Runs ``ray down`` and stops the monitor process.
|
||||
"""
|
||||
if ray.is_initialized:
|
||||
ray.shutdown()
|
||||
|
||||
subprocess.check_call(
|
||||
f"RAY_FAKE_CLUSTER=1 ray down -y {self.config_file}", shell=True
|
||||
)
|
||||
|
||||
self._stop_monitor()
|
||||
self._execution_event.set()
|
||||
|
||||
def _update_nodes(self):
|
||||
with open(self._nodes_file, "rt") as f:
|
||||
self._nodes = json.load(f)
|
||||
|
||||
def _update_status(self):
|
||||
with open(self._status_file, "rt") as f:
|
||||
self._status = json.load(f)
|
||||
|
||||
def _get_node(
|
||||
self,
|
||||
node_id: Optional[str] = None,
|
||||
num: Optional[int] = None,
|
||||
rand: Optional[str] = None,
|
||||
) -> str:
|
||||
self._update_nodes()
|
||||
if node_id:
|
||||
assert (
|
||||
not num and not rand
|
||||
), "Only provide either `node_id`, `num`, or `random`."
|
||||
elif num:
|
||||
assert (
|
||||
not node_id and not rand
|
||||
), "Only provide either `node_id`, `num`, or `random`."
|
||||
base = "fffffffffffffffffffffffffffffffffffffffffffffffffff"
|
||||
node_id = base + str(num).zfill(5)
|
||||
elif rand:
|
||||
assert (
|
||||
not node_id and not num
|
||||
), "Only provide either `node_id`, `num`, or `random`."
|
||||
assert rand in [
|
||||
"worker",
|
||||
"any",
|
||||
], "`random` must be one of ['worker', 'any']"
|
||||
choices = list(self._nodes.keys())
|
||||
if rand == "worker":
|
||||
choices.remove(
|
||||
"fffffffffffffffffffffffffffffffffffffffffffffffffff00000"
|
||||
)
|
||||
# Else: any
|
||||
node_id = random.choice(choices)
|
||||
|
||||
assert node_id in self._nodes, f"Node with ID {node_id} is not in active nodes."
|
||||
return node_id
|
||||
|
||||
def _get_docker_container(self, node_id: str) -> Optional[str]:
|
||||
self._update_status()
|
||||
node_status = self._status.get(node_id)
|
||||
if not node_status:
|
||||
return None
|
||||
|
||||
return node_status["Name"]
|
||||
|
||||
def kill_node(
|
||||
self,
|
||||
node_id: Optional[str] = None,
|
||||
num: Optional[int] = None,
|
||||
rand: Optional[str] = None,
|
||||
):
|
||||
"""Kill node.
|
||||
|
||||
If ``node_id`` is given, kill that node.
|
||||
|
||||
If ``num`` is given, construct node_id from this number, and kill
|
||||
that node.
|
||||
|
||||
If ``rand`` is given (as either ``worker`` or ``any``), kill a random
|
||||
node.
|
||||
"""
|
||||
node_id = self._get_node(node_id=node_id, num=num, rand=rand)
|
||||
container = self._get_docker_container(node_id=node_id)
|
||||
subprocess.check_call(f"docker kill {container}", shell=True)
|
||||
|
||||
|
||||
class RemoteAPI:
|
||||
"""Remote API to control cluster state from within cluster tasks.
|
||||
|
||||
This API uses a Ray queue to interact with an execution thread on the
|
||||
host machine that will execute commands passed to the queue.
|
||||
|
||||
Instances of this class can be serialized and passed to Ray remote actors
|
||||
to interact with cluster state (but they can also be used outside actors).
|
||||
|
||||
The API subset is limited to specific commands.
|
||||
|
||||
Args:
|
||||
queue: Ray queue to push command instructions to.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, queue: Queue):
|
||||
self._queue = queue
|
||||
|
||||
def kill_node(
|
||||
self,
|
||||
node_id: Optional[str] = None,
|
||||
num: Optional[int] = None,
|
||||
rand: Optional[str] = None,
|
||||
):
|
||||
self._queue.put(("kill_node", dict(node_id=node_id, num=num, rand=rand)))
|
||||
Reference in New Issue
Block a user