[Core] Exit with non-zero code on launch/exec/logs/jobs launch/jobs logs (skypilot-org#4846)

* Support return code based on job success/failure

* Return exit code for tailing managed jobs

* Fixes

* lint

* Create JobExitCode enum

* Get JobExitCode from ManagedJobStatus

* lint

* cleanup

* cleanup

* Add tests

* lint

* Managed jobs back compat

* Skylet backward compatibility

* lint

* Update logs --status returncodes

* Update logs --status returncodes

* lint

* fix retcode

* Fix tests

* lint

* Fix --no-follow

* Fix cli docs rendering

* minor

* rename ret_code to returncode

* rename SUCCESS to SUCCEEDED

* Refactor JobExitCode to exceptions

* lint

Author: romilbhardwaj

sky/backends/cloud_vm_ray_backend.py

@@ -3823,6 +3823,10 @@ def tail_logs(self,
             follow: Whether to follow the logs.
             tail: The number of lines to display from the end of the
                 log file. If 0, print all lines.
+
+        Returns:
+            The exit code of the tail command. Returns code 100 if the job has
+            failed. See exceptions.JobExitCode for possible return codes.
         """
         code = job_lib.JobLibCodeGen.tail_logs(job_id,
                                                managed_job_id=managed_job_id,
@@ -3856,7 +3860,7 @@ def tail_managed_job_logs(self,
                               job_id: Optional[int] = None,
                               job_name: Optional[str] = None,
                               controller: bool = False,
-                              follow: bool = True) -> None:
+                              follow: bool = True) -> int:
         # if job_name is not None, job_id should be None
         assert job_name is None or job_id is None, (job_name, job_id)
         code = managed_jobs.ManagedJobCodeGen.stream_logs(
@@ -3869,13 +3873,17 @@ def tail_managed_job_logs(self,
             signal.signal(signal.SIGTSTP, backend_utils.stop_handler)
 
         # Refer to the notes in tail_logs.
-        self.run_on_head(
-            handle,
-            code,
-            stream_logs=True,
-            process_stream=False,
-            ssh_mode=command_runner.SshMode.INTERACTIVE,
-        )
+        try:
+            returncode = self.run_on_head(
+                handle,
+                code,
+                stream_logs=True,
+                process_stream=False,
+                ssh_mode=command_runner.SshMode.INTERACTIVE,
+            )
+        except SystemExit as e:
+            returncode = e.code
+        return returncode
 
     def sync_down_managed_job_logs(
             self,

sky/cli.py

@@ -1227,11 +1227,15 @@ def launch(
             clusters=[handle.get_cluster_name()])
         # job_id will be None if no job was submitted (e.g. no entrypoint
         # provided)
+        returncode = 0
         if not detach_run and job_id is not None:
-            sdk.tail_logs(handle.get_cluster_name(), job_id, follow=True)
+            returncode = sdk.tail_logs(handle.get_cluster_name(),
+                                       job_id,
+                                       follow=True)
         click.secho(
             ux_utils.command_hint_messages(ux_utils.CommandHintType.CLUSTER_JOB,
                                            job_id, handle.get_cluster_name()))
+        sys.exit(returncode)
 
 
 @cli.command(cls=_DocumentedCodeCommand)
@@ -1377,7 +1381,8 @@ def exec(cluster: Optional[str], cluster_option: Optional[str],
     job_id_handle = _async_call_or_wait(request_id, async_call, 'sky.exec')
     if not async_call and not detach_run:
         job_id, _ = job_id_handle
-        sdk.tail_logs(cluster, job_id, follow=True)
+        returncode = sdk.tail_logs(cluster, job_id, follow=True)
+        sys.exit(returncode)
 
 
 def _handle_jobs_queue_request(
@@ -2121,12 +2126,20 @@ def logs(
     one job_id can be provided.
 
     2. If ``--status`` is specified, print the status of the job and exit with
-    returncode 0 if the job succeeded, or 1 otherwise. At most one job_id can
-    be specified.
+    returncode 0 if the job succeeded. At most one job_id can
+    be specified. Other possible return codes:
+
+    - 100: job failed.
+    - 101: job not finished.
+    - 102: job not found.
+    - 103: job was cancelled by the user.
 
     3. If ``--sync-down`` is specified, the logs of the job will be downloaded
     from the cluster and saved to the local machine under
-    ``~/sky_logs``. Mulitple job_ids can be specified.
+    ``~/sky_logs``. Multiple job_ids can be specified.
+
+    4. If the job fails or fetching the logs fails, the command will exit with
+    a non-zero return code.
     """
     if sync_down and status:
         raise click.UsageError(
@@ -2174,17 +2187,18 @@ def logs(
         # it will return {None: None}.
         if job_id is None:
             click.secho(f'No job found on cluster {cluster!r}.', fg='red')
-            sys.exit(1)
+            sys.exit(exceptions.JobExitCode.NOT_FOUND)
         job_status = list(job_statuses.values())[0]
         job_status_str = job_status.value if job_status is not None else 'None'
         click.echo(f'Job {job_id}: {job_status_str}')
         if job_status == job_lib.JobStatus.SUCCEEDED:
             return
         else:
+            returncode = exceptions.JobExitCode.from_job_status(job_status)
             if job_status is None:
                 id_str = '' if job_id is None else f'{job_id} '
                 click.secho(f'Job {id_str}not found', fg='red')
-            sys.exit(1)
+            sys.exit(returncode)
 
     job_str = f'job {job_id}'
     if job_id is None:
@@ -2194,7 +2208,8 @@ def logs(
                 f'{colorama.Style.RESET_ALL}')
 
     # Stream logs from the server.
-    sdk.tail_logs(cluster, job_id, follow, tail=tail)
+    returncode = sdk.tail_logs(cluster, job_id, follow, tail=tail)
+    sys.exit(returncode)
 
 
 @cli.command()
@@ -3893,10 +3908,11 @@ def jobs_launch(
                                         'sky.jobs.launch')
     if not async_call and not detach_run:
         job_id = job_id_handle[0]
-        managed_jobs.tail_logs(name=None,
-                               job_id=job_id,
-                               follow=True,
-                               controller=False)
+        returncode = managed_jobs.tail_logs(name=None,
+                                            job_id=job_id,
+                                            follow=True,
+                                            controller=False)
+        sys.exit(returncode)
 
 
 @jobs.command('queue', cls=_DocumentedCodeCommand)
@@ -4127,11 +4143,12 @@ def jobs_logs(name: Optional[str], job_id: Optional[int], follow: bool,
                 logger.info(f'{fore.CYAN}Job {job} logs{controller_str}: '
                             f'{log_local_path}{style.RESET_ALL}')
         else:
-            managed_jobs.tail_logs(name=name,
-                                   job_id=job_id,
-                                   follow=follow,
-                                   controller=controller,
-                                   refresh=refresh)
+            returncode = managed_jobs.tail_logs(name=name,
+                                                job_id=job_id,
+                                                follow=follow,
+                                                controller=controller,
+                                                refresh=refresh)
+            sys.exit(returncode)
     except exceptions.ClusterNotUpError:
         with ux_utils.print_exception_no_traceback():
             raise

sky/client/sdk.py

@@ -583,7 +583,7 @@ def tail_logs(cluster_name: str,
               job_id: Optional[int],
               follow: bool,
               tail: int = 0,
-              output_stream: Optional['io.TextIOBase'] = None) -> None:
+              output_stream: Optional['io.TextIOBase'] = None) -> int:
     """Tails the logs of a job.
 
     Args:
@@ -596,7 +596,9 @@ def tail_logs(cluster_name: str,
             console.
 
     Returns:
-        None
+        Exit code based on success or failure of the job. 0 if success,
+        100 if the job failed. See exceptions.JobExitCode for possible exit
+        codes.
 
     Request Raises:
         ValueError: if arguments are invalid or the cluster is not supported.
@@ -622,7 +624,7 @@ def tail_logs(cluster_name: str,
         timeout=(client_common.API_SERVER_REQUEST_CONNECTION_TIMEOUT_SECONDS,
                  None))
     request_id = server_common.get_request_id(response)
-    stream_response(request_id, response, output_stream)
+    return stream_response(request_id, response, output_stream)
 
 
 @usage_lib.entrypoint

sky/core.py

@@ -826,7 +826,7 @@ def cancel(
 def tail_logs(cluster_name: str,
               job_id: Optional[int],
               follow: bool = True,
-              tail: int = 0) -> None:
+              tail: int = 0) -> int:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Tails the logs of a job.
 
@@ -842,6 +842,12 @@ def tail_logs(cluster_name: str,
           not the same as the user who created the cluster.
         sky.exceptions.CloudUserIdentityError: if we fail to get the current
           user identity.
+
+    Returns:
+        Return code based on success or failure of the job. 0 if success,
+          100 if the job failed. Note: This is not the return code of the job
+          script.
+
     """
     # Check the status of the cluster.
     handle = backend_utils.check_cluster_available(
@@ -851,7 +857,7 @@ def tail_logs(cluster_name: str,
     backend = backend_utils.get_backend_from_handle(handle)
 
     usage_lib.record_cluster_name_for_current_operation(cluster_name)
-    backend.tail_logs(handle, job_id, follow=follow, tail=tail)
+    return backend.tail_logs(handle, job_id, follow=follow, tail=tail)
 
 
 @usage_lib.entrypoint

sky/exceptions.py

@@ -9,7 +9,9 @@
 from sky.utils import env_options
 
 if typing.TYPE_CHECKING:
+    from sky import jobs as managed_jobs
     from sky.backends import backend
+    from sky.skylet import job_lib
     from sky.utils import status_lib
 
 # Return code for keyboard interruption and SIGTSTP
@@ -236,7 +238,7 @@ def __init__(self, returncode: int, command: str, error_msg: str,
         else:
             if (len(command) > 100 and
                     not env_options.Options.SHOW_DEBUG_INFO.get()):
-                # Chunck the command to avoid overflow.
+                # Chunk the command to avoid overflow.
                 command = command[:100] + '...'
             message = (f'Command {command} failed with return code '
                        f'{returncode}.\n{error_msg}')
@@ -449,3 +451,80 @@ def __init__(self, server_url: str):
             f'Could not connect to SkyPilot API server at {server_url}. '
             f'Please ensure that the server is running. '
             f'Try: curl {server_url}/api/health')
+
+
+class JobExitCode(enum.IntEnum):
+    """Job exit code enum.
+
+    These codes are used as return codes for job-related operations and as
+    process exit codes to indicate job status.
+    """
+
+    SUCCEEDED = 0
+    """The job completed successfully"""
+
+    FAILED = 100
+    """The job failed (due to user code, setup, or driver failure)"""
+
+    NOT_FINISHED = 101
+    """The job has not finished yet"""
+
+    NOT_FOUND = 102
+    """The job was not found"""
+
+    CANCELLED = 103
+    """The job was cancelled by the user"""
+
+    @classmethod
+    def from_job_status(cls,
+                        status: Optional['job_lib.JobStatus']) -> 'JobExitCode':
+        """Convert a job status to an exit code."""
+        # Import here to avoid circular imports
+        # pylint: disable=import-outside-toplevel
+        from sky.skylet import job_lib
+
+        if status is None:
+            return cls.NOT_FOUND
+
+        if not status.is_terminal():
+            return cls.NOT_FINISHED
+
+        if status == job_lib.JobStatus.SUCCEEDED:
+            return cls.SUCCEEDED
+
+        if status == job_lib.JobStatus.CANCELLED:
+            return cls.CANCELLED
+
+        if status in job_lib.JobStatus.user_code_failure_states(
+        ) or status == job_lib.JobStatus.FAILED_DRIVER:
+            return cls.FAILED
+
+        # Should not hit this case, but included to avoid errors
+        return cls.FAILED
+
+    @classmethod
+    def from_managed_job_status(
+            cls,
+            status: Optional['managed_jobs.ManagedJobStatus']) -> 'JobExitCode':
+        """Convert a managed job status to an exit code."""
+        # Import here to avoid circular imports
+        # pylint: disable=import-outside-toplevel
+        from sky import jobs as managed_jobs
+
+        if status is None:
+            return cls.NOT_FOUND
+
+        if not status.is_terminal():
+            return cls.NOT_FINISHED
+
+        if status == managed_jobs.ManagedJobStatus.SUCCEEDED:
+            return cls.SUCCEEDED
+
+        if status == managed_jobs.ManagedJobStatus.CANCELLED:
+            return cls.CANCELLED
+
+        if status.is_failed():
+            return cls.FAILED
+
+        # Should not hit this case, but included to avoid errors
+        return cls.FAILED

sky/jobs/client/sdk.py

@@ -184,7 +184,7 @@ def tail_logs(name: Optional[str] = None,
               follow: bool = True,
               controller: bool = False,
               refresh: bool = False,
-              output_stream: Optional['io.TextIOBase'] = None) -> None:
+              output_stream: Optional['io.TextIOBase'] = None) -> int:
     """Tails logs of managed jobs.
 
     You can provide either a job name or a job ID to tail logs. If both are not
@@ -199,6 +199,11 @@ def tail_logs(name: Optional[str] = None,
         output_stream: The stream to write the logs to. If None, print to the
             console.
 
+    Returns:
+        Exit code based on success or failure of the job. 0 if success,
+        100 if the job failed. See exceptions.JobExitCode for possible exit
+        codes.
+
     Request Raises:
         ValueError: invalid arguments.
         sky.exceptions.ClusterNotUpError: the jobs controller is not up.
@@ -217,7 +222,7 @@ def tail_logs(name: Optional[str] = None,
         timeout=(5, None),
     )
     request_id = server_common.get_request_id(response)
-    sdk.stream_response(request_id, response, output_stream)
+    return sdk.stream_response(request_id, response, output_stream)
 
 
 @usage_lib.entrypoint

sky/jobs/constants.py

@@ -40,7 +40,7 @@
 # The version of the lib files that jobs/utils use. Whenever there is an API
 # change for the jobs/utils, we need to bump this version and update
 # job.utils.ManagedJobCodeGen to handle the version update.
-MANAGED_JOBS_VERSION = 2
+MANAGED_JOBS_VERSION = 3
 
 # The command for setting up the jobs dashboard on the controller. It firstly
 # checks if the systemd services are available, and if not (e.g., Kubernetes

sky/jobs/server/core.py

@@ -460,12 +460,17 @@ def cancel(name: Optional[str] = None,
 
 @usage_lib.entrypoint
 def tail_logs(name: Optional[str], job_id: Optional[int], follow: bool,
-              controller: bool, refresh: bool) -> None:
+              controller: bool, refresh: bool) -> int:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Tail logs of managed jobs.
 
     Please refer to sky.cli.job_logs for documentation.
 
+    Returns:
+        Exit code based on success or failure of the job. 0 if success,
+        100 if the job failed. See exceptions.JobExitCode for possible exit
+        codes.
+
     Raises:
         ValueError: invalid arguments.
         sky.exceptions.ClusterNotUpError: the jobs controller is not up.
@@ -494,11 +499,11 @@ def tail_logs(name: Optional[str], job_id: Optional[int], follow: bool,
     backend = backend_utils.get_backend_from_handle(handle)
     assert isinstance(backend, backends.CloudVmRayBackend), backend
 
-    backend.tail_managed_job_logs(handle,
-                                  job_id=job_id,
-                                  job_name=name,
-                                  follow=follow,
-                                  controller=controller)
+    return backend.tail_managed_job_logs(handle,
+                                         job_id=job_id,
+                                         job_name=name,
+                                         follow=follow,
+                                         controller=controller)
 
 
 def start_dashboard_forwarding(refresh: bool = False) -> Tuple[int, int]: