testsuite/python/tests/test_122_4.py

############################################################################
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
############################################################################
"""Verify srun --jobid can attach a step to a running job array task.

Covers ticket 25328: attaching with --overlap to a running array task
both when all tasks are running and when sibling tasks are still pending.
"""

import pytest

import atf

array_size = 4


@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_version(
        (26, 5),
        "bin/srun",
        reason="Ticket 25328: srun --jobid cannot attach to an array task before this fix",
    )

    atf.require_config_parameter("MaxArraySize", array_size, ">=")
    # test_het_step_on_array_task submits a het step
    atf.require_config_parameter_excludes("SchedulerParameters", "disable_hetjob_steps")
    atf.require_nodes(array_size)
    atf.require_slurm_running()


def test_srun_jobid_attach_to_array_task():
    """Test that srun --jobid=<array_job_id>_<array_task_id> can attach to a
    running task of a job array"""

    # Submit a job array where each task just sleeps, and wait for it to run
    job_id = atf.submit_job_sbatch(
        f"-N1 --array=0-{array_size - 1} -t5 --wrap='sleep infinity'",
        fatal=True,
    )
    for task_id in range(array_size):
        # A task only gets its own record once it is scheduled
        raw_job_id = 0
        for t in atf.timer(fatal=True):
            raw_job_id = atf.get_job_id_from_array_task(job_id, task_id)
            if raw_job_id != 0:
                break
        atf.wait_for_job_state(raw_job_id, "RUNNING", fatal=True)

        # Both the array (job_id_task_id) and the raw task job id are valid
        for jobid in (f"{job_id}_{task_id}", raw_job_id):
            results = atf.run_job(f"--jobid={jobid} --overlap printenv SLURM_JOB_ID")
            assert (
                results["exit_code"] == 0
            ), f"srun --jobid={jobid} failed to attach to the running array task"
            assert results["stdout"].strip() == str(raw_job_id), (
                f"srun --jobid={jobid} attached to the wrong job "
                f"(expected {raw_job_id}, got '{results['stdout'].strip()}')"
            )
            assert (
                "error" not in results["stderr"].lower()
            ), f"srun --jobid={jobid} emitted errors: {results['stderr']}"


def test_srun_jobid_attach_with_pending_sibling():
    """Test that srun --jobid=<array_job_id>_<array_task_id> can attach to a
    running array task while sibling tasks are still pending"""

    # %1 keeps only one task running at a time, so the rest stay pending
    job_id = atf.submit_job_sbatch(
        f"-N1 --array=0-{array_size - 1}%1 -t5 --wrap='sleep infinity'",
        fatal=True,
    )

    # The lowest-index task runs first; wait for its record to appear
    running_task = 0
    raw_job_id = 0
    for t in atf.timer(fatal=True):
        raw_job_id = atf.get_job_id_from_array_task(job_id, running_task)
        if raw_job_id != 0:
            break
    atf.wait_for_job_state(raw_job_id, "RUNNING", fatal=True)
    atf.wait_for_job_state(job_id, "PENDING", fatal=True)

    # Both the array (job_id_task_id) and raw task job id must resolve to the
    # running task even while its siblings are still pending
    for jobid in (f"{job_id}_{running_task}", raw_job_id):
        results = atf.run_job(f"--jobid={jobid} --overlap printenv SLURM_JOB_ID")
        assert (
            results["exit_code"] == 0
        ), f"srun --jobid={jobid} failed to attach with pending siblings"
        assert results["stdout"].strip() == str(raw_job_id), (
            f"srun --jobid={jobid} attached to the wrong job "
            f"(expected {raw_job_id}, got '{results['stdout'].strip()}')"
        )
        assert (
            "error" not in results["stderr"].lower()
        ), f"srun --jobid={jobid} emitted errors: {results['stderr']}"


def test_srun_jobid_attach_to_pending_task_fails():
    """Test that srun --jobid=<array_job_id>_<array_task_id> does not attach to
    a pending array task nor silently resolve to a running sibling"""

    # %1 keeps only the lowest-index task running; higher indexes stay pending
    job_id = atf.submit_job_sbatch(
        f"-N1 --array=0-{array_size - 1}%1 -t5 --wrap='sleep infinity'",
        fatal=True,
    )

    running_task = 0
    running_job_id = 0
    for t in atf.timer(fatal=True):
        running_job_id = atf.get_job_id_from_array_task(job_id, running_task)
        if running_job_id != 0:
            break
    atf.wait_for_job_state(running_job_id, "RUNNING", fatal=True)
    atf.wait_for_job_state(job_id, "PENDING", fatal=True)

    pending_task = array_size - 1
    results = atf.run_job(
        f"--jobid={job_id}_{pending_task} --overlap printenv SLURM_JOB_ID"
    )
    assert (
        results["exit_code"] != 0
    ), f"srun --jobid={job_id}_{pending_task} wrongly attached to a pending task"
    assert results["stdout"].strip() != str(running_job_id), (
        f"srun --jobid={job_id}_{pending_task} wrongly resolved to the running "
        f"task {running_job_id}"
    )


@pytest.mark.skipif(
    atf.get_version("sbin/slurmctld") < (26, 11),
    reason="Issue 50940: slurmctld cannot find an array task by its own job id for a het step before 26.11",
)
def test_het_step_on_array_task():
    """Test that a het step can attach to an array task that runs under its
    own job id"""

    # Every array task but the last to start runs under its own job id. Task 1
    # starts before task 2, so it gets its own.
    job_id = atf.submit_job_sbatch(
        "--array=0-2 -N2 -n2 -t5 --wrap 'sleep infinity'", fatal=True
    )

    task_job_id = 0
    for t in atf.timer(fatal=True):
        task_job_id = atf.get_job_id_from_array_task(job_id, 1)
        if task_job_id != 0:
            break
    atf.wait_for_job_state(task_job_id, "RUNNING", fatal=True)
    assert (
        task_job_id != job_id
    ), f"array task {job_id}_1 should run under its own job id"

    # A het step's components do not share a node, so each takes one of the
    # task's two nodes
    component = f"--jobid={job_id}_1 -N1 -n1 printenv SLURM_JOB_ID"
    result = atf.run_command(f"srun {component} : {component}", quiet=True)
    assert result["exit_code"] == 0, (
        f"the het step on array task {job_id}_1 should run, got exit "
        f"{result['exit_code']}: {result['stderr']!r}"
    )
    assert result["stdout"].split() == [str(task_job_id)] * 2, (
        f"both components should run in job {task_job_id}, got: "
        f"{result['stdout']!r}"
    )