Repository: skypilot-org/skypilot
Pull request: [Core] Fix log buffering issue
<!-- Describe the changes in this PR -->

Fixes #2065.

The original selectors implementation introduces buffer. We changed it to threading based implementation and it now fixes the buffering problem. Tried asyncio, but since the log_lib can be used in a thread and asyncio does not work well within a thread, we now use the threading implementation.

TODO:
- [x] Examine if the previous arguments for `process_subprocess_stream` is still useful

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: `bash format.sh`
- [x] Any manual or new tests for this PR (please specify below)
  - [x] `sky launch --cloud gcp --cpus 2+ -c test-log test.yaml` with test.yaml in the issue.
  - [x] `sky launch -c test-log-2 --num-nodes 2 --cloud gcp --cpus 2+ test.yaml`
  - [x]  Progress bar test in #261 `sky launch -c progress-bar test2.yaml`
```
# test.yaml
run: |
  echo hi
  echo hi2
  echo hi3
  sleep 12000
  echo bye
```
```
# test2.yaml
name: minimal

setup: |
  echo "running setup"

run: |
  conda env list
  python3 -u -c "from tqdm import tqdm; import time; [time.sleep(1) for i in tqdm(range(10))]"
  echo bye
```
- [x] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`


Problem statement:
[Core] Logging buffered for run section
When running the following yaml file, the `hi2` and `hi3` will only be printed after the `sleep` is finished. It can cause trouble for spot jobs debugging as the preemption can cause the latest log not showing up.
```
run: |
  echo hi
  echo hi2
  echo hi3
  sleep 12000
  echo bye
```

Hints from issue discussion:
This problem should be more critical than the progress bar

Commit message: add comment

Patch to review:
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index ee98ac2fafb..c22ec0e8b9d 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -252,6 +252,8 @@ def add_prologue(self, job_id: int, is_local: bool = False) -> None:
             # FIXME: This is a hack to make sure that the functions can be found
             # by ray.remote. This should be removed once we have a better way to
             # specify dependencies for ray.
+            inspect.getsource(log_lib._ProcessingArgs),  # pylint: disable=protected-access
+            inspect.getsource(log_lib._handle_io_stream),  # pylint: disable=protected-access
             inspect.getsource(log_lib.process_subprocess_stream),
             inspect.getsource(log_lib.run_with_log),
             inspect.getsource(log_lib.make_task_bash_script),
diff --git a/sky/skylet/log_lib.py b/sky/skylet/log_lib.py
index b8890c0f023..666ae76c6b5 100644
--- a/sky/skylet/log_lib.py
+++ b/sky/skylet/log_lib.py
@@ -2,9 +2,10 @@
 
 This is a remote utility module that provides logging functionality.
 """
+import copy
 import io
+import multiprocessing.pool
 import os
-import selectors
 import subprocess
 import sys
 import time
@@ -26,84 +27,98 @@
 logger = sky_logging.init_logger(__name__)
 
 
-def process_subprocess_stream(
-    proc,
-    log_path: str,
-    stream_logs: bool,
-    start_streaming_at: str = '',
-    end_streaming_at: Optional[str] = None,
-    skip_lines: Optional[List[str]] = None,
-    replace_crlf: bool = False,
-    line_processor: Optional[log_utils.LineProcessor] = None,
-    streaming_prefix: Optional[str] = None,
-) -> Tuple[str, str]:
-    """Redirect the process's filtered stdout/stderr to both stream and file"""
-    if line_processor is None:
-        line_processor = log_utils.LineProcessor()
-
-    sel = selectors.DefaultSelector()
-    out_io = io.TextIOWrapper(proc.stdout,
+class _ProcessingArgs:
+    """Arguments for processing logs."""
+
+    def __init__(self,
+                 log_path: str,
+                 stream_logs: bool,
+                 start_streaming_at: str = '',
+                 end_streaming_at: Optional[str] = None,
+                 skip_lines: Optional[List[str]] = None,
+                 replace_crlf: bool = False,
+                 line_processor: Optional[log_utils.LineProcessor] = None,
+                 streaming_prefix: Optional[str] = None) -> None:
+        self.log_path = log_path
+        self.stream_logs = stream_logs
+        self.start_streaming_at = start_streaming_at
+        self.end_streaming_at = end_streaming_at
+        self.skip_lines = skip_lines
+        self.replace_crlf = replace_crlf
+        self.line_processor = line_processor
+        self.streaming_prefix = streaming_prefix
+
+
+def _handle_io_stream(io_stream, out_stream, args: _ProcessingArgs):
+    """Process the stream of a process."""
+    out_io = io.TextIOWrapper(io_stream,
                               encoding='utf-8',
                               newline='',
-                              errors='replace')
-    sel.register(out_io, selectors.EVENT_READ)
-    if proc.stderr is not None:
-        err_io = io.TextIOWrapper(proc.stderr,
-                                  encoding='utf-8',
-                                  newline='',
-                                  errors='replace')
-        sel.register(err_io, selectors.EVENT_READ)
-
-    stdout = ''
-    stderr = ''
-
-    if streaming_prefix is None:
-        streaming_prefix = ''
+                              errors='replace',
+                              write_through=True)
 
     start_streaming_flag = False
     end_streaming_flag = False
-    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 (end_streaming_at is not None and
-                            end_streaming_at in line):
-                        # Keep executing the loop, only stop streaming.
-                        # E.g., this is used for `sky bench` to hide the
-                        # redundant messages of `sky launch` while
-                        # saving them in log files.
-                        end_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 and
-                            not end_streaming_flag):
-                        out_stream.write(streaming_prefix + line)
-                        out_stream.flush()
-                    if log_path != '/dev/null':
-                        fout.write(line)
-                        fout.flush()
-                    line_processor.process_line(line)
+    streaming_prefix = args.streaming_prefix if args.streaming_prefix else ''
+    line_processor = (log_utils.LineProcessor()
+                      if args.line_processor is None else args.line_processor)
+
+    out = []
+    with open(args.log_path, 'a') as fout:
+        with line_processor:
+            while True:
+                line = out_io.readline()
+                if not line:
+                    break
+                # start_streaming_at logic in processor.process_line(line)
+                if args.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 (args.skip_lines is not None and
+                        any(skip in line for skip in args.skip_lines)):
+                    continue
+                if args.start_streaming_at in line:
+                    start_streaming_flag = True
+                if (args.end_streaming_at is not None and
+                        args.end_streaming_at in line):
+                    # Keep executing the loop, only stop streaming.
+                    # E.g., this is used for `sky bench` to hide the
+                    # redundant messages of `sky launch` while
+                    # saving them in log files.
+                    end_streaming_flag = True
+                if (args.stream_logs and start_streaming_flag and
+                        not end_streaming_flag):
+                    print(streaming_prefix + line,
+                          end='',
+                          file=out_stream,
+                          flush=True)
+                if args.log_path != '/dev/null':
+                    fout.write(line)
+                    fout.flush()
+                line_processor.process_line(line)
+                out.append(line)
+    return ''.join(out)
+
+
+def process_subprocess_stream(proc, args: _ProcessingArgs) -> Tuple[str, str]:
+    """Redirect the process's filtered stdout/stderr to both stream and file"""
+    if proc.stderr is not None:
+        # Asyncio does not work as the output processing can be executed in a
+        # different thread.
+        # selectors is possible to handle the multiplexing of stdout/stderr,
+        # but it introduces buffering making the output not streaming.
+        with multiprocessing.pool.ThreadPool(processes=1) as pool:
+            err_args = copy.copy(args)
+            err_args.line_processor = None
+            stderr_fut = pool.apply_async(_handle_io_stream,
+                                          args=(proc.stderr, sys.stderr,
+                                                err_args))
+            stdout = _handle_io_stream(proc.stdout, sys.stdout, args)
+            stderr = stderr_fut.get()
+    else:
+        stdout = _handle_io_stream(proc.stdout, sys.stdout, args)
+        stderr = ''
     return stdout, stderr
 
 
@@ -222,10 +237,9 @@ def run_with_log(
             # We need this even if the log_path is '/dev/null' to ensure the
             # progress bar is shown.
             # NOTE: Lines are printed only when '\r' or '\n' is found.
-            stdout, stderr = process_subprocess_stream(
-                proc,
-                log_path,
-                stream_logs,
+            args = _ProcessingArgs(
+                log_path=log_path,
+                stream_logs=stream_logs,
                 start_streaming_at=start_streaming_at,
                 end_streaming_at=end_streaming_at,
                 skip_lines=skip_lines,
@@ -234,6 +248,7 @@ def run_with_log(
                 replace_crlf=with_ray,
                 streaming_prefix=streaming_prefix,
             )
+            stdout, stderr = process_subprocess_stream(proc, args)
         proc.wait()
         if require_outputs:
             return proc.returncode, stdout, stderr
@@ -311,14 +326,16 @@ def run_bash_command_with_log(bash_command: str,
         else:
             subprocess_cmd = inner_command
 
-        return run_with_log(subprocess_cmd,
-                            log_path,
-                            ray_job_id=job_lib.make_ray_job_id(
-                                job_id, job_owner),
-                            stream_logs=stream_logs,
-                            with_ray=with_ray,
-                            use_sudo=use_sudo,
-                            shell=True)
+        return run_with_log(
+            subprocess_cmd,
+            log_path,
+            ray_job_id=job_lib.make_ray_job_id(job_id, job_owner),
+            stream_logs=stream_logs,
+            with_ray=with_ray,
+            use_sudo=use_sudo,
+            # Disable input to avoid blocking.
+            stdin=subprocess.DEVNULL,
+            shell=True)
 
 
 def _follow_job_logs(file,
diff --git a/sky/skylet/log_lib.pyi b/sky/skylet/log_lib.pyi
index 6a85ba6fb8f..8c329885721 100644
--- a/sky/skylet/log_lib.pyi
+++ b/sky/skylet/log_lib.pyi
@@ -11,16 +11,33 @@ from typing import Dict, List, Optional, Tuple, Union
 from typing_extensions import Literal
 
 
-def process_subprocess_stream(
-        proc,
-        log_path: str,
-        stream_logs: bool,
-        start_streaming_at: str = ...,
-        end_streaming_at: Optional[str] = ...,
-        skip_lines: Optional[List[str]] = ...,
-        replace_crlf: bool = ...,
-        line_processor: Optional[log_utils.LineProcessor] = ...,
-        streaming_prefix: Optional[str] = ...) -> Tuple[str, str]:
+class _ProcessingArgs:
+    log_path: str
+    stream_logs: bool
+    start_streaming_at: str = ...
+    end_streaming_at: Optional[str] = ...
+    skip_lines: Optional[List[str]] = ...
+    replace_crlf: bool = ...
+    line_processor: Optional[log_utils.LineProcessor] = ...
+    streaming_prefix: Optional[str] = ...
+
+    def __init__(self,
+                 log_path: str,
+                 stream_logs: bool,
+                 start_streaming_at: str = ...,
+                 end_streaming_at: Optional[str] = ...,
+                 skip_lines: Optional[List[str]] = ...,
+                 replace_crlf: bool = ...,
+                 line_processor: Optional[log_utils.LineProcessor] = ...,
+                 streaming_prefix: Optional[str] = ...) -> None:
+        ...
+
+
+def _handle_io_stream(io_stream, out_stream, args: _ProcessingArgs):
+    ...
+
+
+def process_subprocess_stream(proc, args: _ProcessingArgs) -> Tuple[str, str]:
     ...
 
 


Code review comments:
diff_hunk: @@ -26,84 +27,98 @@
 logger = sky_logging.init_logger(__name__)
 
 
-def process_subprocess_stream(
-    proc,
-    log_path: str,
-    stream_logs: bool,
-    start_streaming_at: str = '',
-    end_streaming_at: Optional[str] = None,
-    skip_lines: Optional[List[str]] = None,
-    replace_crlf: bool = False,
-    line_processor: Optional[log_utils.LineProcessor] = None,
-    streaming_prefix: Optional[str] = None,
-) -> Tuple[str, str]:
-    """Redirect the process's filtered stdout/stderr to both stream and file"""
-    if line_processor is None:
-        line_processor = log_utils.LineProcessor()
-
-    sel = selectors.DefaultSelector()
-    out_io = io.TextIOWrapper(proc.stdout,
+class _ProcessingArgs:
+    """Arguments for processing logs."""
+
+    def __init__(self,
+                 log_path: str,
+                 stream_logs: bool,
+                 start_streaming_at: str = '',
+                 end_streaming_at: Optional[str] = None,
+                 skip_lines: Optional[List[str]] = None,
+                 replace_crlf: bool = False,
+                 line_processor: Optional[log_utils.LineProcessor] = None,
+                 streaming_prefix: Optional[str] = None) -> None:
+        self.log_path = log_path
+        self.stream_logs = stream_logs
+        self.start_streaming_at = start_streaming_at
+        self.end_streaming_at = end_streaming_at
+        self.skip_lines = skip_lines
+        self.replace_crlf = replace_crlf
+        self.line_processor = line_processor
+        self.streaming_prefix = streaming_prefix
+
+
+def _handle_io_stream(io_stream, out_stream, args: _ProcessingArgs):
+    """Process the stream of a process."""
+    out_io = io.TextIOWrapper(io_stream,
                               encoding='utf-8',
                               newline='',
-                              errors='replace')
-    sel.register(out_io, selectors.EVENT_READ)
-    if proc.stderr is not None:
-        err_io = io.TextIOWrapper(proc.stderr,
-                                  encoding='utf-8',
-                                  newline='',
-                                  errors='replace')
-        sel.register(err_io, selectors.EVENT_READ)
-
-    stdout = ''
-    stderr = ''
-
-    if streaming_prefix is None:
-        streaming_prefix = ''
+                              errors='replace',
+                              write_through=True)
 
     start_streaming_flag = False
     end_streaming_flag = False
-    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 (end_streaming_at is not None and
-                            end_streaming_at in line):
-                        # Keep executing the loop, only stop streaming.
-                        # E.g., this is used for `sky bench` to hide the
-                        # redundant messages of `sky launch` while
-                        # saving them in log files.
-                        end_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 and
-                            not end_streaming_flag):
-                        out_stream.write(streaming_prefix + line)
-                        out_stream.flush()
-                    if log_path != '/dev/null':
-                        fout.write(line)
-                        fout.flush()
-                    line_processor.process_line(line)
+    streaming_prefix = args.streaming_prefix if args.streaming_prefix else ''
+    line_processor = (log_utils.LineProcessor()
+                      if args.line_processor is None else args.line_processor)
+
+    out = []
+    with open(args.log_path, 'a') as fout:
+        with line_processor:
+            while True:
+                line = out_io.readline()
+                if not line:
+                    break
+                # start_streaming_at logic in processor.process_line(line)
+                if args.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 (args.skip_lines is not None and
+                        any(skip in line for skip in args.skip_lines)):
+                    continue
+                if args.start_streaming_at in line:
+                    start_streaming_flag = True
+                if (args.end_streaming_at is not None and
+                        args.end_streaming_at in line):
+                    # Keep executing the loop, only stop streaming.
+                    # E.g., this is used for `sky bench` to hide the
+                    # redundant messages of `sky launch` while
+                    # saving them in log files.
+                    end_streaming_flag = True
+                if (args.stream_logs and start_streaming_flag and
+                        not end_streaming_flag):
+                    print(streaming_prefix + line,
+                          end='',
+                          file=out_stream,
+                          flush=True)
+                if args.log_path != '/dev/null':
+                    fout.write(line)
+                    fout.flush()
+                line_processor.process_line(line)
+                out.append(line)
+    return ''.join(out)
+
+
+def process_subprocess_stream(proc, args: _ProcessingArgs) -> Tuple[str, str]:
+    """Redirect the process's filtered stdout/stderr to both stream and file"""
+    if proc.stderr is not None:
+        # Asyncio does not work as the output processing can be executed in a
+        # different thread.
+        # selectors is possible to handle the multiplexing of stdout/stderr,
+        # but it introduces buffering making the output not streaming.
+        with multiprocessing.pool.ThreadPool(processes=1) as pool:
+            err_args = copy.copy(args)
+            err_args.line_processor = None
+            stderr_fut = pool.apply_async(_handle_io_stream,
+                                          args=(proc.stderr, sys.stderr,
+                                                err_args))
+            stdout = _handle_io_stream(proc.stdout, sys.stdout, args)
+            stderr = stderr_fut.get()
path: sky/skylet/log_lib.py
text: @user1:
Do you think we also need to do any exception handling here (or after `while True:` at L70) if something goes wrong in _handle_io_stream? Or will this run normally even if there's an exception?

@author:
I could not think of any possible error that can be raised there. Maybe we can just let the error to be thrown as we did before, and wait for any report of exceptions?