U  g@sddlZddlZddlZddlZddlZddlZddlZddlm Z m Z ej j dddZ edZddd Zd d Zd d ZddZddZdddZdS)N)MPIMixedCMDFGenz mpirun.lsfF) raise_missinglsf_cray_intelcCs@|dkr t}tjdg}tjj|d|d}|}|}dS)z%!Runs the "sync" command as an exe().Nz /bin/syncT)capturelogger) module_loggerprodutilprogRunnerpipelinePipeline to_stringpoll)rsyncpversionstatusrI/lfs/h1/ops/prod/packages/hmon.v3.2.7/ush/produtil/mpi_impl/mpirun_lsf.pyrunsyncs rcCs8|dk r0t|dr||_t|dr4|j|dn|`|S)z!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 integerNthreadsenv)OMP_NUM_THREADS)hasattrrr)argrrrropenmp"s  rcCstdk S)zn!Determines if LSF+IBMPE should be used to run MPI programs by looking for the mpirun.lsf program in $PATHN)mpirun_lsf_pathrrrrdetect2srcCsdS)z@!Does this module represent an MPI implementation? Returns True.Trrrrr can_run_mpi7sr cKstjjt|gf|S)z!Returns an ImmutableRunner that will run the specified program. @returns an empty list @param exe The executable to run on compute nodes. @param kwargs Ignored.)r r ImmutableRunnerstr)exekwargsrrr make_bigexe;sr%c Ks |dkrtd}t|tjjs$t|\}}|j}|js@d}d|}i}| ddr`ddi}|rp|rpt d | dkr|rd d |j t gggd D} tj| j|d djf|S| dkr6t} tjdd |j t gggd D} | jd| | dd d ddjf|} |dk r2|dt| | S|rFtdn|rtj|}dd |j dddD} tjddkrtdtjjt dgtd| fddd|d} | j|d djf|Sd d |j dddD} tjjt gtd!| fd"d#d|d} | j|d djf|SdS)$ar!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 @note LSF does not support modifying the number of MPI ranks to use when running a program. You can only use all provided ranks, or one rank.N mpirun_lsfrz%dlabel_ioFZ MP_LABELIOyeszACannot mix serial and parallel MPI ranks in the same MPI program.cSsg|]}|qSrr.0arrr `szmpirunner..)prebeforebetween1)rMKL_NUM_THREADScSsg|]}|qSrrr)rrrr,lsz{(0)}z 1Zcore)ZLSB_PJL_TASK_GEOMETRYZ LSB_HOSTSZLSB_MCPU_HOSTSZLSB_DJOB_NUMPROCZLSB_MAX_NUM_PROCESSORSZMP_TASK_AFFINITYzCUsing a hack to work around an LSF bug and run a one core program: zWhen using allranks=True, you must provide an mpi program specification with only one MPI rank (to be duplicated across all ranks).cSsg|]}|qSrrr)rrrr,sT)to_shellexpand mpiserialzaAttempting to run a serial program via mpirun.lsf but the mpiserial program is not in your $PATH. serialcmdf SCR_CMDFILEZ SCR_PGMMODEL) cmd_envar model_envar)preruncSsg|]}|qSrrr)rrrr,sZmpirun_lsf_cmdfZ MP_CMDFILEZ MP_PGMMODEL)logging getLogger isinstancer mpiprog MPIRanksBaseAssertionError check_serialrgetrnranks to_arglistrr r rsocket gethostnameinforeprMPIAllRanksErrorcollapsefileopfind_exeMPISerialMissingr) rallranksrr$serialparallelrZsthreadsZmore_envarglisthostrunnerlinesrrr mpirunnerBs            rT)N)FN)osrDr:produtil.fileopr produtil.progprodutil.mpiprogprodutil.pipeline mpi_impl_baserrrJrKrr;r rrrr r%rTrrrrs