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",
|
"PANELS_FILE": "panels.json",
|
||||||
"OAUTH_PRIVATE_KEY_FILE": "oauth-key.pem",
|
"OAUTH_PRIVATE_KEY_FILE": "oauth-key.pem",
|
||||||
"CLOUD_CONFIG_FILE": "cloud.json",
|
"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:
|
# Which failures reach which channel. Beside the flows, not in them:
|
||||||
# alerting is the deployment's concern, not any one flow's.
|
# alerting is the deployment's concern, not any one flow's.
|
||||||
ALERTS_FILE: Path = Path("flow-data/alerts.json")
|
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
|
# 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
|
# alerting is: where a screen hangs is the deployment's concern rather than
|
||||||
# any one dashboard's.
|
# any one dashboard's.
|
||||||
|
|||||||
@@ -297,6 +297,22 @@ class Placer:
|
|||||||
provisioner.provision(cpus, gpus, ram_mb, device)
|
provisioner.provision(cpus, gpus, ram_mb, device)
|
||||||
return
|
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:
|
def _announce(self, node: str, run: str, since: float, reason: str) -> int:
|
||||||
"""Say a node is queued, not stuck.
|
"""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)
|
short = placer.satisfiable(needs)
|
||||||
if short:
|
if short:
|
||||||
missing.append(f"with {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
|
return missing
|
||||||
|
|
||||||
def _waiting(self, run_id: str, missing: list[str]) -> None:
|
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.pipeline import ValueSource
|
||||||
from fluksio.flow.placement import Placer
|
from fluksio.flow.placement import Placer
|
||||||
from fluksio.flow.plugins import load_plugins
|
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.queue import MemoryWorkQueue, RedisWorkQueue, WorkQueue
|
||||||
from fluksio.flow.remote import RemoteWorkerHub
|
from fluksio.flow.remote import RemoteWorkerHub
|
||||||
from fluksio.flow.resources import ResourceAccountant, fair_share_env
|
from fluksio.flow.resources import ResourceAccountant, fair_share_env
|
||||||
@@ -147,6 +148,8 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|||||||
app.state.resources = accountant
|
app.state.resources = accountant
|
||||||
# Every machine a node could run on: this one, and whatever attaches.
|
# Every machine a node could run on: this one, and whatever attaches.
|
||||||
placer = Placer(local=accountant, events=event_bus)
|
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
|
app.state.placer = placer
|
||||||
pool = PythonWorkerPool(
|
pool = PythonWorkerPool(
|
||||||
python=modules.venv_python(),
|
python=modules.venv_python(),
|
||||||
@@ -251,6 +254,10 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|||||||
await run_in_threadpool(run_service.stop)
|
await run_in_threadpool(run_service.stop)
|
||||||
await controller.stop()
|
await controller.stop()
|
||||||
pool.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()
|
close_shared_client()
|
||||||
if settings.MCP_ENABLED:
|
if settings.MCP_ENABLED:
|
||||||
from fluksio.mcp.http import aclose
|
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]
|
||||||
@@ -161,6 +161,57 @@ will never accept it; `pip install -U fluksio-worker` on that host is the whole
|
|||||||
upgrade. Nothing changed in the runner served at `GET /api/v1/workers/runtime`,
|
upgrade. Nothing changed in the runner served at `GET /api/v1/workers/runtime`,
|
||||||
so a host that copies its two files copies the same one as before.
|
so a host that copies its two files copies the same one as before.
|
||||||
|
|
||||||
|
## Machines from a batch scheduler
|
||||||
|
|
||||||
|
A cluster is not a machine that attaches and stays — it is a queue somebody else
|
||||||
|
owns. So Fluksio does not submit *nodes* to Slurm. It submits a job whose payload
|
||||||
|
is an ordinary worker dialling back in, and from there everything works the way
|
||||||
|
it already does: the same protocol, the same artifacts, the same cancellation.
|
||||||
|
|
||||||
|
Write the clusters into `provisioners.json` beside the flows:
|
||||||
|
|
||||||
|
```json
|
||||||
|
[{
|
||||||
|
"type": "slurm",
|
||||||
|
"name": "hpc",
|
||||||
|
"login": "me@login.cluster",
|
||||||
|
"ssh_key": "/secrets/hpc_ed25519",
|
||||||
|
"engine_url": "wss://api.example.com/api/v1/workers/attach",
|
||||||
|
"max_idle_s": 300,
|
||||||
|
"provision_timeout_s": 900,
|
||||||
|
"profiles": [{
|
||||||
|
"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"]
|
||||||
|
}]
|
||||||
|
}]
|
||||||
|
```
|
||||||
|
|
||||||
|
A **profile** is what the scheduler is asked for, where a flavor is what a node
|
||||||
|
asks for. They are separate on purpose, and agree when you set them up to.
|
||||||
|
|
||||||
|
When a node needs a machine nothing attached can give, and a profile would fit,
|
||||||
|
the engine `sbatch`es one over ssh — the system `ssh`, so nothing new is
|
||||||
|
installed — and the run waits meanwhile, saying so. `prerun` owns the
|
||||||
|
environment: a `module load`, a venv with `fluksio-worker` already in it. There
|
||||||
|
is deliberately no `pip install` in the generated script, because what is
|
||||||
|
installed on a cluster is somebody's decision and not this program's.
|
||||||
|
|
||||||
|
One outstanding request per profile, however often it is asked for. A job that
|
||||||
|
never attaches within `provision_timeout_s` is `scancel`led, as is anything
|
||||||
|
outstanding when the engine stops. `--max-idle` is what ends the job at the
|
||||||
|
other end, so an allocation goes back rather than idling to its walltime.
|
||||||
|
|
||||||
|
`GET /workers/resources` reports what is outstanding and what last went wrong.
|
||||||
|
|
||||||
|
!!! note "It needs a route out"
|
||||||
|
|
||||||
|
A compute node must be able to open a connection to the engine. That is
|
||||||
|
true of most clusters and false of air-gapped ones; there is no staging
|
||||||
|
path today.
|
||||||
|
|
||||||
## What a worker is not
|
## What a worker is not
|
||||||
|
|
||||||
It is not a second engine. Subscriptions, schedules, webhooks, the dashboards
|
It is not a second engine. Subscriptions, schedules, webhooks, the dashboards
|
||||||
|
|||||||
Reference in New Issue
Block a user