Repository: skypilot-org/skypilot
Pull request: [Storage] Add storage translation for newly created buckets
Closes #3666.

~We skip the storage name validation because it is already done when the job is submitted from the client and the guardrails are not required for the controller.~

We now translate storage objects which create new buckets to become externally sourced buckets.

E.g., 
```
/output/:
  name: my-new-bucket
  store: s3
  mode: MOUNT
```
is translated to:
```
/output/:
  source: s3://my-new-bucket
  mode: MOUNT
```

Also updated smoke tests to test for writing to an object store. 

Tested (run the relevant ones):

- [ ] `pytest -v tests/test_smoke.py::test_managed_jobs_storage --kubernetes`
- [ ] `pytest -v tests/test_smoke.py::test_managed_jobs_storage --aws`
- [ ] `pytest -v tests/test_smoke.py::test_managed_jobs_storage --gcp`
- [ ] `pytest tests/test_smoke.py::TestStorageWithCredentials` 

Problem statement:
[Storage] Creating new buckets in `sky jobs launch` fails
`sky jobs launch --cloud aws task.yaml`

```
file_mounts:
  /outputs:
    name: romil-sky-outputs1 # This is a new unique bucket
    store: s3
    mode: MOUNT

run: |
  echo "Hello, world!" > /outputs/hello.txt
```

Fails with:

```
$ sky jobs logs --controller 1
...
(sky-4d6a-romilb, pid=30569) sky.exceptions.StorageSpecError: Attempted to mount a non-sky managed bucket 'romil-sky-outputs1' without specifying the storage source. Bucket 'romil-sky-outputs1' already exists.
(sky-4d6a-romilb, pid=30569)     • To create a new bucket, specify a unique name.
(sky-4d6a-romilb, pid=30569)     • To mount an externally created bucket (e.g., created through cloud console or cloud cli), specify the bucket URL in the source field instead of its name. I.e., replace `name: romil-sky-outputs1` with `source: s3://romil-sky-outputs1`.
ERROR: Job 1 failed with return code list: [1]
```

For controllers, we may need to relax this check.

Hints from issue discussion:
Should also update smoke tests to test this case.

Commit message: fix smoke tests

Patch to review:
diff --git a/examples/managed_job_with_storage.yaml b/examples/managed_job_with_storage.yaml
index ecefccd8b3d..f26b409a1cf 100644
--- a/examples/managed_job_with_storage.yaml
+++ b/examples/managed_job_with_storage.yaml
@@ -20,6 +20,12 @@ file_mounts:
     source: ./examples
     persistent: false
     mode: COPY
+
+  /output_path:
+    # Change this to the your own globally unique bucket name.
+    name: sky-output-bucket
+    mode: MOUNT
+
   /imagenet-image:
     source: s3://sky-imagenet-data
   
@@ -55,3 +61,6 @@ run: |
 
   cat ~/tmpfile
   cat ~/a/b/c/tmpfile
+  
+  # Write to a file in the mounted bucket
+  echo "hello world!" > /output_path/output.txt
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 03f644930f4..1b7e7bdae2f 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -2251,7 +2251,7 @@ def check_cluster_available(
 
 # TODO(tian): Refactor to controller_utils. Current blocker: circular import.
 def is_controller_accessible(
-    controller: controller_utils.Controllers,
+    controller: 'controller_utils.Controllers',
     stopped_message: str,
     non_existent_message: Optional[str] = None,
     exit_if_not_accessible: bool = False,
diff --git a/sky/data/storage.py b/sky/data/storage.py
index e43406c3951..5977ad9f218 100644
--- a/sky/data/storage.py
+++ b/sky/data/storage.py
@@ -27,6 +27,7 @@
 from sky.data import storage_utils
 from sky.data.data_utils import Rclone
 from sky.utils import common_utils
+from sky.utils import controller_utils
 from sky.utils import rich_utils
 from sky.utils import schemas
 from sky.utils import ux_utils
@@ -325,6 +326,13 @@ def __deepcopy__(self, memo):
 
     def _validate_existing_bucket(self):
         """Validates the storage fields for existing buckets."""
+        if controller_utils.is_running_on_jobs_controller():
+            # If running on jobs controller, we skip the check because it is
+            # already done when the job is submitted from the client. In that
+            # case, checking again raises an error since the bucket is created
+            # by the client and not the jobs controller, and is thus seen
+            # (incorrectly) by the jobs controller as an external bucket.
+            return
         # Check if 'source' is None, this is only allowed when Storage is in
         # either MOUNT mode or COPY mode with sky-managed storage.
         # Note: In COPY mode, a 'source' being None with non-sky-managed
diff --git a/sky/jobs/dashboard/dashboard.py b/sky/jobs/dashboard/dashboard.py
index 89c97274646..232d1d141aa 100644
--- a/sky/jobs/dashboard/dashboard.py
+++ b/sky/jobs/dashboard/dashboard.py
@@ -7,10 +7,8 @@
 comment).
 """
 import datetime
-import pathlib
 
 import flask
-import yaml
 
 from sky import jobs as managed_jobs
 from sky.utils import common_utils
@@ -19,30 +17,9 @@
 app = flask.Flask(__name__)
 
 
-def _is_running_on_jobs_controller() -> bool:
-    """Am I running on jobs controller?
-
-    Loads ~/.sky/sky_ray.yml and check cluster_name.
-    """
-    if pathlib.Path('~/.sky/sky_ray.yml').expanduser().exists():
-        config = yaml.safe_load(
-            pathlib.Path('~/.sky/sky_ray.yml').expanduser().read_text())
-        cluster_name = config.get('cluster_name', '')
-        candidate_controller_names = (
-            controller_utils.Controllers.JOBS_CONTROLLER.value.
-            candidate_cluster_names)
-        # We use startswith instead of exact match because the cluster name in
-        # the yaml file is cluster_name_on_cloud which may have additional
-        # suffices.
-        return any(
-            cluster_name.startswith(name)
-            for name in candidate_controller_names)
-    return False
-
-
 @app.route('/')
 def home():
-    if not _is_running_on_jobs_controller():
+    if not controller_utils.is_running_on_jobs_controller():
         # Experimental: run on laptop (refresh is very slow).
         all_managed_jobs = managed_jobs.queue(refresh=True, skip_finished=False)
     else:
diff --git a/sky/task.py b/sky/task.py
index 3dd254838f0..952e2288e90 100644
--- a/sky/task.py
+++ b/sky/task.py
@@ -789,7 +789,7 @@ def update_file_mounts(self, file_mounts: Dict[str, str]) -> 'Task':
 
     def set_storage_mounts(
         self,
-        storage_mounts: Optional[Dict[str, storage_lib.Storage]],
+        storage_mounts: Optional[Dict[str, 'storage_lib.Storage']],
     ) -> 'Task':
         """Sets the storage mounts for this task.
 
@@ -855,7 +855,7 @@ def set_storage_mounts(
         return self
 
     def update_storage_mounts(
-            self, storage_mounts: Dict[str, storage_lib.Storage]) -> 'Task':
+            self, storage_mounts: Dict[str, 'storage_lib.Storage']) -> 'Task':
         """Updates the storage mounts for this task.
 
         Different from set_storage_mounts(), this function updates into the
@@ -882,7 +882,7 @@ def update_storage_mounts(
         return self.set_storage_mounts(task_storage_mounts)
 
     def _get_preferred_store(
-            self) -> Tuple[storage_lib.StoreType, Optional[str]]:
+            self) -> Tuple['storage_lib.StoreType', Optional[str]]:
         """Returns the preferred store type and region for this task."""
         # TODO(zhwu, romilb): The optimizer should look at the source and
         #  destination to figure out the right stores to use. For now, we
diff --git a/sky/utils/controller_utils.py b/sky/utils/controller_utils.py
index c1859d52663..3317115ba67 100644
--- a/sky/utils/controller_utils.py
+++ b/sky/utils/controller_utils.py
@@ -4,11 +4,13 @@
 import enum
 import getpass
 import os
+import pathlib
 import tempfile
 import typing
 from typing import Any, Dict, Iterable, List, Optional, Set
 
 import colorama
+import yaml
 
 from sky import check as sky_check
 from sky import clouds
@@ -742,3 +744,24 @@ def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
             store_prefix = store_type.store_prefix()
             storage_obj.source = f'{store_prefix}{storage_obj.name}'
             storage_obj.force_delete = True
+
+
+def is_running_on_jobs_controller() -> bool:
+    """Am I running on jobs controller?
+
+    Loads ~/.sky/sky_ray.yml and check cluster_name.
+    """
+    if pathlib.Path('~/.sky/sky_ray.yml').expanduser().exists():
+        config = yaml.safe_load(
+            pathlib.Path('~/.sky/sky_ray.yml').expanduser().read_text(
+                encoding='utf-8'))
+        cluster_name = config.get('cluster_name', '')
+        candidate_controller_names = (
+            Controllers.JOBS_CONTROLLER.value.candidate_cluster_names)
+        # We use startswith instead of exact match because the cluster name in
+        # the yaml file is cluster_name_on_cloud which may have additional
+        # suffices.
+        return any(
+            cluster_name.startswith(name)
+            for name in candidate_controller_names)
+    return False
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index d70a9fce4cd..175b85744c2 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -2823,7 +2823,9 @@ def test_managed_jobs_storage(generic_cloud: str):
     name = _get_cluster_name()
     yaml_str = pathlib.Path(
         'examples/managed_job_with_storage.yaml').read_text()
-    storage_name = f'sky-test-{int(time.time())}'
+    timestamp = int(time.time())
+    storage_name = f'sky-test-{timestamp}'
+    output_storage_name = f'sky-test-output-{timestamp}'
 
     # Also perform region testing for bucket creation to validate if buckets are
     # created in the correct region and correctly mounted in managed jobs.
@@ -2838,16 +2840,32 @@ def test_managed_jobs_storage(generic_cloud: str):
         region_cmd = TestStorageWithCredentials.cli_region_cmd(
             storage_lib.StoreType.S3, storage_name)
         region_validation_cmd = f'{region_cmd} | grep {region}'
+        s3_check_file_count = TestStorageWithCredentials.cli_count_name_in_bucket(
+            storage_lib.StoreType.S3, output_storage_name, 'output.txt')
+        output_check_cmd = f'{s3_check_file_count} | grep 1'
     elif generic_cloud == 'gcp':
         region = 'us-west2'
         region_flag = f' --region {region}'
         region_cmd = TestStorageWithCredentials.cli_region_cmd(
             storage_lib.StoreType.GCS, storage_name)
         region_validation_cmd = f'{region_cmd} | grep {region}'
+        gcs_check_file_count = TestStorageWithCredentials.cli_count_name_in_bucket(
+            storage_lib.StoreType.GCS, output_storage_name, 'output.txt')
+        output_check_cmd = f'{gcs_check_file_count} | grep 1'
     elif generic_cloud == 'kubernetes':
+        # With Kubernetes, we don't know which object storage provider is used.
+        # Check both S3 and GCS if bucket exists in either.
+        s3_check_file_count = TestStorageWithCredentials.cli_count_name_in_bucket(
+            storage_lib.StoreType.S3, output_storage_name, 'output.txt')
+        s3_output_check_cmd = f'{s3_check_file_count} | grep 1'
+        gcs_check_file_count = TestStorageWithCredentials.cli_count_name_in_bucket(
+            storage_lib.StoreType.GCS, output_storage_name, 'output.txt')
+        gcs_output_check_cmd = f'{gcs_check_file_count} | grep 1'
+        output_check_cmd = f'{s3_output_check_cmd} || {gcs_output_check_cmd}'
         use_spot = ' --no-use-spot'
 
     yaml_str = yaml_str.replace('sky-workdir-zhwu', storage_name)
+    yaml_str = yaml_str.replace('sky-output-bucket', output_storage_name)
     with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w') as f:
         f.write(yaml_str)
         f.flush()
@@ -2860,9 +2878,12 @@ def test_managed_jobs_storage(generic_cloud: str):
                 region_validation_cmd,  # Check if the bucket is created in the correct region
                 'sleep 60',  # Wait the spot queue to be updated
                 f'{_JOB_QUEUE_WAIT}| grep {name} | grep SUCCEEDED',
-                f'[ $(aws s3api list-buckets --query "Buckets[?contains(Name, \'{storage_name}\')].Name" --output text | wc -l) -eq 0 ]'
+                f'[ $(aws s3api list-buckets --query "Buckets[?contains(Name, \'{storage_name}\')].Name" --output text | wc -l) -eq 0 ]',
+                # Check if file was written to the mounted output bucket
+                output_check_cmd
             ],
-            _JOB_CANCEL_WAIT.format(job_name=name),
+            (_JOB_CANCEL_WAIT.format(job_name=name),
+             f'; sky storage delete {output_storage_name} || true'),
             # Increase timeout since sky jobs queue -r can be blocked by other spot tests.
             timeout=20 * 60,
         )


Code review comments:
diff_hunk: @@ -325,6 +326,13 @@ def __deepcopy__(self, memo):
 
     def _validate_existing_bucket(self):
         """Validates the storage fields for existing buckets."""
+        if controller_utils.is_running_on_jobs_controller():
+            # If running on jobs controller, we skip the check because it is
path: sky/data/storage.py
text: @user1:
```suggestion
            # If running on jobs/serve controller, we skip the check because it is
```

diff_hunk: @@ -325,6 +326,13 @@ def __deepcopy__(self, memo):
 
     def _validate_existing_bucket(self):
         """Validates the storage fields for existing buckets."""
+        if controller_utils.is_running_on_jobs_controller():
path: sky/data/storage.py
text: @user1:
Instead of skipping the check on the controller, we may want to translate those storage entries into external ones in `controller_utils.maybe_translate_local_file_mounts_and_sync_up` before submitting to the controller.

Reason: if the storage entry does not contain `store` field, it might be created on one cloud provider locally, but the controller may create another one on a different cloud provider:
```yaml
file_mounts:
  /output_path:
    name: my-bucket
    mode: MOUNT
```