############################################################################
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
############################################################################
"""Verify srun --bcast with hetjobs."""

import os
import re

import pytest

import atf

# srun.1 documents this name for a broadcast file whose destination is a
# directory rather than a file: slurm_bcast_<job_id>.<step_id>_<nodename>.
BCAST_NAME_PATTERN = r"slurm_bcast_\d+\.\d+_\S+"

# Hetjobs only start on a backfill cycle, so allow more than the default
# command timeout. A timeout otherwise reports exit code 110 and no stderr.
SRUN_TIMEOUT = 120

# What detects the regression these tests cover is the srun exit code, not any
# field parsed out of a broadcast name. slurmd names the file it writes from
# the sbcast credential the controller issued, while slurmstepd rebuilds the
# same name from the step it is about to run, so a credential carrying the
# wrong step id sends the file somewhere the task never looks, and a
# credential carrying the wrong het component sends it to the other
# component's node. Either way the exec fails with ENOENT.
#
# The controller side of that was broken in 25.11 and is fixed in 26.11 and in
# 26.05.5, so every test below that depends on the credential is expected to
# fail against a controller older than that. tox.ini sets xfail_strict, so a
# controller in this range that does have the fix is a hard failure rather than
# an XPASS signal: the range has to shrink in step with every backport.
xfail_bcast_hetjob = pytest.mark.xfail(
    (25, 11) <= atf.get_version("sbin/slurmctld") < (26, 5, 5),
    reason="Issue 51136: srun --bcast in a hetjob is broken before 26.05.5",
)

# --bcast-exclude has a separate gate. Its defect is in srun rather than the
# controller, so an upgraded controller in front of an old srun still shows it,
# and it predates 25.11 rather than being a regression in it, so there is no
# lower bound. srun dropped a non-final component's list, leaving that
# component to take its successor's value.
xfail_bcast_exclude_hetjob = pytest.mark.xfail(
    atf.get_version("bin/srun") < (26, 5, 5),
    reason="Issue 51136: srun --bcast-exclude is lost for every hetjob"
    " component but the last before 26.05.5",
)


@pytest.fixture(scope="module", autouse=True)
def setup():
    # Two nodes because each component asks for -N1. Nothing here forces the
    # components onto different nodes: require_nodes() treats CPUs as a
    # minimum rather than an exact count, and heterogeneous_jobs.shtml says
    # components may end up sharing a node. What the tests rely on is that
    # backfill normally has to place each component separately to start the
    # allocation at all.
    atf.require_nodes(2, [("CPUs", 1)])
    # Hetjobs are only ever started by the backfill scheduler.
    atf.require_config_parameter_includes("SchedulerParameters", ("bf_interval", 1))
    # With no dest_path srun.1 prefers BcastParameters DestDir over --chdir.
    # The tests below assert the --chdir form, so require the whole parameter
    # unset rather than report a site setting as if --bcast were broken.
    # SbcastParameters is still accepted as a deprecated alias for the same
    # value (read_config.c only logs a rename error), so it has to be unset
    # too or a site using the old spelling slips a DestDir past this.
    atf.require_config_parameter("BcastParameters", None)
    atf.require_config_parameter("SbcastParameters", None)
    # --bcast-exclude overrides slurm.conf BcastExclude, so the exclude test
    # can only tell an honoured override from an ignored one if BcastExclude is
    # known. Unset is the documented default, which is non-empty.
    atf.require_config_parameter("BcastExclude", None)
    atf.require_slurm_running()


@pytest.fixture(scope="module")
def component_scripts():
    """Return one distinct executable per hetjob component.

    Each script prints its own marker, its own argv[0] and the node it ran
    on. srun rewrites argv[0] to the broadcast destination, so the output
    says which file the component ran, where that file landed, and which
    node ran it.
    """

    scripts = []
    for component in range(2):
        name = f"comp{component}.sh"
        atf.make_bash_script(name, f'echo "COMP{component} $0 $SLURMD_NODENAME"')
        scripts.append(f"{os.getcwd()}/{name}")
    return scripts


def executed_files(output):
    """Map each component marker to its (argv[0], nodename)."""

    executed = {}
    for line in output.splitlines():
        match = re.fullmatch(r"(COMP\d) (\S+) (\S+)", line.strip())
        if match:
            executed[match.group(1)] = (match.group(2), match.group(3))
    return executed


def assert_documented_bcast_name(marker, argv0, directory):
    """Assert argv[0] is the name srun.1 documents, under directory.

    srun.1 says a broadcast file whose destination is a directory is named
    slurm_bcast_<job_id>.<step_id>_<nodename>. The directory half carries the
    weight here: it comes from srun's --chdir, so matching it is what proves
    the documented no-dest_path resolution was used.

    The fields inside the name are deliberately not asserted. slurmstepd
    rebuilds them from the step it is running, so they look correct whatever
    the sbcast credential said; the credential only shows up in whether the
    exec succeeded.
    """

    assert re.fullmatch(f"{re.escape(directory)}/{BCAST_NAME_PATTERN}", argv0), (
        f"{marker} should have run the documented broadcast name under"
        f" {directory}, but ran {argv0}"
    )


@xfail_bcast_hetjob
def test_bcast_hetjob_default_destination(component_scripts):
    """Verify --bcast with no path runs each component's own broadcast file.

    With no path and no BcastParameters DestDir, srun.1 puts the file in
    --chdir under the documented name. slurmd builds that name from the
    sbcast credential while slurmstepd rebuilds it from the step it runs,
    so the two have to agree or the task dies with ENOENT.

    The exit code assertion is therefore the regression detector, not the
    name checks that follow it. Do not weaken it.

    Note: assumes the same file system is visible to us and to the slurmd.
    """

    result = atf.run_command(
        f"srun -t1 --bcast -N1 -n1 {component_scripts[0]}"
        f" : --bcast -N1 -n1 {component_scripts[1]}",
        timeout=SRUN_TIMEOUT,
    )

    assert result["exit_code"] == 0, (
        f"srun --bcast in a hetjob should succeed, but got rc="
        f"{result['exit_code']}: {result['stderr']}"
    )

    executed = executed_files(result["stdout"])
    assert set(executed) == {"COMP0", "COMP1"}, (
        "Each hetjob component should have run its own broadcast file, but got"
        f" {result['stdout']!r}"
    )

    cwd = os.getcwd()
    for marker, (argv0, _) in executed.items():
        assert_documented_bcast_name(marker, argv0, cwd)


@pytest.fixture(scope="module")
def shared_script():
    """Return one executable for both hetjob components to share."""

    atf.make_bash_script("shared.sh", 'echo "SHARED $0 $SLURMD_NODENAME"')
    return f"{os.getcwd()}/shared.sh"


@xfail_bcast_hetjob
def test_bcast_hetjob_shared_application(shared_script):
    """Verify --bcast with one application shared by both components.

    heterogeneous_jobs.shtml says a component that lacks an application
    specification uses the next one provided, so this single command line
    broadcasts one source file on behalf of two components. That is the
    hardest shape to get right and the easiest to get wrong quietly: there is
    only one source path to reason about, so a credential naming the wrong
    component has nothing else to disagree with.

    As above, the exit code is the detector.

    Note: assumes the same file system is visible to us and to the slurmd.
    """

    result = atf.run_command(
        f"srun -t1 -n1 --bcast : -n1 --bcast {shared_script}",
        timeout=SRUN_TIMEOUT,
    )

    assert result["exit_code"] == 0, (
        "srun --bcast with an application shared by both hetjob components"
        f" should succeed, but got rc={result['exit_code']}: {result['stderr']}"
    )

    ran = [
        match.groups()
        for match in (
            re.fullmatch(r"SHARED (\S+) (\S+)", line.strip())
            for line in result["stdout"].splitlines()
        )
        if match
    ]
    assert len(ran) == 2, (
        "Both hetjob components should have run the shared application, but"
        f" got {result['stdout']!r}"
    )

    cwd = os.getcwd()
    for argv0, node in ran:
        assert_documented_bcast_name(f"The component on {node}", argv0, cwd)

    assert len({argv0 for argv0, _ in ran}) == 2, (
        "Each component should have broadcast the shared application to its"
        f" own destination, but both ran {ran[0][0]}"
    )


def test_bcast_hetjob_destination_file(component_scripts):
    """Verify --bcast=<absolute file> is used verbatim for each component.

    srun.1 says a path that does not end in '/' is the destination file name.
    No name is derived from the job or step in that case, so this covers the
    other half of the documented path resolution from the pattern tests.

    Note: assumes the same file system is visible to us and to the slurmd.
    """

    cwd = os.getcwd()
    destinations = [f"{cwd}/bcast_comp0", f"{cwd}/bcast_comp1"]

    result = atf.run_command(
        f"srun -t1 --bcast={destinations[0]} -N1 -n1 {component_scripts[0]}"
        f" : --bcast={destinations[1]} -N1 -n1 {component_scripts[1]}",
        timeout=SRUN_TIMEOUT,
    )

    assert result["exit_code"] == 0, (
        f"srun --bcast=<file> in a hetjob should succeed, but got rc="
        f"{result['exit_code']}: {result['stderr']}"
    )

    executed = executed_files(result["stdout"])
    assert {marker: argv0 for marker, (argv0, _) in executed.items()} == {
        "COMP0": destinations[0],
        "COMP1": destinations[1],
    }, f"Each component should have run its own --bcast file, but got {executed}"


@xfail_bcast_hetjob
def test_bcast_hetjob_destination_dir(component_scripts):
    """Verify --bcast=<dir>/ names the file per the documented pattern.

    srun.1 says a path ending in '/' is a target directory, and that the
    destination file in it is named slurm_bcast_<job_id>.<step_id>_<nodename>.

    Note: assumes the same file system is visible to us and to the slurmd.
    """

    cwd = os.getcwd()
    directories = [f"{cwd}/bcast_dir0", f"{cwd}/bcast_dir1"]
    for directory in directories:
        atf.run_command(f"mkdir -p {directory}", fatal=True)

    result = atf.run_command(
        f"srun -t1 --bcast={directories[0]}/ -N1 -n1 {component_scripts[0]}"
        f" : --bcast={directories[1]}/ -N1 -n1 {component_scripts[1]}",
        timeout=SRUN_TIMEOUT,
    )

    assert result["exit_code"] == 0, (
        f"srun --bcast=<dir>/ in a hetjob should succeed, but got rc="
        f"{result['exit_code']}: {result['stderr']}"
    )

    executed = executed_files(result["stdout"])
    assert set(executed) == {"COMP0", "COMP1"}, (
        "Each hetjob component should have run its own broadcast file, but got"
        f" {result['stdout']!r}"
    )
    for component, marker in enumerate(["COMP0", "COMP1"]):
        argv0, _ = executed[marker]
        assert_documented_bcast_name(marker, argv0, directories[component])


@xfail_bcast_hetjob
def test_bcast_hetjob_chdir_per_component(component_scripts):
    """Verify --bcast with no path resolves each component's own --chdir.

    srun.1 falls back to --chdir when no dest_path is given and no
    BcastParameters DestDir is configured. Each component parses its own
    --chdir, so each broadcast file has to land under that component's
    directory rather than under one the components share.

    The scripts are named by absolute path so --chdir moves the destination
    without also moving the source.

    Note: assumes the same file system is visible to us and to the slurmd.
    """

    cwd = os.getcwd()
    directories = [f"{cwd}/chdir_comp0", f"{cwd}/chdir_comp1"]
    for directory in directories:
        atf.run_command(f"mkdir -p {directory}", fatal=True)

    result = atf.run_command(
        f"srun -t1 --chdir={directories[0]} --bcast -N1 -n1 {component_scripts[0]}"
        f" : --chdir={directories[1]} --bcast -N1 -n1 {component_scripts[1]}",
        timeout=SRUN_TIMEOUT,
    )

    assert result["exit_code"] == 0, (
        f"srun --bcast with a per-component --chdir should succeed, but got rc="
        f"{result['exit_code']}: {result['stderr']}"
    )

    executed = executed_files(result["stdout"])
    assert set(executed) == {"COMP0", "COMP1"}, (
        "Each hetjob component should have run its own broadcast file, but got"
        f" {result['stdout']!r}"
    )
    for component, marker in enumerate(["COMP0", "COMP1"]):
        argv0, _ = executed[marker]
        assert_documented_bcast_name(marker, argv0, directories[component])


@pytest.fixture(scope="module")
def library_probe():
    """Return a dynamically linked executable, its library dirs and a probe.

    --send-libs only has something to send for a dynamically linked ELF, so a
    bash script cannot be the broadcast executable. Broadcasting bash itself
    and handing it a probe script makes the task its own oracle: it reports
    the LD_LIBRARY_PATH it was launched with and the contents of the cache
    directory srun.1 promises to add to it, so nothing here has to know how
    slurmd spells that directory.
    """

    atf.require_tool("bash")
    # srun runs LDD_PATH unconditionally rather than searching PATH
    # (src/common/file_bcast.c), so gate on the path srun uses.
    atf.require_tool("/usr/bin/ldd")

    binary = atf.run_command_output("command -v bash", quiet=True, fatal=True).strip()
    # Not fatal: ldd exits non-zero on a non-dynamic binary, which the skip
    # below is there to handle. srun treats the same failure as non-fatal.
    output = atf.run_command_output(f"/usr/bin/ldd {binary}", quiet=True)

    # Take every absolute path ldd prints: the first '/' in the line up to the
    # next space. Matching on "=>" alone would miss the dynamic linker line,
    # which has no "=>" in it, and leave that directory unexcluded.
    directories = set()
    for line in output.splitlines():
        start = line.find("/")
        if start < 0:
            continue
        directories.add(os.path.dirname(line[start:].split(" ")[0]))
    if not directories:
        # A bare skip rather than an atf.require_*() or skipif: whether the
        # shell is statically linked is only knowable from ldd's output, and
        # with nothing to send there is no exclusion to observe.
        pytest.skip(f"{binary} is not dynamically linked, so nothing is sent")

    # $1 is the component's marker, $2 the directory it broadcast into.
    # srun.1 promises the libraries are placed in a directory alongside the
    # executable and that LD_LIBRARY_PATH is updated to include that cache
    # directory, so report every LD_LIBRARY_PATH entry that is in the
    # component's own directory along with what is in it. The DONE line makes
    # "the task ran and was given no cache directory" distinguishable from
    # "the task never ran".
    atf.make_bash_script(
        "libs_probe.sh",
        """
IFS=:
for entry in $LD_LIBRARY_PATH; do
    case "$entry" in
        "$2"/*) echo "$1 CACHE $entry $(ls -A "$entry" | tr '\\n' ',')" ;;
    esac
done
echo "$1 DONE"
""",
    )

    return binary, sorted(directories), f"{os.getcwd()}/libs_probe.sh"


def probe_caches(output):
    """Map each marker to the (directory, contents) pairs the probe reported.

    A marker present with an empty list means the task ran and was given no
    library cache directory at all.
    """

    caches = {}
    for line in output.splitlines():
        line = line.strip()
        done = re.fullmatch(r"(COMP\d) DONE", line)
        if done:
            caches.setdefault(done.group(1), [])
            continue
        cache = re.fullmatch(r"(COMP\d) CACHE (\S+) ?(.*)", line)
        if cache:
            contents = [name for name in cache.group(3).split(",") if name]
            caches.setdefault(cache.group(1), []).append((cache.group(2), contents))
    return caches


@xfail_bcast_exclude_hetjob
def test_bcast_exclude_applies_per_hetjob_component(library_probe):
    """Verify each hetjob component honours its own --bcast-exclude list.

    srun.1 says --send-libs places the executable's shared objects in a
    directory alongside it and updates LD_LIBRARY_PATH to include that cache
    directory, and that --bcast-exclude (defaulting to the slurm.conf
    BcastExclude value) keeps listed paths out of it. Each component parses
    its own command line, so each must keep its own list rather than share
    one.

    Both components ask for --send-libs, so both must be given a cache
    directory on LD_LIBRARY_PATH; the only difference between them should be
    what is in it.

    The excluding component is deliberately the first one. srun copies each
    component's options while parsing the next, so only a non-final component
    can lose its own value.

    The dangling read left by srun's per-component copy defect does not
    return the original bytes here. Component 0 picks up component 1's NONE
    and is broadcast the shared objects it excluded, on every unpatched srun
    the mixed-version matrix runs. This test is the regression detector for
    that, not just a statement of the functional contract. Do not weaken it.

    Note: assumes the same file system is visible to us and to the slurmd.
    """

    binary, library_dirs, probe = library_probe
    cwd = os.getcwd()
    directories = [f"{cwd}/libs_comp0", f"{cwd}/libs_comp1"]
    for directory in directories:
        atf.run_command(f"mkdir -p {directory}", fatal=True)
    excluded = ",".join(library_dirs)

    result = atf.run_command(
        f"srun -t1 --send-libs --bcast={directories[0]}/probe"
        f" --bcast-exclude={excluded} -N1 -n1"
        f" {binary} {probe} COMP0 {directories[0]}"
        f" : --send-libs --bcast={directories[1]}/probe"
        f" --bcast-exclude=NONE -N1 -n1"
        f" {binary} {probe} COMP1 {directories[1]}",
        timeout=SRUN_TIMEOUT,
    )

    assert result["exit_code"] == 0, (
        f"srun --bcast --send-libs in a hetjob should succeed, but got rc="
        f"{result['exit_code']}: {result['stderr']}"
    )

    caches = probe_caches(result["stdout"])
    assert set(caches) == {"COMP0", "COMP1"}, (
        "Each hetjob component should have run the probe, but got"
        f" {result['stdout']!r}"
    )

    for marker, directory in zip(["COMP0", "COMP1"], directories):
        assert len(caches[marker]) == 1, (
            f"--send-libs should have given {marker} exactly one library cache"
            f" directory under {directory} on its LD_LIBRARY_PATH, but it"
            f" reported {caches[marker]}"
        )

    cache_dir, included = caches["COMP1"][0]
    assert included, (
        "--bcast-exclude=NONE should have broadcast the shared objects of"
        f" component 1, but {cache_dir} is empty"
    )

    cache_dir, withheld = caches["COMP0"][0]
    assert not withheld, (
        f"Component 0 excluded {excluded}, which is every directory ldd"
        f" reported for {binary}, so nothing should have been broadcast, but"
        f" {cache_dir} holds {sorted(withheld)}"
    )


@xfail_bcast_hetjob
def test_bcast_het_group_in_existing_allocation():
    """Verify --het-group picks the component --bcast transfers to.

    heterogeneous_jobs.shtml describes --bcast as transferring files to the
    nodes of the application to be launched, as selected by --het-group, and
    srun.1 says --het-group applies only inside a salloc allocation or an
    sbatch script. That is a different path than srun creating the hetjob.

    This uses the default destination on purpose. An explicit path carries no
    node identity, so a transfer credentialed for the wrong component still
    execs from the shared file system and the test cannot tell. The
    documented name embeds the node, so the wrong component gives ENOENT and
    the srun in the batch script fails.

    Note: assumes the same file system is visible to us and to the slurmd.
    """

    cwd = os.getcwd()
    output = "het_group_bcast.out"
    error = "het_group_bcast.err"
    batch_script = "het_group_bcast.sh"

    for component in range(2):
        atf.make_bash_script(
            f"group{component}.sh", f'echo "COMP{component} $0 $SLURMD_NODENAME"'
        )

    atf.make_bash_script(
        batch_script,
        f"""
#SBATCH -N1 -t1
#SBATCH hetjob
#SBATCH -N1 -t1

set -e

srun --het-group=0 --bcast -n1 {cwd}/group0.sh
srun --het-group=1 --bcast -n1 {cwd}/group1.sh
""",
    )

    job_id = atf.submit_job_sbatch(
        f"-t1 -o {output} -e {error} {batch_script}", fatal=True
    )
    # Same backfill-cycle wait as SRUN_TIMEOUT above: the 45s default polling
    # timeout is not enough for a hetjob to be scheduled and complete.
    atf.wait_for_job_state(job_id, "DONE", timeout=SRUN_TIMEOUT, fatal=True)

    # Check the job state before polling the output. A failed srun leaves the
    # output empty for good, so waiting on it first only burns the timeout and
    # hides the real error, which srun wrote to the error file.
    state = atf.get_job_parameter(job_id, "JobState", fatal=True)
    assert state == "COMPLETED", (
        f"Expected COMPLETED but got {state!r}:"
        f" {atf.run_command_output(f'cat {error}', quiet=True)}"
    )

    # Waits for the file and re-polls until the second component's line lands.
    # The wait includes the start of the name, not just the marker, so a read
    # that catches the line half written cannot be mistaken for a component
    # that never ran.
    atf.assert_file_contents(output, f"COMP1 {cwd}/slurm_bcast_", contains=True)
    contents = atf.run_command_output(f"cat {output}", fatal=True)

    executed = executed_files(contents)
    assert set(executed) == {"COMP0", "COMP1"}, (
        "Each --het-group should have run its own broadcast file, but got"
        f" {contents!r}"
    )

    for marker, (argv0, _) in executed.items():
        assert_documented_bcast_name(marker, argv0, cwd)

    # Not a regression check by itself, and nothing above forces the split:
    # hetjob components are observed to be placed on separate nodes. That
    # separation is what makes the ENOENT above possible, so if it ever stops
    # holding, the test above has quietly stopped testing anything.
    nodes = {node for _, node in executed.values()}
    assert len(nodes) == 2, (
        "The two --het-group components should have run on different nodes,"
        f" each with a broadcast file named for its own node, but got {nodes}"
    )


@xfail_bcast_hetjob
def test_bcast_het_group_spanning_components():
    """Verify --het-group=0,1 broadcasts to every component it names.

    srun.1 gives --het-group a set expression, so one step can span several
    components. That is the only shape where a single sbcast credential has to
    resolve to more than one component's nodes, which is what the het step
    component field selects; with a single component there is only one answer
    it could give.

    The default destination is used for the same reason as the test above: the
    documented name embeds the node, so a component the file never reached
    gives ENOENT rather than silently execing from the shared file system.

    Note: assumes the same file system is visible to us and to the slurmd.
    """

    cwd = os.getcwd()
    output = "het_group_span.out"
    error = "het_group_span.err"
    batch_script = "het_group_span.sh"

    atf.make_bash_script("span.sh", 'echo "SPAN $0 $SLURMD_NODENAME"')
    atf.make_bash_script(
        batch_script,
        f"""
#SBATCH -N1 -t1
#SBATCH hetjob
#SBATCH -N1 -t1

srun --het-group=0,1 --bcast {cwd}/span.sh
""",
    )

    job_id = atf.submit_job_sbatch(
        f"-t1 -o {output} -e {error} {batch_script}", fatal=True
    )
    atf.wait_for_job_state(job_id, "DONE", timeout=SRUN_TIMEOUT, fatal=True)

    state = atf.get_job_parameter(job_id, "JobState", fatal=True)
    assert state == "COMPLETED", (
        f"Expected COMPLETED but got {state!r}:"
        f" {atf.run_command_output(f'cat {error}', quiet=True)}"
    )

    contents = ""
    for _ in atf.timer(fatal=True):
        contents = atf.run_command_output(f"cat {output}", quiet=True, fatal=True)
        if contents.count(f"SPAN {cwd}/slurm_bcast_") == 2:
            break

    ran = [
        match.groups()
        for match in (
            re.fullmatch(r"SPAN (\S+) (\S+)", line.strip())
            for line in contents.splitlines()
        )
        if match
    ]
    assert len(ran) == 2, (
        "--het-group=0,1 should have run a task in each component it names,"
        f" but got {contents!r}"
    )

    for argv0, node in ran:
        assert_documented_bcast_name(f"The component on {node}", argv0, cwd)

    assert len({argv0 for argv0, _ in ran}) == 2, (
        "Each component named by --het-group=0,1 should have been broadcast to"
        f" separately, but both ran {ran[0][0]}"
    )

    # heterogeneous_jobs.shtml says all components of a job step share the step
    # ID, while each component carries its own job id. This is the documented
    # cross-component invariant, not a credential detector.
    ids = [
        re.fullmatch(r"slurm_bcast_(\d+)\.(\d+)_\S+", os.path.basename(argv0)).groups()
        for argv0, _ in ran
    ]
    assert ids[0][1] == ids[1][1], (
        "Components of one --het-group=0,1 step should share a step id, but got"
        f" {ids[0][1]} and {ids[1][1]}"
    )
    assert (
        ids[0][0] != ids[1][0]
    ), f"Each component should carry its own job id, but both used {ids[0][0]}"