From 4ac58d54d6c163d6a94d7cd13031ee16c0766c20 Mon Sep 17 00:00:00 2001 From: Ameer Haj Ali Date: Thu, 24 Sep 2020 11:18:38 -0700 Subject: [PATCH] prepare for head node (#10997) Co-authored-by: Ameer Haj Ali --- python/ray/autoscaler/_private/commands.py | 2 ++ python/ray/autoscaler/node_provider.py | 6 +++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/python/ray/autoscaler/_private/commands.py b/python/ray/autoscaler/_private/commands.py index 9a6a189cf..b9d441cac 100644 --- a/python/ray/autoscaler/_private/commands.py +++ b/python/ray/autoscaler/_private/commands.py @@ -654,6 +654,8 @@ def get_or_create_head_node(config, remote_config["file_mounts"] = new_mounts remote_config["no_restart"] = no_restart + remote_config = provider.prepare_for_head_node(remote_config) + # Now inject the rewritten config and SSH key into the head node remote_config_file = tempfile.NamedTemporaryFile( "w", prefix="ray-bootstrap-") diff --git a/python/ray/autoscaler/node_provider.py b/python/ray/autoscaler/node_provider.py index 6f651a9da..831687946 100644 --- a/python/ray/autoscaler/node_provider.py +++ b/python/ray/autoscaler/node_provider.py @@ -242,7 +242,7 @@ class NodeProvider: process_runner, use_internal_ip, docker_config=None) -> Any: - """ Returns the CommandRunner class used to perform SSH commands. + """Returns the CommandRunner class used to perform SSH commands. Args: log_prefix(str): stores "NodeUpdater: {}: ".format(). Used @@ -271,3 +271,7 @@ class NodeProvider: return DockerCommandRunner(docker_config, **common_args) else: return SSHCommandRunner(**common_args) + + def prepare_for_head_node(self, cluster_config): + """Returns a new cluster config with custom configs for head node.""" + return cluster_config