diff --git a/docs/releases/changelog-dev.md b/docs/releases/changelog-dev.md index 7185c775e..243ab42f5 100644 --- a/docs/releases/changelog-dev.md +++ b/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 diff --git a/src/qililab/slurm.py b/src/qililab/slurm.py index 2b942db6a..90f28d853 100644 --- a/src/qililab/slurm.py +++ b/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,6 +86,7 @@ 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 @@ -92,6 +94,8 @@ def submit_job(line: str, cell: str, local_ns: dict) -> None: 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 diff --git a/tests/test_slurm.py b/tests/test_slurm.py index d7fec0a37..4dfc76fc4 100644 --- a/tests/test_slurm.py +++ b/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,7 +113,7 @@ 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") @@ -110,5 +121,5 @@ def test_setting_parameters(self, ip): 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"}, )