Melvin Strobl
1dc891aa1b
- NodeDef `device`/`device_policy` become `worker`/`worker_policy`, the issue `ignored_device` becomes `ignored_worker`, and the SDK's `@node` takes `worker=` and `worker_policy=` - FlowStore._migrate renames the keys in published and draft documents once, bumping their version; a queued run's `needs` reads the old key - placement passes the label as `label`, client regenerated - docs: worker placement wording, MCP issuer without DOMAIN, device rows in the CLI data directory and the API table, device links in compose comments
240 lines
8.1 KiB
Python
240 lines
8.1 KiB
Python
"""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
|
|
# The token travels in the environment: argv is in `ps` for anyone.
|
|
assert "--token" not in script
|
|
assert "\nexport FLUKSIO_WORKER_TOKEN=ey" 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, "worker": ""}
|
|
|
|
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]
|