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

 ############################################################################
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 ############################################################################
 """Heterogeneous job + stepmgr-on-stepd: verify job-global step ids.

 With --stepmgr on a heterogeneous job, each component's stepmgr asks the
 leader for the next step id (REQUEST_HET_STEP_ID). Step ids must remain
 unique and contiguous across the whole het job, not per-component.

 Without that RPC, sequential srun --het-group calls would reuse ids
 0,1,0,1,... per component instead of 0,1,2,3 across the job.
 """

 import collections
 import re

 import pytest

 import atf

 pytestmark = pytest.mark.slow

 # Echoed after the last srun so we can tell the batch script's whole output
 # has been flushed before parsing it.
 OUTPUT_MARKER = "HETJOB_STEPS_DONE"


 @pytest.fixture(scope="module", autouse=True)
 def setup():
     # The client tools upgrade last, so gating on srun >= 26.11 already
     # implies slurmd and slurmctld are >= 26.11.
     atf.require_version(
         (26, 11),
         component="bin/srun",
         reason="Issue 50976: hetjob + stepmgr requires srun 26.11+",
     )
     atf.require_nodes(3)
     atf.require_config_parameter_includes("PrologFlags", "Contain")
     # Het jobs are only started by backfill; don't wait for its full cycle
     atf.require_config_parameter_includes("SchedulerParameters", ("bf_interval", 1))
     atf.require_config_parameter("SchedulerType", "sched/backfill")
     atf.require_slurm_running()


 def _het_opts(component_nodes):
     """Build het job options with --stepmgr on the leader only."""

     return " : ".join(
         f"-N{nodes} --ntasks-per-node=1 -t2" + (" --stepmgr" if offset == 0 else "")
         for offset, nodes in enumerate(component_nodes)
     )


 @pytest.mark.parametrize(
     "command, component_nodes",
     [
         ("sbatch", [1, 1]),
         ("salloc", [1, 1]),
         ("sbatch", [1, 1, 1]),
         ("sbatch", [2, 1]),
     ],
     ids=["sbatch-2comp", "salloc-2comp", "sbatch-3comp", "sbatch-multinode"],
 )
 def test_hetjob_stepmgr_job_global_step_ids(command, component_nodes):
     """Sequential sruns across all components yield unique, contiguous step ids."""

     test_id = f"{command}_{'-'.join(str(n) for n in component_nodes)}"
     file_in = f"hetjob_stepmgr_{test_id}.in"
     file_out = f"hetjob_stepmgr_{test_id}.out"

     # Two rounds over every component, so a per-component counter would
     # repeat ids instead of continuing the job-global sequence.
     sruns = "\n".join(
         f"srun --het-group={offset} --ntasks=1 "
         f"""bash -c 'echo "STEP=$SLURM_STEP_ID"' >> {file_out}"""
         for _ in range(2)
         for offset in range(len(component_nodes))
     )
     # sbatch also exports a bare SLURM_STEPMGR but salloc does not, so use the
     # per-component form, which both submission paths set.
     stepmgr_echos = "\n".join(
         f'echo "STEPMGR_{offset}=$SLURM_STEPMGR_HET_GROUP_{offset}" >> {file_out}'
         for offset in range(len(component_nodes))
     )
     atf.make_bash_script(
         file_in,
         f"{stepmgr_echos}\n{sruns}\necho {OUTPUT_MARKER} >> {file_out}\n",
     )

     het_opts = _het_opts(component_nodes)
     job_param = het_opts if command == "salloc" else f"-o /dev/null {het_opts}"
     atf.submit_job(command, job_param, f"./{file_in}", wrap_job=False, fatal=True)

     atf.assert_file_contents(file_out, OUTPUT_MARKER, contains=True)

     output = atf.run_command_output(f"cat {file_out}", fatal=True)
     step_ids = [int(m) for m in re.findall(r"^STEP=(\d+)$", output, re.MULTILINE)]

     # A plain het job already numbers its steps job-globally, so the ids below
     # only say something about stepmgr once every component is known to have
     # one.
     stepmgr_offsets = sorted(
         int(offset)
         for offset in re.findall(r"^STEPMGR_(\d+)=\S+$", output, re.MULTILINE)
     )
     assert stepmgr_offsets == list(range(len(component_nodes))), (
         f"Every component must report a stepmgr host for the step ids to "
         f"exercise stepmgr rather than the controller, got offsets "
         f"{stepmgr_offsets}.\nRaw output:\n{output}"
     )

     expected_ids = list(range(2 * len(component_nodes)))
     assert step_ids == expected_ids, (
         f"Expected job-global step ids {expected_ids}, got {step_ids}.\n"
         f"Per-component duplicates indicate stepmgr is not allocating job-globally.\n"
         f"Raw output:\n{output}"
     )


 @pytest.mark.parametrize("follower_stepmgr", [False, True])
 @pytest.mark.parametrize("leader_stepmgr", [False, True])
 @pytest.mark.parametrize("command", ["sbatch", "salloc"])
 def test_hetjob_stepmgr_follower_tracks_leader(
     command, leader_stepmgr, follower_stepmgr
 ):
     """The follower's StepMgrEnabled always matches the leader's.

     Only the leader's --stepmgr decides for the whole job, so passing it on
     a non-leader component is inert. salloc and sbatch document the override
     identically, so both submission paths are covered.
     """

     # --stepmgr only enables stepmgr per-job "if it isn't enabled system wide",
     # and there is no per-job way to opt back out, so under
     # SlurmctldParameters=enable_stepmgr every component is enabled whatever
     # was submitted and leader-wins cannot be observed.
     if atf.config_parameter_includes("SlurmctldParameters", "enable_stepmgr"):
         pytest.skip(
             "Issue 50976: leader-wins is unobservable when stepmgr is enabled"
             " globally"
         )

     leader_opts = "-N1 --ntasks=1 -t2" + (" --stepmgr" if leader_stepmgr else "")
     follower_opts = "-N1 --ntasks=1 -t2" + (" --stepmgr" if follower_stepmgr else "")

     if command == "sbatch":
         file_in = f"hetjob_stepmgr_{int(leader_stepmgr)}{int(follower_stepmgr)}.in"
         atf.make_bash_script(
             file_in,
             f"""
 #SBATCH {leader_opts}
 #SBATCH hetjob
 #SBATCH {follower_opts}

 true
 """,
         )
         leader_job_id = atf.submit_job_sbatch(f"-o /dev/null {file_in}", fatal=True)
     else:
         # --no-shell leaves the allocation in place so StepMgrEnabled is still
         # readable for every component.
         leader_job_id = atf.submit_job_salloc(
             f"--no-shell {leader_opts} : {follower_opts}", fatal=True
         )

     jobs = atf.get_jobs(leader_job_id, fatal=True)
     components = atf.range_to_list(jobs[leader_job_id]["HetJobIdSet"])
     follower_id = next(c for c in components if jobs[c]["HetJobOffset"] == 1)

     # scontrol prints StepMgrEnabled only when the bit is set, so a disabled
     # component has no such key at all.
     leader_flag = jobs[leader_job_id].get("StepMgrEnabled")
     follower_flag = jobs[follower_id].get("StepMgrEnabled")

     assert follower_flag == leader_flag, (
         f"JobId {follower_id}: follower StepMgrEnabled must match the leader "
         f"(leader-wins), got leader={leader_flag!r}, follower={follower_flag!r}"
     )

     expected_flag = "Yes" if leader_stepmgr else None
     assert leader_flag == expected_flag, (
         f"JobId {leader_job_id}: leader submitted with --stepmgr="
         f"{leader_stepmgr} should have StepMgrEnabled={expected_flag!r}, "
         f"got {leader_flag!r}"
     )


 def test_hetjob_stepmgr_multi_component_step_shares_id():
     """A step spanning both components shares one id; skipped ids leave gaps.

     Mirrors the worked example in the heterogeneous_jobs documentation: a
     step on both components, then one on the leader only, then both again.
     The follower must jump straight from id 0 to id 2.
     """

     file_in = "hetjob_stepmgr_multi.in"
     file_out = "hetjob_stepmgr_multi.out"
     atf.make_bash_script(
         file_in,
         f"""
 #SBATCH -N1 --ntasks=1 -t2 --stepmgr
 #SBATCH hetjob
 #SBATCH -N1 --ntasks=1 -t2

 echo "STEPMGR_0=$SLURM_STEPMGR_HET_GROUP_0"
 echo "STEPMGR_1=$SLURM_STEPMGR_HET_GROUP_1"
 srun --het-group=0,1 bash -c 'echo "STEP=$SLURM_STEP_ID"'
 srun --het-group=0 bash -c 'echo "STEP=$SLURM_STEP_ID"'
 srun --het-group=0,1 bash -c 'echo "STEP=$SLURM_STEP_ID"'
 echo {OUTPUT_MARKER}
 """,
     )

     leader_job_id = atf.submit_job_sbatch(f"-o {file_out} {file_in}", fatal=True)
     atf.wait_for_job_state(leader_job_id, "DONE", fatal=True)
     atf.assert_file_contents(file_out, OUTPUT_MARKER, contains=True)

     output = atf.run_command_output(f"cat {file_out}", fatal=True)
     step_ids = [int(m) for m in re.findall(r"^STEP=(\d+)$", output, re.MULTILINE)]

     # A plain het job already numbers its steps job-globally, so the ids below
     # only say something about stepmgr once every component is known to have
     # one.
     stepmgr_offsets = sorted(
         int(offset)
         for offset in re.findall(r"^STEPMGR_(\d+)=\S+$", output, re.MULTILINE)
     )
     assert stepmgr_offsets == [0, 1], (
         f"Every component must report a stepmgr host for the step ids to "
         f"exercise stepmgr rather than the controller, got offsets "
         f"{stepmgr_offsets}.\nRaw output:\n{output}"
     )

     # One task reports per participating component, so the number of lines
     # carrying a step id is how many components ran that step. The middle
     # step is leader-only, which is what leaves the follower a gap at id 1.
     assert collections.Counter(step_ids) == {0: 2, 1: 1, 2: 2}, (
         f"Expected the both-component steps to share ids 0 and 2 (2 lines "
         f"each) and the leader-only step to take id 1 (1 line), leaving the "
         f"follower a gap at 1; got {sorted(step_ids)}.\nRaw output:\n{output}"
     )


 def test_hetjob_stepmgr_env_vars():
     """Each component exports its own SLURM_STEPMGR_HET_GROUP_<N>.

     --exclusive forces the components onto different nodes so that a
     regression broadcasting the leader's host to every suffix is caught.
     """

     file_in = "hetjob_stepmgr_env.in"
     file_out = "hetjob_stepmgr_env.out"
     atf.make_bash_script(
         file_in,
         f"""
 #SBATCH -N1 --ntasks=1 -t2 --exclusive --stepmgr
 #SBATCH hetjob
 #SBATCH -N1 --ntasks=1 -t2 --exclusive

 env
 echo {OUTPUT_MARKER}
 sleep infinity
 """,
     )

     leader_job_id = atf.submit_job_sbatch(f"-o {file_out} {file_in}", fatal=True)

     # The script stays alive after dumping env so that BatchHost is still
     # readable for every component once the dump is complete.
     atf.assert_file_contents(file_out, OUTPUT_MARKER, contains=True)

     jobs = atf.get_jobs(leader_job_id, fatal=True)
     components = sorted(
         atf.range_to_list(jobs[leader_job_id]["HetJobIdSet"]),
         key=lambda job_id: jobs[job_id]["HetJobOffset"],
     )
     batch_hosts = [jobs[job_id]["BatchHost"] for job_id in components]

     output = atf.run_command_output(f"cat {file_out}", fatal=True)

     if batch_hosts[0] == batch_hosts[1]:
         pytest.fail(
             f"The components should be on different nodes for this test to be "
             f"meaningful, got {batch_hosts}"
         )

     for offset, host in enumerate(batch_hosts):
         for group in ("HET_GROUP", "PACK_GROUP"):
             var = f"SLURM_STEPMGR_{group}_{offset}"
             assert (
                 re.search(rf"^{var}={host}$", output, re.MULTILINE) is not None
             ), f"Missing {var}={host} in batch environment"

     assert (
         re.search(r"^SLURM_STEPMGR=\S+$", output, re.MULTILINE) is not None
     ), "Missing bare SLURM_STEPMGR, which older clients still rely on"

     double_suffixed = re.findall(
         r"^(SLURM_\w+_(?:HET|PACK)_GROUP_\d+_(?:HET|PACK)_GROUP_\d+)=",
         output,
         re.MULTILINE,
     )
     assert (
         double_suffixed == []
     ), f"Env vars must be suffixed once per component, got {double_suffixed}"


 def test_hetjob_stepmgr_srun_task_env():
     """srun exports the per-component stepmgr vars in the task environment.

     srun builds the task environment itself, so the batch environment checked
     above comes from a different producer and does not cover it.
     """

     file_in = "hetjob_stepmgr_srun_env.in"
     file_out = "hetjob_stepmgr_srun_env.out"
     atf.make_bash_script(
         file_in,
         f"""
 #SBATCH -N1 --ntasks=1 -t2 --exclusive --stepmgr
 #SBATCH hetjob
 #SBATCH -N1 --ntasks=1 -t2 --exclusive

 srun --het-group=0,1 env >> {file_out}
 echo {OUTPUT_MARKER} >> {file_out}
 """,
     )

     leader_job_id = atf.submit_job_sbatch(f"-o /dev/null {file_in}", fatal=True)
     atf.wait_for_job_state(leader_job_id, "DONE", fatal=True)
     atf.assert_file_contents(file_out, OUTPUT_MARKER, contains=True)

     output = atf.run_command_output(f"cat {file_out}", fatal=True)

     for offset in (0, 1):
         for group in ("HET_GROUP", "PACK_GROUP"):
             var = f"SLURM_STEPMGR_{group}_{offset}"
             assert (
                 re.search(rf"^{var}=\S+$", output, re.MULTILINE) is not None
             ), f"Missing {var} in the srun task environment"

     double_suffixed = re.findall(
         r"^(SLURM_\w+_(?:HET|PACK)_GROUP_\d+_(?:HET|PACK)_GROUP_\d+)=",
         output,
         re.MULTILINE,
     )
     assert (
         double_suffixed == []
     ), f"Task env vars must be suffixed once per component, got {double_suffixed}"
	############################################################################
	# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	############################################################################
	"""Heterogeneous job + stepmgr-on-stepd: verify job-global step ids.

	With --stepmgr on a heterogeneous job, each component's stepmgr asks the
	leader for the next step id (REQUEST_HET_STEP_ID). Step ids must remain
	unique and contiguous across the whole het job, not per-component.

	Without that RPC, sequential srun --het-group calls would reuse ids
	0,1,0,1,... per component instead of 0,1,2,3 across the job.
	"""

	import collections
	import re

	import pytest

	import atf

	pytestmark = pytest.mark.slow

	# Echoed after the last srun so we can tell the batch script's whole output
	# has been flushed before parsing it.
	OUTPUT_MARKER = "HETJOB_STEPS_DONE"


	@pytest.fixture(scope="module", autouse=True)
	def setup():
	# The client tools upgrade last, so gating on srun >= 26.11 already
	# implies slurmd and slurmctld are >= 26.11.
	atf.require_version(
	(26, 11),
	component="bin/srun",
	reason="Issue 50976: hetjob + stepmgr requires srun 26.11+",
	)
	atf.require_nodes(3)
	atf.require_config_parameter_includes("PrologFlags", "Contain")
	# Het jobs are only started by backfill; don't wait for its full cycle
	atf.require_config_parameter_includes("SchedulerParameters", ("bf_interval", 1))
	atf.require_config_parameter("SchedulerType", "sched/backfill")
	atf.require_slurm_running()


	def _het_opts(component_nodes):
	"""Build het job options with --stepmgr on the leader only."""

	return " : ".join(
	f"-N{nodes} --ntasks-per-node=1 -t2" + (" --stepmgr" if offset == 0 else "")
	for offset, nodes in enumerate(component_nodes)
	)


	@pytest.mark.parametrize(
	"command, component_nodes",
	[
	("sbatch", [1, 1]),
	("salloc", [1, 1]),
	("sbatch", [1, 1, 1]),
	("sbatch", [2, 1]),
	],
	ids=["sbatch-2comp", "salloc-2comp", "sbatch-3comp", "sbatch-multinode"],
	)
	def test_hetjob_stepmgr_job_global_step_ids(command, component_nodes):
	"""Sequential sruns across all components yield unique, contiguous step ids."""

	test_id = f"{command}_{'-'.join(str(n) for n in component_nodes)}"
	file_in = f"hetjob_stepmgr_{test_id}.in"
	file_out = f"hetjob_stepmgr_{test_id}.out"

	# Two rounds over every component, so a per-component counter would
	# repeat ids instead of continuing the job-global sequence.
	sruns = "\n".join(
	f"srun --het-group={offset} --ntasks=1 "
	f"""bash -c 'echo "STEP=$SLURM_STEP_ID"' >> {file_out}"""
	for _ in range(2)
	for offset in range(len(component_nodes))
	)
	# sbatch also exports a bare SLURM_STEPMGR but salloc does not, so use the
	# per-component form, which both submission paths set.
	stepmgr_echos = "\n".join(
	f'echo "STEPMGR_{offset}=$SLURM_STEPMGR_HET_GROUP_{offset}" >> {file_out}'
	for offset in range(len(component_nodes))
	)
	atf.make_bash_script(
	file_in,
	f"{stepmgr_echos}\n{sruns}\necho {OUTPUT_MARKER} >> {file_out}\n",
	)

	het_opts = _het_opts(component_nodes)
	job_param = het_opts if command == "salloc" else f"-o /dev/null {het_opts}"
	atf.submit_job(command, job_param, f"./{file_in}", wrap_job=False, fatal=True)

	atf.assert_file_contents(file_out, OUTPUT_MARKER, contains=True)

	output = atf.run_command_output(f"cat {file_out}", fatal=True)
	step_ids = [int(m) for m in re.findall(r"^STEP=(\d+)$", output, re.MULTILINE)]

	# A plain het job already numbers its steps job-globally, so the ids below
	# only say something about stepmgr once every component is known to have
	# one.
	stepmgr_offsets = sorted(
	int(offset)
	for offset in re.findall(r"^STEPMGR_(\d+)=\S+$", output, re.MULTILINE)
	)
	assert stepmgr_offsets == list(range(len(component_nodes))), (
	f"Every component must report a stepmgr host for the step ids to "
	f"exercise stepmgr rather than the controller, got offsets "
	f"{stepmgr_offsets}.\nRaw output:\n{output}"
	)

	expected_ids = list(range(2 * len(component_nodes)))
	assert step_ids == expected_ids, (
	f"Expected job-global step ids {expected_ids}, got {step_ids}.\n"
	f"Per-component duplicates indicate stepmgr is not allocating job-globally.\n"
	f"Raw output:\n{output}"
	)


	@pytest.mark.parametrize("follower_stepmgr", [False, True])
	@pytest.mark.parametrize("leader_stepmgr", [False, True])
	@pytest.mark.parametrize("command", ["sbatch", "salloc"])
	def test_hetjob_stepmgr_follower_tracks_leader(
	command, leader_stepmgr, follower_stepmgr
	):
	"""The follower's StepMgrEnabled always matches the leader's.

	Only the leader's --stepmgr decides for the whole job, so passing it on
	a non-leader component is inert. salloc and sbatch document the override
	identically, so both submission paths are covered.
	"""

	# --stepmgr only enables stepmgr per-job "if it isn't enabled system wide",
	# and there is no per-job way to opt back out, so under
	# SlurmctldParameters=enable_stepmgr every component is enabled whatever
	# was submitted and leader-wins cannot be observed.
	if atf.config_parameter_includes("SlurmctldParameters", "enable_stepmgr"):
	pytest.skip(
	"Issue 50976: leader-wins is unobservable when stepmgr is enabled"
	" globally"
	)

	leader_opts = "-N1 --ntasks=1 -t2" + (" --stepmgr" if leader_stepmgr else "")
	follower_opts = "-N1 --ntasks=1 -t2" + (" --stepmgr" if follower_stepmgr else "")

	if command == "sbatch":
	file_in = f"hetjob_stepmgr_{int(leader_stepmgr)}{int(follower_stepmgr)}.in"
	atf.make_bash_script(
	file_in,
	f"""
	#SBATCH {leader_opts}
	#SBATCH hetjob
	#SBATCH {follower_opts}

	true
	""",
	)
	leader_job_id = atf.submit_job_sbatch(f"-o /dev/null {file_in}", fatal=True)
	else:
	# --no-shell leaves the allocation in place so StepMgrEnabled is still
	# readable for every component.
	leader_job_id = atf.submit_job_salloc(
	f"--no-shell {leader_opts} : {follower_opts}", fatal=True
	)

	jobs = atf.get_jobs(leader_job_id, fatal=True)
	components = atf.range_to_list(jobs[leader_job_id]["HetJobIdSet"])
	follower_id = next(c for c in components if jobs[c]["HetJobOffset"] == 1)

	# scontrol prints StepMgrEnabled only when the bit is set, so a disabled
	# component has no such key at all.
	leader_flag = jobs[leader_job_id].get("StepMgrEnabled")
	follower_flag = jobs[follower_id].get("StepMgrEnabled")

	assert follower_flag == leader_flag, (
	f"JobId {follower_id}: follower StepMgrEnabled must match the leader "
	f"(leader-wins), got leader={leader_flag!r}, follower={follower_flag!r}"
	)

	expected_flag = "Yes" if leader_stepmgr else None
	assert leader_flag == expected_flag, (
	f"JobId {leader_job_id}: leader submitted with --stepmgr="
	f"{leader_stepmgr} should have StepMgrEnabled={expected_flag!r}, "
	f"got {leader_flag!r}"
	)


	def test_hetjob_stepmgr_multi_component_step_shares_id():
	"""A step spanning both components shares one id; skipped ids leave gaps.

	Mirrors the worked example in the heterogeneous_jobs documentation: a
	step on both components, then one on the leader only, then both again.
	The follower must jump straight from id 0 to id 2.
	"""

	file_in = "hetjob_stepmgr_multi.in"
	file_out = "hetjob_stepmgr_multi.out"
	atf.make_bash_script(
	file_in,
	f"""
	#SBATCH -N1 --ntasks=1 -t2 --stepmgr
	#SBATCH hetjob
	#SBATCH -N1 --ntasks=1 -t2

	echo "STEPMGR_0=$SLURM_STEPMGR_HET_GROUP_0"
	echo "STEPMGR_1=$SLURM_STEPMGR_HET_GROUP_1"
	srun --het-group=0,1 bash -c 'echo "STEP=$SLURM_STEP_ID"'
	srun --het-group=0 bash -c 'echo "STEP=$SLURM_STEP_ID"'
	srun --het-group=0,1 bash -c 'echo "STEP=$SLURM_STEP_ID"'
	echo {OUTPUT_MARKER}
	""",
	)

	leader_job_id = atf.submit_job_sbatch(f"-o {file_out} {file_in}", fatal=True)
	atf.wait_for_job_state(leader_job_id, "DONE", fatal=True)
	atf.assert_file_contents(file_out, OUTPUT_MARKER, contains=True)

	output = atf.run_command_output(f"cat {file_out}", fatal=True)
	step_ids = [int(m) for m in re.findall(r"^STEP=(\d+)$", output, re.MULTILINE)]

	# A plain het job already numbers its steps job-globally, so the ids below
	# only say something about stepmgr once every component is known to have
	# one.
	stepmgr_offsets = sorted(
	int(offset)
	for offset in re.findall(r"^STEPMGR_(\d+)=\S+$", output, re.MULTILINE)
	)
	assert stepmgr_offsets == [0, 1], (
	f"Every component must report a stepmgr host for the step ids to "
	f"exercise stepmgr rather than the controller, got offsets "
	f"{stepmgr_offsets}.\nRaw output:\n{output}"
	)

	# One task reports per participating component, so the number of lines
	# carrying a step id is how many components ran that step. The middle
	# step is leader-only, which is what leaves the follower a gap at id 1.
	assert collections.Counter(step_ids) == {0: 2, 1: 1, 2: 2}, (
	f"Expected the both-component steps to share ids 0 and 2 (2 lines "
	f"each) and the leader-only step to take id 1 (1 line), leaving the "
	f"follower a gap at 1; got {sorted(step_ids)}.\nRaw output:\n{output}"
	)


	def test_hetjob_stepmgr_env_vars():
	"""Each component exports its own SLURM_STEPMGR_HET_GROUP_<N>.

	--exclusive forces the components onto different nodes so that a
	regression broadcasting the leader's host to every suffix is caught.
	"""

	file_in = "hetjob_stepmgr_env.in"
	file_out = "hetjob_stepmgr_env.out"
	atf.make_bash_script(
	file_in,
	f"""
	#SBATCH -N1 --ntasks=1 -t2 --exclusive --stepmgr
	#SBATCH hetjob
	#SBATCH -N1 --ntasks=1 -t2 --exclusive

	env
	echo {OUTPUT_MARKER}
	sleep infinity
	""",
	)

	leader_job_id = atf.submit_job_sbatch(f"-o {file_out} {file_in}", fatal=True)

	# The script stays alive after dumping env so that BatchHost is still
	# readable for every component once the dump is complete.
	atf.assert_file_contents(file_out, OUTPUT_MARKER, contains=True)

	jobs = atf.get_jobs(leader_job_id, fatal=True)
	components = sorted(
	atf.range_to_list(jobs[leader_job_id]["HetJobIdSet"]),
	key=lambda job_id: jobs[job_id]["HetJobOffset"],
	)
	batch_hosts = [jobs[job_id]["BatchHost"] for job_id in components]

	output = atf.run_command_output(f"cat {file_out}", fatal=True)

	if batch_hosts[0] == batch_hosts[1]:
	pytest.fail(
	f"The components should be on different nodes for this test to be "
	f"meaningful, got {batch_hosts}"
	)

	for offset, host in enumerate(batch_hosts):
	for group in ("HET_GROUP", "PACK_GROUP"):
	var = f"SLURM_STEPMGR_{group}_{offset}"
	assert (
	re.search(rf"^{var}={host}$", output, re.MULTILINE) is not None
	), f"Missing {var}={host} in batch environment"

	assert (
	re.search(r"^SLURM_STEPMGR=\S+$", output, re.MULTILINE) is not None
	), "Missing bare SLURM_STEPMGR, which older clients still rely on"

	double_suffixed = re.findall(
	r"^(SLURM_\w+_(?:HET\|PACK)_GROUP_\d+_(?:HET\|PACK)_GROUP_\d+)=",
	output,
	re.MULTILINE,
	)
	assert (
	double_suffixed == []
	), f"Env vars must be suffixed once per component, got {double_suffixed}"


	def test_hetjob_stepmgr_srun_task_env():
	"""srun exports the per-component stepmgr vars in the task environment.

	srun builds the task environment itself, so the batch environment checked
	above comes from a different producer and does not cover it.
	"""

	file_in = "hetjob_stepmgr_srun_env.in"
	file_out = "hetjob_stepmgr_srun_env.out"
	atf.make_bash_script(
	file_in,
	f"""
	#SBATCH -N1 --ntasks=1 -t2 --exclusive --stepmgr
	#SBATCH hetjob
	#SBATCH -N1 --ntasks=1 -t2 --exclusive

	srun --het-group=0,1 env >> {file_out}
	echo {OUTPUT_MARKER} >> {file_out}
	""",
	)

	leader_job_id = atf.submit_job_sbatch(f"-o /dev/null {file_in}", fatal=True)
	atf.wait_for_job_state(leader_job_id, "DONE", fatal=True)
	atf.assert_file_contents(file_out, OUTPUT_MARKER, contains=True)

	output = atf.run_command_output(f"cat {file_out}", fatal=True)

	for offset in (0, 1):
	for group in ("HET_GROUP", "PACK_GROUP"):
	var = f"SLURM_STEPMGR_{group}_{offset}"
	assert (
	re.search(rf"^{var}=\S+$", output, re.MULTILINE) is not None
	), f"Missing {var} in the srun task environment"

	double_suffixed = re.findall(
	r"^(SLURM_\w+_(?:HET\|PACK)_GROUP_\d+_(?:HET\|PACK)_GROUP_\d+)=",
	output,
	re.MULTILINE,
	)
	assert (
	double_suffixed == []
	), f"Task env vars must be suffixed once per component, got {double_suffixed}"