testsuite/python/tests/test_107_5.py - SchedMD/slurm - Git at Google

 ############################################################################
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 ############################################################################
 """Test scancel of a het job step."""

 import re

 import pytest

 import atf


 @pytest.fixture(scope="module", autouse=True)
 def setup():
     # Het jobs are only started by backfill; don't wait for its full cycle
     atf.require_config_parameter_includes("SchedulerParameters", ("bf_interval", 1))
     # --stepmgr is silently dropped without it, see PrologFlags in slurm.conf(5)
     atf.require_config_parameter_includes("PrologFlags", "Contain")
     atf.require_nodes(2, [("CPUs", 2)])
     atf.require_slurm_running()


 def submit_hetjob_with_step(het_group, file_out, msg_done, stepmgr=False):
     """Submit a 2 component het job running a step on the given het groups.

     The batch script echoes msg_done once the step returns, so that the file
     contents tell the step being canceled apart from the whole job dying.

     Returns a (leader_job_id, component_job_ids) tuple.
     """
     suffix = "_stepmgr" if stepmgr else ""
     script = f"hetjob_step_{het_group.replace(',', '_')}{suffix}.sh"
     leader_opts = "-n1" + (" --stepmgr" if stepmgr else "")
     atf.make_bash_script(
         script,
         f"""
 #SBATCH {leader_opts}
 #SBATCH hetjob
 #SBATCH -n1

 srun --het-group={het_group} sleep infinity
 echo "{msg_done}"
 sleep infinity
         """,
     )
     leader_job_id = atf.submit_job_sbatch(
         f"-t3 --output={file_out} {script}", fatal=True
     )
     atf.wait_for_job_state(leader_job_id, "RUNNING", fatal=True)

     component_job_ids = atf.range_to_list(
         atf.get_job_parameter(leader_job_id, "HetJobIdSet", fatal=True)
     )
     assert (
         len(component_job_ids) == 2
     ), f"the het job must have 2 components, got {component_job_ids}"

     return leader_job_id, component_job_ids


 @pytest.mark.xfail(
     (25, 11) <= atf.get_version("sbin/slurmctld") < (26, 5, 4),
     reason="Issue 51060: since 4d70133431 (25.11) kill_job_step() propagated a"
     " whole-job cancel to het components instead of a step kill",
 )
 def test_cancel_hetjob_step_keeps_allocation():
     """Issue 51060: scancel of a het job step must not cancel the entire job.

     Canceling <leader_job_id>.<step_id> propagates the SIGKILL to the
     matching step of every het job component, but must only kill the steps:
     the component jobs (the allocations) have to keep running. A regression
     turned the propagation into a whole-job cancel of every component.
     """
     file_out = "hetjob_step.out"
     msg_done = "step_terminated"

     leader_job_id, component_job_ids = submit_hetjob_with_step(
         "0,1", file_out, msg_done
     )

     for job_id in component_job_ids:
         atf.wait_for_step(job_id, 0, fatal=True)

     atf.run_command(f"scancel {leader_job_id}.0", fatal=True)

     # srun returning (and the script moving on) proves both steps died
     atf.assert_file_contents(
         file_out,
         msg_done,
         contains=True,
         message=f"srun did not return after scancel {leader_job_id}.0: either the"
         " step of some het component was not canceled, or the whole component"
         " jobs were killed",
     )

     # Het step ids may be displayed with a +<comp> suffix on some versions
     for _ in atf.timer(fatal=True):
         if all(
             not any(
                 re.fullmatch(rf"{job_id}\.0(\+\d+)?", step_id)
                 for step_id in atf.get_steps(job_id, quiet=True)
             )
             for job_id in component_job_ids
         ):
             break

     for job_id in component_job_ids:
         assert (
             atf.get_job_parameter(job_id, "JobState") == "RUNNING"
         ), f"component job {job_id} must survive the cancellation of its step"


 @pytest.mark.xfail(
     atf.get_version("sbin/slurmctld") < (26, 5, 4),
     reason="Issue 51060: before 25.11 a component without the step made the"
     " whole request report an invalid job id, and since 4d70133431 the"
     " propagation canceled the component jobs instead of their steps",
 )
 @pytest.mark.parametrize("het_group", ["0", "1"])
 def test_cancel_hetjob_step_of_one_component(het_group):
     """Issue 51060: canceling a step only some het job components have.

     A step does not have to exist in every component, so canceling it
     through the leader must cancel it where it exists and report no error
     for the components that do not have it. The leader is one of the
     components that may not have it.
     """
     file_out = f"hetjob_step_one_{het_group}.out"
     msg_done = "step_terminated"

     leader_job_id, component_job_ids = submit_hetjob_with_step(
         het_group, file_out, msg_done
     )

     atf.wait_for_step(component_job_ids[int(het_group)], 0, fatal=True)

     result = atf.run_command(f"scancel {leader_job_id}.0", fatal=True)

     assert "Invalid job id" not in result["stderr"], (
         "canceling a step that only some het job components have must not be"
         f" an error. stderr:\n{result['stderr']}"
     )

     atf.assert_file_contents(
         file_out,
         msg_done,
         contains=True,
         message=f"srun did not return after scancel {leader_job_id}.0, so the"
         " step of the component owning it was not canceled",
     )

     for job_id in component_job_ids:
         assert (
             atf.get_job_parameter(job_id, "JobState") == "RUNNING"
         ), f"component job {job_id} must survive the cancellation of the step"


 @pytest.mark.xfail(
     (25, 11) <= atf.get_version("sbin/slurmctld") < (26, 5, 4),
     reason="Issue 51060: since 4d70133431 (25.11) the propagation to the het"
     " components canceled the component jobs for a step id that none of them"
     " has",
 )
 def test_cancel_hetjob_step_of_no_component():
     """Issue 51060: canceling a step id that no het job component has.

     Canceling a step no component has must leave the existing steps running.
     Without stepmgr that request is reported as an invalid job id; with
     stepmgr the controller forwards the cancel to the components' stepmgr and
     cannot tell synchronously that no component has the step.
     """
     file_out = "hetjob_step_none.out"
     msg_done = "step_terminated"

     leader_job_id, component_job_ids = submit_hetjob_with_step(
         "0,1", file_out, msg_done
     )

     # Keyed off the job, not SlurmctldParameters: before 26.11 a het job does
     # not get a stepmgr even where enable_stepmgr is configured.
     jobs = atf.get_jobs(leader_job_id, fatal=True)
     stepmgr_enabled = jobs[leader_job_id].get("StepMgrEnabled") == "Yes"

     for job_id in component_job_ids:
         atf.wait_for_step(job_id, 0, fatal=True)

     result = atf.run_command(
         f"scancel {leader_job_id}.7", fatal=True, xfail=not stepmgr_enabled
     )

     if stepmgr_enabled:
         assert result["exit_code"] == 0, (
             "with stepmgr the cancel is forwarded to the components' stepmgr"
             " and must be accepted, got exit code"
             f" {result['exit_code']}. stderr:\n{result['stderr']}"
         )
     else:
         assert "Invalid job id" in result["stderr"], (
             "canceling a step id that no het job component has must report an"
             f" invalid job id. stderr:\n{result['stderr']}"
         )

     for job_id in component_job_ids:
         assert any(
             re.fullmatch(rf"{job_id}\.0(\+\d+)?", step_id)
             for step_id in atf.get_steps(job_id)
         ), (
             f"step {job_id}.0 must keep running after canceling a step id that"
             " no het job component has"
         )


 # The client tools upgrade last, so gating on srun >= 26.11 already implies
 # slurmd and slurmctld are >= 26.11.
 @pytest.mark.skipif(
     atf.get_version("bin/srun") < (26, 11),
     reason="Issue 50976: het jobs only run with stepmgr since 26.11",
 )
 def test_cancel_hetjob_step_with_stepmgr():
     """Canceling a het job step of a --stepmgr het job keeps the allocations.

     With stepmgr the controller forwards the cancel to the components' step
     managers rather than killing the steps itself, so the propagation is a
     different path from the controller-managed one.
     """
     file_out = "hetjob_step_stepmgr.out"
     msg_done = "step_terminated"

     leader_job_id, component_job_ids = submit_hetjob_with_step(
         "0,1", file_out, msg_done, stepmgr=True
     )

     jobs = atf.get_jobs(leader_job_id, fatal=True)
     assert jobs[leader_job_id].get("StepMgrEnabled") == "Yes", (
         "the het job must have a stepmgr for this test to exercise the stepmgr"
         " cancel path"
     )

     for job_id in component_job_ids:
         atf.wait_for_step(job_id, 0, fatal=True)

     atf.run_command(f"scancel {leader_job_id}.0", fatal=True)

     atf.assert_file_contents(
         file_out,
         msg_done,
         contains=True,
         message=f"srun did not return after scancel {leader_job_id}.0 on a"
         " stepmgr het job, so the step of some component was not canceled",
     )

     for job_id in component_job_ids:
         assert (
             atf.get_job_parameter(job_id, "JobState", fatal=True) == "RUNNING"
         ), f"component job {job_id} must survive the cancellation of its step"
	############################################################################
	# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	############################################################################
	"""Test scancel of a het job step."""

	import re

	import pytest

	import atf


	@pytest.fixture(scope="module", autouse=True)
	def setup():
	# Het jobs are only started by backfill; don't wait for its full cycle
	atf.require_config_parameter_includes("SchedulerParameters", ("bf_interval", 1))
	# --stepmgr is silently dropped without it, see PrologFlags in slurm.conf(5)
	atf.require_config_parameter_includes("PrologFlags", "Contain")
	atf.require_nodes(2, [("CPUs", 2)])
	atf.require_slurm_running()


	def submit_hetjob_with_step(het_group, file_out, msg_done, stepmgr=False):
	"""Submit a 2 component het job running a step on the given het groups.

	The batch script echoes msg_done once the step returns, so that the file
	contents tell the step being canceled apart from the whole job dying.

	Returns a (leader_job_id, component_job_ids) tuple.
	"""
	suffix = "_stepmgr" if stepmgr else ""
	script = f"hetjob_step_{het_group.replace(',', '_')}{suffix}.sh"
	leader_opts = "-n1" + (" --stepmgr" if stepmgr else "")
	atf.make_bash_script(
	script,
	f"""
	#SBATCH {leader_opts}
	#SBATCH hetjob
	#SBATCH -n1

	srun --het-group={het_group} sleep infinity
	echo "{msg_done}"
	sleep infinity
	""",
	)
	leader_job_id = atf.submit_job_sbatch(
	f"-t3 --output={file_out} {script}", fatal=True
	)
	atf.wait_for_job_state(leader_job_id, "RUNNING", fatal=True)

	component_job_ids = atf.range_to_list(
	atf.get_job_parameter(leader_job_id, "HetJobIdSet", fatal=True)
	)
	assert (
	len(component_job_ids) == 2
	), f"the het job must have 2 components, got {component_job_ids}"

	return leader_job_id, component_job_ids


	@pytest.mark.xfail(
	(25, 11) <= atf.get_version("sbin/slurmctld") < (26, 5, 4),
	reason="Issue 51060: since 4d70133431 (25.11) kill_job_step() propagated a"
	" whole-job cancel to het components instead of a step kill",
	)
	def test_cancel_hetjob_step_keeps_allocation():
	"""Issue 51060: scancel of a het job step must not cancel the entire job.

	Canceling <leader_job_id>.<step_id> propagates the SIGKILL to the
	matching step of every het job component, but must only kill the steps:
	the component jobs (the allocations) have to keep running. A regression
	turned the propagation into a whole-job cancel of every component.
	"""
	file_out = "hetjob_step.out"
	msg_done = "step_terminated"

	leader_job_id, component_job_ids = submit_hetjob_with_step(
	"0,1", file_out, msg_done
	)

	for job_id in component_job_ids:
	atf.wait_for_step(job_id, 0, fatal=True)

	atf.run_command(f"scancel {leader_job_id}.0", fatal=True)

	# srun returning (and the script moving on) proves both steps died
	atf.assert_file_contents(
	file_out,
	msg_done,
	contains=True,
	message=f"srun did not return after scancel {leader_job_id}.0: either the"
	" step of some het component was not canceled, or the whole component"
	" jobs were killed",
	)

	# Het step ids may be displayed with a +<comp> suffix on some versions
	for _ in atf.timer(fatal=True):
	if all(
	not any(
	re.fullmatch(rf"{job_id}\.0(\+\d+)?", step_id)
	for step_id in atf.get_steps(job_id, quiet=True)
	)
	for job_id in component_job_ids
	):
	break

	for job_id in component_job_ids:
	assert (
	atf.get_job_parameter(job_id, "JobState") == "RUNNING"
	), f"component job {job_id} must survive the cancellation of its step"


	@pytest.mark.xfail(
	atf.get_version("sbin/slurmctld") < (26, 5, 4),
	reason="Issue 51060: before 25.11 a component without the step made the"
	" whole request report an invalid job id, and since 4d70133431 the"
	" propagation canceled the component jobs instead of their steps",
	)
	@pytest.mark.parametrize("het_group", ["0", "1"])
	def test_cancel_hetjob_step_of_one_component(het_group):
	"""Issue 51060: canceling a step only some het job components have.

	A step does not have to exist in every component, so canceling it
	through the leader must cancel it where it exists and report no error
	for the components that do not have it. The leader is one of the
	components that may not have it.
	"""
	file_out = f"hetjob_step_one_{het_group}.out"
	msg_done = "step_terminated"

	leader_job_id, component_job_ids = submit_hetjob_with_step(
	het_group, file_out, msg_done
	)

	atf.wait_for_step(component_job_ids[int(het_group)], 0, fatal=True)

	result = atf.run_command(f"scancel {leader_job_id}.0", fatal=True)

	assert "Invalid job id" not in result["stderr"], (
	"canceling a step that only some het job components have must not be"
	f" an error. stderr:\n{result['stderr']}"
	)

	atf.assert_file_contents(
	file_out,
	msg_done,
	contains=True,
	message=f"srun did not return after scancel {leader_job_id}.0, so the"
	" step of the component owning it was not canceled",
	)

	for job_id in component_job_ids:
	assert (
	atf.get_job_parameter(job_id, "JobState") == "RUNNING"
	), f"component job {job_id} must survive the cancellation of the step"


	@pytest.mark.xfail(
	(25, 11) <= atf.get_version("sbin/slurmctld") < (26, 5, 4),
	reason="Issue 51060: since 4d70133431 (25.11) the propagation to the het"
	" components canceled the component jobs for a step id that none of them"
	" has",
	)
	def test_cancel_hetjob_step_of_no_component():
	"""Issue 51060: canceling a step id that no het job component has.

	Canceling a step no component has must leave the existing steps running.
	Without stepmgr that request is reported as an invalid job id; with
	stepmgr the controller forwards the cancel to the components' stepmgr and
	cannot tell synchronously that no component has the step.
	"""
	file_out = "hetjob_step_none.out"
	msg_done = "step_terminated"

	leader_job_id, component_job_ids = submit_hetjob_with_step(
	"0,1", file_out, msg_done
	)

	# Keyed off the job, not SlurmctldParameters: before 26.11 a het job does
	# not get a stepmgr even where enable_stepmgr is configured.
	jobs = atf.get_jobs(leader_job_id, fatal=True)
	stepmgr_enabled = jobs[leader_job_id].get("StepMgrEnabled") == "Yes"

	for job_id in component_job_ids:
	atf.wait_for_step(job_id, 0, fatal=True)

	result = atf.run_command(
	f"scancel {leader_job_id}.7", fatal=True, xfail=not stepmgr_enabled
	)

	if stepmgr_enabled:
	assert result["exit_code"] == 0, (
	"with stepmgr the cancel is forwarded to the components' stepmgr"
	" and must be accepted, got exit code"
	f" {result['exit_code']}. stderr:\n{result['stderr']}"
	)
	else:
	assert "Invalid job id" in result["stderr"], (
	"canceling a step id that no het job component has must report an"
	f" invalid job id. stderr:\n{result['stderr']}"
	)

	for job_id in component_job_ids:
	assert any(
	re.fullmatch(rf"{job_id}\.0(\+\d+)?", step_id)
	for step_id in atf.get_steps(job_id)
	), (
	f"step {job_id}.0 must keep running after canceling a step id that"
	" no het job component has"
	)


	# The client tools upgrade last, so gating on srun >= 26.11 already implies
	# slurmd and slurmctld are >= 26.11.
	@pytest.mark.skipif(
	atf.get_version("bin/srun") < (26, 11),
	reason="Issue 50976: het jobs only run with stepmgr since 26.11",
	)
	def test_cancel_hetjob_step_with_stepmgr():
	"""Canceling a het job step of a --stepmgr het job keeps the allocations.

	With stepmgr the controller forwards the cancel to the components' step
	managers rather than killing the steps itself, so the propagation is a
	different path from the controller-managed one.
	"""
	file_out = "hetjob_step_stepmgr.out"
	msg_done = "step_terminated"

	leader_job_id, component_job_ids = submit_hetjob_with_step(
	"0,1", file_out, msg_done, stepmgr=True
	)

	jobs = atf.get_jobs(leader_job_id, fatal=True)
	assert jobs[leader_job_id].get("StepMgrEnabled") == "Yes", (
	"the het job must have a stepmgr for this test to exercise the stepmgr"
	" cancel path"
	)

	for job_id in component_job_ids:
	atf.wait_for_step(job_id, 0, fatal=True)

	atf.run_command(f"scancel {leader_job_id}.0", fatal=True)

	atf.assert_file_contents(
	file_out,
	msg_done,
	contains=True,
	message=f"srun did not return after scancel {leader_job_id}.0 on a"
	" stepmgr het job, so the step of some component was not canceled",
	)

	for job_id in component_job_ids:
	assert (
	atf.get_job_parameter(job_id, "JobState", fatal=True) == "RUNNING"
	), f"component job {job_id} must survive the cancellation of its step"