diff --git a/autofit/campaign/__init__.py b/autofit/campaign/__init__.py new file mode 100644 index 000000000..32d2bbec9 --- /dev/null +++ b/autofit/campaign/__init__.py @@ -0,0 +1,22 @@ +""" +Campaigns: run many fits on a scheduler from one declaration (hpc-campaign phase 2). + +- ``spec``: ``campaign.yaml`` (stages, manifest, chunks, resources, dependencies, + success checks, propose-only retry policy, the autonerves site); +- ``scheduler``: the ``Scheduler`` protocol and its Fake, Local and Slurm adapters; +- ``lint``: the partition lint (CPU-only work never goes on a GPU partition); +- ``submit``: arrays per stage and chunk, chained per index, with a job ledger. + +Command line: ``python -m autofit.campaign {validate,submit}``; in Python, +``autofit.campaign.submit.submit(campaign, scheduler)``. + +Run states and failure classes come from ``autofit.non_linear.status.states``. +""" +from autofit.campaign.lint import ( + LintResult, + PartitionLintError, + check_partitions, + check_sbatch, +) +from autofit.campaign.spec import Campaign, CampaignException, Resources, Stage +from autofit.campaign.submit import read_ledger diff --git a/autofit/campaign/__main__.py b/autofit/campaign/__main__.py new file mode 100644 index 000000000..a2efe46e5 --- /dev/null +++ b/autofit/campaign/__main__.py @@ -0,0 +1,88 @@ +""" +Command-line entry point: ``python -m autofit.campaign ``. + +Commands +-------- +validate + Check the campaign file against its site; print every problem. +submit [--stage S ...] [--indices SPEC] [--dry-run] + Submit the stages (default: all) over the manifest (or ``--indices``, a Slurm + array spec of campaign indices), chunk by chunk. Every array is linted against + the GPU-partition rule first. ``--dry-run`` prints the batch scripts and submits + nothing. +""" +import argparse +import json +import sys + +from autofit.campaign.lint import PartitionLintError +from autofit.campaign.spec import Campaign, CampaignException +from autofit.campaign.units import parse_array_spec + + +def _validate(arguments) -> int: + try: + campaign = Campaign.load(arguments.campaign) + except CampaignException as exception: + print(exception, file=sys.stderr) + return 1 + chunks = campaign.chunks() + print( + f"{campaign.name}: valid — {len(campaign.indices())} objects, " + f"{len(chunks)} chunk(s) of ≤{campaign.chunk_size()}, " + f"stages {' → '.join(campaign.order())}" + ) + if campaign.site_name and campaign.site is not None and campaign.site.unverified(): + print(f"site {campaign.site.name}: unverified values {campaign.site.unverified()}") + return 0 + + +def _submit(arguments) -> int: + from autofit.campaign.scheduler import scheduler_for + from autofit.campaign.submit import submit + + try: + campaign = Campaign.load(arguments.campaign) + indices = parse_array_spec(arguments.indices)[0] if arguments.indices else None + scheduler = None if arguments.dry_run else scheduler_for(campaign) + entries = submit( + campaign, + scheduler, + stages=arguments.stage, + indices=indices, + dry_run=arguments.dry_run, + ) + except (CampaignException, PartitionLintError) as exception: + print(f"refused: {exception}", file=sys.stderr) + return 1 + if arguments.dry_run: + for entry in entries: + print(f"# {entry['stage']} chunk {entry['offset']} (lint: {entry['lint']})") + print(entry.get("script", json.dumps(entry, indent=1))) + return 0 + for entry in entries: + print(f"{entry['stage']} chunk {entry['offset']}: job {entry['job_id']}") + return 0 + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(prog="python -m autofit.campaign") + commands = parser.add_subparsers(dest="command", required=True) + + validate = commands.add_parser("validate", help="Check a campaign file.") + validate.add_argument("campaign") + validate.set_defaults(handler=_validate) + + submit = commands.add_parser("submit", help="Submit a campaign's stages.") + submit.add_argument("campaign") + submit.add_argument("--stage", action="append", help="Stage to submit (repeatable).") + submit.add_argument("--indices", help="Campaign indices, as a Slurm array spec.") + submit.add_argument("--dry-run", action="store_true", help="Print, do not submit.") + submit.set_defaults(handler=_submit) + + arguments = parser.parse_args(argv) + return arguments.handler(arguments) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/autofit/campaign/lint.py b/autofit/campaign/lint.py new file mode 100644 index 000000000..d096b9042 --- /dev/null +++ b/autofit/campaign/lint.py @@ -0,0 +1,244 @@ +""" +The partition lint: CPU-only work never goes on a GPU partition. + +On 2026-09-30 CPU-only arrays submitted to RAL with ``--partition=ral,gpu`` filled every +CPU of both A100 nodes and left all eight GPUs idle and unschedulable for hours. Every +submit, resubmit and pending-job update in ``autofit.campaign`` passes through +``check_partitions``, which refuses a partition list containing a GPU partition for a +request that asks for no GPU. + +A GPU partition is one the site file marks ``gpu_only`` (``autonerves.site_config``); +with no site file, any partition whose name contains ``gpu``. Whether the request asks +for a GPU is read from its ``--gres`` / ``--gpus`` / TRES request, never from the +partition name. + +The one exemption is a small, bounded CPU timing leg on a GPU partition on its own +(never a mixed list such as ``ral,gpu``), and only when all three hold at submit time: + +- at most 8 CPUs per task; +- an array throttle of at most ``%2`` (or at most two tasks in all); +- no job pending on that partition that requests a GPU. + +When the pending-GPU count is unknown the exemption is refused: the rule fails closed. +""" +import re +import shlex +from dataclasses import dataclass, field +from typing import List, Optional, Sequence, Union + +from autofit.campaign.units import parse_array_spec, requests_gpu + +EXEMPT_MAX_CPUS = 8 +EXEMPT_MAX_THROTTLE = 2 + + +class PartitionLintError(Exception): + pass + + +@dataclass +class LintResult: + ok: bool + partitions: List[str] + gpu_partitions: List[str] + requests_gpu: bool + exempt: bool = False + reasons: List[str] = field(default_factory=list) + + def raise_for_refusal(self): + if not self.ok: + raise PartitionLintError("; ".join(self.reasons)) + + +def split_partitions(partitions: Union[str, Sequence[str], None]) -> List[str]: + if partitions is None: + return [] + if isinstance(partitions, str): + partitions = partitions.split(",") + return [p.strip() for p in partitions if p and p.strip()] + + +def gpu_partitions_of(partitions: Sequence[str], site=None) -> List[str]: + """ + The members of ``partitions`` that are GPU partitions: ``gpu_only`` in ``site`` + when it names the partition, else any whose name contains ``gpu``. + """ + found = [] + for name in partitions: + site_partition = None + if site is not None: + site_partition = getattr(site, "partitions", {}).get(name) + if site_partition is not None: + if site_partition.gpu_only: + found.append(name) + elif "gpu" in name.lower(): + found.append(name) + return found + + +def check_partitions( + partitions: Union[str, Sequence[str], None], + cpus: Optional[int] = None, + gres: Optional[str] = None, + throttle: Optional[int] = None, + n_tasks: Optional[int] = None, + pending_gpu_jobs: Optional[int] = None, + site=None, +) -> LintResult: + """ + Lint one request's partition list against the GPU-partition rule. + + Parameters + ---------- + partitions + A partition name, a Slurm list (``"ral,gpu"``) or a list of names. + cpus + CPUs per task. + gres + The request's GPU / generic-resource request (``gpu:1``, ``gres/gpu=1``). + throttle + The array throttle (``%N``), if any. + n_tasks + How many tasks the request launches (1 for a plain job). + pending_gpu_jobs + How many jobs are pending on the GPU partition requesting a GPU, measured at + submit time; ``None`` when unknown (refuses the exemption). + site + The ``autonerves.site_config.Site`` naming which partitions are ``gpu_only``. + + Returns + ------- + The verdict, with the reasons for a refusal. + """ + names = split_partitions(partitions) + gpu_names = gpu_partitions_of(names, site=site) + gpu_request = requests_gpu(gres) + result = LintResult( + ok=True, partitions=names, gpu_partitions=gpu_names, requests_gpu=gpu_request + ) + if not gpu_names or gpu_request: + return result + + listed = ",".join(names) + if len(names) > len(gpu_names): + result.ok = False + result.reasons.append( + f"partition list {listed!r} mixes GPU partition(s) {gpu_names} with CPU " + f"partitions for a request with no GPU: CPU-only work never goes on a GPU " + f"partition; use {[n for n in names if n not in gpu_names]} alone" + ) + return result + + failures = [] + if cpus is None or cpus > EXEMPT_MAX_CPUS: + failures.append(f"{cpus} CPUs per task (exemption needs ≤{EXEMPT_MAX_CPUS})") + bounded = (throttle is not None and throttle <= EXEMPT_MAX_THROTTLE) or ( + n_tasks is not None and n_tasks <= EXEMPT_MAX_THROTTLE + ) + if not bounded: + failures.append( + f"throttle {throttle if throttle is not None else 'none'} over " + f"{n_tasks if n_tasks is not None else 'unknown'} tasks " + f"(exemption needs %≤{EXEMPT_MAX_THROTTLE})" + ) + if pending_gpu_jobs is None: + failures.append("pending GPU jobs unknown (exemption needs none pending)") + elif pending_gpu_jobs > 0: + failures.append(f"{pending_gpu_jobs} GPU job(s) pending on {gpu_names}") + + if failures: + result.ok = False + result.reasons.append( + f"partition {listed!r} is a GPU partition and the request asks for no GPU; " + f"the timing-leg exemption does not apply: " + "; ".join(failures) + ) + return result + + result.exempt = True + return result + + +_SBATCH = re.compile(r"^\s*#SBATCH\s+(.*)$") + + +@dataclass +class SbatchRequest: + partition: Optional[str] = None + cpus: Optional[int] = None + gres: Optional[str] = None + throttle: Optional[int] = None + n_tasks: Optional[int] = None + + +def _apply_option(request: SbatchRequest, option: str, value: Optional[str]): + if option in ("--partition", "-p"): + request.partition = value + elif option in ("--cpus-per-task", "-c"): + request.cpus = int(value) + elif option in ("--gres", "--tres-per-node", "--tres-per-task"): + request.gres = value if request.gres is None else f"{request.gres},{value}" + elif option in ("--gpus", "-G", "--gpus-per-node", "--gpus-per-task"): + gres = f"gpu:{value}" + request.gres = gres if request.gres is None else f"{request.gres},{gres}" + elif option in ("--array", "-a"): + indices, throttle = parse_array_spec(value) + request.n_tasks = len(indices) + request.throttle = throttle + + +def _parse_options(tokens: List[str], request: SbatchRequest): + i = 0 + while i < len(tokens): + token = tokens[i] + if token.startswith("--") and "=" in token: + option, value = token.split("=", 1) + elif token.startswith("-"): + option = token + value = tokens[i + 1] if i + 1 < len(tokens) else None + i += 1 + else: + i += 1 + continue + _apply_option(request, option, value) + i += 1 + + +def sbatch_request(script: str, arguments: Sequence[str] = ()) -> SbatchRequest: + """ + The partition, CPUs, GPU request and array of a batch script's ``#SBATCH`` + directives, with command-line ``arguments`` (``sbatch -p gpu ...``) overriding + them as ``sbatch`` does. + """ + request = SbatchRequest(n_tasks=1) + for line in script.splitlines(): + stripped = line.strip() + if stripped and not stripped.startswith("#"): + # sbatch reads directives only up to the first command. + break + match = _SBATCH.match(line) + if match: + _parse_options(shlex.split(match.group(1), comments=True), request) + _parse_options(list(arguments), request) + return request + + +def check_sbatch( + script: str, + arguments: Sequence[str] = (), + pending_gpu_jobs: Optional[int] = None, + site=None, +) -> LintResult: + """ + Lint a batch script (and any ``sbatch`` command-line overrides) against the + GPU-partition rule; see ``check_partitions``. + """ + request = sbatch_request(script, arguments) + return check_partitions( + request.partition, + cpus=request.cpus if request.cpus is not None else 1, + gres=request.gres, + throttle=request.throttle, + n_tasks=request.n_tasks, + pending_gpu_jobs=pending_gpu_jobs, + site=site, + ) diff --git a/autofit/campaign/scheduler/__init__.py b/autofit/campaign/scheduler/__init__.py new file mode 100644 index 000000000..6aa5b0dce --- /dev/null +++ b/autofit/campaign/scheduler/__init__.py @@ -0,0 +1,29 @@ +""" +Scheduler adapters for ``autofit.campaign``: the ``Scheduler`` protocol, an in-memory +``FakeScheduler`` for tests, a subprocess ``LocalScheduler`` and the ``SlurmScheduler``. +None of them runs a command at import. +""" +from autofit.non_linear.status.states import RunState +from autofit.campaign.scheduler.base import ( + CANCELLED, + PREEMPTED, + REQUEUED, + UNKNOWN, + PartitionCapacity, + Scheduler, + SubmitRequest, + TaskRecord, + run_state_of, +) +from autofit.campaign.scheduler.fake import FakeScheduler +from autofit.campaign.scheduler.local import LocalScheduler +from autofit.campaign.scheduler.slurm import SlurmScheduler, sbatch_script + + +def scheduler_for(campaign, runner=None): + """ + The scheduler a campaign's ``scheduler:`` names (``slurm`` or ``local``). + """ + if campaign.scheduler == "local": + return LocalScheduler() + return SlurmScheduler(runner=runner, site=campaign.site if campaign.site_name else None) diff --git a/autofit/campaign/scheduler/base.py b/autofit/campaign/scheduler/base.py new file mode 100644 index 000000000..0a496e849 --- /dev/null +++ b/autofit/campaign/scheduler/base.py @@ -0,0 +1,209 @@ +""" +The ``Scheduler`` protocol every campaign adapter implements, and the records it +returns. + +A scheduler state maps onto the shared run-state vocabulary +(``autofit.non_linear.status.states.RunState``): ``TIMEOUT``, ``OUT_OF_MEMORY`` and +``NODE_FAIL`` become the collector-only states ``timeout``, ``oom`` and ``node_fail``, +which a fit cannot write for itself. A few scheduler outcomes have no run state, because +no fit ran to an end: they keep a ``scheduler_state`` of ``cancelled``, ``requeued`` or +``preempted``. +""" +from dataclasses import dataclass, field +from datetime import datetime +from typing import Dict, List, Optional, Protocol, Sequence, runtime_checkable + +from autofit.non_linear.status.states import RunState + +# Scheduler outcomes outside the run-state vocabulary. +CANCELLED = "cancelled" +REQUEUED = "requeued" +PREEMPTED = "preempted" +UNKNOWN = "unknown" + +_SLURM_STATES = { + "PENDING": RunState.PENDING, + "CONFIGURING": RunState.PENDING, + "REQUEUE_HOLD": RunState.PENDING, + "REQUEUE_FED": RunState.PENDING, + "RESIZING": RunState.RUNNING, + "RUNNING": RunState.RUNNING, + "COMPLETING": RunState.RUNNING, + "SUSPENDED": RunState.RUNNING, + "STOPPED": RunState.RUNNING, + "SIGNALING": RunState.RUNNING, + "STAGE_OUT": RunState.RUNNING, + "COMPLETED": RunState.COMPLETED, + "FAILED": RunState.FAILED, + "TIMEOUT": RunState.TIMEOUT, + "DEADLINE": RunState.TIMEOUT, + "OUT_OF_MEMORY": RunState.OOM, + "NODE_FAIL": RunState.NODE_FAIL, + "BOOT_FAIL": RunState.NODE_FAIL, + "CANCELLED": CANCELLED, + "REVOKED": CANCELLED, + "SPECIAL_EXIT": RunState.FAILED, + "REQUEUED": REQUEUED, + "PREEMPTED": PREEMPTED, +} + +_SQUEUE_COMPACT = { + "PD": "PENDING", + "R": "RUNNING", + "CG": "COMPLETING", + "CF": "CONFIGURING", + "S": "SUSPENDED", + "ST": "STOPPED", + "CD": "COMPLETED", + "F": "FAILED", + "TO": "TIMEOUT", + "OOM": "OUT_OF_MEMORY", + "NF": "NODE_FAIL", + "BF": "BOOT_FAIL", + "CA": "CANCELLED", + "PR": "PREEMPTED", + "RQ": "REQUEUED", + "RH": "REQUEUE_HOLD", + "SE": "SPECIAL_EXIT", + "DL": "DEADLINE", +} + + +def run_state_of(scheduler_state: Optional[str]): + """ + The run state (or ``cancelled`` / ``requeued`` / ``preempted`` / ``unknown``) of a + Slurm state, long (``CANCELLED by 16309``) or compact (``CA``). + """ + if not scheduler_state: + return UNKNOWN + word = scheduler_state.strip().split()[0].rstrip("+").upper() + word = _SQUEUE_COMPACT.get(word, word) + return _SLURM_STATES.get(word, UNKNOWN) + + +@dataclass +class TaskRecord: + """ + One scheduler task (an array task or a plain job), as accounting reports it. + Memory is in MB, times in seconds. + """ + + job_id: str + task_id: Optional[int] = None + name: Optional[str] = None + scheduler_state: Optional[str] = None + exit_code: Optional[str] = None + partition: Optional[str] = None + node: Optional[str] = None + submit: Optional[datetime] = None + start: Optional[datetime] = None + end: Optional[datetime] = None + elapsed_s: Optional[int] = None + time_limit_s: Optional[int] = None + cpus: Optional[int] = None + req_mem_mb: Optional[float] = None + max_rss_mb: Optional[float] = None + total_cpu_s: Optional[float] = None + workdir: Optional[str] = None + req_tres: Optional[str] = None + alloc_tres: Optional[str] = None + raw_job_id: Optional[str] = None + attempts: int = 1 + + @property + def key(self) -> str: + return self.job_id if self.task_id is None else f"{self.job_id}_{self.task_id}" + + @property + def state(self): + return run_state_of(self.scheduler_state) + + @property + def gres(self) -> Optional[str]: + """ + The GPU part of the request, read from TRES (never from the partition name). + """ + for tres in (self.req_tres, self.alloc_tres): + if tres and "gres/gpu" in tres: + return ",".join( + item for item in tres.split(",") if item.startswith("gres/gpu") + ) + return None + + @property + def cpu_efficiency(self) -> Optional[float]: + if not self.elapsed_s or not self.cpus or self.total_cpu_s is None: + return None + return self.total_cpu_s / (self.elapsed_s * self.cpus) + + @property + def cpu_hours(self) -> float: + if not self.elapsed_s or not self.cpus: + return 0.0 + return self.elapsed_s * self.cpus / 3600.0 + + +@dataclass +class PartitionCapacity: + """ + A snapshot of one partition: CPUs by state, and how many jobs are pending on it + requesting a GPU (``None`` when not measured). + """ + + name: str + cpus_idle: int = 0 + cpus_allocated: int = 0 + cpus_other: int = 0 + cpus_total: int = 0 + nodes: int = 0 + nodes_unreachable: int = 0 + free_mem_mb: Optional[float] = None + pending_gpu_jobs: Optional[int] = None + + +@dataclass +class SubmitRequest: + """ + One array for one stage: campaign indices ``indices`` (a chunk starting at + ``offset``), dependencies on upstream arrays as ``{job_id: kind}``. + """ + + campaign: str + stage: str + script: str + offset: int + indices: List[int] + resources: "object" + project_root: str + manifest: Optional[str] = None + env: Dict[str, str] = field(default_factory=dict) + after: Dict[str, str] = field(default_factory=dict) + log_dir: Optional[str] = None + + @property + def task_ids(self) -> List[int]: + return [index - self.offset for index in self.indices] + + @property + def job_name(self) -> str: + return f"{self.campaign}_{self.stage}" + + +@runtime_checkable +class Scheduler(Protocol): + def capacity( + self, partitions: Optional[Sequence[str]] = None + ) -> Dict[str, PartitionCapacity]: + ... + + def query(self, job_ids: Sequence[str]) -> List[TaskRecord]: + ... + + def submit(self, request: SubmitRequest) -> str: + ... + + def update_pending(self, job_id: str, **fields) -> None: + ... + + def cancel(self, job_id: str) -> None: + ... diff --git a/autofit/campaign/scheduler/fake.py b/autofit/campaign/scheduler/fake.py new file mode 100644 index 000000000..f948512ad --- /dev/null +++ b/autofit/campaign/scheduler/fake.py @@ -0,0 +1,170 @@ +""" +``FakeScheduler``: an in-memory scheduler for tests, with scripted outcomes. + +Submitted tasks wait as ``PENDING`` until ``run()``, which settles them in submit order: +a task whose dependency can never be met is cancelled (as ``--kill-on-invalid-dep=yes`` +does), every other task takes the outcome its script gives. It lints like the Slurm +adapter, using its own ``capacity`` snapshot. +""" +import copy +from datetime import datetime, timedelta +from typing import Callable, Dict, List, Optional, Sequence, Union + +from autofit.campaign.lint import check_partitions, split_partitions +from autofit.campaign.scheduler.base import ( + PartitionCapacity, + SubmitRequest, + TaskRecord, +) + +Outcome = Dict[str, object] +OutcomeScript = Union[Callable[[str, int, int], Optional[Outcome]], Dict, None] + +_EPOCH = datetime(2026, 1, 1) + + +class FakeScheduler: + def __init__( + self, + outcomes: OutcomeScript = None, + capacity: Optional[Dict[str, PartitionCapacity]] = None, + site=None, + ): + """ + Parameters + ---------- + outcomes + ``outcomes(stage, index, attempt)`` or a dict keyed ``(stage, index)`` + giving the task's ``TaskRecord`` fields, e.g. ``{"scheduler_state": + "TIMEOUT", "elapsed_s": 3600, "max_rss_mb": 900}``. Missing fields default + to a completed task using a tenth of its time and half its memory. + capacity + The ``capacity()`` snapshot, including ``pending_gpu_jobs``. + """ + self.outcomes = outcomes + self._capacity = capacity or {} + self.site = site + self.requests: Dict[str, SubmitRequest] = {} + self.records: Dict[str, List[TaskRecord]] = {} + self.updates: List[tuple] = [] + self.cancelled: List[str] = [] + self.attempts: Dict[tuple, int] = {} + self._next_id = 1000 + + def capacity(self, partitions: Optional[Sequence[str]] = None) -> Dict[str, PartitionCapacity]: + if partitions is None: + return dict(self._capacity) + return {name: self._capacity[name] for name in partitions if name in self._capacity} + + def _pending_gpu(self, names) -> Optional[int]: + known = [self._capacity[name].pending_gpu_jobs for name in names if name in self._capacity] + # Unmeasured counts as unknown, which refuses the lint's exemption. + if not known or any(value is None for value in known): + return None + return sum(known) + + def submit(self, request: SubmitRequest) -> str: + resources = request.resources + names = split_partitions(resources.partition) + check_partitions( + names, + cpus=resources.cpus, + gres=resources.gres, + throttle=resources.throttle, + n_tasks=len(request.indices), + pending_gpu_jobs=self._pending_gpu(names) if names else 0, + site=self.site, + ).raise_for_refusal() + + job_id = str(self._next_id) + self._next_id += 1 + self.requests[job_id] = request + self.records[job_id] = [ + TaskRecord( + job_id=job_id, + task_id=task_id, + name=request.job_name, + scheduler_state="PENDING", + partition=resources.partition, + cpus=resources.cpus, + req_mem_mb=resources.mem_mb, + time_limit_s=resources.time_s, + workdir=str(request.project_root), + req_tres=f"cpu={resources.cpus}" + + (f",gres/{resources.gres.replace(':', '=')}" if resources.gres else ""), + ) + for task_id in request.task_ids + ] + return job_id + + def _outcome(self, stage: str, index: int, attempt: int) -> Outcome: + if callable(self.outcomes): + outcome = self.outcomes(stage, index, attempt) + elif isinstance(self.outcomes, dict): + outcome = self.outcomes.get((stage, index)) + else: + outcome = None + return dict(outcome or {}) + + def _dependency_met(self, request: SubmitRequest, task_id: int) -> Optional[bool]: + for upstream_id, kind in request.after.items(): + upstream = self.records.get(str(upstream_id), []) + if kind == "aftercorr": + matching = [r for r in upstream if r.task_id == task_id] + if not matching or matching[0].scheduler_state != "COMPLETED": + return False + elif kind == "afterok": + if any(r.scheduler_state != "COMPLETED" for r in upstream): + return False + return True + + def run(self): + """ + Settle every pending task, in submit order. + """ + for job_id, records in self.records.items(): + request = self.requests[job_id] + for record in records: + if record.scheduler_state != "PENDING": + continue + index = request.offset + record.task_id + key = (request.stage, index) + self.attempts[key] = self.attempts.get(key, 0) + 1 + if not self._dependency_met(request, record.task_id): + record.scheduler_state = "CANCELLED" + record.elapsed_s = 0 + continue + limit = record.time_limit_s or 3600 + outcome = self._outcome(request.stage, index, self.attempts[key]) + record.scheduler_state = str(outcome.pop("scheduler_state", "COMPLETED")) + record.elapsed_s = int(outcome.pop("elapsed_s", limit // 10)) + record.max_rss_mb = outcome.pop( + "max_rss_mb", (record.req_mem_mb or 1000) / 2 + ) + record.total_cpu_s = outcome.pop( + "total_cpu_s", record.elapsed_s * (record.cpus or 1) * 0.9 + ) + record.start = _EPOCH + record.end = _EPOCH + timedelta(seconds=record.elapsed_s) + record.exit_code = "0:0" if record.scheduler_state == "COMPLETED" else "1:0" + for name, value in outcome.items(): + setattr(record, name, value) + + def query(self, job_ids: Sequence[str]) -> List[TaskRecord]: + return [copy.copy(record) for job_id in job_ids for record in self.records.get(str(job_id), [])] + + def update_pending(self, job_id: str, **fields) -> None: + self.updates.append((str(job_id), fields)) + for record in self.records.get(str(job_id), []): + if record.scheduler_state != "PENDING": + continue + if "time" in fields: + record.time_limit_s = int(fields["time"]) + if "mem_per_cpu_mb" in fields and record.cpus: + record.req_mem_mb = float(fields["mem_per_cpu_mb"]) * record.cpus + + def cancel(self, job_id: str) -> None: + self.cancelled.append(str(job_id)) + for record in self.records.get(str(job_id), []): + if record.scheduler_state == "PENDING": + record.scheduler_state = "CANCELLED" diff --git a/autofit/campaign/scheduler/local.py b/autofit/campaign/scheduler/local.py new file mode 100644 index 000000000..396059bed --- /dev/null +++ b/autofit/campaign/scheduler/local.py @@ -0,0 +1,188 @@ +""" +``LocalScheduler``: run a campaign's tasks as subprocesses on this machine (a laptop or +a single workstation), with the same environment a Slurm array task sees. + +Each task runs the stage's command in ``project_root`` with ``CAMPAIGN``, ``STAGE``, +``CAMPAIGN_OFFSET``, ``CAMPAIGN_INDEX``, ``SLURM_ARRAY_TASK_ID`` (the task id inside +its array, so Slurm-style scripts work unchanged) and ``MANIFEST`` set. ``aftercorr`` +waits for the same task of the upstream array and cancels the task if it did not +complete; ``afterok`` waits for the whole upstream array. +""" +import os +import shlex +import subprocess +import sys +import threading +import time +from concurrent.futures import Future, ThreadPoolExecutor +from datetime import datetime +from pathlib import Path +from typing import Dict, List, Optional, Sequence + +from autofit.campaign.scheduler.base import ( + PartitionCapacity, + SubmitRequest, + TaskRecord, +) + + +class LocalScheduler: + def __init__(self, max_workers: Optional[int] = None): + """ + Parameters + ---------- + max_workers + Tasks run at once (default: the number of CPUs). + """ + self.max_workers = max_workers or os.cpu_count() or 1 + self._pool = ThreadPoolExecutor(max_workers=self.max_workers) + self._lock = threading.Lock() + self._next_id = 1 + self.records: Dict[str, List[TaskRecord]] = {} + self._futures: Dict[str, Dict[int, Future]] = {} + self._processes: Dict[str, Dict[int, subprocess.Popen]] = {} + + def capacity(self, partitions: Optional[Sequence[str]] = None) -> Dict[str, PartitionCapacity]: + running = sum( + 1 + for records in self.records.values() + for record in records + if record.scheduler_state == "RUNNING" + ) + total = os.cpu_count() or 1 + return { + "local": PartitionCapacity( + name="local", + cpus_allocated=running, + cpus_idle=max(total - running, 0), + cpus_total=total, + nodes=1, + pending_gpu_jobs=0, + ) + } + + def _command(self, request: SubmitRequest) -> List[str]: + command = shlex.split(request.script) + if command and command[0].endswith(".py"): + command = [sys.executable] + command + return command + + def _wait_upstream(self, request: SubmitRequest, task_id: int) -> bool: + for upstream_id, kind in request.after.items(): + futures = self._futures.get(str(upstream_id), {}) + if kind == "aftercorr": + future = futures.get(task_id) + if future is None or future.result() != "COMPLETED": + return False + elif kind == "afterok": + if any(future.result() != "COMPLETED" for future in futures.values()): + return False + else: + for future in futures.values(): + future.result() + return True + + def _run_task(self, job_id: str, request: SubmitRequest, record: TaskRecord) -> str: + if not self._wait_upstream(request, record.task_id): + record.scheduler_state = "CANCELLED" + return record.scheduler_state + if record.scheduler_state == "CANCELLED": + return record.scheduler_state + + env = dict(os.environ) + env.update(request.env) + env.update( + { + "PYTHONUNBUFFERED": "1", + "CAMPAIGN": request.campaign, + "STAGE": request.stage, + "CAMPAIGN_OFFSET": str(request.offset), + "CAMPAIGN_INDEX": str(request.offset + record.task_id), + "SLURM_ARRAY_TASK_ID": str(record.task_id), + } + ) + if request.manifest: + env["MANIFEST"] = request.manifest + + log = None + if request.log_dir: + log_dir = Path(request.project_root) / request.log_dir + log_dir.mkdir(parents=True, exist_ok=True) + log = open(log_dir / f"{request.job_name}_{job_id}_{record.task_id}.out", "w") + + record.start = datetime.now() + record.scheduler_state = "RUNNING" + started = time.monotonic() + try: + process = subprocess.Popen( + self._command(request), + cwd=request.project_root, + env=env, + stdout=log or subprocess.DEVNULL, + stderr=subprocess.STDOUT if log else subprocess.DEVNULL, + ) + self._processes.setdefault(job_id, {})[record.task_id] = process + timeout = request.resources.time_s + try: + code = process.wait(timeout=timeout) + record.scheduler_state = "COMPLETED" if code == 0 else "FAILED" + record.exit_code = f"{code}:0" + except subprocess.TimeoutExpired: + process.terminate() + process.wait() + record.scheduler_state = "TIMEOUT" + finally: + if log: + log.close() + record.end = datetime.now() + record.elapsed_s = int(time.monotonic() - started) + return record.scheduler_state + + def submit(self, request: SubmitRequest) -> str: + with self._lock: + job_id = str(self._next_id) + self._next_id += 1 + records = [ + TaskRecord( + job_id=job_id, + task_id=task_id, + name=request.job_name, + scheduler_state="PENDING", + partition="local", + cpus=request.resources.cpus, + req_mem_mb=request.resources.mem_mb, + time_limit_s=request.resources.time_s, + workdir=str(request.project_root), + ) + for task_id in request.task_ids + ] + self.records[job_id] = records + self._futures[job_id] = { + record.task_id: self._pool.submit(self._run_task, job_id, request, record) + for record in records + } + return job_id + + def wait(self, job_ids: Optional[Sequence[str]] = None): + """ + Block until every task of ``job_ids`` (default: all) has finished. + """ + for job_id in job_ids or list(self._futures): + for future in self._futures.get(str(job_id), {}).values(): + future.result() + + def query(self, job_ids: Sequence[str]) -> List[TaskRecord]: + return [record for job_id in job_ids for record in self.records.get(str(job_id), [])] + + def update_pending(self, job_id: str, **fields) -> None: + for record in self.records.get(str(job_id), []): + if record.scheduler_state == "PENDING" and "time" in fields: + record.time_limit_s = int(fields["time"]) + + def cancel(self, job_id: str) -> None: + for record in self.records.get(str(job_id), []): + if record.scheduler_state == "PENDING": + record.scheduler_state = "CANCELLED" + for process in self._processes.get(str(job_id), {}).values(): + if process.poll() is None: + process.terminate() diff --git a/autofit/campaign/scheduler/sacct.py b/autofit/campaign/scheduler/sacct.py new file mode 100644 index 000000000..9fbfcf410 --- /dev/null +++ b/autofit/campaign/scheduler/sacct.py @@ -0,0 +1,214 @@ +""" +Parsing ``sacct --parsable2 --units=M`` output into ``TaskRecord``s. + +``sacct`` prints one *allocation* line per job or array task and one line per *step* +(``.batch``, ``.extern``, ``.0``...). Memory high-water marks (``MaxRSS``) are only on +the step lines, so each step is folded into the allocation it belongs to (the last +allocation line with the same ``JobID`` above it, which is how ``sacct`` orders its +output, including under ``--duplicates``). + +With ``--duplicates`` a requeued task appears once per attempt; ``parse`` keeps every +attempt (they all used CPU time) and ``latest_per_task`` reduces them to one record per +task with ``attempts`` counted. +""" +import csv +import io +from datetime import datetime +from typing import Dict, Iterable, List, Optional, Sequence + +from autofit.campaign.scheduler.base import TaskRecord +from autofit.campaign.units import memory_per, parse_memory_mb, parse_time + +FIELDS = ( + "JobID", + "JobIDRaw", + "JobName%60", + "Partition", + "State", + "ExitCode", + "Submit", + "Start", + "End", + "Elapsed", + "Timelimit", + "AllocCPUS", + "ReqMem", + "MaxRSS", + "AveRSS", + "TotalCPU", + "CPUTime", + "NNodes", + "NodeList", + "ReqTRES%80", + "AllocTRES%80", + "WorkDir%200", +) + + +def field_names(fields: Sequence[str] = FIELDS) -> List[str]: + return [name.split("%", 1)[0] for name in fields] + + +def _datetime(value: Optional[str]) -> Optional[datetime]: + if not value or value in ("Unknown", "None", "N/A"): + return None + try: + return datetime.fromisoformat(value) + except ValueError: + return None + + +def _int(value: Optional[str]) -> Optional[int]: + try: + return int(value) if value not in (None, "") else None + except ValueError: + return None + + +def _memory(value: Optional[str]) -> Optional[float]: + try: + return parse_memory_mb(value) if value else None + except ValueError: + return None + + +def _seconds(value: Optional[str]) -> Optional[float]: + try: + return parse_time(value) if value else None + except ValueError: + return None + + +def _total_cpu(value: Optional[str]) -> Optional[float]: + # TotalCPU carries a fractional part ("12-01:02:31" or "00:01.234"). + if not value: + return None + whole, _, fraction = value.partition(".") + try: + seconds = float(parse_time(whole) if ":" in whole else int(whole)) + return seconds + (float(f"0.{fraction}") if fraction.isdigit() else 0.0) + except ValueError: + return None + + +def _split_job_id(job_id: str): + """ + ``(base, step)`` of a ``JobID``: ``"342027_1.batch"`` → ``("342027_1", "batch")``. + """ + if "." in job_id and not job_id.endswith("]"): + base, step = job_id.split(".", 1) + return base, step + return job_id, None + + +def _array_parts(base: str): + """ + ``(job_id, task_id)``: ``"342027_1"`` → ``("342027", 1)``; a plain job or a pending + array placeholder (``"344645_[586,588-594]"``) has no task id. + """ + if "_" in base: + job, task = base.split("_", 1) + if task.isdigit(): + return job, int(task) + return job, None + return base, None + + +def _rows(text: str, fields: Optional[Sequence[str]] = None) -> Iterable[Dict[str, str]]: + lines = [line for line in text.splitlines() if line.strip()] + if not lines: + return [] + if fields is None: + reader = csv.DictReader(io.StringIO("\n".join(lines)), delimiter="|", quoting=csv.QUOTE_NONE) + return list(reader) + names = field_names(fields) + reader = csv.reader(io.StringIO("\n".join(lines)), delimiter="|", quoting=csv.QUOTE_NONE) + return [dict(zip(names, row)) for row in reader] + + +def parse(text: str, fields: Optional[Sequence[str]] = None) -> List[TaskRecord]: + """ + Every allocation (job or array-task attempt) in ``sacct --parsable2`` output, with + its steps folded in. + + Parameters + ---------- + text + The output; with a header line when ``fields`` is ``None`` (``sacct`` without + ``--noheader``), else exactly the columns of ``fields``. + """ + records: List[TaskRecord] = [] + latest: Dict[str, TaskRecord] = {} + for row in _rows(text, fields): + job_id = row.get("JobID", "") + if not job_id: + continue + base, step = _split_job_id(job_id) + if step is not None: + record = latest.get(base) + if record is None: + continue + rss = _memory(row.get("MaxRSS")) + if rss is not None: + record.max_rss_mb = max(record.max_rss_mb or 0.0, rss) + if record.total_cpu_s is None: + record.total_cpu_s = _total_cpu(row.get("TotalCPU")) + if not record.node and row.get("NodeList"): + record.node = row["NodeList"] + continue + + job, task = _array_parts(base) + cpus = _int(row.get("AllocCPUS")) + req_mem = row.get("ReqMem") + req_mem_mb = _memory(req_mem) + if req_mem_mb is not None and memory_per(req_mem) == "cpu" and cpus: + req_mem_mb *= cpus + record = TaskRecord( + job_id=job, + task_id=task, + name=row.get("JobName") or None, + scheduler_state=row.get("State") or None, + exit_code=row.get("ExitCode") or None, + partition=row.get("Partition") or None, + node=row.get("NodeList") or None, + submit=_datetime(row.get("Submit")), + start=_datetime(row.get("Start")), + end=_datetime(row.get("End")), + elapsed_s=_seconds(row.get("Elapsed")), + time_limit_s=_seconds(row.get("Timelimit")), + cpus=cpus, + req_mem_mb=req_mem_mb, + max_rss_mb=_memory(row.get("MaxRSS")), + total_cpu_s=_total_cpu(row.get("TotalCPU")), + workdir=row.get("WorkDir") or None, + req_tres=row.get("ReqTRES") or None, + alloc_tres=row.get("AllocTRES") or None, + raw_job_id=row.get("JobIDRaw") or None, + ) + if base.endswith("]"): + # A pending array's placeholder line: one record for the whole range. + record.job_id = base + records.append(record) + latest[base] = record + return records + + +def latest_per_task(records: Iterable[TaskRecord]) -> List[TaskRecord]: + """ + One record per task, the latest attempt (by start, then submit, then position), + with ``attempts`` counting every attempt seen. + """ + by_key: Dict[str, List[TaskRecord]] = {} + for record in records: + by_key.setdefault(record.key, []).append(record) + result = [] + for attempts in by_key.values(): + def order(item): + position, record = item + stamp = record.start or record.submit or datetime.min + return stamp, position + + _, last = max(enumerate(attempts), key=order) + last.attempts = len(attempts) + result.append(last) + return result diff --git a/autofit/campaign/scheduler/slurm.py b/autofit/campaign/scheduler/slurm.py new file mode 100644 index 000000000..40b789cca --- /dev/null +++ b/autofit/campaign/scheduler/slurm.py @@ -0,0 +1,299 @@ +""" +The Slurm adapter: batched ``sacct`` / ``squeue`` / ``sinfo`` queries and ``sbatch`` +script generation. + +Nothing here runs at import. Every command goes through a *runner*, +``runner(arguments, input=None) -> stdout``, which by default calls ``subprocess`` and +in tests is a fake; a login node reached over ssh is just a runner that prefixes +``ssh ``. + +Every ``submit`` and every ``update_pending`` that changes a partition is linted +(``autofit.campaign.lint``) against a live count of GPU jobs pending on the target +partition: CPU-only work never goes on a GPU partition. +""" +import logging +import shlex +import subprocess +from typing import Callable, Dict, List, Optional, Sequence + +from autofit.campaign.lint import check_partitions, split_partitions +from autofit.campaign.scheduler import sacct +from autofit.campaign.scheduler.base import ( + PartitionCapacity, + SubmitRequest, + TaskRecord, +) +from autofit.campaign.units import array_spec, format_time, requests_gpu + +logger = logging.getLogger(__name__) + +Runner = Callable[..., str] + +SIGNAL = "B:USR1@600" + + +def subprocess_runner(arguments: Sequence[str], input: Optional[str] = None) -> str: + completed = subprocess.run( + list(arguments), input=input, capture_output=True, text=True, check=True + ) + return completed.stdout + + +def _export(key: str, value: str) -> str: + # A value naming a variable (``$SLURM_CPUS_PER_TASK``) is expanded by the job. + if "$" in value: + return f'export {key}="{value}"' + return f"export {key}={shlex.quote(value)}" + + +def sbatch_script(request: SubmitRequest) -> str: + """ + The batch script of one array. + + The job runs the stage's command in the background so the batch shell can forward + the ``--signal=B:USR1@600`` pre-timeout warning to it (the fit flushes its + checkpoint and writes ``timeout_pending``); the shell exits with the command's + status. + """ + resources = request.resources + lines = ["#!/bin/bash", f"#SBATCH --job-name={request.job_name}"] + if resources.partition: + lines.append(f"#SBATCH --partition={resources.partition}") + lines.append(f"#SBATCH --array={array_spec(request.task_ids, resources.throttle)}") + lines.append(f"#SBATCH --cpus-per-task={resources.cpus}") + if resources.mem_per_cpu_mb is not None: + lines.append(f"#SBATCH --mem-per-cpu={int(resources.mem_per_cpu_mb)}") + if resources.time_s is not None: + lines.append(f"#SBATCH --time={format_time(resources.time_s)}") + if resources.gres: + lines.append(f"#SBATCH --gres={resources.gres}") + if resources.nice is not None: + lines.append(f"#SBATCH --nice={resources.nice}") + if resources.qos: + lines.append(f"#SBATCH --qos={resources.qos}") + if resources.account: + lines.append(f"#SBATCH --account={resources.account}") + if request.after: + dependency = ",".join(f"{kind}:{job_id}" for job_id, kind in request.after.items()) + lines.append(f"#SBATCH --dependency={dependency}") + lines.append("#SBATCH --kill-on-invalid-dep=yes") + lines.append(f"#SBATCH --signal={SIGNAL}") + if request.log_dir: + lines.append(f"#SBATCH --output={request.log_dir}/%x_%A_%a.out") + lines.append(f"#SBATCH --error={request.log_dir}/%x_%A_%a.err") + + env = { + "PYTHONUNBUFFERED": "1", + "CAMPAIGN": request.campaign, + "STAGE": request.stage, + "CAMPAIGN_OFFSET": str(request.offset), + } + if request.manifest: + env["MANIFEST"] = request.manifest + env.update(request.env) + + command = shlex.split(request.script) + if command and command[0].endswith(".py"): + command = ["python"] + command + + lines += [ + "", + "set -euo pipefail", + *(_export(key, value) for key, value in env.items()), + 'export CAMPAIGN_INDEX=$((CAMPAIGN_OFFSET + SLURM_ARRAY_TASK_ID))', + f"cd {shlex.quote(str(request.project_root))}", + "", + f"{' '.join(shlex.quote(part) for part in command)} &", + "child=$!", + "trap 'kill -USR1 \"$child\" 2>/dev/null || true' USR1", + "trap 'kill -TERM \"$child\" 2>/dev/null || true' TERM", + "status=0", + 'until wait "$child"; do', + " status=$?", + ' kill -0 "$child" 2>/dev/null || break', + "done", + 'exit "$status"', + "", + ] + return "\n".join(lines) + + +class SlurmScheduler: + def __init__(self, runner: Optional[Runner] = None, user: Optional[str] = None, site=None): + """ + A ``Scheduler`` over the Slurm command-line tools. + + Parameters + ---------- + runner + ``runner(arguments, input=None) -> stdout``; ``subprocess`` by default. + user + Whose jobs ``squeue`` lists (default: the current user, ``--me``). + site + The ``autonerves.site_config.Site`` naming the GPU-only partitions. + """ + self.runner = runner or subprocess_runner + self.user = user + self.site = site + + def _run(self, arguments: List[str], input: Optional[str] = None) -> str: + logger.debug(f"slurm: {' '.join(arguments)}") + return self.runner(arguments, input=input) if input is not None else self.runner(arguments) + + def capacity(self, partitions: Optional[Sequence[str]] = None) -> Dict[str, PartitionCapacity]: + """ + CPUs by state per partition (one ``sinfo``) and GPU-requesting jobs pending on + each (one ``squeue``). + """ + arguments = ["sinfo", "-N", "-h", "-o", "%P|%n|%t|%C|%m|%e"] + if partitions: + arguments += ["-p", ",".join(partitions)] + result: Dict[str, PartitionCapacity] = {} + seen = set() + for line in self._run(arguments).splitlines(): + parts = line.strip().split("|") + if len(parts) < 4: + continue + partition, node, state, cpus = parts[:4] + partition = partition.rstrip("*") + if (partition, node) in seen: + continue + seen.add((partition, node)) + capacity = result.setdefault(partition, PartitionCapacity(name=partition)) + capacity.nodes += 1 + # `*` marks an unreachable node: its CPUs are not free. + if state.endswith("*") or state.startswith(("down", "drain", "fail")): + capacity.nodes_unreachable += 1 + try: + allocated, idle, other, total = (int(v) for v in cpus.split("/")) + except ValueError: + continue + if state.endswith("*"): + other, idle = other + idle, 0 + capacity.cpus_allocated += allocated + capacity.cpus_idle += idle + capacity.cpus_other += other + capacity.cpus_total += total + if len(parts) > 5 and parts[5].strip().isdigit(): + capacity.free_mem_mb = (capacity.free_mem_mb or 0) + float(parts[5]) + + pending = self._run(["squeue", "-h", "-r", "-t", "PD", "-o", "%P|%b|%i"]) + counts: Dict[str, int] = {name: 0 for name in result} + for line in pending.splitlines(): + parts = line.strip().split("|") + if len(parts) < 2: + continue + if requests_gpu(parts[1]): + for partition in split_partitions(parts[0]): + counts[partition] = counts.get(partition, 0) + 1 + for name, capacity in result.items(): + capacity.pending_gpu_jobs = counts.get(name, 0) + return result + + def pending_gpu_jobs(self, partitions: Sequence[str]) -> int: + capacity = self.capacity(partitions) + return sum(capacity[name].pending_gpu_jobs or 0 for name in partitions if name in capacity) + + def query(self, job_ids: Sequence[str]) -> List[TaskRecord]: + """ + Every allocation of ``job_ids`` from one ``sacct`` call (steps folded in, + attempts kept; see ``sacct.latest_per_task``). + """ + if not job_ids: + return [] + arguments = [ + "sacct", + "-j", + ",".join(str(job_id) for job_id in job_ids), + "--parsable2", + "--noheader", + "--units=M", + "--duplicates", + "-o", + ",".join(sacct.FIELDS), + ] + return sacct.parse(self._run(arguments), fields=sacct.FIELDS) + + def queue(self) -> List[Dict[str, str]]: + """ + The user's queued tasks, one row per array task (``squeue -r``). + """ + arguments = ["squeue", "-h", "-r", "-o", "%i|%F|%K|%j|%T|%r|%P|%C|%m|%b|%Q|%y"] + arguments += ["-u", self.user] if self.user else ["--me"] + names = ["id", "array_job", "array_task", "name", "state", "reason", "partition", + "cpus", "min_memory", "gres", "priority", "nice"] + return [ + dict(zip(names, line.strip().split("|"))) + for line in self._run(arguments).splitlines() + if line.strip() + ] + + def check(self, request: SubmitRequest): + resources = request.resources + names = split_partitions(resources.partition) + pending = self.pending_gpu_jobs(names) if names else 0 + verdict = check_partitions( + names, + cpus=resources.cpus, + gres=resources.gres, + throttle=resources.throttle, + n_tasks=len(request.indices), + pending_gpu_jobs=pending, + site=self.site, + ) + verdict.raise_for_refusal() + return verdict + + def submit(self, request: SubmitRequest) -> str: + """ + Lint, then ``sbatch --parsable`` the array's script on stdin; returns the job + id. + + Raises + ------ + PartitionLintError + If the request would put CPU-only work on a GPU partition. + """ + self.check(request) + output = self._run(["sbatch", "--parsable"], input=sbatch_script(request)) + return output.strip().split(";")[0] + + def update_pending(self, job_id: str, **fields) -> None: + """ + ``scontrol update`` of a job's still-pending tasks: ``time`` (Slurm time or + seconds), ``mem_per_cpu_mb``, ``throttle``, ``nice``, ``partition`` (linted). + Started tasks keep their request (Slurm splits them off). + """ + mapping = { + "time": "TimeLimit", + "mem_per_cpu_mb": "MinMemoryCPU", + "throttle": "ArrayTaskThrottle", + "nice": "Nice", + "partition": "Partition", + } + # `cpus` and `gres` describe the job for the lint; they are not updated. + cpus = fields.pop("cpus", None) + gres = fields.pop("gres", None) + unknown = set(fields) - set(mapping) + if unknown: + raise ValueError(f"Cannot update {sorted(unknown)}; allowed {sorted(mapping)}.") + if "partition" in fields: + names = split_partitions(fields["partition"]) + verdict = check_partitions( + names, + cpus=cpus, + gres=gres, + throttle=fields.get("throttle"), + pending_gpu_jobs=self.pending_gpu_jobs(names), + site=self.site, + ) + verdict.raise_for_refusal() + arguments = ["scontrol", "update", f"JobId={job_id}"] + for key, value in fields.items(): + if key == "time" and not isinstance(value, str): + value = format_time(value) + arguments.append(f"{mapping[key]}={int(value) if key == 'mem_per_cpu_mb' else value}") + self._run(arguments) + + def cancel(self, job_id: str) -> None: + self._run(["scancel", str(job_id)]) diff --git a/autofit/campaign/spec.py b/autofit/campaign/spec.py new file mode 100644 index 000000000..f64563c6c --- /dev/null +++ b/autofit/campaign/spec.py @@ -0,0 +1,467 @@ +""" +``campaign.yaml``: the single declaration of a modelling campaign. + +.. code-block:: yaml + + campaign: euclid_dr1_next10k + project_root: /mnt/ral/jnightin/euclid_dr1 + manifest: hpc/run_manifests/next10k.txt # 0-based line = campaign index + manifest_fields: [sample, tile] # optional names for a line's fields + chunk: 1000 # tasks per array, ≤ the site's max_array_size + site: ral # autonerves site layer + scheduler: slurm # slurm (default) or local + status_root: output # where fits write *.status.json + log_dir: hpc/logs + stages: + vis_lp: + script: scripts/initial_lens_model.py --stage=vis_lp + env: {JAX_PLATFORMS: cpu} + resources: {cpus: 4, mem_per_cpu_mb: 3072, time: "04:00:00", + partition: ral, throttle: 400, nice: 1000} + success: {zip_complete: "output/{sample}/{tile}/initial_lens_model/vis_lp", + members: [files/samples_summary.json]} + retry: {timeout: {time: "p95+25%", max: 1}, oom: {mem: "peak*1.5", max: 1}} + vis_pix: + after: {vis_lp: aftercorr} + script: scripts/pixelization.py + +Each manifest line is one object; its 0-based line number is its *campaign index* and +is the same in every stage. The manifest is split into arrays of at most ``chunk`` +tasks; inside an array the Slurm task id runs from 0 and the batch script exports +``CAMPAIGN_INDEX = CAMPAIGN_OFFSET + SLURM_ARRAY_TASK_ID``, so per-index dependencies +(``aftercorr``) line up between stages chunk by chunk. + +``retry`` is a policy for *proposals* only: the collector proposes resubmits sized +from the failed tasks' own measurements and a human approves every one. +""" +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import yaml + +from autofit.campaign.lint import check_partitions, split_partitions +from autofit.campaign.units import parse_memory_mb, parse_time + +CAMPAIGN_FIELDS = { + "campaign", + "project_root", + "manifest", + "manifest_fields", + "manifest_separator", + "chunk", + "site", + "scheduler", + "status_root", + "log_dir", + "stages", + "note", +} +STAGE_FIELDS = {"script", "env", "resources", "after", "seed", "success", "retry", "note"} +RESOURCE_FIELDS = { + "cpus", + "mem_per_cpu_mb", + "mem_mb", + "time", + "partition", + "gres", + "throttle", + "nice", + "qos", + "account", +} +DEPENDENCY_KINDS = {"aftercorr", "afterok", "afterany"} +RETRY_CLASSES = { + "timeout", + "oom", + "node_fail", + "exception", + "bad_input", + "missing_seed", + "upstream_incomplete", +} +SCHEDULERS = {"slurm", "local"} + + +class CampaignException(Exception): + pass + + +@dataclass +class Resources: + cpus: int = 1 + mem_per_cpu_mb: Optional[int] = None + time_s: Optional[int] = None + partition: Optional[str] = None + gres: Optional[str] = None + throttle: Optional[int] = None + nice: Optional[int] = None + qos: Optional[str] = None + account: Optional[str] = None + + @classmethod + def from_dict(cls, raw: Optional[Dict]) -> "Resources": + raw = dict(raw or {}) + cpus = int(raw.get("cpus", 1)) + mem_per_cpu = raw.get("mem_per_cpu_mb") + if mem_per_cpu is None and raw.get("mem_mb") is not None: + # One knob: memory follows the CPU count (`--mem-per-cpu`), so a total is + # turned into a per-CPU value. + mem_per_cpu = -(-parse_memory_mb(raw["mem_mb"]) // cpus) + return cls( + cpus=cpus, + mem_per_cpu_mb=None + if mem_per_cpu is None + else int(parse_memory_mb(mem_per_cpu)), + time_s=parse_time(raw.get("time")), + partition=raw.get("partition"), + gres=raw.get("gres"), + throttle=None if raw.get("throttle") is None else int(raw["throttle"]), + nice=None if raw.get("nice") is None else int(raw["nice"]), + qos=raw.get("qos"), + account=raw.get("account"), + ) + + @property + def mem_mb(self) -> Optional[int]: + if self.mem_per_cpu_mb is None: + return None + return self.mem_per_cpu_mb * self.cpus + + +@dataclass +class Success: + zip_complete: Optional[str] = None + members: List[str] = field(default_factory=list) + + +@dataclass +class Stage: + name: str + script: str + resources: Resources + env: Dict[str, str] = field(default_factory=dict) + after: Dict[str, str] = field(default_factory=dict) + seed: Optional[str] = None + success: Success = field(default_factory=Success) + retry: Dict[str, Dict] = field(default_factory=dict) + + +class Campaign: + def __init__(self, raw: Dict, source: Optional[Path] = None): + """ + A ``campaign.yaml`` (``Campaign.load``). Construction does not validate; call + ``validate`` (``load`` does) for the list of problems. + """ + self.raw = raw + self.source = source + self.name = raw.get("campaign") + root = raw.get("project_root") + if root is None and source is not None: + root = Path(source).parent + self.project_root = Path(root) if root is not None else Path.cwd() + self.manifest_path = raw.get("manifest") + self.manifest_fields = list(raw.get("manifest_fields") or []) + self.manifest_separator = raw.get("manifest_separator") + self.chunk = raw.get("chunk") + self.site_name = raw.get("site") + self.scheduler = raw.get("scheduler", "slurm") + self.status_root = raw.get("status_root", "output") + self.log_dir = raw.get("log_dir", "hpc/logs") + self.stages: Dict[str, Stage] = {} + for name, stage in (raw.get("stages") or {}).items(): + stage = stage or {} + if not isinstance(stage, dict): + continue + success = stage.get("success") or {} + self.stages[name] = Stage( + name=name, + script=stage.get("script", ""), + resources=Resources.from_dict(stage.get("resources")), + env={k: str(v) for k, v in (stage.get("env") or {}).items()}, + after=dict(stage.get("after") or {}), + seed=stage.get("seed"), + success=Success( + zip_complete=success.get("zip_complete"), + members=list(success.get("members") or []), + ), + retry=dict(stage.get("retry") or {}), + ) + self._site = None + self._manifest = None + + @classmethod + def load(cls, path: Union[str, Path], site=None) -> "Campaign": + """ + Load and validate the campaign file at ``path``. + + Parameters + ---------- + site + The ``autonerves.site_config.Site`` to validate against; by default the + file's ``site:`` is loaded through the autonerves site layer. + + Raises + ------ + CampaignException + If the file cannot be read or fails ``validate``. + """ + path = Path(path) + try: + raw = yaml.safe_load(path.read_text()) + except OSError as exception: + raise CampaignException(f"Campaign file {path} cannot be read: {exception}") + except yaml.YAMLError as exception: + raise CampaignException(f"Campaign file {path} is not valid YAML: {exception}") + if not isinstance(raw, dict): + raise CampaignException(f"Campaign file {path} must be a mapping.") + campaign = cls(raw, source=path) + if site is not None: + campaign._site = site + errors = campaign.validate() + if errors: + raise CampaignException( + f"Campaign file {path} is invalid:\n " + "\n ".join(errors) + ) + return campaign + + @property + def site(self): + """ + The site layer named by ``site:`` (``None`` when the file names none). + """ + if self._site is None and self.site_name: + from autonerves import conf + from autonerves.site_config import load_site, site_path + + paths = getattr(conf.instance, "paths", ()) + self._site = load_site(site_path(str(self.site_name), config_paths=paths)) + return self._site + + @property + def directory(self) -> Path: + """ + Where the campaign keeps its job ledger and digest: + ``/campaign//``. + """ + return self.project_root / "campaign" / str(self.name) + + def path(self, value: Union[str, Path]) -> Path: + value = Path(value) + return value if value.is_absolute() else self.project_root / value + + def manifest(self) -> List[str]: + """ + The manifest's lines (blank lines and ``#`` comments are not objects but keep + their line number, so an index is always a line number). + """ + if self._manifest is None: + text = self.path(self.manifest_path).read_text() + self._manifest = text.splitlines() + return self._manifest + + def indices(self) -> List[int]: + return [ + index + for index, line in enumerate(self.manifest()) + if line.strip() and not line.lstrip().startswith("#") + ] + + def fields(self, index: int) -> Dict[str, str]: + """ + The template fields of manifest line ``index``: ``index``, ``line`` and, with + ``manifest_fields``, each named field of the split line. + """ + line = self.manifest()[index].strip() + values = {"index": str(index), "line": line} + if self.manifest_fields: + parts = line.split(self.manifest_separator) + for name, part in zip(self.manifest_fields, parts): + values[name] = part.strip() + return values + + def chunk_size(self) -> int: + size = self.chunk + maximum = self.max_array_size() + if size is None: + size = maximum or 1000 + return int(size) + + def max_array_size(self) -> Optional[int]: + site = self._site_or_none() + if site is None: + return None + sizes = [ + site.partitions[name].max_array_size + for stage in self.stages.values() + for name in split_partitions(stage.resources.partition) + if name in site.partitions + and site.partitions[name].max_array_size is not None + ] + return min(sizes) if sizes else None + + def chunks(self, indices: Optional[List[int]] = None) -> List[Tuple[int, List[int]]]: + """ + ``(offset, indices)`` per array: campaign indices grouped by + ``offset = (index // chunk) * chunk``, so one index is always in the same array + in every stage and its Slurm task id is ``index - offset``. + """ + size = self.chunk_size() + indices = self.indices() if indices is None else sorted(set(indices)) + chunks: Dict[int, List[int]] = {} + for index in indices: + chunks.setdefault((index // size) * size, []).append(index) + return sorted(chunks.items()) + + def order(self) -> List[str]: + """ + The stage names in dependency order (a stage after every stage it waits on). + + Raises + ------ + CampaignException + On a dependency cycle or an unknown stage. + """ + ordered: List[str] = [] + visiting = set() + + def visit(name, trail): + if name in ordered: + return + if name not in self.stages: + raise CampaignException(f"Stage {trail[-1]!r} waits on unknown stage {name!r}.") + if name in visiting: + raise CampaignException( + f"Stage dependency cycle: {' -> '.join(trail + [name])}." + ) + visiting.add(name) + for upstream in self.stages[name].after: + visit(upstream, trail + [name]) + visiting.discard(name) + ordered.append(name) + + for name in self.stages: + visit(name, [name]) + return ordered + + def _site_or_none(self): + try: + return self.site + except Exception: + return None + + def validate(self) -> List[str]: + """ + Every problem with the campaign, as messages (empty when valid). + """ + raw = self.raw + errors = [] + unknown = set(raw) - CAMPAIGN_FIELDS + if unknown: + errors.append(f"unknown campaign keys {sorted(unknown)}") + for key in ("campaign", "manifest", "stages"): + if not raw.get(key): + errors.append(f"{key}: required") + if self.scheduler not in SCHEDULERS: + errors.append(f"scheduler: {self.scheduler!r} is not one of {sorted(SCHEDULERS)}") + if self.chunk is not None and (not isinstance(self.chunk, int) or self.chunk < 1): + errors.append(f"chunk: must be a positive integer, not {self.chunk!r}") + + site = None + if self.site_name: + try: + site = self.site + except Exception as exception: + errors.append(f"site: {exception}") + + maximum = self.max_array_size() if site is not None else None + if maximum is not None and self.chunk is not None and self.chunk > maximum: + errors.append( + f"chunk: {self.chunk} is above the site's max_array_size {maximum}" + ) + + for name, raw_stage in (raw.get("stages") or {}).items(): + where = f"stages.{name}" + if not isinstance(raw_stage, dict): + errors.append(f"{where}: must be a mapping") + continue + unknown = set(raw_stage) - STAGE_FIELDS + if unknown: + errors.append(f"{where}: unknown keys {sorted(unknown)}") + if not raw_stage.get("script"): + errors.append(f"{where}.script: required") + resources = raw_stage.get("resources") or {} + unknown = set(resources) - RESOURCE_FIELDS + if unknown: + errors.append(f"{where}.resources: unknown keys {sorted(unknown)}") + try: + Resources.from_dict(resources) + except (ValueError, TypeError) as exception: + errors.append(f"{where}.resources: {exception}") + continue + for upstream, kind in (raw_stage.get("after") or {}).items(): + if upstream not in (raw.get("stages") or {}): + errors.append(f"{where}.after: unknown stage {upstream!r}") + if kind not in DEPENDENCY_KINDS: + errors.append( + f"{where}.after.{upstream}: {kind!r} is not one of " + f"{sorted(DEPENDENCY_KINDS)}" + ) + for failure, policy in (raw_stage.get("retry") or {}).items(): + if failure not in RETRY_CLASSES: + errors.append(f"{where}.retry: unknown failure class {failure!r}") + elif policy not in (None, "none") and not isinstance(policy, dict): + errors.append(f"{where}.retry.{failure}: must be a mapping or none") + + if name in self.stages and self.scheduler == "slurm": + errors += self._resource_errors(where, self.stages[name], site) + + if not errors: + try: + self.order() + except CampaignException as exception: + errors.append(str(exception)) + return errors + + def _resource_errors(self, where: str, stage: Stage, site) -> List[str]: + errors = [] + resources = stage.resources + if not resources.partition: + errors.append(f"{where}.resources.partition: required for a Slurm campaign") + return errors + names = split_partitions(resources.partition) + if site is not None: + for partition_name in names: + if partition_name not in site.partitions: + errors.append( + f"{where}.resources.partition: site {site.name!r} has no " + f"partition {partition_name!r}" + ) + continue + partition = site.partitions[partition_name] + max_time = parse_time(partition.max_time) if partition.max_time else None + if max_time and resources.time_s and resources.time_s > max_time: + errors.append( + f"{where}.resources.time: above {partition_name!r} max_time " + f"{partition.max_time}" + ) + mem_node = partition.mem_per_node_mb + if mem_node and resources.mem_mb and resources.mem_mb > mem_node: + errors.append( + f"{where}.resources: {resources.mem_mb} MB per task is above " + f"{partition_name!r} mem_per_node_mb {mem_node}" + ) + # Static half of the partition lint: a mixed or oversized request can never + # qualify for the exemption; the pending-GPU check runs again at submit time. + verdict = check_partitions( + names, + cpus=resources.cpus, + gres=resources.gres, + throttle=resources.throttle, + n_tasks=None, + pending_gpu_jobs=0, + site=site, + ) + if not verdict.ok: + errors.append(f"{where}.resources.partition: " + "; ".join(verdict.reasons)) + return errors diff --git a/autofit/campaign/submit.py b/autofit/campaign/submit.py new file mode 100644 index 000000000..dc1fdeef7 --- /dev/null +++ b/autofit/campaign/submit.py @@ -0,0 +1,176 @@ +""" +Submitting a campaign: one array per stage per manifest chunk, chained chunk by chunk, +recorded in the campaign's job ledger ``/campaign//jobs.json``. + +The ledger is how the collector finds the campaign's jobs later; it is written +atomically after every successful submit, so an interrupted submit leaves a ledger that +names exactly the arrays that exist. +""" +import json +import logging +from datetime import datetime, timezone +from pathlib import Path +from typing import Dict, List, Optional, Sequence + +from autofit.campaign.lint import check_partitions +from autofit.campaign.scheduler.base import SubmitRequest +from autofit.campaign.spec import Campaign, CampaignException +from autofit.campaign.units import array_spec, parse_array_spec +from autofit.tools.util import open_atomic + +logger = logging.getLogger(__name__) + +LEDGER = "jobs.json" + + +def ledger_path(campaign: Campaign) -> Path: + return campaign.directory / LEDGER + + +def read_ledger(campaign: Campaign) -> List[Dict]: + """ + The campaign's submitted arrays: ``{stage, job_id, offset, indices, submitted, + resources}`` per array, oldest first (empty before the first submit). + """ + path = ledger_path(campaign) + if not path.exists(): + return [] + return json.loads(path.read_text()) + + +def _write_ledger(campaign: Campaign, entries: List[Dict]): + path = ledger_path(campaign) + path.parent.mkdir(parents=True, exist_ok=True) + with open_atomic(path, "w") as f: + json.dump(entries, f, indent=1) + + +def ledger_indices(entry: Dict) -> List[int]: + return parse_array_spec(entry["indices"])[0] + + +def requests_for( + campaign: Campaign, + stages: Optional[Sequence[str]] = None, + indices: Optional[Sequence[int]] = None, + resources: Optional[Dict[str, object]] = None, +) -> List[SubmitRequest]: + """ + The arrays a submit of ``stages`` (default: all, in dependency order) over + ``indices`` (default: the whole manifest) would create, before dependencies are + resolved. + """ + order = campaign.order() + names = [name for name in order if stages is None or name in stages] + unknown = set(stages or ()) - set(order) + if unknown: + raise CampaignException(f"Unknown stages {sorted(unknown)}.") + requests = [] + for name in names: + stage = campaign.stages[name] + for offset, chunk in campaign.chunks(indices): + requests.append( + SubmitRequest( + campaign=str(campaign.name), + stage=name, + script=stage.script, + offset=offset, + indices=chunk, + resources=(resources or {}).get(name, stage.resources), + project_root=str(campaign.project_root), + manifest=str(campaign.path(campaign.manifest_path)), + env=dict(stage.env), + log_dir=campaign.log_dir, + ) + ) + return requests + + +def submit( + campaign: Campaign, + scheduler, + stages: Optional[Sequence[str]] = None, + indices: Optional[Sequence[int]] = None, + resources: Optional[Dict[str, object]] = None, + dry_run: bool = False, +) -> List[Dict]: + """ + Submit ``stages`` over ``indices``, chaining each array to the array of its + upstream stages that covers the same chunk (submitted now, or earlier and found + in the ledger). + + Parameters + ---------- + resources + Per-stage ``Resources`` overriding the campaign's (an approved resubmit). + dry_run + Return the requests (with ``"script"`` holding the batch script for a Slurm + scheduler) without submitting or writing the ledger. + + Returns + ------- + The ledger entries of the arrays submitted (or that would be). + """ + ledger = [] if dry_run else read_ledger(campaign) + submitted: List[Dict] = [] + # (stage, offset) -> job id, newest wins: an upstream resubmit supersedes the + # original array for the dependency of a later downstream submit. + jobs = {(entry["stage"], entry["offset"]): entry["job_id"] for entry in ledger} + + for request in requests_for(campaign, stages, indices, resources): + stage = campaign.stages[request.stage] + for upstream, kind in stage.after.items(): + job_id = jobs.get((upstream, request.offset)) + if job_id is None: + if dry_run: + job_id = f"<{upstream}@{request.offset}>" + else: + raise CampaignException( + f"Stage {request.stage!r} chunk {request.offset} waits on " + f"{upstream!r}, which has no array for that chunk; submit " + f"{upstream!r} first." + ) + request.after[str(job_id)] = kind + + entry = { + "stage": request.stage, + "offset": request.offset, + "indices": array_spec(request.indices), + "after": dict(request.after), + "resources": { + key: value + for key, value in vars(request.resources).items() + if value is not None + }, + } + if dry_run: + from autofit.campaign.scheduler.slurm import SlurmScheduler, sbatch_script + + if isinstance(scheduler, SlurmScheduler) or scheduler is None: + entry["script"] = sbatch_script(request) + # A dry run never queries the cluster, so the pending-GPU count is + # unknown: only the lint's static half can pass here. + resources = request.resources + verdict = check_partitions( + resources.partition, + cpus=resources.cpus, + gres=resources.gres, + throttle=resources.throttle, + n_tasks=len(request.indices), + pending_gpu_jobs=None, + site=campaign.site if campaign.site_name else None, + ) + entry["lint"] = "ok" if verdict.ok else "; ".join(verdict.reasons) + entry["job_id"] = f"<{request.stage}@{request.offset}>" + else: + entry["job_id"] = scheduler.submit(request) + entry["submitted"] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + ledger.append(entry) + _write_ledger(campaign, ledger) + logger.info( + f"Submitted {request.stage} chunk {request.offset} " + f"({len(request.indices)} tasks) as {entry['job_id']}." + ) + jobs[(request.stage, request.offset)] = entry["job_id"] + submitted.append(entry) + return submitted diff --git a/autofit/campaign/units.py b/autofit/campaign/units.py new file mode 100644 index 000000000..6d4181f3f --- /dev/null +++ b/autofit/campaign/units.py @@ -0,0 +1,177 @@ +""" +Parsing and formatting of the time, memory and array-index values a scheduler speaks. + +Slurm writes times as ``minutes``, ``MM:SS``, ``HH:MM:SS``, ``D-HH``, ``D-HH:MM`` or +``D-HH:MM:SS`` (with an optional ``.fff`` fraction from ``sacct``), and memory as an +integer number of MB, optionally suffixed ``K``/``M``/``G``/``T`` and, in older +``ReqMem`` output, ``n`` (per node) or ``c`` (per CPU). +""" +import re +from typing import Iterable, List, Optional, Tuple, Union + +_MEMORY = re.compile(r"^\s*([0-9]*\.?[0-9]+)\s*([KMGT]?)B?([nc]?)\s*$", re.IGNORECASE) +_MEMORY_SCALE = {"": 1.0, "K": 1 / 1024, "M": 1.0, "G": 1024.0, "T": 1024.0**2} + + +def parse_time(value: Union[str, int, float, None]) -> Optional[int]: + """ + Seconds in a Slurm time value, or ``None`` for an empty, ``UNLIMITED``, + ``Partition_Limit`` or ``INVALID`` value. A bare number is minutes (Slurm's + convention for ``--time``). + + Raises + ------ + ValueError + If ``value`` is not a Slurm time. + """ + if value is None: + return None + if isinstance(value, (int, float)): + return int(value * 60) + text = value.strip() + if not text or text.upper() in ("UNLIMITED", "PARTITION_LIMIT", "INVALID", "NONE"): + return None + + days = 0 + if "-" in text: + day_text, text = text.split("-", 1) + days = int(day_text) + parts = text.split(":") + # D-HH, D-HH:MM, D-HH:MM:SS + parts += ["0"] * (3 - len(parts)) + hours, minutes, seconds = parts + else: + parts = text.split(":") + if len(parts) == 1: + return int(float(parts[0]) * 60) + if len(parts) == 2: + hours, (minutes, seconds) = "0", parts + elif len(parts) == 3: + hours, minutes, seconds = parts + else: + raise ValueError(f"Not a Slurm time: {value!r}") + try: + return int( + days * 86400 + int(hours) * 3600 + int(minutes) * 60 + float(seconds) + ) + except ValueError: + raise ValueError(f"Not a Slurm time: {value!r}") + + +def format_time(seconds: Optional[float]) -> Optional[str]: + """ + ``seconds`` as Slurm's ``[D-]HH:MM:SS``, rounded up to the whole second. + """ + if seconds is None: + return None + total = int(-(-seconds // 1)) + days, rest = divmod(total, 86400) + hours, rest = divmod(rest, 3600) + minutes, secs = divmod(rest, 60) + clock = f"{hours:02d}:{minutes:02d}:{secs:02d}" + return f"{days}-{clock}" if days else clock + + +def parse_memory_mb(value: Union[str, int, float, None]) -> Optional[float]: + """ + MB in a memory value (``12288``, ``12288M``, ``12G``, ``29823.80M``, ``3072Mc``), + or ``None`` for an empty value. A bare number is MB. + + Raises + ------ + ValueError + If ``value`` is not a memory value. + """ + if value is None: + return None + if isinstance(value, (int, float)): + return float(value) + text = value.strip() + if not text: + return None + match = _MEMORY.match(text) + if match is None: + raise ValueError(f"Not a memory value: {value!r}") + number, unit, _ = match.groups() + return float(number) * _MEMORY_SCALE[unit.upper()] + + +def memory_per(value: Optional[str]) -> Optional[str]: + """ + ``"node"``, ``"cpu"`` or ``None``: what an old-style ``ReqMem`` value is per. + """ + if not value: + return None + match = _MEMORY.match(value.strip()) + if match is None or not match.group(3): + return None + return {"n": "node", "c": "cpu"}[match.group(3).lower()] + + +def array_spec(indices: Iterable[int], throttle: Optional[int] = None) -> str: + """ + A Slurm ``--array`` value for ``indices``: sorted runs collapsed into ranges + (``0-9,12,15-20``), with ``%throttle`` when given. + """ + values = sorted(set(int(index) for index in indices)) + if not values: + raise ValueError("An array needs at least one index.") + runs: List[Tuple[int, int]] = [] + start = previous = values[0] + for value in values[1:]: + if value == previous + 1: + previous = value + continue + runs.append((start, previous)) + start = previous = value + runs.append((start, previous)) + spec = ",".join(f"{a}" if a == b else f"{a}-{b}" for a, b in runs) + if throttle: + spec += f"%{int(throttle)}" + return spec + + +def parse_array_spec(spec: str) -> Tuple[List[int], Optional[int]]: + """ + The indices and throttle of a Slurm ``--array`` value (``0-9:2,12%4``). + """ + throttle = None + if "%" in spec: + spec, throttle_text = spec.split("%", 1) + throttle = int(throttle_text) + indices = [] + for part in spec.split(","): + part = part.strip() + if not part: + continue + step = 1 + if ":" in part: + part, step_text = part.split(":", 1) + step = int(step_text) + if "-" in part: + a, b = part.split("-", 1) + indices.extend(range(int(a), int(b) + 1, step)) + else: + indices.append(int(part)) + return indices, throttle + + +def requests_gpu(gres: Optional[str]) -> bool: + """ + Whether a ``--gres`` / ``--gpus`` / TRES request asks for a GPU (``gpu:1``, + ``gres/gpu=1``, ``gpu:a100:2``). A partition name never tells you the device: + only the request does. + """ + if not gres: + return False + for item in re.split(r"[,\s]+", str(gres)): + item = item.strip().lower() + if not item: + continue + name = item.split("=", 1)[0].split(":", 1)[0] + if name in ("gpu", "gres/gpu", "gpus") or name.startswith("gres/gpu"): + count = item.split("=", 1)[1] if "=" in item else None + if count is not None and count.strip() in ("0", ""): + continue + return True + return False diff --git a/test_autofit/campaign/__init__.py b/test_autofit/campaign/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test_autofit/campaign/conftest.py b/test_autofit/campaign/conftest.py new file mode 100644 index 000000000..02a2e7263 --- /dev/null +++ b/test_autofit/campaign/conftest.py @@ -0,0 +1,52 @@ +import textwrap +from pathlib import Path + +import pytest + +from autonerves.site_config import load_site, site_path + + +@pytest.fixture(name="ral") +def make_ral(): + return load_site(site_path("ral")) + + +CAMPAIGN = """ +campaign: demo +project_root: {root} +manifest: manifest.txt +manifest_fields: [sample, tile] +chunk: 4 +site: ral +log_dir: logs +stages: + vis_lp: + script: scripts/fit.py --stage=vis_lp + env: {{JAX_PLATFORMS: cpu, OMP_NUM_THREADS: "$SLURM_CPUS_PER_TASK"}} + resources: {{cpus: 4, mem_per_cpu_mb: 3072, time: "04:00:00", partition: ral, throttle: 400, nice: 1000}} + success: {{zip_complete: "output/{{sample}}/{{tile}}/vis_lp", members: [files/samples_summary.json]}} + retry: {{timeout: {{time: "p95+25%", max: 1}}, oom: {{mem: "peak*1.5", max: 1}}}} + vis_pix: + after: {{vis_lp: aftercorr}} + script: scripts/fit.py --stage=vis_pix + resources: {{cpus: 4, mem_mb: 16384, time: "12:00:00", partition: ral}} + sed: + after: {{vis_lp: aftercorr}} + script: scripts/fit.py --stage=sed + resources: {{cpus: 2, mem_per_cpu_mb: 4096, time: "06:00:00", partition: ral}} +""" + + +def write_campaign(root: Path, text: str = CAMPAIGN, n: int = 10) -> Path: + root.mkdir(parents=True, exist_ok=True) + (root / "manifest.txt").write_text( + "\n".join(f"dr1 Tile{index:03d}" for index in range(n)) + "\n" + ) + path = root / "campaign.yaml" + path.write_text(textwrap.dedent(text.format(root=root))) + return path + + +@pytest.fixture(name="campaign_file") +def make_campaign_file(tmp_path): + return write_campaign(tmp_path / "project") diff --git a/test_autofit/campaign/test_cli.py b/test_autofit/campaign/test_cli.py new file mode 100644 index 000000000..5d0789bbe --- /dev/null +++ b/test_autofit/campaign/test_cli.py @@ -0,0 +1,21 @@ +from autofit.campaign.__main__ import main + + +def test_validate(campaign_file, capsys): + assert main(["validate", str(campaign_file)]) == 0 + out = capsys.readouterr().out + assert "10 objects, 3 chunk(s) of ≤4, stages vis_lp → vis_pix → sed" in out + assert "unverified" in out + + +def test_validate_refuses(campaign_file, capsys): + campaign_file.write_text(campaign_file.read_text().replace("partition: ral, throttle", "partition: \"gpu,ral\", throttle")) + assert main(["validate", str(campaign_file)]) == 1 + assert "mixes GPU partition" in capsys.readouterr().err + + +def test_submit_dry_run(campaign_file, capsys): + assert main(["submit", str(campaign_file), "--dry-run", "--stage", "vis_lp", "--indices", "0-5"]) == 0 + out = capsys.readouterr().out + assert out.count("#!/bin/bash") == 2 + assert "#SBATCH --array=0-1%400" in out diff --git a/test_autofit/campaign/test_lint.py b/test_autofit/campaign/test_lint.py new file mode 100644 index 000000000..c31122625 --- /dev/null +++ b/test_autofit/campaign/test_lint.py @@ -0,0 +1,121 @@ +""" +The GPU-partition rule (human, 2026-09-30): every branch. +""" +import pytest + +from autofit.campaign.lint import PartitionLintError, check_partitions, check_sbatch + + +def _check(partitions, ral, **kwargs): + defaults = dict(cpus=4, gres=None, throttle=None, n_tasks=1000, pending_gpu_jobs=0) + defaults.update(kwargs) + return check_partitions(partitions, site=ral, **defaults) + + +def test_cpu_partition_is_fine(ral): + result = _check("ral", ral) + assert result.ok and not result.exempt and result.gpu_partitions == [] + + +@pytest.mark.parametrize("partitions", ["ral,gpu", "gpu,ral", ["ral", "gpu"], " gpu , ral "]) +def test_mixed_list_always_refused(partitions, ral): + # Even a request that would meet the exemption: the incident form is never allowed. + result = _check(partitions, ral, cpus=1, throttle=1, n_tasks=1, pending_gpu_jobs=0) + assert not result.ok + assert "mixes GPU partition" in result.reasons[0] + with pytest.raises(PartitionLintError): + result.raise_for_refusal() + + +def test_bulk_cpu_array_on_gpu_refused(ral): + result = _check("gpu", ral, cpus=8, throttle=400) + assert not result.ok + assert "throttle 400" in result.reasons[0] + + +def test_timing_leg_exemption(ral): + result = _check("gpu", ral, cpus=8, throttle=2, pending_gpu_jobs=0) + assert result.ok and result.exempt + + +def test_exemption_small_plain_job(ral): + assert _check("gpu", ral, cpus=2, throttle=None, n_tasks=1).exempt + + +def test_exemption_too_many_cpus(ral): + result = _check("gpu", ral, cpus=9, throttle=2) + assert not result.ok and "9 CPUs per task" in result.reasons[0] + + +def test_exemption_gpu_jobs_pending(ral): + result = _check("gpu", ral, cpus=8, throttle=2, pending_gpu_jobs=3) + assert not result.ok and "3 GPU job(s) pending" in result.reasons[0] + + +def test_exemption_fails_closed_when_pending_unknown(ral): + result = _check("gpu", ral, cpus=8, throttle=2, pending_gpu_jobs=None) + assert not result.ok and "unknown" in result.reasons[0] + + +def test_gpu_request_may_use_gpu(ral): + result = _check("gpu", ral, cpus=32, gres="gpu:1", throttle=None, pending_gpu_jobs=5) + assert result.ok and result.requests_gpu and not result.exempt + + +def test_gpu_request_on_mixed_list_is_not_linted(ral): + assert _check("ral,gpu", ral, gres="gpu:a100:1").ok + + +def test_without_site_names_containing_gpu_are_gpu_partitions(): + result = check_partitions("cpu,gpu_a100", cpus=4, n_tasks=10, pending_gpu_jobs=0) + assert not result.ok and result.gpu_partitions == ["gpu_a100"] + + +def test_site_decides_over_name(ral): + # `ral` is gpu_only: false in the site file; `gpu` is gpu_only: true. + assert _check("ral", ral).gpu_partitions == [] + assert _check("gpu", ral, throttle=2).gpu_partitions == ["gpu"] + + +SCRIPT = """#!/bin/bash +#SBATCH --job-name=euclid_dr1_vis_lp +#SBATCH --partition=ral +#SBATCH --array=0-999%400 +#SBATCH --cpus-per-task=4 +python fit.py +""" + + +def test_sbatch_script_lint(ral): + assert check_sbatch(SCRIPT, site=ral, pending_gpu_jobs=0).ok + + +@pytest.mark.parametrize("arguments", [["-p", "gpu"], ["--partition=ral,gpu"], ["--partition", "gpu,ral"]]) +def test_sbatch_command_line_override_is_linted(arguments, ral): + assert not check_sbatch(SCRIPT, arguments, site=ral, pending_gpu_jobs=0).ok + + +def test_sbatch_directive_mixed_list_refused(ral): + script = SCRIPT.replace("--partition=ral", "--partition=ral,gpu") + assert not check_sbatch(script, site=ral, pending_gpu_jobs=0).ok + + +def test_sbatch_timing_leg(ral): + script = SCRIPT.replace("--partition=ral", "-p gpu").replace("0-999%400", "0-5%2").replace( + "--cpus-per-task=4", "-c 8" + ) + result = check_sbatch(script, site=ral, pending_gpu_jobs=0) + assert result.ok and result.exempt + + +def test_sbatch_gpu_job(ral): + script = SCRIPT.replace("--partition=ral", "--partition=gpu\n#SBATCH --gpus=1") + result = check_sbatch(script, site=ral, pending_gpu_jobs=7) + assert result.ok and result.requests_gpu + + +def test_sbatch_directive_after_command_is_ignored(ral): + # sbatch ignores directives after the first command, so a late --gpus does not + # make a CPU job a GPU job. + script = SCRIPT.replace("--partition=ral", "--partition=gpu") + "#SBATCH --gpus=1\n" + assert not check_sbatch(script, site=ral, pending_gpu_jobs=0).ok diff --git a/test_autofit/campaign/test_local.py b/test_autofit/campaign/test_local.py new file mode 100644 index 000000000..7fa5ce43c --- /dev/null +++ b/test_autofit/campaign/test_local.py @@ -0,0 +1,55 @@ +import sys +import textwrap + +from autofit.campaign import Campaign +from autofit.campaign.submit import submit +from autofit.campaign.scheduler import LocalScheduler +from test_autofit.campaign.conftest import write_campaign + +TEXT = """ +campaign: local_demo +project_root: {root} +manifest: manifest.txt +scheduler: local +chunk: 2 +stages: + first: + script: task.py first + resources: {{time: "00:01:00"}} + second: + after: {{first: aftercorr}} + script: task.py second +""" + + +def test_local_campaign(tmp_path): + root = tmp_path / "p" + path = write_campaign(root, TEXT, n=4) + (root / "task.py").write_text( + textwrap.dedent( + """ + import os, sys + index = int(os.environ["CAMPAIGN_INDEX"]) + assert index == int(os.environ["CAMPAIGN_OFFSET"]) + int(os.environ["SLURM_ARRAY_TASK_ID"]) + if sys.argv[1] == "first" and index == 2: + sys.exit(3) + open(f"{sys.argv[1]}_{index}.done", "w").write(os.environ["MANIFEST"]) + """ + ) + ) + campaign = Campaign.load(path) + scheduler = LocalScheduler(max_workers=2) + entries = submit(campaign, scheduler) + scheduler.wait() + states = { + (record.name, record.task_id + entry["offset"]): record.scheduler_state + for entry in entries + for record in scheduler.query([entry["job_id"]]) + } + assert states[("local_demo_first", 2)] == "FAILED" + assert states[("local_demo_second", 2)] == "CANCELLED" + assert states[("local_demo_second", 3)] == "COMPLETED" + assert sorted(p.name for p in root.glob("*.done")) == [ + "first_0.done", "first_1.done", "first_3.done", + "second_0.done", "second_1.done", "second_3.done", + ] diff --git a/test_autofit/campaign/test_sacct.py b/test_autofit/campaign/test_sacct.py new file mode 100644 index 000000000..912cfb5a6 --- /dev/null +++ b/test_autofit/campaign/test_sacct.py @@ -0,0 +1,68 @@ +from autofit.campaign.scheduler import RunState, run_state_of +from autofit.campaign.scheduler.sacct import latest_per_task, parse + +HEADER = "JobID|JobIDRaw|JobName|Partition|State|ExitCode|Submit|Start|End|Elapsed|Timelimit|AllocCPUS|ReqMem|MaxRSS|AveRSS|TotalCPU|CPUTime|NNodes|NodeList|ReqTRES|AllocTRES|WorkDir" + +DUMP = "\n".join( + [ + HEADER, + "342027_1|342029|vis_lp|ral|COMPLETED|0:0|2026-08-30T02:54:31|2026-08-30T02:54:31|2026-09-01T00:04:55|1-21:10:24|2-00:00:00|8|98304M|||12-01:02:31|15-01:23:12|1|n21|billing=8,cpu=8,mem=98304M,node=1|billing=8,cpu=8,mem=98304M,node=1|/w/euclid_dr1/hpc", + "342027_1.batch|342029.batch|batch||COMPLETED|0:0|2026-08-30T02:54:31|2026-08-30T02:54:31|2026-09-01T00:04:55|1-21:10:24||8||29823.80M|29823.80M|12-01:02:31|15-01:23:12|1|n21||cpu=8,mem=98304M,node=1|", + "342027_1.extern|342029.extern|extern||COMPLETED|0:0|2026-08-30T02:54:31|2026-08-30T02:54:31|2026-09-01T00:04:55|1-21:10:24||8||0.50M|0.50M|00:00.001|15-01:23:12|1|n21||cpu=8,mem=98304M,node=1|", + "342027_2|342030|vis_lp|ral|REQUEUED|0:0|2026-08-30T02:54:31|2026-08-30T02:55:00|2026-08-30T03:00:00|00:05:00|2-00:00:00|8|98304M|||00:30:00|00:40:00|1|n3|billing=8|billing=8|/w/euclid_dr1/hpc", + "342027_2.batch|342030.batch|batch||CANCELLED|0:15|2026-08-30T02:54:31|2026-08-30T02:55:00|2026-08-30T03:00:00|00:05:00||8||100M|100M|00:30:00|00:40:00|1|n3||cpu=8|", + "342027_2|342030|vis_lp|ral|TIMEOUT|0:0|2026-08-30T03:01:00|2026-08-30T03:02:00|2026-09-01T03:02:00|2-00:00:00|2-00:00:00|8|98304M|||15-00:00:00|16-00:00:00|1|n4|billing=8|billing=8|/w/euclid_dr1/hpc", + "342027_2.batch|342030.batch|batch||CANCELLED|0:15|2026-08-30T03:01:00|2026-08-30T03:02:00|2026-09-01T03:02:00|2-00:00:00||8||50000M|50000M|15-00:00:00|16-00:00:00|1|n4||cpu=8|", + "350000|350000|sed|gpu|OUT_OF_MEMORY|0:125|2026-09-30T10:00:00|2026-09-30T10:00:01|2026-09-30T10:10:01|00:10:00|08:00:00|4|3072Mc|||00:30:00|00:40:00|1|g1|billing=4,cpu=4,gres/gpu=1|billing=4,cpu=4|/w/euclid_dr1/hpc", + "350000.batch|350000.batch|batch||OUT_OF_MEMORY|0:125|2026-09-30T10:00:00|2026-09-30T10:00:01|2026-09-30T10:10:01|00:10:00||4||12500M|12500M|00:30:00|00:40:00|1|g1||cpu=4|", + "344645_[586,588-594,...]|344645|vis_pix|ral|PENDING|0:0|2026-09-02T17:39:47|Unknown|Unknown|00:00:00|2-00:00:00|8|98304M|||00:00:00|00:00:00|1|None assigned|billing=8|billing=8|/w/euclid_dr1/hpc", + ] +) + + +def test_parse_folds_steps(): + records = parse(DUMP) + assert len(records) == 5 + first = records[0] + assert (first.job_id, first.task_id, first.name) == ("342027", 1, "vis_lp") + assert first.max_rss_mb == 29823.8 + assert first.elapsed_s == 162624 and first.time_limit_s == 172800 + assert first.total_cpu_s == 12 * 86400 + 3751 + assert first.cpu_efficiency == first.total_cpu_s / (162624 * 8) + assert first.state == RunState.COMPLETED + + +def test_requeue_attempts_and_latest(): + tasks = {record.key: record for record in latest_per_task(parse(DUMP))} + task = tasks["342027_2"] + assert task.attempts == 2 + assert task.state == RunState.TIMEOUT + assert task.max_rss_mb == 50000 + + +def test_per_cpu_memory_and_gpu_tres(): + record = parse(DUMP)[3] + assert record.req_mem_mb == 3072 * 4 + assert record.state == RunState.OOM + assert record.gres == "gres/gpu=1" + assert record.task_id is None + + +def test_pending_placeholder(): + record = parse(DUMP)[-1] + assert record.job_id.startswith("344645_[") and record.state == RunState.PENDING + assert record.start is None + + +def test_noheader_with_fields(): + fields = ["JobID", "JobName", "State", "Elapsed", "AllocCPUS"] + records = parse("7_0|fit|COMPLETED|00:10:00|2\n7_0.batch|batch|COMPLETED|00:10:00|2\n", fields=fields) + assert len(records) == 1 and records[0].cpu_hours == 20 / 60 + + +def test_state_mapping(): + assert run_state_of("CANCELLED by 16309") == "cancelled" + assert run_state_of("NODE_FAIL") == RunState.NODE_FAIL + assert run_state_of("PD") == RunState.PENDING + assert run_state_of("COMPLETED+") == RunState.COMPLETED + assert run_state_of("WEIRD") == "unknown" diff --git a/test_autofit/campaign/test_slurm.py b/test_autofit/campaign/test_slurm.py new file mode 100644 index 000000000..7e59aed4e --- /dev/null +++ b/test_autofit/campaign/test_slurm.py @@ -0,0 +1,162 @@ +""" +The Slurm adapter over an injected runner: no command ever reaches a cluster. +""" +import subprocess +import sys + +import pytest + +from autofit.campaign import PartitionLintError +from autofit.campaign.scheduler import SlurmScheduler, SubmitRequest, sbatch_script +from autofit.campaign.spec import Resources + + +class Runner: + def __init__(self, outputs=None): + self.calls = [] + self.outputs = outputs or {} + + def __call__(self, arguments, input=None): + self.calls.append((list(arguments), input)) + return self.outputs.get(arguments[0], "") + + +SINFO = "\n".join( + [ + "ral*|n1|mix|200/52/0/252|920000|300000", + "ral*|n2|idle|0/252/0/252|920000|900000", + "ral*|n3|down*|0/252/0/252|920000|0", + "gpu|g1|mix|60/2/0/62|500000|1000", + "gpu|g2|mix|62/0/0/62|500000|1000", + ] +) +SQUEUE_PENDING = "gpu|gres/gpu:1|413161_[0-5%16]\nral|N/A|413200_[0-99]\ngpu|N/A|413300\n" + + +def _request(**resources): + defaults = dict(cpus=4, mem_per_cpu_mb=3072, time="04:00:00", partition="ral", throttle=400, nice=1000) + defaults.update(resources) + return SubmitRequest( + campaign="dr1", + stage="vis_lp", + script="scripts/fit.py --stage=vis_lp", + offset=1000, + indices=list(range(1000, 1010)), + resources=Resources.from_dict(defaults), + project_root="/mnt/ral/x/euclid_dr1", + manifest="/mnt/ral/x/euclid_dr1/m.txt", + env={"OMP_NUM_THREADS": "$SLURM_CPUS_PER_TASK", "NAME": "a b"}, + after={"999": "aftercorr"}, + log_dir="hpc/logs", + ) + + +def test_sbatch_script(): + script = sbatch_script(_request()) + for line in [ + "#SBATCH --job-name=dr1_vis_lp", + "#SBATCH --partition=ral", + "#SBATCH --array=0-9%400", + "#SBATCH --cpus-per-task=4", + "#SBATCH --mem-per-cpu=3072", + "#SBATCH --time=04:00:00", + "#SBATCH --nice=1000", + "#SBATCH --dependency=aftercorr:999", + "#SBATCH --kill-on-invalid-dep=yes", + "#SBATCH --signal=B:USR1@600", + "#SBATCH --output=hpc/logs/%x_%A_%a.out", + "set -euo pipefail", + "export PYTHONUNBUFFERED=1", + "export CAMPAIGN_OFFSET=1000", + 'export OMP_NUM_THREADS="$SLURM_CPUS_PER_TASK"', + "export NAME='a b'", + "export CAMPAIGN_INDEX=$((CAMPAIGN_OFFSET + SLURM_ARRAY_TASK_ID))", + "python scripts/fit.py --stage=vis_lp &", + "trap 'kill -USR1 \"$child\" 2>/dev/null || true' USR1", + ]: + assert line in script.splitlines(), line + assert "--gres" not in script + + +def test_sbatch_script_passes_its_own_lint(ral): + from autofit.campaign.lint import check_sbatch + + assert check_sbatch(sbatch_script(_request()), site=ral, pending_gpu_jobs=0).ok + assert not check_sbatch(sbatch_script(_request(partition="ral,gpu")), site=ral, pending_gpu_jobs=0).ok + + +def test_capacity(): + scheduler = SlurmScheduler(runner=Runner({"sinfo": SINFO, "squeue": SQUEUE_PENDING})) + capacity = scheduler.capacity() + ral, gpu = capacity["ral"], capacity["gpu"] + assert (ral.cpus_idle, ral.cpus_allocated, ral.cpus_total, ral.nodes, ral.nodes_unreachable) == (304, 200, 756, 3, 1) + assert ral.cpus_other == 252 + assert (gpu.cpus_idle, gpu.pending_gpu_jobs) == (2, 1) + assert ral.pending_gpu_jobs == 0 + + +def test_submit_lints_then_sbatches(ral): + runner = Runner({"sinfo": SINFO, "squeue": SQUEUE_PENDING, "sbatch": "413500;cluster\n"}) + scheduler = SlurmScheduler(runner=runner, site=ral) + assert scheduler.submit(_request()) == "413500" + arguments, script = runner.calls[-1] + assert arguments == ["sbatch", "--parsable"] + assert "#SBATCH --array=0-9%400" in script + + +@pytest.mark.parametrize("partition", ["ral,gpu", "gpu,ral", "gpu"]) +def test_submit_refuses_cpu_on_gpu(partition, ral): + runner = Runner({"sinfo": SINFO, "squeue": SQUEUE_PENDING, "sbatch": "1"}) + with pytest.raises(PartitionLintError): + SlurmScheduler(runner=runner, site=ral).submit(_request(partition=partition)) + assert not any(call[0][0] == "sbatch" for call in runner.calls) + + +def test_timing_leg_refused_while_gpu_jobs_pend(ral): + runner = Runner({"sinfo": SINFO, "squeue": SQUEUE_PENDING, "sbatch": "1"}) + request = _request(partition="gpu", cpus=8, throttle=2) + with pytest.raises(PartitionLintError, match="1 GPU job"): + SlurmScheduler(runner=runner, site=ral).submit(request) + runner.outputs["squeue"] = "ral|N/A|1\n" + assert SlurmScheduler(runner=runner, site=ral).submit(request) == "1" + + +def test_query_is_one_sacct_call(): + runner = Runner({"sacct": "5_0|5|fit|ral|COMPLETED|0:0|||||01:00:00|2|4096M||||||||||\n"}) + records = SlurmScheduler(runner=runner).query(["5", "6"]) + (arguments, _), = runner.calls + assert arguments[:3] == ["sacct", "-j", "5,6"] + assert "--parsable2" in arguments and "--units=M" in arguments and "--duplicates" in arguments + assert records[0].key == "5_0" and records[0].cpus == 2 + + +def test_update_pending(ral): + runner = Runner({"sinfo": SINFO, "squeue": SQUEUE_PENDING}) + scheduler = SlurmScheduler(runner=runner, site=ral) + scheduler.update_pending("413500", time=5 * 3600, mem_per_cpu_mb=4096.4, throttle=200) + assert runner.calls[-1][0] == [ + "scontrol", "update", "JobId=413500", "TimeLimit=05:00:00", "MinMemoryCPU=4096", "ArrayTaskThrottle=200" + ] + with pytest.raises(PartitionLintError): + scheduler.update_pending("413500", partition="ral,gpu", cpus=4) + with pytest.raises(ValueError): + scheduler.update_pending("413500", gpus=4) + + +def test_cancel(): + runner = Runner() + SlurmScheduler(runner=runner).cancel("413500") + assert runner.calls == [(["scancel", "413500"], None)] + + +def test_no_command_at_import(): + # Import autofit first (it is not ours to police), then forbid any command while + # the campaign package and its adapters import. + code = ( + "import subprocess, autofit\n" + "def refuse(*a, **k): raise SystemExit('a command ran at import')\n" + "subprocess.run = subprocess.Popen = subprocess.check_output = refuse\n" + "import autofit.campaign, autofit.campaign.scheduler, autofit.campaign.__main__\n" + ) + completed = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) + assert completed.returncode == 0, completed.stderr[-2000:] diff --git a/test_autofit/campaign/test_spec.py b/test_autofit/campaign/test_spec.py new file mode 100644 index 000000000..eaa5d107a --- /dev/null +++ b/test_autofit/campaign/test_spec.py @@ -0,0 +1,81 @@ +import pytest + +from autofit.campaign.spec import Campaign, CampaignException +from test_autofit.campaign.conftest import CAMPAIGN, write_campaign + + +def test_load(campaign_file): + campaign = Campaign.load(campaign_file) + assert campaign.name == "demo" + assert campaign.order() == ["vis_lp", "vis_pix", "sed"] + vis_lp = campaign.stages["vis_lp"].resources + assert (vis_lp.cpus, vis_lp.mem_per_cpu_mb, vis_lp.time_s, vis_lp.mem_mb) == (4, 3072, 14400, 12288) + # A total `mem_mb` becomes one per-CPU knob. + assert campaign.stages["vis_pix"].resources.mem_per_cpu_mb == 4096 + assert campaign.stages["vis_lp"].success.members == ["files/samples_summary.json"] + assert campaign.site.name == "ral" + + +def test_chunks_and_fields(campaign_file): + campaign = Campaign.load(campaign_file) + assert campaign.chunks() == [(0, [0, 1, 2, 3]), (4, [4, 5, 6, 7]), (8, [8, 9])] + assert campaign.chunks([9, 1]) == [(0, [1]), (8, [9])] + assert campaign.fields(3) == {"index": "3", "line": "dr1 Tile003", "sample": "dr1", "tile": "Tile003"} + + +def test_comments_keep_line_numbers(tmp_path): + path = write_campaign(tmp_path / "p") + (tmp_path / "p" / "manifest.txt").write_text("# header\na\n\nb\n") + assert Campaign.load(path).indices() == [1, 3] + + +def _errors(tmp_path, old, new): + path = write_campaign(tmp_path / "p", CAMPAIGN.replace(old, new)) + with pytest.raises(CampaignException) as info: + Campaign.load(path) + return str(info.value) + + +def test_cpu_stage_on_gpu_refused(tmp_path): + message = _errors(tmp_path, "time: \"12:00:00\", partition: ral}", "time: \"12:00:00\", partition: \"ral,gpu\"}") + assert "mixes GPU partition" in message + + +def test_bulk_cpu_stage_on_gpu_alone_refused(tmp_path): + message = _errors(tmp_path, "partition: ral, throttle: 400", "partition: gpu, throttle: 400") + assert "stages.vis_lp.resources.partition" in message and "throttle 400" in message + + +def test_unknown_partition(tmp_path): + assert "no partition 'cpu'" in _errors(tmp_path, "partition: ral, throttle", "partition: cpu, throttle") + + +def test_unknown_keys_and_dependencies(tmp_path): + message = _errors(tmp_path, "after: {{vis_lp: aftercorr}}\n script: scripts/fit.py --stage=vis_pix", + "after: {{nope: afterwards}}\n script: scripts/fit.py --stage=vis_pix\n walltime: 3") + assert "unknown stage 'nope'" in message + assert "'afterwards' is not one of" in message + assert "unknown keys ['walltime']" in message + + +def test_chunk_above_max_array_size(tmp_path): + assert "max_array_size 1001" in _errors(tmp_path, "chunk: 4", "chunk: 5000") + + +def test_unknown_retry_class(tmp_path): + assert "unknown failure class 'segfault'" in _errors(tmp_path, "retry: {{timeout", "retry: {{segfault: none, timeout") + + +def test_cycle(tmp_path): + text = CAMPAIGN.replace(" script: scripts/fit.py --stage=vis_lp", " after: {{sed: aftercorr}}\n script: scripts/fit.py --stage=vis_lp") + message = _errors(tmp_path, CAMPAIGN, text) + assert "cycle" in message + + +def test_local_campaign_needs_no_partition(tmp_path): + text = ( + "campaign: local_demo\nproject_root: {root}\nmanifest: manifest.txt\nscheduler: local\n" + "stages:\n fit:\n script: fit.py\n" + ) + campaign = Campaign.load(write_campaign(tmp_path / "p", text)) + assert campaign.chunks() == [(0, list(range(10)))] diff --git a/test_autofit/campaign/test_submit.py b/test_autofit/campaign/test_submit.py new file mode 100644 index 000000000..8ad95e345 --- /dev/null +++ b/test_autofit/campaign/test_submit.py @@ -0,0 +1,73 @@ +""" +Submitting a campaign end to end on the FakeScheduler: arrays per chunk, per-index +dependencies, the job ledger, and the lint. +""" +import pytest + +from autofit.campaign import Campaign, PartitionLintError, read_ledger +from autofit.campaign.submit import submit +from autofit.campaign.scheduler import FakeScheduler, PartitionCapacity +from autofit.campaign.spec import CampaignException, Resources + + +def _capacity(pending=0): + return { + "ral": PartitionCapacity("ral", cpus_idle=100, cpus_total=256, pending_gpu_jobs=0), + "gpu": PartitionCapacity("gpu", cpus_idle=10, cpus_total=124, pending_gpu_jobs=pending), + } + + +def test_submit_chains_chunks(campaign_file, ral): + campaign = Campaign.load(campaign_file) + scheduler = FakeScheduler( + outcomes={("vis_lp", 5): {"scheduler_state": "FAILED"}}, + capacity=_capacity(), + site=ral, + ) + entries = submit(campaign, scheduler) + assert [(e["stage"], e["offset"]) for e in entries] == [ + ("vis_lp", 0), ("vis_lp", 4), ("vis_lp", 8), + ("vis_pix", 0), ("vis_pix", 4), ("vis_pix", 8), + ("sed", 0), ("sed", 4), ("sed", 8), + ] + jobs = {(e["stage"], e["offset"]): e["job_id"] for e in entries} + assert entries[4]["after"] == {jobs[("vis_lp", 4)]: "aftercorr"} + assert read_ledger(campaign) == entries + + scheduler.run() + records = {(r.name, int(r.job_id), r.task_id): r.scheduler_state for job in jobs.values() for r in scheduler.query([job])} + # Index 5 is task 1 of the offset-4 arrays: its vis_lp failed, so the dependent + # vis_pix and sed tasks are cancelled (kill-on-invalid-dep); every other one ran. + assert records[("demo_vis_lp", int(jobs[("vis_lp", 4)]), 1)] == "FAILED" + assert records[("demo_vis_pix", int(jobs[("vis_pix", 4)]), 1)] == "CANCELLED" + assert records[("demo_sed", int(jobs[("sed", 4)]), 1)] == "CANCELLED" + assert records[("demo_vis_pix", int(jobs[("vis_pix", 4)]), 0)] == "COMPLETED" + + +def test_downstream_alone_uses_ledger(campaign_file, ral): + campaign = Campaign.load(campaign_file) + scheduler = FakeScheduler(capacity=_capacity(), site=ral) + with pytest.raises(CampaignException, match="submit 'vis_lp' first"): + submit(campaign, scheduler, stages=["vis_pix"]) + first = submit(campaign, scheduler, stages=["vis_lp"], indices=[1, 9]) + second = submit(campaign, scheduler, stages=["vis_pix"], indices=[1, 9]) + assert [e["after"] for e in second] == [{e["job_id"]: "aftercorr"} for e in first] + assert len(read_ledger(campaign)) == 4 + + +def test_override_resources_is_linted(campaign_file, ral): + campaign = Campaign.load(campaign_file) + scheduler = FakeScheduler(capacity=_capacity(), site=ral) + bad = Resources.from_dict({"cpus": 4, "partition": "ral,gpu"}) + with pytest.raises(PartitionLintError): + submit(campaign, scheduler, stages=["vis_lp"], resources={"vis_lp": bad}) + assert read_ledger(campaign) == [] + + +def test_dry_run(campaign_file): + campaign = Campaign.load(campaign_file) + entries = submit(campaign, None, dry_run=True) + assert len(entries) == 9 + assert "#SBATCH --dependency=aftercorr:" in entries[4]["script"] + assert entries[0]["lint"] == "ok" + assert not campaign.directory.exists() diff --git a/test_autofit/campaign/test_units.py b/test_autofit/campaign/test_units.py new file mode 100644 index 000000000..c4f5dd6ca --- /dev/null +++ b/test_autofit/campaign/test_units.py @@ -0,0 +1,81 @@ +import pytest + +from autofit.campaign import units + + +@pytest.mark.parametrize( + "text, seconds", + [ + ("04:00:00", 14400), + ("1-21:10:24", 162624), + ("2-00:00:00", 172800), + ("2-12", 216000), + ("2-12:30", 217800), + ("10:30", 630), + ("90", 5400), + ("UNLIMITED", None), + ("", None), + ("Partition_Limit", None), + ], +) +def test_parse_time(text, seconds): + assert units.parse_time(text) == seconds + + +def test_parse_time_rejects_garbage(): + with pytest.raises(ValueError): + units.parse_time("soon") + + +def test_format_time_round_trip(): + assert units.format_time(14400) == "04:00:00" + assert units.format_time(172800) == "2-00:00:00" + assert units.format_time(59.2) == "00:01:00" + assert units.parse_time(units.format_time(162624)) == 162624 + + +@pytest.mark.parametrize( + "text, mb", + [ + ("98304M", 98304), + ("29823.80M", 29823.8), + ("12G", 12288), + ("12288", 12288), + ("3072Mc", 3072), + ("1024K", 1), + ("", None), + ], +) +def test_parse_memory(text, mb): + assert units.parse_memory_mb(text) == pytest.approx(mb) if mb else units.parse_memory_mb(text) is None + + +def test_memory_per(): + assert units.memory_per("3072Mc") == "cpu" + assert units.memory_per("98304Mn") == "node" + assert units.memory_per("98304M") is None + + +def test_array_spec_round_trip(): + spec = units.array_spec([0, 1, 2, 3, 7, 9, 10], throttle=400) + assert spec == "0-3,7,9-10%400" + assert units.parse_array_spec(spec) == ([0, 1, 2, 3, 7, 9, 10], 400) + assert units.parse_array_spec("0-6:2") == ([0, 2, 4, 6], None) + + +@pytest.mark.parametrize( + "gres, expected", + [ + ("gpu:1", True), + ("gpu:a100:2", True), + ("billing=8,cpu=8,gres/gpu=1,mem=1G", True), + ("gres/gpu:a100=1", True), + ("billing=8,cpu=8,mem=98304M,node=1", False), + ("gpu:0", True), + ("gres/gpu=0", False), + (None, False), + ("", False), + ], +) +def test_requests_gpu(gres, expected): + assert units.requests_gpu(gres) is expected