lsf_cray_intel.py

## @namespace produtil.mpi_impl.lsf_cray_intel
# Adds support for LSF+aprun with the Intel OpenMP to produtil.run
#
# This module is part of the mpi_impl package -- see produtil.mpi_impl
# for details.  This implements the bizarre combination of LSF, Cray
# aprun with Intel OpenMP.  

import os, socket, logging
import produtil.fileop,produtil.prog,produtil.mpiprog

from .mpi_impl_base import MPIMixed,CMDFGen

module_logger=logging.getLogger('lsf_cray_intel')

##@var mpirun_lsf_path
# Path to the mpirun.lsf program, or None if it isn't found.
aprun_path=produtil.fileop.find_exe('aprun',raise_missing=False)

class Syncer(object):
    """!Calls sync and aprun -B sync.  This is meant to be run before 
    aprunning another program."""
    def __init__(self):
        super(Syncer,self).__init__()
    def __call__(self,runner=None,logger=None):
        if logger is None: logger=module_logger
        logger.info("Running sync...")
        produtil.pipeline.simple_run(["/bin/sync"],logger=logger)
        logger.info('Running aprun -B sync...')
        produtil.pipeline.simple_run([aprun_path,'-B','/bin/sync'],logger=logger)

def runsync(logger=None):
    """!Runs the "sync" command as an exe()."""
    s=Syncer()
    s(logger=logger)

def openmp(arg,threads):
    """!Adds OpenMP support to the provided object

    @param arg An produtil.prog.Runner or
    produtil.mpiprog.MPIRanksBase object tree
    @param threads the number of threads, or threads per rank, an
    integer"""
    
    if threads is None:
        try:
            ont=os.environ.get('OMP_NUM_THREADS','')
            ont=int(ont)
            if ont>0:
                threads=ont
        except (KeyError,TypeError,ValueError) as e:
            pass

    if threads is None:
        nodesize=os.environ.get('PRODUTIL_RUN_NODESIZE','24')
        nodesize=int(nodesize)
        threads=max(1,nodesize-1)
        
    assert(threads>0)
    threads=int(threads)
    module_logger.info('Threaded with %s threads so add -cc depth.'%(
            repr(threads),))
    for a in reversed(['-cc','depth','-d',str(threads)]):
        arg=arg.argins(1,a)

    arg=arg.env(KMP_AFFINITY='disabled',OMP_NUM_THREADS=threads)
    arg.threads=threads
    
    return arg

def aprun_ln_sf(source,target,content,logger=None):
    produtil.pipeline.simple_run([
        aprun_path,"-N","1","-n","1","-j","1","/bin/ln","-sf",
        str(source),str(target)],   logger=logger)
    content=os.readlink(target)
    if content==source: return True
    raise produtil.fileop.WrongSymlink(
        "FILESYSTEM FAILURE: Python and aprun ln -sf both cannot symlink \"%s\" -> \"%s\".  Instead, the symlink is to \"%s\"."%(
            target,source,content),target)

def detect():
    """!Determines if Cray aprun should be used to run MPI programs by
    looking for the aprun program in $PATH"""
    detected=aprun_path is not None
    if detected:
        produtil.fileop.ln_sf = aprun_ln_sf
    return detected

def can_run_mpi():
    """!Does this module represent an MPI implementation? Returns True."""
    return True

def make_bigexe(exe,**kwargs): 
    """!Returns an ImmutableRunner that will run the specified program.
    @returns an empty list
    @note This function does NOT search $PATH.  That ensures the $PATH
      will be expanded on the compute node instead.  Use
      produtil.fileop.find_exe() if you want to explicitly search the
      PATH before execution
    @param exe The executable to run on compute nodes.
    @param kwargs Ignored."""
    exe=str(exe)    
    inside_aprun=int(os.environ.get('INSIDE_APRUN','0'))
    inside_aprun+=1
    r=produtil.prog.Runner(['aprun','-N','1','-j','1','-n','1',exe])\
        .env(INSIDE_APRUN=str(inside_aprun))
    r.prerun(Syncer())
    return produtil.prog.ImmutableRunner(r)

def mpirunner(arg,allranks=False,logger=None,**kwargs):
    if logger is None:
        logger=module_logger
    m=mpirunner2(arg,allranks=allranks,logger=logger,**kwargs)
    logger.info('mpirunner: %s => %s'%(repr(arg),repr(m)))
    return m

def mpirunner2(arg,allranks=False,logger=None,**kwargs):
    """!Turns a produtil.mpiprog.MPIRanksBase tree into a produtil.prog.Runner
    @param arg a tree of produtil.mpiprog.MPIRanksBase objects
    @param allranks if True, and only one rank is requested by arg, then
      all MPI ranks will be used
    @param logger a logging.Logger for log messages
    @param kwargs passed to produtil.mpi_impl.mpi_impl_base.CMDFGen
      when mpiserial is in use.
    @returns a produtil.prog.Runner that will run the selected MPI program"""
    if logger is None:
        logger=logging.getLogger('lsf_cray_intel')
    assert(isinstance(arg,produtil.mpiprog.MPIRanksBase))
    (serial,parallel)=arg.check_serial()
    if serial and parallel:
        raise MPIMixed(
            'Cannot mix serial and parallel MPI ranks in the same '
            'MPI program.')
    nodesize=os.environ.get('PRODUTIL_RUN_NODESIZE','24')
    nodesize=int(nodesize)
    hyperthreads=os.environ.get('PRODUTIL_RUN_HYPERTHREADS','1')
    hyperthreads=int(hyperthreads)
    maxtasks=os.environ['TOTAL_TASKS']
    maxtasks=int(maxtasks)

    # The returned runner object.  We'll add to this below:
    runner=produtil.prog.Runner([aprun_path]).env(KMP_AFFINITY='disabled')
    threads=arg.threads
    logger.info('Decide what to do with -cc option.')
    if threads is not None and threads>1:
        logger.info('Threaded with %s threads so add -cc depth.'%(
                repr(threads),))
        runner['-cc','depth','-d',str(threads)]
    else:
        logger.info('No threads (threads=%s).'%(repr(threads),))

    # Set up the INSIDE_APRUN variable so the executed MPI program
    # will be able to run serial programs.
    inside_aprun=int(os.environ.get('INSIDE_APRUN','0'))
    inside_aprun+=1
    runner=runner.env(INSIDE_APRUN=str(inside_aprun))

    if arg.nranks()==1 and allranks:
        for rank,count in arg.expand_iter(expand=False):
            pernode=min(maxtasks,nodesize)
            runner=runner['-j','%d'%(hyperthreads),
                          '-N','%d'%(pernode),
                          '-n','%d'%(maxtasks)][rank.args()]
            runner.prerun(Syncer())
            return runner
    elif allranks:
        raise MPIAllRanksError(
            "When using allranks=True, you must provide an mpi program "
            "specification with only one MPI rank (to be duplicated across "
            "all ranks).")
    elif serial:
        lines=[a for a in arg.to_arglist(to_shell=True,expand=True)]
        if produtil.fileop.find_exe('mpiserial') is None:
            raise MPISerialMissing(
                'Attempting to run a serial program via aprun mpiserial but '
                'the mpiserial program is not in your $PATH.')
        pernode=min(arg.nranks(),nodesize)
        return runner['-n%s'%(arg.nranks()),
                      '-j','1','-N','%d'%pernode,'mpiserial']\
                      .prerun(CMDFGen('serialcmdf',lines,**kwargs)) \
                      .prerun(Syncer())
    else:
        first=True
        
        maxtasks=1
        for rank,count in arg.expand_iter(expand=False):
            maxtasks=max(count,maxtasks)
            
        for rank,count in arg.expand_iter(expand=False):
            pernode=min(count,maxtasks,nodesize)
            if not first: runner=runner[':']
            runner=runner['-j','%d'%(hyperthreads),
                          '-N','%d'%(pernode),
                          '-n','%d'%(count),
                          rank.to_shell()]
            first=False
        runner.prerun(Syncer())
        return runner