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

 ############################################################################
 # Copyright (C) SchedMD LLC.
 ############################################################################
 import logging
 import re
 import time

 import pytest

 import atf

 port_range_size = 9


 # Setup
 @pytest.fixture(scope="module", autouse=True)
 def setup(safe_port_range):
     lo, hi = safe_port_range
     atf.require_auto_config("wants to edit SrunPortRange and create a node")
     atf.require_config_parameter("SrunPortRange", f"{lo}-{hi}")
     atf.require_nodes(145, [("CPUs", 2), ("RealMemory", 2)])
     atf.require_slurm_running()


 @pytest.fixture(scope="module")
 def safe_port_range():
     """Returns (lo, hi) for a free port range outside ip_local_port_range.

     Placing the range outside the ephemeral range and ensuring that none is
     blocked avoids future interferences and ensures current availability.
     """
     range_str = atf.run_command_output(
         "cat /proc/sys/net/ipv4/ip_local_port_range", fatal=True, quiet=True
     )
     ephem_lo, ephem_hi = (int(x) for x in range_str.split())

     # Prefer just above the ephemeral range; fall back to just below.
     candidates = (ephem_hi + 100, ephem_lo - 100 - port_range_size)
     for lo in candidates:
         hi = lo + port_range_size - 1
         if not (1024 <= lo and hi <= 65535 and (hi < ephem_lo or lo > ephem_hi)):
             continue
         if is_port_range_available(lo, hi):
             return lo, hi
         logging.debug(f"[PORT_PICK] Skipping candidate {lo}-{hi}: already blocked")

     pytest.fail(
         f"Cannot find a {port_range_size}-port range outside ephemeral ({ephem_lo}-{ephem_hi}) with no active binders"
     )


 def is_port_range_available(lo, hi):
     """Return True if every port in [lo, hi] is available/bindable.

     Any TCP socket on a local port in the range whose state is NOT TIME-WAIT
     will make srun's bind() fail with EADDRINUSE, even though srun sets
     SO_REUSEADDR.
     """
     logging.debug(f"[PORT_MONITOR] Checking ports {lo}-{hi}")

     ss_result = atf.run_command(
         f'ss -tanH "sport >= :{lo} and sport <= :{hi}"',
         timeout=10,
         quiet=True,
     )
     if ss_result["exit_code"] != 0:
         logging.debug(f"[PORT_MONITOR] ss command failed: {ss_result['stderr']}")
         return False

     # ss -tanH columns: State Recv-Q Send-Q Local-Addr:Port Peer-Addr:Port
     for line in ss_result["stdout"].splitlines():
         line = line.strip()
         if not line:
             continue
         state = line.split()[0]
         if state == "TIME-WAIT":
             continue
         logging.debug(f"[PORT_MONITOR] Blocking socket ({state}): {line}")
         return False

     return True


 def wait_for_srun_ports_clear(lo, hi, timeout=90):
     """Wait until all ports in [lo, hi] are released.

     Timeout defaults to 90s so FIN-WAIT-2 sockets from a prior srun can
     drain (bounded by net.ipv4.tcp_fin_timeout, default 60s).
     """
     logging.debug(
         f"[PORT_WAIT] Waiting for ports {lo}-{hi} to be released (timeout: {timeout}s)"
     )

     for t in atf.timer(timeout=timeout, fatal=True, quiet=True):
         if is_port_range_available(lo, hi):
             logging.debug("[PORT_WAIT] All ports available")
             break

         logging.debug(
             f"[PORT_WAIT] Still waiting... some ports not yet available ({t:.1f}s remaining)"
         )


 @pytest.mark.parametrize("nodes", [1, 10, 48, 49, 96, 100, 144])
 def test_srun_ports_in_range(nodes, safe_port_range):
     """Test srun uses the right SrunPortRange"""

     lo, hi = safe_port_range

     # Wait for ports to be open before running test
     wait_for_srun_ports_clear(lo, hi)

     command = """bash -c '
         echo "[DEBUG] Starting port check" >&2
         echo "[DEBUG] Environment: SLURM_JOBID=$SLURM_JOBID SLURM_STEPID=$SLURM_STEPID" >&2
         echo "[DEBUG] Hostname: $(hostname)" >&2
         echo "[DEBUG] PID: $$" >&2

         echo "[DEBUG] Running scontrol show step..." >&2
         step_info=$(scontrol show step $SLURM_JOBID.$SLURM_STEPID 2>&1)
         scontrol_exit_code=$?
         echo "[DEBUG] scontrol exit code: $scontrol_exit_code" >&2
         echo "[DEBUG] scontrol output length: ${#step_info}" >&2
         echo "[DEBUG] scontrol output: $step_info" >&2

         if [[ $scontrol_exit_code -ne 0 ]]; then
             echo "[ERROR] scontrol failed with exit code $scontrol_exit_code" >&2
             exit 1
         fi

         echo "[DEBUG] Extracting SrunHost info..." >&2
         srun_host_line=$(echo "$step_info" | grep SrunHost)
         echo "[DEBUG] SrunHost line: $srun_host_line" >&2

         if [[ -z "$srun_host_line" ]]; then
             echo "[ERROR] No SrunHost found in step info" >&2
             exit 2
         fi

         task_id=$(echo "$srun_host_line" | awk -F: '"'"'{print $3}'"'"')
         echo "[DEBUG] Extracted task_id: $task_id" >&2

         if [[ -z "$task_id" ]]; then
             echo "[ERROR] Could not extract task_id" >&2
             exit 3
         fi

         echo "[DEBUG] Checking if process $task_id exists..." >&2
         if ! ps -p $task_id > /dev/null 2>&1; then
             echo "[ERROR] Process $task_id does not exist" >&2
             echo "[DEBUG] Current srun processes:" >&2
             ps aux | grep srun | grep -v grep >&2
             exit 4
         fi

         echo "[DEBUG] Running ss to find listening ports for PID $task_id..." >&2
         ss_output=$(ss -tlnp 2>&1)
         ss_exit_code=$?
         echo "[DEBUG] ss exit code: $ss_exit_code" >&2

         if [[ $ss_exit_code -ne 0 ]]; then
             echo "[ERROR] ss failed with exit code $ss_exit_code" >&2
             echo "[DEBUG] ss error output: $ss_output" >&2
             exit 5
         fi

         echo "[DEBUG] Filtering ss output for PID $task_id..." >&2
         filtered_output=$(echo "$ss_output" | grep "pid=$task_id,")
         echo "[DEBUG] Filtered output: $filtered_output" >&2

         echo "[DEBUG] Using filtered output as listen ports..." >&2
         listen_ports="$filtered_output"
         echo "[DEBUG] LISTEN lines: $listen_ports" >&2

         echo "[DEBUG] Extracting port numbers..." >&2
         port_numbers=$(echo "$listen_ports" | awk '"'"'{print $4}'"'"' | awk -F: '"'"'{print $NF}'"'"')
         echo "[DEBUG] Port numbers: $port_numbers" >&2

         echo "[DEBUG] Final output:" >&2
         echo "$port_numbers"

         echo "[DEBUG] Completed" >&2
     \'    """

     start_time = time.time()

     # Log existing srun processes before test
     ps_result = atf.run_command("ps aux | grep srun | grep -v grep")
     # Logging for Ticket 19089
     logging.debug(f"[TEST] Existing srun processes before test: {ps_result['stdout']}")

     result = atf.run_job(f"-N{nodes} {command}", timeout=180)
     end_time = time.time()

     # Logging for Ticket 19089
     logging.debug(f"[TEST] Command took {end_time - start_time:.3f} seconds")

     output = result["stdout"].split("\n")
     count = 0
     for port_string in output:
         # Ignore blank lines
         if len(port_string) < 4:
             continue
         count += 1
         port_int = int(port_string)
         assert lo <= port_int <= hi, f"Port {port_int} is not in range {lo}-{hi}"

     # From the docs:
     # "A single srun opens 4 listening ports plus 2 more for every 48 hosts
     # beyond the first 48."
     ports = nodes * (4 + 2 * ((nodes - 1) // 48))
     assert count == ports, f"srun with -N{nodes} should use {ports} ports, not {count}"


 @pytest.mark.parametrize("nodes", [145])
 def test_srun_ports_out_of_range(nodes, safe_port_range):
     """Test sruns with too many nodes, so with not enough SrunPortRange"""

     lo, hi = safe_port_range

     # Wait for ports to be open before running test
     wait_for_srun_ports_clear(lo, hi)

     result = atf.run_job_error(f"-t1 -N{nodes} sleep 1", fatal=True, xfail=True)

     regex = rf"all ports in range .{lo}, {hi}. exhausted"
     assert (
         re.search(regex, result) is not None
     ), "srun's stderr should contain the 'all ports in range exhausted' message"


 def test_out_of_srun_ports(safe_port_range):
     """Test exhausted ports"""

     lo, hi = safe_port_range

     # Wait for ports to be open before running test
     wait_for_srun_ports_clear(lo, hi)

     job_id1 = atf.submit_job_sbatch('-N1 -o/dev/null --wrap="srun sleep 30"')
     job_id2 = atf.submit_job_sbatch('-N1 -o/dev/null --wrap="srun sleep 30"')

     atf.wait_for_step(job_id1, 0, fatal=True)
     atf.wait_for_step(job_id2, 0, fatal=True)

     result = atf.run_job_error("-t1 -N1 sleep 1", fatal=True, xfail=True)

     regex = rf"all ports in range .{lo}, {hi}. exhausted"
     assert (
         re.search(regex, result) is not None
     ), "srun's stderr should contain the 'all ports in range exhausted' message"
	############################################################################
	# Copyright (C) SchedMD LLC.
	############################################################################
	import logging
	import re
	import time

	import pytest

	import atf

	port_range_size = 9


	# Setup
	@pytest.fixture(scope="module", autouse=True)
	def setup(safe_port_range):
	lo, hi = safe_port_range
	atf.require_auto_config("wants to edit SrunPortRange and create a node")
	atf.require_config_parameter("SrunPortRange", f"{lo}-{hi}")
	atf.require_nodes(145, [("CPUs", 2), ("RealMemory", 2)])
	atf.require_slurm_running()


	@pytest.fixture(scope="module")
	def safe_port_range():
	"""Returns (lo, hi) for a free port range outside ip_local_port_range.

	Placing the range outside the ephemeral range and ensuring that none is
	blocked avoids future interferences and ensures current availability.
	"""
	range_str = atf.run_command_output(
	"cat /proc/sys/net/ipv4/ip_local_port_range", fatal=True, quiet=True
	)
	ephem_lo, ephem_hi = (int(x) for x in range_str.split())

	# Prefer just above the ephemeral range; fall back to just below.
	candidates = (ephem_hi + 100, ephem_lo - 100 - port_range_size)
	for lo in candidates:
	hi = lo + port_range_size - 1
	if not (1024 <= lo and hi <= 65535 and (hi < ephem_lo or lo > ephem_hi)):
	continue
	if is_port_range_available(lo, hi):
	return lo, hi
	logging.debug(f"[PORT_PICK] Skipping candidate {lo}-{hi}: already blocked")

	pytest.fail(
	f"Cannot find a {port_range_size}-port range outside ephemeral ({ephem_lo}-{ephem_hi}) with no active binders"
	)


	def is_port_range_available(lo, hi):
	"""Return True if every port in [lo, hi] is available/bindable.

	Any TCP socket on a local port in the range whose state is NOT TIME-WAIT
	will make srun's bind() fail with EADDRINUSE, even though srun sets
	SO_REUSEADDR.
	"""
	logging.debug(f"[PORT_MONITOR] Checking ports {lo}-{hi}")

	ss_result = atf.run_command(
	f'ss -tanH "sport >= :{lo} and sport <= :{hi}"',
	timeout=10,
	quiet=True,
	)
	if ss_result["exit_code"] != 0:
	logging.debug(f"[PORT_MONITOR] ss command failed: {ss_result['stderr']}")
	return False

	# ss -tanH columns: State Recv-Q Send-Q Local-Addr:Port Peer-Addr:Port
	for line in ss_result["stdout"].splitlines():
	line = line.strip()
	if not line:
	continue
	state = line.split()[0]
	if state == "TIME-WAIT":
	continue
	logging.debug(f"[PORT_MONITOR] Blocking socket ({state}): {line}")
	return False

	return True


	def wait_for_srun_ports_clear(lo, hi, timeout=90):
	"""Wait until all ports in [lo, hi] are released.

	Timeout defaults to 90s so FIN-WAIT-2 sockets from a prior srun can
	drain (bounded by net.ipv4.tcp_fin_timeout, default 60s).
	"""
	logging.debug(
	f"[PORT_WAIT] Waiting for ports {lo}-{hi} to be released (timeout: {timeout}s)"
	)

	for t in atf.timer(timeout=timeout, fatal=True, quiet=True):
	if is_port_range_available(lo, hi):
	logging.debug("[PORT_WAIT] All ports available")
	break

	logging.debug(
	f"[PORT_WAIT] Still waiting... some ports not yet available ({t:.1f}s remaining)"
	)


	@pytest.mark.parametrize("nodes", [1, 10, 48, 49, 96, 100, 144])
	def test_srun_ports_in_range(nodes, safe_port_range):
	"""Test srun uses the right SrunPortRange"""

	lo, hi = safe_port_range

	# Wait for ports to be open before running test
	wait_for_srun_ports_clear(lo, hi)

	command = """bash -c '
	echo "[DEBUG] Starting port check" >&2
	echo "[DEBUG] Environment: SLURM_JOBID=$SLURM_JOBID SLURM_STEPID=$SLURM_STEPID" >&2
	echo "[DEBUG] Hostname: $(hostname)" >&2
	echo "[DEBUG] PID: $$" >&2

	echo "[DEBUG] Running scontrol show step..." >&2
	step_info=$(scontrol show step $SLURM_JOBID.$SLURM_STEPID 2>&1)
	scontrol_exit_code=$?
	echo "[DEBUG] scontrol exit code: $scontrol_exit_code" >&2
	echo "[DEBUG] scontrol output length: ${#step_info}" >&2
	echo "[DEBUG] scontrol output: $step_info" >&2

	if [[ $scontrol_exit_code -ne 0 ]]; then
	echo "[ERROR] scontrol failed with exit code $scontrol_exit_code" >&2
	exit 1
	fi

	echo "[DEBUG] Extracting SrunHost info..." >&2
	srun_host_line=$(echo "$step_info" \| grep SrunHost)
	echo "[DEBUG] SrunHost line: $srun_host_line" >&2

	if [[ -z "$srun_host_line" ]]; then
	echo "[ERROR] No SrunHost found in step info" >&2
	exit 2
	fi

	task_id=$(echo "$srun_host_line" \| awk -F: '"'"'{print $3}'"'"')
	echo "[DEBUG] Extracted task_id: $task_id" >&2

	if [[ -z "$task_id" ]]; then
	echo "[ERROR] Could not extract task_id" >&2
	exit 3
	fi

	echo "[DEBUG] Checking if process $task_id exists..." >&2
	if ! ps -p $task_id > /dev/null 2>&1; then
	echo "[ERROR] Process $task_id does not exist" >&2
	echo "[DEBUG] Current srun processes:" >&2
	ps aux \| grep srun \| grep -v grep >&2
	exit 4
	fi

	echo "[DEBUG] Running ss to find listening ports for PID $task_id..." >&2
	ss_output=$(ss -tlnp 2>&1)
	ss_exit_code=$?
	echo "[DEBUG] ss exit code: $ss_exit_code" >&2

	if [[ $ss_exit_code -ne 0 ]]; then
	echo "[ERROR] ss failed with exit code $ss_exit_code" >&2
	echo "[DEBUG] ss error output: $ss_output" >&2
	exit 5
	fi

	echo "[DEBUG] Filtering ss output for PID $task_id..." >&2
	filtered_output=$(echo "$ss_output" \| grep "pid=$task_id,")
	echo "[DEBUG] Filtered output: $filtered_output" >&2

	echo "[DEBUG] Using filtered output as listen ports..." >&2
	listen_ports="$filtered_output"
	echo "[DEBUG] LISTEN lines: $listen_ports" >&2

	echo "[DEBUG] Extracting port numbers..." >&2
	port_numbers=$(echo "$listen_ports" \| awk '"'"'{print $4}'"'"' \| awk -F: '"'"'{print $NF}'"'"')
	echo "[DEBUG] Port numbers: $port_numbers" >&2

	echo "[DEBUG] Final output:" >&2
	echo "$port_numbers"

	echo "[DEBUG] Completed" >&2
	\' """

	start_time = time.time()

	# Log existing srun processes before test
	ps_result = atf.run_command("ps aux \| grep srun \| grep -v grep")
	# Logging for Ticket 19089
	logging.debug(f"[TEST] Existing srun processes before test: {ps_result['stdout']}")

	result = atf.run_job(f"-N{nodes} {command}", timeout=180)
	end_time = time.time()

	# Logging for Ticket 19089
	logging.debug(f"[TEST] Command took {end_time - start_time:.3f} seconds")

	output = result["stdout"].split("\n")
	count = 0
	for port_string in output:
	# Ignore blank lines
	if len(port_string) < 4:
	continue
	count += 1
	port_int = int(port_string)
	assert lo <= port_int <= hi, f"Port {port_int} is not in range {lo}-{hi}"

	# From the docs:
	# "A single srun opens 4 listening ports plus 2 more for every 48 hosts
	# beyond the first 48."
	ports = nodes * (4 + 2 * ((nodes - 1) // 48))
	assert count == ports, f"srun with -N{nodes} should use {ports} ports, not {count}"


	@pytest.mark.parametrize("nodes", [145])
	def test_srun_ports_out_of_range(nodes, safe_port_range):
	"""Test sruns with too many nodes, so with not enough SrunPortRange"""

	lo, hi = safe_port_range

	# Wait for ports to be open before running test
	wait_for_srun_ports_clear(lo, hi)

	result = atf.run_job_error(f"-t1 -N{nodes} sleep 1", fatal=True, xfail=True)

	regex = rf"all ports in range .{lo}, {hi}. exhausted"
	assert (
	re.search(regex, result) is not None
	), "srun's stderr should contain the 'all ports in range exhausted' message"


	def test_out_of_srun_ports(safe_port_range):
	"""Test exhausted ports"""

	lo, hi = safe_port_range

	# Wait for ports to be open before running test
	wait_for_srun_ports_clear(lo, hi)

	job_id1 = atf.submit_job_sbatch('-N1 -o/dev/null --wrap="srun sleep 30"')
	job_id2 = atf.submit_job_sbatch('-N1 -o/dev/null --wrap="srun sleep 30"')

	atf.wait_for_step(job_id1, 0, fatal=True)
	atf.wait_for_step(job_id2, 0, fatal=True)

	result = atf.run_job_error("-t1 -N1 sleep 1", fatal=True, xfail=True)

	regex = rf"all ports in range .{lo}, {hi}. exhausted"
	assert (
	re.search(regex, result) is not None
	), "srun's stderr should contain the 'all ports in range exhausted' message"