From aba9288615a5769541a082689a1d08e9e08c553f Mon Sep 17 00:00:00 2001 From: Nikita Vemuri Date: Mon, 9 Nov 2020 20:03:15 -0800 Subject: [PATCH] [Autoscaler] Introduce callback system (#11674) Co-authored-by: Nikita Vemuri Co-authored-by: Xiayue Charles Lin Co-authored-by: Richard Liaw --- python/ray/autoscaler/_private/aws/config.py | 5 + python/ray/autoscaler/_private/commands.py | 15 +++ .../ray/autoscaler/_private/event_system.py | 100 ++++++++++++++++++ python/ray/autoscaler/_private/updater.py | 18 ++++ python/ray/autoscaler/sdk.py | 21 +++- 5 files changed, 158 insertions(+), 1 deletion(-) create mode 100644 python/ray/autoscaler/_private/event_system.py diff --git a/python/ray/autoscaler/_private/aws/config.py b/python/ray/autoscaler/_private/aws/config.py index 960e2f9b4..79fc57896 100644 --- a/python/ray/autoscaler/_private/aws/config.py +++ b/python/ray/autoscaler/_private/aws/config.py @@ -17,6 +17,8 @@ from ray.autoscaler._private.providers import _PROVIDER_PRETTY_NAMES from ray.autoscaler._private.aws.utils import LazyDefaultDict, \ handle_boto_error from ray.autoscaler._private.cli_logger import cli_logger, cf +from ray.autoscaler._private.event_system import (CreateClusterEvent, + global_event_system) logger = logging.getLogger(__name__) @@ -191,6 +193,9 @@ def bootstrap_aws(config): # Configure SSH access, using an existing key pair if possible. config = _configure_key_pair(config) + global_event_system.execute_callback( + CreateClusterEvent.ssh_keypair_downloaded, + {"ssh_key_path": config["auth"]["ssh_private_key"]}) # Pick a reasonable subnet if not specified by the user. config = _configure_subnet(config) diff --git a/python/ray/autoscaler/_private/commands.py b/python/ray/autoscaler/_private/commands.py index b4cfb93e2..63cd5f9aa 100644 --- a/python/ray/autoscaler/_private/commands.py +++ b/python/ray/autoscaler/_private/commands.py @@ -36,6 +36,8 @@ from ray.autoscaler._private.cli_logger import cli_logger, cf from ray.autoscaler._private.updater import NodeUpdaterThread from ray.autoscaler._private.command_runner import set_using_login_shells, \ set_rsync_silent +from ray.autoscaler._private.event_system import (CreateClusterEvent, + global_event_system) from ray.autoscaler._private.log_timer import LogTimer from ray.worker import global_worker # type: ignore from ray.util.debug import log_once @@ -167,6 +169,8 @@ def create_or_update_cluster(config_file: str, except yaml.scanner.ScannerError as e: handle_yaml_error(e) raise + global_event_system.execute_callback(CreateClusterEvent.up_started, + {"cluster_config": config}) # todo: validate file_mounts, ssh keys, etc. @@ -476,6 +480,8 @@ def get_or_create_head_node(config: Dict[str, Any], _provider: Optional[NodeProvider] = None, _runner: ModuleType = subprocess) -> None: """Create the cluster head node, which in turn creates the workers.""" + global_event_system.execute_callback( + CreateClusterEvent.cluster_booting_started) provider = (_provider or _get_node_provider(config["provider"], config["cluster_name"])) @@ -536,6 +542,8 @@ def get_or_create_head_node(config: Dict[str, Any], if head_node is None or provider.node_tags(head_node).get( TAG_RAY_LAUNCH_CONFIG) != launch_hash: with cli_logger.group("Acquiring an up-to-date head node"): + global_event_system.execute_callback( + CreateClusterEvent.acquiring_new_head_node) if head_node is not None: cli_logger.print( "Currently running head node is out-of-date with " @@ -571,6 +579,8 @@ def get_or_create_head_node(config: Dict[str, Any], time.sleep(POLL_INTERVAL) cli_logger.newline() + global_event_system.execute_callback(CreateClusterEvent.head_node_acquired) + with cli_logger.group( "Setting up head node", _numbered=("<>", 1, 1), @@ -664,6 +674,11 @@ def get_or_create_head_node(config: Dict[str, Any], cli_logger.abort("Failed to setup head node.") sys.exit(1) + global_event_system.execute_callback( + CreateClusterEvent.cluster_booting_completed, { + "head_node_id": head_node, + }) + monitor_str = "tail -n 100 -f /tmp/ray/session_latest/logs/monitor*" if override_cluster_name: modifiers = " --cluster-name={}".format(quote(override_cluster_name)) diff --git a/python/ray/autoscaler/_private/event_system.py b/python/ray/autoscaler/_private/event_system.py new file mode 100644 index 000000000..19b5da886 --- /dev/null +++ b/python/ray/autoscaler/_private/event_system.py @@ -0,0 +1,100 @@ +from enum import Enum, auto +from typing import Any, Callable, Dict, List, Union + +from ray.autoscaler._private.cli_logger import cli_logger + + +class CreateClusterEvent(Enum): + """Events to track in ray.autoscaler.sdk.create_or_update_cluster. + + Attributes: + up_started : Invoked at the beginning of create_or_update_cluster. + ssh_keypair_downloaded : Invoked when the ssh keypair is downloaded. + cluster_booting_started : Invoked when when the cluster booting starts. + acquiring_new_head_node : Invoked before the head node is acquired. + head_node_acquired : Invoked after the head node is acquired. + ssh_control_acquired : Invoked when the node is being updated. + run_initialization_cmd : Invoked before all initialization + commands are called and again before each initialization command. + run_setup_cmd : Invoked before all setup commands are + called and again before each setup command. + start_ray_runtime : Invoked before ray start commands are run. + start_ray_runtime_completed : Invoked after ray start commands + are run. + cluster_booting_completed : Invoked after cluster booting + is completed. + """ + + up_started = auto() + ssh_keypair_downloaded = auto() + cluster_booting_started = auto() + acquiring_new_head_node = auto() + head_node_acquired = auto() + ssh_control_acquired = auto() + run_initialization_cmd = auto() + run_setup_cmd = auto() + start_ray_runtime = auto() + start_ray_runtime_completed = auto() + cluster_booting_completed = auto() + + +class _EventSystem: + """Event system that handles storing and calling callbacks for events. + + Attributes: + callback_map (Dict[str, List[Callable]]) : Stores list of callbacks + for events when registered. + """ + + def __init__(self): + self.callback_map = {} + + def add_callback_handler( + self, + event: str, + callback: Union[Callable[[Dict], None], List[Callable[[Dict], + None]]], + ): + """Stores callback handler for event. + + Args: + event (str): Event that callback should be called on. See + CreateClusterEvent for details on the events available to be + registered against. + callback (Callable[[Dict], None]): Callable object that is invoked + when specified event occurs. + """ + if event not in CreateClusterEvent.__members__.values(): + cli_logger.warning(f"{event} is not currently tracked, and this" + " callback will not be invoked.") + + self.callback_map.setdefault( + event, + []).extend([callback] if type(callback) is not list else callback) + + def execute_callback(self, event: str, event_data: Dict[str, Any] = {}): + """Executes all callbacks for event. + + Args: + event (str): Event that is invoked. See CreateClusterEvent + for details on the available events. + event_data (Dict[str, Any]): Argument that is passed to each + callable object stored for this particular event. + """ + event_data["event_name"] = event + if event in self.callback_map: + for callback in self.callback_map[event]: + callback(event_data) + + def clear_callbacks_for_event(self, event: str): + """Clears stored callable objects for event. + + Args: + event (str): Event that has callable objects stored in map. + See CreateClusterEvent for details on the available events. + """ + if event in self.callback_map: + del self.callback_map[event] + + +global_event_system = _EventSystem() diff --git a/python/ray/autoscaler/_private/updater.py b/python/ray/autoscaler/_private/updater.py index 942635426..d32c75364 100644 --- a/python/ray/autoscaler/_private/updater.py +++ b/python/ray/autoscaler/_private/updater.py @@ -17,6 +17,8 @@ from ray.autoscaler._private.cli_logger import cli_logger, cf import ray.autoscaler._private.subprocess_output_util as cmd_output_util from ray.autoscaler._private.constants import \ RESOURCES_ENVIRONMENT_VARIABLE +from ray.autoscaler._private.event_system import (CreateClusterEvent, + global_event_system) logger = logging.getLogger(__name__) @@ -271,6 +273,8 @@ class NodeUpdater: deadline = time.time() + NODE_START_WAIT_S self.wait_ready(deadline) + global_event_system.execute_callback( + CreateClusterEvent.ssh_control_acquired) node_tags = self.provider.node_tags(self.node_id) logger.debug("Node tags: {}".format(str(node_tags))) @@ -317,10 +321,15 @@ class NodeUpdater: with cli_logger.group( "Running initialization commands", _numbered=("[]", 3, 5)): + global_event_system.execute_callback( + CreateClusterEvent.run_initialization_cmd) with LogTimer( self.log_prefix + "Initialization commands", show_status=True): for cmd in self.initialization_commands: + global_event_system.execute_callback( + CreateClusterEvent.run_initialization_cmd, + {"command": cmd}) try: # Overriding the existing SSHOptions class # with a new SSHOptions class that uses @@ -352,12 +361,17 @@ class NodeUpdater: "Running setup commands", # todo: fix command numbering _numbered=("[]", 4, 6)): + global_event_system.execute_callback( + CreateClusterEvent.run_setup_cmd) with LogTimer( self.log_prefix + "Setup commands", show_status=True): total = len(self.setup_commands) for i, cmd in enumerate(self.setup_commands): + global_event_system.execute_callback( + CreateClusterEvent.run_setup_cmd, + {"command": cmd}) if cli_logger.verbosity == 0 and len(cmd) > 30: cmd_to_print = cf.bold(cmd[:30]) + "..." else: @@ -385,6 +399,8 @@ class NodeUpdater: with cli_logger.group( "Starting the Ray runtime", _numbered=("[]", 6, 6)): + global_event_system.execute_callback( + CreateClusterEvent.start_ray_runtime) with LogTimer( self.log_prefix + "Ray start commands", show_status=True): for cmd in self.ray_start_commands: @@ -409,6 +425,8 @@ class NodeUpdater: cli_logger.error("See above for stderr.") raise click.ClickException("Start command failed.") + global_event_system.execute_callback( + CreateClusterEvent.start_ray_runtime_completed) def rsync_up(self, source, target, docker_mount_if_possible=False): options = {} diff --git a/python/ray/autoscaler/sdk.py b/python/ray/autoscaler/sdk.py index c3ba09d51..c6f165ded 100644 --- a/python/ray/autoscaler/sdk.py +++ b/python/ray/autoscaler/sdk.py @@ -1,12 +1,15 @@ """IMPORTANT: this is an experimental interface and not currently stable.""" from contextlib import contextmanager -from typing import Any, Dict, Iterator, List, Optional, Union +from typing import Any, Callable, Dict, Iterator, List, Optional, Union import json import os import tempfile from ray.autoscaler._private import commands +from ray.autoscaler._private.event_system import ( # noqa: F401 + CreateClusterEvent, # noqa: F401 + global_event_system) def create_or_update_cluster(cluster_config: Union[dict, str], @@ -224,6 +227,22 @@ def fillout_defaults(config: Dict[str, Any]) -> Dict[str, Any]: return fillout_defaults(config) +def register_callback_handler( + event_name: str, + callback: Union[Callable[[Dict], None], List[Callable[[Dict], None]]], +) -> None: + """Registers a callback handler for autoscaler events. + + Args: + event_name (str): Event that callback should be called on. See + CreateClusterEvent for details on the events available to be + registered against. + callback (Callable): Callable object that is invoked + when specified event occurs. + """ + global_event_system.add_callback_handler(event_name, callback) + + def get_docker_host_mount_location(cluster_name: str) -> str: """Return host path that Docker mounts attach to.""" docker_mount_prefix = "/tmp/ray_tmp_mount/{cluster_name}"