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

 ############################################################################
 # Copyright (C) SchedMD LLC.
 ############################################################################
 """GPU selection for jobs and steps when GPUs are bound to specific cores."""

 import re

 import pytest

 import atf

 # One GPU per socket, so a single-socket step forces pass 2 of the GRES
 # selection to fill more than one remaining device.
 _NUM_GPUS = 4
 # Spare sockets that hold no GPU of their own, so a request that cannot be
 # placed on a GPU-bound socket is one --gres-flags=enforce-binding has to
 # reject rather than one the node simply cannot fit.
 _NUM_SOCKETS = _NUM_GPUS + 2
 _ALL_GPUS = ",".join(str(i) for i in range(_NUM_GPUS))

 _OVERLAP_GRES_UNFIXED = atf.get_version("sbin/slurmctld") < (26, 5, 3)
 _OVERLAP_GRES_REASON = "Ticket 24862: --overlap GRES double-select fixed in 26.05.3"

 node_name = None


 @pytest.fixture(scope="module", autouse=True)
 def setup():
     global node_name
     atf.require_config_parameter("SelectType", "select/cons_tres")
     atf.require_config_parameter("SelectTypeParameters", "CR_Socket_Memory")
     atf.require_config_parameter_includes("GresTypes", "gpu")
     atf.require_nodes(
         1,
         [
             ("Sockets", _NUM_SOCKETS),
             ("CoresPerSocket", 1),
             ("ThreadsPerCore", 1),
             ("Gres", f"gpu:{_NUM_GPUS}"),
         ],
     )
     # GPU's need to point to existing device files
     for i in range(_NUM_GPUS):
         atf.require_tty(i)
     gpu_nodes = [
         name
         for name, params in atf.get_nodes(live=False, quiet=True, fatal=True).items()
         if name != "DEFAULT" and "gpu:" in str(params.get("Gres", ""))
     ]
     assert gpu_nodes, "Expected require_nodes() to configure a node with GPUs"
     node_name = gpu_nodes[0]
     # Bind GPU i to socket i (core i), so a step on one socket owns just one
     # GPU and pass 2 must supply the rest.
     gpu_lines = []
     for i in range(_NUM_GPUS):
         prefix = node_name if i == 0 else f"NodeName={node_name}"
         gpu_lines.append(f"{prefix} Name=gpu Cores={i} File=/dev/tty{i}")
     atf.require_config_parameter("NodeName", "\n".join(gpu_lines), source="gres")
     atf.require_slurm_running()


 @pytest.fixture(scope="function")
 def single_socket_gpu_job():
     """Job holding every GPU, with its cores confined to a single socket.

     The preconditions are asserted here so that every test sharing the
     fixture checks them in one place, and so that a topology that does not
     reproduce the condition says which precondition it failed rather than
     only that the step came up short on GPUs.
     """

     job_id = atf.submit_job_sbatch(
         f"-N1 -n1 --gpus-per-node={_NUM_GPUS} --wrap 'sleep infinity'",
         fatal=True,
     )
     atf.wait_for_job_state(job_id, "RUNNING", fatal=True)

     job_gres = atf.get_job_parameter(job_id, "GRES", fatal=True, quiet=True)
     assert job_gres and f"gpu:{_NUM_GPUS}(IDX:" in job_gres, (
         f"Expected job to hold every GPU (gpu:{_NUM_GPUS}(IDX:...)), "
         f"got GRES={job_gres!r} - topology did not produce the "
         f"expected no-enforce-binding allocation"
     )

     num_cpus = atf.get_job_parameter(job_id, "NumCPUs", fatal=True, quiet=True)
     assert num_cpus == 1, (
         f"Expected job to hold a single socket (1 CPU with CoresPerSocket=1), "
         f"got NumCPUs={num_cpus} - topology did not confine the job to one "
         f"socket, so the step already covers every GPU-bound socket"
     )

     return job_id


 def test_gpu_socket_sharing():
     """Test allocating multiple core-bound gpus with enforce-binding"""

     output = atf.run_command_output(
         f"srun --gres-flags=enforce-binding --ntasks-per-socket=1 \
                     --cpus-per-task=1 --ntasks-per-node=2 -N1 \
                     --gpus-per-task=1 scontrol show nodes {node_name} -d",
         fatal=True,
     )
     assert (
         re.search(r"GresUsed=gpu:(?:[^:]*:)?2\(", output) is not None
     ), "Verify that job allocated 2 gpus"


 def test_gpu_socket_sharing_no_alloc():
     """Test allocating multiple core-bound gpus with enforce-binding without enough resources"""

     output = atf.run_command(
         f"srun --gres-flags=enforce-binding --ntasks-per-socket=1 \
                     --cpus-per-task=2 --ntasks-per-node=2 -N1 \
                     --gpus-per-task=1 scontrol show nodes {node_name} -d",
         xfail=True,
         fatal=True,
     )
     assert output["exit_code"] != 0, "Verify that srun command failed"

     expected_msg = r"srun: error: .+ Requested node configuration is not available"
     assert (
         re.search(
             expected_msg,
             str(output["stderr"]),
         )
         is not None
     ), "Verify that job is rejected with the right reason."


 @pytest.mark.xfail(_OVERLAP_GRES_UNFIXED, reason=_OVERLAP_GRES_REASON)
 def test_overlap_step_gets_all_gpus(single_socket_gpu_job):
     """An --overlap step must receive every GPU it requests.

     srun --overlap shares all of the job's GRES with the step.  The batch
     job requests -n1 so CR_Socket_Memory confines its cores to one socket,
     while --gpus-per-node (no enforce-binding) allocates every GPU across
     all sockets.  The --overlap step lands on that single socket but, per the
     srun default, inherits every GPU the job holds, so it must be granted
     all of them even though most are bound to sockets its cores do not
     cover.  A short allocation shows up as a denied device in the step's
     cgroup if task/cgroup and ConstrainDevices=yes is configured
     (ticket 24862 regression).
     """

     job_id = single_socket_gpu_job

     result = atf.run_job(
         f"--jobid={job_id} --overlap -n1 --mem=0 "
         f"bash -c 'echo $SLURM_STEP_GPUS; echo $CUDA_VISIBLE_DEVICES'",
         fatal=True,
         quiet=True,
     )
     output = result["stdout"].splitlines()
     assert len(output) >= 2, (
         f"Expected two output lines (SLURM_STEP_GPUS and "
         f"CUDA_VISIBLE_DEVICES), got {result['stdout']!r}"
     )
     step_gpus = output[0].strip()
     visible_devices = output[1].strip()

     assert step_gpus == _ALL_GPUS, (
         f"Step SLURM_STEP_GPUS should be {_ALL_GPUS!r} (every GPU), "
         f"got {step_gpus!r} - the --overlap step was not granted all of "
         f"the job's GPUs (ticket 24862 regression)"
     )
     assert visible_devices == _ALL_GPUS, (
         f"Step CUDA_VISIBLE_DEVICES should be exported and agree with "
         f"SLURM_STEP_GPUS ({_ALL_GPUS}), got {visible_devices!r} - the "
         f"step's GPU allocation is not rendered consistently across both "
         f"variables (ticket 24862 regression)"
     )


 @pytest.mark.xfail(_OVERLAP_GRES_UNFIXED, reason=_OVERLAP_GRES_REASON)
 @pytest.mark.parametrize("first_step_overlaps", [True, False])
 def test_concurrent_steps_each_get_all_gpus(first_step_overlaps, single_socket_gpu_job):
     """Concurrent steps sharing an allocation each receive every GPU.

     srun --overlap shares all resources with all other steps, even those
     that did not specify --overlap.  Two steps running at the same time
     against the same allocation must therefore both be granted every GPU
     the job holds, whichever of them asked to overlap.
     """

     job_id = single_socket_gpu_job

     first_step_gpus_file = "first_step_gpus"
     overlap = "--overlap " if first_step_overlaps else ""
     atf.run_command(
         f"srun --jobid={job_id} {overlap}-n1 --mem=0 bash -c "
         f"'echo $SLURM_STEP_GPUS > {first_step_gpus_file}; sleep infinity'",
         background=True,
     )
     atf.assert_file_contents(
         first_step_gpus_file,
         _ALL_GPUS,
         message=(
             f"First step SLURM_STEP_GPUS should be {_ALL_GPUS!r} (every GPU) "
             f"- the step was not granted all of the job's GPUs"
         ),
     )

     result = atf.run_job(
         f"--jobid={job_id} --overlap -n1 --mem=0 bash -c 'echo $SLURM_STEP_GPUS'",
         fatal=True,
         quiet=True,
     )
     second_step_gpus = result["stdout"].strip()

     assert second_step_gpus == _ALL_GPUS, (
         f"Second step SLURM_STEP_GPUS should be {_ALL_GPUS!r} (every GPU), "
         f"got {second_step_gpus!r} - an --overlap step running alongside "
         f"another step was not granted all of the job's GPUs"
     )


 @pytest.mark.xfail(_OVERLAP_GRES_UNFIXED, reason=_OVERLAP_GRES_REASON)
 @pytest.mark.parametrize("requested_gpus", range(2, _NUM_GPUS))
 def test_overlap_step_gets_requested_gpu_subset(requested_gpus, single_socket_gpu_job):
     """An --overlap step requesting a subset gets that many distinct GPUs.

     A job step may explicitly request a different generic resource count
     than the job, so an --overlap step asking for fewer GPUs than the job
     holds must be granted exactly that many, and no device twice.
     """

     job_id = single_socket_gpu_job

     result = atf.run_job(
         f"--jobid={job_id} --overlap -n1 --mem=0 --gres=gpu:{requested_gpus} "
         f"bash -c 'echo $SLURM_STEP_GPUS'",
         fatal=True,
         quiet=True,
     )
     step_gpus = result["stdout"].strip()
     indices = step_gpus.split(",") if step_gpus else []

     assert len(indices) == requested_gpus, (
         f"Step requesting gpu:{requested_gpus} should be granted "
         f"{requested_gpus} GPUs, got SLURM_STEP_GPUS={step_gpus!r}"
     )
     assert len(set(indices)) == len(indices), (
         f"Step requesting gpu:{requested_gpus} should be granted distinct "
         f"GPUs, got SLURM_STEP_GPUS={step_gpus!r}"
     )


 def test_non_overlap_step_gets_all_gpus():
     """A default (non --overlap) step also receives every GPU it requests.

     This is not a reproducer of ticket 24862 (that bug was specific to
     --overlap); it guards the shared non-overlap path against collateral
     breakage from the fix.  srun's own step lands on a single socket
     (CR_Socket_Memory with -n1) while requesting every GPU across all
     sockets, so it must still be granted all of them.
     """

     result = atf.run_job(
         f"-N1 -n1 --gpus-per-node={_NUM_GPUS} bash -c "
         f"'echo $SLURM_CPUS_ON_NODE; echo $SLURM_STEP_GPUS'",
         fatal=True,
         quiet=True,
     )
     output = result["stdout"].splitlines()
     assert len(output) >= 2, (
         f"Expected two output lines (SLURM_CPUS_ON_NODE and "
         f"SLURM_STEP_GPUS), got {result['stdout']!r}"
     )
     cpus_on_node = output[0].strip()
     step_gpus = output[1].strip()

     assert cpus_on_node == "1", (
         f"Expected the step to hold a single socket (1 CPU with "
         f"CoresPerSocket=1), got SLURM_CPUS_ON_NODE={cpus_on_node!r} - "
         f"topology did not confine the step to one socket, so it already "
         f"covers every GPU-bound socket"
     )
     assert step_gpus == _ALL_GPUS, (
         f"Non-overlap step SLURM_STEP_GPUS should be {_ALL_GPUS!r} "
         f"(every GPU), got {step_gpus!r}"
     )
	############################################################################
	# Copyright (C) SchedMD LLC.
	############################################################################
	"""GPU selection for jobs and steps when GPUs are bound to specific cores."""

	import re

	import pytest

	import atf

	# One GPU per socket, so a single-socket step forces pass 2 of the GRES
	# selection to fill more than one remaining device.
	_NUM_GPUS = 4
	# Spare sockets that hold no GPU of their own, so a request that cannot be
	# placed on a GPU-bound socket is one --gres-flags=enforce-binding has to
	# reject rather than one the node simply cannot fit.
	_NUM_SOCKETS = _NUM_GPUS + 2
	_ALL_GPUS = ",".join(str(i) for i in range(_NUM_GPUS))

	_OVERLAP_GRES_UNFIXED = atf.get_version("sbin/slurmctld") < (26, 5, 3)
	_OVERLAP_GRES_REASON = "Ticket 24862: --overlap GRES double-select fixed in 26.05.3"

	node_name = None


	@pytest.fixture(scope="module", autouse=True)
	def setup():
	global node_name
	atf.require_config_parameter("SelectType", "select/cons_tres")
	atf.require_config_parameter("SelectTypeParameters", "CR_Socket_Memory")
	atf.require_config_parameter_includes("GresTypes", "gpu")
	atf.require_nodes(
	1,
	[
	("Sockets", _NUM_SOCKETS),
	("CoresPerSocket", 1),
	("ThreadsPerCore", 1),
	("Gres", f"gpu:{_NUM_GPUS}"),
	],
	)
	# GPU's need to point to existing device files
	for i in range(_NUM_GPUS):
	atf.require_tty(i)
	gpu_nodes = [
	name
	for name, params in atf.get_nodes(live=False, quiet=True, fatal=True).items()
	if name != "DEFAULT" and "gpu:" in str(params.get("Gres", ""))
	]
	assert gpu_nodes, "Expected require_nodes() to configure a node with GPUs"
	node_name = gpu_nodes[0]
	# Bind GPU i to socket i (core i), so a step on one socket owns just one
	# GPU and pass 2 must supply the rest.
	gpu_lines = []
	for i in range(_NUM_GPUS):
	prefix = node_name if i == 0 else f"NodeName={node_name}"
	gpu_lines.append(f"{prefix} Name=gpu Cores={i} File=/dev/tty{i}")
	atf.require_config_parameter("NodeName", "\n".join(gpu_lines), source="gres")
	atf.require_slurm_running()


	@pytest.fixture(scope="function")
	def single_socket_gpu_job():
	"""Job holding every GPU, with its cores confined to a single socket.

	The preconditions are asserted here so that every test sharing the
	fixture checks them in one place, and so that a topology that does not
	reproduce the condition says which precondition it failed rather than
	only that the step came up short on GPUs.
	"""

	job_id = atf.submit_job_sbatch(
	f"-N1 -n1 --gpus-per-node={_NUM_GPUS} --wrap 'sleep infinity'",
	fatal=True,
	)
	atf.wait_for_job_state(job_id, "RUNNING", fatal=True)

	job_gres = atf.get_job_parameter(job_id, "GRES", fatal=True, quiet=True)
	assert job_gres and f"gpu:{_NUM_GPUS}(IDX:" in job_gres, (
	f"Expected job to hold every GPU (gpu:{_NUM_GPUS}(IDX:...)), "
	f"got GRES={job_gres!r} - topology did not produce the "
	f"expected no-enforce-binding allocation"
	)

	num_cpus = atf.get_job_parameter(job_id, "NumCPUs", fatal=True, quiet=True)
	assert num_cpus == 1, (
	f"Expected job to hold a single socket (1 CPU with CoresPerSocket=1), "
	f"got NumCPUs={num_cpus} - topology did not confine the job to one "
	f"socket, so the step already covers every GPU-bound socket"
	)

	return job_id


	def test_gpu_socket_sharing():
	"""Test allocating multiple core-bound gpus with enforce-binding"""

	output = atf.run_command_output(
	f"srun --gres-flags=enforce-binding --ntasks-per-socket=1 \
	--cpus-per-task=1 --ntasks-per-node=2 -N1 \
	--gpus-per-task=1 scontrol show nodes {node_name} -d",
	fatal=True,
	)
	assert (
	re.search(r"GresUsed=gpu:(?:[^:]*:)?2\(", output) is not None
	), "Verify that job allocated 2 gpus"


	def test_gpu_socket_sharing_no_alloc():
	"""Test allocating multiple core-bound gpus with enforce-binding without enough resources"""

	output = atf.run_command(
	f"srun --gres-flags=enforce-binding --ntasks-per-socket=1 \
	--cpus-per-task=2 --ntasks-per-node=2 -N1 \
	--gpus-per-task=1 scontrol show nodes {node_name} -d",
	xfail=True,
	fatal=True,
	)
	assert output["exit_code"] != 0, "Verify that srun command failed"

	expected_msg = r"srun: error: .+ Requested node configuration is not available"
	assert (
	re.search(
	expected_msg,
	str(output["stderr"]),
	)
	is not None
	), "Verify that job is rejected with the right reason."


	@pytest.mark.xfail(_OVERLAP_GRES_UNFIXED, reason=_OVERLAP_GRES_REASON)
	def test_overlap_step_gets_all_gpus(single_socket_gpu_job):
	"""An --overlap step must receive every GPU it requests.

	srun --overlap shares all of the job's GRES with the step. The batch
	job requests -n1 so CR_Socket_Memory confines its cores to one socket,
	while --gpus-per-node (no enforce-binding) allocates every GPU across
	all sockets. The --overlap step lands on that single socket but, per the
	srun default, inherits every GPU the job holds, so it must be granted
	all of them even though most are bound to sockets its cores do not
	cover. A short allocation shows up as a denied device in the step's
	cgroup if task/cgroup and ConstrainDevices=yes is configured
	(ticket 24862 regression).
	"""

	job_id = single_socket_gpu_job

	result = atf.run_job(
	f"--jobid={job_id} --overlap -n1 --mem=0 "
	f"bash -c 'echo $SLURM_STEP_GPUS; echo $CUDA_VISIBLE_DEVICES'",
	fatal=True,
	quiet=True,
	)
	output = result["stdout"].splitlines()
	assert len(output) >= 2, (
	f"Expected two output lines (SLURM_STEP_GPUS and "
	f"CUDA_VISIBLE_DEVICES), got {result['stdout']!r}"
	)
	step_gpus = output[0].strip()
	visible_devices = output[1].strip()

	assert step_gpus == _ALL_GPUS, (
	f"Step SLURM_STEP_GPUS should be {_ALL_GPUS!r} (every GPU), "
	f"got {step_gpus!r} - the --overlap step was not granted all of "
	f"the job's GPUs (ticket 24862 regression)"
	)
	assert visible_devices == _ALL_GPUS, (
	f"Step CUDA_VISIBLE_DEVICES should be exported and agree with "
	f"SLURM_STEP_GPUS ({_ALL_GPUS}), got {visible_devices!r} - the "
	f"step's GPU allocation is not rendered consistently across both "
	f"variables (ticket 24862 regression)"
	)


	@pytest.mark.xfail(_OVERLAP_GRES_UNFIXED, reason=_OVERLAP_GRES_REASON)
	@pytest.mark.parametrize("first_step_overlaps", [True, False])
	def test_concurrent_steps_each_get_all_gpus(first_step_overlaps, single_socket_gpu_job):
	"""Concurrent steps sharing an allocation each receive every GPU.

	srun --overlap shares all resources with all other steps, even those
	that did not specify --overlap. Two steps running at the same time
	against the same allocation must therefore both be granted every GPU
	the job holds, whichever of them asked to overlap.
	"""

	job_id = single_socket_gpu_job

	first_step_gpus_file = "first_step_gpus"
	overlap = "--overlap " if first_step_overlaps else ""
	atf.run_command(
	f"srun --jobid={job_id} {overlap}-n1 --mem=0 bash -c "
	f"'echo $SLURM_STEP_GPUS > {first_step_gpus_file}; sleep infinity'",
	background=True,
	)
	atf.assert_file_contents(
	first_step_gpus_file,
	_ALL_GPUS,
	message=(
	f"First step SLURM_STEP_GPUS should be {_ALL_GPUS!r} (every GPU) "
	f"- the step was not granted all of the job's GPUs"
	),
	)

	result = atf.run_job(
	f"--jobid={job_id} --overlap -n1 --mem=0 bash -c 'echo $SLURM_STEP_GPUS'",
	fatal=True,
	quiet=True,
	)
	second_step_gpus = result["stdout"].strip()

	assert second_step_gpus == _ALL_GPUS, (
	f"Second step SLURM_STEP_GPUS should be {_ALL_GPUS!r} (every GPU), "
	f"got {second_step_gpus!r} - an --overlap step running alongside "
	f"another step was not granted all of the job's GPUs"
	)


	@pytest.mark.xfail(_OVERLAP_GRES_UNFIXED, reason=_OVERLAP_GRES_REASON)
	@pytest.mark.parametrize("requested_gpus", range(2, _NUM_GPUS))
	def test_overlap_step_gets_requested_gpu_subset(requested_gpus, single_socket_gpu_job):
	"""An --overlap step requesting a subset gets that many distinct GPUs.

	A job step may explicitly request a different generic resource count
	than the job, so an --overlap step asking for fewer GPUs than the job
	holds must be granted exactly that many, and no device twice.
	"""

	job_id = single_socket_gpu_job

	result = atf.run_job(
	f"--jobid={job_id} --overlap -n1 --mem=0 --gres=gpu:{requested_gpus} "
	f"bash -c 'echo $SLURM_STEP_GPUS'",
	fatal=True,
	quiet=True,
	)
	step_gpus = result["stdout"].strip()
	indices = step_gpus.split(",") if step_gpus else []

	assert len(indices) == requested_gpus, (
	f"Step requesting gpu:{requested_gpus} should be granted "
	f"{requested_gpus} GPUs, got SLURM_STEP_GPUS={step_gpus!r}"
	)
	assert len(set(indices)) == len(indices), (
	f"Step requesting gpu:{requested_gpus} should be granted distinct "
	f"GPUs, got SLURM_STEP_GPUS={step_gpus!r}"
	)


	def test_non_overlap_step_gets_all_gpus():
	"""A default (non --overlap) step also receives every GPU it requests.

	This is not a reproducer of ticket 24862 (that bug was specific to
	--overlap); it guards the shared non-overlap path against collateral
	breakage from the fix. srun's own step lands on a single socket
	(CR_Socket_Memory with -n1) while requesting every GPU across all
	sockets, so it must still be granted all of them.
	"""

	result = atf.run_job(
	f"-N1 -n1 --gpus-per-node={_NUM_GPUS} bash -c "
	f"'echo $SLURM_CPUS_ON_NODE; echo $SLURM_STEP_GPUS'",
	fatal=True,
	quiet=True,
	)
	output = result["stdout"].splitlines()
	assert len(output) >= 2, (
	f"Expected two output lines (SLURM_CPUS_ON_NODE and "
	f"SLURM_STEP_GPUS), got {result['stdout']!r}"
	)
	cpus_on_node = output[0].strip()
	step_gpus = output[1].strip()

	assert cpus_on_node == "1", (
	f"Expected the step to hold a single socket (1 CPU with "
	f"CoresPerSocket=1), got SLURM_CPUS_ON_NODE={cpus_on_node!r} - "
	f"topology did not confine the step to one socket, so it already "
	f"covers every GPU-bound socket"
	)
	assert step_gpus == _ALL_GPUS, (
	f"Non-overlap step SLURM_STEP_GPUS should be {_ALL_GPUS!r} "
	f"(every GPU), got {step_gpus!r}"
	)