Merge pull request #588 from shreyb/issue-505

Add --gpu flag to submit jobs to GPU resources

lib/get_parser.py

@@ -305,7 +305,9 @@ def get_parser() -> argparse.ArgumentParser:
         default=os.environ.get("CMTCONFIG", ""),
         help=" Set up minervasoft release built with cmt configuration. default is $CMTCONFIG",
     )
-    parser.add_argument("--cpu", help="request worker nodes have at least NUMBER cpus")
+    parser.add_argument(
+        "--cpu", metavar="NUMBER", help="Request worker nodes have at least NUMBER cpus"
+    )
     parser.add_argument(
         "--dag",
         help="submit and run a dagNabbit input file",
@@ -413,6 +415,9 @@ def get_parser() -> argparse.ArgumentParser:
         help="generate and mail a summary report of completed/failed/removed"
         " jobs in a DAG",
     )
+    parser.add_argument(
+        "--gpu", metavar="NUMBER", help="request worker nodes have at least NUMBER cpus"
+    )
     parser.add_argument(
         "-L", "--log-file", "--log_file", help="Log file to hold log output from job."
     )

templates/simple/simple.cmd

@@ -26,6 +26,10 @@ transfer_input_files = {{executable|basename}}
 transfer_output_files = .empty_file
 when_to_transfer_output = ON_EXIT_OR_EVICT
 {%if    cpu is defined and cpu %}request_cpus = {{cpu}}{%endif%}
+{%if    gpu is defined and gpu %}
+request_GPUs = {{gpu}}
++RequestGPUs = {{gpu}}
+{%endif%}
 {%if memory is defined and memory %}request_memory = {{memory}}{%endif%}
 {%if   disk is defined and disk %}request_disk = {{disk}}KB{%endif%}
 {%if     OS is defined and OS %}+DesiredOS="{{OS}}"{%endif%}

tests/test_get_parser_unit.py

@@ -176,6 +176,8 @@ def all_test_args():
         "--generate-email-summary",
         "--global-pool",
         "dune",
+        "--gpu",
+        "xxgpuxx",
         "--group",
         "xxgroupxx",
         "--job-info",

tests/test_jobsub_submit_unit.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 import sys
 
 import pytest
@@ -115,3 +116,46 @@ def test_do_dataset_defaults_1(self):
         submit_support.do_dataset_defaults(varg)
         for var in ["PROJECT", "DATASET", "USER", "GROUP", "STATION"]:
             assert repr(varg["environment"]).find("SAM_%s" % var) > 0
+
+    @pytest.mark.unit
+    def test_do_gpus(self, tmp_path, monkeypatch):
+        """make sure --gpu NUMBER sets arguments it's supposed to"""
+        # Test parameters
+        num_gpus = 42
+        wanted_strings = ("request_GPUs = 42", "+RequestGPUs = 42")
+
+        # Set up our temp test directories and copy the template file
+        temp_prefix = (
+            tmp_path / "indir"
+        )  # Root directory of template files.  Will be mock for submit_support.PREFIX
+        indir = (
+            temp_prefix / "templates" / "simple"
+        )  # Directory where the template file will be copied
+        indir.mkdir(parents=True)
+        shutil.copy(f"{submit_support.PREFIX}/templates/simple/simple.cmd", indir)
+
+        outdir = (
+            tmp_path / "outdir"
+        )  # Directory where the rendered template will be written
+        outdir.mkdir()
+
+        # Set up the vargs for our template
+        varg = TestUnit.test_vargs.copy()
+        varg.update(TestUnit.test_extra_template_args)
+        varg["gpu"] = num_gpus
+        varg["no_submit"] = True
+        varg["outdir"] = outdir
+
+        # Extra required arg for template to render
+        varg["mail"] = "Never"
+
+        # Render the template like in a real job submission, with PREFIX mocked
+        with monkeypatch.context() as m:
+            m.setattr(submit_support, "PREFIX", str(temp_prefix))
+            submit_support.jobsub_submit_simple(varg, "fakeschedd.example.org")
+
+        with open(outdir / "simple.cmd", "r") as f:
+            cmd_text = f.read()
+
+        for s in wanted_strings:
+            assert s in cmd_text