Polish workdir/file_mounts validation and logging. (skypilot-org#495)

* Polish workdir/file_mounts validation and logging.

* Fix cloud URIs being displayed with an extra slash

Previous:
   gs://cloud-tpu-test-datasets/fake_imagenet/train-00001-of-01024/ -> /train-00001-of-01024
Now:
   gs://cloud-tpu-test-datasets/fake_imagenet/train-00001-of-01024 -> /train-00001-of-01024

* Fail early for non-existent local file mount sources.

* !r

Author: concretevitamin

sky/backends/backend.py

@@ -1,13 +1,12 @@
 """Sky backend interface."""
-from typing import Any, Callable, Dict, Optional
+from typing import Dict, Optional
 
 from sky import resources
 from sky import task as task_lib
 
 Task = task_lib.Task
 Resources = resources.Resources
 Path = str
-PostSetupFn = Callable[[str], Any]
 
 
 class Backend:

sky/backends/cloud_vm_ray_backend.py

@@ -13,7 +13,7 @@
 import tempfile
 import textwrap
 import time
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import colorama
 from rich import console as rich_console
@@ -34,12 +34,10 @@
 
 Dag = dag_lib.Dag
 OptimizeTarget = optimizer.OptimizeTarget
+Path = str
 Resources = resources_lib.Resources
 Task = task_lib.Task
 
-Path = str
-PostSetupFn = Callable[[str], Any]
-SKY_DIRSIZE_WARN_THRESHOLD = 100
 SKY_REMOTE_APP_DIR = backend_utils.SKY_REMOTE_APP_DIR
 SKY_REMOTE_WORKDIR = backend_utils.SKY_REMOTE_WORKDIR
 SKY_LOGS_DIRECTORY = job_lib.SKY_LOGS_DIRECTORY
@@ -49,6 +47,8 @@
 logger = sky_logging.init_logger(__name__)
 console = rich_console.Console()
 
+_PATH_SIZE_MEGABYTES_WARN_THRESHOLD = 256
+
 
 def _check_cluster_name_is_valid(cluster_name: str) -> None:
     """Errors out on invalid cluster names not supported by cloud providers.
@@ -104,7 +104,7 @@ def _remove_cluster_from_ssh_config(cluster_ip: str,
 
 
 def _path_size_megabytes(path: str) -> int:
-    # Returns the size of files occupied in path in megabytes
+    """Returns the size of 'path' (directory or file) in megabytes."""
     return int(
         subprocess.check_output(['du', '-sh', '-m',
                                  path]).split()[0].decode('utf-8'))
@@ -1128,22 +1128,24 @@ def sync_workdir(self, handle: ResourceHandle, workdir: Path) -> None:
         ip_list = self._get_node_ips(handle.cluster_yaml, handle.launched_nodes)
         full_workdir = os.path.abspath(os.path.expanduser(workdir))
 
-        # Raise warning if directory is too large
-        if not os.path.exists(full_workdir):
-            logger.error(f'{fore.RED}Workdir {workdir} does not exist.'
-                         f'{style.RESET_ALL}')
-            sys.exit(1)
-        dir_size = _path_size_megabytes(full_workdir)
-        if dir_size >= SKY_DIRSIZE_WARN_THRESHOLD:
-            logger.warning(f'{fore.YELLOW}The size of workdir {workdir} '
-                           f'is {dir_size} MB. Try to keep workdir small, as '
-                           f'large sizes will slowdown rsync.{style.RESET_ALL}')
+        # These asserts have been validated at Task construction time.
+        assert os.path.exists(full_workdir), f'{full_workdir} does not exist'
         if os.path.islink(full_workdir):
             logger.warning(
-                f'{fore.YELLOW}Workdir {workdir} is a symlink. '
+                f'{fore.YELLOW}Workdir {workdir!r} is a symlink. '
                 f'Symlink contents are not uploaded.{style.RESET_ALL}')
         else:
-            workdir = f'{workdir}/'
+            assert os.path.isdir(
+                full_workdir), f'{full_workdir} should be a directory.'
+            workdir = os.path.join(workdir, '')  # Adds trailing / if needed.
+
+        # Raise warning if directory is too large
+        dir_size = _path_size_megabytes(full_workdir)
+        if dir_size >= _PATH_SIZE_MEGABYTES_WARN_THRESHOLD:
+            logger.warning(
+                f'{fore.YELLOW}The size of workdir {workdir!r} '
+                f'is {dir_size} MB. Try to keep workdir small, as '
+                f'large sizes will slow down rsync.{style.RESET_ALL}')
 
         def _sync_workdir_node(ip):
             self._rsync_up(handle,
@@ -1157,8 +1159,11 @@ def _sync_workdir_node(ip):
 
         num_nodes = handle.launched_nodes
         plural = 's' if num_nodes > 1 else ''
-        logger.info(f'{fore.CYAN}Syncing (on {num_nodes} node{plural}): '
-                    f'{style.BRIGHT}workdir ({workdir}){style.RESET_ALL}.')
+        logger.info(
+            f'{fore.CYAN}Syncing workdir (to {num_nodes} node{plural}): '
+            f'{style.BRIGHT}{workdir}{style.RESET_ALL}'
+            f' -> '
+            f'{style.BRIGHT}{SKY_REMOTE_WORKDIR}{style.RESET_ALL}')
         with console.status('[bold cyan]Syncing[/]'):
             backend_utils.run_in_parallel(_sync_workdir_node, ip_list)
 
@@ -1192,9 +1197,13 @@ def sync_to_all_nodes(src: str,
                               dst: str,
                               command: Optional[str] = None,
                               run_rsync: Optional[bool] = False):
-            full_src = os.path.abspath(os.path.expanduser(src))
-            if not os.path.islink(full_src) and not os.path.isfile(full_src):
-                src = f'{src}/'
+            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(src))
+                if not os.path.islink(full_src) and not os.path.isfile(
+                        full_src):
+                    src = os.path.join(src, '')  # Adds trailing / if needed.
 
             def _sync_node(ip):
                 if command is not None:
@@ -1225,28 +1234,27 @@ def _sync_node(ip):
 
             num_nodes = handle.launched_nodes
             plural = 's' if num_nodes > 1 else ''
-            logger.info(f'{fore.CYAN}Syncing (on {num_nodes} node{plural}): '
-                        f'{style.BRIGHT}{src} -> {dst}{style.RESET_ALL}')
+            logger.info(f'{fore.CYAN}Syncing (to {num_nodes} node{plural}): '
+                        f'{style.BRIGHT}{src}{style.RESET_ALL} -> '
+                        f'{style.BRIGHT}{dst}{style.RESET_ALL}')
             with console.status('[bold cyan]Syncing[/]'):
                 backend_utils.run_in_parallel(_sync_node, ip_list)
 
         # Check the files and warn
         for dst, src in mounts.items():
             if not task_lib.is_cloud_store_url(src):
                 full_src = os.path.abspath(os.path.expanduser(src))
-                if not os.path.exists(full_src):
-                    logger.error(f'{fore.RED}Directory "{src}" does not exist.'
-                                 f'{style.RESET_ALL}')
-                    sys.exit(1)
+                # Checked during Task.set_file_mounts().
+                assert os.path.exists(full_src), f'{full_src} does not exist.'
                 src_size = _path_size_megabytes(full_src)
-                if src_size >= SKY_DIRSIZE_WARN_THRESHOLD:
+                if src_size >= _PATH_SIZE_MEGABYTES_WARN_THRESHOLD:
                     logger.warning(
-                        f'{fore.YELLOW}The size of file mount src {src} '
+                        f'{fore.YELLOW}The size of file mount src {src!r} '
                         f'is {src_size} MB. Try to keep src small, as '
                         f'large sizes will slow down rsync.{style.RESET_ALL}')
                 if os.path.islink(full_src):
                     logger.warning(
-                        f'{fore.YELLOW}Source path {src} is a symlink. '
+                        f'{fore.YELLOW}Source path {src!r} is a symlink. '
                         f'Symlink contents are not uploaded.{style.RESET_ALL}')
 
         for dst, src in mounts.items():

sky/backends/local_docker_backend.py

@@ -1,7 +1,7 @@
 """Local docker backend for sky"""
 import subprocess
 import tempfile
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Dict, Optional, Union
 
 import colorama
 
@@ -17,7 +17,6 @@
 Task = task_lib.Task
 Resources = resources.Resources
 Path = str
-PostSetupFn = Callable[[str], Any]
 
 logger = sky_logging.init_logger(__name__)
 

sky/task.py

@@ -132,6 +132,7 @@ def __init__(
         dag.add(self)
 
     def _validate(self):
+        """Checks if the Task fields are valid."""
         if not _is_valid_name(self.name):
             raise ValueError(f'Invalid task name {self.name}. Valid name: '
                              f'{_VALID_NAME_DESCR}')
@@ -169,6 +170,13 @@ def _validate(self):
                              f'a command generator ({CommandGen}). '
                              f'Got {type(self.run)}')
 
+        # Workdir.
+        if self.workdir is not None and not os.path.isdir(self.workdir):
+            # Symlink to a dir is legal (isdir() follows symlinks).
+            raise ValueError(
+                'Workdir must exist and must be a directory (or '
+                f'a symlink to a directory). Found: {self.workdir}')
+
     @staticmethod
     def from_yaml(yaml_path):
         with open(os.path.expanduser(yaml_path), 'r') as f:
@@ -478,6 +486,13 @@ def set_file_mounts(self, file_mounts: Optional[Dict[str, str]]) -> None:
             if is_cloud_store_url(target):
                 raise ValueError(
                     'File mount destination paths cannot be cloud storage')
+            if not is_cloud_store_url(source):
+                if not os.path.exists(
+                        os.path.abspath(os.path.expanduser(source))):
+                    raise ValueError(
+                        f'File mount source {source!r} does not exist locally. '
+                        'To fix: check if it exists, and correct the path.')
+
         self.file_mounts = file_mounts
         return self