Repository: skypilot-org/skypilot
Pull request: [Storage] Continue storage deletion in `sky storage delete` when some fail
<!-- Describe the changes in this PR -->

Fix #4050.

Make `sky storage delete` continue deleting remaining storages when some fail, instead of aborting the entire operation. Failed deletions are reported individually.

We Implemented this feature via adding a `continue_on_error` parameter in `run_in_parallel`

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `conda deactivate; bash -i tests/backward_compatibility_tests.sh`


Problem statement:
[Storage] `sky storage delete -a` aborted when deletion of one storage failed
<!-- Describe the bug report / feature request here -->

`sky storage delete -a` aborted for all storage when deletion of one storage failed. Expected behaviour: print out error for corresponding storage while still delete all other storages.

```bash
sky storage delete -a
Deleting 7 storages: skypilot-workdir-txia-60a9b7bc, skypilot-workdir-txia-f8997026, skypilot-workdir-txia-cf42367e, skypilot-workdir-txia-4e6768c9, skypilot_test_compress, comp-289a68bd-c3ed-42bd-93b8-ee3735cdb835-0, comp-3525982f-f669-4c43-b2e7-f5a1e9789a2b-1. Proceed? [Y/n]: ^[[A
Error: invalid input
Deleting 7 storages: skypilot-workdir-txia-60a9b7bc, skypilot-workdir-txia-f8997026, skypilot-workdir-txia-cf42367e, skypilot-workdir-txia-4e6768c9, skypilot_test_compress, comp-289a68bd-c3ed-42bd-93b8-ee3735cdb835-0, comp-3525982f-f669-4c43-b2e7-f5a1e9789a2b-1. Proceed? [Y/n]: 
botocore.exceptions.ClientError: An error occurred (403) when calling the HeadBucket operation: Forbidden

The above exception was the direct cause of the following exception:

sky.exceptions.StorageBucketGetError: Failed to access existing bucket 'skypilot-workdir-txia-60a9b7bc'. This is likely because it is a private bucket you do not have access to.
To fix: 
  1. If you are trying to create a new bucket: use a different name.
  2. If you are trying to connect to an existing bucket: make sure your cloud credentials have access to it. To debug, consider running `aws s3 ls skypilot-workdir-txia-60a9b7bc`.
```

Hints from issue discussion:

Commit message: chore: remove TODO

Patch to review:
diff --git a/sky/cli.py b/sky/cli.py
index 1faf0003ff9..afbcec5e13e 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -3525,11 +3525,10 @@ def storage_delete(names: List[str], all: bool, yes: bool):  # pylint: disable=r
     if sum([len(names) > 0, all]) != 1:
         raise click.UsageError('Either --all or a name must be specified.')
     if all:
-        storages = sky.storage_ls()
-        if not storages:
+        names = global_user_state.get_glob_storage_name('*')
+        if not names:
             click.echo('No storage(s) to delete.')
             return
-        names = [s['name'] for s in storages]
     else:
         names = _get_glob_storages(names)
     if names:
@@ -3543,7 +3542,18 @@ def storage_delete(names: List[str], all: bool, yes: bool):  # pylint: disable=r
                 abort=True,
                 show_default=True)
 
-    subprocess_utils.run_in_parallel(sky.storage_delete, names)
+    def delete_storage(name: str) -> None:
+        try:
+            sky.storage_delete(name)
+        except (exceptions.StorageBucketDeleteError, PermissionError) as e:
+            click.secho(f'Error deleting storage {name}: {e}', fg='red')
+        except ValueError as e:
+            click.secho(f'Error deleting storage {name}: {e}', fg='red')
+        except Exception as e:  # pylint: disable=broad-except
+            with ux_utils.print_exception_no_traceback():
+                raise e
+
+    subprocess_utils.run_in_parallel(delete_storage, names)
 
 
 @cli.group(cls=_NaturalOrderGroup)
diff --git a/sky/core.py b/sky/core.py
index 9f1288d7fb6..732d7202bc5 100644
--- a/sky/core.py
+++ b/sky/core.py
@@ -915,8 +915,11 @@ def storage_delete(name: str) -> None:
     handle = global_user_state.get_handle_from_storage_name(name)
     if handle is None:
         raise ValueError(f'Storage name {name!r} not found.')
-    else:
-        storage_object = data.Storage(name=handle.storage_name,
-                                      source=handle.source,
-                                      sync_on_reconstruction=False)
-        storage_object.delete()
+
+    assert handle.storage_name == name, (
+        f'In global_user_state, storage name {name!r} does not match '
+        f'handle.storage_name {handle.storage_name!r}')
+    storage_object = data.Storage(name=handle.storage_name,
+                                  source=handle.source,
+                                  sync_on_reconstruction=False)
+    storage_object.delete()
diff --git a/sky/data/storage.py b/sky/data/storage.py
index 897f2f96b94..63544e1c353 100644
--- a/sky/data/storage.py
+++ b/sky/data/storage.py
@@ -950,18 +950,16 @@ def delete(self, store_type: Optional[StoreType] = None) -> None:
         if not self.stores:
             logger.info('No backing stores found. Deleting storage.')
             global_user_state.remove_storage(self.name)
-        if store_type:
+        if store_type is not None:
             store = self.stores[store_type]
-            is_sky_managed = store.is_sky_managed
             # We delete a store from the cloud if it's sky managed. Else just
             # remove handle and return
-            if is_sky_managed:
+            if store.is_sky_managed:
                 self.handle.remove_store(store)
                 store.delete()
                 # Check remaining stores - if none is sky managed, remove
                 # the storage from global_user_state.
-                delete = all(
-                    s.is_sky_managed is False for s in self.stores.values())
+                delete = all(not s.is_sky_managed for s in self.stores.values())
                 if delete:
                     global_user_state.remove_storage(self.name)
                 else:
@@ -1491,6 +1489,9 @@ def _delete_s3_bucket(self, bucket_name: str) -> bool:
 
         Returns:
          bool; True if bucket was deleted, False if it was deleted externally.
+
+        Raises:
+            StorageBucketDeleteError: If deleting the bucket fails.
         """
         # Deleting objects is very slow programatically
         # (i.e. bucket.objects.all().delete() is slow).
@@ -1934,6 +1935,11 @@ def _delete_gcs_bucket(self, bucket_name: str) -> bool:
 
         Returns:
          bool; True if bucket was deleted, False if it was deleted externally.
+
+        Raises:
+            StorageBucketDeleteError: If deleting the bucket fails.
+            PermissionError: If the bucket is external and the user is not
+                allowed to delete it.
         """
 
         with rich_utils.safe_status(
@@ -3096,6 +3102,9 @@ def _delete_r2_bucket(self, bucket_name: str) -> bool:
 
         Returns:
          bool; True if bucket was deleted, False if it was deleted externally.
+
+        Raises:
+            StorageBucketDeleteError: If deleting the bucket fails.
         """
         # Deleting objects is very slow programatically
         # (i.e. bucket.objects.all().delete() is slow).
@@ -3532,7 +3541,7 @@ def _create_cos_bucket(self,
 
         return self.bucket
 
-    def _delete_cos_bucket(self):
+    def _delete_cos_bucket(self) -> None:
         bucket = self.s3_resource.Bucket(self.name)
         try:
             bucket_versioning = self.s3_resource.BucketVersioning(self.name)
diff --git a/sky/global_user_state.py b/sky/global_user_state.py
index 2a5cbc7eb3f..2c1ecc33575 100644
--- a/sky/global_user_state.py
+++ b/sky/global_user_state.py
@@ -827,7 +827,7 @@ def get_storage_names_start_with(starts_with: str) -> List[str]:
 
 
 def get_storage() -> List[Dict[str, Any]]:
-    rows = _DB.cursor.execute('select * from storage')
+    rows = _DB.cursor.execute('SELECT * FROM storage')
     records = []
     for name, launched_at, handle, last_use, status in rows:
         # TODO: use namedtuple instead of dict
diff --git a/sky/utils/subprocess_utils.py b/sky/utils/subprocess_utils.py
index 992c6bbe3ff..e38ed3ec87a 100644
--- a/sky/utils/subprocess_utils.py
+++ b/sky/utils/subprocess_utils.py
@@ -101,8 +101,6 @@ def run_in_parallel(func: Callable,
                     num_threads: Optional[int] = None) -> List[Any]:
     """Run a function in parallel on a list of arguments.
 
-    The function 'func' should raise a CommandError if the command fails.
-
     Args:
         func: The function to run in parallel
         args: Iterable of arguments to pass to func
@@ -111,14 +109,16 @@ def run_in_parallel(func: Callable,
 
     Returns:
       A list of the return values of the function func, in the same order as the
-      arguments.
+        arguments.
+
+    Raises:
+        Exception: The first exception encountered.
     """
-    # Reference: https://stackoverflow.com/questions/25790279/python-multiprocessing-early-termination # pylint: disable=line-too-long
-    processes = num_threads if num_threads is not None else get_parallel_threads(
-    )
+    processes = (num_threads
+                 if num_threads is not None else get_parallel_threads())
     with pool.ThreadPool(processes=processes) as p:
-        # Run the function in parallel on the arguments, keeping the order.
-        return list(p.imap(func, args))
+        ordered_iterators = p.imap(func, args)
+        return list(ordered_iterators)
 
 
 def handle_returncode(returncode: int,


Code review comments:
diff_hunk: @@ -3543,7 +3542,18 @@ def storage_delete(names: List[str], all: bool, yes: bool):  # pylint: disable=r
                 abort=True,
                 show_default=True)
 
-    subprocess_utils.run_in_parallel(sky.storage_delete, names)
+    def delete_storage(name: str) -> None:
+        try:
+            sky.storage_delete(name)
+        except (exceptions.StorageBucketDeleteError, PermissionError) as e:
+            click.secho(f'Error deleting storage {name}: {e}', fg='red')
+        except ValueError as e:
+            click.secho(f'Error deleting storage {name}: {e}', fg='red')
path: sky/cli.py
text: @user1:
```suggestion
        except (exceptions.StorageBucketDeleteError, PermissionError, ValueError) as e:
            click.secho(f'Error deleting storage {name}: {e}', fg='red')
```
can we merge them?

@user1:
Actually, seems like here is one exception that is not being caught. After a second thought, for simplicity, we can probably just do broad except here..

```bash
(skypilot) (base) ➜  skypilot git:(remove_token) ✗ sky storage delete -a                                                   
Deleting 15 storages: skypilot-workdir-windsey-afa32d67, skypilot-workdir-windsey-81d2115f, skypilot-workdir-windsey-b8fd0382, skypilot-workdir-windsey-499178bb, skypilot-workdir-windsey-dd64cc00, skypilot-workdir-windsey-d04e17a8, skypilot-workdir-windsey-43925bf5, skypilot-workdir-windsey-3375424a, skypilot-workdir-windsey-5e824f17, skypilot-workdir-windsey-db755382, skypilot-workdir-windsey-25270ec6, skypilot-workdir-windsey-1cb4689a, skypilot-workdir-windsey-f40df433, skypilot-workdir-windsey-2ef8a0bf, skypilot-workdir-windsey-58878e24. Proceed? [Y/n]: 
No backing stores found. Deleting storage.
botocore.exceptions.ClientError: An error occurred (403) when calling the HeadBucket operation: Forbidden

The above exception was the direct cause of the following exception:

sky.exceptions.StorageBucketGetError: Failed to access existing bucket 'skypilot-workdir-windsey-afa32d67'. This is likely because it is a private bucket you do not have access to.
To fix: 
  1. If you are trying to create a new bucket: use a different name.
  2. If you are trying to connect to an existing bucket: make sure your cloud credentials have access to it. To debug, consider running `aws s3 ls skypilot-workdir-windsey-afa32d67`.
```

@author:
Yeah... I agree

diff_hunk: @@ -3525,11 +3525,10 @@ def storage_delete(names: List[str], all: bool, yes: bool):  # pylint: disable=r
     if sum([len(names) > 0, all]) != 1:
         raise click.UsageError('Either --all or a name must be specified.')
     if all:
-        storages = sky.storage_ls()
-        if not storages:
+        names = global_user_state.get_glob_storage_name('*')
path: sky/cli.py
text: @user1:
```suggestion
        # Use '*' to get all storages.
        names = global_user_state.get_glob_storage_name(cluster_name='*')
```
nit: use kwargs for literal.

diff_hunk: @@ -3543,7 +3542,18 @@ def storage_delete(names: List[str], all: bool, yes: bool):  # pylint: disable=r
                 abort=True,
                 show_default=True)
 
-    subprocess_utils.run_in_parallel(sky.storage_delete, names)
+    def delete_storage(name: str) -> None:
+        try:
+            sky.storage_delete(name)
+        except (exceptions.StorageBucketDeleteError, PermissionError) as e:
+            click.secho(f'Error deleting storage {name}: {e}', fg='red')
+        except ValueError as e:
+            click.secho(f'Error deleting storage {name}: {e}', fg='red')
+        except Exception as e:  # pylint: disable=broad-except
+            with ux_utils.print_exception_no_traceback():
+                raise e
+
path: sky/cli.py
text: @user1:
```suggestion

```
It seems like we dont need it? We can simply not `except` it. And any exceptions raised by `sky.storage_delete` should already have the surpress traceback context?

@author:
Why they should have the traceback context supressed?

@user1:
Cuz this is a user-facing API and we want to keep the error message simple. The content in the error message should be self-explanatory.