run_unfold_genie_multiprocess.py

Go to the documentation of this file.

#!/usr/bin/env python

import subprocess
import multiprocessing
import os
import argparse
def make_cmd(nnodes, ranks_per_node, procid, nuniv_per_rank, nfiles_per_rank, genie_batch_idx, totuniv, totfiles):

    cmd = []
    name_cmd = []
    covered = False
    # ncmd = nuniv*nfiles
    nfiles_last_batch = totfiles%nfiles_per_rank

    if nnodes > 1 and int(procid) == nnodes - 1:
        ifilefirst = totfiles - nfiles_last_batch
        ifilelast = ifilefirst + nfiles_last_batch -1
    else:
        ifilefirst = int(procid) * nfiles_per_rank 
        ifilelast = ifilefirst + nfiles_per_rank - 1
    # if int(procid) == ranks_per_node:
    #     ifilefirst = int(procid) * nfiles_per_rank * ranks_per_node + irank * nfiles_per_rank
    #     ifilelast = ifilefirst + nfiles_per_rank - 1
    # else:
    #     ifilefirst = totfiles - nfiles_last_batch
    #     ifilelast = totfiles - 1

    for irank in range(ranks_per_node):
        iunivfirst = genie_batch_idx * totuniv + irank * nuniv_per_rank
        iunivlast = iunivfirst + nuniv_per_rank - 1

        icmd = "cafe -bq -nr --numuccinc /development/NDAna/numucc_inc/specprod_numuccinc.C genie {} {} -1 -1 0 {} {}".format(
                iunivfirst, iunivlast, ifilefirst, ifilelast)
        iname_cmd = "ProcID: {} , Genie batch index: {}, cafe -bq -nr --numuccinc /development/NDAna/numucc_inc/specprod_numuccinc.C genie {} {} -1 -1 0 {} {}".format(
                procid, genie_batch_idx, iunivfirst, iunivlast, ifilefirst, ifilelast)
        cmd.append(icmd)
        name_cmd.append(iname_cmd)

    return (name_cmd, cmd)

class Consumer(multiprocessing.Process):
    def __init__(self, task_queue, result_queue):
        multiprocessing.Process.__init__(self)
        self.task_queue = task_queue
        self.result_queue = result_queue

    def run(self):
        proc_name = self.name
        while True:
            next_task = self.task_queue.get()
            if next_task is None:
                print('{}: Exiting'.format(proc_name))
                break
            print('{}: {}'.format(proc_name, next_task))
            answer = next_task()
            self.result_queue.put(answer)
        return

class Task(object):
    def __init__(self, a, b):
        # a is for command name, b is for command
        self.a = a
        self.b = b

    def __call__(self):
        print('Starting -- {} : {}'.format(self.a, self.b))
        subprocess.call(self.b, shell=True)
        print('Finishing -- {} : {}'.format(self.a, self.b))
        return

    def __str__(self):
        return '%s : %s' % (self.a, self.b)


def RunMultiProcess(nCMD, lCMD, num_consumers=0):
    # Establish communication queues
    tasks = multiprocessing.Queue()
    results = multiprocessing.Queue()

    # Start consumers
    if num_consumers == 0:
        num_consumers = multiprocessing.cpu_count()
    print('Creating {} consumers'.format(num_consumers))
    consumers = [Consumer(tasks, results) for i in range(num_consumers)]
    for w in consumers:
        w.start()

    num_jobs = len(lCMD)
    for i in range(num_jobs):
        tasks.put(Task(nCMD[i], lCMD[i]))

    # Add a poison pill for each consumer
    for i in range(num_consumers):
        tasks.put(None)
    # Start printing results
    while num_jobs:
        result = results.get()
        print('Result: {}'.format(result))
        num_jobs -= 1
    return


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog='make_cafe_cmd',
        description="""Form cafe cmd in shifter jobs. """,
        epilog="Questions and Comments to dingpf@fnal.gov")

    # required arguments
    parser.add_argument('--nnodes', default=1, type=int,
                        help="number of nodes;")
    parser.add_argument('--nuniv_per_rank', default=5, type=int,
                        help="number of universes per rank;")
    parser.add_argument('--nfiles_per_rank', default=100, type=int,
                        help="number of files per rank;")
    # parser.add_argument('--nuniv-rank', default=20, type=int,
    #                     help="number of universes per rank;")
    parser.add_argument('--totuniv', default=1000, type=int,
                        help="total number of universes;")
    parser.add_argument('--genie_batch_idx', default=0, type=int,
                        help="Batch index for Genie universes;")
    parser.add_argument('--totfiles', default=56324, type=int,
                        help="total number of files;")
    parser.add_argument('--nranks_per_node', default=25, type=int,
                        help="number of ranks per node;")

    args = parser.parse_args()

    if args.nranks_per_node == -1:
        nranks_per_node = multiprocessing.cpu_count()
    else:
        nranks_per_node = args.nranks_per_node

    if 'JOBSUBJOBID' in os.environ:
        jsub_id = os.environ['JOBSUBJOBID']
        jid = jsub_id.split('@')[0].split('.')[0]
        procid = jsub_id.split('@')[0].split('.')[1]
    else:
        jid = os.environ["SLURM_JOB_ID"]
        procid = os.environ["SLURM_PROCID"]

    (ncmd, cmd) = make_cmd(args.nnodes, nranks_per_node, procid, args.nuniv_per_rank, args.nfiles_per_rank, args.genie_batch_idx, args.totuniv, args.totfiles)

    for idx in range(len(cmd)):
        print(ncmd[idx], cmd[idx])
    RunMultiProcess(ncmd, cmd, nranks_per_node)

    # subprocess.call("mkdir -p /output/{}/{}/{}".format(args.mode, jid, procid), shell=True)
    # subprocess.call("mv /dev/shm/{}/{}/*.root /output/{}/{}/{}".format(jid, procid,  args.mode, jid, procid), shell=True)
    # subprocess.call("rm -rf /dev/shm/{}".format(jid), shell=True)