Repository: skypilot-org/skypilot
Pull request: Fix delayed logs of `sky logs` locally vs `tail -f` remotely
Closes #451 

The previous delay can be as large as 5 minutes. I did the following changes to optimize the situation.
1. ~~Using the `ray job logs` instead of our self-implemented busy for loop for tailing the log file. (`ray job` API has a shorter calling path)~~ I have to bring back our own log tailing function, since the `ray job logs` does not handle progress bar correctly.
2. Change the SshMode for log tailing to `Interactive`, so that the `ssh` buffer is disabled.
3. Remove the log redirecting and parsing for the `subprocess` for the log tailing.

The first two modifications reduce the logging delay from 5 min to around 6 seconds. The last modification reduces the logging delay to within 1 second by eyes. After changing back the log tailing function to our own, the time delay is still within 1 second.

Tested:
`sky launch -c test-log examples/resnet_app.yaml` (changed the `train_steps` to 25000)
- [x]  Compare `sky logs test-log 1` vs `ssh test-log; tail -f ./sky_logs/.../run.log`
- [x]  Compare `sky exec test-log "tail -f ./sky_logs/.../run.log"` vs `ssh test-log; tail -f ./sky_logs/.../run.log`

Problem statement:
`sky logs` has large delay compares to the real logs
This is an observation from @dhroth where the `sky logs` has a large delay of the actual log written in the `sky_logs/` folder. By experiments with the `resnet_app.yaml` with the `train_steps=25000`. It has up to 5-minute delay of the logs between `sky logs ` and the real logs. It seems like the problem is caused by the buffering or other text holding reason caused by `ssh`.

Hints from issue discussion:

Commit message: Move back to run_with_log

Patch to review:
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index c37e9ad2327..8a92bc045c2 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -599,6 +599,8 @@ def run_command_on_ip_via_ssh(
     port_forward: Optional[List[int]] = None,
     # Advanced options.
     log_path: str = '/dev/null',
+    # Do not redirect stdout/stderr to optimize performance.
+    output_only: bool = False,
     stream_logs: bool = True,
     check: bool = False,
     ssh_mode: SshMode = SshMode.NON_INTERACTIVE,
@@ -652,7 +654,11 @@ def run_command_on_ip_via_ssh(
         shlex.quote(f'true && source ~/.bashrc && export OMP_NUM_THREADS=1 '
                     f'PYTHONWARNINGS=ignore && ({cmd})'),
     ]
-    return log_lib.run_with_log(command, log_path, stream_logs, check=check)
+    return log_lib.run_with_log(command,
+                                log_path,
+                                stream_logs,
+                                check=check,
+                                output_only=output_only)
 
 
 def run(cmd, **kwargs):
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 8e6e0fc0e75..3d90c40efef 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -1031,6 +1031,7 @@ def _set_tpu_name(self, cluster_config_file: str, num_nodes: int,
                 cmd,
                 ssh_user=ssh_user,
                 ssh_private_key=ssh_private_key,
+                log_path=os.path.join(self.log_dir, 'tpu_setup.log'),
                 check=True)
 
     def provision(self,
@@ -1378,7 +1379,7 @@ def _exec_code_on_head(
         job_submit_cmd = (
             f'mkdir -p {remote_log_dir} && ray job submit '
             f'--address=127.0.0.1:8265 --job-id {job_id} --no-wait '
-            f'-- "{executable} -u {script_path} > {remote_log_path} 2>&1"')
+            f'-- "{executable} -u {script_path} 2>&1 | tee {remote_log_path}"')
 
         self.run_on_head(handle,
                          f'{cd} && {job_submit_cmd}',
@@ -1419,7 +1420,12 @@ def tail_logs(self, handle: ResourceHandle, job_id: int) -> None:
         code = codegen.build()
         click.secho('Start streaming logs...', fg='yellow')
         try:
-            self.run_on_head(handle, code, stream_logs=True, check=False)
+            self.run_on_head(handle,
+                             code,
+                             stream_logs=True,
+                             check=False,
+                             output_only=True,
+                             ssh_mode=backend_utils.SshMode.INTERACTIVE)
         except KeyboardInterrupt:
             # Do nothing. When receiving ctrl-c.
             pass
@@ -1804,6 +1810,7 @@ def run_on_head(
         *,
         port_forward: Optional[List[str]] = None,
         log_path: str = '/dev/null',
+        output_only: bool = False,
         stream_logs: bool = False,
         use_cached_head_ip: bool = True,
         check: bool = False,
@@ -1824,6 +1831,7 @@ def run_on_head(
             ssh_private_key=ssh_private_key,
             port_forward=port_forward,
             log_path=log_path,
+            output_only=output_only,
             stream_logs=stream_logs,
             check=check,
             ssh_mode=ssh_mode,
diff --git a/sky/cli.py b/sky/cli.py
index 3df3a0207e5..013eb58558c 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -651,7 +651,8 @@ def exec(cluster: str, entrypoint: str, detach_run: bool,
                         entrypoint,
                         stream_logs=True,
                         ssh_mode=backend_utils.SshMode.INTERACTIVE,
-                        under_remote_workdir=True)
+                        under_remote_workdir=True,
+                        output_only=True)
                     return
 
         # Override.
diff --git a/sky/skylet/log_lib.py b/sky/skylet/log_lib.py
index 5e0d5dffe9d..f290b86f5eb 100644
--- a/sky/skylet/log_lib.py
+++ b/sky/skylet/log_lib.py
@@ -8,9 +8,8 @@
 import subprocess
 import sys
 import textwrap
-import time
 import tempfile
-from typing import Iterator, List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 from sky.skylet import job_lib
 
@@ -89,6 +88,7 @@ def run_with_log(
     check: bool = False,
     shell: bool = False,
     with_ray: bool = False,
+    output_only: bool = False,
     **kwargs,
 ) -> Union[None, Tuple[subprocess.Popen, str, str]]:
     """Runs a command and logs its output to a file.
@@ -96,11 +96,13 @@ def run_with_log(
     Retruns the process, stdout and stderr of the command.
       Note that the stdout and stderr is already decoded.
     """
+    assert not (output_only and log_path != '/dev/null')
     # Redirect stderr to stdout when using ray, to preserve the order of
     # stdout and stderr.
+    stdout = None if output_only else subprocess.PIPE
     stderr = subprocess.PIPE if not with_ray else subprocess.STDOUT
     with subprocess.Popen(cmd,
-                          stdout=subprocess.PIPE,
+                          stdout=stdout,
                           stderr=stderr,
                           start_new_session=True,
                           shell=shell,
@@ -125,23 +127,28 @@ def run_with_log(
             stdout=subprocess.DEVNULL,
             stderr=subprocess.DEVNULL,
         )
-        # We need this even if the log_path is '/dev/null' to ensure the
-        # progress bar is shown.
-        stdout, stderr = redirect_process_output(
-            proc,
-            log_path,
-            stream_logs,
-            start_streaming_at=start_streaming_at,
-            # Skip these lines caused by `-i` option of bash. Failed to find
-            # other way to turn off these two warning.
-            # https://stackoverflow.com/questions/13300764/how-to-tell-bash-not-to-issue-warnings-cannot-set-terminal-process-group-and # pylint: disable=line-too-long
-            skip_lines=[
-                'bash: cannot set terminal process group',
-                'bash: no job control in this shell',
-            ],
-            # Replace CRLF when the output is logged to driver by ray.
-            replace_crlf=with_ray,
-        )
+        if output_only:
+            # Do not redirect stdout/stderr to improve performance.
+            stdout = ''
+            stderr = ''
+        else:
+            # We need this even if the log_path is '/dev/null' to ensure the
+            # progress bar is shown.
+            stdout, stderr = redirect_process_output(
+                proc,
+                log_path,
+                stream_logs,
+                start_streaming_at=start_streaming_at,
+                # Skip these lines caused by `-i` option of bash. Failed to find
+                # other way to turn off these two warning.
+                # https://stackoverflow.com/questions/13300764/how-to-tell-bash-not-to-issue-warnings-cannot-set-terminal-process-group-and # pylint: disable=line-too-long
+                skip_lines=[
+                    'bash: cannot set terminal process group',
+                    'bash: no job control in this shell',
+                ],
+                # Replace CRLF when the output is logged to driver by ray.
+                replace_crlf=with_ray,
+            )
         proc.wait()
         if proc.returncode and check:
             if stderr:
@@ -196,48 +203,8 @@ def run_bash_command_with_log(bash_command: str,
         )
 
 
-def _follow_job_logs(file,
-                     job_id: int,
-                     sleep_sec: float = 0.5,
-                     start_streaming_at: str = '') -> Iterator[str]:
-    """Yield each line from a file as they are written.
-
-    `sleep_sec` is the time to sleep after empty reads. """
-    line = ''
-    status = job_lib.query_job_status([job_id])[0]
-    start_streaming = False
-    while True:
-        tmp = file.readline()
-        if tmp is not None and tmp != '':
-            line += tmp
-            if '\n' in line or '\r' in line:
-                if start_streaming_at in line:
-                    start_streaming = True
-                if start_streaming:
-                    # TODO(zhwu): Consider using '\33[2K' to clear the
-                    # line when line endswith '\r' (to avoid previous line
-                    # to long problem). `colorama.ansi.clear_line`
-                    yield line
-                line = ''
-        else:
-            # Reach the end of the file, check the status or sleep and
-            # retry.
-
-            # Auto-exit the log tailing, if the job has finished. Check
-            # the job status before query again to avoid unfinished logs.
-            if status not in [
-                    job_lib.JobStatus.RUNNING, job_lib.JobStatus.PENDING
-            ]:
-                return
-
-            if sleep_sec:
-                time.sleep(sleep_sec)
-            status = job_lib.query_job_status([job_id])[0]
-
-
 def tail_logs(job_id: int, log_dir: Optional[str],
-              status: Optional[job_lib.JobStatus]):
-    # TODO(zhwu): Maybe switch to `ray job logs` according to the performance.
+              status: Optional[job_lib.JobStatus]) -> None:
     if log_dir is None:
         print(f'Job {job_id} not found (see `sky queue`).', file=sys.stderr)
         return
@@ -246,14 +213,14 @@ def tail_logs(job_id: int, log_dir: Optional[str],
     log_path = os.path.expanduser(log_path)
     if status in [job_lib.JobStatus.RUNNING, job_lib.JobStatus.PENDING]:
         try:
-            with open(log_path, 'r', newline='') as log_file:
-                # Using `_follow` instead of `tail -f` to streaming the whole
-                # log and creating a new process for tail.
-                for line in _follow_job_logs(
-                        log_file,
-                        job_id=job_id,
-                        start_streaming_at='SKY INFO: Reserving task slots on'):
-                    print(line, end='', flush=True)
+            tail_cmd = [
+                'ray', 'job', 'logs', '--address', '127.0.0.1:8265', '--follow',
+                f'{job_id}'
+            ]
+            run_with_log(tail_cmd,
+                         log_path='/dev/null',
+                         stream_logs=True,
+                         output_only=True)
         except KeyboardInterrupt:
             return
     else:


Code review comments:
diff_hunk: @@ -89,18 +88,21 @@ def run_with_log(
     check: bool = False,
     shell: bool = False,
     with_ray: bool = False,
+    output_only: bool = False,
     **kwargs,
 ) -> Union[None, Tuple[subprocess.Popen, str, str]]:
     """Runs a command and logs its output to a file.
 
     Retruns the process, stdout and stderr of the command.
       Note that the stdout and stderr is already decoded.
     """
+    assert not (output_only and log_path != '/dev/null')
     # Redirect stderr to stdout when using ray, to preserve the order of
     # stdout and stderr.
+    stdout = None if output_only else subprocess.PIPE
path: sky/skylet/log_lib.py
text: @user1:
Don't we want to do something to stderr, according to L131?

@author:
Good catch! Fixed. Thanks!

diff_hunk: @@ -1419,7 +1420,12 @@ def tail_logs(self, handle: ResourceHandle, job_id: int) -> None:
         code = codegen.build()
         click.secho('Start streaming logs...', fg='yellow')
         try:
-            self.run_on_head(handle, code, stream_logs=True, check=False)
+            self.run_on_head(handle,
+                             code,
+                             stream_logs=True,
+                             check=False,
+                             output_only=True,
+                             ssh_mode=backend_utils.SshMode.INTERACTIVE)
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
Comment on L1428: Allocate a pseudo-terminal to disable output buffering.

diff_hunk: @@ -246,14 +213,14 @@ def tail_logs(job_id: int, log_dir: Optional[str],
     log_path = os.path.expanduser(log_path)
     if status in [job_lib.JobStatus.RUNNING, job_lib.JobStatus.PENDING]:
         try:
-            with open(log_path, 'r', newline='') as log_file:
-                # Using `_follow` instead of `tail -f` to streaming the whole
-                # log and creating a new process for tail.
-                for line in _follow_job_logs(
-                        log_file,
-                        job_id=job_id,
-                        start_streaming_at='SKY INFO: Reserving task slots on'):
-                    print(line, end='', flush=True)
+            tail_cmd = [
+                'ray', 'job', 'logs', '--address', '127.0.0.1:8265', '--follow',
+                f'{job_id}'
+            ]
+            run_with_log(tail_cmd,
+                         log_path='/dev/null',
+                         stream_logs=True,
+                         output_only=True)
path: sky/skylet/log_lib.py
text: @user1:
Just reading this call site without looking inside, I do not know what `output_only=True` means.  How about `redirect_stdout_stderr=False` for this func and related funcs?

@author:
Good point. Fixed. Thanks!