testsuite/python/tests/test_105_10.py

############################################################################
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
############################################################################
from pathlib import Path

import pytest

import atf

oom_started = "Allocating"
oom_finished = "Done."
sleeper_survived = "SLEEPER SURVIVED OOM_KILL_STEP"
post_oom_success = "POST OOM STEP SUCCEEDED"

step_memory_mib = 32
job_memory_mib = 128
# Large enough to exceed step_memory_mib, but still fit inside job_memory_mib.
oom_allocation_mib = 92

# Ticket 22370: before 25.11, a TLS slurmstepd sent its Slurm API calls through
# slurmd, which crashed relaying the step info for the OOM-killed step.
tls_slurmd_crash = atf.get_config_parameter(
    "TLSType", default="tls/none", live=False, quiet=True
).casefold() != "tls/none" and atf.get_version("sbin/slurmd") < (25, 11)
ticket_22370 = "Ticket 22370: slurmd crashed relaying the step info for the OOM-killed step of a TLS slurmstepd, fixed in 25.11"


@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_config_parameter("SelectType", ["select/cons_tres", "select/linear"])
    atf.require_config_parameter_includes(
        "SelectTypeParameters", ["CR_Core_Memory", "CR_Memory"]
    )

    atf.require_config_parameter_includes("TaskPlugin", "cgroup")

    atf.require_config_parameter("ConstrainRAMSpace", "yes", source="cgroup")
    atf.require_config_parameter("AllowedRAMSpace", 100, source="cgroup")

    # Make OOM deterministic
    atf.require_config_parameter("ConstrainSwapSpace", "yes", source="cgroup")
    atf.require_config_parameter("AllowedSwapSpace", 0, source="cgroup")

    # Ensure step_memory_mib is not clipped
    atf.require_config_parameter("MinRAMSpace", 30, source="cgroup")
    atf.require_config_parameter("CgroupPlugin", "autodetect", source="cgroup")

    atf.require_nodes(3, [("CPUs", 3), ("RealMemory", job_memory_mib)])
    atf.require_slurm_running()


# The first task is a return 0 and we do the OOM on the one with ID 1 as we want
# to make sure that we check what happens on a regular task, as the first one,
# as well as the last one, thus the sleep 5 on the else, tends to have special
# treatment.
@pytest.fixture(scope="module")
def step_script(use_memory_program):
    path = Path("step.sh").absolute()
    atf.make_bash_script(
        path,
        f"""
if [ "$SLURM_PROCID" -eq 0 ]; then
    exit 0
elif [ "$SLURM_PROCID" -eq 1 ]; then
    sleep 1
    {use_memory_program} {oom_allocation_mib} 1
else
    sleep 5
    echo "{sleeper_survived} $SLURM_PROCID"
fi
""",
    )

    return str(path)


@pytest.mark.parametrize(
    "num_nodes",
    [
        pytest.param(
            1,
            id="single-node",
            marks=pytest.mark.xfail(
                atf.get_version("sbin/slurmd") < (26, 11),
                reason="Ticket 25503: OOM detection in a non-first-finishing task fixed in 26.11",
            ),
        ),
        pytest.param(3, id="multi-node"),
    ],
)
@pytest.mark.parametrize(
    "oom_kill_step",
    [
        pytest.param(0, id="oom-kill-step-disabled"),
        pytest.param(1, id="oom-kill-step-enabled"),
    ],
)
def test_oom_kill_step(request, num_nodes, oom_kill_step, step_script):
    """Test that sbatch's --oom-kill-step flag controls OOM step cleanup."""

    # Seen in the multi-node case; the single-node one is xfailed for Ticket 25503
    if tls_slurmd_crash and num_nodes > 1 and oom_kill_step:
        request.applymarker(pytest.mark.xfail(reason=ticket_22370))

    ntasks_per_node = 3 if num_nodes == 1 else 1
    output_file = f"oom_kill_{num_nodes}_{oom_kill_step}.out"
    batch_script = f"oom_kill_{num_nodes}_{oom_kill_step}.sh"
    atf.make_bash_script(
        batch_script,
        f"""
#SBATCH --output={output_file}
#SBATCH --nodes={num_nodes}
#SBATCH --ntasks=3
#SBATCH --ntasks-per-node={ntasks_per_node}
#SBATCH --mem={job_memory_mib}M

srun --mem={step_memory_mib}M --kill-on-bad-exit=0 {step_script}
srun --nodes=1 --ntasks=1 --ntasks-per-node=1 echo "{post_oom_success}"
""",
    )

    job_id = atf.submit_job_sbatch(
        f"--oom-kill-step={oom_kill_step} {batch_script}", fatal=True
    )
    atf.wait_for_job_state(job_id, "DONE", fatal=True)

    for t in atf.timer():
        output = atf.run_command_output(f"cat {output_file}")
        if post_oom_success in output:
            break
    else:
        assert False, f"Output file should contain the last line: {post_oom_success}"

    assert oom_started in output, "OOM script task should start"
    assert (
        "step tasks have been OOM Killed" in output
    ), "Step OOM killer should be triggered"
    if oom_kill_step:
        assert (
            sleeper_survived not in output
        ), "--oom-kill-step=1 should kill all step tasks after the OOM event"
    else:
        assert (
            sleeper_survived in output
        ), "--oom-kill-step=0 should NOT kill healthy step tasks after the OOM event"
    assert (
        oom_finished not in output
    ), "OOM script finished allocation instead of being killed"