mpi_commands.py 2.98 KB
Newer Older
Michael Salim's avatar
Michael Salim committed
1
class DEFAULTMPICommand(object):
2
    '''Single node OpenMPI: ppn == num_ranks'''
3 4 5
    def __init__(self):
        self.mpi = 'mpirun'
        self.nproc = '-n'
6 7
        self.ppn = '-npernode'
        self.env = '-x'
8 9 10 11 12 13 14
        self.cpu_binding = None
        self.threads_per_rank = None
        self.threads_per_core = None

    def worker_str(self, workers):
        return ""

15
    def env_str(self, envs):
16
        envstrs = (f'{self.env} {var}="{val}"' for var,val in envs.items())
17
        return " ".join(envstrs)
18

19
    def threads(self, thread_per_rank, thread_per_core):
20 21 22 23
        result= ""
        if self.cpu_binding:
            result += f"{self.cpu_binding} "
        if self.threads_per_rank:
24
            result += f"{self.threads_per_rank} {thread_per_rank} "
25
        if self.threads_per_core:
26
            result += f"{self.threads_per_core} {thread_per_core} "
27 28
        return result

29 30
    def __call__(self, workers, *, app_cmd, num_ranks, ranks_per_node, envs,threads_per_rank=1,threads_per_core=1):
        '''Build the mpirun/aprun/runjob command line string'''
31
        workers = self.worker_str(workers)
32 33 34
        envs = self.env_str(envs)
        thread_str = self.threads(threads_per_rank, threads_per_core)
        result =  (f"{self.mpi} {self.nproc} {num_ranks} {self.ppn} "
35
                   f"{ranks_per_node} {envs} {workers} {thread_str} {app_cmd}")
36 37 38
        return result


Michael Salim's avatar
Michael Salim committed
39
class BGQMPICommand(DEFAULTMPICommand):
40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55
    def __init__(self):
        self.mpi = 'runjob'
        self.nproc = '--np'
        self.ppn = '-p'
        self.env = '--envs' # VAR1=val1:VAR2=val2
        self.cpu_binding = None
        self.threads_per_rank = None
        self.threads_per_core = None
    
    def worker_str(self, workers):
        if len(workers) != 1:
            raise BalsamRunnerException("BGQ requires exactly 1 worker (sub-block)")
        worker = workers[0]
        shape, block, corner = worker.shape, worker.block, worker.corner
        return f"--shape {shape} --block {block} --corner {corner} "

Michael Salim's avatar
Michael Salim committed
56
class CRAYMPICommand(DEFAULTMPICommand):
57 58 59 60 61
    def __init__(self):
        # 64 independent jobs, 1 per core of a KNL node: -n64 -N64 -d1 -j1
        self.mpi = 'aprun'
        self.nproc = '-n'
        self.ppn = '-N'
62
        self.env = '-e'
63 64 65 66 67 68 69
        self.cpu_binding = '-cc depth'
        self.threads_per_rank = '-d'
        self.threads_per_core = '-j'
    
    def worker_str(self, workers):
        if not workers:
            return ""
Michael Salim's avatar
Michael Salim committed
70
        return f"-L {','.join(str(worker.id) for worker in workers)}"
71 72 73 74 75 76 77

class COOLEYMPICommand(DEFAULTMPICommand):
    def __init__(self):
        # 64 independent jobs, 1 per core of a KNL node: -n64 -N64 -d1 -j1
        self.mpi = 'mpirun'
        self.nproc = '-n'
        self.ppn = '--ppn'
78
        self.env = '--env'
79 80 81 82 83 84 85 86
        self.cpu_binding = None
        self.threads_per_rank = None
        self.threads_per_core = None
    
    def worker_str(self, workers):
        if not workers:
            return ""
        return f"--hosts {','.join(str(worker.id) for worker in workers)} "