Repository: skypilot-org/skypilot
Pull request: [Storage/File Mount] Check `workdir` and Filemount `src` Size
Solves Issue #391. Checks the size of `workdir` (without recursive tracking) and warns the user if the work directory is too big

<img width="1269" alt="Screen Shot 2022-02-24 at 12 20 09 AM" src="https://user-images.githubusercontent.com/20640975/155485642-e0eab110-a628-4d2d-8eff-8d4f3ad66d1c.png">
.

Problem statement:
[Storage] Large Dataset inside `workdir`
In some cases, users will have the entire dataset e.g. `/tmp/imagenet` in their work directory! This is extremely bad since syncing the work directory will hanging if the dataset is big. 

Potential Fix: Some mechanism to determine the size of the workdir and throw warning if workdir is too big. Sky can also automatically manage big folders and move them onto S3 Store.

Hints from issue discussion:
The first suggested fix seems reasonable for now (we should also explicitly mention in documentation the limitations of the `workdir` in terms of size and to never put large datasets there).

The second suggestion (Sky automatically managing big folders and moving to S3/GCS) is interesting! I can imagine that this will remove the need for explicitly specifying `file_mounts` for some users and then only power users (i.e. Daniel et al) will use the explicit Storage spec/API.
Let's just stick with the first solution for now.... (a cap of 100 MB)

Commit message: Done

Patch to review:
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index fb512b20bd5..e0c284f19cc 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -103,6 +103,12 @@ def _remove_cluster_from_ssh_config(cluster_ip: str,
     backend_utils.SSHConfigHelper.remove_cluster(cluster_ip, auth_config)
 
 
+def _path_size(path: str) -> int:
+    return int(
+        subprocess.check_output(['du', '-sh', '-m',
+                                 path]).split()[0].decode('utf-8'))
+
+
 class RayCodeGen:
     """Code generator of a Ray program that executes a sky.Task.
 
@@ -1104,18 +1110,28 @@ def sync_workdir(self, handle: ResourceHandle, workdir: Path) -> None:
         style = colorama.Style
         ip_list = self._get_node_ips(handle.cluster_yaml, handle.launched_nodes)
         full_workdir = os.path.abspath(os.path.expanduser(workdir))
+
+        # If work directory exceeds 100 MB, raise warning
+        dir_size = _path_size(full_workdir)
+        if dir_size >= 100:
+            logger.warning(
+                f'{fore.YELLOW}The size of workdir {workdir} '
+                f'is {dir_size} MB, exceeding 100 MB. '
+                f'Try to keep workdir small, as large sizes will slow '
+                f'down rsync.{style.RESET_ALL}')
+        if os.path.islink(full_workdir):
+            logger.warning(
+                f'{fore.YELLOW}Workdir {workdir} is a symlink. '
+                f'Symlink contents are not uploaded.{style.RESET_ALL}')
+        else:
+            workdir = f'{workdir}/'
+
         # TODO(zhwu): make this in parallel
         for i, ip in enumerate(ip_list):
             node_name = f'worker{i}' if i > 0 else 'head'
             logger.info(
                 f'{fore.CYAN}Syncing: {style.BRIGHT}workdir ({workdir}) -> '
                 f'{node_name}{style.RESET_ALL}.')
-            if os.path.islink(full_workdir):
-                logger.warning(
-                    f'{fore.RED}Workdir {workdir} is a symlink. '
-                    f'Symlink contents are not uploaded.{style.RESET_ALL}')
-            else:
-                workdir = f'{workdir}/'
             self._rsync_up(handle,
                            ip=ip,
                            source=workdir,
@@ -1150,20 +1166,14 @@ 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}/'
             # TODO(zhwu): make this in parallel
             for i, ip in enumerate(ip_list):
                 node_name = f'worker{i}' if i > 0 else 'head'
                 logger.info(f'{fore.CYAN}Syncing (on {node_name}): '
                             f'{style.BRIGHT}{src} -> {dst}{style.RESET_ALL}')
-
-                full_src = os.path.abspath(os.path.expanduser(src))
-                if os.path.islink(full_src):
-                    logger.warning(
-                        f'{fore.RED}Source path {src} is a symlink. '
-                        f'Symlink contents are not uploaded.{style.RESET_ALL}')
-                elif not os.path.isfile(full_src):
-                    src = f'{src}/'
-
                 if command is not None:
                     backend_utils.run_command_on_ip_via_ssh(
                         ip,
@@ -1209,6 +1219,18 @@ def sync_to_all_nodes(src: str,
                 else:
                     mkdir_for_wrapped_dst = f'mkdir -p {wrapped_dst}'
 
+                src_size = _path_size(full_src)
+                if src_size >= 100:
+                    logger.warning(
+                        f'{fore.YELLOW}The size of file mount src {src} '
+                        f'is {src_size} MB, exceeding 100 MB. '
+                        f'Try to keep src small, as large sizes will slow '
+                        f'down rsync.{style.RESET_ALL}')
+                if os.path.islink(full_src):
+                    logger.warning(
+                        f'{fore.YELLOW}Source path {src} is a symlink. '
+                        f'Symlink contents are not uploaded.{style.RESET_ALL}')
+
                 # TODO(mluo): Fix method so that mkdir and rsync run together
                 sync_to_all_nodes(src=src,
                                   dst=wrapped_dst,


Code review comments:
diff_hunk: @@ -1104,18 +1110,28 @@ def sync_workdir(self, handle: ResourceHandle, workdir: Path) -> None:
         style = colorama.Style
         ip_list = self._get_node_ips(handle.cluster_yaml, handle.launched_nodes)
         full_workdir = os.path.abspath(os.path.expanduser(workdir))
+
+        # If work directory exceeds 100 MB, raise warning
+        dir_size = _path_size(full_workdir)
+        if dir_size >= 100:
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
perhaps create a global var (`SKY_DIRSIZE_WARN_THRESHOLD=100` or some better name) and use that instead of hardcoding `100` here

diff_hunk: @@ -1104,18 +1110,28 @@ def sync_workdir(self, handle: ResourceHandle, workdir: Path) -> None:
         style = colorama.Style
         ip_list = self._get_node_ips(handle.cluster_yaml, handle.launched_nodes)
         full_workdir = os.path.abspath(os.path.expanduser(workdir))
+
+        # If work directory exceeds 100 MB, raise warning
+        dir_size = _path_size(full_workdir)
+        if dir_size >= 100:
+            logger.warning(
+                f'{fore.YELLOW}The size of workdir {workdir} '
+                f'is {dir_size} MB, exceeding 100 MB. '
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
Remove the `, exceeding 100 MB.` part as it's an arbitrary threshold we determine. Same w/ L1226.

diff_hunk: @@ -1104,18 +1110,28 @@ def sync_workdir(self, handle: ResourceHandle, workdir: Path) -> None:
         style = colorama.Style
         ip_list = self._get_node_ips(handle.cluster_yaml, handle.launched_nodes)
         full_workdir = os.path.abspath(os.path.expanduser(workdir))
+
+        # If work directory exceeds 100 MB, raise warning
+        dir_size = _path_size(full_workdir)
+        if dir_size >= 100:
+            logger.warning(
+                f'{fore.YELLOW}The size of workdir {workdir} '
+                f'is {dir_size} MB, exceeding 100 MB. '
+                f'Try to keep workdir small, as large sizes will slow '
path: sky/backends/cloud_vm_ray_backend.py
text: @user2:
this string doesn't need `f` (surprised this passed lint; maybe the upcoming flake8 will catch it @user1)

diff_hunk: @@ -103,6 +103,12 @@ def _remove_cluster_from_ssh_config(cluster_ip: str,
     backend_utils.SSHConfigHelper.remove_cluster(cluster_ip, auth_config)
 
 
+def _path_size(path: str) -> int:
+    return int(
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
Add a docstring saying the returned value is in megabytes(?)

@user2:
+1, even better change the func name to `_path_size_megabytes` so it's clear at call sites.