Limits yaml schema

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

roles/slurm/templates/update_config.j2 → 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)

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