Ask a cluster for a machine when nothing here will do
Slurm is not a machine that attaches and stays; it is a queue somebody else owns. So nothing here submits a node to it. It submits a job whose payload is an ordinary worker dialling back in, and everything downstream — the protocol, the artifacts, cancellation, the books — already worked and did not have to learn what Slurm is. The alternative, which Covalent takes, is to stage a serialized call and a runner onto the login node, poll squeue and copy the result back: a second way of running a node beside the one that exists. The cost of not doing that is one assumption, that a compute node can open a connection outward. Where that is false, _payload is the single method a staged variant would replace. Clusters are configured in provisioners.json beside the alerts, since this is infrastructure an operator writes rather than anything a flow says. The script is generated with the system ssh and no new dependency, and prerun owns the environment — deliberately no pip install, because what is on a cluster is somebody's decision. One outstanding request per profile, cancelled if it never attaches and on the way out. Nothing autoscales. The run gate needed the same hook: a run held before it starts never reaches the placer's own wait, so it would have queued forever on a machine nothing had asked for. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01A6HeySA27EkGANZN95QySW
This commit is contained in:
@@ -24,6 +24,7 @@ DERIVED_PATHS = {
|
||||
"PANELS_FILE": "panels.json",
|
||||
"OAUTH_PRIVATE_KEY_FILE": "oauth-key.pem",
|
||||
"CLOUD_CONFIG_FILE": "cloud.json",
|
||||
"PROVISIONERS_FILE": "provisioners.json",
|
||||
}
|
||||
|
||||
|
||||
@@ -65,6 +66,10 @@ class Settings(BaseSettings):
|
||||
# Which failures reach which channel. Beside the flows, not in them:
|
||||
# alerting is the deployment's concern, not any one flow's.
|
||||
ALERTS_FILE: Path = Path("flow-data/alerts.json")
|
||||
# Where machines can be started from when a node needs one and nothing that
|
||||
# could take it is attached. Operator-authored, like the alerts beside it,
|
||||
# and absent on an installation that has nowhere to start one.
|
||||
PROVISIONERS_FILE: Path = Path("flow-data/provisioners.json")
|
||||
# Which dashboards each device shows. Beside the flows for the same reason
|
||||
# alerting is: where a screen hangs is the deployment's concern rather than
|
||||
# any one dashboard's.
|
||||
|
||||
@@ -297,6 +297,22 @@ class Placer:
|
||||
provisioner.provision(cpus, gpus, ram_mb, device)
|
||||
return
|
||||
|
||||
def provision_for(self, needs: dict[str, Any] | None) -> None:
|
||||
"""Ask for a machine a queued run is waiting on.
|
||||
|
||||
A run held before it starts never reaches the wait above, so this is
|
||||
where the asking happens for it — otherwise it would sit waiting for a
|
||||
machine that nothing ever requested.
|
||||
"""
|
||||
if not needs:
|
||||
return
|
||||
self._provision(
|
||||
int(needs.get("cpus") or 1),
|
||||
int(needs.get("gpus") or 0),
|
||||
int(needs.get("ram_mb") or 0),
|
||||
needs.get("device") or None,
|
||||
)
|
||||
|
||||
def _announce(self, node: str, run: str, since: float, reason: str) -> int:
|
||||
"""Say a node is queued, not stuck.
|
||||
|
||||
|
||||
@@ -0,0 +1,378 @@
|
||||
"""Machines started on demand, for a node that has nowhere to run.
|
||||
|
||||
A cluster is not a machine that attaches and stays. It is a queue somebody else
|
||||
owns, and what you get from it is an allocation for as long as your job holds
|
||||
one. So a provisioner does not run nodes: it *asks* for a machine, and what
|
||||
comes back is an ordinary worker dialling in on the ordinary socket. Everything
|
||||
downstream of that — the protocol, the artifacts, cancellation, the books —
|
||||
already works, and none of it had to learn what Slurm is.
|
||||
|
||||
That is the whole design. The alternative, which Covalent takes, is to stage a
|
||||
serialized call and a runner script onto the login node, submit *that*, poll
|
||||
``squeue``, and copy the result back: a second way of running a node, beside
|
||||
the one that already exists, with its own transport and its own failure modes.
|
||||
The cost of doing it this way instead is one assumption — that a compute node
|
||||
can open a connection outward — which is true of most clusters and false of
|
||||
air-gapped ones. Where it is false, :meth:`SlurmProvisioner._payload` is the
|
||||
one method a staged variant would need to replace.
|
||||
|
||||
Nothing here autoscales. One outstanding request per profile, dropped when the
|
||||
machine attaches or when it has taken too long, and the job is cancelled on the
|
||||
way out.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any, Protocol
|
||||
|
||||
from fluksio.core import security
|
||||
from fluksio.flow.events import EventBus
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: How long a submission is given before it is treated as failed.
|
||||
SUBMIT_TIMEOUT_S = 30.0
|
||||
#: How long a worker's credential is good for. Short: it sits in a scheduler's
|
||||
#: spool file, and the machine it is for is expected within the hour.
|
||||
TOKEN_HOURS = 24
|
||||
|
||||
|
||||
class Provisioner(Protocol):
|
||||
"""Somewhere machines can be asked for.
|
||||
|
||||
Implemented once, for Slurm. The shape is deliberately small — asking is
|
||||
not scheduling, and the placer already does the scheduling.
|
||||
"""
|
||||
|
||||
name: str
|
||||
|
||||
def covers(
|
||||
self, cpus: int, gpus: int, ram_mb: int, label: str | None = None
|
||||
) -> bool:
|
||||
"""Whether a machine this could start would fit this node."""
|
||||
|
||||
def provision(
|
||||
self, cpus: int, gpus: int, ram_mb: int, label: str | None = None
|
||||
) -> None:
|
||||
"""Ask for one. Returns at once; the machine arrives by attaching."""
|
||||
|
||||
def shapes(self, label: str | None = None) -> list[tuple[int, int, int]]:
|
||||
"""The sizes it can start, as cpus, gpus and MB."""
|
||||
|
||||
def reconcile(self, attached: set[str]) -> None:
|
||||
"""Forget requests that arrived, and give up on ones that did not."""
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Cancel whatever is outstanding."""
|
||||
|
||||
def status(self) -> dict[str, Any]:
|
||||
"""What it is called, what is outstanding, and what last went wrong."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SlurmProfile:
|
||||
"""One kind of machine this cluster can be asked for.
|
||||
|
||||
Named like a flavor and deliberately not the same thing: a flavor is what a
|
||||
node asks for, a profile is what a scheduler is asked for. They agree when
|
||||
somebody sets them up to.
|
||||
"""
|
||||
|
||||
name: str
|
||||
cpus: int = 1
|
||||
gpus: int = 0
|
||||
ram_mb: int = 0
|
||||
#: What a worker started this way advertises, so a node bound to a device
|
||||
#: can be the reason one is started.
|
||||
labels: list[str] = field(default_factory=list)
|
||||
#: ``#SBATCH`` lines, verbatim: partition, gres, walltime, account.
|
||||
sbatch: list[str] = field(default_factory=list)
|
||||
#: Shell run before the worker starts. Where the environment comes from —
|
||||
#: `module load`, a venv with fluksio-worker already in it. Deliberately
|
||||
#: not pip: what is installed on a cluster is somebody's decision, not this.
|
||||
prerun: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Job:
|
||||
job_id: str
|
||||
worker: str
|
||||
since: float
|
||||
|
||||
|
||||
class SlurmProvisioner:
|
||||
"""A Slurm cluster, asked for machines over ssh.
|
||||
|
||||
Submission is an ``sbatch`` whose payload is a worker dialling back here.
|
||||
The engine needs no route to the compute node and no share of its
|
||||
filesystem; what it needs is for the compute node to reach the engine.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
login: str,
|
||||
engine_url: str,
|
||||
profiles: list[SlurmProfile],
|
||||
ssh_key: str = "",
|
||||
artifact_url: str = "",
|
||||
max_idle_s: float = 300.0,
|
||||
provision_timeout_s: float = 900.0,
|
||||
events: EventBus | None = None,
|
||||
) -> None:
|
||||
self.name = name
|
||||
self.login = login
|
||||
self.engine_url = engine_url
|
||||
self.profiles = sorted(profiles, key=lambda p: (p.gpus, p.cpus, p.ram_mb))
|
||||
self.ssh_key = ssh_key
|
||||
self.artifact_url = artifact_url
|
||||
self.max_idle_s = max_idle_s
|
||||
self.provision_timeout_s = provision_timeout_s
|
||||
self.events = events
|
||||
self.last_error = ""
|
||||
self._outstanding: dict[str, _Job] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
# -- what it can be asked for ----------------------------------------------
|
||||
|
||||
def _matching(
|
||||
self, cpus: int, gpus: int, ram_mb: int, label: str | None
|
||||
) -> SlurmProfile | None:
|
||||
"""The smallest profile that would fit, or None."""
|
||||
for profile in self.profiles:
|
||||
if label and label not in profile.labels and label != profile.name:
|
||||
continue
|
||||
if profile.cpus >= cpus and profile.gpus >= gpus:
|
||||
if not ram_mb or not profile.ram_mb or profile.ram_mb >= ram_mb:
|
||||
return profile
|
||||
return None
|
||||
|
||||
def covers(
|
||||
self, cpus: int, gpus: int, ram_mb: int, label: str | None = None
|
||||
) -> bool:
|
||||
return self._matching(cpus, gpus, ram_mb, label) is not None
|
||||
|
||||
def shapes(self, label: str | None = None) -> list[tuple[int, int, int]]:
|
||||
return [
|
||||
(profile.cpus, profile.gpus, profile.ram_mb)
|
||||
for profile in self.profiles
|
||||
if not label or label in profile.labels or label == profile.name
|
||||
]
|
||||
|
||||
# -- asking ----------------------------------------------------------------
|
||||
|
||||
def provision(
|
||||
self, cpus: int, gpus: int, ram_mb: int, label: str | None = None
|
||||
) -> None:
|
||||
profile = self._matching(cpus, gpus, ram_mb, label)
|
||||
if profile is None:
|
||||
return
|
||||
with self._lock:
|
||||
if profile.name in self._outstanding:
|
||||
# One at a time per profile. The placer asks on every pass of
|
||||
# its wait, and a queue full of jobs nobody needed is worse
|
||||
# than a node waiting a little longer.
|
||||
return
|
||||
self._outstanding[profile.name] = _Job("", "", time.monotonic())
|
||||
# Off the caller's thread: it is holding the placer's condition, and
|
||||
# ssh to a login node is not something to hold a lock across.
|
||||
threading.Thread(
|
||||
target=self._submit,
|
||||
args=(profile,),
|
||||
name=f"provision-{self.name}",
|
||||
daemon=True,
|
||||
).start()
|
||||
|
||||
def _submit(self, profile: SlurmProfile) -> None:
|
||||
worker = f"{self.name}-{profile.name}-{uuid.uuid4().hex[:8]}"
|
||||
token = security.create_worker_token(worker, timedelta(hours=TOKEN_HOURS))
|
||||
try:
|
||||
result = self._ssh(["sbatch", "--parsable"], self._payload(profile, token))
|
||||
except (OSError, subprocess.SubprocessError) as exc:
|
||||
self._failed(profile, str(exc))
|
||||
return
|
||||
if result.returncode != 0:
|
||||
self._failed(profile, (result.stderr or result.stdout).strip()[:500])
|
||||
return
|
||||
job_id = (result.stdout or "").strip().splitlines()[0] if result.stdout else ""
|
||||
with self._lock:
|
||||
self._outstanding[profile.name] = _Job(job_id, worker, time.monotonic())
|
||||
logger.info(
|
||||
"Asked %s for a '%s' machine: job %s, worker '%s'",
|
||||
self.name,
|
||||
profile.name,
|
||||
job_id or "?",
|
||||
worker,
|
||||
)
|
||||
self._publish("worker_provisioned", profile, job=job_id, worker=worker)
|
||||
|
||||
def _payload(self, profile: SlurmProfile, token: str) -> str:
|
||||
"""The batch script: a worker that dials in, and stops when it is done.
|
||||
|
||||
A staged executor — for a cluster whose compute nodes cannot reach the
|
||||
engine — is this method and nothing else.
|
||||
"""
|
||||
worker = [
|
||||
"exec fluksio-worker",
|
||||
f"--url {self.engine_url}",
|
||||
f"--token {token}",
|
||||
f"--cpus {profile.cpus}",
|
||||
f"--gpus {profile.gpus}",
|
||||
f"--max-idle {self.max_idle_s:.0f}",
|
||||
]
|
||||
if profile.ram_mb:
|
||||
worker.append(f"--ram-mb {profile.ram_mb}")
|
||||
if profile.labels:
|
||||
worker.append(f"--labels {','.join(profile.labels)}")
|
||||
if self.artifact_url:
|
||||
worker.append(f"--artifact-url {self.artifact_url}")
|
||||
lines = [
|
||||
"#!/bin/bash",
|
||||
f"#SBATCH --job-name=fluksio-{profile.name}",
|
||||
*[f"#SBATCH {option}" for option in profile.sbatch],
|
||||
"",
|
||||
*profile.prerun,
|
||||
" ".join(worker),
|
||||
"",
|
||||
]
|
||||
return "\n".join(lines)
|
||||
|
||||
# -- keeping the outstanding list honest -----------------------------------
|
||||
|
||||
def reconcile(self, attached: set[str]) -> None:
|
||||
expired = []
|
||||
with self._lock:
|
||||
for name, job in list(self._outstanding.items()):
|
||||
if job.worker and job.worker in attached:
|
||||
# It arrived. Asking again is somebody else's decision.
|
||||
del self._outstanding[name]
|
||||
elif time.monotonic() - job.since > self.provision_timeout_s:
|
||||
del self._outstanding[name]
|
||||
if job.job_id:
|
||||
expired.append(job.job_id)
|
||||
for job_id in expired:
|
||||
logger.warning(
|
||||
"%s: job %s never attached, cancelling it", self.name, job_id
|
||||
)
|
||||
self._cancel(job_id)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
with self._lock:
|
||||
jobs = [job.job_id for job in self._outstanding.values() if job.job_id]
|
||||
self._outstanding.clear()
|
||||
for job_id in jobs:
|
||||
self._cancel(job_id)
|
||||
|
||||
def _cancel(self, job_id: str) -> None:
|
||||
try:
|
||||
self._ssh(["scancel", job_id])
|
||||
except (OSError, subprocess.SubprocessError) as exc:
|
||||
logger.warning("%s: could not cancel job %s: %s", self.name, job_id, exc)
|
||||
|
||||
def status(self) -> dict[str, Any]:
|
||||
with self._lock:
|
||||
outstanding = [
|
||||
{
|
||||
"profile": name,
|
||||
"job": job.job_id,
|
||||
"worker": job.worker,
|
||||
"seconds": round(time.monotonic() - job.since, 1),
|
||||
}
|
||||
for name, job in self._outstanding.items()
|
||||
]
|
||||
return {
|
||||
"name": self.name,
|
||||
"kind": "slurm",
|
||||
"profiles": [profile.name for profile in self.profiles],
|
||||
"outstanding": outstanding,
|
||||
"last_error": self.last_error,
|
||||
}
|
||||
|
||||
# -- the one thing that talks to the cluster -------------------------------
|
||||
|
||||
def _ssh(
|
||||
self, command: list[str], stdin: str = ""
|
||||
) -> subprocess.CompletedProcess[str]:
|
||||
"""Run a command on the login node. The system ssh, and no new library."""
|
||||
argv = ["ssh", "-o", "BatchMode=yes"]
|
||||
if self.ssh_key:
|
||||
argv += ["-i", self.ssh_key]
|
||||
argv += [self.login, *command]
|
||||
return subprocess.run(
|
||||
argv,
|
||||
input=stdin,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=SUBMIT_TIMEOUT_S,
|
||||
check=False,
|
||||
)
|
||||
|
||||
def _failed(self, profile: SlurmProfile, detail: str) -> None:
|
||||
self.last_error = detail
|
||||
with self._lock:
|
||||
self._outstanding.pop(profile.name, None)
|
||||
logger.error(
|
||||
"%s: could not start a '%s' machine: %s", self.name, profile.name, detail
|
||||
)
|
||||
self._publish("worker_provision_failed", profile, detail=detail)
|
||||
|
||||
def _publish(self, kind: str, profile: SlurmProfile, **extra: Any) -> None:
|
||||
if self.events is None:
|
||||
return
|
||||
self.events.publish(
|
||||
{
|
||||
"type": kind,
|
||||
"provisioner": self.name,
|
||||
"profile": profile.name,
|
||||
"ts": time.time(),
|
||||
**extra,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def load_provisioners(path: Path, events: EventBus | None = None) -> list[Any]:
|
||||
"""Read the configured clusters. An installation with none has no file."""
|
||||
if not path.exists():
|
||||
return []
|
||||
try:
|
||||
entries = json.loads(path.read_text())
|
||||
except (OSError, ValueError):
|
||||
logger.exception("Could not read %s; no machines can be started", path)
|
||||
return []
|
||||
|
||||
found = []
|
||||
for entry in entries if isinstance(entries, list) else []:
|
||||
kind = str(entry.get("type") or "slurm")
|
||||
if kind != "slurm":
|
||||
logger.error("%s: no provisioner of kind '%s'", path, kind)
|
||||
continue
|
||||
try:
|
||||
found.append(
|
||||
SlurmProvisioner(
|
||||
name=str(entry["name"]),
|
||||
login=str(entry["login"]),
|
||||
engine_url=str(entry["engine_url"]),
|
||||
profiles=[
|
||||
SlurmProfile(**profile) for profile in entry.get("profiles", [])
|
||||
],
|
||||
ssh_key=str(entry.get("ssh_key") or ""),
|
||||
artifact_url=str(entry.get("artifact_url") or ""),
|
||||
max_idle_s=float(entry.get("max_idle_s", 300)),
|
||||
provision_timeout_s=float(entry.get("provision_timeout_s", 900)),
|
||||
events=events,
|
||||
)
|
||||
)
|
||||
except (KeyError, TypeError, ValueError):
|
||||
logger.exception("%s: could not read a provisioner", path)
|
||||
if found:
|
||||
logger.info("%d provisioner(s) can start machines", len(found))
|
||||
return found
|
||||
@@ -929,6 +929,10 @@ class RunService:
|
||||
short = placer.satisfiable(needs)
|
||||
if short:
|
||||
missing.append(f"with {short}")
|
||||
# Held here rather than at the placer's own wait, so this is
|
||||
# where the machine has to be asked for. Asked again on every
|
||||
# pass; one outstanding request is the provisioner's business.
|
||||
placer.provision_for(needs)
|
||||
return missing
|
||||
|
||||
def _waiting(self, run_id: str, missing: list[str]) -> None:
|
||||
|
||||
@@ -31,6 +31,7 @@ from fluksio.flow.nodes.http import close_shared_client
|
||||
from fluksio.flow.pipeline import ValueSource
|
||||
from fluksio.flow.placement import Placer
|
||||
from fluksio.flow.plugins import load_plugins
|
||||
from fluksio.flow.provision import load_provisioners
|
||||
from fluksio.flow.queue import MemoryWorkQueue, RedisWorkQueue, WorkQueue
|
||||
from fluksio.flow.remote import RemoteWorkerHub
|
||||
from fluksio.flow.resources import ResourceAccountant, fair_share_env
|
||||
@@ -147,6 +148,8 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
||||
app.state.resources = accountant
|
||||
# Every machine a node could run on: this one, and whatever attaches.
|
||||
placer = Placer(local=accountant, events=event_bus)
|
||||
# Where more machines can be asked for when nothing attached will do.
|
||||
placer.provisioners = load_provisioners(settings.PROVISIONERS_FILE, event_bus)
|
||||
app.state.placer = placer
|
||||
pool = PythonWorkerPool(
|
||||
python=modules.venv_python(),
|
||||
@@ -251,6 +254,10 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
||||
await run_in_threadpool(run_service.stop)
|
||||
await controller.stop()
|
||||
pool.stop()
|
||||
# A machine asked for and not yet arrived would hold an allocation
|
||||
# nobody is going to use.
|
||||
for provisioner in placer.provisioners:
|
||||
await run_in_threadpool(provisioner.shutdown)
|
||||
close_shared_client()
|
||||
if settings.MCP_ENABLED:
|
||||
from fluksio.mcp.http import aclose
|
||||
|
||||
@@ -0,0 +1,219 @@
|
||||
"""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
|
||||
|
||||
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)
|
||||
|
||||
hpc.provision(cpus=2, gpus=0, ram_mb=0)
|
||||
assert ssh.ran.wait(5)
|
||||
hpc.reconcile(set())
|
||||
|
||||
assert hpc.status()["outstanding"] == []
|
||||
assert ["scancel", "4711"] == ssh.calls[-1][0][-2:]
|
||||
|
||||
|
||||
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_installation_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]
|
||||
Reference in New Issue
Block a user