Commit 3d9eac84 authored by Ambra Di Piano's avatar Ambra Di Piano
Browse files

unify make_sh and make_sbatch and slurm_submission

parent d8ea673c

astrort/configure/slurmjobs.py

+21 −17
Original line number Diff line number Diff line
@@ -7,8 +7,10 @@ 
# *****************************************************************************



from yaml import dump

from os import system

from os import system, makedirs

from os.path import join, dirname, abspath

from astrort.utils.wrap import load_yaml_conf

from astrort.configure.logging import set_logger, get_log_level



def make_configuration(jobname_conf, configuration, node_number):

    # simulator
@@ -41,21 +43,23 @@ def make_sh(jobname, slurmconf, jobname_conf, jobname_sh, jobname_log, mode='sim 
        else:

            raise ValueError(f"Invalid 'mode' {mode}")



def make_simulator_sbatch(jobname, configuration, node_number):

    output = configuration['simulator']['output']

    jobname_sh = join(output, f"{jobname}_simulator.sh")

    jobname_log = join(output, f"{jobname}_simulator.slurm")

    jobname_conf = join(output, f"{jobname}_simulator.yml")

    make_configuration(jobname_conf, configuration, node_number, mode='simulator')

def make_sbatch(jobname, configuration, node_number, mode):

    output = configuration[mode]['output']

    jobname_sh = join(output, f"{jobname}_{mode}.sh")

    jobname_log = join(output, f"{jobname}_{mode}.slurm")

    jobname_conf = join(output, f"{jobname}_{mode}.yml")

    make_configuration(jobname_conf, configuration, node_number, mode=mode)

    make_sh(jobname, configuration['slurm'], jobname_conf, jobname_sh, jobname_log)

    system(f"sbatch {jobname_sh}")



def make_mapper_sbatch(jobname, configuration, node_number):

    output = configuration['mapper']['output']

    jobname_sh = join(output, f"{jobname}_mapper.sh")

    jobname_log = join(output, f"{jobname}_mapper.slurm")

    jobname_conf = join(output, f"{jobname}_mapper.yml")

    make_configuration(jobname_conf, configuration, node_number, mode='mapper')

    make_sh(jobname, configuration['slurm'], jobname_conf, jobname_sh, jobname_log)

    system(f"sbatch {jobname_sh}")

    return
 
 No newline at end of file
 
def slurm_submission(configuration_file, nodes, mode):

    configuration = load_yaml_conf(configuration_file)

    log = set_logger(get_log_level(configuration['logging']['level']))

    # create output dir

    log.info(f"Creating {configuration['simulator']['output']}")

    makedirs(configuration['simulator']['output'], exist_ok=True)

    # sbatch jobs per each nodes

    configuration['slurm']['nodes'] = nodes

    for node_number in range(configuration['slurm']['nodes']):

        jobname = f"{configuration['slurm']['name']}_{node_number+1}"

        make_sbatch(jobname, configuration, node_number, mode=mode)

astrort/simulator/base_mapper.py

+2 −16
Original line number Diff line number Diff line
@@ -13,7 +13,7 @@ from os.path import join 
from astrort.utils.wrap import load_yaml_conf, write_mapping_info, execute_mapper_no_visibility

from astrort.utils.utils import get_all_seeds

from astrort.configure.logging import set_logger, get_log_level

from astrort.configure.slurmjobs import make_mapper_sbatch

from astrort.configure.slurmjobs import slurm_submission



def base_mapper(configuration_file, seeds=None):

    clock = time()
@@ -46,25 +46,11 @@ def base_mapper(configuration_file, seeds=None): 
    log.info(f"\n {'-'*15} \n| STOP MAPPER | \n {'-'*15} \n")

    log.info(f"Process complete, took {time() - clock} s")



def slurm_submission(configuration_file, nodes):

    configuration = load_yaml_conf(configuration_file)

    log = set_logger(get_log_level(configuration['logging']['level']))

    # create output dir

    log.info(f"Creating {configuration['simulator']['output']}")

    makedirs(configuration['simulator']['output'], exist_ok=True)

    # sbatch jobs per each nodes

    configuration['slurm']['nodes'] = nodes

    for node_number in range(configuration['slurm']['nodes']):

        jobname = f"{configuration['slurm']['name']}_{node_number+1}"

        make_mapper_sbatch(jobname, configuration, node_number)

    return





def main(configuration, nodes):

    if nodes == 0:

        base_mapper()(configuration)

    else:

        slurm_submission(configuration, nodes)

        slurm_submission(configuration, nodes, mode='mapper')



if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='')

astrort/simulator/base_simulator.py

+2 −16
Original line number Diff line number Diff line
@@ -13,7 +13,7 @@ from os.path import join 
from rtasci.lib.RTACtoolsSimulation import RTACtoolsSimulation

from astrort.utils.wrap import load_yaml_conf, configure_simulator_no_visibility, write_simulation_info, set_pointing

from astrort.configure.logging import set_logger, get_log_level

from astrort.configure.slurmjobs import make_simulator_sbatch

from astrort.configure.slurmjobs import slurm_submission



def base_simulator(configuration_file):

    clock = time()
@@ -46,25 +46,11 @@ def base_simulator(configuration_file): 
    log.info(f"\n {'-'*17} \n| STOP SIMULATOR | \n {'-'*17} \n")

    log.info(f"Process complete, took {time() - clock} s")





def slurm_submission(configuration_file, nodes):

    configuration = load_yaml_conf(configuration_file)

    log = set_logger(get_log_level(configuration['logging']['level']))

    # create output dir

    log.info(f"Creating {configuration['simulator']['output']}")

    makedirs(configuration['simulator']['output'], exist_ok=True)

    # sbatch jobs per each nodes

    configuration['slurm']['nodes'] = nodes

    for node_number in range(configuration['slurm']['nodes']):

        jobname = f"{configuration['slurm']['name']}_{node_number+1}"

        make_simulator_sbatch(jobname, configuration, node_number)

    return



def main(configuration, nodes):

    if nodes == 0:

        base_simulator(configuration)

    else:

        slurm_submission(configuration, nodes)

        slurm_submission(configuration, nodes, mode='simulator')



if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='')

astrort/testing/test_configure/test_slurmjobs.py

+9 −8
Original line number Diff line number Diff line
@@ -10,7 +10,7 @@ import pytest 
from shutil import rmtree

from os import listdir, makedirs

from os.path import isfile, join

from astrort.configure.slurmjobs import make_configuration, make_sh, make_simulator_sbatch

from astrort.configure.slurmjobs import make_configuration, make_sh

from astrort.utils.wrap import load_yaml_conf



@pytest.mark.test_conf_file
@@ -34,7 +34,8 @@ def test_make_configuration(test_conf_file): 
    assert found_configurations == expected_configurations, f"Expected {expected_configurations} simulations, found {found_configurations}"



@pytest.mark.test_conf_file

def test_make_sh(test_conf_file):

@pytest.mark.parametrize('mode', ['simulator', 'mapper'])

def test_make_sh(test_conf_file, mode):



    # clean output

    conf = load_yaml_conf(test_conf_file)
@@ -46,13 +47,13 @@ def test_make_sh(test_conf_file): 
    for node_number in range(conf['slurm']['nodes']):

        node_number += 1

        jobname = f"{conf['slurm']['name']}_{node_number}"

        jobname_sh = join(output, f"job_{jobname}.sh")

        jobname_log = join(output, f"job_{jobname}.log")

        jobname_conf = join(output, f"job_{jobname}.yml")

        make_sh(jobname, conf['slurm'], jobname_conf, jobname_sh, jobname_log, mode='simulator')

        jobname_sh = join(output, f"job_{jobname}_{mode}.sh")

        jobname_log = join(output, f"job_{jobname}_{mode}.log")

        jobname_conf = join(output, f"job_{jobname}_{mode}.yml")

        make_sh(jobname, conf['slurm'], jobname_conf, jobname_sh, jobname_log, mode=mode)



    # check output

    expected_sh = conf['slurm']['nodes']

    found_sh = len([f for f in listdir(conf['simulator']['output']) if isfile(join(conf['simulator']['output'], f)) and '.sh' in f and conf['slurm']['name'] in f])

    assert found_sh == expected_sh, f"Expected {expected_sh} simulations, found {found_sh}"

    found_sh = len([f for f in listdir(conf['simulator']['output']) if isfile(join(conf['simulator']['output'], f)) and '.sh' in f and conf['slurm']['name'] in f and mode in f])

    assert found_sh == expected_sh, f"Expected {expected_sh} files for {mode}, found {found_sh}"