Commit ccd056bf authored by David M. Rogers's avatar David M. Rogers
Browse files

Added tgt: alternative to top-level, parsed SLURM_JOB_CPUS_PER_NODE.

parent b8e76347

graph.py

+7 −4
Original line number Diff line number Diff line
@@ -23,17 +23,20 @@ def gen_targets(params): 
        yield ttype, tname

    if 'loop' in params:

        assert len(params['loop']) == 2, "Loop must have a single variable."

        assert 'out' in params['loop'], "Loop must contain 'out:'"

        tgtname = 'tgt'

        if tgtname not in params['loop']:

            tgtname = 'in'

        assert tgtname in params['loop'], "Loop must contain 'tgt:/in:'"

        for var, r in params['loop'].items():

            if var != 'out':

            if var is not tgtname:

                break

        assert isinstance(r, str), "Loop must have a single variable holding an expr."

        if '{' in r:

            r = r.format(**params)

        for n in eval(r, {}, {}):

            for ttype, out in params['loop']['out'].items():

            for ttype, tgt in params['loop'][tgtname].items():

                fmt = {var: n}

                yield ttype, out.format(**fmt)

                yield ttype, tgt.format(**fmt)



# Determine if dirname / tname must be generated by pmake

def must_generate(dirname, tname, types):

machine.py

+9 −4
Original line number Diff line number Diff line
@@ -61,10 +61,15 @@ class ResourceSet: 
        self.stop_time = tick()



def machine(time):

    # Rhea

    # Generic Slurm

    if 'SLURM_JOB_NUM_NODES' in os.environ:

        log.info("Running on rhea")

        return Machine(time, int(os.environ['SLURM_JOB_NUM_NODES']), 0, 32, "srun")

        # FIXME: determine HT count

        # "grep 'physical id' /proc/cpuinfo | sort -u | wc -l"

        #import multiprocessing as m

        #cpu = m.cpu_count() // 2

        cpu = int( os.environ['SLURM_JOB_CPUS_PER_NODE'] )

        return Machine(time, int(os.environ['SLURM_JOB_NUM_NODES']),

                       0, cpu, "srun")



    # Summit

    if 'LSB_MAX_NUM_PROCESSORS' in os.environ:
@@ -91,7 +96,7 @@ class Machine: 
        if self.batch == 'srun':

            nodes = R.min_nodes(self.CpN, self.GpN)

            assert nodes is not None

            srun = "srun --exclusive -N %d --ntasks %d -c %d -G %d" % (

            srun = "srun --exclusive -N %d --cpu-bind=cores --ntasks %d -c %d -G %d" % (

                         nodes, R.nrs*R.tasks, R.nrs*R.cpu//nodes, R.nrs*R.gpu)

            if R.srun_attr is not None:

                srun += " %s" % R.srun_attr