#!/usr/bin/env python3
#
# This script is a system-wide prediff for use with slurm-based
# launchers.
#
import sys, re

# These match the message blocks we want to remove.
msgs = (
b"""srun: error: .+: task [0-9]+: Exited with exit code [0-9]+
""",
b"""srun: error: .+: task [0-9]+: (Killed|Terminated)
""",
b"""srun: error: .+: tasks [0-9]+-[0-9]+: (Killed|Terminated)
""",
b"""srun: .*(Force Terminated|Terminating) (job step |StepId=)[0-9.]+
""",
b"""srun: Job step aborted: Waiting up to [0-9]+ seconds for job step .*
""",
rb"""slurmstepd: error: \*\*\* STEP [0-9.]+ ON .+ CANCELLED AT [-0-9T.:]+ \*\*\*
""",
rb"""\[[-0-9T.:]+] error: \*\*\* STEP [0-9.]+ ON .+ CANCELLED AT [-0-9T.:]+ DUE TO TASK FAILURE \*\*\*
""",
rb"""slurmstepd: error: .+ \[[0-9]+] .*
""",
b"""srun: error: spank: /opt/cray/.*: Plugin file not found
""",
b"""cp: cannot create regular file .*/[.]module/PrgEnv-.*: File exists
""",
b"""cp: failed to close .*/[.]module/PrgEnv-.*: Stale file handle
""",
b"""diff: .*/[.]module/PrgEnv-.*: No such file or directory
""",
)

outfname = sys.argv[2]
with open(outfname, "rb") as f:
    outText = f.read()
for m in msgs:
    outText = re.sub(m, b"", outText, flags = re.MULTILINE)
with open(outfname, "wb") as f:
    f.write(outText)
