Add line processor abstraction and fix gitignored path size (skypilot-org#615)

Author: gmittal

sky/backends/cloud_vm_ray_backend.py

@@ -32,7 +32,7 @@
 from sky import task as task_lib
 from sky.backends import backend_utils
 from sky.backends import wheel_utils
-from sky.skylet import autostop_lib, job_lib, log_lib
+from sky.skylet import autostop_lib, job_lib, log_lib, log_utils
 
 if typing.TYPE_CHECKING:
     from sky import dag
@@ -56,6 +56,10 @@
 _LOCK_FILENAME = '~/.sky/.{}.lock'
 _FILELOCK_TIMEOUT_SECONDS = 10
 
+_RSYNC_DISPLAY_OPTION = '-Pavz'
+_RSYNC_FILTER_OPTION = '--filter=\'dir-merge,- .gitignore\''
+_RSYNC_EXCLUDE_OPTION = '--exclude-from=.git/info/exclude'
+
 
 def _check_cluster_name_is_valid(cluster_name: str) -> None:
     """Errors out on invalid cluster names not supported by cloud providers.
@@ -100,43 +104,18 @@ def _get_task_demands_dict(
     return accelerator_dict
 
 
-def _path_size_megabytes(path: str, exclude_gitignore: bool = False) -> int:
-    """Returns the size of 'path' (directory or file) in megabytes.
-
-    Args:
-        path: The path to check.
-        exclude_gitignore: If True, excludes files matched in .gitignore.
-
-    Returns:
-        The size of 'path' in megabytes.
-    """
-    if exclude_gitignore:
-        try:
-            # FIXME: add git index size (du -hsc .git) in this computation.
-            awk_program = '{ sum += $1 } END { print sum }'
-            return int(
-                subprocess.check_output(
-                    f'( git status --short {path} | '
-                    'grep "^?" | cut -d " " -f2- '
-                    f'&& git ls-files {path} ) | '
-                    'xargs -n 1000 du -hsk | '
-                    f'awk {awk_program!r}',
-                    shell=True,
-                    stderr=subprocess.DEVNULL)) // (2**10)
-        except (subprocess.CalledProcessError, ValueError):
-            # If git is not installed, or if the user is not in a git repo.
-            # Fall back to du -shk if it is not a git repo (size does not
-            # consider .gitignore).
-            logger.debug('Failed to get size with .gitignore exclusion, '
-                         'falling back to du -shk')
-            pass
-    return int(
-        subprocess.check_output([
-            'du',
-            '-sh',
-            '-k',
-            path,
-        ]).split()[0].decode('utf-8')) // (2**10)
+def _path_size_megabytes(path: str) -> int:
+    """Returns the size of 'path' (directory or file) in megabytes."""
+    git_exclude_filter = ''
+    if (pathlib.Path(path) / '.git/info/exclude').exists():
+        git_exclude_filter = f' {_RSYNC_EXCLUDE_OPTION}'
+    rsync_output = str(
+        subprocess.check_output(
+            f'rsync {_RSYNC_DISPLAY_OPTION} {_RSYNC_FILTER_OPTION}'
+            f'{git_exclude_filter} --dry-run {path}',
+            shell=True).splitlines()[-1])
+    total_bytes = rsync_output.split(' ')[3].replace(',', '')
+    return int(total_bytes) // 10**6
 
 
 class RayCodeGen:
@@ -201,7 +180,7 @@ def add_prologue(self, job_id: int) -> None:
             import ray
             import ray.util as ray_util
 
-            from sky.skylet import job_lib
+            from sky.skylet import job_lib, log_utils
 
             SKY_REMOTE_WORKDIR = {log_lib.SKY_REMOTE_WORKDIR!r}
             job_lib.set_status({job_id!r}, job_lib.JobStatus.PENDING)
@@ -946,7 +925,7 @@ def ray_up(start_streaming_at):
                 log_abs_path,
                 stream_logs=False,
                 start_streaming_at=start_streaming_at,
-                parse_ray_up_logs=True,
+                line_processor=log_utils.RayUpLineProcessor(),
                 # Reduce BOTO_MAX_RETRIES from 12 to 5 to avoid long hanging
                 # time during 'ray up' if insufficient capacity occurs.
                 env=dict(os.environ, BOTO_MAX_RETRIES='5'),
@@ -1397,15 +1376,13 @@ def sync_workdir(self, handle: ResourceHandle, workdir: Path) -> None:
             workdir = os.path.join(workdir, '')  # Adds trailing / if needed.
 
         # Raise warning if directory is too large
-        dir_size = _path_size_megabytes(full_workdir, exclude_gitignore=True)
+        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 or use '
-                '.gitignore to exclude large files, as '
-                'large sizes will slow down rsync. If you use .gitignore but '
-                'the path is not initialized in git, you can ignore this '
-                f'warning.{style.RESET_ALL}')
+                '.gitignore to exclude large files, as large sizes will slow '
+                'down rsync. {style.RESET_ALL}')
 
         log_path = os.path.join(self.log_dir, 'workdir_sync.log')
 
@@ -1512,16 +1489,13 @@ def _sync_node(ip):
                 full_src = os.path.abspath(os.path.expanduser(src))
                 # 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,
-                                                exclude_gitignore=True)
+                src_size = _path_size_megabytes(full_src)
                 if src_size >= _PATH_SIZE_MEGABYTES_WARN_THRESHOLD:
                     logger.warning(
                         f'{fore.YELLOW}The size of file mount src {src!r} '
                         f'is {src_size} MB. Try to keep src small or use '
-                        '.gitignore to exclude large files, as '
-                        'large sizes will slow down rsync. If you use '
-                        '.gitignore but the path is not initialized in git, you'
-                        f' can ignore this warning.{style.RESET_ALL}')
+                        '.gitignore to exclude large files, as large sizes '
+                        f'will slow down rsync. {style.RESET_ALL}')
                 if os.path.islink(full_src):
                     logger.warning(
                         f'{fore.YELLOW}Source path {src!r} is a symlink. '
@@ -2171,19 +2145,19 @@ def _rsync_up(
         # shooting a lot of messages to the output. --info=progress2 is used
         # to get a total progress bar, but it requires rsync>=3.1.0 and Mac
         # OS has a default rsync==2.6.9 (16 years old).
-        rsync_command = ['rsync', '-Pavz']
+        rsync_command = ['rsync', _RSYNC_DISPLAY_OPTION]
         # Legend
         #   dir-merge: ignore file can appear in any subdir, applies to that
         #     subdir downwards
         # Note that "-" is mandatory for rsync and means all patterns in the
         # ignore files are treated as *exclude* patterns.  Non-exclude
         # patterns, e.g., "!  do_not_exclude" doesn't work, even though git
         # allows it.
-        rsync_command.append('--filter=\'dir-merge,- .gitignore\'')
+        rsync_command.append(_RSYNC_FILTER_OPTION)
         git_exclude = '.git/info/exclude'
         if (pathlib.Path(source) / git_exclude).exists():
             # Ensure file exists; otherwise, rsync will error out.
-            rsync_command.append('--exclude-from=.git/info/exclude')
+            rsync_command.append(_RSYNC_EXCLUDE_OPTION)
 
         ssh_options = ' '.join(
             backend_utils.ssh_options_list(ssh_key,

sky/skylet/log_lib.py

@@ -1,8 +1,7 @@
-"""Sky logging utils.
+"""Sky logging library.
 
 This is a remote utility module that provides logging functionality.
 """
-import enum
 import io
 import os
 import selectors
@@ -14,47 +13,32 @@
 from typing import Iterator, List, Optional, Tuple, Union
 
 import colorama
-import rich.status
 
-from sky.skylet import job_lib
+from sky.skylet import job_lib, log_utils
 from sky import sky_logging
 
 SKY_REMOTE_WORKDIR = '~/sky_workdir'
 
 logger = sky_logging.init_logger(__name__)
 
 
-class ProvisionStatus(enum.Enum):
-    LAUNCH = 0
-    RUNTIME_SETUP = 1
-
-
-def _update_ray_up_status(log_line, ray_up_state: str,
-                          status_display: rich.status.Status) -> None:
-    if 'Shared connection to' in log_line and ray_up_state[
-            'state'] == ProvisionStatus.LAUNCH:
-        status_display.stop()
-        logger.info(f'{colorama.Fore.GREEN}Head node is up.'
-                    f'{colorama.Style.RESET_ALL}')
-        status_display.start()
-        status_display.update('[bold cyan] Preparing Sky runtime')
-        ray_up_state['state'] = ProvisionStatus.RUNTIME_SETUP
-
-
-def redirect_process_output(proc,
-                            log_path: str,
-                            stream_logs: bool,
-                            start_streaming_at: str = '',
-                            skip_lines: Optional[List[str]] = None,
-                            parse_ray_up_logs: bool = False,
-                            replace_crlf: bool = False) -> Tuple[str, str]:
+def redirect_process_output(
+    proc,
+    log_path: str,
+    stream_logs: bool,
+    start_streaming_at: str = '',
+    skip_lines: Optional[List[str]] = None,
+    replace_crlf: bool = False,
+    line_processor: Optional[log_utils.LineProcessor] = None
+) -> Tuple[str, str]:
     """Redirect the process's filtered stdout/stderr to both stream and file"""
-    # FIXME(gmittal): Remove hard-coded `parse_ray_up_logs` flag and add general
-    # LineParser: https://github.com/sky-proj/sky/pull/565#discussion_r826615923
     log_path = os.path.expanduser(log_path)
     dirname = os.path.dirname(log_path)
     os.makedirs(dirname, exist_ok=True)
 
+    if line_processor is None:
+        line_processor = log_utils.LineProcessor()
+
     sel = selectors.DefaultSelector()
     out_io = io.TextIOWrapper(proc.stdout,
                               encoding='utf-8',
@@ -72,44 +56,40 @@ def redirect_process_output(proc,
     stderr = ''
 
     start_streaming_flag = False
-    if parse_ray_up_logs:
-        ray_up_state = {'state': ProvisionStatus.LAUNCH}
-        provision_status = rich.status.Status('[bold cyan]Launching')
-        provision_status.start()
-    with open(log_path, 'a') as fout:
-        while len(sel.get_map()) > 0:
-            events = sel.select()
-            for key, _ in events:
-                line = key.fileobj.readline()
-                if not line:
-                    # Unregister the io when EOF reached
-                    sel.unregister(key.fileobj)
-                    continue
-                if replace_crlf and line.endswith('\r\n'):
-                    # Replace CRLF with LF to avoid ray logging to the same line
-                    # due to separating lines with '\n'.
-                    line = line[:-2] + '\n'
-                if (skip_lines is not None and
-                        any(skip in line for skip in skip_lines)):
-                    continue
-                if start_streaming_at in line:
-                    start_streaming_flag = True
-                if key.fileobj is out_io:
-                    stdout += line
-                    out_stream = sys.stdout
-                else:
-                    stderr += line
-                    out_stream = sys.stderr
-                if stream_logs and start_streaming_flag:
-                    out_stream.write(line)
-                    out_stream.flush()
-                if log_path != '/dev/null':
-                    fout.write(line)
-                    fout.flush()
-                if parse_ray_up_logs:
-                    _update_ray_up_status(line, ray_up_state, provision_status)
-    if parse_ray_up_logs:
-        provision_status.stop()
+    with line_processor:
+        with open(log_path, 'a') as fout:
+            while len(sel.get_map()) > 0:
+                events = sel.select()
+                for key, _ in events:
+                    line = key.fileobj.readline()
+                    if not line:
+                        # Unregister the io when EOF reached
+                        sel.unregister(key.fileobj)
+                        continue
+                    # TODO(zhwu,gmittal): Put replace_crlf, skip_lines, and
+                    # start_streaming_at logic in processor.process_line(line)
+                    if replace_crlf and line.endswith('\r\n'):
+                        # Replace CRLF with LF to avoid ray logging to the same
+                        # line due to separating lines with '\n'.
+                        line = line[:-2] + '\n'
+                    if (skip_lines is not None and
+                            any(skip in line for skip in skip_lines)):
+                        continue
+                    if start_streaming_at in line:
+                        start_streaming_flag = True
+                    if key.fileobj is out_io:
+                        stdout += line
+                        out_stream = sys.stdout
+                    else:
+                        stderr += line
+                        out_stream = sys.stderr
+                    if stream_logs and start_streaming_flag:
+                        out_stream.write(line)
+                        out_stream.flush()
+                    if log_path != '/dev/null':
+                        fout.write(line)
+                        fout.flush()
+                    line_processor.process_line(line)
     return stdout, stderr
 
 
@@ -122,7 +102,7 @@ def run_with_log(
     shell: bool = False,
     with_ray: bool = False,
     redirect_stdout_stderr: bool = True,
-    parse_ray_up_logs: bool = False,
+    line_processor: Optional[log_utils.LineProcessor] = None,
     **kwargs,
 ) -> Union[int, Tuple[int, str, str]]:
     """Runs a command and logs its output to a file.
@@ -181,7 +161,7 @@ def run_with_log(
                     'bash: cannot set terminal process group',
                     'bash: no job control in this shell',
                 ],
-                parse_ray_up_logs=parse_ray_up_logs,
+                line_processor=line_processor,
                 # Replace CRLF when the output is logged to driver by ray.
                 replace_crlf=with_ray,
             )

sky/skylet/log_utils.py

@@ -0,0 +1,51 @@
+"""Sky logging utils."""
+import enum
+
+import colorama
+import rich.status
+
+from sky import sky_logging
+
+logger = sky_logging.init_logger(__name__)
+
+
+class LineProcessor(object):
+    """A processor for log lines."""
+
+    def __enter__(self):
+        pass
+
+    def process_line(self, log_line):
+        pass
+
+    def __exit__(self, except_type, except_value, traceback):
+        del except_type, except_value, traceback  # unused
+        pass
+
+
+class RayUpLineProcessor(LineProcessor):
+    """A processor for `ray up` log lines."""
+
+    class ProvisionStatus(enum.Enum):
+        LAUNCH = 0
+        RUNTIME_SETUP = 1
+
+    def __enter__(self):
+        self.state = self.ProvisionStatus.LAUNCH
+        self.status_display = rich.status.Status('[bold cyan]Launching')
+        self.status_display.start()
+
+    def process_line(self, log_line):
+        if ('Shared connection to' in log_line and
+                self.state == self.ProvisionStatus.LAUNCH):
+            self.status_display.stop()
+            logger.info(f'{colorama.Fore.GREEN}Head node is up.'
+                        f'{colorama.Style.RESET_ALL}')
+            self.status_display.start()
+            self.status_display.update(
+                '[bold cyan]Launching - Preparing Sky runtime')
+            self.state = self.ProvisionStatus.RUNTIME_SETUP
+
+    def __exit__(self, except_type, except_value, traceback):
+        del except_type, except_value, traceback  # unused
+        self.status_display.stop()