import json
import logging
import os
import random
import shutil
import subprocess
import sys
import tempfile
import time
from typing import Dict, Any, Optional
import yaml
import ray
from ray.autoscaler._private.fake_multi_node.node_provider import (
FAKE_DOCKER_DEFAULT_CLIENT_PORT,
FAKE_DOCKER_DEFAULT_GCS_PORT,
)
from ray.util.ml_utils.dict import deep_update
logger = logging.getLogger(__name__)
DEFAULT_DOCKER_IMAGE = "rayproject/ray:nightly-py{major}{minor}-cpu"
class ResourcesNotReadyError(RuntimeError):
pass
[docs]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
@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
)
[docs] 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 (bool): If True, uses Ray client to connect to the
cluster. If False, uses GCS to connect to the cluster.
timeout (int): 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://{host}:{port}"
else:
port = self.gcs_port
address = f"{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}")
[docs] @staticmethod
def wait_for_resources(resources: Dict[str, float], timeout: int = 60):
"""Wait until Ray cluster resources are available
Args:
resources (Dict[str, float]): Minimum resources needed before
this function returns.
timeout (int): 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()
[docs] 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 (Dict[str, Any]): 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_output(
f"docker pull {self._docker_image}", shell=True
)
except Exception as e:
logger.error(f"Error pulling image {self._docker_image}: {e}")
[docs] 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()
[docs] def teardown(self, keep_dir: bool = False):
"""Tear down docker compose cluster environment.
Args:
keep_dir (bool): 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(
["python", 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
[docs] def start(self):
"""Start docker compose cluster.
Starts the monitor process and runs ``ray up``.
"""
self._start_monitor()
subprocess.check_output(
f"RAY_FAKE_CLUSTER=1 ray up -y {self.config_file}", shell=True
)
[docs] def stop(self):
"""Stop docker compose cluster.
Runs ``ray down`` and stops the monitor process.
"""
if ray.is_initialized:
ray.shutdown()
subprocess.check_output(
f"RAY_FAKE_CLUSTER=1 ray down -y {self.config_file}", shell=True
)
self._stop_monitor()
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"]
[docs] 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_output(f"docker kill {container}", shell=True)
