| ############################################################################ |
| # Copyright (C) SchedMD LLC. |
| ############################################################################ |
| import atf |
| import pytest |
| import re |
| |
| |
| @pytest.fixture(scope="module", autouse=True) |
| def setup(): |
| atf.require_accounting() |
| atf.require_nodes(3, [("CPUs", 8), ("RealMemory", 80)]) |
| |
| # only to make the test faster |
| atf.require_config_parameter("SchedulerParameters", "bf_interval=1") |
| |
| atf.require_slurm_running() |
| |
| |
| def test_hetjob(): |
| file_in = "hetjob.in" |
| atf.make_bash_script( |
| file_in, |
| """ |
| #SBATCH --cpus-per-task=4 --mem-per-cpu=10 --ntasks=1 |
| #SBATCH hetjob |
| #SBATCH --cpus-per-task=2 --mem-per-cpu=2 --ntasks=1 -t1 |
| #SBATCH hetjob |
| #SBATCH --cpus-per-task=1 --mem-per-cpu=6 --ntasks=1 -t1 |
| |
| srun sleep 60 |
| """, |
| ) |
| leader_job_id = atf.submit_job_sbatch(f"-t1 {file_in}", fatal=True) |
| atf.wait_for_job_state(leader_job_id, "RUNNING", fatal=True) |
| |
| jobs_dict = atf.get_jobs() |
| |
| # Verify details about leader job |
| assert leader_job_id in jobs_dict |
| leader_job_dict = jobs_dict[leader_job_id] |
| assert leader_job_dict["JobId"] == leader_job_id |
| assert leader_job_dict["HetJobId"] == leader_job_id |
| assert leader_job_dict["HetJobOffset"] == 0 |
| het_job_expression = leader_job_dict["HetJobIdSet"] |
| assert re.search(r"[\d,-]+", het_job_expression) is not None |
| assert leader_job_dict["CPUs/Task"] == 4 |
| assert leader_job_dict["MinMemoryCPU"] == "10M" |
| het_job_component_list = atf.range_to_list(het_job_expression) |
| het_job_component_list.remove(leader_job_id) |
| assert len(het_job_component_list) == 2 |
| |
| # Verify details about first component job |
| component1_job_id = het_job_component_list[0] |
| assert component1_job_id in jobs_dict |
| component1_job_dict = jobs_dict[component1_job_id] |
| assert component1_job_dict["JobId"] == component1_job_id |
| assert component1_job_dict["HetJobId"] == leader_job_id |
| assert component1_job_dict["HetJobOffset"] == 1 |
| assert re.search(r"[\d,-]+", component1_job_dict["HetJobIdSet"]) is not None |
| assert component1_job_dict["CPUs/Task"] == 2 |
| assert component1_job_dict["MinMemoryCPU"] == "2M" |
| |
| # Verify details about second component job |
| component2_job_id = het_job_component_list[1] |
| assert component2_job_id in jobs_dict |
| component2_job_dict = jobs_dict[component2_job_id] |
| assert component2_job_dict["JobId"] == component2_job_id |
| assert component2_job_dict["HetJobId"] == leader_job_id |
| assert component2_job_dict["HetJobOffset"] == 2 |
| assert re.search(r"[\d,-]+", component2_job_dict["HetJobIdSet"]) is not None |
| assert component2_job_dict["CPUs/Task"] == 1 |
| assert component2_job_dict["MinMemoryCPU"] == "6M" |
