Fix a crash in cstr.to_gres_dict when the input is just "gres:gpu" (#334)

Slurm allows to pass things like --gres=gpu, which in this case the
default value will be just one. However, the char *gres value also
really just contains "gres:gpu" when displayed via scontrol, with no
indication of the count. The code previously assumed that Slurm
internally modifies the char* explicitly to reflect the default count of
1.

Author: tazend

pyslurm/utils/cstr.pyx

@@ -214,6 +214,9 @@ cpdef dict to_gres_dict(char *gres):
         return {}
 
     for item in re.split(",(?=[^,]+?:)", gres_str):
+        # char *gres might contain just "gres:gpu", without any count.
+        # If not given, the count is always 1, so default to it.
+        cnt = typ = "1"
 
         # Remove the additional "gres" specifier if it exists
         if gres_delim in item:
@@ -223,15 +226,22 @@ cpdef dict to_gres_dict(char *gres):
             ":(?=[^:]+?)",
             item.replace("(", ":", 1).replace(")", "")
         )
+        gres_splitted_len = len(gres_splitted)
 
-        name, typ, cnt = gres_splitted[0], gres_splitted[1], 0
+        name = gres_splitted[0]
+        if gres_splitted_len > 1:
+            typ = gres_splitted[1]
 
         # Check if we have a gres type.
         if typ.isdigit():
             cnt = typ
             typ = None
-        else:
+        elif gres_splitted_len > 2:
             cnt = gres_splitted[2]
+        else:
+            # String is somehow malformed, should never happen when the input
+            # comes from the slurmctld. Ignore if it happens.
+            continue
 
         # Dict Key-Name depends on if we have a gres type or not
         name_and_typ = f"{name}:{typ}" if typ else name

tests/unit/test_common.py

@@ -152,12 +152,32 @@ def test_dict_to_gres_str(self):
         assert cstr.from_gres_dict("tesla:3,a100:5", "gpu") == expected_str
 
     def test_str_to_gres_dict(self):
-        input_str = "gpu:nvidia-a100:1(IDX:0,1)"
-        expected = {"gpu:nvidia-a100":{"count": 1, "indexes": "0,1"}}
+        input_str = "gpu:nvidia-a100:1(IDX:0)"
+        expected = {"gpu:nvidia-a100": {"count": 1, "indexes": "0"}}
         assert cstr.to_gres_dict(input_str) == expected
 
-        input_str = "gpu:nvidia-a100:1"
-        expected = {"gpu:nvidia-a100": 1}
+        input_str = "gpu:nvidia-a100:2(IDX:0,1)"
+        expected = {"gpu:nvidia-a100": {"count": 2, "indexes": "0,1"}}
+        assert cstr.to_gres_dict(input_str) == expected
+
+        input_str = "gpu:nvidia-a100:5"
+        expected = {"gpu:nvidia-a100": 5}
+        assert cstr.to_gres_dict(input_str) == expected
+
+        input_str = "gpu:nvidia-a100:5,gres:gpu:nvidia-v100:10"
+        expected = {"gpu:nvidia-a100": 5, "gpu:nvidia-v100": 10}
+        assert cstr.to_gres_dict(input_str) == expected
+
+        input_str = "gres:gpu:2"
+        expected = {"gpu": 2}
+        assert cstr.to_gres_dict(input_str) == expected
+
+        input_str = "gres:gpu"
+        expected = {"gpu": 1}
+        assert cstr.to_gres_dict(input_str) == expected
+
+        input_str = "gres:gpu:INVALID_COUNT"
+        expected = {}
         assert cstr.to_gres_dict(input_str) == expected
 
     def test_gres_from_tres_dict(self):