Law slurm submission test

slurm_test.py

# coding: utf-8


import time

import law


def test_htcondor():
    # load the htcondor contrib package, making all exported objects available under law.htcondor
    law.contrib.load("htcondor")

    # create a job manager and job file factory
    manager = law.htcondor.HTCondorJobManager()
    factory = law.htcondor.HTCondorJobFileFactory()

    # htcondor expects an executable file, so create one that just does a simple directory listing
    executable = law.LocalFileTarget(is_tmp=True)
    executable.touch("ls")
    executable.chmod(0o0770)  # rwxrwx---, might not be needed though

    # create a job file
    job_file = factory.create(executable=executable.path)
    print("create job file {}".format(job_file))
    print("content:\n{}\n".format(open(job_file, "r").read()))

    # submit the job
    # the return value will always be a list of job ids as htcondor supports cluster submission,
    # i.e., a single job file can result in multiple jobs
    job_id = manager.submit(job_file)[0]
    print("submitted job, id is {}\n".format(job_id))

    # start querying the status
    print("start status queries")
    for _ in range(1000):
        time.sleep(5)
        response = manager.query(job_id)
        print(response)

        if response["status"] not in ("pending", "running"):
            print("stop querying")
            break


def test_slurm():
    law.contrib.load("slurm")

    manager = law.slurm.SlurmJobManager()
    factory = law.slurm.SlurmJobFileFactory()

    executable = None  # TODO

    job_file = factory.create(executable=executable)
    print("create job file {}".format(job_file))
    print("content:\n{}\n".format(open(job_file, "r").read()))

    job_id = manager.submit(job_file)[0]
    print("submitted job, id is {}\n".format(job_id))

    print("start status queries")
    for _ in range(1000):
        time.sleep(5)
        response = manager.query(job_id)
        print(response)

        if response["status"] not in ("pending", "running"):
            print("stop querying")
            break


if __name__ == "__main__":
    test_htcondor()
    # test_slurm()