Repository: skypilot-org/skypilot
Pull request: Fix storage deletion for spot
Closes #959.

The previous storage class cannot delete the non-sky-managed store, which is required for managed spot.

Main changes:
1. Add an internal option `force_delete` to the storage class.
2. Make controller delete the bucket after finishing.
3. Modify `./tests/run_smoke_tests.sh test_spot_storage` for testing the leakage as well.

Tested:
- [x] `sky spot launch -n test-spot-storage ./examples/test.yaml` with the following yaml and check the bucket correctly deleted:
```
resources:
  cloud: aws
  use_spot: true
  spot_recovery: failover

file_mounts:
  ~/sky_workdir:
    # Change this to the your own globally unique bucket name.
    name: leakage-test-zhwu
    source: ./examples/docker
    persistent: false
    mode: COPY
  /imagenet-image:
    source: s3://sky-imagenet-data

run: |
  set -ex
  ls ~/sky_workdir/detectron2
  ls -l /imagenet-image/datasets
```
- [x] `./tests/run_smoke_tests.sh test_spot_storage`

Problem statement:
Spot storage leakage
Although `examples/managed_spot_with_storage.yaml` set the persistent to false, our following smoke test will still have leakage in the `sky storage ls`.
```
def test_spot_storage():
    """Test storage with managed spot"""
    name = _get_cluster_name()
    yaml_str = pathlib.Path(
        'examples/managed_spot_with_storage.yaml').read_text()
    yaml_str = yaml_str.replace('sky-workdir-zhwu',
                                f'sky-test-{int(time.time())}')
    with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w') as f:
        f.write(yaml_str)
        f.flush()
        file_path = f.name
        test = Test(
            'managed-spot-storage',
            [
                f'sky spot launch -n {name} {file_path} -y',
                'sleep 60',  # Wait the spot status to be updated
                f'sky spot status | grep {name} | grep SUCCEEDED',
            ],
            f'sky spot cancel -y -n {name}',
        )
        run_one_test(test)
```

Hints from issue discussion:

Commit message: format

Patch to review:
diff --git a/examples/managed_spot_with_storage.yaml b/examples/managed_spot_with_storage.yaml
index 627efa7f0ec..53c3b402df3 100644
--- a/examples/managed_spot_with_storage.yaml
+++ b/examples/managed_spot_with_storage.yaml
@@ -15,7 +15,7 @@ file_mounts:
   ~/sky_workdir:
     # Change this to the your own globally unique bucket name.
     name: sky-workdir-zhwu
-    source: .
+    source: ./examples/docker
     persistent: false
     mode: COPY
   /imagenet-image:
@@ -23,5 +23,5 @@ file_mounts:
 
 run: |
   set -ex
-  ls ~/sky_workdir/sky
+  ls ~/sky_workdir/detectron2
   ls -l /imagenet-image/datasets
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 398cd8591e5..12cca0ef806 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -1078,7 +1078,7 @@ def _process_cli_query(
 
 
 @contextlib.contextmanager
-def subpress_output():
+def suppress_output():
     """Suppress stdout and stderr."""
     with open(os.devnull, 'w') as devnull:
         with contextlib.redirect_stdout(devnull), contextlib.redirect_stderr(
@@ -1095,7 +1095,7 @@ def _ray_launch_hash(cluster_name: str, ray_config: Dict[str, Any]) -> Set[str]:
     if ray_launch_hashes is not None:
         logger.debug('Using cached launch_caches')
         return set(ray_launch_hashes)
-    with subpress_output():
+    with suppress_output():
         ray_config = ray_commands._bootstrap_config(ray_config)  # pylint: disable=protected-access
     # Adopted from https://github.com/ray-project/ray/blob/ray-1.10.0/python/ray/autoscaler/_private/node_launcher.py#L46-L54
     # TODO(zhwu): this logic is duplicated from the ray code above (keep in sync).
@@ -1109,7 +1109,7 @@ def _ray_launch_hash(cluster_name: str, ray_config: Dict[str, Any]) -> Set[str]:
         launch_config = copy.deepcopy(launch_config)
 
         launch_config.update(node_config['node_config'])
-        with subpress_output():
+        with suppress_output():
             current_hash = ray_util.hash_launch_conf(launch_config,
                                                      ray_config['auth'])
         launch_hashes.add(current_hash)
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index fdd73cb8577..c2c595ee7eb 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -1818,7 +1818,7 @@ def rsync_down(ip: str, local_log_dir: str,
             for i, ip in enumerate(ips):
                 try:
                     # Disable the output of rsync.
-                    with backend_utils.subpress_output():
+                    with backend_utils.suppress_output():
                         rsync_down(ip, local_log_dir, remote_log_dir)
                     logger.info(
                         f'{fore.CYAN}Job {job_id} logs: downloaded from '
diff --git a/sky/cli.py b/sky/cli.py
index f58f2b8057b..5972df68592 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -2153,7 +2153,9 @@ def spot_launch(
     # Copy the local source to a bucket. The task will not be executed locally,
     # so we need to copy the files to the bucket manually here before sending to
     # the remote spot controller.
-    task.add_storage_mounts()
+    click.echo('Copying files to the bucket.')
+    with backend_utils.suppress_output():
+        task.add_storage_mounts()
 
     # Replace the source field that is local path in all storage_mounts with
     # bucket URI and remove the name field.
@@ -2175,6 +2177,7 @@ def spot_launch(
                 with ux_utils.print_exception_no_traceback():
                     raise ValueError(f'Unsupported store type: {store_type}')
             storage_obj.name = None
+            storage_obj.force_delete = True
 
     with tempfile.NamedTemporaryFile(prefix=f'sky-spot-task-{name}-',
                                      mode='w') as f:
diff --git a/sky/data/storage.py b/sky/data/storage.py
index d471639885f..1ff2f29f01b 100644
--- a/sky/data/storage.py
+++ b/sky/data/storage.py
@@ -263,8 +263,8 @@ class StorageMetadata(object):
         def __init__(
             self,
             *,
-            storage_name: str,
-            source: str,
+            storage_name: Optional[str],
+            source: Optional[str],
             sky_stores: Optional[Dict[StoreType,
                                       AbstractStore.StoreMetadata]] = None):
             assert storage_name is not None or source is not None
@@ -340,6 +340,8 @@ def __init__(self,
         assert mode in StorageMode
         self.sync_on_reconstruction = sync_on_reconstruction
 
+        self.force_delete = False
+
         # Validate and correct inputs if necessary
         self._validate_storage_spec()
 
@@ -608,6 +610,8 @@ def delete(self, store_type: Optional[StoreType] = None) -> None:
                     global_user_state.remove_storage(self.name)
                 else:
                     global_user_state.set_storage_handle(self.name, self.handle)
+            elif self.force_delete:
+                store.delete()
             # Remove store from bookkeeping
             del self.stores[store_type]
         else:
@@ -615,6 +619,8 @@ def delete(self, store_type: Optional[StoreType] = None) -> None:
                 if store.is_sky_managed:
                     self.handle.remove_store(store)
                     store.delete()
+                elif self.force_delete:
+                    store.delete()
             self.stores = {}
             # Remove storage from global_user_state if present
             global_user_state.remove_storage(self.name)
@@ -649,6 +655,7 @@ def from_yaml_config(cls, config: Dict[str, str]) -> 'Storage':
         source = config.pop('source', None)
         store = config.pop('store', None)
         mode_str = config.pop('mode', None)
+        force_delete = config.pop('_force_delete', False)
 
         if isinstance(mode_str, str):
             # Make mode case insensitive, if specified
@@ -667,6 +674,9 @@ def from_yaml_config(cls, config: Dict[str, str]) -> 'Storage':
                           mode=mode)
         if store is not None:
             storage_obj.add_store(StoreType(store.upper()))
+
+        # Add force deletion flag
+        storage_obj.force_delete = force_delete
         return storage_obj
 
     def to_yaml_config(self) -> Dict[str, str]:
@@ -690,6 +700,8 @@ def add_if_not_none(key, value):
         add_if_not_none('store', stores)
         add_if_not_none('persistent', self.persistent)
         add_if_not_none('mode', self.mode.value)
+        if self.force_delete:
+            config['_force_delete'] = True
         return config
 
 
@@ -784,12 +796,18 @@ def sync_local_dir(self) -> None:
         with backend_utils.safe_console_status(
                 f'[bold cyan]Syncing '
                 f'[green]{self.source} to s3://{self.name}/'):
-            with subprocess.Popen(sync_command.split(' '),
-                                  stderr=subprocess.PIPE) as process:
+            # TODO(zhwu): Use log_lib.run_with_log() and redirect the output
+            # to a log file.
+            with subprocess.Popen(
+                    sync_command.split(' '),
+                    stderr=subprocess.PIPE,
+            ) as process:
+                stderr = []
                 while True:
                     line = process.stderr.readline()
                     if not line:
                         break
+                    stderr.append(line)
                     str_line = line.decode('utf-8')
                     logger.info(str_line)
                     if 'Access Denied' in str_line:
@@ -802,7 +820,9 @@ def sync_local_dir(self) -> None:
                                 'the bucket is public.')
                 returncode = process.wait()
                 if returncode != 0:
+                    stderr = '\n'.join(stderr)
                     with ux_utils.print_exception_no_traceback():
+                        logger.error(f'{process.stdout}\n{stderr}')
                         raise exceptions.StorageUploadError(
                             f'Upload to S3 failed for store {self.name} and '
                             f'source {self.source}. Please check the logs.')
diff --git a/sky/spot/recovery_strategy.py b/sky/spot/recovery_strategy.py
index ad2ac6967d9..c63e94a1756 100644
--- a/sky/spot/recovery_strategy.py
+++ b/sky/spot/recovery_strategy.py
@@ -128,6 +128,7 @@ def terminate_cluster(self):
         handle = global_user_state.get_handle_from_cluster_name(
             self.cluster_name)
         if handle is not None:
+            # Do not teardown the bucket, as it may be reused after recovery.
             self.backend.teardown(handle, terminate=True)
 
 
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
index 48519eca9a7..325fff4ffe5 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -96,6 +96,9 @@ def get_storage_schema():
                 'case_insensitive_enum': [
                     mode.value for mode in storage.StorageMode
                 ]
+            },
+            '_force_delete': {
+                'type': 'boolean',
             }
         }
     }
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index d688f55691d..0b7bf0b58b1 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -536,8 +536,8 @@ def test_spot_storage():
     name = _get_cluster_name()
     yaml_str = pathlib.Path(
         'examples/managed_spot_with_storage.yaml').read_text()
-    yaml_str = yaml_str.replace('sky-workdir-zhwu',
-                                f'sky-test-{int(time.time())}')
+    storage_name = f'sky-test-{int(time.time())}'
+    yaml_str = yaml_str.replace('sky-workdir-zhwu', storage_name)
     with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w') as f:
         f.write(yaml_str)
         f.flush()
@@ -548,6 +548,7 @@ def test_spot_storage():
                 f'sky spot launch -n {name} {file_path} -y',
                 'sleep 60',  # Wait the spot status to be updated
                 f'sky spot status | grep {name} | grep SUCCEEDED',
+                f'[ $(aws s3api list-buckets --query "Buckets[?contains(Name, \'{storage_name}\')].Name" --output text | wc -l) -eq 0 ]'
             ],
             f'sky spot cancel -y -n {name}',
         )


Code review comments:
diff_hunk: @@ -784,12 +796,18 @@ def sync_local_dir(self) -> None:
         with backend_utils.safe_console_status(
                 f'[bold cyan]Syncing '
                 f'[green]{self.source} to s3://{self.name}/'):
-            with subprocess.Popen(sync_command.split(' '),
-                                  stderr=subprocess.PIPE) as process:
+            # TODO(zhwu): Use log_lib.run_with_log() and redirect the output
+            # to a log file.
+            with subprocess.Popen(
+                    sync_command.split(' '),
+                    stderr=subprocess.PIPE,
+            ) as process:
path: sky/data/storage.py
text: @user1:
curious why the formatter split this into separate lines even though nothing is changed here :)

@author:
Oh, I was going to remove the output with `stdout=subprocess.PIPE`, but removed it for debugging. Added back. Thanks!

diff_hunk: @@ -2153,7 +2153,9 @@ def spot_launch(
     # Copy the local source to a bucket. The task will not be executed locally,
     # so we need to copy the files to the bucket manually here before sending to
     # the remote spot controller.
-    task.add_storage_mounts()
+    click.echo('Copying files to the bucket.')
+    with backend_utils.suppress_output():
+        task.add_storage_mounts()
path: sky/cli.py
text: @user1:
How about using `with console.status(...):` here?

@author:
Good idea! Added. Thanks!