Repository: skypilot-org/skypilot
Pull request: Fix copy mounting when storage already exists
Closes #856.

Tested:
- [x] `pytest test_smoke.py::TestStorageWithCredentials`
- [x] Manually created a `MOUNT` storage in minimal.yaml and wrote to it, then used the storage in `COPY` mode in a copy of minimal.yaml to read the file earlier written.

Problem statement:
[Storage] `COPY` mode mounting fails for existing Storage objects
From @concretevitamin:

I’m trying to test this out:
```
file_mounts:
  /checkpoints:
    name: my-sky-ckpts
    mode: MOUNT
```
this works.

Then I want to add a variant of COPY, so I intuitively change mode
```
file_mounts:
  /checkpoints:
    name: my-sky-ckpts
    mode: COPY 
```
which prints
`sky.exceptions.StorageSourceError: Storage source must be specified when using COPY mode.`
I think you explained once why this is the case.  But from a UX perspective, this seems very surprising to me.

Hints from issue discussion:
+1 I was surprised by this a couple days ago.
+1 I experienced the exactly same thing as well.

Commit message: address comments

Patch to review:
diff --git a/sky/data/storage.py b/sky/data/storage.py
index da44484db45..5ffe4a083a1 100644
--- a/sky/data/storage.py
+++ b/sky/data/storage.py
@@ -441,14 +441,20 @@ def _validate_storage_spec(self) -> None:
         if self.source is None:
             # If the mode is COPY, the source must be specified
             if self.mode == StorageMode.COPY:
-                # TODO(romilb): What about when a Storage object without source
-                #  already exists in global_user_state (e.g. used for scratch)
-                #  and now the user wants to mount it in COPY mode? We should
-                #  perhaps check for existence in global_user_state here.
-                raise exceptions.StorageSourceError(
-                    'Storage source must be specified when using COPY mode.')
+                # Check if a Storage object already exists in global_user_state
+                # (e.g. used as scratch previously). Such storage objects can be
+                # mounted in copy mode even though they have no source in the
+                # yaml spec (the name is the source).
+                handle = global_user_state.get_handle_from_storage_name(
+                    self.name)
+                if handle is not None:
+                    return
+                else:
+                    raise exceptions.StorageSourceError(
+                        'Storage source must be specified when '
+                        'using COPY mode.')
             else:
-                # If source is not specified in mount mode, the intent is to
+                # If source is not specified in COPY mode, the intent is to
                 # create a bucket and use it as scratch disk. Name must be
                 # specified to create bucket.
                 if not self.name:
diff --git a/sky/task.py b/sky/task.py
index e6739b04578..945fe7d2bb9 100644
--- a/sky/task.py
+++ b/sky/task.py
@@ -475,7 +475,8 @@ def add_storage_mounts(self) -> None:
                 store_type = storage_plans[storage]
                 if store_type is storage_lib.StoreType.S3:
                     # TODO: allow for Storage mounting of different clouds
-                    if storage.source.startswith('s3://'):
+                    if storage.source is not None and \
+                            storage.source.startswith('s3://'):
                         blob_path = storage.source
                     else:
                         blob_path = 's3://' + storage.name


Code review comments:
diff_hunk: @@ -475,7 +475,8 @@ def add_storage_mounts(self) -> None:
                 store_type = storage_plans[storage]
                 if store_type is storage_lib.StoreType.S3:
                     # TODO: allow for Storage mounting of different clouds
-                    if storage.source.startswith('s3://'):
+                    if storage.source is not None and \
path: sky/task.py
text: @user1:
Nit: use parenthesis rather than backslashes, so yapf auto-formats it
```
                    if (storage.source is not None and
                            storage.source.startswith('s3://')):
```

https://google.github.io/styleguide/pyguide.html#32-line-length "Make use of Python’s [implicit line joining inside parentheses, brackets and braces](http://docs.python.org/reference/lexical_analysis.html#implicit-line-joining). If necessary, you can add an extra pair of parentheses around an expression."