diff --git a/roles/slurm/files/test_update_config.py b/roles/slurm/files/test_update_config.py new file mode 100644 index 00000000..99f437b2 --- /dev/null +++ b/roles/slurm/files/test_update_config.py @@ -0,0 +1,47 @@ +import pytest +import requests_mock +import yaml + +import update_config + + +@pytest.mark.parametrize( + "data", + [ + (""" + VM.Standard2.1: + 1: 1 + 2: 1 + 3: 1 + VM.Standard2.2: + 1: 1 + 2: 1 + 3: 1 + """), + ] +) +def test_get_limits(mocker, data): + mocker.patch("update_config.load_yaml", return_value=yaml.safe_load(data)) + assert isinstance(update_config.get_limits(), dict) + + +@pytest.mark.parametrize( + "data,error", + [ + (""" + VM.Standard2.1: + 1: 1 + 2: 1 + 3: 1 + VM.Standard2.2: + 1: 1 + 2: 1 + 3: 1 + """, + SyntaxError), + ] +) +def test_get_limits_errors(mocker, data, error): + mocker.patch("update_config.load_yaml", return_value=yaml.safe_load(data)) + with pytest.raises(error): + update_config.get_limits() diff --git a/roles/slurm/templates/update_config.j2 b/roles/slurm/files/update_config.py similarity index 68% rename from roles/slurm/templates/update_config.j2 rename to roles/slurm/files/update_config.py index 39a21842..a361e0b5 100644 --- a/roles/slurm/templates/update_config.j2 +++ b/roles/slurm/files/update_config.py @@ -16,7 +16,19 @@ def get_limits() -> Dict[str, Dict[str, str]]: Until OCI has an API to fetch service limits, we have to hard-code them in a file. """ - return load_yaml("limits.yaml") + limits = load_yaml("limits.yaml") + for mappings in limits.values(): + if not isinstance(mappings, dict): + raise SyntaxError + for ad, count in mappings.items(): + if not isinstance(ad, int): + raise SyntaxError + if not isinstance(count, int): + raise SyntaxError + for shape in limits: + if re.match(r"\w{2}\.", shape) is None: + raise ValueError + return limits def get_shapes() -> Dict[str, Dict[str, str]]: @@ -73,18 +85,24 @@ def get_node_configs(limits, shapes, mgmt_info): yield create_slurmconf_line(i, shape_info, shape, ad) -# TODO Make sure that any nodes which are no longer managed due to service limit reductions are terminated. +if __name__ == "__main__": + # TODO Make sure that any nodes which are no longer managed due to service limit reductions are terminated. -slurm_conf_filename = "/mnt/shared/etc/slurm/slurm.conf" + slurm_conf_filename = "/mnt/shared/etc/slurm/slurm.conf" -node_config = "\n".join(get_node_configs(get_limits(), get_shapes(), get_mgmt_info())) + try: + limits = get_limits() + except SyntaxError: + print("ERROR: Syntax error in `limits.yaml`.") + exit(1) + node_config = "\n".join(get_node_configs(limits, get_shapes(), get_mgmt_info())) -chop = re.compile('(?<=# STARTNODES\n)(.*?)(?=\n?# ENDNODES)', re.DOTALL) + chop = re.compile('(?<=# STARTNODES\n)(.*?)(?=\n?# ENDNODES)', re.DOTALL) -with open(slurm_conf_filename) as f: - all_config = f.read() + with open(slurm_conf_filename) as f: + all_config = f.read() -new_config = chop.sub('{}'.format(node_config), all_config) + new_config = chop.sub('{}'.format(node_config), all_config) -with open(slurm_conf_filename, "w") as f: - f.write(new_config) + with open(slurm_conf_filename, "w") as f: + f.write(new_config) diff --git a/roles/slurm/tasks/elastic.yml b/roles/slurm/tasks/elastic.yml index c37c5e4b..58d86d6b 100644 --- a/roles/slurm/tasks/elastic.yml +++ b/roles/slurm/tasks/elastic.yml @@ -59,9 +59,9 @@ dest: /usr/local/bin/stopnode mode: 0755 -- name: configure update_config script - template: - src: update_config.j2 +- name: install update_config script + copy: + src: update_config.py dest: /usr/local/bin/update_config mode: 0755