lstmcpipe : script to repair a workflow with failed (e.g. TIMEOUT) jobs

lstmcpipe_repair_workflow.py

#!/usr/bin/env python3
"""
Repair a broken lstmcpipe production.

Input: a command-yaml where each stage maps job_id -> full sbatch command
(the kind produced alongside logs_reduced_*.yml, containing --dependency=
afterok:... chains and, for r0_to_dl1, --array=... chains). NOT the
description-yaml (job_id -> free text).

For every job with at least one FAILED/TIMEOUT/CANCELLED/etc task:
  - resubmit ONLY the failed array indices, not the whole array
    (verified: task index maps positionally to the sublist file list in
    the wrap command, so re-running index N alone reprocesses exactly
    what index N processed before -- nothing else)
  - for non-array jobs, resubmit the whole thing (no array to slice)
  - rewrite --dependency=afterok:... swapping any old id that was itself
    resubmitted for its new id
  - a job that's still RUNNING/PENDING is left alone, not touched,
    not resubmitted
  - a job that's COMPLETED but depends on something resubmitted is
    flagged as possibly stale, not touched -- verify by hand

Only touches stages matching STAGE_PATTERNS. check_full_workflow (or any
stage whose keys aren't job ids) is skipped -- rerun that check by hand
once this finishes.
"""

import argparse
import re
import subprocess
import sys
import yaml
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed

# Substrings, matched case-insensitively against actual stage keys in the
# yaml. Order here is the DAG order (upstream first) -- it controls the
# order stages are repaired in, which matters for dependency propagation.
STAGE_PATTERNS = ["r0_dl1", "r0_to_dl1", "merge_dl1", "merge_and_copy_dl1",
                   "train_pipe", "plot_rf_feat"]

TERMINAL_BAD_STATES = {"FAILED", "CANCELLED", "TIMEOUT", "OUT_OF_MEMORY",
                        "NODE_FAIL", "PREEMPTED", "BOOT_FAIL", "DEADLINE",
                        "REVOKED", "SPECIAL_EXIT"}
IN_PROGRESS_STATES = {"PENDING", "RUNNING", "COMPLETING", "REQUEUED",
                       "RESIZING", "SUSPENDED", "CONFIGURING"}

DEP_RE = re.compile(r"--dependency=afterok:([\d:]+)")
ARRAY_RE = re.compile(r"--array=(\S+)")


def resolve_stages(data):
    """Match yaml stage keys against STAGE_PATTERNS, preserving DAG order.
    A stage key must contain one of the patterns (case-insensitive)."""
    matched = []
    for pattern in STAGE_PATTERNS:
        for key in data:
            if pattern.lower() in key.lower() and key not in matched:
                matched.append(key)
    return matched


def batched(seq, n=200):
    seq = list(seq)
    for i in range(0, len(seq), n):
        yield seq[i:i + n]


def sacct_states(job_ids, start_date, batch_size=400, workers=6):
    """Parallel batched sacct calls.
    Returns {base_job_id: {task_id: state}}. task_id is None for
    non-array jobs, or the array index string ("0", "1", ...) for array
    jobs. -X drops per-step rows (.batch/.extern), keeping one row per
    job or per array task. -S is mandatory: sacct defaults to showing
    only today's jobs, so without it, older jobs silently vanish and
    look identical to a job that was never found."""
    per_task = defaultdict(dict)
    chunks = list(batched(job_ids, batch_size))
    print(f"  {len(chunks)} sacct calls, {workers} in parallel")

    def query(chunk):
        cmd = ["sacct", "-j", ",".join(chunk), "-S", start_date, "-X",
               "--format=JobID,State", "--noheader", "--parsable2"]
        return subprocess.run(cmd, capture_output=True, text=True)

    with ThreadPoolExecutor(max_workers=workers) as pool:
        futures = [pool.submit(query, c) for c in chunks]
        for i, fut in enumerate(as_completed(futures), 1):
            result = fut.result()
            print(f"  batch {i}/{len(chunks)} done")
            if result.returncode != 0:
                print(f"WARNING: sacct failed for a batch: {result.stderr.strip()}", file=sys.stderr)
                continue
            for line in result.stdout.strip().splitlines():
                if "|" not in line:
                    continue
                jid_field, state = line.split("|", 1)
                state = state.split()[0]  # "CANCELLED by 12345" -> "CANCELLED"
                if "_" in jid_field:
                    base_id, task_id = jid_field.split("_", 1)
                else:
                    base_id, task_id = jid_field, None
                per_task[base_id][task_id] = state
    return per_task


def classify(task_states):
    """(status, failed_tasks) for one job's {task_id: state} dict.
    status in {"COMPLETED", "IN_PROGRESS", "FAILED", "UNKNOWN"}.
    failed_tasks is a list of task_id (None for non-array jobs, or
    array index strings) that need resubmitting."""
    states = set(task_states.values())
    if states == {"COMPLETED"}:
        return "COMPLETED", []
    failed = [t for t, s in task_states.items() if s in TERMINAL_BAD_STATES]
    if failed:
        return "FAILED", failed
    in_progress = [t for t, s in task_states.items() if s in IN_PROGRESS_STATES]
    if in_progress:
        return "IN_PROGRESS", []
    return "UNKNOWN", []


def dependency_ids(command):
    m = DEP_RE.search(command)
    return m.group(1).split(":") if m else []


def rewrite_dependencies(command, id_map):
    m = DEP_RE.search(command)
    if not m:
        return command
    old_ids = m.group(1).split(":")
    new_ids = [id_map.get(i, i) for i in old_ids]
    return command[:m.start()] + "--dependency=afterok:" + ":".join(new_ids) + command[m.end():]


def rewrite_array(command, failed_tasks):
    """Restrict --array=... to just the failed indices, keeping any
    %throttle suffix. No-op if this command has no --array (not an
    array job) or failed_tasks is [None] (non-array failure marker)."""
    if failed_tasks == [None]:
        return command
    m = ARRAY_RE.search(command)
    if not m:
        return command
    throttle = ""
    if "%" in m.group(1):
        throttle = "%" + m.group(1).split("%", 1)[1]
    new_spec = ",".join(sorted(failed_tasks, key=int)) + throttle
    return command[:m.start()] + f"--array={new_spec}" + command[m.end():]


def submit(command):
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"sbatch failed: {result.stderr.strip()}")
    return result.stdout.strip().split(";")[0].strip()


def main():
    ap = argparse.ArgumentParser(description="Find incomplete jobs, resubmit only what failed, with fixed dependencies.")
    ap.add_argument("filename", help="command-yaml (job_id -> sbatch command)")
    ap.add_argument("--dry-run", action="store_true", help="plan only, no sbatch calls")
    ap.add_argument("--output", default=None, help="path for repaired yaml")
    ap.add_argument("--start-date", default="2026-01-01",
                     help="sacct -S value. Must predate your earliest job in this "
                          "yaml, or sacct silently drops it. Default 2026-01-01 -- "
                          "override if this production started before that.")
    ap.add_argument("--sacct-batch-size", type=int, default=400,
                     help="job ids per sacct call")
    ap.add_argument("--sacct-workers", type=int, default=6,
                     help="parallel sacct calls; lower this if slurmdbd complains")
    args = ap.parse_args()

    with open(args.filename) as f:
        data = yaml.safe_load(f)

    stages = resolve_stages(data)
    skipped = [s for s in data if s not in stages]
    if skipped:
        print(f"Skipping unmatched stages: {skipped}")
    print(f"Matched stages, in repair order: {stages}")

    all_ids = []
    for stage in stages:
        for job_id in data[stage]:
            if not str(job_id).isdigit():
                print(f"WARNING: non-numeric key '{job_id}' in stage {stage}, skipping")
                continue
            all_ids.append(str(job_id))

    print(f"Checking {len(all_ids)} jobs via sacct")
    states = sacct_states(all_ids, args.start_date,
                           batch_size=args.sacct_batch_size,
                           workers=args.sacct_workers)
    unknown_ids = [j for j in all_ids if j not in states]
    if unknown_ids:
        print(f"WARNING: {len(unknown_ids)} job ids not found in sacct "
              f"(purged, typo, or predates --start-date). Flagged for manual review, not touched.")

    id_map = {}
    new_data = {s: {} for s in stages}
    n_resubmitted = 0
    n_stale = 0
    n_in_progress = 0
    n_manual_review = 0

    for stage in stages:
        stage_bad = 0
        for job_id, command in data[stage].items():
            job_id = str(job_id)
            if not job_id.isdigit():
                continue

            task_states = states.get(job_id)
            deps = dependency_ids(command)
            deps_changed = any(d in id_map for d in deps)

            if task_states is None:
                print(f"MANUAL REVIEW: job {job_id} ({stage}) not found in sacct, leaving untouched")
                new_data[stage][job_id] = command
                n_manual_review += 1
                continue

            status, failed_tasks = classify(task_states)

            if status == "COMPLETED" and not deps_changed:
                new_data[stage][job_id] = command
                continue

            if status == "COMPLETED" and deps_changed:
                print(f"NOTE: job {job_id} ({stage}) COMPLETED but a dependency "
                      f"was resubmitted. Output may be stale, not touching it. "
                      f"Verify manually.")
                new_data[stage][job_id] = command
                n_stale += 1
                continue

            if status == "IN_PROGRESS":
                print(f"job {job_id} ({stage}) still PENDING/RUNNING, leaving alone")
                new_data[stage][job_id] = command
                n_in_progress += 1
                continue

            if status == "UNKNOWN":
                print(f"MANUAL REVIEW: job {job_id} ({stage}) states {set(task_states.values())} "
                      f"don't classify cleanly, leaving untouched")
                new_data[stage][job_id] = command
                n_manual_review += 1
                continue

            # status == "FAILED"
            stage_bad += 1
            fixed_command = rewrite_dependencies(command, id_map)
            fixed_command = rewrite_array(fixed_command, failed_tasks)

            if failed_tasks == [None]:
                scope = "whole job"
            else:
                scope = f"{len(failed_tasks)}/{len(task_states)} array tasks (idx {','.join(sorted(failed_tasks, key=int))})"
            print(f"job {job_id} stage {stage}: {scope} failed -> resubmitting")

            if args.dry_run:
                new_id = f"DRYRUN-{job_id}"
            else:
                new_id = submit(fixed_command)

            id_map[job_id] = new_id
            new_data[stage][new_id] = fixed_command
            n_resubmitted += 1

        if stage_bad == 0:
            print(f"STAGE {stage}: complete ({len(data[stage])} jobs)")
        else:
            print(f"STAGE {stage}: {stage_bad} jobs resubmitted")

    print(f"\nTotal resubmitted: {n_resubmitted}")
    if n_in_progress:
        print(f"Still running, left alone: {n_in_progress}")
    if n_stale:
        print(f"Flagged stale (completed but upstream changed): {n_stale}")
    if n_manual_review:
        print(f"Flagged for manual review (unknown/not found): {n_manual_review}")

    out_path = args.output or re.sub(r"\.ya?ml$", "_repaired.yml", args.filename)
    with open(out_path, "w") as f:
        yaml.dump(new_data, f, default_flow_style=False)
    print(f"Repaired yaml: {out_path}")


if __name__ == "__main__":
    main()