[Jobs] Allowing to specify intermediate bucket for file upload

sky/data/storage.py

@@ -189,6 +189,29 @@ def get_endpoint_url(cls, store: 'AbstractStore', path: str) -> str:
         return bucket_endpoint_url
 
 
+class StorePrefix(enum.Enum):
+    """Enum for the prefix of different stores."""
+    S3 = 's3://'
+    GCS = 'gs://'
+    AZURE = 'https://'
+    R2 = 'r2://'
+    IBM = 'cos://'
+
+    def to_store_type(self) -> StoreType:
+        if self == StorePrefix.S3:
+            return StoreType.S3
+        elif self == StorePrefix.GCS:
+            return StoreType.GCS
+        elif self == StorePrefix.AZURE:
+            return StoreType.AZURE
+        elif self == StorePrefix.R2:
+            return StoreType.R2
+        elif self == StorePrefix.IBM:
+            return StoreType.IBM
+        else:
+            raise ValueError(f'Unknown store prefix: {self}')
+
+
 class StorageMode(enum.Enum):
     MOUNT = 'MOUNT'
     COPY = 'COPY'

sky/skylet/constants.py

@@ -241,9 +241,7 @@
 # Used for translate local file mounts to cloud storage. Please refer to
 # sky/execution.py::_maybe_translate_local_file_mounts_and_sync_up for
 # more details.
-WORKDIR_BUCKET_NAME = 'skypilot-workdir-{username}-{id}'
-FILE_MOUNTS_BUCKET_NAME = 'skypilot-filemounts-folder-{username}-{id}'
-FILE_MOUNTS_FILE_ONLY_BUCKET_NAME = 'skypilot-filemounts-files-{username}-{id}'
+FILE_MOUNTS_BUCKET_NAME = 'skypilot-filemounts-{username}-{id}'
 FILE_MOUNTS_LOCAL_TMP_DIR = 'skypilot-filemounts-files-{id}'
 FILE_MOUNTS_REMOTE_TMP_DIR = '/tmp/sky-{}-filemounts-files'
 

sky/utils/controller_utils.py

@@ -6,7 +6,7 @@
 import os
 import tempfile
 import typing
-from typing import Any, Dict, Iterable, List, Optional, Set
+from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
 
 import colorama
 
@@ -676,6 +676,21 @@ def replace_skypilot_config_path_in_file_mounts(
                      f'with the real path in file mounts: {file_mounts}')
 
 
+def _get_bucket_name_and_store_type_from_job_config(
+) -> Tuple[Optional[str], Optional[str]]:
+    bucket_wth_prefix = skypilot_config.get_nested(('jobs', 'bucket'), None)
+    if bucket_wth_prefix is None:
+        return None, None
+
+    for prefix in storage_lib.StorePrefix:
+        if bucket_wth_prefix.startswith(prefix.value):
+            bucket_name = bucket_wth_prefix[len(prefix.value):]
+            store = prefix.to_store_type().value
+            return bucket_name, store
+
+    raise ValueError(f'Invalid bucket name with prefix: {bucket_wth_prefix}')
+
+
 def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
                                                   path: str) -> None:
     """Translates local->VM mounts into Storage->VM, then syncs up any Storage.
@@ -717,11 +732,16 @@ def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
             ux_utils.spinner_message(
                 f'Translating {msg} to SkyPilot Storage...'))
 
+    # Get the bucket name for the workdir and file mounts,
+    # we stores all these files in same bucket from config.
+    bucket_name, store = _get_bucket_name_and_store_type_from_job_config()
+    if bucket_name is None:
+        bucket_name = constants.FILE_MOUNTS_BUCKET_NAME.format(
+            username=common_utils.get_cleaned_username(), id=run_id)
+
     # Step 1: Translate the workdir to SkyPilot storage.
     new_storage_mounts = {}
     if task.workdir is not None:
-        bucket_name = constants.WORKDIR_BUCKET_NAME.format(
-            username=common_utils.get_cleaned_username(), id=run_id)
         workdir = task.workdir
         task.workdir = None
         if (constants.SKY_REMOTE_WORKDIR in original_file_mounts or
@@ -736,6 +756,7 @@ def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
                 'source': workdir,
                 'persistent': False,
                 'mode': 'COPY',
+                'store': store,
             })
         # Check of the existence of the workdir in file_mounts is done in
         # the task construction.
@@ -754,15 +775,12 @@ def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
         if os.path.isfile(os.path.abspath(os.path.expanduser(src))):
             copy_mounts_with_file_in_src[dst] = src
             continue
-        bucket_name = constants.FILE_MOUNTS_BUCKET_NAME.format(
-            username=common_utils.get_cleaned_username(),
-            id=f'{run_id}-{i}',
-        )
         new_storage_mounts[dst] = storage_lib.Storage.from_yaml_config({
             'name': bucket_name,
             'source': src,
             'persistent': False,
             'mode': 'COPY',
+            'store': store,
         })
         logger.info(f'  {colorama.Style.DIM}Folder : {src!r} '
                     f'-> storage: {bucket_name!r}.{colorama.Style.RESET_ALL}')
@@ -773,8 +791,6 @@ def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
         tempfile.gettempdir(),
         constants.FILE_MOUNTS_LOCAL_TMP_DIR.format(id=run_id))
     os.makedirs(local_fm_path, exist_ok=True)
-    file_bucket_name = constants.FILE_MOUNTS_FILE_ONLY_BUCKET_NAME.format(
-        username=common_utils.get_cleaned_username(), id=run_id)
     file_mount_remote_tmp_dir = constants.FILE_MOUNTS_REMOTE_TMP_DIR.format(
         path)
     if copy_mounts_with_file_in_src:
@@ -786,10 +802,11 @@ def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
 
         new_storage_mounts[
             file_mount_remote_tmp_dir] = storage_lib.Storage.from_yaml_config({
-                'name': file_bucket_name,
+                'name': bucket_name,
                 'source': local_fm_path,
                 'persistent': False,
                 'mode': 'MOUNT',
+                'store': store,
             })
         if file_mount_remote_tmp_dir in original_storage_mounts:
             with ux_utils.print_exception_no_traceback():
@@ -800,7 +817,7 @@ def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
         sources = list(src_to_file_id.keys())
         sources_str = '\n    '.join(sources)
         logger.info(f'  {colorama.Style.DIM}Files (listed below) '
-                    f' -> storage: {file_bucket_name}:'
+                    f' -> storage: {bucket_name}:'
                     f'\n    {sources_str}{colorama.Style.RESET_ALL}')
     rich_utils.force_update_status(
         ux_utils.spinner_message('Uploading translated local files/folders'))
@@ -849,7 +866,7 @@ def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
         store_type = list(storage_obj.stores.keys())[0]
         store_object = storage_obj.stores[store_type]
         bucket_url = storage_lib.StoreType.get_endpoint_url(
-            store_object, file_bucket_name)
+            store_object, bucket_name)
         for dst, src in copy_mounts_with_file_in_src.items():
             file_id = src_to_file_id[src]
             new_file_mounts[dst] = bucket_url + f'/file-{file_id}'

sky/utils/schemas.py

@@ -3,6 +3,7 @@
 Schemas conform to the JSON Schema specification as defined at
 https://json-schema.org/
 """
+import copy
 import enum
 from typing import Any, Dict, List, Tuple
 
@@ -707,6 +708,12 @@ def get_config_schema():
             },
         }
     }
+    jobs_configs = copy.deepcopy(controller_resources_schema)
+    jobs_configs['properties']['bucket'] = {
+        'type': 'string',
+        'pattern': '^(https|s3|gs|r2|cos)://.+',
+        'required': []
+    }
     cloud_configs = {
         'aws': {
             'type': 'object',
@@ -927,7 +934,7 @@ def get_config_schema():
         'required': [],
         'additionalProperties': False,
         'properties': {
-            'jobs': controller_resources_schema,
+            'jobs': jobs_configs,
             'spot': controller_resources_schema,
             'serve': controller_resources_schema,
             'allowed_clouds': allowed_clouds,