testsuite/python/tests/test_162_4.py

############################################################################
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
############################################################################
"""Verify pending synchronous step id-at-submit holds under stepmgr.

Issue 50938: a synchronous step that has to queue is assigned its StepId at
submission, is visible in PENDING under that id, launches under it, and can
be cancelled by it. test_116_61.py covers that against the controller's own
step manager; under SlurmctldParameters=enable_stepmgr the step create is
served by the batch host's slurmstepd instead, so the same promises are
pinned here against that path.

StepId ordering itself (ids tracking submission order, with no gaps) is a
property of the shared _set_step_id()/_build_pending_step() assignment and is
pinned once, by test_116_61. Its scenario needs two nodes; this module needs
only one, so it is not duplicated here. test_162_11 covers the stepmgr's FIFO
*launch* ordering, which is the part that is genuinely stepmgr specific.
"""

from pathlib import Path

import pytest

import atf


@pytest.fixture(scope="module", autouse=True)
def setup():
    # Daemons are always >= every client command, so requiring srun >= 26.11
    # already implies the stepmgr serving the step create is new enough.
    atf.require_version(
        (26, 11),
        "bin/srun",
        reason="Issue 50938: pending-step id-at-submit requires a 26.11+ srun",
    )
    # stepmgr is enabled when both SlurmctldParameters=enable_stepmgr and
    # PrologFlags=Contain are set.
    atf.require_config_parameter_includes("SlurmctldParameters", "enable_stepmgr")
    atf.require_config_parameter_includes("PrologFlags", "Contain")
    # The hog and the queued step are both --exclusive, so they contend
    # whatever allocation unit the site's SelectTypeParameters names.
    atf.require_nodes(1, [("CPUs", 2)])
    atf.require_slurm_running()


def test_pending_sync_step_shows_real_id_and_reuses_it_under_stepmgr():
    """A synchronous step queued behind a CPU-hogging step shows a real StepId
    while PENDING and launches with that same id once the hog is freed."""

    ready = Path("stepmgr_hog_ready")
    script = Path("stepmgr_hog_and_pending.sh")
    atf.make_bash_script(
        script,
        atf.hog_allocation_snippet(ready)
        + "srun --exclusive -n2 sleep infinity &\n"
        + "wait\n",
    )
    job_id = atf.submit_job(
        "sbatch", "-N1 -n2 -t5", str(script), wrap_job=False, fatal=True
    )
    assert job_id != 0, "sbatch should submit the job"

    hog_id, pending_id = f"{job_id}.0", f"{job_id}.1"
    atf.wait_for_step(job_id, 0, fatal=True)
    atf.wait_for_step(job_id, 1, state="PENDING", fatal=True)

    step_show = atf.run_command_output(
        f"scontrol -o show step {pending_id}", quiet=True, fatal=True
    )
    assert (
        f"StepId={pending_id}" in step_show
    ), f"the queued step should be addressable as {pending_id}, got: {step_show!r}"
    assert (
        "State=PENDING" in step_show
    ), f"the queued step should show State=PENDING, got: {step_show!r}"
    assert "TBD" not in atf.run_command_output(
        f"scontrol -o show step {job_id}", quiet=True, fatal=True
    ), "no step should render as TBD once a real id is assigned at submit"

    # Free the resources; the pending step launches under the SAME id.
    atf.run_command_exit(f"scancel {hog_id}", quiet=True, fatal=True)
    atf.wait_for_step(job_id, 1, fatal=True)


def test_queued_step_reports_queued_and_started_under_stepmgr():
    """srun prints both the queued and the started notice on this path.

    Under stepmgr the step manager creates the woken step and pushes it, so
    srun gets its step back inside its first create attempt instead of
    re-sending. The launch notice must therefore not be keyed off the retry
    count, or a stepmgr user sees "queued" followed by silence while the
    controller path prints the pair (srun.1).
    """

    ready = Path("stepmgr_notice_ready")
    pend_err = Path("stepmgr_notice_pend.err")
    out = Path("stepmgr_notice.out")
    script = Path("stepmgr_notice.sh")
    atf.make_bash_script(
        script,
        atf.hog_allocation_snippet(ready)
        + f"srun --exclusive -n2 true 2>'{pend_err}' &\n"
        + "pend=$!\n"
        + atf.wait_for_step_snippet(1)
        + "scancel $SLURM_JOB_ID.0\n"
        + "wait $pend\n"
        + "echo PEND_RC=$?\n",
    )
    job_id = atf.submit_job(
        "sbatch",
        f"-N1 -n2 -t5 --output={out} --error={out}",
        str(script),
        wrap_job=False,
        fatal=True,
    )
    assert job_id != 0, "sbatch should submit the job"
    atf.wait_for_job_state(job_id, "DONE", fatal=True)

    atf.assert_file_contents(pend_err, f"StepId={job_id}.1 queued", contains=True)
    atf.assert_file_contents(pend_err, f"StepId={job_id}.1 started", contains=True)


def test_scancel_pending_sync_step_by_id_under_stepmgr():
    """scancel <jobid>.<stepid> reaps a queued synchronous step before it
    launches and its waiting srun aborts non-zero, leaving the hog running."""

    out = Path("stepmgr_cancel.out")
    ready = Path("stepmgr_cancel_ready")
    script = Path("stepmgr_cancel.sh")
    # The pending srun runs in the foreground so its exit status reaches the
    # job output. The trailing wait blocks on the backgrounded hog, which
    # never ends on its own, so the job outlives the cancelled step and the
    # hog-survival check below cannot race the script exiting.
    atf.make_bash_script(
        script,
        atf.hog_allocation_snippet(ready)
        + "srun --exclusive -n2 sleep infinity\n"
        + 'echo "PENDING_SRUN_RC=$?"\n'
        + "wait\n",
    )
    job_id = atf.submit_job(
        "sbatch",
        f"-N1 -n2 -t2 --output={out} --error={out}",
        str(script),
        wrap_job=False,
        fatal=True,
    )
    assert job_id != 0, "sbatch should submit the job"

    hog_id, pending_id = f"{job_id}.0", f"{job_id}.1"
    # The hog must be running before the queued step can pend on it, and
    # before the post-cancel check below can claim it survived.
    atf.wait_for_step(job_id, 0, fatal=True)
    atf.wait_for_step(job_id, 1, state="PENDING", fatal=True)

    assert (
        atf.run_command_exit(f"scancel {pending_id}", quiet=True) == 0
    ), f"scancel {pending_id} should succeed against a pending step"

    # The placeholder is reaped outright, not merely moved out of PENDING.
    for _ in atf.timer(fatal=True):
        if f"StepId={pending_id}" not in atf.run_command_output(
            f"scontrol -o show step {job_id}", quiet=True
        ):
            break
    assert f"StepId={hog_id}" in atf.run_command_output(
        f"scontrol -o show step {job_id}", quiet=True, fatal=True
    ), f"the hogging step {hog_id} should still be running after the step cancel"

    # The job is still running (held by the hog), so poll for the aborted
    # srun's line rather than waiting for the job to finish.
    atf.assert_file_contents(out, "PENDING_SRUN_RC=", contains=True)
    text = atf.run_command_output(f"cat {out}", quiet=True, fatal=True)
    assert (
        "PENDING_SRUN_RC=0" not in text
    ), f"the cancelled pending srun should exit non-zero, got: {text}"
    assert (
        "Pending job step cancelled" in text
    ), f"srun should log the pending-step cancellation, got: {text}"