slurm jobs configuration utils and tests

# *****************************************************************************

# Copyright (C) 2023 INAF

# This software is distributed under the terms of the BSD-3-Clause license

#

# Authors:

# Ambra Di Piano <ambra.dipiano@inaf.it>

# *****************************************************************************

from yaml import dump

from os import system

from os.path import join, dirname, abspath

def make_configuration(jobname_conf, configuration, node_number):

    configuration['simulator']['seed'] = node_number*configuration['simulator']['samples']

    # write new configuration

    with open(jobname_conf, 'w+') as f:

        new_configuration = dump(configuration, f, default_flow_style=False)

def make_sh(jobname, slurmconf, jobname_conf, jobname_sh, jobname_log):

    # write sbatch

    with open(jobname_sh, 'w+') as f:

        f.write("#!/bin/bash")

        #f.write("\n#SBATCH --nodes=1")

        #f.write("\n#SBATCH --ntasks-per-node=1")

        #f.write("\n#SBATCH --cpus-per-task=2")

        f.write(f"\n#SBATCH --job-name={jobname}")

        #f.write(f"\n#SBATCH --mem={slurmconf['memory']}")

        f.write(f"\n#SBATCH --output={jobname_log}")

        f.write(f"\n#SBATCH --account={slurmconf['account']}")

        f.write(f"\n#SBATCH --partition={slurmconf['partition']}")

        f.write(f"\n")

        f.write(f"\nsource activate {slurmconf['environment']}")

        f.write(f"\npython {join(dirname(abspath(__file__)).replace('configure', 'simulator'), 'base_simulator.py')} -f {jobname_conf}\n")

def make_sbatch(jobname, configuration, node_number):

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

    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_configuration(jobname_conf, configuration, node_number)

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

    system(f"sbatch {jobname_sh}")

 No newline at end of file

# *****************************************************************************

# Copyright (C) 2023 INAF

# This software is distributed under the terms of the BSD-3-Clause license

#

# Authors:

# Ambra Di Piano <ambra.dipiano@inaf.it>

# *****************************************************************************

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_sbatch

from astrort.utils.wrap import load_yaml_conf

@pytest.mark.test_conf_file

def test_make_configuration(test_conf_file):

    # clean output

    conf = load_yaml_conf(test_conf_file)

    rmtree(conf['simulator']['output'])

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

    # make configurations

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

        node_number += 1

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

        jobname_conf = join(conf['simulator']['output'], f"job_{jobname}.yml")

        make_configuration(jobname_conf, conf, node_number)

    # check output

    expected_configurations = conf['slurm']['nodes']

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

    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):

    # clean output

    conf = load_yaml_conf(test_conf_file)

    rmtree(conf['simulator']['output'])

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

    # make configurations

    output = conf['simulator']['output']

    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)

    # 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}"