testsuite/python/tests/test_128_1.py

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

import pytest

import atf

pytestmark = pytest.mark.slow


@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_config_parameter("PreemptType", "preempt/partition_prio")
    atf.require_config_parameter("PreemptMode", "CANCEL,GANG")

    # We need to define nodes memory and DefMem to support SelectTypeParameters with Memory
    atf.require_nodes(1, [("CPUs", 2), ("RealMemory", 512)])
    atf.require_config_parameter("DefMemPerNode", "128")

    atf.require_slurm_running()


@pytest.fixture(scope="module")
def partition_node():
    """Obtain a node that we can use in the partitions"""
    return atf.run_job_nodes("-N1 -t1 --exclusive true", fatal=True)[0]


@pytest.fixture(scope="module")
def partition1(partition_node):
    """Create partition 1"""
    partition_name = "partition1"
    atf.run_command(
        f"scontrol create partitionname={partition_name} nodes={partition_node} priority=1 preemptmode=cancel",
        user=atf.properties["slurm-user"],
        fatal=True,
    )
    return partition_name


@pytest.fixture(scope="module")
def partition2(partition_node):
    """Create partition 2"""
    partition_name = "partition2"
    atf.run_command(
        f"scontrol create partitionname={partition_name} nodes={partition_node} priority=2 preemptmode=off",
        user=atf.properties["slurm-user"],
        fatal=True,
    )
    return partition_name


def test_preempt_cancel(partition1, partition2):
    """Test preempt cancel"""

    job_id1 = atf.submit_job_sbatch(
        f'-N1 -t1 -o /dev/null --exclusive -p {partition1} --wrap "sleep 120"',
        fatal=True,
    )
    assert atf.wait_for_job_state(
        job_id1, "RUNNING"
    ), f"Job 1 ({job_id1}) did not start"
    job_id2 = atf.submit_job_sbatch(
        f'-N1 -t1 -o /dev/null --exclusive -p {partition2} --wrap "sleep 30"',
        fatal=True,
    )
    assert atf.wait_for_job_state(
        job_id2, "RUNNING"
    ), f"Job 2 ({job_id2}) did not start"
    assert atf.wait_for_job_state(
        job_id1, "PREEMPTED"
    ), f"Job 1 ({job_id1}) did not get preempted"


def test_preempt_suspend(partition1, partition2):
    """Test preempt suspend"""

    atf.run_command(
        f"scontrol update partitionname={partition1} preemptmode=suspend",
        user=atf.properties["slurm-user"],
        fatal=True,
    )

    job_id1 = atf.submit_job_sbatch(
        f'-N1 -t1 -o /dev/null --exclusive -p {partition1} --wrap "sleep 120"',
        fatal=True,
    )
    assert atf.wait_for_job_state(
        job_id1, "RUNNING"
    ), f"Job 1 ({job_id1}) did not start"
    job_id2 = atf.submit_job_sbatch(
        f'-N1 -t1 -o /dev/null --exclusive -p {partition2} --wrap "sleep 30"',
        fatal=True,
    )
    assert atf.wait_for_job_state(job_id2, "RUNNING"), f"Job 2 {job_id2} did not start"
    assert atf.wait_for_job_state(
        job_id1, "SUSPENDED"
    ), f"Job 1 ({job_id1}) did not get suspended"
    assert atf.wait_for_job_state(
        job_id2, "DONE", timeout=60, poll_interval=1
    ), f"Job 2 ({job_id2}) did not complete"
    assert atf.wait_for_job_state(
        job_id1, "RUNNING"
    ), f"Job 1 ({job_id1}) did not start running again"


@pytest.mark.parametrize(
    "job_options,batch_flag",
    [
        ('--wrap "sleep infinity"', 1),
        pytest.param(
            "--external",
            0,
            marks=pytest.mark.skipif(
                atf.get_version("sbin/slurmctld") < (26, 11),
                reason="Issue 166 (slurm-bridge): external job requeue added in 26.11",
            ),
        ),
    ],
    ids=["batch", "external"],
)
def test_preempt_requeue(partition1, partition2, job_options, batch_flag):
    """Test preempt requeue for batch and external jobs."""

    atf.run_command(
        f"scontrol update partitionname={partition1} preemptmode=requeue",
        user=atf.properties["slurm-user"],
        fatal=True,
    )

    job_id1 = atf.submit_job_sbatch(
        f"-N1 -t1 -o /dev/null --exclusive --requeue -p {partition1} {job_options}",
        fatal=True,
    )
    assert atf.wait_for_job_state(
        job_id1, "RUNNING"
    ), f"Job  1 ({job_id1}) did not start"
    job_id2 = atf.submit_job_sbatch(
        f'-N1 -t1 -o /dev/null --exclusive -p {partition2} --wrap "sleep 30"',
        fatal=True,
    )
    assert atf.wait_for_job_state(
        job_id2, "RUNNING"
    ), f"Job 2 ({job_id2}) did not start"
    assert atf.wait_for_job_state(
        job_id1, "PENDING"
    ), f"Job 1 ({job_id1}) did not return to pending"
    job = atf.get_job(job_id1)
    assert job["Restarts"] == 1, job
    assert job["BatchFlag"] == batch_flag, job
    assert job.get("NodeList") in (None, ""), job
    assert atf.wait_for_job_state(
        job_id2, "DONE", timeout=60, poll_interval=1
    ), f"Job 2 ({job_id2}) did not complete"
    assert atf.wait_for_job_state(
        job_id1, "RUNNING", timeout=150, poll_interval=1
    ), f"Job 1 ({job_id1}) did not start running again"
    time.sleep(5)
    assert (
        atf.get_job_parameter(job_id1, "JobState") == "RUNNING"
    ), f"Job 1 ({job_id1}) was not still running after 5 seconds"