wait command returns prematurely for array jobs (parse_status only inspects first sacct line)

[ultimatile]

Summary

hpc wait <array_job_id> returns long before the array job actually finishes. For a Slurm array job, parse_status reads sacct -j <id> --format=State --noheader and returns the status from only the first output line, ignoring every other array task.

Affected code

src/hpc/scheduler.py (Slurm branch):

def status_cmd(self, job_id: str) -> list[str]:
    return ["sacct", "-j", job_id, "--format=State", "--noheader"]

def parse_status(self, output: str) -> JobStatus:
    lines = output.strip().splitlines()
    status_str = lines[0].strip().rstrip("+") if lines else ""
    return _STATUS_MAP.get(status_str, JobStatus.FAILED)

wait_for_job in src/hpc/job.py checks status in terminal_states and returns as soon as parse_status reports a terminal state.

Repro

For an array job with mixed state (some tasks COMPLETED, others still RUNNING/PENDING):

$ sacct -j 7701965 --format=State --noheader -X
 COMPLETED 
 COMPLETED 
 COMPLETED 
 ...
 RUNNING
 RUNNING
 PENDING

parse_status reads only the first line (COMPLETED), so wait_for_job returns JobStatus.COMPLETED while ~150 array tasks are still pending or running.

Expected

hpc wait should block until all array tasks reach a terminal state.

Suggested fix

Aggregate the statuses over all sacct rows in Slurm.parse_status (and analogously for PJM.parse_status). Sketch:

def parse_status(self, output: str) -> JobStatus:
    lines = [ln.strip().rstrip("+") for ln in output.strip().splitlines() if ln.strip()]
    if not lines:
        return JobStatus.FAILED
    statuses = [_STATUS_MAP.get(s, JobStatus.FAILED) for s in lines]
    if any(s == JobStatus.PENDING for s in statuses):
        return JobStatus.PENDING
    if any(s == JobStatus.RUNNING for s in statuses):
        return JobStatus.RUNNING
    for terminal in (JobStatus.FAILED, JobStatus.CANCELLED, JobStatus.TIMEOUT):
        if any(s == terminal for s in statuses):
            return terminal
    return JobStatus.COMPLETED

Notes:

• Adding -X to status_cmd is helpful too — it suppresses jobsteps so each row is one array task and the aggregation is over a clean set.
• For non-array jobs, sacct returns one row anyway, so the aggregated logic still works.
• The PJM branch has the same shape (lines[1].strip()) and likely the same issue for PJM step jobs; worth a parallel fix.

Workaround

Until this is fixed, fall back to a manual loop:

hpc exec "until [ \$(squeue -j <id> -h 2>/dev/null | wc -l) -eq 0 ]; do sleep 60; done"

Environment

• hpc 0.4.0
• Scheduler: Slurm
• Array job submitted with [slurm.options].array = "1-N%K"

[ultimatile]

Implementation plan

After investigation, plan to fix this in three coordinated changes in src/hpc/scheduler.py (Slurm only — see deferred section for PJM).

Changes

1. Slurm.status_cmd — add -X flag

def status_cmd(self, job_id: str) -> list[str]:
    return ["sacct", "-j", job_id, "--format=State", "--noheader", "-X"]

-X suppresses jobsteps (.batch, .extern) so each row is exactly one allocation. For array jobs, this gives one row per task; for non-array jobs, behavior is unchanged (still one row).

2. Slurm.parse_status — aggregate over all rows

def parse_status(self, output: str) -> JobStatus:
    lines = [
        ln.strip().rstrip("+")
        for ln in output.strip().splitlines()
        if ln.strip()
    ]
    if not lines:
        return JobStatus.FAILED
    statuses = [_STATUS_MAP.get(s, JobStatus.FAILED) for s in lines]
    for status in (
        JobStatus.RUNNING,
        JobStatus.PENDING,
        JobStatus.FAILED,
        JobStatus.CANCELLED,
        JobStatus.TIMEOUT,
    ):
        if status in statuses:
            return status
    return JobStatus.COMPLETED

Priority: RUNNING > PENDING > FAILED > CANCELLED > TIMEOUT > COMPLETED. Non-terminal states win, ensuring wait_for_job keeps polling until every task is terminal. Terminal-failure states win over COMPLETED so the array's worst-case outcome is reported.

rstrip("+") is applied to every row (handles CANCELLED+, COMPLETED+ etc. uniformly). Single-row output collapses to identity (aggregate({s}) = s), preserving backward compat with mocked single-line tests.

3. Expand _STATUS_MAP to cover known Slurm states

_STATUS_MAP = {
    # Non-terminal — pending/queued
    "PENDING": JobStatus.PENDING,
    "CONFIGURING": JobStatus.PENDING,
    "REQUEUED": JobStatus.PENDING,
    # Non-terminal — running/active
    "RUNNING": JobStatus.RUNNING,
    "RESIZING": JobStatus.RUNNING,
    "SUSPENDED": JobStatus.RUNNING,
    "COMPLETING": JobStatus.RUNNING,
    # Terminal — success
    "COMPLETED": JobStatus.COMPLETED,
    # Terminal — failure
    "FAILED": JobStatus.FAILED,
    "BOOT_FAIL": JobStatus.FAILED,
    "NODE_FAIL": JobStatus.FAILED,
    "OUT_OF_MEMORY": JobStatus.FAILED,
    # Terminal — cancelled
    "CANCELLED": JobStatus.CANCELLED,
    "PREEMPTED": JobStatus.CANCELLED,
    "REVOKED": JobStatus.CANCELLED,
    # Terminal — timeout
    "TIMEOUT": JobStatus.TIMEOUT,
    "DEADLINE": JobStatus.TIMEOUT,
}

Without this expansion, the aggregation amplifies a pre-existing bug: any task in a state like REQUEUED or SUSPENDED would fall back to FAILED and terminate the wait prematurely. Unknown states still fall back to FAILED (conservative for forward-compat with future Slurm versions).

Tests (tests/test_scheduler.py)

Add TestSlurmParseStatus:

• Single-line per status (COMPLETED, PENDING, FAILED, REQUEUED, SUSPENDED, BOOT_FAIL, OUT_OF_MEMORY, DEADLINE, PREEMPTED)
• Multi-line all COMPLETED → COMPLETED
• Multi-line one PENDING others COMPLETED → PENDING
• Multi-line one RUNNING others PENDING/COMPLETED → RUNNING (priority)
• Multi-line one REQUEUED others COMPLETED → PENDING
• All terminal one FAILED others COMPLETED → FAILED
• Mixed FAILED/CANCELLED/COMPLETED → FAILED
• COMPLETED+ suffix handled by rstrip("+")
• Empty/whitespace input → FAILED
• Unknown status → FAILED (fallback preserved)

Add TestSlurmStatusCmd:

• status_cmd includes -X

Impact

• JobManager.get_job_status calls parse_status — transparent change, returns JobStatus.
• JobManager.wait_for_job checks status in terminal_states — no code change needed; the new aggregation guarantees a non-terminal aggregate as long as any task is non-terminal.
• JobManager.get_job_output uses JobStatus.PENDING/RUNNING to render "not yet available" — non-terminal aggregate preserves current UX.
• Existing mock tests pass single-line outputs — aggregation collapses to identity for n=1.
• Side-effect: non-array Slurm jobs also benefit, since REQUEUED/SUSPENDED etc. are now correctly non-terminal.

Invariant

aggregate(S) ∈ terminal_states ⟺ ∀ s_i ∈ S, s_i ∈ terminal_states where terminal_states = {COMPLETED, FAILED, CANCELLED, TIMEOUT}. Single-element S satisfies aggregate({s}) = s.

Inconclusive / Deferred items

• deferred — PJM.parse_status likely has the same fixed-index bug
  reason: requires verification of pjstat --choose st <id> output format for array/step jobs in a real PJM environment, which is not currently accessible.
  resolution-point: file a follow-up issue after this fix lands; address once PJM environment access is available.

• deferred — empty sacct output → FAILED (pre-existing)
  reason: not a regression introduced by this change; wait_for_job may exit prematurely if sacct returns empty output for a freshly submitted job not yet visible in accounting.
  resolution-point: separate issue if observed in practice; would likely involve treating empty output as transient and retrying.