Merge pull request GoogleCloudPlatform#4546 from sharabiani/gres-gpu-type

Add GPU type to Slurm's gres.config

Author: sharabiani

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/conf.py

@@ -367,28 +367,36 @@ def install_jobsubmit_lua(lkp: util.Lookup) -> None:
     util.chown_slurm(conf_file, 0o600)
 
 
-def gen_cloud_gres_conf(lkp: util.Lookup) -> None:
-    """generate cloud_gres.conf"""
+def gen_cloud_gres_conf_lines(lkp: util.Lookup) -> str:
+    """generate cloud_gres.conf's content"""
 
     gpu_nodes = defaultdict(list)
     for nodeset in lkp.cfg.nodeset.values():
         ti = lkp.template_info(nodeset.instance_template)
         gpu_count = ti.gpu.count if ti.gpu  else 0
+        gpu_type = ti.gpu.type if ti.gpu  else None
         if gpu_count:
-            gpu_nodes[gpu_count].append(lkp.nodelist(nodeset))
+            gpu_nodes[(gpu_count, gpu_type)].append(lkp.nodelist(nodeset))
 
     lines = [
         dict_to_conf(
             {
                 "NodeName": names,
                 "Name": "gpu",
-                "File": "/dev/nvidia{}".format(f"[0-{i-1}]" if i > 1 else "0"),
+                "Type": gpu_type,
+                "File": "/dev/nvidia{}".format(f"[0-{gpu_count-1}]" if gpu_count > 1 else "0"),
             }
         )
-        for i, names in gpu_nodes.items()
+        for (gpu_count, gpu_type), names in gpu_nodes.items()
     ]
     lines.append("\n")
-    content = FILE_PREAMBLE + "\n".join(lines)
+    return "\n".join(lines)
+
+
+def gen_cloud_gres_conf(lkp: util.Lookup) -> None:
+    """create cloud_gres.conf file"""
+
+    content = FILE_PREAMBLE + gen_cloud_gres_conf_lines(lkp)
 
     conf_file = lkp.etc_dir / "cloud_gres.conf"
     conf_file.write_text(content)

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/tests/test_conf.py

@@ -190,3 +190,25 @@ def test_conflines(cfg, want):
 
     cfg.cloud_parameters = addict.Dict(cfg.cloud_parameters)
     assert conf.conflines(util.Lookup(cfg)) == want
+
+
+@pytest.mark.parametrize(
+    "cfg,gputype,gpucount,want",
+    [
+        (TstCfg(),
+        "",
+        0,
+         "\n"),
+        (TstCfg(
+            nodeset={"turbo": TstNodeset("turbo")}
+        ), 
+        "Popov",
+        8,
+         "Name=gpu Type=Popov File=/dev/nvidia[0-7]\n\n"),
+    ])
+def test_gen_cloud_gres_conf_lines(cfg, gputype, gpucount, want):
+    lkp = util.Lookup(cfg)
+    lkp.template_info = Mock(return_value=TstTemplateInfo(
+        gpu=util.AcceleratorInfo(type=gputype, count=gpucount)
+    ))
+    assert conf.gen_cloud_gres_conf_lines(lkp) == want