stroblmeandClaude Opus 5
37a7df9d24
Docs / docs (push) Successful in 29s
Playwright Tests / test-playwright (1, 2) (push) Successful in 3m33s
Playwright Tests / test-playwright (2, 2) (push) Successful in 2m3s
pre-commit / pre-commit (push) Failing after 3m9s
Test Backend / test-backend (push) Successful in 2m46s
Compose Smoke Test / test-compose (push) Successful in 39s
Playwright Tests / merge-reports (push) Successful in 1m47s
Concurrent runs sat at 4 whatever FLOW_MAX_CASCADES said: that setting bounds cascades, and the run drivers read a hardcoded MAX_PARALLEL nobody could reach. FLOW_MAX_RUNS is the knob they read now, --max-runs/--max-cascades/--max-workers are the same three as flags on serve, and the engine says which numbers it started with — which is the only way to tell that a settings file was read. Events keep the run they happened in. The payload always carried it and the persist path dropped it, so reading one run's failures meant filtering the engine-wide list; a batch run's id reaches those events now too, since a run has no journaled item to name itself by. Also: a provisioner's 0 means "no deadline" rather than "cancel on the next reconcile", and a command that reaches no engine says how to start one. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015sbYeYaVgYQqm1sbx7wPdL
387 lines
14 KiB
Python
387 lines
14 KiB
Python
"""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 = "",
|
|
#: How long a machine sits idle before it gives itself back; 0 keeps it
|
|
#: for as long as the job runs, which is what a queue paid for in hours
|
|
#: of wall time wants.
|
|
max_idle_s: float = 300.0,
|
|
#: How long a submitted job may take to attach before it is cancelled;
|
|
#: 0 waits for as long as the scheduler makes it wait.
|
|
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 (
|
|
self.provision_timeout_s > 0
|
|
and 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
|