Repository: skypilot-org/skypilot
Pull request: Add logging for execution_v2
Added the logging for the execution_v2 (as in #25 ). The behavior would be as follows:
1. All the provision information/task output will be written to files so that users can tail the files for the output.
2. If the `stream_logs` is set to true, the output provision logs will be streamed to the console if the `Shared connection to` appears in the output, so that only the succeeded `ray up`'s output will be shown in the console.

Problem statement:
Stream logs to files during VM provisioning and setup
Is it expected behavior that there is no output for VM setup during provisioning? I think it might be hard for the user to debug for the setup commends. Also, I failed to find the output anywhere as a file.

I am running `python ./examples/resnet_app.py`, and the output is as below.
```
I 11-12 10:57:38 cloud_vm_ray_backend.py:213] If this takes longer than ~30 seconds, provisioning is likely successful. Setting up may take a few minutes.
I 11-12 10:57:38 backend_utils.py:38] Created or updated file config/aws-ray.yml
I 11-12 11:02:52 cloud_vm_ray_backend.py:238] Successfully provisioned or found existing VM(s). Setup completed.
```

Hints from issue discussion:
Yes, expected to not show, although we can change.  The rationale is if you print ray up's stdout/stderr, and if you have retries, that will be very crowded.  Give it a try by setting a p3.16x requirement.

We could try to write the stdout/stderr (captured in variables) to files.  This is similar to V1's logging.
We should at least stream the logs to files and tell user to `tail -F path/to/logfile` if they're curious about progress. When I try it now I have to stare at a blank screen for 10 minutes with no progress:

```
I 11-16 19:37:15 cloud_vm_ray_backend.py:210] Launching on AWS us-west-2 (us-west-2a,us-west-2b,us-west-2c,us-west-2d).
I 11-16 19:37:15 cloud_vm_ray_backend.py:213] If this takes longer than ~30 seconds, provisioning is likely successful. Setting up may take a few minutes.
I 11-16 19:37:15 backend_utils.py:38] Created or updated file config/aws-ray.yml

```
+1 to stream to a file and hint on tail -f.

Commit message: Fix tail -f information for provision

Patch to review:
diff --git a/prototype/sky/backends/backend.py b/prototype/sky/backends/backend.py
index 4a37b0ec685..0ec810d27dd 100644
--- a/prototype/sky/backends/backend.py
+++ b/prototype/sky/backends/backend.py
@@ -17,8 +17,8 @@ class Backend(object):
     # Examples: 'cluster.yaml'; 'ray://...', 'k8s://...'.
     ResourceHandle = Any
 
-    def provision(self, task: App, to_provision: Resources,
-                  dryrun: bool) -> ResourceHandle:
+    def provision(self, task: App, to_provision: Resources, dryrun: bool,
+                  stream_logs: bool) -> ResourceHandle:
         raise NotImplementedError
 
     def sync_workdir(self, handle: ResourceHandle, workdir: Path) -> None:
diff --git a/prototype/sky/backends/backend_utils.py b/prototype/sky/backends/backend_utils.py
index 6f9105c0145..7a0e6de051c 100644
--- a/prototype/sky/backends/backend_utils.py
+++ b/prototype/sky/backends/backend_utils.py
@@ -1,6 +1,8 @@
 """Util constants/functions for the backends."""
 import datetime
+import io
 import os
+import selectors
 import subprocess
 import tempfile
 import textwrap
@@ -24,6 +26,7 @@
 # NOTE: keep in sync with the cluster template 'file_mounts'.
 SKY_REMOTE_WORKDIR = '/tmp/workdir'
 IP_ADDR_REGEX = r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}'
+SKY_LOGS_DIRECTORY = './sky_logs'
 
 
 def _fill_template(template_path: str,
@@ -224,3 +227,38 @@ def run_command_on_ip_via_ssh(ip: str,
         if errs:
             logger.error(errs)
         raise subprocess.CalledProcessError(proc.returncode, cmd)
+
+
+def redirect_process_output(proc, log_path, stream_logs, start_streaming_at=''):
+    """Redirect the process's filtered stdout/stderr to both stream and file"""
+    dirname = os.path.dirname(log_path)
+    os.makedirs(dirname, exist_ok=True)
+
+    out_io = io.TextIOWrapper(proc.stdout, encoding='utf-8', newline='')
+    err_io = io.TextIOWrapper(proc.stderr, encoding='utf-8', newline='')
+    sel = selectors.DefaultSelector()
+    sel.register(out_io, selectors.EVENT_READ)
+    sel.register(err_io, selectors.EVENT_READ)
+
+    stdout = ''
+    stderr = ''
+
+    start_streaming_flag = False
+    with open(log_path, 'a') as fout:
+        while True:
+            for key, _ in sel.select():
+                line = key.fileobj.readline()
+                if not line:
+                    return stdout, stderr
+                if start_streaming_at in line:
+                    start_streaming_flag = True
+                if key.fileobj is out_io:
+                    stdout += line
+                    fout.write(line)
+                    fout.flush()
+                else:
+                    stderr += line
+                    fout.write(line)
+                    fout.flush()
+                if stream_logs and start_streaming_flag:
+                    print(line, end='')
diff --git a/prototype/sky/backends/cloud_vm_ray_backend.py b/prototype/sky/backends/cloud_vm_ray_backend.py
index 7a01f0e0192..6306ba0f39c 100644
--- a/prototype/sky/backends/cloud_vm_ray_backend.py
+++ b/prototype/sky/backends/cloud_vm_ray_backend.py
@@ -26,9 +26,57 @@
 Path = str
 PostSetupFn = Callable[[str], Any]
 SKY_REMOTE_WORKDIR = backend_utils.SKY_REMOTE_WORKDIR
+SKY_LOGS_DIRECTORY = backend_utils.SKY_LOGS_DIRECTORY
 
 logger = logging.init_logger(__name__)
 
+TASK_LAUNCH_CODE_GENERATOR = """
+        import io
+        import os
+        import ray
+        import selectors
+        import subprocess
+        ray.init('auto', namespace='__sky__', log_to_driver={stream_logs})
+        print('cluster_resources:', ray.cluster_resources())
+        print('available_resources:', ray.available_resources())
+        print('live nodes:', ray.state.node_ids())
+        
+        def redirect_process_output(proc, log_path, stream_logs, start_streaming_at=''):
+            dirname = os.path.dirname(log_path)
+            os.makedirs(dirname, exist_ok=True)
+
+            out_io = io.TextIOWrapper(proc.stdout, encoding='utf-8', newline='')
+            err_io = io.TextIOWrapper(proc.stderr, encoding='utf-8', newline='')
+            sel = selectors.DefaultSelector()
+            sel.register(out_io, selectors.EVENT_READ)
+            sel.register(err_io, selectors.EVENT_READ)
+
+            stdout = ''
+            stderr = ''
+
+            start_streaming_flag = False
+            with open(log_path, 'a') as fout:
+                while True:
+                    for key, _ in sel.select():
+                        line = key.fileobj.readline()
+                        if not line:
+                            return stdout, stderr
+                        if start_streaming_at in line:
+                            start_streaming_flag = True
+                        if key.fileobj is out_io:
+                            stdout += line
+                            fout.write(line)
+                            fout.flush()
+                        else:
+                            stderr += line
+                            fout.write(line)
+                            fout.flush()
+                        if stream_logs and start_streaming_flag:
+                            print(line, end='')
+        
+        futures = []
+    """
+
 
 def _run(cmd, **kwargs):
     return subprocess.run(cmd, shell=True, check=True, **kwargs)
@@ -38,6 +86,21 @@ def _run_no_outputs(cmd, **kwargs):
     return _run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, **kwargs)
 
 
+def _run_with_log(cmd,
+                  log_path,
+                  stream_logs=False,
+                  start_streaming_at='',
+                  **kwargs):
+    proc = subprocess.Popen(cmd,
+                            stdout=subprocess.PIPE,
+                            stderr=subprocess.PIPE,
+                            **kwargs)
+    stdout, stderr = backend_utils.redirect_process_output(
+        proc, log_path, stream_logs, start_streaming_at=start_streaming_at)
+    proc.wait()
+    return proc, stdout, stderr
+
+
 def _get_cluster_config_template(task):
     _CLOUD_TO_TEMPLATE = {
         clouds.AWS: 'config/aws-ray.yml.j2',
@@ -64,9 +127,10 @@ def _to_accelerator_and_count(resources: Optional[Resources]
 class RetryingVmProvisioner(object):
     """A provisioner that retries different regions/zones within a cloud."""
 
-    def __init__(self):
+    def __init__(self, log_dir):
         self._blocked_regions = set()
         self._blocked_zones = set()
+        self.log_dir = log_dir
         colorama.init()
 
     def _in_blocklist(self, cloud, region, zones):
@@ -86,7 +150,6 @@ def _update_blocklist_on_gcp_error(self, region, zones, stdout, stderr):
         Style = colorama.Style
         assert len(zones) == 1, zones
         zone = zones[0]
-        stderr = stderr.decode()
         splits = stderr.split('\n')
         exception_str = [s for s in splits if s.startswith('Exception: ')]
         if len(exception_str) == 1:
@@ -116,7 +179,7 @@ def _update_blocklist_on_gcp_error(self, region, zones, stdout, stderr):
                 self._blocked_zones.add(zone.name)
             else:
                 logger.info('====== stdout ======')
-                for s in stdout.decode().split('\n'):
+                for s in stdout.split('\n'):
                     print(s)
                 logger.info('====== stderr ======')
                 for s in splits:
@@ -128,8 +191,8 @@ def _update_blocklist_on_aws_error(self, region, zones, stdout, stderr):
         # The underlying ray autoscaler / boto3 will try all zones of a region
         # at once.
         Style = colorama.Style
-        stdout_splits = stdout.decode().split('\n')
-        stderr_splits = stderr.decode().split('\n')
+        stdout_splits = stdout.split('\n')
+        stderr_splits = stderr.split('\n')
         errors = [
             s.strip()
             for s in stdout_splits + stderr_splits
@@ -200,8 +263,12 @@ def _yield_region_zones(self, task: App, cloud: clouds.Cloud):
             yield (region, zones)
 
     def provision_with_retries(self, task: App, to_provision: Resources,
-                               dryrun: bool):
+                               dryrun: bool, stream_logs: bool):
         """The provision retry loop."""
+        # Get log_path name
+        log_path = os.path.join(self.log_dir, 'provision.log')
+        log_abs_path = os.path.abspath(log_path)
+
         self._clear_blocklist()
         Style = colorama.Style
         for region, zones in self._yield_region_zones(task, to_provision.cloud):
@@ -243,11 +310,17 @@ def provision_with_retries(self, task: App, to_provision: Resources,
                         logger.error(stderr)
                         raise e
             cluster_config_file = config_dict['ray']
-            # Captures stdout/err.  Otherwise, too many repeated messages.
-            proc = subprocess.Popen(['ray', 'up', '-y', cluster_config_file],
-                                    stdout=subprocess.PIPE,
-                                    stderr=subprocess.PIPE)
-            stdout, stderr = proc.communicate()
+
+            tail_cmd = f'tail -n100 -f {log_path}'
+            logger.info(
+                f'To view progress: {Style.BRIGHT}{tail_cmd}{Style.RESET_ALL}')
+            # Redirect stdout/err to the file and streaming (if stream_logs).
+            proc, stdout, stderr = _run_with_log(
+                ['ray', 'up', '-y', cluster_config_file],
+                log_abs_path,
+                stream_logs,
+                start_streaming_at='Shared connection to')
+
             if proc.returncode != 0:
                 self._update_blocklist_on_error(to_provision.cloud, region,
                                                 zones, stdout, stderr)
@@ -290,14 +363,17 @@ class CloudVmRayBackend(backends.Backend):
     def __init__(self):
         # TODO: should include this as part of the handle.
         self._managed_tpu = None
+        run_id = backend_utils.get_run_id()
+        self.log_dir = os.path.join(SKY_LOGS_DIRECTORY, run_id)
+        os.makedirs(self.log_dir, exist_ok=True)
 
-    def provision(self, task: App, to_provision: Resources,
-                  dryrun: bool) -> ResourceHandle:
+    def provision(self, task: App, to_provision: Resources, dryrun: bool,
+                  stream_logs: bool) -> ResourceHandle:
         """Provisions using 'ray up'."""
         # ray up: the VMs.
-        provisioner = RetryingVmProvisioner()
+        provisioner = RetryingVmProvisioner(self.log_dir)
         config_dict = provisioner.provision_with_retries(
-            task, to_provision, dryrun)
+            task, to_provision, dryrun, stream_logs)
         if dryrun:
             return
         cluster_config_file = config_dict['ray']
@@ -371,16 +447,11 @@ def _execute_par_task(self, handle: ResourceHandle,
         # manipulating the futures).
         #
         # TODO: possible to open the port in the yaml?  Run Ray inside docker?
+        log_dir = os.path.join(f'{SKY_REMOTE_WORKDIR}', f'{self.log_dir}',
+                               'tasks')
         codegen = [
-            textwrap.dedent(f"""\
-        import ray
-        import subprocess
-        ray.init('auto', namespace='__sky__', log_to_driver={stream_logs})
-        print('cluster_resources:', ray.cluster_resources())
-        print('available_resources:', ray.available_resources())
-        print('live nodes:', ray.state.node_ids())
-        futures = []
-        """)
+            textwrap.dedent(
+                TASK_LAUNCH_CODE_GENERATOR.format(stream_logs=stream_logs))
         ]
         for i, t in enumerate(par_task.tasks):
             # '. $(conda info --base)/etc/profile.d/conda.sh || true' is used to initialize
@@ -404,10 +475,17 @@ def _execute_par_task(self, handle: ResourceHandle,
                 resources_str = ''
                 num_gpus_str = ''
             name = f'task-{i}' if t.name is None else t.name
+            log_path = os.path.join(f'{log_dir}', f'{name}.log')
+
             task_i_codegen = textwrap.dedent(f"""\
-        futures.append(ray.remote(lambda: subprocess.run(
-            {cmd},
-              shell=True, check=True, executable='/bin/bash')) \\
+        def start_task():
+            proc = subprocess.Popen({cmd},
+                                    stdout=subprocess.PIPE,
+                                    stderr=subprocess.PIPE,
+                                    shell=True,
+                                    executable='/bin/bash')
+            redirect_process_output(proc, '{log_path}', {stream_logs})
+        futures.append(ray.remote(start_task) \\
               .options(name='{name}'{resources_str}{num_gpus_str}) \\
               .remote())
         """)
@@ -424,10 +502,17 @@ def _execute_par_task(self, handle: ResourceHandle,
         if not stream_logs:
             logger.info(
                 f'{Fore.CYAN}Logs will not be streamed (stream_logs=False).'
-                f'{Style.RESET_ALL} Hint: in the run command, redirect each'
-                ' task\'s output to a file, and use `tail -f` to monitor.\n')
+                f'{Style.RESET_ALL} Hint: task\'s output will be redirect to'
+                f'{Style.BRIGHT} {log_dir} {Style.RESET_ALL}. Please use'
+                '`ray attach` and `tail -f` to monitor.\n')
 
         self._exec_code_on_head(handle, codegen)
+        task_log_dir = os.path.join(f'{self.log_dir}', 'tasks')
+        logger.info(
+            f'Syncing down the loggings to {Style.BRIGHT}{task_log_dir}{Style.RESET_ALL}.'
+        )
+        os.makedirs(task_log_dir, exist_ok=True)
+        _run(f'ray rsync_down {handle} {log_dir}/* {task_log_dir}')
 
     def _exec_code_on_head(self,
                            handle: ResourceHandle,
@@ -447,15 +532,14 @@ def _exec_code_on_head(self,
             executable = 'python3'
         cd = f'cd {SKY_REMOTE_WORKDIR}'
         cmd = f'ray exec {handle} \'{cd} && {executable} /tmp/{basename}\''
+        log_path = os.path.join(self.log_dir, f'run.log')
         if not stream_logs:
-            out = tempfile.NamedTemporaryFile('w', prefix='sky_',
-                                              suffix='.out').name
-            cmd += f' >{out}'
             colorama.init()
             Style = colorama.Style
             logger.info(f'Redirecting stdout, to monitor: '
-                        f'{Style.BRIGHT}tail -f {out}{Style.RESET_ALL}')
-        _run(cmd)
+                        f'{Style.BRIGHT}tail -f {log_path}{Style.RESET_ALL}')
+
+        _run_with_log(cmd, log_path, stream_logs, shell=True)
 
     def execute(self, handle: ResourceHandle, task: App,
                 stream_logs: bool) -> None:
@@ -485,7 +569,7 @@ def _execute_task_one_node(self, handle: ResourceHandle, task: App,
             f'Task(run=...) should be a string (found {type(task.run)}).'
         codegen = textwrap.dedent(f"""\
             #!/bin/bash
-            . $(conda info --base)/etc/profile.d/conda.sh
+            . $(conda info --base)/etc/profile.d/conda.sh || true
             {task.run}
         """)
         self._exec_code_on_head(handle, codegen, stream_logs, executable='bash')
@@ -496,16 +580,11 @@ def _execute_task_n_nodes(self, handle: ResourceHandle, task: App,
         #   ray.init(..., log_to_driver=False); otherwise too many logs.
         #   for node:
         #     submit _run_cmd(cmd) with resource {node_i: 1}
+        log_dir = os.path.join(f'{SKY_REMOTE_WORKDIR}', f'{self.log_dir}',
+                               'tasks')
         codegen = [
-            textwrap.dedent(f"""\
-        import subprocess
-        import ray
-        ray.init('auto', namespace='__sky__', log_to_driver={stream_logs})
-        print('cluster_resources:', ray.cluster_resources())
-        print('available_resources:', ray.available_resources())
-        print('live nodes:', ray.state.node_ids())
-        futures = []
-        """)
+            textwrap.dedent(
+                TASK_LAUNCH_CODE_GENERATOR.format(stream_logs=stream_logs))
         ]
         acc, acc_count = _to_accelerator_and_count(task.best_resources)
         # Get private ips here as Ray internally uses 'node:private_ip' as
@@ -533,13 +612,21 @@ def _execute_task_n_nodes(self, handle: ResourceHandle, task: App,
                 num_gpus_str = f', num_gpus={acc_count}'
             # Set the executable to /bin/bash, so that the 'source ~/.bashrc'
             # and 'source activate conda_env' can be used.
+            name = f'task-{ip}'
+            log_path = os.path.join(f'{log_dir}', f'{name}.log')
+
             codegen.append(
                 textwrap.dedent(f"""\
-        futures.append(ray.remote(lambda: subprocess.run(
-            {cmd},
-                shell=True, check=True, executable='/bin/bash')) \\
-                .options(name='task-{ip}'{resources_str}{num_gpus_str}) \\
-                .remote())
+        def start_task():
+            proc = subprocess.Popen({cmd},
+                                    stdout=subprocess.PIPE,
+                                    stderr=subprocess.PIPE,
+                                    shell=True,
+                                    executable='/bin/bash')
+            redirect_process_output(proc, '{log_path}', {stream_logs})
+        futures.append(ray.remote(start_task) \\
+              .options(name='{name}'{resources_str}{num_gpus_str}) \\
+              .remote())
         """))
         # Block.
         codegen.append('ray.get(futures)\n')
@@ -552,9 +639,17 @@ def _execute_task_n_nodes(self, handle: ResourceHandle, task: App,
         if not stream_logs:
             logger.info(
                 f'{Fore.CYAN}Logs will not be streamed (stream_logs=False).'
-                f'{Style.RESET_ALL} Hint: in the run command, redirect each'
-                ' task\'s output to a file, and use `tail -f` to monitor.\n')
+                f'{Style.RESET_ALL} Hint: task\'s output will be redirect to'
+                f'{Style.BRIGHT} {log_dir} {Style.RESET_ALL}. Please use'
+                '`ray attach` and `tail -f` to monitor.\n')
+
         self._exec_code_on_head(handle, codegen)
+        task_log_dir = os.path.join(f'{self.log_dir}', 'tasks')
+        logger.info(
+            f'Syncing down the loggings to {Style.BRIGHT}{task_log_dir}{Style.RESET_ALL}.'
+        )
+        os.makedirs(task_log_dir, exist_ok=True)
+        _run(f'ray rsync_down {handle} {log_dir}/* {task_log_dir}')
 
     def post_execute(self, handle: ResourceHandle, teardown: bool) -> None:
         colorama.init()
diff --git a/prototype/sky/execution.py b/prototype/sky/execution.py
index 12185b93c89..011e7f1c5a9 100644
--- a/prototype/sky/execution.py
+++ b/prototype/sky/execution.py
@@ -341,7 +341,10 @@ def execute_v2(dag: sky.Dag,
     # Future backends: K8S, SLURM, VM, LOCAL, LOCAL_DOCKER, etc.
     backend = backends.CloudVmRayBackend()
 
-    handle = backend.provision(task, best_resources, dryrun=dryrun)
+    handle = backend.provision(task,
+                               best_resources,
+                               dryrun=dryrun,
+                               stream_logs=stream_logs)
     if dryrun:
         logger.info('Dry run finished.')
         return
diff --git a/prototype/sky/logging.py b/prototype/sky/logging.py
index 6e230e02cb7..41aee47208d 100644
--- a/prototype/sky/logging.py
+++ b/prototype/sky/logging.py
@@ -1,12 +1,12 @@
-import sys
 import logging
+import sys
 
 FORMAT = '%(levelname).1s %(asctime)s %(filename)s:%(lineno)d] %(message)s'
 DATE_FORMAT = '%m-%d %H:%M:%S'
 
 
 class NewLineFormatter(logging.Formatter):
-
+    """Add logging prefix to the newlines to align multi-line logging messages"""
     def __init__(self, fmt, datefmt=None):
         logging.Formatter.__init__(self, fmt, datefmt)
 


Code review comments:
diff_hunk: @@ -424,10 +502,17 @@ def _execute_par_task(self, handle: ResourceHandle,
         if not stream_logs:
             logger.info(
                 f'{Fore.CYAN}Logs will not be streamed (stream_logs=False).'
-                f'{Style.RESET_ALL} Hint: in the run command, redirect each'
-                ' task\'s output to a file, and use `tail -f` to monitor.\n')
+                f'{Style.RESET_ALL} Hint: task\'s output will be redirect to'
+                f'{Style.BRIGHT} {log_dir} {Style.RESET_ALL}. Please use'
+                '`ray attach` and `tail -f` to monitor.\n')
path: prototype/sky/backends/cloud_vm_ray_backend.py
text: @user1:
Can now be explicit as we have control of the logs, something like:
```
                f'{Style.RESET_ALL} Hint: task outputs are redirected to'
                f'{Style.BRIGHT} {log_dir} {Style.RESET_ALL} on the cluster. To monitor:'
                f'ray exec {handle} "tail -f {log_dir}/*.log"\n')
```

@author:
Yes! This would be what we want.

diff_hunk: @@ -26,9 +26,57 @@
 Path = str
 PostSetupFn = Callable[[str], Any]
 SKY_REMOTE_WORKDIR = backend_utils.SKY_REMOTE_WORKDIR
+SKY_LOGS_DIRECTORY = backend_utils.SKY_LOGS_DIRECTORY
 
 logger = logging.init_logger(__name__)
 
+TASK_LAUNCH_CODE_GENERATOR = """
path: prototype/sky/backends/cloud_vm_ray_backend.py
text: @user1:
Nits: Add a `\` at the end of this line, so that the codegen'd file doesn't have its first line blank.

Add `_` to the beginning of the variable to indicate module-local constants: https://google.github.io/styleguide/pyguide.html#s3.16.2-naming-conventions