modify magic method + unit tests (#828)

* modify magic method + unit tests

* imporove error message

* code quality

* update changelog

* solve space typo

docs/releases/changelog-dev.md

@@ -2,6 +2,10 @@
 
 ### New features since last release
 
+- Support GRES in %%submit_job magic method
+
+[#828](https://github.com/qilimanjaro-tech/qililab/pull/828)
+
 - Added intermediate frequency to single input lines on qm. The default is 0 (this prevents some bugs from qua-qm). Now it is possible to use the set_parameter IF and qm.set_frequency for buses with single_input.
 
 [#807](https://github.com/qilimanjaro-tech/qililab/pull/807)
@@ -101,7 +105,7 @@ instruments:
     ...
 ```
 
-  [#826](https://github.com/qilimanjaro-tech/qililab/pull/826)
+[#826](https://github.com/qilimanjaro-tech/qililab/pull/826)
 
 ### Improvements
 

src/qililab/slurm.py

@@ -45,6 +45,7 @@ def is_variable_used(code, variable):
     " the results of the job, you need to call `variable.result()`.",
 )
 @argument("-p", "--partition", help="Name of the partition where you want to execute the SLURM job.")
+@argument("-g", "--gres", default=None, help="GRES (chip) where you want to execute the SLURM job.")
 @argument("-n", "--name", default="submitit", help="Name of the slurm job.")
 @argument("-t", "--time", default=120, help="Time limit (in minutes) of the job.")
 @argument(
@@ -85,13 +86,16 @@ def submit_job(line: str, cell: str, local_ns: dict) -> None:
     args = parse_argstring(submit_job, line)
     output = args.output
     partition = args.partition
+    gres = args.gres
     job_name = args.name
     time_limit = int(args.time)
     folder_path = args.logs
     execution_env = args.execution_environment
     begin_time = args.begin
     low_priority = args.low_priority
 
+    if gres is None:
+        raise ValueError("GRES needs to be provided! See the available ones typing 'sinfo -o '%G'' in the terminal")
     nice_factor = 0
     if low_priority in ["True", "true"]:
         nice_factor = 1000000  # this ensures Lab jobs have 0 priority, same as QaaS jobs
@@ -108,7 +112,7 @@ def submit_job(line: str, cell: str, local_ns: dict) -> None:
         slurm_partition=partition,
         name=job_name,
         timeout_min=time_limit,
-        slurm_additional_parameters={"begin": begin_time, "nice": nice_factor},
+        slurm_additional_parameters={"begin": begin_time, "nice": nice_factor, "gres": f"{gres}:1"},
     )
 
     # Compile the code defined above

tests/test_slurm.py

@@ -4,9 +4,8 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
-from IPython.testing.globalipapp import start_ipython
-
 import qililab as ql
+from IPython.testing.globalipapp import start_ipython
 
 slurm_job_data_test = "slurm_job_data_test"
 
@@ -34,7 +33,7 @@ def test_submit_job(self, ip):
         ip.run_cell(raw_cell="a=1\nb=1")
         ip.run_cell_magic(
             magic_name="submit_job",
-            line=f"-o results -p debug -l {slurm_job_data_test} -n unit_test -e local",
+            line=f"-o results -g aQPU1 -l {slurm_job_data_test} -n unit_test -e local",
             cell="results = a+b ",
         )
         time.sleep(4)
@@ -48,7 +47,19 @@ def test_submit_job_output_not_assigned(self, ip):
         ):
             ip.run_cell_magic(
                 magic_name="submit_job",
-                line=f"-o results -p debug -l {slurm_job_data_test} -n unit_test -e local",
+                line=f"-o results -g aQPU1 -l {slurm_job_data_test} -n unit_test -e local",
+                cell="a+b",
+            )
+
+    def test_submit_job_no_gres_provided(self, ip):
+        """Check ValueError is raised in case GRES is not provided."""
+        ip.run_cell(raw_cell="a=1\nb=1")
+        with pytest.raises(
+            ValueError, match="GRES needs to be provided! See the available ones typing 'sinfo -o '%G'' in the terminal"
+        ):
+            ip.run_cell_magic(
+                magic_name="submit_job",
+                line=f"-o results -l {slurm_job_data_test} -n unit_test -e local",
                 cell="a+b",
             )
 
@@ -59,7 +70,7 @@ def test_submit_job_with_random_file_in_logs_folder(self, ip):
         )
         ip.run_cell_magic(
             magic_name="submit_job",
-            line=f"-o results -p debug -l {slurm_job_data_test} -n unit_test -e local",
+            line=f"-o results -g aQPU1 -l {slurm_job_data_test} -n unit_test -e local",
             cell="results=a+b",
         )
         time.sleep(4)  # give time to ensure processes are finished
@@ -72,7 +83,7 @@ def test_submit_job_delete_info_from_past_jobs(self, ip):
         for _ in range(int(ql.slurm.num_files_to_keep / 4)):
             ip.run_cell_magic(
                 magic_name="submit_job",
-                line=f"-o results -p debug -l {slurm_job_data_test} -n unit_test -e local",
+                line=f"-o results -g aQPU1 -l {slurm_job_data_test} -n unit_test -e local",
                 cell="results=a+b",
             )
             time.sleep(2)  # give time submitit to create the files
@@ -83,7 +94,7 @@ def test_submit_job_delete_info_from_past_jobs(self, ip):
         )
         ip.run_cell_magic(
             magic_name="submit_job",
-            line=f"-o results -p debug -l {slurm_job_data_test} -n unit_test -e local",
+            line=f"-o results -g aQPU1 -l {slurm_job_data_test} -n unit_test -e local",
             cell="results=a+b",
         )
         time.sleep(2)
@@ -102,13 +113,13 @@ def test_setting_parameters(self, ip):
             with patch("qililab.slurm.os"):
                 ip.run_cell_magic(
                     magic_name="submit_job",
-                    line=f"-o results -p debug -l {slurm_job_data_test} -n unit_test -e local -t 5 --begin now+1hour -lp true",
+                    line=f"-o results -p debug -l {slurm_job_data_test} -n unit_test -e local -t 5 --begin now+1hour -lp true -g aQPU1",
                     cell="results=a+b",
                 )
         executor.assert_called_once_with(folder=slurm_job_data_test, cluster="local")
         executor_instance.update_parameters.assert_called_once_with(
             slurm_partition="debug",
             name="unit_test",
             timeout_min=5,
-            slurm_additional_parameters={"begin": "now+1hour", "nice": 1000000},
+            slurm_additional_parameters={"begin": "now+1hour", "nice": 1000000, "gres":"aQPU1:1"},
         )