Repository: skypilot-org/skypilot
Pull request: Support glob patterns for jobs for `sky logs -s`
Closes #625 by supporting globbing on `sky logs -s`

Example output:
<img width="1132" alt="image" src="https://user-images.githubusercontent.com/8518898/161404721-6963ba25-6ea3-4295-95e6-e25a54b2c1ba.png">

Problem statement:
Make `sky logs -s` support glob pattern on the job IDs
Feature request from Daniel:

> I have 32 jobs and I want to search all of them for certain keywords (to match script outputs with wandb runs based on the wandb run name). Right now I'm running
for j in `seq 56 87`; do sky logs -s myclus $j; done
and then I'm planning to grep the logs locally, but **syncing each log takes a long time**. I tried sky logs -s myclus "*" but got an error. It's not a huge deal to wait, but is there any way to sync a bunch of logs all at once?

It'd be ideal to support globing job IDs, e.g., `sky logs -s myclus "*"`, when `-s` is specified.

(Also, the current `sky logs -s` code path does one roundtrip to get the log dir, then does another for the actual rsync.  We may want to optimize this a bit.)

Hints from issue discussion:

Commit message: removed whitespace

Patch to review:
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index cd1cbe306b9..011aaa45290 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -1679,29 +1679,36 @@ def get_job_status(self, handle: ResourceHandle,
             return None
         return job_lib.JobStatus(result.split(' ')[-1])
 
-    def sync_down_logs(self, handle: ResourceHandle, job_id: int) -> None:
-        code = job_lib.JobLibCodeGen.get_log_path(job_id)
-        returncode, log_dir, stderr = self.run_on_head(handle,
-                                                       code,
-                                                       stream_logs=False,
-                                                       require_outputs=True)
+    def sync_down_logs(self, handle: ResourceHandle, job_ids: str) -> None:
+        code = job_lib.JobLibCodeGen.get_log_path_with_globbing(job_ids)
+        returncode, log_dirs, stderr = self.run_on_head(handle,
+                                                        code,
+                                                        stream_logs=False,
+                                                        require_outputs=True)
         backend_utils.handle_returncode(returncode, code,
                                         'Failed to sync logs.', stderr)
-        log_dir = log_dir.strip()
-
-        local_log_dir = os.path.expanduser(log_dir)
-        remote_log_dir = log_dir
+        log_dirs = ast.literal_eval(log_dirs)
+        job_ids = [log_dir[0] for log_dir in log_dirs]
+        local_log_dirs = [
+            os.path.expanduser(log_dir[1]) for log_dir in log_dirs
+        ]
+        remote_log_dirs = [log_dir[1] for log_dir in log_dirs]
 
         style = colorama.Style
         fore = colorama.Fore
-        logger.info(f'{fore.CYAN}Logs Directory: '
-                    f'{style.BRIGHT}{local_log_dir}{style.RESET_ALL}')
+        logger.info(f'{fore.CYAN}Logs Directories: {style.RESET_ALL}')
+        for job_id, log_dir in zip(job_ids, local_log_dirs):
+            logger.info(f'{fore.CYAN}'
+                        f'{style.BRIGHT}Job ID: {style.NORMAL}{job_id}'
+                        f' {style.BRIGHT}Path: {style.NORMAL}{log_dir}'
+                        f'{style.RESET_ALL}')
 
         ips = backend_utils.get_node_ips(handle.cluster_yaml,
                                          handle.launched_nodes,
                                          handle=handle)
 
-        def rsync_down(ip: str) -> None:
+        def rsync_down(ip: str, local_log_dir: str,
+                       remote_log_dir: str) -> None:
             from ray.autoscaler import sdk  # pylint: disable=import-outside-toplevel
             sdk.rsync(
                 handle.cluster_yaml,
@@ -1714,21 +1721,27 @@ def rsync_down(ip: str) -> None:
             )
 
         # Call the ray sdk to rsync the logs back to local.
-        for i, ip in enumerate(ips):
-            try:
-                # Disable the output of rsync.
-                with open('/dev/null', 'a') as f, contextlib.redirect_stdout(
-                        f), contextlib.redirect_stderr(f):
-                    rsync_down(ip)
-                logger.info(f'{fore.CYAN}Job {job_id} logs: Downloaded from '
-                            f'node-{i} ({ip}){style.RESET_ALL}')
-            except click.exceptions.ClickException as e:
-                # Raised by rsync_down. Remote log dir may not exist, since
-                # the job can be run on some part of the nodes.
-                if 'SSH command failed' in str(e):
-                    logger.debug(f'node-{i} ({ip}) does not have the tasks/*.')
-                else:
-                    raise e
+        for job_id, local_log_dir, remote_log_dir in zip(
+                job_ids, local_log_dirs, remote_log_dirs):
+            for i, ip in enumerate(ips):
+                try:
+                    # Disable the output of rsync.
+                    with open('/dev/null',
+                              'a') as f, contextlib.redirect_stdout(
+                                  f), contextlib.redirect_stderr(f):
+                        rsync_down(ip, local_log_dir, remote_log_dir)
+                        pass
+                    logger.info(
+                        f'{fore.CYAN}Job {job_id} logs: Downloaded from '
+                        f'node-{i} ({ip}){style.RESET_ALL}')
+                except click.exceptions.ClickException as e:
+                    # Raised by rsync_down. Remote log dir may not exist, since
+                    # the job can be run on some part of the nodes.
+                    if 'SSH command failed' in str(e):
+                        logger.debug(
+                            f'node-{i} ({ip}) does not have the tasks/*.')
+                    else:
+                        raise e
 
     def _exec_code_on_head(
         self,
diff --git a/sky/skylet/job_lib.py b/sky/skylet/job_lib.py
index 29bde4a2542..8a3da5f92cd 100644
--- a/sky/skylet/job_lib.py
+++ b/sky/skylet/job_lib.py
@@ -333,11 +333,29 @@ def log_dir(job_id: int) -> Optional[str]:
             WHERE job_id=(?)""", (job_id,))
     row = _CURSOR.fetchone()
     if row is None:
-        return None, None
+        return None
     run_timestamp = row[JobInfoLoc.RUN_TIMESTAMP.value]
     return os.path.join(SKY_LOGS_DIRECTORY, run_timestamp)
 
 
+def log_dirs_with_globbing(job_id: int) -> Optional[str]:
+    """Returns the relative paths to the log files for job with globbing."""
+    _CURSOR.execute(
+        """\
+            SELECT * FROM jobs
+            WHERE job_id GLOB (?)""", (job_id,))
+    rows = _CURSOR.fetchall()
+    if len(rows) == 0:
+        return None
+    log_paths = []
+    for row in rows:
+        job_id = row[JobInfoLoc.JOB_ID.value]
+        run_timestamp = row[JobInfoLoc.RUN_TIMESTAMP.value]
+        log_path = os.path.join(SKY_LOGS_DIRECTORY, run_timestamp)
+        log_paths.append((job_id, log_path))
+    return log_paths
+
+
 class JobLibCodeGen:
     """Code generator for job utility functions.
 
@@ -407,6 +425,14 @@ def get_log_path(cls, job_id: int) -> str:
         ]
         return cls._build(code)
 
+    @classmethod
+    def get_log_path_with_globbing(cls, job_id: int) -> str:
+        code = [
+            f'log_dirs = job_lib.log_dirs_with_globbing(\'{job_id}\')',
+            'print(log_dirs, flush=True)',
+        ]
+        return cls._build(code)
+
     @classmethod
     def _build(cls, code: List[str]) -> str:
         code = cls._PREFIX + code


Code review comments:
diff_hunk: @@ -407,6 +425,14 @@ def get_log_path(cls, job_id: int) -> str:
         ]
         return cls._build(code)
 
+    @classmethod
+    def get_log_path_with_globbing(cls, job_id: int) -> str:
path: sky/skylet/job_lib.py
text: @user1:
nit: int -> str?

diff_hunk: @@ -1714,21 +1721,27 @@ def rsync_down(ip: str) -> None:
             )
 
         # Call the ray sdk to rsync the logs back to local.
-        for i, ip in enumerate(ips):
-            try:
-                # Disable the output of rsync.
-                with open('/dev/null', 'a') as f, contextlib.redirect_stdout(
-                        f), contextlib.redirect_stderr(f):
-                    rsync_down(ip)
-                logger.info(f'{fore.CYAN}Job {job_id} logs: Downloaded from '
-                            f'node-{i} ({ip}){style.RESET_ALL}')
-            except click.exceptions.ClickException as e:
-                # Raised by rsync_down. Remote log dir may not exist, since
-                # the job can be run on some part of the nodes.
-                if 'SSH command failed' in str(e):
-                    logger.debug(f'node-{i} ({ip}) does not have the tasks/*.')
-                else:
-                    raise e
+        for job_id, local_log_dir, remote_log_dir in zip(
+                job_ids, local_log_dirs, remote_log_dirs):
+            for i, ip in enumerate(ips):
+                try:
+                    # Disable the output of rsync.
+                    with open('/dev/null',
+                              'a') as f, contextlib.redirect_stdout(
+                                  f), contextlib.redirect_stderr(f):
+                        rsync_down(ip, local_log_dir, remote_log_dir)
+                        pass
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
nit: remnant?

diff_hunk: @@ -333,11 +333,29 @@ def log_dir(job_id: int) -> Optional[str]:
             WHERE job_id=(?)""", (job_id,))
     row = _CURSOR.fetchone()
     if row is None:
-        return None, None
+        return None
     run_timestamp = row[JobInfoLoc.RUN_TIMESTAMP.value]
     return os.path.join(SKY_LOGS_DIRECTORY, run_timestamp)
 
 
+def log_dirs_with_globbing(job_id: int) -> Optional[str]:
path: sky/skylet/job_lib.py
text: @user1:
nit: type hint to str?