| ############################################################################ |
| # Copyright (C) SchedMD LLC. |
| ############################################################################ |
| import atf |
| import pytest |
| import re |
| |
| |
| # 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_tty(0) |
| atf.require_config_parameter("Name", {"gpu": {"File": "/dev/tty0"}}, source="gres") |
| atf.require_config_parameter_includes("GresTypes", "gpu") |
| atf.require_nodes(1, [("Gres", "gpu:1"), ("RealMemory", 1)]) |
| atf.require_slurm_running() |
| |
| |
| # Global variables set via init_gpu_vars |
| node_count = 1 |
| gpus_per_node = 1 |
| gpu_count = 1 |
| cpus_per_gpu = 1 |
| sockets_per_node = 1 |
| task_count = 1 |
| gpus_per_task = 1 |
| memory_per_gpu = 1 |
| |
| |
| @pytest.fixture(scope="module") |
| def init_gpu_vars(): |
| global node_count, gpus_per_node, gpu_count, task_count, gpus_per_task, mem_per_gpu |
| |
| nodes_with_gpus = 0 |
| min_gpus_per_node = 1024 |
| min_cpus_per_node = 1024 |
| min_sockets_per_node = 1024 |
| min_memory_per_node = 1024 |
| |
| for node_name in atf.nodes: |
| node_dict = atf.nodes[node_name] |
| if "gres" in node_dict and node_dict["gres"] is not None: |
| if match := re.search(r"gpu:(\d+)", node_dict["gres"]): |
| nodes_with_gpus += 1 |
| node_gpu_count = int(match.group(1)) |
| node_cpu_count = node_dict["cpus"] |
| node_memory = node_dict["real_memory"] |
| node_socket_count = node_dict["sockets"] |
| if node_cpu_count < node_gpu_count: |
| node_gpu_count = node_cpu_count |
| if node_gpu_count < min_gpus_per_node: |
| min_gpus_per_node = node_gpu_count |
| if node_cpu_count < min_cpus_per_node: |
| min_cpus_per_node = node_cpu_count |
| if node_socket_count < min_sockets_per_node: |
| min_sockets_per_node = node_socket_count |
| if node_memory < min_memory_per_node: |
| min_memory_per_node = node_memory |
| |
| node_count = nodes_with_gpus |
| gpus_per_node = min_gpus_per_node |
| gpu_count = gpus_per_node * node_count |
| # sockets_per_node = min_sockets_per_node |
| if gpus_per_node % 2 == 0 and min_cpus_per_node > 1: |
| task_count = node_count * 2 |
| else: |
| task_count = node_count |
| gpus_per_task = int(gpu_count / task_count) |
| memory_per_gpu = int(min_memory_per_node / min_gpus_per_node) |
| if memory_per_gpu < 1: |
| pytest.skip( |
| "This test requires at least one node with {min_gpus_per_node} memory" |
| ) |
| |
| |
| def test_gpus_per_cpu(init_gpu_vars): |
| """Test a batch job with various gpu options including ---gpus""" |
| |
| gpu_bind = "closest" |
| gpu_freq = "medium" |
| |
| job_id = atf.submit_job_sbatch( |
| f'--cpus-per-gpu={cpus_per_gpu} --gpu-bind={gpu_bind} --gpu-freq={gpu_freq} --gpus={gpu_count} --gpus-per-node={gpus_per_node} --gpus-per-task={gpus_per_task} --mem-per-gpu={memory_per_gpu} --nodes={node_count} --ntasks={task_count} -t1 --wrap "true"', |
| fatal=True, |
| ) |
| job_dict = atf.get_job(job_id) |
| |
| assert job_dict["CpusPerTres"] == f"gres/gpu:{cpus_per_gpu}" |
| assert job_dict["MemPerTres"] == f"gres/gpu:{memory_per_gpu}" |
| assert job_dict["TresBind"] == f"gres/gpu:{gpu_bind}" |
| assert job_dict["TresFreq"] == f"gpu:{gpu_freq}" |
| assert job_dict["TresPerJob"] == f"gres/gpu:{gpu_count}" |
| assert job_dict["TresPerNode"] == f"gres/gpu:{gpus_per_node}" |
| if atf.get_version("bin/scontrol") >= (24, 5): |
| assert job_dict["TresPerTask"] == f"gres/gpu={gpus_per_task}" |
| else: |
| assert job_dict["TresPerTask"] == f"gres/gpu:{gpus_per_task}" |
| |
| |
| def test_gpus_per_socket(init_gpu_vars): |
| """Test a batch job with various gpu options including --gpus-per-socket""" |
| gpus_per_socket = 1 |
| |
| job_id = atf.submit_job_sbatch( |
| f'--cpus-per-gpu={cpus_per_gpu} --gpus-per-socket={gpus_per_socket} --sockets-per-node={sockets_per_node} --nodes={node_count} --ntasks={task_count} -t1 --wrap "true"', |
| fatal=True, |
| ) |
| job_dict = atf.get_job(job_id) |
| |
| assert job_dict["TresPerSocket"] == f"gres/gpu:{gpus_per_socket}" |
