[autoscaler] Refactor multi node type autoscaler config (#10190)

python/ray/autoscaler/commands.py

@@ -5,6 +5,7 @@ import logging
 import os
 import random
 import sys
+import subprocess
 import tempfile
 import time
 from typing import Any, Dict, Optional
@@ -25,7 +26,7 @@ from ray.autoscaler.node_provider import get_node_provider, NODE_PROVIDERS, \
     PROVIDER_PRETTY_NAMES, try_get_log_state, try_logging_config, \
     try_reload_log_state
 from ray.autoscaler.tags import TAG_RAY_NODE_TYPE, TAG_RAY_LAUNCH_CONFIG, \
-    TAG_RAY_NODE_NAME, NODE_TYPE_WORKER, NODE_TYPE_HEAD
+    TAG_RAY_NODE_NAME, NODE_TYPE_WORKER, NODE_TYPE_HEAD, TAG_RAY_INSTANCE_TYPE
 
 from ray.ray_constants import AUTOSCALER_RESOURCE_REQUEST_CHANNEL
 from ray.autoscaler.updater import NodeUpdaterThread
@@ -454,11 +455,19 @@ def warn_about_bad_start_command(start_commands):
             "to ray start in the head_start_ray_commands section.")
 
 
-def get_or_create_head_node(config, config_file, no_restart, restart_only, yes,
-                            override_cluster_name):
+def get_or_create_head_node(config,
+                            config_file,
+                            no_restart,
+                            restart_only,
+                            yes,
+                            override_cluster_name,
+                            _provider=None,
+                            _runner=subprocess):
     """Create the cluster head node, which in turn creates the workers."""
-    provider = get_node_provider(config["provider"], config["cluster_name"])
+    provider = (_provider or get_node_provider(config["provider"],
+                                               config["cluster_name"]))
 
+    config = copy.deepcopy(config)
     raw_config_file = config_file  # used for printing to the user
     config_file = os.path.abspath(config_file)
     try:
@@ -508,7 +517,14 @@ def get_or_create_head_node(config, config_file, no_restart, restart_only, yes,
                     _abort=True)
         cli_logger.newline()
 
-        launch_hash = hash_launch_conf(config["head_node"], config["auth"])
+        # TODO(ekl) this logic is duplicated in node_launcher.py (keep in sync)
+        head_node_config = copy.deepcopy(config["head_node"])
+        if "head_node_type" in config:
+            head_node_tags[TAG_RAY_INSTANCE_TYPE] = config["head_node_type"]
+            head_node_config.update(config["available_node_types"][config[
+                "head_node_type"]]["node_config"])
+
+        launch_hash = hash_launch_conf(head_node_config, config["auth"])
         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"):
@@ -540,7 +556,7 @@ def get_or_create_head_node(config, config_file, no_restart, restart_only, yes,
                 head_node_tags[TAG_RAY_LAUNCH_CONFIG] = launch_hash
                 head_node_tags[TAG_RAY_NODE_NAME] = "ray-{}-head".format(
                     config["cluster_name"])
-                provider.create_node(config["head_node"], head_node_tags, 1)
+                provider.create_node(head_node_config, head_node_tags, 1)
                 cli_logger.print("Launched a new head node")
 
                 start = time.time()
@@ -633,6 +649,7 @@ def get_or_create_head_node(config, config_file, no_restart, restart_only, yes,
                 initialization_commands=config["initialization_commands"],
                 setup_commands=init_commands,
                 ray_start_commands=ray_start_commands,
+                process_runner=_runner,
                 runtime_hash=runtime_hash,
                 file_mounts_contents_hash=file_mounts_contents_hash,
                 docker_config=config.get("docker"))