Repository: skypilot-org/skypilot
Pull request #1242 (2022-10-13 19:00:17)

Usage lib: cover all programmatic APIs.
Tested:
- manually called all programmatic APIs in a script and observed expected behaviors
- ensured that the cluster names are the same between CLI and programmatic


Patch:
diff --git a/docs/source/examples/syncing-code-artifacts.rst b/docs/source/examples/syncing-code-artifacts.rst
index 46d2acb4d..e76348728 100644
--- a/docs/source/examples/syncing-code-artifacts.rst
+++ b/docs/source/examples/syncing-code-artifacts.rst
@@ -64,6 +64,7 @@ scripts, access checkpoints, etc.).
     $ # cluster1 will get the updated code.
     $ sky exec cluster1 task.yaml
 
+.. _file-mounts-example:
 
 Uploading files outside of workdir
 --------------------------------------
diff --git a/docs/source/reference/storage.rst b/docs/source/reference/storage.rst
index e7808cbae..8ec6992f7 100644
--- a/docs/source/reference/storage.rst
+++ b/docs/source/reference/storage.rst
@@ -4,7 +4,8 @@ SkyPilot Storage
 =================
 
 A SkyPilot Storage object represents an abstract data store containing large data
-files required by the task. Compared to file_mounts, storage is faster and
+files required by the task. Think of it as a bucket of files that can be attached
+to your task. Compared to file_mounts, storage is faster and
 can persist across runs, requiring fewer uploads from your local machine.
 Behind the scenes, storage automatically uploads all data in the source
 to a backing object store in a particular cloud (S3/GCS/Azure Blob).
@@ -175,6 +176,11 @@ and storage mounting:
     For mounts not using SkyPilot Storage (e.g., those using rsync) the symbolic links are directly copied, not their target data.
     The targets must be separately mounted or else the symlinks may break.
 
+.. note::
+    Storage only supports uploading directories (i.e., :code:`source` cannot be a file).
+    To upload a single file to a bucket, please put in a directory and specify the directory as the source.
+    To directly copy a file to a VM, please use regular :ref:`file mounts <file-mounts-example>`.
+
 Creating a shared file system
 -----------------------------
 
diff --git a/examples/storage_demo.yaml b/examples/storage_demo.yaml
index 35fac797e..b6d048064 100644
--- a/examples/storage_demo.yaml
+++ b/examples/storage_demo.yaml
@@ -94,7 +94,7 @@ file_mounts:
   # When the VM is initialized, the contents of the bucket are copied to
   # /datasets-storage. If the bucket already exists, it is fetched and re-used.
   /datasets-storage:
-    name: sky-dataset-romil # Make sure this name is unique or you own this bucket
+    name: sky-dataset-romilzz # Make sure this name is unique or you own this bucket
     source: ~/datasets
     store: s3 # Could be either of [s3, gcs]; default: None
     persistent: True  # Defaults to True, can be set to false.
@@ -110,7 +110,7 @@ file_mounts:
   # other storage mounts using the same bucket with MOUNT mode. Note that the
   # source is synced with the remote bucket everytime this task is run.
   /dataset-storage-mount:
-    name: sky-dataset-romil
+    name: sky-dataset-romilzz
     source: ~/datasets
     mode: MOUNT
 
@@ -137,7 +137,7 @@ file_mounts:
   # this approach can also be used to create a shared filesystem across workers.
   # See examples/storage/pingpong.yaml for an example.
   /outputs-mount:
-    name: romil-output-bucket
+    name: romil-output-bucketzz
     mode: MOUNT
 
 run: |
diff --git a/examples/using_file_mounts.yaml b/examples/using_file_mounts.yaml
index 75f1ab3af..328e0b398 100644
--- a/examples/using_file_mounts.yaml
+++ b/examples/using_file_mounts.yaml
@@ -30,6 +30,7 @@ file_mounts:
 
   ~/tmpfile: ~/tmpfile
   ~/a/b/c/tmpfile: ~/tmpfile
+  ~/tmpfile2: ~/tmp-workdir/tmp file
 
   # Intermediate dirs are created and chown-ed to $USER if they don't already
   # exist.  Thus they can be used too, e.g., 'mkdir -p /data/logs'.  See 'run'.
diff --git a/sky/backends/backend.py b/sky/backends/backend.py
index 4dcff6e5d..272ffb7e3 100644
--- a/sky/backends/backend.py
+++ b/sky/backends/backend.py
@@ -44,7 +44,7 @@ class Backend:
                   retry_until_up: bool = False) -> ResourceHandle:
         if cluster_name is None:
             cluster_name = sky.backends.backend_utils.generate_cluster_name()
-        usage_lib.messages.usage.update_cluster_name(cluster_name)
+        usage_lib.record_cluster_name_for_current_operation(cluster_name)
         usage_lib.messages.usage.update_actual_task(task)
         return self._provision(task, to_provision, dryrun, stream_logs,
                                cluster_name, retry_until_up)
@@ -76,7 +76,8 @@ class Backend:
     @usage_lib.messages.usage.update_runtime('execute')
     def execute(self, handle: ResourceHandle, task: 'task_lib.Task',
                 detach_run: bool) -> None:
-        usage_lib.messages.usage.update_cluster_name(handle.get_cluster_name())
+        usage_lib.record_cluster_name_for_current_operation(
+            handle.get_cluster_name())
         usage_lib.messages.usage.update_actual_task(task)
         return self._execute(handle, task, detach_run)
 
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 51acef19e..235c2ed9e 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -31,6 +31,7 @@ from ray.autoscaler._private import commands as ray_commands
 from ray.autoscaler._private import util as ray_util
 import rich.console as rich_console
 import rich.progress as rich_progress
+import yaml
 
 import sky
 from sky import authentication as auth
@@ -109,12 +110,33 @@ CLUSTER_STATUS_LOCK_TIMEOUT_SECONDS = 20
 # Remote dir that holds our runtime files.
 _REMOTE_RUNTIME_FILES_DIR = '~/.sky/.runtime_files'
 
+# Include the fields that will be used for generating tags that  distinguishes the
+# cluster in ray, to avoid the stopped cluster being discarded due to updates in
+# the yaml template.
+# Some notes on the fields:
+# - 'provider' fields will be used for bootstrapping and insert more new items in
+#   'node_config'.
+# - keeping the auth is not enough becuase the content of the key file will be used
+#   for calculating the hash.
+# TODO(zhwu): Keep in sync with the fields used in https://github.com/ray-project/ray/blob/e4ce38d001dbbe09cd21c497fedd03d692b2be3e/python/ray/autoscaler/_private/commands.py#L687-L701
+_RAY_YAML_KEYS_TO_RESTORE_FOR_BACK_COMPATIBILITY = {
+    'cluster_name', 'provider', 'auth', 'node_config'
+}
+
 
 def is_ip(s: str) -> bool:
     """Returns whether this string matches IP_ADDR_REGEX."""
     return len(re.findall(IP_ADDR_REGEX, s)) == 1
 
 
+def _get_yaml_path_from_cluster_name(cluster_name: str,
+                                     prefix: str = SKY_USER_FILE_PATH) -> str:
+    output_path = pathlib.Path(
+        prefix).expanduser().resolve() / f'{cluster_name}.yml'
+    os.makedirs(output_path.parents[0], exist_ok=True)
+    return str(output_path)
+
+
 def fill_template(template_name: str,
                   variables: Dict,
                   output_path: Optional[str] = None,
@@ -129,10 +151,8 @@ def fill_template(template_name: str,
     if output_path is None:
         assert ('cluster_name' in variables), ('cluster_name is required.')
         cluster_name = variables.get('cluster_name')
-        output_path = pathlib.Path(
-            output_prefix).expanduser() / f'{cluster_name}.yml'
-        os.makedirs(output_path.parents[0], exist_ok=True)
-        output_path = str(output_path)
+        output_path = _get_yaml_path_from_cluster_name(cluster_name,
+                                                       output_prefix)
     output_path = os.path.abspath(output_path)
 
     # Add yaml file path to the template variables.
@@ -222,8 +242,11 @@ def _optimize_file_mounts(yaml_path: str) -> None:
 
     # (For local) Move all runtime files, including the just-written yaml, to
     # local_runtime_files_dir/.
-    all_local_sources = ' '.join(
-        local_src for local_src in file_mounts.values())
+    all_local_sources = ''
+    for local_src in file_mounts.values():
+        full_local_src = str(pathlib.Path(local_src).expanduser())
+        # Add quotes for paths containing spaces.
+        all_local_sources += f'{full_local_src!r} '
     # Takes 10-20 ms on laptop incl. 3 clouds' credentials.
     subprocess.run(f'cp -r {all_local_sources} {local_runtime_files_dir}/',
                    shell=True,
@@ -244,7 +267,7 @@ def path_size_megabytes(path: str) -> int:
         subprocess.check_output(
             f'rsync {command_runner.RSYNC_DISPLAY_OPTION} '
             f'{command_runner.RSYNC_FILTER_OPTION} '
-            f'{git_exclude_filter} --dry-run {path}',
+            f'{git_exclude_filter} --dry-run {path!r}',
             shell=True).splitlines()[-1])
     total_bytes = rsync_output.split(' ')[3].replace(',', '')
     return int(total_bytes) // 10**6
@@ -655,18 +678,46 @@ class SSHConfigHelper(object):
                 break
 
 
+def _replace_yaml_dicts(new_yaml: str, old_yaml: str,
+                        key_names: Set[str]) -> str:
+    """Replaces 'new' with 'old' for all keys in key_names.
+
+    The replacement will be applied recursively and only for the blocks
+    with the key in key_names, and have the same ancestors in both 'new'
+    and 'old' YAML tree.
+    """
+
+    def _restore_block(new_block: Dict[str, Any], old_block: Dict[str, Any]):
+        for key, value in new_block.items():
+            if key in key_names:
+                if key in old_block:
+                    new_block[key] = old_block[key]
+                else:
+                    del new_block[key]
+            elif isinstance(value, dict):
+                if key in old_block:
+                    _restore_block(value, old_block[key])
+
+    new_config = yaml.safe_load(new_yaml)
+    old_config = yaml.safe_load(old_yaml)
+    _restore_block(new_config, old_config)
+    return common_utils.dump_yaml_str(new_config)
+
+
 # TODO: too many things happening here - leaky abstraction. Refactor.
 @timeline.event
-def write_cluster_config(to_provision: 'resources.Resources',
-                         num_nodes: int,
-                         cluster_config_template: str,
-                         cluster_name: str,
-                         local_wheel_path: pathlib.Path,
-                         wheel_hash: str,
-                         region: Optional[clouds.Region] = None,
-                         zones: Optional[List[clouds.Zone]] = None,
-                         auth_config: Optional[Dict[str, str]] = None,
-                         dryrun: bool = False) -> Dict[str, str]:
+def write_cluster_config(
+        to_provision: 'resources.Resources',
+        num_nodes: int,
+        cluster_config_template: str,
+        cluster_name: str,
+        local_wheel_path: pathlib.Path,
+        wheel_hash: str,
+        region: Optional[clouds.Region] = None,
+        zones: Optional[List[clouds.Zone]] = None,
+        auth_config: Optional[Dict[str, str]] = None,
+        dryrun: bool = False,
+        keep_launch_fields_in_existing_config: bool = True) -> Dict[str, str]:
     """Fills in cluster configuration templates and writes them out.
 
     Returns: {provisioner: path to yaml, the provisioning spec}.
@@ -700,6 +751,12 @@ def write_cluster_config(to_provision: 'resources.Resources',
         auth_config = onprem_utils.get_local_auth_config(cluster_name)
     region_name = resources_vars.get('region')
 
+    yaml_path = _get_yaml_path_from_cluster_name(cluster_name)
+    old_yaml_content = None
+    if os.path.exists(yaml_path) and keep_launch_fields_in_existing_config:
+        with open(yaml_path, 'r') as f:
+            old_yaml_content = f.read()
+
     yaml_path = fill_template(
         cluster_config_template,
         dict(
@@ -750,6 +807,16 @@ def write_cluster_config(to_provision: 'resources.Resources',
         # been fully tested yet.
         _optimize_file_mounts(yaml_path)
 
+    # Restore the old yaml content for backward compatibility.
+    if old_yaml_content is not None:
+        with open(yaml_path, 'r') as f:
+            new_yaml_content = f.read()
+        restored_yaml_content = _replace_yaml_dicts(
+            old_yaml_content, new_yaml_content,
+            _RAY_YAML_KEYS_TO_RESTORE_FOR_BACK_COMPATIBILITY)
+        with open(yaml_path, 'w') as f:
+            f.write(restored_yaml_content)
+
     usage_lib.messages.usage.update_ray_yaml(yaml_path)
     # For TPU nodes. TPU VMs do not need TPU_NAME.
     if (resources_vars.get('tpu_type') is not None and
@@ -966,12 +1033,6 @@ def parallel_data_transfer_to_nodes(
     style = colorama.Style
 
     origin_source = source
-    if run_rsync:
-        # Do this for local src paths, not for cloud store URIs
-        # (otherwise we have '<abs path to cwd>/gs://.../object/').
-        full_src = os.path.abspath(os.path.expanduser(origin_source))
-        if not os.path.islink(full_src) and not os.path.isfile(full_src):
-            source = os.path.join(full_src, '')
 
     def _sync_node(runner: 'command_runner.SSHCommandRunner') -> None:
         if cmd is not None:
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 404305119..efbca949d 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -924,7 +924,9 @@ class RetryingVmProvisioner(object):
                 self._wheel_hash,
                 region=region,
                 zones=zones,
-                dryrun=dryrun)
+                dryrun=dryrun,
+                keep_launch_fields_in_existing_config=prev_cluster_status
+                is not None)
             if dryrun:
                 return
             cluster_config_file = config_dict['ray']
@@ -1868,8 +1870,6 @@ class CloudVmRayBackend(backends.Backend):
         else:
             assert os.path.isdir(
                 full_workdir), f'{full_workdir} should be a directory.'
-            # FIXME(zongheng): audit; why not give users control to add '/'?
-            workdir = os.path.join(workdir, '')  # Adds trailing / if needed.
 
         # Raise warning if directory is too large
         dir_size = backend_utils.path_size_megabytes(full_workdir)
diff --git a/sky/cli.py b/sky/cli.py
index 3be2ae374..8fe8076e0 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -1391,7 +1391,7 @@ def logs(
         job_status_str = job_status.value if job_status is not None else 'None'
         click.echo(f'Job {job_id}: {job_status_str}')
         if job_status == job_lib.JobStatus.SUCCEEDED:
-            sys.exit(0)
+            return
         else:
             if job_status is None:
                 id_str = '' if job_id is None else f'{job_id} '
@@ -1874,7 +1874,7 @@ def _terminate_or_stop_clusters(
             # should've been printed by _get_glob_clusters() above.
             continue
         clusters.append(name)
-    usage_lib.messages.usage.update_cluster_name(clusters)
+    usage_lib.record_cluster_name_for_current_operation(clusters)
 
     if not clusters:
         click.echo('\nCluster(s) not found (tip: see `sky status`).')
@@ -2398,7 +2398,7 @@ def admin_deploy(clusterspec_yaml: str):
     if not isinstance(ips, list):
         ips = [ips]
     local_cluster_name = yaml_config['cluster']['name']
-    usage_lib.messages.usage.update_cluster_name(local_cluster_name)
+    usage_lib.record_cluster_name_for_current_operation(local_cluster_name)
     usage_lib.messages.usage.update_cluster_resources(
         len(ips), sky.Resources(sky.Local()))
 
diff --git a/sky/core.py b/sky/core.py
index 87d2a9846..df64637df 100644
--- a/sky/core.py
+++ b/sky/core.py
@@ -15,6 +15,7 @@ from sky.backends import backend_utils
 from sky.backends import onprem_utils
 from sky.skylet import constants
 from sky.skylet import job_lib
+from sky.usage import usage_lib
 from sky.utils import tpu_utils
 from sky.utils import ux_utils
 from sky.utils import subprocess_utils
@@ -25,8 +26,10 @@ logger = sky_logging.init_logger(__name__)
 # = Cluster Management =
 # ======================
 
-
 # pylint: disable=redefined-builtin
+
+
+@usage_lib.entrypoint
 def status(all: bool, refresh: bool) -> List[Dict[str, Any]]:
     """Get the cluster status in dict.
 
@@ -57,6 +60,8 @@ def _start(cluster_name: str,
 
     cluster_status, handle = backend_utils.refresh_cluster_status_handle(
         cluster_name)
+    if handle is None:
+        raise ValueError(f'Cluster {cluster_name!r} does not exist.')
     if cluster_status == global_user_state.ClusterStatus.UP:
         print(f'Cluster {cluster_name!r} is already up.')
         return
@@ -69,6 +74,11 @@ def _start(cluster_name: str,
         raise exceptions.NotSupportedError(
             f'Starting cluster {cluster_name!r} with backend {backend.NAME} '
             'is not supported.')
+
+    # NOTE: if spot_status() calls _start() and hits here, that entrypoint
+    # would have a cluster name (the controller) filled in.
+    usage_lib.record_cluster_name_for_current_operation(cluster_name)
+
     with dag.Dag():
         dummy_task = task.Task().set_resources(handle.launched_resources)
         dummy_task.num_nodes = handle.launched_nodes
@@ -83,6 +93,7 @@ def _start(cluster_name: str,
     return handle
 
 
+@usage_lib.entrypoint
 def start(cluster_name: str,
           idle_minutes_to_autostop: Optional[int] = None,
           retry_until_up: bool = False):
@@ -91,13 +102,15 @@ def start(cluster_name: str,
     Please refer to the sky.cli.start for the document.
 
     Raises:
+        ValueError: cluster does not exist.
         sky.exceptions.NotSupportedError: the cluster is not supported.
     """
     _start(cluster_name, idle_minutes_to_autostop, retry_until_up)
 
 
+@usage_lib.entrypoint
 def stop(cluster_name: str, purge: bool = False):
-    """Stop the cluster
+    """Stop the cluster.
 
     Please refer to the sky.cli.stop for the document.
 
@@ -132,11 +145,13 @@ def stop(cluster_name: str, purge: bool = False):
             'disks will be lost.\n'
             '  To terminate the cluster instead, run: '
             f'{colorama.Style.BRIGHT}sky down {cluster_name}')
+    usage_lib.record_cluster_name_for_current_operation(cluster_name)
     backend.teardown(handle, terminate=False, purge=purge)
 
 
+@usage_lib.entrypoint
 def down(cluster_name: str, purge: bool = False):
-    """Down the cluster
+    """Down the cluster.
 
     Please refer to the sky.cli.down for the document.
 
@@ -152,10 +167,12 @@ def down(cluster_name: str, purge: bool = False):
     if handle is None:
         raise ValueError(f'Cluster {cluster_name!r} does not exist.')
 
+    usage_lib.record_cluster_name_for_current_operation(cluster_name)
     backend = backend_utils.get_backend_from_handle(handle)
     backend.teardown(handle, terminate=True, purge=purge)
 
 
+@usage_lib.entrypoint
 def autostop(cluster_name: str, idle_minutes_to_autostop: int):
     """Set the autostop time of the cluster.
 
@@ -202,6 +219,7 @@ def autostop(cluster_name: str, idle_minutes_to_autostop: int):
                 f'{colorama.Style.RESET_ALL}'
                 '\n  Auto-stop can only be set/unset for '
                 f'{global_user_state.ClusterStatus.UP.value} clusters.')
+    usage_lib.record_cluster_name_for_current_operation(cluster_name)
     backend.set_autostop(handle, idle_minutes_to_autostop)
 
 
@@ -243,6 +261,7 @@ def _check_cluster_available(cluster_name: str,
     return handle
 
 
+@usage_lib.entrypoint
 def queue(cluster_name: str,
           skip_finished: bool = False,
           all_users: bool = False) -> List[dict]:
@@ -291,6 +310,7 @@ def queue(cluster_name: str,
     return jobs
 
 
+@usage_lib.entrypoint
 # pylint: disable=redefined-builtin
 def cancel(cluster_name: str,
            all: bool = False,
@@ -331,6 +351,7 @@ def cancel(cluster_name: str,
     backend.cancel_jobs(handle, job_ids)
 
 
+@usage_lib.entrypoint
 def tail_logs(cluster_name: str, job_id: Optional[str], follow: bool = True):
     """Tail the logs of a job.
 
@@ -352,9 +373,11 @@ def tail_logs(cluster_name: str, job_id: Optional[str], follow: bool = True):
           f'Tailing logs of {job_str} on cluster {cluster_name!r}...'
           f'{colorama.Style.RESET_ALL}')
 
+    usage_lib.record_cluster_name_for_current_operation(cluster_name)
     backend.tail_logs(handle, job_id, follow=follow)
 
 
+@usage_lib.entrypoint
 def download_logs(
         cluster_name: str,
         job_ids: Optional[List[str]],
@@ -374,6 +397,7 @@ def download_logs(
     if job_ids is not None and len(job_ids) == 0:
         return []
 
+    usage_lib.record_cluster_name_for_current_operation(cluster_name)
     print(f'{colorama.Fore.YELLOW}'
           'Syncing down logs to local...'
           f'{colorama.Style.RESET_ALL}')
@@ -381,6 +405,7 @@ def download_logs(
     return local_log_dirs
 
 
+@usage_lib.entrypoint
 def job_status(
         cluster_name: str,
         job_ids: Optional[List[str]],
@@ -407,6 +432,7 @@ def job_status(
           'Getting job status...'
           f'{colorama.Style.RESET_ALL}')
 
+    usage_lib.record_cluster_name_for_current_operation(cluster_name)
     statuses = backend.get_job_status(handle, job_ids, stream_logs=stream_logs)
     return statuses
 
@@ -436,6 +462,7 @@ def _is_spot_controller_up(
     return controller_status, handle
 
 
+@usage_lib.entrypoint
 def spot_status(refresh: bool) -> List[Dict[str, Any]]:
     """Get statuses of managed spot jobs.
 
@@ -503,6 +530,7 @@ def spot_status(refresh: bool) -> List[Dict[str, Any]]:
     return jobs
 
 
+@usage_lib.entrypoint
 # pylint: disable=redefined-builtin
 def spot_cancel(name: Optional[str] = None,
                 job_ids: Optional[Tuple[int]] = None,
@@ -556,6 +584,7 @@ def spot_cancel(name: Optional[str] = None,
                 'Please specify the job ID instead of the job name.')
 
 
+@usage_lib.entrypoint
 def spot_tail_logs(name: Optional[str], job_id: Optional[int], follow: bool):
     """Tail logs of managed spot jobs.
 
@@ -586,6 +615,7 @@ def spot_tail_logs(name: Optional[str], job_id: Optional[int], follow: bool):
 # ======================
 # = Storage Management =
 # ======================
+@usage_lib.entrypoint
 def storage_ls() -> List[Dict[str, Any]]:
     """Get the storages.
 
@@ -606,7 +636,8 @@ def storage_ls() -> List[Dict[str, Any]]:
     return storages
 
 
-def storage_delete(name: Optional[str] = None):
+@usage_lib.entrypoint
+def storage_delete(name: str):
     """Delete a storage.
 
     Raises:
diff --git a/sky/data/storage.py b/sky/data/storage.py
index 72504041b..b3e279559 100644
--- a/sky/data/storage.py
+++ b/sky/data/storage.py
@@ -459,15 +459,26 @@ class Storage(object):
                         '(try "/mydir: /mydir" or "/myfile: /myfile"). '
                         f'Found source={source}')
             # Local path, check if it exists
-            source = os.path.abspath(os.path.expanduser(source))
+            full_src = os.path.abspath(os.path.expanduser(source))
             # Only check if local source exists if it is synced to the bucket
-            if not os.path.exists(source) and sync_on_reconstruction:
+            if not os.path.exists(full_src) and sync_on_reconstruction:
                 with ux_utils.print_exception_no_traceback():
                     raise exceptions.StorageSourceError(
                         'Local source path does not'
                         f' exist: {source}')
+            # Check if source is a file - throw error if it is
+            if os.path.isfile(full_src):
+                with ux_utils.print_exception_no_traceback():
+                    raise exceptions.StorageSourceError(
+                        'Storage source path cannot be a file - only '
+                        'directories are supported as a source. '
+                        'To upload a single file, either\n'
+                        '1) Use regular file mounts by writing '
+                        f'<destination_path>: {source} in the file_mounts '
+                        'section of your YAML, or\n2) Place the file in a '
+                        'directory and specify the directory as the source.')
             # Raise warning if user's path is a symlink
-            elif os.path.islink(source):
+            elif os.path.islink(full_src):
                 logger.warning(f'Source path {source} is a symlink. '
                                'Referenced contents are uploaded, matching '
                                'the default behavior for S3 and GCS syncing.')
diff --git a/sky/execution.py b/sky/execution.py
index 038555bc5..bad89e0f9 100644
--- a/sky/execution.py
+++ b/sky/execution.py
@@ -340,6 +340,7 @@ def exec(  # pylint: disable=redefined-builtin
              detach_run=detach_run)
 
 
+@usage_lib.entrypoint
 def spot_launch(
     dag: sky.Dag,
     name: Optional[str] = None,
diff --git a/sky/skylet/job_lib.py b/sky/skylet/job_lib.py
index aef17542d..ddd8ca957 100644
--- a/sky/skylet/job_lib.py
+++ b/sky/skylet/job_lib.py
@@ -248,12 +248,12 @@ def get_latest_job_id() -> Optional[int]:
         return job_id
 
 
-def get_job_time(job_id: int, is_end: bool) -> Optional[int]:
+def get_job_time_payload(job_id: int, is_end: bool) -> Optional[int]:
     field = 'end_at' if is_end else 'start_at'
     rows = _CURSOR.execute(f'SELECT {field} FROM jobs WHERE job_id=(?)',
                            (job_id,))
     for (timestamp,) in rows:
-        return timestamp
+        return common_utils.encode_payload(timestamp)
 
 
 def _get_records_from_rows(rows) -> List[Dict[str, Any]]:
@@ -632,13 +632,13 @@ class JobLibCodeGen:
         return cls._build(code)
 
     @classmethod
-    def get_job_time(cls,
-                     job_id: Optional[int] = None,
-                     is_end: bool = False) -> str:
+    def get_job_time_payload(cls,
+                             job_id: Optional[int] = None,
+                             is_end: bool = False) -> str:
         code = [
             f'job_id = {job_id} if {job_id} is not None '
             'else job_lib.get_latest_job_id()',
-            f'job_time = job_lib.get_job_time(job_id, {is_end})',
+            f'job_time = job_lib.get_job_time_payload(job_id, {is_end})',
             'print(job_time, flush=True)',
         ]
         return cls._build(code)
diff --git a/sky/spot/spot_utils.py b/sky/spot/spot_utils.py
index 991853b38..1784d986a 100644
--- a/sky/spot/spot_utils.py
+++ b/sky/spot/spot_utils.py
@@ -114,7 +114,8 @@ def update_spot_job_status(job_id: Optional[int] = None):
 def get_job_timestamp(backend: 'backends.CloudVmRayBackend', cluster_name: str,
                       get_end_time: bool) -> float:
     """Get the started/ended time of the job."""
-    code = job_lib.JobLibCodeGen.get_job_time(job_id=None, is_end=get_end_time)
+    code = job_lib.JobLibCodeGen.get_job_time_payload(job_id=None,
+                                                      is_end=get_end_time)
     handle = global_user_state.get_handle_from_cluster_name(cluster_name)
     returncode, stdout, stderr = backend.run_on_head(handle,
                                                      code,
@@ -123,6 +124,7 @@ def get_job_timestamp(backend: 'backends.CloudVmRayBackend', cluster_name: str,
     subprocess_utils.handle_returncode(returncode, code,
                                        'Failed to get job time.',
                                        stdout + stderr)
+    stdout = common_utils.decode_payload(stdout)
     return float(stdout)
 
 
diff --git a/sky/task.py b/sky/task.py
index 00874079b..f4207cded 100644
--- a/sky/task.py
+++ b/sky/task.py
@@ -10,6 +10,7 @@ import yaml
 import sky
 from sky import check
 from sky import clouds
+from sky import exceptions
 from sky import global_user_state
 from sky.backends import backend_utils
 from sky.data import storage as storage_lib
@@ -260,7 +261,13 @@ class Task:
             mount_path = storage[0]
             assert mount_path, \
                 'Storage mount path cannot be empty.'
-            storage_obj = storage_lib.Storage.from_yaml_config(storage[1])
+            try:
+                storage_obj = storage_lib.Storage.from_yaml_config(storage[1])
+            except exceptions.StorageSourceError as e:
+                # Patch the error message to include the mount path, if included
+                e.args = (e.args[0].replace('<destination_path>',
+                                            mount_path),) + e.args[1:]
+                raise e
             task_storage_mounts[mount_path] = storage_obj
         task.set_storage_mounts(task_storage_mounts)
 
diff --git a/sky/templates/aws-ray.yml.j2 b/sky/templates/aws-ray.yml.j2
index da3fac071..5732d61dc 100644
--- a/sky/templates/aws-ray.yml.j2
+++ b/sky/templates/aws-ray.yml.j2
@@ -26,10 +26,15 @@ available_node_types:
     node_config:
       InstanceType: {{instance_type}}
       ImageId: {{image_id}}  # Deep Learning AMI (Ubuntu 18.04); see aws.py.
+      # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ec2.html#EC2.ServiceResource.create_instances
       BlockDeviceMappings:
         - DeviceName: /dev/sda1
           Ebs:
             VolumeSize: {{disk_size}}
+            # use default Iops for gp3
+            VolumeType: gp3
+            Iops: 3000
+            Throughput: 125 
       {% if use_spot %}
       InstanceMarketOptions:
           MarketType: spot
@@ -49,6 +54,9 @@ available_node_types:
         - DeviceName: /dev/sda1
           Ebs:
             VolumeSize: {{disk_size}}
+            VolumeType: gp3
+            Iops: 3000
+            Throughput: 125 
       {% if use_spot %}
       InstanceMarketOptions:
           MarketType: spot
@@ -121,13 +129,13 @@ head_start_ray_commands:
   # Line "which prlimit ..": increase the limit of the number of open files for the raylet process, as the `ulimit` may not take effect at this point, because it requires
   # all the sessions to be reloaded. This is a workaround.
   - ((ps aux | grep "-m sky.skylet.skylet" | grep -q python3) || nohup python3 -m sky.skylet.skylet >> ~/.sky/skylet.log 2>&1 &);
-    ray stop; ray start --disable-usage-stats --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}};
-    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id; done;
+    ray stop; ray start --disable-usage-stats --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} || exit 1;
+    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 
 {%- if num_nodes > 1 %}
 worker_start_ray_commands:
-  - ray stop; ray start --disable-usage-stats --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}};
-    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id; done;
+  - ray stop; ray start --disable-usage-stats --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} || exit 1;
+    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 {%- else %}
 worker_start_ray_commands: []
 {%- endif %}
diff --git a/sky/templates/azure-ray.yml.j2 b/sky/templates/azure-ray.yml.j2
index 0e136f058..88f442a4d 100644
--- a/sky/templates/azure-ray.yml.j2
+++ b/sky/templates/azure-ray.yml.j2
@@ -125,13 +125,13 @@ head_start_ray_commands:
   # Start skylet daemon. (Should not place it in the head_setup_commands, otherwise it will run before skypilot is installed.)
   # NOTE: --disable-usage-stats in `ray start` saves 10 seconds of idle wait.
   - ((ps aux | grep "-m sky.skylet.skylet" | grep -q python3) || nohup python3 -m sky.skylet.skylet >> ~/.sky/skylet.log 2>&1 &);
-    ray stop; ray start --disable-usage-stats --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}};
-    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id; done;
+    ray stop; ray start --disable-usage-stats --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} || exit 1;
+    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 
 {%- if num_nodes > 1 %}
 worker_start_ray_commands:
-  - ray stop; ray start --disable-usage-stats --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}};
-    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id; done;
+  - ray stop; ray start --disable-usage-stats --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} || exit 1;
+    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 {%- else %}
 worker_start_ray_commands: []
 {%- endif %}
diff --git a/sky/templates/gcp-ray.yml.j2 b/sky/templates/gcp-ray.yml.j2
index 1994094a6..892b4986b 100644
--- a/sky/templates/gcp-ray.yml.j2
+++ b/sky/templates/gcp-ray.yml.j2
@@ -164,15 +164,15 @@ head_start_ray_commands:
   # all the sessions to be reloaded. This is a workaround.
   - ((ps aux | grep "-m sky.skylet.skylet" | grep -q python3) || nohup python3 -m sky.skylet.skylet >> ~/.sky/skylet.log 2>&1 &);
     export SKY_NUM_GPUS=0 && which nvidia-smi > /dev/null && SKY_NUM_GPUS=$(nvidia-smi --query-gpu=index,name --format=csv,noheader | wc -l);
-    ray stop; ray start --disable-usage-stats --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --num-gpus=$SKY_NUM_GPUS;
-    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id; done;
+    ray stop; ray start --disable-usage-stats --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --num-gpus=$SKY_NUM_GPUS || exit 1;
+    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 
 # Worker commands are needed for TPU VM Pods
 {%- if num_nodes > 1 or tpu_vm %}
 worker_start_ray_commands:
   - SKY_NUM_GPUS=0 && which nvidia-smi > /dev/null && SKY_NUM_GPUS=$(nvidia-smi --query-gpu=index,name --format=csv,noheader | wc -l);
-    ray stop; ray start --disable-usage-stats --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} --num-gpus=$SKY_NUM_GPUS;
-    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id; done;
+    ray stop; ray start --disable-usage-stats --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} --num-gpus=$SKY_NUM_GPUS || exit 1;
+    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 {%- else %}
 worker_start_ray_commands: []
 {%- endif %}
diff --git a/sky/usage/usage_lib.py b/sky/usage/usage_lib.py
index 493226a17..403e6c984 100644
--- a/sky/usage/usage_lib.py
+++ b/sky/usage/usage_lib.py
@@ -430,3 +430,24 @@ def entrypoint(name_or_fn: str, fallback: bool = False):
     return common_utils.make_decorator(entrypoint_context,
                                        name_or_fn,
                                        fallback=fallback)
+
+
+# Convenience methods below.
+
+
+def record_cluster_name_for_current_operation(
+        cluster_name: Union[List[str], str]) -> None:
+    """Records cluster name(s) for the current operation.
+
+    Usage:
+
+       def op():  # CLI or programmatic API
+
+           ...validate errors...
+
+           usage_lib.record_cluster_name_for_current_operation(
+              <actual clusters being operated on>)
+
+           do_actual_op()
+    """
+    messages.usage.update_cluster_name(cluster_name)
diff --git a/sky/utils/command_runner.py b/sky/utils/command_runner.py
index cd018570e..71652ccc9 100644
--- a/sky/utils/command_runner.py
+++ b/sky/utils/command_runner.py
@@ -304,18 +304,24 @@ class SSHCommandRunner:
             rsync_command.append(
                 RSYNC_EXCLUDE_OPTION.format(str(resolved_source / GIT_EXCLUDE)))
 
+        # rsync doesn't support '~' in a quoted target path. need to expand it.
+        full_source_str = str(resolved_source)
+        if resolved_source.is_dir():
+            full_source_str = os.path.join(full_source_str, '')
+
         ssh_options = ' '.join(
             ssh_options_list(self.ssh_private_key, self.ssh_control_name))
         rsync_command.append(f'-e "ssh {ssh_options}"')
+        # To support spaces in the path, we need to quote source and target.
         if up:
             rsync_command.extend([
-                source,
-                f'{self.ssh_user}@{self.ip}:{target}',
+                f'{full_source_str!r}',
+                f'{self.ssh_user}@{self.ip}:{target!r}',
             ])
         else:
             rsync_command.extend([
-                f'{self.ssh_user}@{self.ip}:{source}',
-                target,
+                f'{self.ssh_user}@{self.ip}:{full_source_str!r}',
+                f'{target!r}',
             ])
         command = ' '.join(rsync_command)
 
diff --git a/sky/utils/common_utils.py b/sky/utils/common_utils.py
index 6e5f77986..93e3b23e2 100644
--- a/sky/utils/common_utils.py
+++ b/sky/utils/common_utils.py
@@ -12,7 +12,7 @@ import socket
 import sys
 import time
 import uuid
-from typing import Dict, List, Union
+from typing import Any
 import yaml
 
 from sky import sky_logging
@@ -223,11 +223,17 @@ def retry(method, max_retries=3, initial_backoff=1):
     return method_with_retries
 
 
-def encode_payload(payload: Union[List, Dict]) -> str:
+def encode_payload(payload: Any) -> str:
     """Encode a payload to make it more robust for parsing.
 
+    The make the message transfer more robust to any additional
+    strings added to the message during transfering.
+
+    An example message that is polluted by the system warning:
+    "LC_ALL: cannot change locale (en_US.UTF-8)\n<sky-payload>hello, world</sky-payload>" # pylint: disable=line-too-long
+
     Args:
-        payload: A dict or list to be encoded.
+        payload: A str, dict or list to be encoded.
 
     Returns:
         A string that is encoded from the payload.
@@ -237,15 +243,18 @@ def encode_payload(payload: Union[List, Dict]) -> str:
     return payload_str
 
 
-def decode_payload(payload_str: str) -> Union[List, Dict]:
+def decode_payload(payload_str: str) -> Any:
     """Decode a payload string.
 
     Args:
         payload_str: A string that is encoded from a payload.
 
     Returns:
-        A dict or list that is decoded from the payload string.
+        A str, dict or list that is decoded from the payload string.
     """
-    payload_str = _PAYLOAD_PATTERN.match(payload_str).group(1)
+    matched = _PAYLOAD_PATTERN.findall(payload_str)
+    if not matched:
+        raise ValueError(f'Invalid payload string: \n{payload_str}')
+    payload_str = matched[0]
     payload = json.loads(payload_str)
     return payload


Test patch:
diff --git a/tests/backward_comaptibility_tests.sh b/tests/backward_comaptibility_tests.sh
index d18b2bc05..72baf0bc5 100644
--- a/tests/backward_comaptibility_tests.sh
+++ b/tests/backward_comaptibility_tests.sh
@@ -1,71 +1,78 @@
 #!/bin/bash
-set -ex
+set -ev
 
-CLUSTER_NAME="test-back-compat"
-. $(conda info --base 2> /dev/null)/etc/profile.d/conda.sh
+source ~/.bashrc 
+CLUSTER_NAME="test-back-compat-$USER"
+source $(conda info --base 2> /dev/null)/etc/profile.d/conda.sh
 
 git clone https://github.com/skypilot-org/skypilot.git sky-master || true
 
 
 # Create environment for compatibility tests
-conda create -n sky-back-compat-master -y --clone sky-dev
-conda create -n sky-back-compat-current -y --clone sky-dev
-
-conda activate sky-back-compat-master
-rm -r  ~/.sky/sky_wheels || true
-pip uninstall skypilot -y
-
+mamba > /dev/null 2>&1 || conda install -n base -c conda-forge mamba -y
+mamba init
+source ~/.bashrc
+mamba env list | grep sky-back-compat-master || mamba create -n sky-back-compat-master -y python=3.9
+
+mamba activate sky-back-compat-master
+mamba install -c conda-forge google-cloud-sdk -y
+rm -r  ~/.sky/wheels || true
 cd sky-master
 git pull origin master
 pip install -e ".[all]"
 cd -
 
-conda activate sky-back-compat-current
-rm -r  ~/.sky/sky_wheels || true
+mamba env list | grep sky-back-compat-current || mamba create -n sky-back-compat-current -y --clone sky-back-compat-master
+mamba activate sky-back-compat-current
+rm -r  ~/.sky/wheels || true
 pip uninstall -y skypilot
 pip install -e ".[all]"
 
 
 # exec + launch
-conda activate sky-back-compat-master
-rm -r  ~/.sky/sky_wheels || true
+mamba activate sky-back-compat-master
+rm -r  ~/.sky/wheels || true
 which sky
-sky launch --cloud gcp -c ${CLUSTER_NAME} examples/minimal.yaml
+sky launch --cloud gcp -y -c ${CLUSTER_NAME} examples/minimal.yaml
 
-conda activate sky-back-compat-current
-rm -r  ~/.sky/sky_wheels || true
+mamba activate sky-back-compat-current
+rm -r  ~/.sky/wheels || true
 sky exec --cloud gcp ${CLUSTER_NAME} examples/minimal.yaml
-sky launch --cloud gcp -c ${CLUSTER_NAME} examples/minimal.yaml
-sky queue
+s=$(sky launch --cloud gcp -d -c ${CLUSTER_NAME} examples/minimal.yaml)
+echo $s
+echo $s | grep "Job ID: 3"
+sky queue ${CLUSTER_NAME}
 
 # sky stop + sky start + sky exec
-conda activate sky-back-compat-master
-rm -r  ~/.sky/sky_wheels || true
+mamba activate sky-back-compat-master
+rm -r  ~/.sky/wheels || true
 sky launch --cloud gcp -y -c ${CLUSTER_NAME}-2 examples/minimal.yaml
-conda activate sky-back-compat-current
-rm -r  ~/.sky/sky_wheels || true
+mamba activate sky-back-compat-current
+rm -r  ~/.sky/wheels || true
 sky stop -y ${CLUSTER_NAME}-2
 sky start -y ${CLUSTER_NAME}-2
-sky exec --cloud gcp ${CLUSTER_NAME}-2 examples/minimal.yaml
+s=$(sky exec --cloud gcp -d ${CLUSTER_NAME}-2 examples/minimal.yaml)
+echo $s
+echo $s | grep "Job ID: 2"
 
 # `sky autostop` + `sky status -r`
-conda activate sky-back-compat-master
-rm -r  ~/.sky/sky_wheels || true
+mamba activate sky-back-compat-master
+rm -r  ~/.sky/wheels || true
 sky launch --cloud gcp -y -c ${CLUSTER_NAME}-3 examples/minimal.yaml
-conda activate sky-back-compat-current
-rm -r  ~/.sky/sky_wheels || true
+mamba activate sky-back-compat-current
+rm -r  ~/.sky/wheels || true
 sky autostop -y -i0 ${CLUSTER_NAME}-3
 sleep 100
 sky status -r | grep ${CLUSTER_NAME}-3 | grep STOPPED
 
 
 # (1 node) sky launch + sky exec + sky queue + sky logs
-conda activate sky-back-compat-master
-rm -r  ~/.sky/sky_wheels || true
+mamba activate sky-back-compat-master
+rm -r  ~/.sky/wheels || true
 sky launch --cloud gcp -y -c ${CLUSTER_NAME}-4 examples/minimal.yaml
 sky stop -y ${CLUSTER_NAME}-4
-conda activate sky-back-compat-current
-rm -r  ~/.sky/sky_wheels || true
+mamba activate sky-back-compat-current
+rm -r  ~/.sky/wheels || true
 sky launch --cloud gcp -y -c ${CLUSTER_NAME}-4 examples/minimal.yaml
 sky queue ${CLUSTER_NAME}-4
 sky logs ${CLUSTER_NAME}-4 1 --status
@@ -74,12 +81,12 @@ sky logs ${CLUSTER_NAME}-4 1
 sky logs ${CLUSTER_NAME}-4 2
 
 # (1 node) sky start + sky exec + sky queue + sky logs
-conda activate sky-back-compat-master
-rm -r  ~/.sky/sky_wheels || true
+mamba activate sky-back-compat-master
+rm -r  ~/.sky/wheels || true
 sky launch --cloud gcp -y -c ${CLUSTER_NAME}-5 examples/minimal.yaml
 sky stop -y ${CLUSTER_NAME}-5
-conda activate sky-back-compat-current
-rm -r  ~/.sky/sky_wheels || true
+mamba activate sky-back-compat-current
+rm -r  ~/.sky/wheels || true
 sky start -y ${CLUSTER_NAME}-5
 sky queue ${CLUSTER_NAME}-5
 sky logs ${CLUSTER_NAME}-5 1 --status
@@ -90,12 +97,12 @@ sky logs ${CLUSTER_NAME}-5 2 --status
 sky logs ${CLUSTER_NAME}-5 2
 
 # (2 nodes) sky launch + sky exec + sky queue + sky logs
-conda activate sky-back-compat-master
-rm -r  ~/.sky/sky_wheels || true
+mamba activate sky-back-compat-master
+rm -r  ~/.sky/wheels || true
 sky launch --cloud gcp -y -c ${CLUSTER_NAME}-6 examples/multi_hostname.yaml
 sky stop -y ${CLUSTER_NAME}-6
-conda activate sky-back-compat-current
-rm -r  ~/.sky/sky_wheels || true
+mamba activate sky-back-compat-current
+rm -r  ~/.sky/wheels || true
 sky start -y ${CLUSTER_NAME}-6
 sky queue ${CLUSTER_NAME}-6
 sky logs ${CLUSTER_NAME}-6 1 --status
@@ -104,3 +111,5 @@ sky exec --cloud gcp ${CLUSTER_NAME}-6 examples/multi_hostname.yaml
 sky queue ${CLUSTER_NAME}-6
 sky logs ${CLUSTER_NAME}-6 2 --status
 sky logs ${CLUSTER_NAME}-6 2
+
+sky down ${CLUSTER_NAME}* -y
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index ba7f1e376..38973e1d2 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -243,6 +243,7 @@ def test_file_mounts():
         [
             'touch ~/tmpfile',
             'mkdir -p ~/tmp-workdir',
+            'touch ~/tmp-workdir/tmp\ file',
             'touch ~/tmp-workdir/foo',
             'ln -f -s ~/tmp-workdir/ ~/tmp-workdir/circle-link',
             f'sky launch -y -c {name} examples/using_file_mounts.yaml',
diff --git a/tests/test_storage.py b/tests/test_storage.py
index b38663594..b579e6d54 100644
--- a/tests/test_storage.py
+++ b/tests/test_storage.py
@@ -1,6 +1,7 @@
 import time
 
 import pytest
+import tempfile
 
 from sky import exceptions
 from sky.data import storage as storage_lib
@@ -20,6 +21,12 @@ class TestStorageSpecLocalSource:
             storage_lib.Storage(name='test', source='/bin/')
         assert 'Storage source paths cannot end with a slash' in str(e)
 
+    def test_source_single_file(self):
+        with pytest.raises(exceptions.StorageSourceError) as e:
+            with tempfile.NamedTemporaryFile() as f:
+                storage_lib.Storage(name='test', source=f.name)
+        assert 'Storage source path cannot be a file' in str(e)
+
 
 class TestStorageSpecValidation:
     """Storage specification validation tests"""