skypilot-nightly 1.0.0.dev20250401__py3-none-any.whl → 1.0.0.dev20250402__py3-none-any.whl

sky/data/mounting_utils.py

@@ -1,10 +1,13 @@
 """Helper functions for object store mounting in Sky Storage"""
+import hashlib
+import os
 import random
 import shlex
 import textwrap
 from typing import Optional
 
 from sky import exceptions
+from sky.skylet import constants
 from sky.utils import command_runner
 
 # Values used to construct mounting commands
@@ -14,11 +17,17 @@ _TYPE_CACHE_TTL = '5s'
 _RENAME_DIR_LIMIT = 10000
 # https://github.com/GoogleCloudPlatform/gcsfuse/releases
 GCSFUSE_VERSION = '2.2.0'
+# Creates a fusermount3 soft link on older (<22) Ubuntu systems to utilize
+# Rclone's mounting utility.
+FUSERMOUNT3_SOFT_LINK_CMD = ('[ ! -f /bin/fusermount3 ] && '
+                             'sudo ln -s /bin/fusermount /bin/fusermount3 || '
+                             'true')
 # https://github.com/Azure/azure-storage-fuse/releases
 BLOBFUSE2_VERSION = '2.2.0'
 _BLOBFUSE_CACHE_ROOT_DIR = '~/.sky/blobfuse2_cache'
 _BLOBFUSE_CACHE_DIR = ('~/.sky/blobfuse2_cache/'
                        '{storage_account_name}_{container_name}')
+# https://github.com/rclone/rclone/releases
 RCLONE_VERSION = 'v1.68.2'
 
 
@@ -112,7 +121,12 @@ def get_az_mount_install_cmd() -> str:
         'sudo apt-get update; '
         'sudo apt-get install -y '
         '-o Dpkg::Options::="--force-confdef" '
-        'fuse3 libfuse3-dev && '
+        'fuse3 libfuse3-dev || { '
+        '  echo "fuse3 not available, falling back to fuse"; '
+        '  sudo apt-get install -y '
+        '  -o Dpkg::Options::="--force-confdef" '
+        '  fuse libfuse-dev; '
+        '} && '
         'ARCH=$(uname -m) && '
         'if [ "$ARCH" = "aarch64" ] || [ "$ARCH" = "arm64" ]; then '
         '  echo "blobfuse2 is not supported on $ARCH" && '
@@ -203,31 +217,17 @@ def get_r2_mount_cmd(r2_credentials_path: str,
     return mount_cmd
 
 
-def get_cos_mount_install_cmd() -> str:
-    """Returns a command to install IBM COS mount utility rclone."""
-    install_cmd = ('rclone version >/dev/null 2>&1 || '
-                   '(curl https://rclone.org/install.sh | '
-                   'sudo bash)')
-    return install_cmd
-
-
-def get_cos_mount_cmd(rclone_config_data: str,
-                      rclone_config_path: str,
-                      bucket_rclone_profile: str,
+def get_cos_mount_cmd(rclone_config: str,
+                      rclone_profile_name: str,
                       bucket_name: str,
                       mount_path: str,
                       _bucket_sub_path: Optional[str] = None) -> str:
     """Returns a command to mount an IBM COS bucket using rclone."""
-    # creates a fusermount soft link on older (<22) Ubuntu systems for
-    # rclone's mount utility.
-    set_fuser3_soft_link = ('[ ! -f /bin/fusermount3 ] && '
-                            'sudo ln -s /bin/fusermount /bin/fusermount3 || '
-                            'true')
     # stores bucket profile in rclone config file at the cluster's nodes.
-    configure_rclone_profile = (f'{set_fuser3_soft_link}; '
-                                'mkdir -p ~/.config/rclone/ && '
-                                f'echo "{rclone_config_data}" >> '
-                                f'{rclone_config_path}')
+    configure_rclone_profile = (f'{FUSERMOUNT3_SOFT_LINK_CMD}; '
+                                f'mkdir -p {constants.RCLONE_CONFIG_DIR} && '
+                                f'echo "{rclone_config}" >> '
+                                f'{constants.RCLONE_CONFIG_PATH}')
     if _bucket_sub_path is None:
         sub_path_arg = f'{bucket_name}/{_bucket_sub_path}'
     else:
@@ -235,11 +235,68 @@ def get_cos_mount_cmd(rclone_config_data: str,
     # --daemon will keep the mounting process running in the background.
     mount_cmd = (f'{configure_rclone_profile} && '
                  'rclone mount '
-                 f'{bucket_rclone_profile}:{sub_path_arg} {mount_path} '
+                 f'{rclone_profile_name}:{sub_path_arg} {mount_path} '
                  '--daemon')
     return mount_cmd
 
 
+def get_mount_cached_cmd(rclone_config: str, rclone_profile_name: str,
+                         bucket_name: str, mount_path: str) -> str:
+    """Returns a command to mount a bucket using rclone with vfs cache."""
+    # stores bucket profile in rclone config file at the remote nodes.
+    configure_rclone_profile = (f'{FUSERMOUNT3_SOFT_LINK_CMD}; '
+                                f'mkdir -p {constants.RCLONE_CONFIG_DIR} && '
+                                f'echo {shlex.quote(rclone_config)} >> '
+                                f'{constants.RCLONE_CONFIG_PATH}')
+    # Assume mount path is unique. We use a hash of mount path as
+    # various filenames related to the mount.
+    # This is because the full path may be longer than
+    # the filename length limit.
+    # The hash is a non-negative integer in string form.
+    hashed_mount_path = hashlib.md5(mount_path.encode()).hexdigest()
+    log_file_path = os.path.join(constants.RCLONE_LOG_DIR,
+                                 f'{hashed_mount_path}.log')
+    create_log_cmd = (f'mkdir -p {constants.RCLONE_LOG_DIR} && '
+                      f'touch {log_file_path}')
+    # when mounting multiple directories with vfs cache mode, it's handled by
+    # rclone to create separate cache directories at ~/.cache/rclone/vfs. It is
+    # not necessary to specify separate cache directories.
+    mount_cmd = (
+        f'{create_log_cmd} && '
+        f'{configure_rclone_profile} && '
+        'rclone mount '
+        f'{rclone_profile_name}:{bucket_name} {mount_path} '
+        # '--daemon' keeps the mounting process running in the background.
+        # fail in 10 seconds if mount cannot complete by then,
+        # which should be plenty of time.
+        '--daemon --daemon-wait 10 '
+        f'--log-file {log_file_path} --log-level INFO '
+        # '--dir-cache-time' sets how long directory listings are cached before
+        # rclone checks the remote storage for changes again. A shorter
+        # interval allows for faster detection of new or updated files on the
+        # remote, but increases the frequency of metadata lookups.
+        '--allow-other --vfs-cache-mode full --dir-cache-time 10s '
+        # '--transfers 1' guarantees the files written at the local mount point
+        # to be uploaded to the backend storage in the order of creation.
+        # '--vfs-cache-poll-interval' specifies the frequency of how often
+        # rclone checks the local mount point for stale objects in cache.
+        # '--vfs-write-back' defines the time to write files on remote storage
+        # after last use of the file in local mountpoint.
+        '--transfers 1 --vfs-cache-poll-interval 10s --vfs-write-back 1s '
+        # Have rclone evict files if the cache size exceeds 10G.
+        # This is to prevent cache from growing too large and
+        # using up all the disk space. Note that files that opened
+        # by a process is not evicted from the cache.
+        '--vfs-cache-max-size 10G '
+        # give each mount its own cache directory
+        f'--cache-dir {constants.RCLONE_CACHE_DIR}/{hashed_mount_path} '
+        # This command produces children processes, which need to be
+        # detached from the current process's terminal. The command doesn't
+        # produce any output, so we aren't dropping any logs.
+        '> /dev/null 2>&1')
+    return mount_cmd
+
+
 def get_rclone_install_cmd() -> str:
     """ RClone installation for both apt-get and rpm.
     This would be common command.

sky/data/storage.py

@@ -30,7 +30,6 @@ from sky.data import data_transfer
 from sky.data import data_utils
 from sky.data import mounting_utils
 from sky.data import storage_utils
-from sky.data.data_utils import Rclone
 from sky.skylet import constants
 from sky.utils import common_utils
 from sky.utils import rich_utils
@@ -266,6 +265,13 @@ class StoreType(enum.Enum):
 class StorageMode(enum.Enum):
     MOUNT = 'MOUNT'
     COPY = 'COPY'
+    MOUNT_CACHED = 'MOUNT_CACHED'
+
+
+MOUNTABLE_STORAGE_MODES = [
+    StorageMode.MOUNT,
+    StorageMode.MOUNT_CACHED,
+]
 
 
 class AbstractStore:
@@ -451,13 +457,27 @@ class AbstractStore:
     def mount_command(self, mount_path: str) -> str:
         """Returns the command to mount the Store to the specified mount_path.
 
-        Includes the setup commands to install mounting tools.
+        This command is used for MOUNT mode. Includes the setup commands to
+        install mounting tools.
 
         Args:
           mount_path: str; Mount path on remote server
         """
         raise NotImplementedError
 
+    def mount_cached_command(self, mount_path: str) -> str:
+        """Returns the command to mount the Store to the specified mount_path.
+
+        This command is used for MOUNT_CACHED mode. Includes the setup commands
+        to install mounting tools.
+
+        Args:
+          mount_path: str; Mount path on remote server
+        """
+        raise exceptions.NotSupportedError(
+            f'{StorageMode.MOUNT_CACHED.value} is '
+            f'not supported for {self.name}.')
+
     def __deepcopy__(self, memo):
         # S3 Client and GCS Client cannot be deep copied, hence the
         # original Store object is returned
@@ -1725,6 +1745,17 @@ class S3Store(AbstractStore):
         return mounting_utils.get_mounting_command(mount_path, install_cmd,
                                                    mount_cmd)
 
+    def mount_cached_command(self, mount_path: str) -> str:
+        install_cmd = mounting_utils.get_rclone_install_cmd()
+        rclone_profile_name = (
+            data_utils.Rclone.RcloneStores.S3.get_profile_name(self.name))
+        rclone_config = data_utils.Rclone.RcloneStores.S3.get_config(
+            rclone_profile_name=rclone_profile_name)
+        mount_cached_cmd = mounting_utils.get_mount_cached_cmd(
+            rclone_config, rclone_profile_name, self.bucket.name, mount_path)
+        return mounting_utils.get_mounting_command(mount_path, install_cmd,
+                                                   mount_cached_cmd)
+
     def _create_s3_bucket(self,
                           bucket_name: str,
                           region=_DEFAULT_REGION) -> StorageHandle:
@@ -2252,6 +2283,17 @@ class GcsStore(AbstractStore):
         return mounting_utils.get_mounting_command(mount_path, install_cmd,
                                                    mount_cmd, version_check_cmd)
 
+    def mount_cached_command(self, mount_path: str) -> str:
+        install_cmd = mounting_utils.get_rclone_install_cmd()
+        rclone_profile_name = (
+            data_utils.Rclone.RcloneStores.GCS.get_profile_name(self.name))
+        rclone_config = data_utils.Rclone.RcloneStores.GCS.get_config(
+            rclone_profile_name=rclone_profile_name)
+        mount_cached_cmd = mounting_utils.get_mount_cached_cmd(
+            rclone_config, rclone_profile_name, self.bucket.name, mount_path)
+        return mounting_utils.get_mounting_command(mount_path, install_cmd,
+                                                   mount_cached_cmd)
+
     def _download_file(self, remote_path: str, local_path: str) -> None:
         """Downloads file from remote to local on GS bucket
 
@@ -3126,6 +3168,19 @@ class AzureBlobStore(AbstractStore):
         return mounting_utils.get_mounting_command(mount_path, install_cmd,
                                                    mount_cmd)
 
+    def mount_cached_command(self, mount_path: str) -> str:
+        install_cmd = mounting_utils.get_rclone_install_cmd()
+        rclone_profile_name = (
+            data_utils.Rclone.RcloneStores.AZURE.get_profile_name(self.name))
+        rclone_config = data_utils.Rclone.RcloneStores.AZURE.get_config(
+            rclone_profile_name=rclone_profile_name,
+            storage_account_name=self.storage_account_name,
+            storage_account_key=self.storage_account_key)
+        mount_cached_cmd = mounting_utils.get_mount_cached_cmd(
+            rclone_config, rclone_profile_name, self.container_name, mount_path)
+        return mounting_utils.get_mounting_command(mount_path, install_cmd,
+                                                   mount_cached_cmd)
+
     def _create_az_bucket(self, container_name: str) -> StorageHandle:
         """Creates AZ Container.
 
@@ -3562,6 +3617,17 @@ class R2Store(AbstractStore):
         return mounting_utils.get_mounting_command(mount_path, install_cmd,
                                                    mount_cmd)
 
+    def mount_cached_command(self, mount_path: str) -> str:
+        install_cmd = mounting_utils.get_rclone_install_cmd()
+        rclone_profile_name = (
+            data_utils.Rclone.RcloneStores.R2.get_profile_name(self.name))
+        rclone_config = data_utils.Rclone.RcloneStores.R2.get_config(
+            rclone_profile_name=rclone_profile_name)
+        mount_cached_cmd = mounting_utils.get_mount_cached_cmd(
+            rclone_config, rclone_profile_name, self.bucket.name, mount_path)
+        return mounting_utils.get_mounting_command(mount_path, install_cmd,
+                                                   mount_cached_cmd)
+
     def _create_r2_bucket(self,
                           bucket_name: str,
                           region='auto') -> StorageHandle:
@@ -3681,11 +3747,10 @@ class IBMCosStore(AbstractStore):
                  _bucket_sub_path: Optional[str] = None):
         self.client: 'storage.Client'
         self.bucket: 'StorageHandle'
+        self.rclone_profile_name = (
+            data_utils.Rclone.RcloneStores.IBM.get_profile_name(self.name))
         super().__init__(name, source, region, is_sky_managed,
                          sync_on_reconstruction, _bucket_sub_path)
-        self.bucket_rclone_profile = \
-          Rclone.generate_rclone_bucket_profile_name(
-            self.name, Rclone.RcloneClouds.IBM)
 
     def _validate(self):
         if self.source is not None and isinstance(self.source, str):
@@ -3897,11 +3962,10 @@ class IBMCosStore(AbstractStore):
             # .git directory is excluded from the sync
             # wrapping src_dir_path with "" to support path with spaces
             src_dir_path = shlex.quote(src_dir_path)
-            sync_command = (
-                'rclone copy --exclude ".git/*" '
-                f'{src_dir_path} '
-                f'{self.bucket_rclone_profile}:{self.name}{sub_path}'
-                f'/{dest_dir_name}')
+            sync_command = ('rclone copy --exclude ".git/*" '
+                            f'{src_dir_path} '
+                            f'{self.rclone_profile_name}:{self.name}{sub_path}'
+                            f'/{dest_dir_name}')
             return sync_command
 
         def get_file_sync_command(base_dir_path, file_names) -> str:
@@ -3927,10 +3991,9 @@ class IBMCosStore(AbstractStore):
                 for file_name in file_names
             ])
             base_dir_path = shlex.quote(base_dir_path)
-            sync_command = (
-                'rclone copy '
-                f'{includes} {base_dir_path} '
-                f'{self.bucket_rclone_profile}:{self.name}{sub_path}')
+            sync_command = ('rclone copy '
+                            f'{includes} {base_dir_path} '
+                            f'{self.rclone_profile_name}:{self.name}{sub_path}')
             return sync_command
 
         # Generate message for upload
@@ -3976,7 +4039,8 @@ class IBMCosStore(AbstractStore):
                 'sky storage delete' or 'sky start'
         """
 
-        bucket_profile_name = Rclone.RcloneClouds.IBM.value + self.name
+        bucket_profile_name = (data_utils.Rclone.RcloneStores.IBM.value +
+                               self.name)
         try:
             bucket_region = data_utils.get_ibm_cos_bucket_region(self.name)
         except exceptions.StorageBucketGetError as e:
@@ -4011,9 +4075,9 @@ class IBMCosStore(AbstractStore):
                     '`rclone lsd <remote>` on relevant remotes returned '
                     'via `rclone listremotes` to debug.')
 
-        Rclone.store_rclone_config(
+        data_utils.Rclone.store_rclone_config(
             self.name,
-            Rclone.RcloneClouds.IBM,
+            data_utils.Rclone.RcloneStores.IBM,
             self.region,  # type: ignore
         )
 
@@ -4053,18 +4117,18 @@ class IBMCosStore(AbstractStore):
           mount_path: str; Path to mount the bucket to.
         """
         # install rclone if not installed.
-        install_cmd = mounting_utils.get_cos_mount_install_cmd()
-        rclone_config_data = Rclone.get_rclone_config(
-            self.bucket.name,
-            Rclone.RcloneClouds.IBM,
-            self.region,  # type: ignore
-        )
-        mount_cmd = mounting_utils.get_cos_mount_cmd(rclone_config_data,
-                                                     Rclone.RCLONE_CONFIG_PATH,
-                                                     self.bucket_rclone_profile,
-                                                     self.bucket.name,
-                                                     mount_path,
-                                                     self._bucket_sub_path)
+        install_cmd = mounting_utils.get_rclone_install_cmd()
+        rclone_config = data_utils.Rclone.RcloneStores.IBM.get_config(
+            rclone_profile_name=self.rclone_profile_name,
+            region=self.region)  # type: ignore
+        mount_cmd = (
+            mounting_utils.get_cos_mount_cmd(
+                rclone_config,
+                self.rclone_profile_name,
+                self.bucket.name,
+                mount_path,
+                self._bucket_sub_path,  # type: ignore
+            ))
         return mounting_utils.get_mounting_command(mount_path, install_cmd,
                                                    mount_cmd)
 
@@ -4128,7 +4192,8 @@ class IBMCosStore(AbstractStore):
         except ibm.ibm_botocore.exceptions.ClientError as e:
             if e.__class__.__name__ == 'NoSuchBucket':
                 logger.debug('bucket already removed')
-        Rclone.delete_rclone_bucket_profile(self.name, Rclone.RcloneClouds.IBM)
+        data_utils.Rclone.delete_rclone_bucket_profile(
+            self.name, data_utils.Rclone.RcloneStores.IBM)
 
 
 class OciStore(AbstractStore):

sky/global_user_state.py

@@ -186,6 +186,11 @@ def get_user(user_id: str) -> models.User:
     return models.User(id=row[0], name=row[1])
 
 
+def get_all_users() -> List[models.User]:
+    rows = _DB.cursor.execute('SELECT id, name FROM users').fetchall()
+    return [models.User(id=row[0], name=row[1]) for row in rows]
+
+
 def add_or_update_cluster(cluster_name: str,
                           cluster_handle: 'backends.ResourceHandle',
                           requested_resources: Optional[Set[Any]],

sky/models.py

@@ -2,7 +2,7 @@
 
 import collections
 import dataclasses
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 
 @dataclasses.dataclass
@@ -12,6 +12,9 @@ class User:
     # Display name of the user
     name: Optional[str] = None
 
+    def to_dict(self) -> Dict[str, Any]:
+        return {'id': self.id, 'name': self.name}
+
 
 RealtimeGpuAvailability = collections.namedtuple(
     'RealtimeGpuAvailability', ['gpu', 'counts', 'capacity', 'available'])

sky/server/requests/payloads.py

@@ -115,8 +115,8 @@ class RequestBody(pydantic.BaseModel):
 
 class CheckBody(RequestBody):
     """The request body for the check endpoint."""
-    clouds: Optional[Tuple[str, ...]]
-    verbose: bool
+    clouds: Optional[Tuple[str, ...]] = None
+    verbose: bool = False
 
 
 class DagRequestBody(RequestBody):
@@ -340,8 +340,8 @@ class JobsQueueBody(RequestBody):
 
 class JobsCancelBody(RequestBody):
     """The request body for the jobs cancel endpoint."""
-    name: Optional[str]
-    job_ids: Optional[List[int]]
+    name: Optional[str] = None
+    job_ids: Optional[List[int]] = None
     all: bool = False
     all_users: bool = False
 

sky/server/server.py

@@ -12,7 +12,7 @@ import pathlib
 import re
 import shutil
 import sys
-from typing import Dict, List, Literal, Optional, Set, Tuple
+from typing import Any, Dict, List, Literal, Optional, Set, Tuple
 import uuid
 import zipfile
 
@@ -675,6 +675,13 @@ async def logs(
     )
 
 
+@app.get('/users')
+async def users() -> List[Dict[str, Any]]:
+    """Gets all users."""
+    user_list = global_user_state.get_all_users()
+    return [user.to_dict() for user in user_list]
+
+
 @app.post('/download_logs')
 async def download_logs(
         request: fastapi.Request,

sky/skylet/constants.py

@@ -51,7 +51,7 @@ SKY_RAY_CMD = (f'{SKY_PYTHON_CMD} $([ -s {SKY_RAY_PATH_FILE} ] && '
                f'cat {SKY_RAY_PATH_FILE} 2> /dev/null || which ray)')
 # Separate env for SkyPilot runtime dependencies.
 SKY_REMOTE_PYTHON_ENV_NAME = 'skypilot-runtime'
-SKY_REMOTE_PYTHON_ENV = f'~/{SKY_REMOTE_PYTHON_ENV_NAME}'
+SKY_REMOTE_PYTHON_ENV: str = f'~/{SKY_REMOTE_PYTHON_ENV_NAME}'
 ACTIVATE_SKY_REMOTE_PYTHON_ENV = f'source {SKY_REMOTE_PYTHON_ENV}/bin/activate'
 # uv is used for venv and pip, much faster than python implementations.
 SKY_UV_INSTALL_DIR = '"$HOME/.local/bin"'
@@ -60,7 +60,7 @@ SKY_UV_CMD = f'UV_SYSTEM_PYTHON=false {SKY_UV_INSTALL_DIR}/uv'
 SKY_UV_INSTALL_CMD = (f'{SKY_UV_CMD} -V >/dev/null 2>&1 || '
                       'curl -LsSf https://astral.sh/uv/install.sh '
                       f'| UV_INSTALL_DIR={SKY_UV_INSTALL_DIR} sh')
-SKY_UV_PIP_CMD = f'VIRTUAL_ENV={SKY_REMOTE_PYTHON_ENV} {SKY_UV_CMD} pip'
+SKY_UV_PIP_CMD: str = (f'VIRTUAL_ENV={SKY_REMOTE_PYTHON_ENV} {SKY_UV_CMD} pip')
 # Deleting the SKY_REMOTE_PYTHON_ENV_NAME from the PATH to deactivate the
 # environment. `deactivate` command does not work when conda is used.
 DEACTIVATE_SKY_REMOTE_PYTHON_ENV = (
@@ -331,6 +331,12 @@ SKYPILOT_NODE_RANK = f'{SKYPILOT_ENV_VAR_PREFIX}NODE_RANK'
 # known after provisioning.
 SKY_SSH_USER_PLACEHOLDER = 'skypilot:ssh_user'
 
+RCLONE_CONFIG_DIR = '~/.config/rclone'
+RCLONE_CONFIG_PATH = f'{RCLONE_CONFIG_DIR}/rclone.conf'
+RCLONE_LOG_DIR = '~/.sky/rclone_log'
+RCLONE_CACHE_DIR = '~/.cache/rclone'
+RCLONE_CACHE_REFRESH_INTERVAL = 10
+
 # The keys that can be overridden in the `~/.sky/config.yaml` file. The
 # overrides are specified in task YAMLs.
 OVERRIDEABLE_CONFIG_KEYS_IN_TASK: List[Tuple[str, ...]] = [

sky/task.py

@@ -1128,7 +1128,7 @@ class Task:
                         assert storage.name is not None, storage
                         # extract region from rclone.conf
                         cos_region = data_utils.Rclone.get_region_from_rclone(
-                            storage.name, data_utils.Rclone.RcloneClouds.IBM)
+                            storage.name, data_utils.Rclone.RcloneStores.IBM)
                         blob_path = f'cos://{cos_region}/{storage.name}'
                     blob_path = storage.get_bucket_sub_path_prefix(blob_path)
                     self.update_file_mounts({mnt_path: blob_path})
@@ -1268,7 +1268,7 @@ class Task:
 
         # Storage mounting
         for _, storage_mount in self.storage_mounts.items():
-            if storage_mount.mode == storage_lib.StorageMode.MOUNT:
+            if storage_mount.mode in storage_lib.MOUNTABLE_STORAGE_MODES:
                 required_features.add(
                     clouds.CloudImplementationFeatures.STORAGE_MOUNTING)
                 break

sky/utils/controller_utils.py

@@ -1025,7 +1025,7 @@ def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
     # it was handled in step 6.
     updated_mount_storages = {}
     for storage_path, storage_obj in task.storage_mounts.items():
-        if (storage_obj.mode == storage_lib.StorageMode.MOUNT and
+        if (storage_obj.mode in storage_lib.MOUNTABLE_STORAGE_MODES and
                 not storage_obj.source):
             # Construct source URL with first store type and storage name
             # E.g., s3://my-storage-name
@@ -1043,7 +1043,7 @@ def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
             new_storage = storage_lib.Storage.from_yaml_config({
                 'source': source,
                 'persistent': storage_obj.persistent,
-                'mode': storage_lib.StorageMode.MOUNT.value,
+                'mode': storage_obj.mode.value,
                 # We enable force delete to allow the controller to delete
                 # the object store in case persistent is set to False.
                 '_force_delete': True