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

 ############################################################################
 # Purpose: Test of Slurm functionality
 #          Test scheduling of gres/gpu and gres/mps
 ############################################################################
 # Copyright (C) SchedMD LLC.
 ############################################################################
 import atf
 import pytest
 import re

 mps_cnt = 100 * 2
 job_mps = int(mps_cnt * 0.5)
 step_mps = int(job_mps * 0.5)


 # Setup
 @pytest.fixture(scope="module", autouse=True)
 def setup():
     atf.require_config_parameter("SelectType", "select/cons_tres")
     atf.require_config_parameter("SelectTypeParameters", "CR_CPU")
     atf.require_config_parameter_includes("GresTypes", "gpu")
     atf.require_config_parameter_includes("GresTypes", "mps")
     atf.require_tty(1)
     atf.require_config_parameter(
         "Name",
         {"gpu": {"File": "/dev/tty[0-1]"}, "mps": {"Count": f"{mps_cnt}"}},
         source="gres",
     )
     atf.require_nodes(2, [("Gres", f"gpu:2,mps:{mps_cnt}"), ("CPUs", 6)])
     atf.require_slurm_running()


 @pytest.fixture(scope="module")
 def mps_nodes():
     return atf.run_job_nodes(f"--gres=mps:{job_mps} -N2 hostname")


 # Makes a commonly re-occuring batch script for file_in1 to allow for independent function tests
 @pytest.fixture(scope="function")
 def file_in_1a():
     file_in1 = atf.module_tmp_path / "input1"
     atf.make_bash_script(
         file_in1,
         """
     echo HOST:$SLURMD_NODENAME
     echo CUDA_VISIBLE_DEVICES:$CUDA_VISIBLE_DEVICES
     echo CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:$CUDA_MPS_ACTIVE_THREAD_PERCENTAGE
     sleep 5
     """,
     )
     return file_in1


 # Makes a commonly re-occuring batch script for file_in2 to allow for independent function tests
 @pytest.fixture(scope="function")
 def file_in_2a(file_in_1a):
     file_in2 = atf.module_tmp_path / "input2"
     atf.make_bash_script(
         file_in2,
         f"""
     srun --mem=0 --overlap --gres=mps:{job_mps}  {file_in_1a} &
     wait
     date
     srun --mem=0 --overlap --gres=mps:{step_mps} {file_in_1a} &
     srun --mem=0 --overlap --gres=mps:{step_mps} {file_in_1a} &
     wait
     date
     """,
     )
     return file_in2


 def test_environment_vars(mps_nodes):
     """Simple MPS request, check environment variables"""

     file_in1 = atf.module_tmp_path / "input1"

     atf.make_bash_script(
         file_in1,
         """
     echo HOST:$SLURMD_NODENAME
     echo CUDA_VISIBLE_DEVICES:$CUDA_VISIBLE_DEVICES
     echo CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:$CUDA_MPS_ACTIVE_THREAD_PERCENTAGE
     """,
     )

     results = atf.run_job(f"--gres=mps:{job_mps} -w {mps_nodes[0]} -n1 -t1 {file_in1}")
     assert results["exit_code"] == 0, "Job failed"
     assert (
         re.search(rf"HOST:{mps_nodes[0]}", results["stdout"]) is not None
     ), "HOST environmental variable not correct value"
     assert (
         match := re.search(r"CUDA_VISIBLE_DEVICES:(\d+)", results["stdout"])
     ) is not None and int(match.group(1)) == 0, "CUDA_VISIBLE_DEVICES != 0"
     assert (
         re.search(rf"CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:{job_mps}", results["stdout"])
         is not None
     ), "CUDA_MPS_ACTIVE_THREAD_PERCENTAGE environmental variable not correct value"


 def test_two_parallel_consumption_sbatch(mps_nodes, file_in_2a):
     """Run two steps in parallel to consume gres/mps using sbatch"""

     file_out1 = atf.module_tmp_path / "output1"

     job_id = atf.submit_job_sbatch(
         f"--gres=mps:{job_mps} -w {mps_nodes[0]} -n1 -t1 -o {file_out1} {file_in_2a}"
     )
     assert job_id != 0, "Job failed to submit"
     atf.wait_for_job_state(job_id, "DONE", timeout=30, fatal=True)
     atf.wait_for_file(file_out1)
     file_output = atf.run_command_output(f"cat {file_out1}")

     assert file_output is not None, "No output file"
     assert (
         len(re.findall(r"HOST:\w+", file_output)) == 3
     ), "HOST not found 3 times, once per job step, in output file"
     assert (
         re.search(r"CUDA_VISIBLE_DEVICES:\d+", file_output) is not None
     ), "CUDA_VISIBLE_DEVICES not found in output file"
     match = re.findall(r"(?s)CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:(\d+)", file_output)
     assert len(match) == 3, "Bad CUDA information about job (match != 3)"
     assert (
         sum(map(int, match)) == job_mps + step_mps * 2
     ), f"Bad CUDA percentage information about job (sum(map(int, match)) != {job_mps + step_mps * 2})"


 def test_two_parallel_consumption_salloc(mps_nodes, file_in_2a):
     """Run two steps in parallel to consume gres/mps using salloc"""

     output = atf.run_command_output(
         f"salloc --gres=mps:{job_mps} -w {mps_nodes[0]} -n1 -t1 {file_in_2a}",
         fatal=True,
     )

     assert (
         len(re.findall(r"HOST:\w+", output)) == 3
     ), "HOST not found 3 times, once per job step, in output file"
     assert (
         re.search(r"CUDA_VISIBLE_DEVICES:\d+", output) is not None
     ), "CUDA_VISIBLE_DEVICES not found in output"
     match = re.findall(r"(?s)CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:(\d+)", output)
     assert len(match) == 3, "Bad CUDA information about job (match != 3)"
     assert (
         sum(map(int, match)) == job_mps + step_mps * 2
     ), f"Bad CUDA percentage information about job ({sum(map(int, match))} != {job_mps + step_mps * 2})"


 def test_three_parallel_consumption_sbatch(mps_nodes, file_in_1a):
     """Run three steps in parallel to make sure steps get delayed as needed to avoid oversubscribing consumed MPS resources"""

     file_in2 = atf.module_tmp_path / "input2"
     file_out1 = atf.module_tmp_path / "output1"

     # Using -c6 for the job and -c2 for the steps to avoid issues if nodes have
     # HT. As Slurm allocates by Cores, with HT steps will allocate 2 CPUs even
     # if only 1 is requested, so the 3 steps won't run in parallel due lack of
     # CPUs instead of lack of MPS (that it's what we want to test).
     atf.make_bash_script(
         file_in2,
         f"""
     srun -vv --mem=0 -c2 --exact --gres=mps:{step_mps} {file_in_1a} &
     srun -vv --mem=0 -c2 --exact --gres=mps:{step_mps} {file_in_1a} &
     srun -vv --mem=0 -c2 --exact --gres=mps:{step_mps} {file_in_1a} &
     wait
     date
     """,
     )

     job_id = atf.submit_job_sbatch(
         f"--gres=mps:{job_mps} -w {mps_nodes[0]} -c6 -n1 -t1 -o {file_out1} {file_in2}"
     )

     assert job_id != 0, "Job failed to submit"
     atf.wait_for_job_state(job_id, "DONE", timeout=40, fatal=True)
     atf.wait_for_file(file_out1)
     file_output = atf.run_command_output(f"cat {file_out1}")
     assert file_output is not None, "No output file"
     assert (
         len(re.findall(r"HOST:\w+", file_output)) == 3
     ), "HOST not found 3 times, once per job, in output file"
     assert (
         re.search(r"CUDA_VISIBLE_DEVICES:\d+", file_output) is not None
     ), "CUDA_VISIBLE_DEVICES not found in output file"
     match = re.findall(r"(?s)CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:(\d+)", file_output)
     assert len(match) == 3, "Bad CUDA information about job (match != 3)"
     assert (
         sum(map(int, match)) == step_mps * 3
     ), f"Bad CUDA percentage information about job ({sum(map(int, match))} != {step_mps * 3})"
     assert atf.check_steps_delayed(
         job_id, file_output, 1
     ), "Failed to delay step for sufficient MPS resources (match != 1)"


 def test_consume_more_gresMps_than_allocated(mps_nodes, file_in_1a):
     """Run step to try to consume more gres/mps than allocated to the job"""

     file_in2 = atf.module_tmp_path / "input2"
     file_out1 = atf.module_tmp_path / "output1"
     job_mps2 = int(mps_cnt / 2)
     step_mps2 = job_mps2 + 1

     atf.make_bash_script(
         file_in2,
         f"""
     srun --mem=0 --overlap --gres=mps:{step_mps2} {file_in_1a}
     """,
     )

     job_id = atf.submit_job_sbatch(
         f"--gres=mps:{job_mps2} -w {mps_nodes[0]} -n1 -t1 -o {file_out1} {file_in2}"
     )

     assert job_id != 0, "Job failed to submit"
     atf.wait_for_job_state(job_id, "DONE", timeout=20, fatal=True)
     atf.wait_for_file(file_out1)
     file_output = atf.run_command_output(f"cat {file_out1}")
     assert file_output is not None, "No output file"
     assert (
         re.search(r"Unable to create step", file_output) is not None
     ), "Did not give expected 'Unable to create step' output in file"
     assert (
         re.search(r"CUDA_VISIBLE_DEVICES:\d+", file_output) is None
     ), "Failed to reject bad step (match != 1)"


 def test_run_multi_node_job(mps_nodes, file_in_1a):
     """Run multi-node job"""

     job_mps2 = int(mps_cnt / 2)
     node_cnt = len(mps_nodes)
     nodes_str = ",".join(mps_nodes)

     results = atf.run_job(
         f"--gres=mps:{job_mps2} -N{node_cnt} -w {nodes_str} -t1 {file_in_1a}"
     )

     assert results["exit_code"] == 0, "Job failed"
     host_match = re.findall(r"(?s)HOST:(\w+)", results["stdout"])
     assert len(host_match) is not None, "HOST not found"
     assert len(host_match) == len(
         mps_nodes
     ), f"Failed to get data from all nodes ({len(host_match)} != {len(mps_nodes)})"
     assert (
         host_match[0] != host_match[1]
     ), f"Two tasks ran on same node {host_match.group(0)}"


 def test_gresGPU_gresMPS_GPU_sharing(mps_nodes):
     """Make sure that gres/gpu and gres/mps jobs either do not share the same GPU or run at different times"""

     file_in1 = atf.module_tmp_path / "input1"
     file_in2 = atf.module_tmp_path / "input2"
     file_out1 = atf.module_tmp_path / "output1"
     file_out2 = atf.module_tmp_path / "output2"
     job_mps2 = int(job_mps / 2)

     atf.make_bash_script(
         file_in1,
         f"""
     echo HOST:$SLURMD_NODENAME CUDA_VISIBLE_DEVICES:$CUDA_VISIBLE_DEVICES CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:$CUDA_MPS_ACTIVE_THREAD_PERCENTAGE
     scontrol -dd show job $SLURM_JOB_ID
     sbatch --gres=mps:{job_mps2} -w $SLURMD_NODENAME -n1 -t1 -o {file_out2} -J test_job {file_in2}
     sleep 30
     """,
     )

     atf.make_bash_script(
         file_in2,
         """
     echo HOST:$SLURMD_NODENAME CUDA_VISIBLE_DEVICES:$CUDA_VISIBLE_DEVICES CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:$CUDA_MPS_ACTIVE_THREAD_PERCENTAGE
     scontrol -dd show job $SLURM_JOB_ID
     squeue --name=test_job --noheader --state=r --format=\"jobid=%i state=%T\"
     """,
     )

     job_id = atf.submit_job_sbatch(
         f"--gres=gpu:1 -w {mps_nodes[0]} -n1 -t1 -o {file_out1} -J 'test_job' {file_in1}"
     )

     assert job_id != 0, "Job failed to submit"
     atf.wait_for_job_state(job_id, "DONE", timeout=60, fatal=True)
     atf.wait_for_file(file_out1)
     file_output = atf.run_command_output(f"cat {file_out1}")
     assert file_output is not None, "No output file"
     assert (
         re.search(r"CUDA_VISIBLE_DEVICES:\d+", file_output) is not None
     ), "CUDA_VISIBLE_DEVICES not found in output file"
     assert (
         re.search(r"gpu:\d+\(IDX:\d+\)", file_output) is not None
     ), "GPU device index not found in output file"

     job_id2 = int(re.search(r"Submitted batch job (\d+)", file_output).group(1))

     assert job_id2 != 0, "Failed to submit second job"
     atf.wait_for_job_state(job_id2, "DONE", fatal=True)
     atf.wait_for_file(file_out2)
     file_output2 = atf.run_command_output(f"cat {file_out2}")
     assert file_output2 is not None, "No job2 output file"
     assert (
         re.search(r"CUDA_VISIBLE_DEVICES:\d+", file_output2) is not None
     ), "CUDA_VISIBLE_DEVICES not found in output2 file"

     assert (
         re.search(rf"CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:{job_mps2}", file_output2)
         is not None
     ), f"CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:{job_mps2} not found in output2 file"
     assert (
         re.search(rf"mps:{job_mps2}\(0/100,{job_mps2}/100\)", file_output2) is not None
     ), "Shared mps distribution across GPU devices not found in output2 file"

     assert (
         re.search(rf"jobid={job_id2} state=RUNNING", file_output2) is not None
         and re.search(rf"jobid={job_id} state=RUNNING", file_output2) is not None
     ), "Both jobs should be running at the same time"
	############################################################################
	# Purpose: Test of Slurm functionality
	# Test scheduling of gres/gpu and gres/mps
	############################################################################
	# Copyright (C) SchedMD LLC.
	############################################################################
	import atf
	import pytest
	import re

	mps_cnt = 100 * 2
	job_mps = int(mps_cnt * 0.5)
	step_mps = int(job_mps * 0.5)


	# Setup
	@pytest.fixture(scope="module", autouse=True)
	def setup():
	atf.require_config_parameter("SelectType", "select/cons_tres")
	atf.require_config_parameter("SelectTypeParameters", "CR_CPU")
	atf.require_config_parameter_includes("GresTypes", "gpu")
	atf.require_config_parameter_includes("GresTypes", "mps")
	atf.require_tty(1)
	atf.require_config_parameter(
	"Name",
	{"gpu": {"File": "/dev/tty[0-1]"}, "mps": {"Count": f"{mps_cnt}"}},
	source="gres",
	)
	atf.require_nodes(2, [("Gres", f"gpu:2,mps:{mps_cnt}"), ("CPUs", 6)])
	atf.require_slurm_running()


	@pytest.fixture(scope="module")
	def mps_nodes():
	return atf.run_job_nodes(f"--gres=mps:{job_mps} -N2 hostname")


	# Makes a commonly re-occuring batch script for file_in1 to allow for independent function tests
	@pytest.fixture(scope="function")
	def file_in_1a():
	file_in1 = atf.module_tmp_path / "input1"
	atf.make_bash_script(
	file_in1,
	"""
	echo HOST:$SLURMD_NODENAME
	echo CUDA_VISIBLE_DEVICES:$CUDA_VISIBLE_DEVICES
	echo CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:$CUDA_MPS_ACTIVE_THREAD_PERCENTAGE
	sleep 5
	""",
	)
	return file_in1


	# Makes a commonly re-occuring batch script for file_in2 to allow for independent function tests
	@pytest.fixture(scope="function")
	def file_in_2a(file_in_1a):
	file_in2 = atf.module_tmp_path / "input2"
	atf.make_bash_script(
	file_in2,
	f"""
	srun --mem=0 --overlap --gres=mps:{job_mps} {file_in_1a} &
	wait
	date
	srun --mem=0 --overlap --gres=mps:{step_mps} {file_in_1a} &
	srun --mem=0 --overlap --gres=mps:{step_mps} {file_in_1a} &
	wait
	date
	""",
	)
	return file_in2


	def test_environment_vars(mps_nodes):
	"""Simple MPS request, check environment variables"""

	file_in1 = atf.module_tmp_path / "input1"

	atf.make_bash_script(
	file_in1,
	"""
	echo HOST:$SLURMD_NODENAME
	echo CUDA_VISIBLE_DEVICES:$CUDA_VISIBLE_DEVICES
	echo CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:$CUDA_MPS_ACTIVE_THREAD_PERCENTAGE
	""",
	)

	results = atf.run_job(f"--gres=mps:{job_mps} -w {mps_nodes[0]} -n1 -t1 {file_in1}")
	assert results["exit_code"] == 0, "Job failed"
	assert (
	re.search(rf"HOST:{mps_nodes[0]}", results["stdout"]) is not None
	), "HOST environmental variable not correct value"
	assert (
	match := re.search(r"CUDA_VISIBLE_DEVICES:(\d+)", results["stdout"])
	) is not None and int(match.group(1)) == 0, "CUDA_VISIBLE_DEVICES != 0"
	assert (
	re.search(rf"CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:{job_mps}", results["stdout"])
	is not None
	), "CUDA_MPS_ACTIVE_THREAD_PERCENTAGE environmental variable not correct value"


	def test_two_parallel_consumption_sbatch(mps_nodes, file_in_2a):
	"""Run two steps in parallel to consume gres/mps using sbatch"""

	file_out1 = atf.module_tmp_path / "output1"

	job_id = atf.submit_job_sbatch(
	f"--gres=mps:{job_mps} -w {mps_nodes[0]} -n1 -t1 -o {file_out1} {file_in_2a}"
	)
	assert job_id != 0, "Job failed to submit"
	atf.wait_for_job_state(job_id, "DONE", timeout=30, fatal=True)
	atf.wait_for_file(file_out1)
	file_output = atf.run_command_output(f"cat {file_out1}")

	assert file_output is not None, "No output file"
	assert (
	len(re.findall(r"HOST:\w+", file_output)) == 3
	), "HOST not found 3 times, once per job step, in output file"
	assert (
	re.search(r"CUDA_VISIBLE_DEVICES:\d+", file_output) is not None
	), "CUDA_VISIBLE_DEVICES not found in output file"
	match = re.findall(r"(?s)CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:(\d+)", file_output)
	assert len(match) == 3, "Bad CUDA information about job (match != 3)"
	assert (
	sum(map(int, match)) == job_mps + step_mps * 2
	), f"Bad CUDA percentage information about job (sum(map(int, match)) != {job_mps + step_mps * 2})"


	def test_two_parallel_consumption_salloc(mps_nodes, file_in_2a):
	"""Run two steps in parallel to consume gres/mps using salloc"""

	output = atf.run_command_output(
	f"salloc --gres=mps:{job_mps} -w {mps_nodes[0]} -n1 -t1 {file_in_2a}",
	fatal=True,
	)

	assert (
	len(re.findall(r"HOST:\w+", output)) == 3
	), "HOST not found 3 times, once per job step, in output file"
	assert (
	re.search(r"CUDA_VISIBLE_DEVICES:\d+", output) is not None
	), "CUDA_VISIBLE_DEVICES not found in output"
	match = re.findall(r"(?s)CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:(\d+)", output)
	assert len(match) == 3, "Bad CUDA information about job (match != 3)"
	assert (
	sum(map(int, match)) == job_mps + step_mps * 2
	), f"Bad CUDA percentage information about job ({sum(map(int, match))} != {job_mps + step_mps * 2})"


	def test_three_parallel_consumption_sbatch(mps_nodes, file_in_1a):
	"""Run three steps in parallel to make sure steps get delayed as needed to avoid oversubscribing consumed MPS resources"""

	file_in2 = atf.module_tmp_path / "input2"
	file_out1 = atf.module_tmp_path / "output1"

	# Using -c6 for the job and -c2 for the steps to avoid issues if nodes have
	# HT. As Slurm allocates by Cores, with HT steps will allocate 2 CPUs even
	# if only 1 is requested, so the 3 steps won't run in parallel due lack of
	# CPUs instead of lack of MPS (that it's what we want to test).
	atf.make_bash_script(
	file_in2,
	f"""
	srun -vv --mem=0 -c2 --exact --gres=mps:{step_mps} {file_in_1a} &
	srun -vv --mem=0 -c2 --exact --gres=mps:{step_mps} {file_in_1a} &
	srun -vv --mem=0 -c2 --exact --gres=mps:{step_mps} {file_in_1a} &
	wait
	date
	""",
	)

	job_id = atf.submit_job_sbatch(
	f"--gres=mps:{job_mps} -w {mps_nodes[0]} -c6 -n1 -t1 -o {file_out1} {file_in2}"
	)

	assert job_id != 0, "Job failed to submit"
	atf.wait_for_job_state(job_id, "DONE", timeout=40, fatal=True)
	atf.wait_for_file(file_out1)
	file_output = atf.run_command_output(f"cat {file_out1}")
	assert file_output is not None, "No output file"
	assert (
	len(re.findall(r"HOST:\w+", file_output)) == 3
	), "HOST not found 3 times, once per job, in output file"
	assert (
	re.search(r"CUDA_VISIBLE_DEVICES:\d+", file_output) is not None
	), "CUDA_VISIBLE_DEVICES not found in output file"
	match = re.findall(r"(?s)CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:(\d+)", file_output)
	assert len(match) == 3, "Bad CUDA information about job (match != 3)"
	assert (
	sum(map(int, match)) == step_mps * 3
	), f"Bad CUDA percentage information about job ({sum(map(int, match))} != {step_mps * 3})"
	assert atf.check_steps_delayed(
	job_id, file_output, 1
	), "Failed to delay step for sufficient MPS resources (match != 1)"


	def test_consume_more_gresMps_than_allocated(mps_nodes, file_in_1a):
	"""Run step to try to consume more gres/mps than allocated to the job"""

	file_in2 = atf.module_tmp_path / "input2"
	file_out1 = atf.module_tmp_path / "output1"
	job_mps2 = int(mps_cnt / 2)
	step_mps2 = job_mps2 + 1

	atf.make_bash_script(
	file_in2,
	f"""
	srun --mem=0 --overlap --gres=mps:{step_mps2} {file_in_1a}
	""",
	)

	job_id = atf.submit_job_sbatch(
	f"--gres=mps:{job_mps2} -w {mps_nodes[0]} -n1 -t1 -o {file_out1} {file_in2}"
	)

	assert job_id != 0, "Job failed to submit"
	atf.wait_for_job_state(job_id, "DONE", timeout=20, fatal=True)
	atf.wait_for_file(file_out1)
	file_output = atf.run_command_output(f"cat {file_out1}")
	assert file_output is not None, "No output file"
	assert (
	re.search(r"Unable to create step", file_output) is not None
	), "Did not give expected 'Unable to create step' output in file"
	assert (
	re.search(r"CUDA_VISIBLE_DEVICES:\d+", file_output) is None
	), "Failed to reject bad step (match != 1)"


	def test_run_multi_node_job(mps_nodes, file_in_1a):
	"""Run multi-node job"""

	job_mps2 = int(mps_cnt / 2)
	node_cnt = len(mps_nodes)
	nodes_str = ",".join(mps_nodes)

	results = atf.run_job(
	f"--gres=mps:{job_mps2} -N{node_cnt} -w {nodes_str} -t1 {file_in_1a}"
	)

	assert results["exit_code"] == 0, "Job failed"
	host_match = re.findall(r"(?s)HOST:(\w+)", results["stdout"])
	assert len(host_match) is not None, "HOST not found"
	assert len(host_match) == len(
	mps_nodes
	), f"Failed to get data from all nodes ({len(host_match)} != {len(mps_nodes)})"
	assert (
	host_match[0] != host_match[1]
	), f"Two tasks ran on same node {host_match.group(0)}"


	def test_gresGPU_gresMPS_GPU_sharing(mps_nodes):
	"""Make sure that gres/gpu and gres/mps jobs either do not share the same GPU or run at different times"""

	file_in1 = atf.module_tmp_path / "input1"
	file_in2 = atf.module_tmp_path / "input2"
	file_out1 = atf.module_tmp_path / "output1"
	file_out2 = atf.module_tmp_path / "output2"
	job_mps2 = int(job_mps / 2)

	atf.make_bash_script(
	file_in1,
	f"""
	echo HOST:$SLURMD_NODENAME CUDA_VISIBLE_DEVICES:$CUDA_VISIBLE_DEVICES CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:$CUDA_MPS_ACTIVE_THREAD_PERCENTAGE
	scontrol -dd show job $SLURM_JOB_ID
	sbatch --gres=mps:{job_mps2} -w $SLURMD_NODENAME -n1 -t1 -o {file_out2} -J test_job {file_in2}
	sleep 30
	""",
	)

	atf.make_bash_script(
	file_in2,
	"""
	echo HOST:$SLURMD_NODENAME CUDA_VISIBLE_DEVICES:$CUDA_VISIBLE_DEVICES CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:$CUDA_MPS_ACTIVE_THREAD_PERCENTAGE
	scontrol -dd show job $SLURM_JOB_ID
	squeue --name=test_job --noheader --state=r --format=\"jobid=%i state=%T\"
	""",
	)

	job_id = atf.submit_job_sbatch(
	f"--gres=gpu:1 -w {mps_nodes[0]} -n1 -t1 -o {file_out1} -J 'test_job' {file_in1}"
	)

	assert job_id != 0, "Job failed to submit"
	atf.wait_for_job_state(job_id, "DONE", timeout=60, fatal=True)
	atf.wait_for_file(file_out1)
	file_output = atf.run_command_output(f"cat {file_out1}")
	assert file_output is not None, "No output file"
	assert (
	re.search(r"CUDA_VISIBLE_DEVICES:\d+", file_output) is not None
	), "CUDA_VISIBLE_DEVICES not found in output file"
	assert (
	re.search(r"gpu:\d+\(IDX:\d+\)", file_output) is not None
	), "GPU device index not found in output file"

	job_id2 = int(re.search(r"Submitted batch job (\d+)", file_output).group(1))

	assert job_id2 != 0, "Failed to submit second job"
	atf.wait_for_job_state(job_id2, "DONE", fatal=True)
	atf.wait_for_file(file_out2)
	file_output2 = atf.run_command_output(f"cat {file_out2}")
	assert file_output2 is not None, "No job2 output file"
	assert (
	re.search(r"CUDA_VISIBLE_DEVICES:\d+", file_output2) is not None
	), "CUDA_VISIBLE_DEVICES not found in output2 file"

	assert (
	re.search(rf"CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:{job_mps2}", file_output2)
	is not None
	), f"CUDA_MPS_ACTIVE_THREAD_PERCENTAGE:{job_mps2} not found in output2 file"
	assert (
	re.search(rf"mps:{job_mps2}\(0/100,{job_mps2}/100\)", file_output2) is not None
	), "Shared mps distribution across GPU devices not found in output2 file"

	assert (
	re.search(rf"jobid={job_id2} state=RUNNING", file_output2) is not None
	and re.search(rf"jobid={job_id} state=RUNNING", file_output2) is not None
	), "Both jobs should be running at the same time"