def test_get_next_expid(self): a = obs.get_next_expid() b = obs.get_next_expid() c = obs.get_next_expid() self.assertNotEqual(a, b) self.assertNotEqual(b, c)
def batch_newexp(batchfile, flavors, nspec=5000, night=None, expids=None, nodes=None, pixprod=None, desi_spectro_sim=None, tileids=None, seed=None): ''' Write a slurm batch script for run newexp-desi for the list of flavors ''' nexp = len(flavors) timestr = '00:30:00' logfile = '{}.%j.log'.format(batchfile) np.random.seed(seed) seeds = np.random.randint(2**32, size=nexp) if night is None: night = obs.get_night() if expids is None: expids = obs.get_next_expid(nexp) if nodes is None: nodes = calc_nodes(nexp, tasktime=1.5, maxtime=20) if tileids is None: tileids = list() for flavor in flavors: flavor = flavor.lower() t = obs.get_next_tileid(program=flavor) tileids.append(t) if flavor in ('arc', 'flat'): obs.update_obslog(obstype=flavor, program='calib', tileid=t) elif flavor in ('bright', 'bgs', 'mws'): obs.update_obslog(obstype='science', program='bright', tileid=t) elif flavor in ('gray', 'grey'): obs.update_obslog(obstype='science', program='gray', tileid=t) else: obs.update_obslog(obstype='science', program='dark', tileid=t) if pixprod is None: if 'PIXPROD' in os.environ: pixprod = os.environ['PIXPROD'] else: raise ValueError('must provide pixprod or set $PIXPROD') if desi_spectro_sim is None: if 'DESI_SPECTRO_SIM' in os.environ: desi_spectro_sim = os.environ['DESI_SPECTRO_SIM'] else: raise ValueError( 'must provide desi_spectro_sim or set $DESI_SPECTRO_SIM') log.info('output dir {}/{}/{}'.format(desi_spectro_sim, pixprod, night)) assert len(expids) == len(flavors) cmd = "srun -n 1 -N 1 -c $nproc /usr/bin/time newexp-desi --night {night} --nspec {nspec} --flavor {flavor} --expid {expid} --tileid {tileid} --seed {seed}" with open(batchfile, 'w') as fx: fx.write("#!/bin/bash -l\n\n") fx.write("#SBATCH --partition=debug\n") fx.write("#SBATCH --account=desi\n") fx.write("#SBATCH --nodes={}\n".format(nodes)) fx.write("#SBATCH --time={}\n".format(timestr)) fx.write("#SBATCH --job-name=newexp\n") fx.write("#SBATCH --output={}\n".format(logfile)) fx.write("if [ ${NERSC_HOST} = edison ]; then\n") fx.write(" nproc=24\n") fx.write("else\n") fx.write(" nproc=32\n") fx.write("fi\n\n") fx.write('export DESI_SPECTRO_SIM={}\n'.format(desi_spectro_sim)) fx.write('export PIXPROD={}\n'.format(pixprod)) fx.write('\n') fx.write('echo Starting at `date`\n\n') fx.write('mkdir -p $DESI_SPECTRO_SIM/$PIXPROD/etc\n') fx.write('mkdir -p $DESI_SPECTRO_SIM/$PIXPROD/{}\n'.format(night)) fx.write('\n') for expid, flavor, tileid, seed in zip(expids, flavors, tileids, seeds): fx.write( cmd.format(nspec=nspec, night=night, expid=expid, flavor=flavor, tileid=tileid, seed=seed) + ' &\n') fx.write('\nwait\n') fx.write('\necho Done at `date`\n') return expids
def batch_pixsim(batchfile, flavors, nspec=5000, night=None, expids=None, nodes=None, pixprod=None, desi_spectro_sim=None, seed=None): ''' Write a slurm batch script for run newexp-desi for the list of flavors ''' nexp = len(flavors) nspectrographs = (nspec - 1) // 500 + 1 ntasks = nexp * 3 timestr = '00:30:00' logfile = '{}.%j.log'.format(batchfile) np.random.seed(seed) seeds = np.random.randint(2**32, size=nexp) if night is None: night = obs.get_night() if expids is None: expids = obs.get_next_expid(nexp) assert len(expids) == len(flavors) if pixprod is None: if 'PIXPROD' in os.environ: pixprod = os.environ['PIXPROD'] else: raise ValueError('must provide pixprod or set $PIXPROD') if desi_spectro_sim is None: if 'DESI_SPECTRO_SIM' in os.environ: desi_spectro_sim = os.environ['DESI_SPECTRO_SIM'] else: raise ValueError( 'must provide desi_spectro_sim or set $DESI_SPECTRO_SIM') outdir = os.path.normpath(os.path.join(desi_spectro_sim, pixprod, night)) log.info('output dir {}'.format(outdir)) if nodes is None: # nodes = calc_nodes(ntasks, tasktime=5, maxtime=20) nodes = nexp * nspectrographs cmd = "srun -n {nspectrographs} -N {nspectrographs} -c $nproc /usr/bin/time pixsim-desi --mpi --verbose --cosmics --night {night} --expid {expid} --seed {seed}" with open(batchfile, 'w') as fx: fx.write("#!/bin/bash -l\n\n") fx.write("#SBATCH --partition=debug\n") fx.write("#SBATCH --account=desi\n") fx.write("#SBATCH --nodes={}\n".format(nodes)) fx.write("#SBATCH --time={}\n".format(timestr)) fx.write("#SBATCH --job-name=pixsim\n") fx.write("#SBATCH --output={}\n".format(logfile)) fx.write("if [ ${NERSC_HOST} = edison ]; then\n") fx.write(" nproc=24\n") fx.write("else\n") fx.write(" nproc=32\n") fx.write("fi\n\n") fx.write('export DESI_SPECTRO_SIM={}\n'.format(desi_spectro_sim)) fx.write('export PIXPROD={}\n'.format(pixprod)) fx.write('\necho Starting at `date`\n') for expid, flavor, seed in zip(expids, flavors, seeds): fx.write('\n#- Exposure {} ({})\n'.format(expid, flavor)) cx = cmd.format(night=night, expid=expid, nspectrographs=nspectrographs, seed=seed) fx.write(cx + ' &\n') fx.write('\nwait\n') fx.write('\n#- Preprocess raw data\n') cmd = 'srun -n 1 -N 1 desi_preproc --infile {infile} --outdir {outdir} --cameras {cameras}' outdir = os.path.join(desi_spectro_sim, pixprod, night) for expid in expids: infile = desispec.io.findfile('raw', night=night, expid=expid, outdir=outdir) for i in range(nspectrographs): cameras = 'b{},r{},z{}'.format(i, i, i) cx = cmd.format(infile=infile, outdir=outdir, cameras=cameras) fx.write(cx + ' &\n') fx.write('\nwait\n') fx.write('echo Done at `date`\n')
] ntask = len(flavors) * len(nights) comm = MPI.COMM_WORLD if comm.size < ntask: if comm.rank == 0: print("Communicator size ({}) too small for {} tasks".format(comm.size, ntask), flush=True) comm.Abort() np.random.seed(123456) seeds = np.random.randint(2**32, size=ntask) expids = None if comm.rank == 0: expids = obs.get_next_expid(ntask) expids = comm.bcast(expids, root=0) tileids = list() if comm.rank == 0: for nt in nights: for fl in flavors: flavor = fl.lower() t = obs.get_next_tileid(program=flavor) tileids.append(t) if flavor in ('arc', 'flat'): obs.update_obslog(obstype=flavor, program='calib', tileid=t) elif flavor in ('bright', 'bgs', 'mws'): obs.update_obslog(obstype='science', program='bright', tileid=t) elif flavor in ('gray', 'grey'): obs.update_obslog(obstype='science', program='gray', tileid=t)