prepare for head node (#10997)

Co-authored-by: Ameer Haj Ali <ameerhajali@ameers-mbp.lan>

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-")

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(<node_id>). 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