[autoscaler] run setup commands with restart_only=True (#13836)

python/ray/autoscaler/_private/commands.py

@@ -646,7 +646,12 @@ def get_or_create_head_node(config: Dict[str, Any],
             cli_logger.print("Prepared bootstrap config")
 
         if restart_only:
-            setup_commands = []
+            # Docker may re-launch nodes, requiring setup
+            # commands to be rerun.
+            if config.get("docker", {}).get("container_name"):
+                setup_commands = config["head_setup_commands"]
+            else:
+                setup_commands = []
             ray_start_commands = config["head_start_ray_commands"]
         elif no_restart:
             setup_commands = config["head_setup_commands"]
@@ -678,7 +683,8 @@ def get_or_create_head_node(config: Dict[str, Any],
                 "rsync_exclude": config.get("rsync_exclude"),
                 "rsync_filter": config.get("rsync_filter")
             },
-            docker_config=config.get("docker"))
+            docker_config=config.get("docker"),
+            restart_only=restart_only)
         updater.start()
         updater.join()
 

python/ray/autoscaler/_private/updater.py

@@ -48,6 +48,7 @@ class NodeUpdater:
         use_internal_ip: Wwhether the node_id belongs to an internal ip
             or external ip.
         docker_config: Docker section of autoscaler yaml
+        restart_only: Whether to skip setup commands & just restart ray
     """
 
     def __init__(self,
@@ -68,7 +69,8 @@ class NodeUpdater:
                  rsync_options=None,
                  process_runner=subprocess,
                  use_internal_ip=False,
-                 docker_config=None):
+                 docker_config=None,
+                 restart_only=False):
 
         self.log_prefix = "NodeUpdater: {}: ".format(node_id)
         use_internal_ip = (use_internal_ip
@@ -106,6 +108,7 @@ class NodeUpdater:
         self.auth_config = auth_config
         self.is_head_node = is_head_node
         self.docker_config = docker_config
+        self.restart_only = restart_only
 
     def run(self):
         if cmd_output_util.does_allow_interactive(
@@ -298,6 +301,11 @@ class NodeUpdater:
                 sync_run_yet=False)
             if init_required:
                 node_tags[TAG_RAY_RUNTIME_CONFIG] += "-invalidate"
+                # This ensures that `setup_commands` are not removed
+                self.restart_only = False
+
+        if self.restart_only:
+            self.setup_commands = []
 
         # runtime_hash will only change whenever the user restarts
         # or updates their cluster with `get_or_create_head_node`