Pull and store process exit status from jobs (#2615)

* Pull the exit status from the job, store in JobModel
* Display non-zero statuses in `dstack apply` and `dstack ps -v`

Closes: #2570

Author: un-def

runner/internal/executor/executor.go

@@ -181,11 +181,16 @@ func (ex *RunExecutor) Run(ctx context.Context) (err error) {
 
 		// todo fail reason?
 		log.Error(ctx, "Exec failed", "err", err)
-		ex.SetJobState(ctx, types.JobStateFailed)
+		var exitError *exec.ExitError
+		if errors.As(err, &exitError) {
+			ex.SetJobStateWithExitStatus(ctx, types.JobStateFailed, exitError.ExitCode())
+		} else {
+			ex.SetJobState(ctx, types.JobStateFailed)
+		}
 		return gerrors.Wrap(err)
 	}
 
-	ex.SetJobState(ctx, types.JobStateDone)
+	ex.SetJobStateWithExitStatus(ctx, types.JobStateDone, 0)
 	return nil
 }
 
@@ -224,6 +229,22 @@ func (ex *RunExecutor) SetJobStateWithTerminationReason(
 	log.Info(ctx, "Job state changed", "new", state)
 }
 
+func (ex *RunExecutor) SetJobStateWithExitStatus(
+	ctx context.Context, state types.JobState, exitStatus int,
+) {
+	ex.mu.Lock()
+	ex.jobStateHistory = append(
+		ex.jobStateHistory,
+		schemas.JobStateEvent{
+			State:      state,
+			Timestamp:  ex.timestamp.Next(),
+			ExitStatus: &exitStatus,
+		},
+	)
+	ex.mu.Unlock()
+	log.Info(ctx, "Job state changed", "new", state)
+}
+
 func (ex *RunExecutor) SetRunnerState(state string) {
 	ex.state = state
 }

runner/internal/executor/executor_test.go

@@ -54,10 +54,15 @@ func TestExecutor_HomeDir(t *testing.T) {
 
 func TestExecutor_NonZeroExit(t *testing.T) {
 	ex := makeTestExecutor(t)
-	ex.jobSpec.Commands = append(ex.jobSpec.Commands, "ehco 1") // note: intentional misspelling
+	ex.jobSpec.Commands = append(ex.jobSpec.Commands, "exit 100")
+	makeCodeTar(t, ex.codePath)
 
-	err := ex.execJob(context.TODO(), io.Discard)
+	err := ex.Run(context.TODO())
 	assert.Error(t, err)
+	assert.NotEmpty(t, ex.jobStateHistory)
+	exitStatus := ex.jobStateHistory[len(ex.jobStateHistory)-1].ExitStatus
+	assert.NotNil(t, exitStatus, ex.jobStateHistory)
+	assert.Equal(t, 100, *exitStatus)
 }
 
 func TestExecutor_SSHCredentials(t *testing.T) {

runner/internal/schemas/schemas.go

@@ -11,6 +11,7 @@ type JobStateEvent struct {
 	Timestamp          int64                   `json:"timestamp"`
 	TerminationReason  types.TerminationReason `json:"termination_reason"`
 	TerminationMessage string                  `json:"termination_message"`
+	ExitStatus         *int                    `json:"exit_status"`
 }
 
 type LogEvent struct {

src/dstack/_internal/cli/services/configurators/run.py

@@ -560,7 +560,9 @@ def print_finished_message(run: Run):
         console.print("[code]Done[/]")
         return
 
-    termination_reason, termination_reason_message = _get_run_termination_reason(run)
+    termination_reason, termination_reason_message, exit_status = (
+        _get_run_termination_reason_and_exit_status(run)
+    )
     message = "Run failed due to unknown reason. Check CLI, server, and run logs."
     if run.status == RunStatus.TERMINATED:
         message = "Run terminated due to unknown reason. Check CLI, server, and run logs."
@@ -572,13 +574,15 @@ def print_finished_message(run: Run):
             "Check CLI and server logs for more details."
         )
     elif termination_reason is not None:
+        exit_status_details = f"Exit status: {exit_status}.\n" if exit_status else ""
         error_details = (
             f"Error: {termination_reason_message}\n" if termination_reason_message else ""
         )
         message = (
             f"Run failed with error code {termination_reason.name}.\n"
+            f"{exit_status_details}"
             f"{error_details}"
-            "Check CLI, server, and run logs for more details."
+            f"Check [bold]dstack logs -d {run.name}[/bold] for more details."
         )
     console.print(f"[error]{message}[/]")
 
@@ -589,14 +593,20 @@ def get_run_exit_code(run: Run) -> int:
     return 1
 
 
-def _get_run_termination_reason(run: Run) -> Tuple[Optional[JobTerminationReason], Optional[str]]:
+def _get_run_termination_reason_and_exit_status(
+    run: Run,
+) -> Tuple[Optional[JobTerminationReason], Optional[str], Optional[int]]:
     if len(run._run.jobs) == 0:
-        return None, None
+        return None, None, None
     job = run._run.jobs[0]
     if len(job.job_submissions) == 0:
-        return None, None
+        return None, None, None
     job_submission = job.job_submissions[0]
-    return job_submission.termination_reason, job_submission.termination_reason_message
+    return (
+        job_submission.termination_reason,
+        job_submission.termination_reason_message,
+        job_submission.exit_status,
+    )
 
 
 def _run_resubmitted(run: Run, current_job_submission: Optional[JobSubmission]) -> bool:

src/dstack/_internal/cli/utils/run.py

@@ -218,6 +218,11 @@ def _get_run_error(run: Run) -> str:
 
 
 def _get_job_error(job: Job) -> str:
-    if job.job_submissions[-1].termination_reason is None:
+    job_submission = job.job_submissions[-1]
+    termination_reason = job_submission.termination_reason
+    exit_status = job_submission.exit_status
+    if termination_reason is None:
         return ""
-    return job.job_submissions[-1].termination_reason.name
+    if exit_status:
+        return f"{termination_reason.name} {exit_status}"
+    return termination_reason.name

src/dstack/_internal/core/models/runs.py

@@ -283,6 +283,7 @@ class JobSubmission(CoreModel):
     status: JobStatus
     termination_reason: Optional[JobTerminationReason]
     termination_reason_message: Optional[str]
+    exit_status: Optional[int]
     job_provisioning_data: Optional[JobProvisioningData]
     job_runtime_data: Optional[JobRuntimeData]
 
@@ -508,21 +509,30 @@ def _get_run_error(
         return ""
     if len(run_jobs) > 1:
         return run_termination_reason.name
-    run_job_termination_reason = _get_run_job_termination_reason(run_jobs)
+    run_job_termination_reason, exit_status = _get_run_job_termination_reason_and_exit_status(
+        run_jobs
+    )
     # For failed runs, also show termination reason to provide more context.
     # For other run statuses, the job termination reason will duplicate run status.
     if run_job_termination_reason is not None and run_termination_reason in [
         RunTerminationReason.JOB_FAILED,
         RunTerminationReason.SERVER_ERROR,
         RunTerminationReason.RETRY_LIMIT_EXCEEDED,
     ]:
+        if exit_status:
+            return (
+                f"{run_termination_reason.name}\n({run_job_termination_reason.name} {exit_status})"
+            )
         return f"{run_termination_reason.name}\n({run_job_termination_reason.name})"
     return run_termination_reason.name
 
 
-def _get_run_job_termination_reason(run_jobs: List[Job]) -> Optional[JobTerminationReason]:
+def _get_run_job_termination_reason_and_exit_status(
+    run_jobs: List[Job],
+) -> tuple[Optional[JobTerminationReason], Optional[int]]:
     for job in run_jobs:
         if len(job.job_submissions) > 0:
-            if job.job_submissions[-1].termination_reason is not None:
-                return job.job_submissions[-1].termination_reason
-    return None
+            job_submission = job.job_submissions[-1]
+            if job_submission.termination_reason is not None:
+                return job_submission.termination_reason, job_submission.exit_status
+    return None, None

src/dstack/_internal/server/background/tasks/process_running_jobs.py

@@ -651,6 +651,10 @@ def _process_running(
                 )
             if latest_state_event.termination_message:
                 job_model.termination_reason_message = latest_state_event.termination_message
+        if (exit_status := latest_state_event.exit_status) is not None:
+            job_model.exit_status = exit_status
+            if exit_status != 0:
+                logger.info("%s: non-zero exit status %s", fmt(job_model), exit_status)
     else:
         _terminate_if_inactivity_duration_exceeded(run_model, job_model, resp.no_connections_secs)
     if job_model.status != previous_status:

src/dstack/_internal/server/migrations/versions/6c1a9d6530ee_add_jobmodel_exit_status.py

@@ -0,0 +1,26 @@
+"""Add JobModel.exit_status
+
+Revision ID: 6c1a9d6530ee
+Revises: 7ba3b59d7ca6
+Create Date: 2025-05-09 10:25:19.715852
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "6c1a9d6530ee"
+down_revision = "7ba3b59d7ca6"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    with op.batch_alter_table("jobs", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("exit_status", sa.Integer(), nullable=True))
+
+
+def downgrade() -> None:
+    with op.batch_alter_table("jobs", schema=None) as batch_op:
+        batch_op.drop_column("exit_status")

src/dstack/_internal/server/models.py

@@ -382,6 +382,7 @@ class JobModel(BaseModel):
         Enum(JobTerminationReason)
     )
     termination_reason_message: Mapped[Optional[str]] = mapped_column(Text)
+    exit_status: Mapped[Optional[int]] = mapped_column(Integer)
     job_spec_data: Mapped[str] = mapped_column(Text)
     job_provisioning_data: Mapped[Optional[str]] = mapped_column(Text)
     runner_timestamp: Mapped[Optional[int]] = mapped_column(BigInteger)

src/dstack/_internal/server/schemas/runner.py

@@ -16,6 +16,7 @@ class JobStateEvent(CoreModel):
     state: JobStatus
     termination_reason: Optional[str] = None
     termination_message: Optional[str] = None
+    exit_status: Optional[int] = None
 
 
 class LogEvent(CoreModel):