save a Slurm job allocation (nodes and cpus) into a text file

slurm_nodefile.py

#!/usr/bin/env python

import os, sys, argparse

def main():
    args = parse_command_line()
    try:
        nodesCompressed = os.environ.get("SLURM_JOB_NODELIST").split(",")
    except:
        print("ERROR: environment variable SLURM_JOB_NODELIST does not exist")
        sys.exit(-1)

    try:
        tasksCompressed = os.environ["SLURM_TASKS_PER_NODE"].split(",")
    except:
        print("ERROR: environment variable SLURM_TASKS_PER_NODE does not exist")
        sys.exit(-1)

    nodes = []
    for ni, n in enumerate(nodesCompressed):
        pos = n.find("[")
        if pos == -1:
            nodes.append(n)
        else:
            prefix = n[:pos]
            n1, n2 = n[pos + 1: n.find("]")].split("-")
            n1 = int(n1)
            n2 = int(n2)
            for nj in range(n1, n2 + 1):
                nodes.append(prefix + str(nj))

    tasks = [0] * len(nodes)
    index = 0
    for ti, t in enumerate(tasksCompressed):
        pos = t.find("(x")
        if pos == -1:
            tasks[index] = int(t)
            index += 1
        else:
            n = int(t[: pos])
            repeat = int(t[pos + 2: t.find(")")])
            for ri in range(repeat):
                tasks[index] = n
                index += 1

    with open(args.outputFile, "wt") as fp:
        for ni, node in enumerate(nodes):
            for ti in range(tasks[ni]):
                fp.write("%s\n" % (node))

def parse_command_line():
    description = "save a Slurm job allocation (nodes and cpus) into a text file"
    parser = argparse.ArgumentParser(description=description)

    parser.add_argument('outputFile', metavar="<output filename>", help='save the node/cpu list to this file')

    args = parser.parse_args()

    return args

if __name__ == "__main__":
    main()