| ############################################################################ |
| # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| ############################################################################ |
| """Test PMIx abort propagation across hetjob components.""" |
| |
| import errno |
| import re |
| |
| import pytest |
| |
| import atf |
| |
| |
| @pytest.fixture(scope="module", autouse=True) |
| def setup(): |
| atf.require_lmod() |
| atf.module_load("openmpi") |
| atf.require_accounting() |
| atf.require_config_parameter("MpiDefault", "pmix") |
| atf.require_config_parameter("KillWait", "5") |
| # Hetjobs are only started by backfill; don't wait for its full cycle |
| atf.require_config_parameter_includes("SchedulerParameters", ("bf_interval", 1)) |
| atf.require_nodes(2, [("CPUs", 2)]) |
| atf.require_slurm_running() |
| |
| |
| @pytest.mark.skipif( |
| atf.get_version("sbin/slurmd") < (26, 5, 4), |
| reason="Issue 51060: PMIx hetjob abort propagation fixed in 26.05.4", |
| ) |
| @pytest.mark.parametrize("mpi_program", ["mpi_signal_test"], indirect=True) |
| @pytest.mark.parametrize( |
| "abort_arg,expected_status,expected_state", |
| [ |
| ("abort", 42, "CANCELLED"), |
| ("abort0", 0, "COMPLETED"), |
| ], |
| ) |
| def test_hetjob_abort_propagation( |
| mpi_program, abort_arg, expected_status, expected_state |
| ): |
| """Issue 51060: MPI_Abort() in one het component terminates all of them. |
| |
| Launches a two-component hetjob where global rank 0 (in component 0) |
| calls MPI_Abort() while every other rank (including all ranks of |
| component 1) blocks forever in pause(). The abort must terminate the |
| steps of *all* components, otherwise srun never returns: |
| - abort (status 42): srun must exit 42. |
| - abort0 (status 0): KILL_NO_SIG_FAIL must survive the propagation |
| to the other components, so the whole run must report success |
| (exit code 0) exactly as it does for a non-hetjob. |
| """ |
| # OpenMPI's sm BTL names its /dev/shm segments after the hostname, which |
| # collides under multi-slurmd, so force TCP for the MPI transport |
| result = atf.run_command( |
| f"srun -t2 --mpi=pmix -n2 {mpi_program} {abort_arg}" |
| f" : -t2 -n2 {mpi_program} {abort_arg}", |
| env_vars="OMPI_MCA_btl=self,tcp", |
| xfail=expected_status != 0, |
| # Bound a possible hang: not propagating the abort blocks srun forever |
| timeout=90, |
| ) |
| |
| # Checked before the timeout so that a job that never ran is not reported |
| # as a propagation failure |
| assert ( |
| f"rank0_calling_abort_{expected_status}" in result["stdout"] |
| ), f"rank 0 never reached MPI_Abort(). stdout:\n{result['stdout']}" |
| |
| assert result["exit_code"] != errno.ETIMEDOUT, ( |
| "srun hung: the abort in component 0 did not terminate the tasks" |
| f" of the other het components. stderr:\n{result['stderr']}" |
| ) |
| |
| # Stepd emits the message once per node, so count it per owning job id: |
| # two distinct ones prove the SIG_TERM_KILL reached the step of every |
| # component, rather than srun returning for some other reason |
| killed_job_ids = set( |
| re.findall( |
| r"STEP (\d+)\.\d+(?:\+\d+)? ON \S+ CANCELLED AT \S+ DUE TO TASK FAILURE", |
| result["stderr"], |
| ) |
| ) |
| assert len(killed_job_ids) == 2, ( |
| "the abort must terminate the step of every het component, but only" |
| f" {sorted(killed_job_ids)} got it. stderr:\n{result['stderr']}" |
| ) |
| |
| assert result["exit_code"] == expected_status, ( |
| f"MPI_Abort({expected_status}) in a hetjob must report the abort" |
| " status as the srun exit code, exactly as it does for a non-hetjob" |
| f" (KILL_NO_SIG_FAIL lost?). stderr:\n{result['stderr']}" |
| ) |
| |
| # srun's exit code does not prove KILL_NO_SIG_FAIL survived all the way |
| # into what the controller hands slurmdbd, which is what users see later. |
| # sacct reports het steps as <leader>+<component>.<step>, so collect them |
| # by job id rather than looking each one up as <job_id>.0 |
| for _ in atf.timer(fatal=True): |
| step_states = {} |
| for job_id in killed_job_ids: |
| for row in atf.run_command_output( |
| f"sacct -j {job_id} --noheader -P -o JobID,State", |
| fatal=True, |
| quiet=True, |
| ).splitlines(): |
| step_id, _, state = row.partition("|") |
| if step_id.endswith(".0"): |
| step_states[step_id] = state |
| # slurmdbd holds the step from launch time, so waiting for the rows to |
| # exist would assert against a state that is not the final one yet |
| if len(step_states) == 2 and not any( |
| state.startswith(("PENDING", "RUNNING")) for state in step_states.values() |
| ): |
| break |
| |
| for step_id, state in sorted(step_states.items()): |
| assert state.startswith(expected_state), ( |
| f"step {step_id} of the hetjob must be accounted as" |
| f" {expected_state} after MPI_Abort({expected_status})," |
| f" got State={state!r}" |
| ) |
| |
| |
| def test_hetjob_kill_on_bad_exit(): |
| """Test --kill-on-bad-exit kills the tasks of every hetjob component. |
| |
| --kill-on-bad-exit is propagated to the subsequent components of a |
| hetjob, so a task exiting non-zero in one component must terminate the |
| tasks of all of them instead of leaving them running. |
| |
| srun signals the step of each component itself, so this is a contract |
| test rather than a regression test of the hetjob signal propagation |
| that slurmctld does for a step kill aimed at the job leader. |
| """ |
| marker = "SHOULD_NOT_BE_HERE" |
| failing_script = "hetjob_bad_exit.sh" |
| surviving_script = "hetjob_survivor.sh" |
| # The sleep is the detection window, not a keep-alive: a task that escapes |
| # the kill has to outlive it to print the marker, and has to print it well |
| # before the command timeout below |
| atf.make_bash_script( |
| failing_script, |
| f""" |
| if [ "$SLURM_PROCID" = "0" ]; then |
| exit 2 |
| fi |
| sleep 60 |
| echo "{marker}" |
| """, |
| ) |
| atf.make_bash_script( |
| surviving_script, |
| f""" |
| sleep 60 |
| echo "{marker}" |
| """, |
| ) |
| |
| result = atf.run_command( |
| f"srun -t2 --kill-on-bad-exit -n2 ./{failing_script}" |
| f" : -t2 -n2 ./{surviving_script}", |
| xfail=True, |
| # Bound a possible hang: surviving tasks sleep until the time limit |
| timeout=90, |
| ) |
| |
| assert result["exit_code"] != errno.ETIMEDOUT, ( |
| "srun hung: the task failure in component 0 did not terminate the" |
| f" tasks of the other het components. stderr:\n{result['stderr']}" |
| ) |
| |
| # Absent this, a step that never launched also leaves the marker unprinted |
| assert "task 0: Exited with exit code 2" in result["stderr"], ( |
| "rank 0 of component 0 did not exit non-zero, so --kill-on-bad-exit" |
| f" was never exercised. stderr:\n{result['stderr']}" |
| ) |
| |
| assert marker not in result["stdout"], ( |
| "--kill-on-bad-exit must terminate the tasks of every het component." |
| f" stdout:\n{result['stdout']}" |
| ) |
