app/backend/tests/flow/test_provision.py

"""Asking a batch scheduler for a machine, and what comes back.

The point of doing it this way is that almost nothing here is new: what a
cluster gives back is an ordinary worker on the ordinary socket. So what is
worth checking is the asking — that it happens once, that the script says what
it should, and that a job nobody ever attached does not sit in the queue.
"""

import subprocess
import threading
import time

from fluksio.flow.placement import Placer
from fluksio.flow.provision import SlurmProfile, SlurmProvisioner, load_provisioners
from fluksio.flow.resources import ResourceAccountant
from fluksio.flow.schemas import Resources


class FakeSsh:
    """Stands in for the login node: records argv, answers like sbatch."""

    def __init__(self, returncode: int = 0, stdout: str = "4711\n") -> None:
        self.calls: list[tuple[list[str], str]] = []
        self.returncode = returncode
        self.stdout = stdout
        self.ran = threading.Event()

    def __call__(self, argv, **kwargs):
        self.calls.append((argv, kwargs.get("input") or ""))
        self.ran.set()
        return subprocess.CompletedProcess(
            argv, self.returncode, stdout=self.stdout, stderr="sbatch: bad partition"
        )


def cluster(events=None, **kwargs) -> SlurmProvisioner:
    return SlurmProvisioner(
        name="hpc",
        login="me@login.cluster",
        engine_url="wss://engine.example.com/api/v1/workers/attach",
        profiles=[
            SlurmProfile(name="cpu", cpus=16, ram_mb=32768),
            SlurmProfile(
                name="gpu-small",
                cpus=8,
                gpus=1,
                ram_mb=65536,
                labels=["gpu"],
                sbatch=["--partition=gpu", "--gres=gpu:1", "--time=04:00:00"],
                prerun=["module load cuda/12", "source ~/venvs/flux/bin/activate"],
            ),
        ],
        events=events,
        **kwargs,
    )


def test_the_smallest_machine_that_would_do_is_the_one_asked_for(monkeypatch):
    ssh = FakeSsh()
    monkeypatch.setattr(subprocess, "run", ssh)
    hpc = cluster()

    hpc.provision(cpus=2, gpus=0, ram_mb=0)
    assert ssh.ran.wait(5)
    argv, script = ssh.calls[0]
    assert argv[:3] == ["ssh", "-o", "BatchMode=yes"]
    assert argv[-2:] == ["sbatch", "--parsable"]
    assert "--job-name=fluksio-cpu" in script


def test_the_script_starts_a_worker_that_dials_back_and_stops_when_done(monkeypatch):
    ssh = FakeSsh()
    monkeypatch.setattr(subprocess, "run", ssh)
    hpc = cluster(max_idle_s=120)

    hpc.provision(cpus=4, gpus=1, ram_mb=0, label="gpu")
    assert ssh.ran.wait(5)
    _, script = ssh.calls[0]

    assert "#SBATCH --gres=gpu:1" in script
    # The environment is the cluster's business, not this module's.
    assert "module load cuda/12" in script
    assert "pip install" not in script
    assert "exec fluksio-worker" in script
    assert "--url wss://engine.example.com/api/v1/workers/attach" in script
    assert "--gpus 1" in script
    assert "--labels gpu" in script
    # Without this the job holds its allocation until the walltime.
    assert "--max-idle 120" in script


def test_one_job_per_profile_however_often_it_is_asked_for(monkeypatch):
    ssh = FakeSsh()
    monkeypatch.setattr(subprocess, "run", ssh)
    hpc = cluster()

    for _ in range(5):
        hpc.provision(cpus=2, gpus=0, ram_mb=0)
    assert ssh.ran.wait(5)
    assert len(ssh.calls) == 1
    assert hpc.status()["outstanding"][0]["job"] == "4711"


def test_a_machine_that_arrives_clears_the_way_for_the_next_ask(monkeypatch):
    ssh = FakeSsh()
    monkeypatch.setattr(subprocess, "run", ssh)
    hpc = cluster()

    hpc.provision(cpus=2, gpus=0, ram_mb=0)
    assert ssh.ran.wait(5)
    arrived = hpc.status()["outstanding"][0]["worker"]

    hpc.reconcile({arrived})
    assert hpc.status()["outstanding"] == []
    # No scancel: it did what it was asked to do.
    assert [argv for argv, _ in ssh.calls if "scancel" in argv] == []


def test_a_job_that_never_attaches_is_cancelled(monkeypatch):
    ssh = FakeSsh()
    monkeypatch.setattr(subprocess, "run", ssh)
    hpc = cluster(provision_timeout_s=0.01)

    hpc.provision(cpus=2, gpus=0, ram_mb=0)
    assert ssh.ran.wait(5)
    # Slept rather than set to zero: zero is what says "no deadline" now.
    time.sleep(0.02)
    hpc.reconcile(set())

    assert hpc.status()["outstanding"] == []
    assert ["scancel", "4711"] == ssh.calls[-1][0][-2:]


def test_no_deadline_waits_for_as_long_as_the_queue_does(monkeypatch):
    """0 is "never give up", not "give up now": a cluster can queue for days."""
    ssh = FakeSsh()
    monkeypatch.setattr(subprocess, "run", ssh)
    hpc = cluster(provision_timeout_s=0)

    hpc.provision(cpus=2, gpus=0, ram_mb=0)
    assert ssh.ran.wait(5)
    hpc.reconcile(set())

    assert hpc.status()["outstanding"][0]["job"] == "4711"
    assert [argv for argv, _ in ssh.calls if "scancel" in argv] == []


def test_a_refused_submission_says_what_the_cluster_said(monkeypatch):
    ssh = FakeSsh(returncode=1)
    monkeypatch.setattr(subprocess, "run", ssh)
    events: list[dict] = []
    hpc = cluster(
        events=type("Bus", (), {"publish": lambda self, e: events.append(e)})()
    )

    hpc.provision(cpus=2, gpus=0, ram_mb=0)
    assert ssh.ran.wait(5)
    for _ in range(100):
        if events:
            break
        threading.Event().wait(0.02)

    assert "bad partition" in hpc.last_error
    assert events[-1]["type"] == "worker_provision_failed"
    # And it is not left outstanding, so asking again actually asks.
    assert hpc.status()["outstanding"] == []


def test_a_node_waiting_for_a_machine_asks_for_one_once(monkeypatch):
    """The placer asks on every pass of its wait; the cluster is asked once."""
    ssh = FakeSsh()
    monkeypatch.setattr(subprocess, "run", ssh)
    hpc = cluster()
    placer = Placer(local=ResourceAccountant(cpus=4, gpus=0))
    placer.provisioners = [hpc]
    ran = threading.Event()

    def wants_a_card() -> None:
        with placer.claim(Resources(cpus=2, gpus=1), node="study.fit"):
            ran.set()

    thread = threading.Thread(target=wants_a_card, daemon=True)
    thread.start()
    # It waits rather than being cut down to what is here, because something
    # can start a machine that would take it.
    assert not ran.wait(0.5)
    assert len(ssh.calls) == 1

    placer.provisioners = []
    placer.wake()
    thread.join(timeout=5)
    assert ran.is_set()


def test_an_instance_with_nowhere_to_start_one_has_no_file(tmp_path):
    assert load_provisioners(tmp_path / "provisioners.json") == []


def test_a_configured_cluster_is_read(tmp_path):
    path = tmp_path / "provisioners.json"
    path.write_text(
        """[{"type": "slurm", "name": "hpc", "login": "me@login",
             "engine_url": "wss://e/api/v1/workers/attach",
             "profiles": [{"name": "gpu-small", "cpus": 8, "gpus": 1,
                           "labels": ["gpu"], "sbatch": ["--gres=gpu:1"]}]}]"""
    )
    found = load_provisioners(path)

    assert [p.name for p in found] == ["hpc"]
    assert found[0].covers(cpus=4, gpus=1, ram_mb=0, label="gpu")
    assert not found[0].covers(cpus=4, gpus=2, ram_mb=0)


def test_a_file_that_cannot_be_read_is_not_a_failure_to_start(tmp_path):
    path = tmp_path / "provisioners.json"
    path.write_text("{ not json")
    assert load_provisioners(path) == []


def test_a_run_held_before_it_starts_still_gets_a_machine_asked_for(monkeypatch):
    """The gate holds the run above the placer, so the asking has to happen there.

    Without this the run waits on a machine nothing ever requested, which is a
    queue entry that never moves and no error anywhere.
    """
    ssh = FakeSsh()
    monkeypatch.setattr(subprocess, "run", ssh)
    hpc = cluster()
    placer = Placer(local=ResourceAccountant(cpus=4, gpus=0))
    placer.provisioners = [hpc]
    needs = {"cpus": 2, "gpus": 1, "ram_mb": 0, "device": ""}

    assert placer.satisfiable(needs) == "1 gpu(s) and 2 cpu(s)"
    placer.provision_for(needs)
    assert ssh.ran.wait(5)
    assert "--gpus 1" in ssh.calls[0][1]