Repository: skypilot-org/skypilot
Pull request: Fix the scale down problem with sky launch
This PR aims to fix #270, where `sky launch` with a task with smaller `num_nodes` on a cluster will scale down the cluster. It also, fix the problem of `sky logs --sync-down`, where previously the log_path gotten from the remote machine has an additional `\n`.

Tested:
- [x] `sky launch -c sc examples/source_check.yaml`; `sky launch -c sc examples/source_check_1node.yaml`.
- [x] `sky logs sc 1 --sync-down`

Problem statement:
`sky launch` task with fewer num_nodes on existing cluster scales down the cluster
`sky launch -c test launch.yaml` creates 4 nodes, but if we do `sky launch -c test task.yaml`, the cluster will be scaled down to 1 node. That may not be the expected behavior, since the user may only want to run the setup on the 4-node cluster. 

Proposed behavior: when launching a task with fewer num_nodes on an existing cluster, do not scale the cluster.

The two solutions I could think of all require API change:
1. Include the num_nodes in the `sky.Resources()`, so that the provision function can use the `num_nodes` in the `to_provision` to launch the cluster, rather than using the one in the task.
2. Let the provision function take the `num_nodes` argument.

```yaml
# launch.yaml
resources:
  cloud: aws
  accelerators:
    V100: 1
num_nodes: 4
```

```yaml
# task.yaml
resources:
  cloud: aws
  accelerators:
    V100: 0.5
num_nodes: 1

setup: # task specific setup
```

Hints from issue discussion:

Commit message: Format

Patch to review:
diff --git a/prototype/sky/backends/backend.py b/prototype/sky/backends/backend.py
index d2d4a52517e..305d0ba3572 100644
--- a/prototype/sky/backends/backend.py
+++ b/prototype/sky/backends/backend.py
@@ -26,6 +26,7 @@ def get_cluster_name(self) -> str:
     def provision(self,
                   task: Task,
                   to_provision: Resources,
+                  num_nodes: int,
                   dryrun: bool,
                   stream_logs: bool,
                   cluster_name: Optional[str] = None) -> ResourceHandle:
@@ -45,8 +46,8 @@ def sync_file_mounts(
     def add_storage_objects(self, task: Task) -> None:
         raise NotImplementedError
 
-    def run_post_setup(self, handle: ResourceHandle, post_setup_fn: PostSetupFn,
-                       task: Task) -> None:
+    def run_post_setup(self, handle: ResourceHandle,
+                       post_setup_fn: PostSetupFn) -> None:
         raise NotImplementedError
 
     def execute(self, handle: ResourceHandle, task: Task,
diff --git a/prototype/sky/backends/backend_utils.py b/prototype/sky/backends/backend_utils.py
index 38b39ae7df6..2827f9a4577 100644
--- a/prototype/sky/backends/backend_utils.py
+++ b/prototype/sky/backends/backend_utils.py
@@ -357,6 +357,7 @@ def _build_sky_wheel() -> pathlib.Path:
 # TODO: too many things happening here - leaky abstraction. Refactor.
 def write_cluster_config(task: task_lib.Task,
                          to_provision: Resources,
+                         num_nodes: int,
                          cluster_config_template: str,
                          cluster_name: str,
                          region: Optional[clouds.Region] = None,
@@ -454,7 +455,7 @@ def write_cluster_config(task: task_lib.Task,
                 'setup_sh_path': setup_sh_path,
                 'workdir': task.workdir,
                 'docker_image': task.docker_image,
-                'num_nodes': task.num_nodes,
+                'num_nodes': num_nodes,
                 # File mounts handling.
                 'file_mounts': wrapped_file_mounts,
                 'initialization_commands': initialization_commands or None,
diff --git a/prototype/sky/backends/cloud_vm_ray_backend.py b/prototype/sky/backends/cloud_vm_ray_backend.py
index d63ab75b382..3326a454958 100644
--- a/prototype/sky/backends/cloud_vm_ray_backend.py
+++ b/prototype/sky/backends/cloud_vm_ray_backend.py
@@ -589,7 +589,8 @@ def _try_provision_tpu(self, to_provision: Resources, config_dict) -> bool:
             raise e
 
     def _retry_region_zones(self, task: Task, to_provision: Resources,
-                            dryrun: bool, stream_logs: bool, cluster_name: str,
+                            num_nodes: int, dryrun: bool, stream_logs: bool,
+                            cluster_name: str,
                             prev_cluster_config: Optional[Dict[str, Any]]):
         """The provision retry loop."""
         style = colorama.Style
@@ -612,6 +613,7 @@ def _retry_region_zones(self, task: Task, to_provision: Resources,
             config_dict = backend_utils.write_cluster_config(
                 task,
                 to_provision,
+                num_nodes,
                 _get_cluster_config_template(to_provision.cloud),
                 region=region,
                 zones=zones,
@@ -632,7 +634,7 @@ def _retry_region_zones(self, task: Task, to_provision: Resources,
             handle = CloudVmRayBackend.ResourceHandle(
                 cluster_name=cluster_name,
                 cluster_yaml=cluster_config_file,
-                launched_nodes=task.num_nodes,
+                launched_nodes=num_nodes,
                 # OK for this to be shown in CLI as status == INIT.
                 launched_resources=to_provision,
                 tpu_delete_script=config_dict.get('tpu-delete-script'))
@@ -641,8 +643,8 @@ def _retry_region_zones(self, task: Task, to_provision: Resources,
                                                     ready=False)
 
             gang_failed, proc, stdout, stderr = self._gang_schedule_ray_up(
-                task, to_provision.cloud, cluster_config_file, log_abs_path,
-                stream_logs, prev_cluster_config)
+                to_provision.cloud, num_nodes, cluster_config_file,
+                log_abs_path, stream_logs, prev_cluster_config)
 
             if gang_failed or proc.returncode != 0:
                 if gang_failed:
@@ -676,7 +678,7 @@ def _retry_region_zones(self, task: Task, to_provision: Resources,
                 self._ensure_cluster_ray_started(handle, log_abs_path)
 
                 cluster_name = config_dict['cluster_name']
-                plural = '' if task.num_nodes == 1 else 's'
+                plural = '' if num_nodes == 1 else 's'
                 logger.info(
                     f'{style.BRIGHT}Successfully provisioned or found'
                     f' existing VM{plural}. Setup completed.{style.RESET_ALL}')
@@ -687,10 +689,9 @@ def _retry_region_zones(self, task: Task, to_provision: Resources,
         logger.error(message)
         raise exceptions.ResourcesUnavailableError()
 
-    def _gang_schedule_ray_up(self, task: Task,
-                              to_provision_cloud: clouds.Cloud,
-                              cluster_config_file: str, log_abs_path: str,
-                              stream_logs: bool,
+    def _gang_schedule_ray_up(self, to_provision_cloud: clouds.Cloud,
+                              num_nodes: int, cluster_config_file: str,
+                              log_abs_path: str, stream_logs: bool,
                               prev_cluster_config: Optional[Dict[str, Any]]):
         """Provisions a cluster via 'ray up' with gang scheduling.
 
@@ -738,13 +739,13 @@ def is_cluster_yaml_identical():
 
         ray_up_on_full_confg_only = False
         # Fast paths for "no need to gang schedule":
-        #  (1) task.num_nodes == 1, no need to pre-provision.
+        #  (1) num_nodes == 1, no need to pre-provision.
         #  (2) otherwise, if cluster configs have not changed, no need either.
         #    TODO: can relax further: if (num_nodes x requested_resources) now
         #    == existing, then skip Step 1.  Requested resources =
         #    cloud(instance_type, acc, acc_args).  If requested clouds are
         #    different, we still need to gang schedule on the new cloud.
-        if task.num_nodes == 1:
+        if num_nodes == 1:
             # ray up the full yaml.
             proc, stdout, stderr = ray_up(
                 start_streaming_at='Shared connection to')
@@ -797,7 +798,7 @@ def is_cluster_yaml_identical():
         # TODO: if requesting a large amount (say 32) of expensive VMs, this
         # may loop for a long time.  Use timeouts and treat as gang_failed.
         cluster_ready = backend_utils.wait_until_ray_cluster_ready(
-            to_provision_cloud, cluster_config_file, task.num_nodes)
+            to_provision_cloud, cluster_config_file, num_nodes)
         gang_failed = not cluster_ready
         if gang_failed or ray_up_on_full_confg_only:
             # Head OK; gang scheduling failure.
@@ -848,7 +849,7 @@ def _ensure_cluster_ray_started(self,
             stream_logs=True)
 
     def provision_with_retries(self, task: Task, to_provision: Resources,
-                               dryrun: bool, stream_logs: bool,
+                               num_nodes: int, dryrun: bool, stream_logs: bool,
                                cluster_name: str):
         """Provision with retries for all launchable resources."""
         assert cluster_name is not None, 'cluster_name must be specified.'
@@ -873,6 +874,7 @@ def provision_with_retries(self, task: Task, to_provision: Resources,
                 config_dict = self._retry_region_zones(
                     task,
                     to_provision,
+                    num_nodes,
                     dryrun=dryrun,
                     stream_logs=stream_logs,
                     cluster_name=cluster_name,
@@ -1000,6 +1002,7 @@ def _check_task_resources_smaller_than_cluster(self, handle: ResourceHandle,
                 f'existing cluster first: sky down {cluster_name}')
 
     def _check_existing_cluster(self, task: Task, to_provision: Resources,
+                                num_nodes: int,
                                 cluster_name: str) -> Tuple[str, Resources]:
         handle = global_user_state.get_handle_from_cluster_name(cluster_name)
         if handle is not None:
@@ -1007,21 +1010,22 @@ def _check_existing_cluster(self, task: Task, to_provision: Resources,
             self._check_task_resources_smaller_than_cluster(handle, task)
             # Use the existing cluster.
             assert handle.launched_resources is not None, (cluster_name, handle)
-            return cluster_name, handle.launched_resources
+            return (cluster_name, handle.launched_resources,
+                    handle.launched_nodes)
 
         logger.info(
             f'{colorama.Fore.CYAN}Creating a new cluster: "{cluster_name}" '
-            f'[{task.num_nodes}x {to_provision}].{colorama.Style.RESET_ALL}\n'
+            f'[{num_nodes}x {to_provision}].{colorama.Style.RESET_ALL}\n'
             'Tip: to reuse an existing cluster, '
             'specify --cluster-name (-c) in the CLI or use '
             'sky.launch(.., cluster_name=..) in the Python API. '
             'Run `sky status` to see existing clusters.')
-        return cluster_name, to_provision
+        return cluster_name, to_provision, num_nodes
 
-    def _set_tpu_name(self, task: Task, cluster_config_file: str,
+    def _set_tpu_name(self, cluster_config_file: str, num_nodes: int,
                       tpu_name: str) -> None:
         """Sets TPU_NAME on all nodes."""
-        ip_list = self._get_node_ips(cluster_config_file, task.num_nodes)
+        ip_list = self._get_node_ips(cluster_config_file, num_nodes)
         ssh_user, ssh_private_key = self._get_ssh_credential(
             cluster_config_file)
 
@@ -1038,6 +1042,7 @@ def _set_tpu_name(self, task: Task, cluster_config_file: str,
     def provision(self,
                   task: Task,
                   to_provision: Resources,
+                  num_nodes: int,
                   dryrun: bool,
                   stream_logs: bool,
                   cluster_name: Optional[str] = None):
@@ -1053,11 +1058,13 @@ def provision(self,
                                             self._optimize_target)
 
         if not dryrun:  # dry run doesn't need to check existing cluster.
-            cluster_name, to_provision = self._check_existing_cluster(
-                task, to_provision, cluster_name)
+            cluster_name, to_provision, num_nodes = (
+                self._check_existing_cluster(task, to_provision, num_nodes,
+                                             cluster_name))
         try:
             config_dict = provisioner.provision_with_retries(
-                task, to_provision, dryrun, stream_logs, cluster_name)
+                task, to_provision, num_nodes, dryrun, stream_logs,
+                cluster_name)
         except exceptions.ResourcesUnavailableError as e:
             # Do not remove the stopped cluster from the global state if failed
             # to start.
@@ -1067,7 +1074,7 @@ def provision(self,
             global_user_state.remove_cluster(cluster_name, terminate=True)
             raise exceptions.ResourcesUnavailableError(
                 'Failed to provision all possible launchable resources. '
-                f'Relax the task\'s resource requirements:\n {task.num_nodes}x '
+                f'Relax the task\'s resource requirements:\n {num_nodes}x '
                 f'{task.resources}') from e
         if dryrun:
             return
@@ -1077,14 +1084,14 @@ def provision(self,
         # Set TPU environment variables
         tpu_name = config_dict.get('tpu_name')
         if tpu_name is not None:
-            self._set_tpu_name(task, cluster_config_file, tpu_name)
+            self._set_tpu_name(cluster_config_file, num_nodes, tpu_name)
 
         handle = self.ResourceHandle(
             cluster_name=cluster_name,
             cluster_yaml=cluster_config_file,
             # Cache head ip in the handle to speed up ssh operations.
-            head_ip=self._get_node_ips(cluster_config_file, task.num_nodes)[0],
-            launched_nodes=task.num_nodes,
+            head_ip=self._get_node_ips(cluster_config_file, num_nodes)[0],
+            launched_nodes=num_nodes,
             launched_resources=provisioned_resources,
             # TPU.
             tpu_delete_script=config_dict.get('tpu-delete-script'))
@@ -1199,10 +1206,10 @@ def sync_to_all_nodes(src: str, dst: str, command: str):
             sync_to_all_nodes(src, dst, command)
 
     def run_post_setup(self, handle: ResourceHandle,
-                       post_setup_fn: Optional[PostSetupFn],
-                       task: Task) -> None:
+                       post_setup_fn: Optional[PostSetupFn]) -> None:
         if post_setup_fn is not None:
-            ip_list = self._get_node_ips(handle.cluster_yaml, task.num_nodes)
+            ip_list = self._get_node_ips(handle.cluster_yaml,
+                                         handle.launched_nodes)
             ssh_user, ssh_private_key = self._get_ssh_credential(
                 handle.cluster_yaml)
 
@@ -1224,12 +1231,10 @@ def sync_down_logs(self, handle: ResourceHandle, job_id: int) -> None:
         codegen = backend_utils.JobLibCodeGen()
         codegen.get_log_path(job_id)
         code = codegen.build()
-        log_dir = self.run_on_head(handle, code, stream_logs=False)[1]
+        log_dir = self.run_on_head(handle, code, stream_logs=False)[1].strip()
 
         local_log_dir = log_dir
         remote_log_dir = os.path.join(SKY_REMOTE_LOGS_ROOT, log_dir)
-        local_tasks_dir = os.path.join(local_log_dir, 'tasks')
-        os.makedirs(local_tasks_dir, exist_ok=True)
 
         style = colorama.Style
         fore = colorama.Fore
@@ -1242,7 +1247,7 @@ def rsync_down(ip: str) -> None:
             sdk.rsync(
                 handle.cluster_yaml,
                 source=f'{remote_log_dir}/*',
-                target=f'{local_log_dir}/',
+                target=f'{local_log_dir}',
                 down=True,
                 ip_address=ip,
                 use_internal_ip=False,
@@ -1253,10 +1258,10 @@ def rsync_down(ip: str) -> None:
         for i, ip in enumerate(ips):
             try:
                 # Disable the output of rsync.
-                with open(os.devnull, 'w') as f, contextlib.redirect_stderr(
-                        f), contextlib.redirect_stdout(f):
+                with contextlib.redirect_stdout(None):
                     rsync_down(ip)
-                logger.info(f'Downloaded logs from node-{i} ({ip})')
+                logger.info(f'{fore.CYAN}Downloaded logs from node-{i} ({ip})'
+                            f'{style.RESET_ALL}')
             except click.exceptions.ClickException as e:
                 # Raised by rsync_down. Remote log dir may not exist, since
                 # the job can be run on some part of the nodes.
diff --git a/prototype/sky/backends/local_docker_backend.py b/prototype/sky/backends/local_docker_backend.py
index f9163dc64f9..7848486e6bd 100644
--- a/prototype/sky/backends/local_docker_backend.py
+++ b/prototype/sky/backends/local_docker_backend.py
@@ -141,6 +141,7 @@ def _update_state(self):
     def provision(self,
                   task: Task,
                   to_provision: Resources,
+                  num_nodes: int,
                   dryrun: bool,
                   stream_logs: bool,
                   cluster_name: Optional[str] = None) -> ResourceHandle:
@@ -150,6 +151,7 @@ def provision(self,
         Since resource demands are ignored, There's no provisioning in local
         docker.
         """
+        assert num_nodes == 1, 'Local docker backend only supports 1 node'
         assert task.name is not None, 'Task name cannot be None - have you ' \
                                       'specified a task name?'
         if cluster_name is None:
@@ -203,8 +205,7 @@ def sync_file_mounts(
         self.volume_mounts[handle] = docker_mounts
 
     def run_post_setup(self, handle: ResourceHandle,
-                       post_setup_fn: Optional[PostSetupFn],
-                       task: Task) -> None:
+                       post_setup_fn: Optional[PostSetupFn]) -> None:
         """
         Launches a container and runs a sleep command on it.
 
diff --git a/prototype/sky/cli.py b/prototype/sky/cli.py
index b3ee4afd3f4..5a207cf8c37 100644
--- a/prototype/sky/cli.py
+++ b/prototype/sky/cli.py
@@ -306,6 +306,7 @@ def _create_and_ssh_into_node(
 
         handle = backend.provision(task,
                                    to_provision=to_provision,
+                                   num_nodes=task.num_nodes,
                                    dryrun=False,
                                    stream_logs=True,
                                    cluster_name=cluster_name)
@@ -946,7 +947,8 @@ def _filter(name, all_clusters):
             dummy_task.num_nodes = handle.launched_nodes
         click.secho(f'Starting cluster {name}...', bold=True)
         backend.provision(dummy_task,
-                          to_provision=None,
+                          to_provision=handle.launched_resources,
+                          num_nodes=handle.launched_nodes,
                           dryrun=False,
                           stream_logs=True,
                           cluster_name=name)
diff --git a/prototype/sky/execution.py b/prototype/sky/execution.py
index 72a31b15cdc..add66ad1797 100644
--- a/prototype/sky/execution.py
+++ b/prototype/sky/execution.py
@@ -113,6 +113,7 @@ def _execute(dag: sky.Dag,
                 # would not be updated correctly.
                 handle = backend.provision(task,
                                            task.best_resources,
+                                           task.num_nodes,
                                            dryrun=dryrun,
                                            stream_logs=stream_logs,
                                            cluster_name=cluster_name)
@@ -130,7 +131,7 @@ def _execute(dag: sky.Dag,
                                      task.get_cloud_to_remote_file_mounts())
 
         if stages is None or Stage.PRE_EXEC in stages:
-            backend.run_post_setup(handle, task.post_setup_fn, task)
+            backend.run_post_setup(handle, task.post_setup_fn)
 
         if stages is None or Stage.EXEC in stages:
             try:


Code review comments:
diff_hunk: @@ -946,7 +947,8 @@ def _filter(name, all_clusters):
             dummy_task.num_nodes = handle.launched_nodes
         click.secho(f'Starting cluster {name}...', bold=True)
         backend.provision(dummy_task,
-                          to_provision=None,
+                          to_provision=handle.launched_resources,
+                          num_nodes=handle.launched_nodes,
path: prototype/sky/cli.py
text: @user1:
It seems to me, here and the only other caller of the new signature, have the invariance that `num_nodes == task.num_nodes` where `task` is the 1st arg.  e.g., L947. If that holds, why add a new `num_nodes` arg everywhere vs. accessing `task.num_nodes`?

@author:
Good point. Thanks! I removed the new num_nodes from the provision function

diff_hunk: @@ -1053,11 +1058,13 @@ def provision(self,
                                             self._optimize_target)
 
         if not dryrun:  # dry run doesn't need to check existing cluster.
-            cluster_name, to_provision = self._check_existing_cluster(
-                task, to_provision, cluster_name)
+            cluster_name, to_provision, num_nodes = (
+                self._check_existing_cluster(task, to_provision, num_nodes,
+                                             cluster_name))
path: prototype/sky/backends/cloud_vm_ray_backend.py
text: @user1:
Would this be simpler:
1. Don't add `num_nodes` everywhere, since it's equal to `task.num_nodes` for all call sites.
2. Here, after L1063, explicitly overwrite `num_nodes = handle.launched_nodes` or rename `num_nodes -> launched_nodes` to make it more clear.  

These changes should be simpler but still solve the problem, it seems.

@author:
Good point, fixed. Thanks!