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