From 323f6d5cd94cd33a86471f4119b8f55763053c81 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 28 Oct 2026 21:49:54 +0000 Subject: [PATCH] feat: per-fit status record, failure record and run-state vocabulary hpc-campaign phase 1 (PyAutoLabs/PyAutoFit#1735). - autofit/non_linear/status/: states.py (RunState, FailureClass, Product: the single vocabulary), writer.py (StatusWriter: atomic .status.json capped at 2 KB, append-only .status.jsonl history, .traceback.txt), probe.py (fail-soft resources: cgroup v2 memory.peak inside a job else getrusage, CPU seconds self+children, NVML only when importable and a GPU is allocated, CPU model, scheduler job from a generic SLURM_* env mapping), signals.py (SIGUSR1 -> checkpoint flush + timeout_pending; SIGTERM same then exit 143; main thread of the top-level process only, previous handlers restored, forked workers pass signals on unchanged), record.py (FitStatus lifecycle, cadence thread, status_recorded / status_notify hooks). - fit() runs inside status_recorded: an exception writes `failed` + failure_class + traceback and is re-raised unchanged. Hooks at start_resume_fit (running), FitContext.update (progress), _complete_fit (completed) and result_via_completed_fit (skipped, not a new attempt). - progress() for Nautilus, Dynesty, Emcee and BFGS/LBFGS. - general.yaml hpc.status (default true) and hpc.status_interval_s (60); hpc.status: false writes nothing and installs no handlers. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_016nhQ8N4RdvbZz9cstLYtNF --- autofit/config/general.yaml | 2 + autofit/non_linear/search/abstract_search.py | 13 + autofit/non_linear/search/fit_context.py | 5 + .../non_linear/search/mcmc/emcee/search.py | 12 + autofit/non_linear/search/mle/bfgs/search.py | 14 +- .../search/nest/dynesty/search/abstract.py | 18 + .../non_linear/search/nest/nautilus/search.py | 22 + autofit/non_linear/status/__init__.py | 36 ++ autofit/non_linear/status/probe.py | 275 +++++++++++ autofit/non_linear/status/record.py | 450 ++++++++++++++++++ autofit/non_linear/status/signals.py | 103 ++++ autofit/non_linear/status/states.py | 110 +++++ autofit/non_linear/status/writer.py | 148 ++++++ test_autofit/non_linear/status/__init__.py | 0 .../non_linear/status/test_fit_status.py | 325 +++++++++++++ test_autofit/non_linear/status/test_units.py | 230 +++++++++ 16 files changed, 1762 insertions(+), 1 deletion(-) create mode 100644 autofit/non_linear/status/__init__.py create mode 100644 autofit/non_linear/status/probe.py create mode 100644 autofit/non_linear/status/record.py create mode 100644 autofit/non_linear/status/signals.py create mode 100644 autofit/non_linear/status/states.py create mode 100644 autofit/non_linear/status/writer.py create mode 100644 test_autofit/non_linear/status/__init__.py create mode 100644 test_autofit/non_linear/status/test_fit_status.py create mode 100644 test_autofit/non_linear/status/test_units.py diff --git a/autofit/config/general.yaml b/autofit/config/general.yaml index a0cf64428..177f93ab7 100644 --- a/autofit/config/general.yaml +++ b/autofit/config/general.yaml @@ -7,6 +7,8 @@ hpc: hpc_mode: false # If True, use HPC mode, which disables GUI visualization, logging to screen and other settings which are not suited to running on a super computer. iterations_per_quick_update: 1e99 # Non-linear search iterations between every quick update, which just displays the maximum likelihood model fit. iterations_per_full_update: 1e99 # Non-linear search iterations between every full update, which outputs all visuals and result fits (e.g. model.result, search.summary), this exits the search and can be slow. + status: true # If True, every fit writes a status record beside its output folder (.status.json, at most 2 KB; .status.jsonl attempt history; .traceback.txt on failure) and flushes its checkpoint on SIGUSR1 / SIGTERM. If False, nothing is written and no signal handlers are installed. + status_interval_s: 60 # Seconds between the status record's cadence writes while a fit runs (it is also written at start, at every full update, on a signal, on completion and on failure). inversion: check_reconstruction: true # If True, the inversion's reconstruction is checked to ensure the solution of a meshs's mapper is not an invalid solution where the values are all the same. reconstruction_vmax_factor: 0.5 # Plots of an Inversion's reconstruction use the reconstructed data's bright value multiplied by this factor. diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index 04ac88544..5da76bff1 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -44,6 +44,7 @@ from autofit.non_linear.search.execution import SearchExecutionMixin, _SearchMeta from autofit.non_linear.search.internal_state import SearchInternalStateMixin from autofit.non_linear.search.updates import SearchUpdatesMixin +from autofit.non_linear.status import status_notify, status_recorded from autofit.non_linear.search.test_mode_bypass import TestModeBypass, TestModeMixin from autofit.non_linear.test_mode import test_mode_level, skip_fit_output @@ -215,6 +216,7 @@ def __getstate__(self): state.pop("_logger", None) state.pop("paths", None) state.pop("_fit_pools", None) + state.pop("_fit_status", None) return state @property @@ -286,6 +288,11 @@ class represented by model M and gives a score for their fitness: the search updated model whose priors represent the beliefs produced by this fit. Raises ``AssertionError`` if the model has 0 dimensions. """ + with status_recorded(self): + return self._fit_and_output(model=model, analysis=analysis, info=info) + + def _fit_and_output(self, model, analysis, info=None): + """The body of ``fit``, inside its status record (``status_recorded``).""" self.check_model(model=model) # Per-fit: set by `start_resume_fit` when the fit's pools are resolved, and read @@ -394,6 +401,8 @@ def start_resume_fit(self, analysis: Analysis, model: AbstractPriorModel) -> Res if self.paths.has_timer: self.timer.start() + status_notify(self, "start") + mode = test_mode_level() if mode >= 2: return TestModeBypass(self).run( @@ -494,6 +503,8 @@ def _complete_fit(self, search_internal): self.paths.completed() + status_notify(self, "completed", search_internal) + def result_via_completed_fit( self, analysis: Analysis, @@ -511,6 +522,8 @@ def result_via_completed_fit( recreates the output files of the model-fit. """ + status_notify(self, "skipped") + model.freeze() samples_summary = self.paths.load_samples_summary() try: diff --git a/autofit/non_linear/search/fit_context.py b/autofit/non_linear/search/fit_context.py index d7029e1ae..14c8d376d 100644 --- a/autofit/non_linear/search/fit_context.py +++ b/autofit/non_linear/search/fit_context.py @@ -28,6 +28,7 @@ from autofit.non_linear.samples.adapter import RawSamples, samples_from_raw from autofit.non_linear.search.seed import seed_sequence +from autofit.non_linear.status import status_notify class UpdateSchedule: @@ -141,6 +142,10 @@ def update(self, internal: Any): if resume_state is not None: resume_state.save(self.paths, internal) + # The fit's status record takes the search's `progress` and keeps `internal` + # for a SIGUSR1 / SIGTERM checkpoint flush (`autofit.non_linear.status`). + status_notify(self._search, "update", internal) + return self._search.perform_update( model=self.model, analysis=self._analysis, diff --git a/autofit/non_linear/search/mcmc/emcee/search.py b/autofit/non_linear/search/mcmc/emcee/search.py index 6dcf23476..41e802ae1 100644 --- a/autofit/non_linear/search/mcmc/emcee/search.py +++ b/autofit/non_linear/search/mcmc/emcee/search.py @@ -139,6 +139,18 @@ def __init__( self.logger.debug("Creating Emcee Search") + def progress(self, search_internal) -> Optional[Dict]: + """ + The run's progress for the status record (``autofit.non_linear.status``): the + steps taken of ``nsteps`` (a run that converges early stops before). + """ + iteration = int(search_internal.iteration) + return { + "n_like": iteration * self.nwalkers, + "n_like_max": self.nsteps * self.nwalkers, + "fraction": min(1.0, iteration / self.nsteps) if self.nsteps else None, + } + def run(self, ctx): """ Run emcee (the ``run(ctx)`` hook; see ``docs/design/run_ctx.md``) and return the diff --git a/autofit/non_linear/search/mle/bfgs/search.py b/autofit/non_linear/search/mle/bfgs/search.py index d1c7fef7e..c50cfd46f 100644 --- a/autofit/non_linear/search/mle/bfgs/search.py +++ b/autofit/non_linear/search/mle/bfgs/search.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Optional, TYPE_CHECKING +from typing import Dict, Optional, TYPE_CHECKING from autofit import exc @@ -224,6 +224,18 @@ def fitness_overrides(self, analysis: Analysis) -> dict: """ return {"store_history": self.should_plot_start_point} + def progress(self, search_internal) -> Optional[Dict]: + """ + The run's progress for the status record (``autofit.non_linear.status``): the + iterations taken of ``maxiter`` (scipy stops earlier once converged). + """ + total = search_internal.get("total_iterations", search_internal.nit) + return { + "n_like": int(total), + "n_like_max": int(self.maxiter), + "fraction": min(1.0, total / self.maxiter) if self.maxiter else None, + } + def run(self, ctx): """ Run scipy's ``minimize`` (the ``run(ctx)`` hook; see ``docs/design/run_ctx.md``) diff --git a/autofit/non_linear/search/nest/dynesty/search/abstract.py b/autofit/non_linear/search/nest/dynesty/search/abstract.py index cef8a86f3..bac777866 100644 --- a/autofit/non_linear/search/nest/dynesty/search/abstract.py +++ b/autofit/non_linear/search/nest/dynesty/search/abstract.py @@ -451,6 +451,24 @@ def search_internal(self): """ return self.resume_state.load(self.paths) + def progress(self, search_internal) -> Optional[Dict]: + """ + The run's progress for the status record (``autofit.non_linear.status``): the + likelihood calls so far, the evidence and, when ``maxcall`` bounds the run, the + share of it spent (``fraction``). + """ + progress = { + "n_like": self.total_calls_from(search_internal=search_internal), + "n_like_max": self.maxcall, + } + try: + progress["log_z"] = float(search_internal.results.logz[-1]) + except Exception: + pass + if self.maxcall: + progress["fraction"] = min(1.0, progress["n_like"] / self.maxcall) + return progress + def iterations_from( self, search_internal: "Union[NestedSampler, DynamicNestedSampler]" ) -> Tuple[int, int]: diff --git a/autofit/non_linear/search/nest/nautilus/search.py b/autofit/non_linear/search/nest/nautilus/search.py index 00152e2c2..2d1244274 100644 --- a/autofit/non_linear/search/nest/nautilus/search.py +++ b/autofit/non_linear/search/nest/nautilus/search.py @@ -857,6 +857,28 @@ def call_search(self, search_internal, model, analysis, fitness, update=None): return search_internal + def progress(self, search_internal) -> Optional[Dict]: + """ + The run's progress for the status record (``autofit.non_linear.status``): the + sampler's likelihood calls, live points, evidence, live-set fraction and + effective sample size; ``fraction`` is the share of ``n_like_max`` spent (none + when it is unbounded, since Nautilus stops on ``n_eff`` / ``f_live``). + """ + progress = { + "n_like": int(search_internal.n_like), + "n_like_max": None, + "n_live": int(search_internal.n_live), + } + for key in ("log_z", "f_live", "n_eff"): + try: + progress[key] = float(getattr(search_internal, key)) + except Exception: + pass + if self.n_like_max is not None and self.n_like_max != float("inf"): + progress["n_like_max"] = int(self.n_like_max) + progress["fraction"] = min(1.0, progress["n_like"] / self.n_like_max) + return progress + def iterations_from( self, search_internal, iterations_per_full_update: Optional[float] = None ) -> Tuple[int, int]: diff --git a/autofit/non_linear/status/__init__.py b/autofit/non_linear/status/__init__.py new file mode 100644 index 000000000..efaddf31b --- /dev/null +++ b/autofit/non_linear/status/__init__.py @@ -0,0 +1,36 @@ +""" +The per-fit status record (hpc-campaign phase 1). + +Every fit with an output folder writes ``.status.json`` beside that folder +(so it survives ``output.remove_files``): a compact (at most 2 KB), atomically replaced +document with the fit's state, attempt, progress, resources and scheduler job, written +at start, every ``general.hpc.status_interval_s`` seconds, at each full update, on +SIGUSR1 / SIGTERM, on completion and on failure. ``.status.jsonl`` keeps +the attempt history and ``.traceback.txt`` the last failure's traceback. +``general.hpc.status: false`` turns all of it off. + +The ``progress`` protocol: a search may define ``progress(search_internal)``, returning +a dict of whichever of ``n_like``, ``n_like_max``, ``n_live``, ``log_z``, +``max_log_likelihood``, ``f_live``, ``n_eff`` and ``fraction`` (0-1) it knows, or +``None``. It is called at each full update (``FitContext.update``); the record adds +``eta_s`` from the fraction's rate. + +``states`` is the single vocabulary of run states, failure classes and product names. +""" +from autofit.non_linear.status.states import ( + COLLECTOR_STATES, + TERMINAL_STATES, + WRITTEN_STATES, + FailureClass, + Product, + RunState, + failure_class_of, +) +from autofit.non_linear.status.writer import MAX_BYTES, SCHEMA, StatusWriter +from autofit.non_linear.status.record import ( + FitStatus, + status_enabled, + status_notify, + status_recorded, + status_interval_s, +) diff --git a/autofit/non_linear/status/probe.py b/autofit/non_linear/status/probe.py new file mode 100644 index 000000000..7aa88023d --- /dev/null +++ b/autofit/non_linear/status/probe.py @@ -0,0 +1,275 @@ +""" +Cheap, fail-soft probes of the resources a fit is using and of the scheduler job it runs +in, for the per-fit status record. + +Every probe returns ``None`` (or omits a field) when its source is unavailable -- macOS, +Windows, a laptop outside any scheduler, a container without cgroup v2 -- and never +raises. Nothing here is a hard dependency: NVML (``pynvml``) is used only when it is +importable and a GPU is allocated to the process. +""" +import logging +import os +import time +from pathlib import Path +from typing import Dict, Optional + +logger = logging.getLogger(__name__) + +# Scheduler environment variables, per scheduler, mapped onto the status record's +# ``job`` fields. A scheduler is detected by the first variable of its mapping +# (``detect``); a new scheduler (PBS, LSF) is one more entry. +SCHEDULER_ENV = { + "slurm": { + "detect": "SLURM_JOB_ID", + "job_id": ("SLURM_ARRAY_JOB_ID", "SLURM_JOB_ID"), + "task_id": ("SLURM_ARRAY_TASK_ID",), + "node": ("SLURMD_NODENAME",), + "cpus": ("SLURM_CPUS_PER_TASK", "SLURM_CPUS_ON_NODE"), + "mem_mb": ("SLURM_MEM_PER_NODE",), + "mem_per_cpu_mb": ("SLURM_MEM_PER_CPU",), + "restart_count": ("SLURM_RESTART_COUNT",), + "start_time": ("SLURM_JOB_START_TIME",), + "end_time": ("SLURM_JOB_END_TIME",), + "gpus": ("SLURM_JOB_GPUS", "SLURM_STEP_GPUS"), + }, +} + + +def _first(env, names) -> Optional[str]: + for name in names: + value = env.get(name) + if value not in (None, ""): + return value + return None + + +def _int(value) -> Optional[int]: + try: + return int(str(value).split("(")[0]) + except (TypeError, ValueError): + return None + + +def job_from_env(env=None) -> Optional[Dict]: + """ + The scheduler job this process runs in, from the environment (``SCHEDULER_ENV``), + or ``None`` outside any known scheduler. + + ``time_limit_s`` is derived from the job's start and end times when the scheduler + exports them (Slurm >= 21.08); ``mem_mb`` falls back to the per-CPU memory times + the CPUs. + """ + env = os.environ if env is None else env + + for scheduler, mapping in SCHEDULER_ENV.items(): + if not env.get(mapping["detect"]): + continue + + def field(name): + return _first(env, mapping.get(name, ())) + + cpus = _int(field("cpus")) + mem_mb = _int(field("mem_mb")) + if mem_mb is None: + per_cpu = _int(field("mem_per_cpu_mb")) + if per_cpu is not None and cpus is not None: + mem_mb = per_cpu * cpus + + start, end = _int(field("start_time")), _int(field("end_time")) + time_limit_s = end - start if start is not None and end is not None else None + + return { + "scheduler": scheduler, + "job_id": field("job_id"), + "task_id": _int(field("task_id")), + "node": field("node"), + "cpus": cpus, + "mem_mb": mem_mb, + "time_limit_s": time_limit_s, + "restart_count": _int(field("restart_count")) or 0, + } + + return None + + +def gpu_allocated(env=None) -> bool: + """ + Whether a GPU is allocated to this process: a scheduler GPU list or a non-empty + ``CUDA_VISIBLE_DEVICES`` that does not hide every device. + """ + env = os.environ if env is None else env + for mapping in SCHEDULER_ENV.values(): + if _first(env, mapping.get("gpus", ())): + return True + visible = env.get("CUDA_VISIBLE_DEVICES") + return visible not in (None, "", "-1", "NoDevFiles") + + +def available_cpus() -> Optional[int]: + try: + return len(os.sched_getaffinity(0)) + except (AttributeError, OSError): + return os.cpu_count() + + +_CPU_MODEL = [] + + +def cpu_model() -> Optional[str]: + """ + The CPU model name from ``/proc/cpuinfo`` (Linux), cached; ``None`` elsewhere. + """ + if not _CPU_MODEL: + model = None + try: + with open("/proc/cpuinfo") as f: + for line in f: + if line.lower().startswith("model name"): + model = line.split(":", 1)[1].strip() + break + except OSError: + pass + _CPU_MODEL.append(model) + return _CPU_MODEL[0] + + +def cpu_seconds() -> Optional[float]: + """ + User + system CPU seconds of this process and its reaped children + (``resource.getrusage``). Children still running (e.g. a live pool's workers) are + counted only once joined. + """ + try: + import resource + + total = 0.0 + for who in (resource.RUSAGE_SELF, resource.RUSAGE_CHILDREN): + usage = resource.getrusage(who) + total += usage.ru_utime + usage.ru_stime + return total + except Exception: + try: + times = os.times() + return times.user + times.system + times.children_user + times.children_system + except Exception: + return None + + +def _cgroup_memory_peak() -> Optional[int]: + """ + The cgroup v2 ``memory.peak`` of this process's cgroup, in bytes (Linux only). + """ + try: + with open("/proc/self/cgroup") as f: + for line in f: + if line.startswith("0::"): + relative = line.strip()[3:].lstrip("/") + peak = Path("/sys/fs/cgroup") / relative / "memory.peak" + return int(peak.read_text().strip()) + except (OSError, ValueError): + pass + return None + + +def peak_rss(use_cgroup: bool = True) -> Dict: + """ + The peak resident memory, in MB, and where it came from: ``cgroup`` (cgroup v2 + ``memory.peak``, which covers every process of the job step) else ``getrusage`` + (the larger of this process and its largest reaped child). + + ``use_cgroup=False`` skips the cgroup: outside a scheduler job a process's cgroup + is a login session or a whole WSL / container instance, whose peak is not the + fit's. + """ + peak = _cgroup_memory_peak() if use_cgroup else None + if peak is not None: + return {"peak_rss_mb": round(peak / 2**20), "peak_rss_source": "cgroup"} + + try: + import resource + import sys + + maxrss = max( + resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, + resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss, + ) + # Kilobytes on Linux, bytes on macOS. + scale = 2**20 if sys.platform == "darwin" else 2**10 + return {"peak_rss_mb": round(maxrss / scale), "peak_rss_source": "getrusage"} + except Exception: + return {"peak_rss_mb": None, "peak_rss_source": None} + + +_NVML = [] + + +def gpu() -> Optional[Dict]: + """ + The first visible GPU's model, used memory (MB) and utilisation (%), through NVML; + ``None`` when no GPU is allocated or ``pynvml`` is not importable. + """ + if not gpu_allocated(): + return None + try: + import pynvml + + if not _NVML: + pynvml.nvmlInit() + _NVML.append(True) + handle = pynvml.nvmlDeviceGetHandleByIndex(0) + name = pynvml.nvmlDeviceGetName(handle) + if isinstance(name, bytes): + name = name.decode() + memory = pynvml.nvmlDeviceGetMemoryInfo(handle) + utilisation = pynvml.nvmlDeviceGetUtilizationRates(handle) + return { + "model": name, + "mem_used_mb": round(memory.used / 2**20), + "util": utilisation.gpu, + } + except Exception: + return None + + +class ResourceProbe: + def __init__(self, cpus: Optional[int] = None, in_job: bool = False): + """ + Resource use of one fit attempt: wall and CPU seconds are measured from the + probe's creation (so earlier work in the same process is not counted), peak + memory is the process's (or the cgroup's) lifetime peak. + + Parameters + ---------- + cpus + The CPUs the job was given, for ``cpu_efficiency``; defaults to the CPUs + this process may run on. + in_job + Whether the process runs in a scheduler job, whose cgroup is the job's + (so its ``memory.peak`` is used). + """ + self.cpus = cpus or available_cpus() + self.in_job = in_job + self._wall_start = time.monotonic() + self._cpu_start = cpu_seconds() + + def resources(self) -> Dict: + wall_s = time.monotonic() - self._wall_start + cpu_now = cpu_seconds() + cpu_s = ( + cpu_now - self._cpu_start + if cpu_now is not None and self._cpu_start is not None + else None + ) + efficiency = ( + round(cpu_s / (wall_s * self.cpus), 3) + if cpu_s is not None and self.cpus and wall_s > 0 + else None + ) + return { + "wall_s": round(wall_s, 1), + "cpu_s": None if cpu_s is None else round(cpu_s, 1), + "cpu_efficiency": efficiency, + **peak_rss(use_cgroup=self.in_job), + "gpu": gpu(), + "cpu_model": cpu_model(), + } diff --git a/autofit/non_linear/status/record.py b/autofit/non_linear/status/record.py new file mode 100644 index 000000000..e8e5fdb20 --- /dev/null +++ b/autofit/non_linear/status/record.py @@ -0,0 +1,450 @@ +""" +``FitStatus``: the lifecycle of one fit's status record. + +``NonLinearSearch.fit`` creates one per fit (``FitStatus.for_search``, ``None`` when +``general.hpc.status`` is off or the paths store nothing on disk) and the search calls +it at each transition: + +- ``start`` when the search starts or resumes (``start_resume_fit``): a new attempt, + state ``running``, the cadence thread and the signal handlers; +- ``update`` at every ``FitContext.update`` (the sampler's full update): the sampler's + ``progress`` and the internal state a signal flushes; +- ``completed`` when the fit is marked complete (``_complete_fit``); +- ``skipped`` when the fit was already complete; +- ``failed`` when an exception leaves ``fit`` (the exception is re-raised unchanged); +- ``timeout_pending`` on SIGUSR1 / SIGTERM (``autofit.non_linear.status.signals``); +- ``close`` on every exit of ``fit``. + +Every write is best effort: a status record that cannot be written logs a warning and +never fails the fit. +""" +import logging +import threading +import time +import traceback +import zipfile +from contextlib import contextmanager +from datetime import datetime, timezone +from pathlib import Path +from typing import Dict, Optional + +from autonerves import conf + +from autofit.non_linear.status import probe, signals +from autofit.non_linear.status.states import ( + FailureClass, + Product, + RunState, + failure_class_of, +) +from autofit.non_linear.status.writer import SCHEMA, StatusWriter + +logger = logging.getLogger(__name__) + +DEFAULT_INTERVAL_S = 60.0 + + +def _hpc_setting(name: str, default): + try: + value = conf.instance["general"]["hpc"][name] + except Exception: + return default + return default if value is None else value + + +def status_enabled() -> bool: + """ + Whether fits write a status record (``general.hpc.status``, default on). + """ + return bool(_hpc_setting("status", True)) + + +def status_interval_s() -> float: + """ + Seconds between cadence writes (``general.hpc.status_interval_s``, default 60). + """ + return float(_hpc_setting("status_interval_s", DEFAULT_INTERVAL_S)) + + +def _now() -> str: + return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + + +class FitStatus: + @classmethod + def for_search(cls, search) -> Optional["FitStatus"]: + """ + The status record of ``search``'s next fit, or ``None`` when status records are + off (``general.hpc.status: false``) or the search writes no output folder + (``NullPaths``, ``DatabasePaths``). + """ + from autofit.non_linear.paths.directory import DirectoryPaths + + if not status_enabled(): + return None + if not isinstance(getattr(search, "paths", None), DirectoryPaths): + return None + return cls(search) + + def __init__(self, search, interval_s: Optional[float] = None): + """ + The status record of one fit of ``search``. Nothing is written until the first + transition, since the output path is only known once ``fit`` has set the + model on the paths. + """ + self.search = search + self.interval_s = status_interval_s() if interval_s is None else interval_s + # A failure before `fit` sets this fit's model on the paths must not be + # recorded: the paths would still name a previous fit's output. + self._model_at_entry = getattr(search.paths, "model", None) + + self.writer: Optional[StatusWriter] = None + self.state: Optional[RunState] = None + self.attempt: Optional[int] = None + self.failure_class: Optional[FailureClass] = None + self.progress: Optional[Dict] = None + self.last_internal = None + self.job = None + self.started: Optional[str] = None + self.n_writes = 0 + + self._probe = None + self._progress_origin = None + self._stop = threading.Event() + self._thread = None + self._restore_signals = None + self._keep_existing = False + self._dirty = False + # Set while a signal handler runs: writes then wait at most this long for the + # writer's lock, which the interrupted main thread may hold. + self._lock_timeout = None + + @property + def paths(self): + return self.search.paths + + def _begin(self): + """ + Open the record for a new attempt: its number (from the history), the job and + the resource baseline, and the history's ``start`` line. + """ + if self.writer is not None: + return + self.writer = StatusWriter(self.paths.output_path) + self.attempt = self.writer.next_attempt() + self.job = probe.job_from_env() + self._probe = probe.ResourceProbe( + cpus=(self.job or {}).get("cpus"), in_job=self.job is not None + ) + self.started = _now() + self.writer.append( + {"event": "start", "attempt": self.attempt, "time": self.started, "job": self.job} + ) + + def start(self): + """ + The search starts (or resumes): write ``running`` and start the cadence writes + and the signal handlers. Idempotent within one fit. + """ + if self.state is not None: + return + try: + self._begin() + except Exception as exception: + logger.warning(f"Status record not started ({exception}).") + return + + self.state = RunState.RUNNING + self.write() + + self._restore_signals = signals.install(self) + + if self.interval_s and self.interval_s > 0: + self._thread = threading.Thread( + target=self._beat, name="autofit-status", daemon=True + ) + self._thread.start() + + def _beat(self): + while not self._stop.wait(self.interval_s): + self.write() + + def update(self, search_internal): + """ + A full update of the run: keep ``search_internal`` (what a signal flushes) and + the sampler's ``progress`` (``None`` when the search does not implement it), + then write. + """ + if self.state is None: + return + self.last_internal = search_internal + self.progress = self._progress_of(search_internal) + self.write() + + def _transition(self, state: RunState, **entry): + self.state = state + self.write() + if self.writer is not None: + try: + self.writer.append( + {"event": str(state), "attempt": self.attempt, "time": _now(), **entry}, + timeout=self._lock_timeout, + ) + except Exception as exception: + logger.warning(f"Status history not written ({exception}).") + + def _progress_of(self, search_internal): + progress = getattr(self.search, "progress", None) + if progress is None or search_internal is None: + return self.progress + try: + return progress(search_internal) + except Exception: + return None + + def completed(self, search_internal=None): + """ + The fit is marked complete; ``search_internal`` gives the final progress. + """ + if self.writer is None: + return + self.progress = self._progress_of(search_internal) + if self.progress is not None and self.progress.get("fraction") is not None: + self.progress["fraction"] = 1.0 + self.progress["eta_s"] = 0 + self._stop.set() + self._transition(RunState.COMPLETED) + + def skipped(self): + """ + The fit was already complete, so the search did not run. This is not a new + attempt: the history gains a ``skipped`` line under the last attempt's number. + A record that already says ``completed`` is kept (it holds the run's + resources); any other record (e.g. one written before status records existed) + becomes ``skipped``. + """ + try: + self.writer = StatusWriter(self.paths.output_path) + self.attempt = max(self.writer.next_attempt() - 1, 1) + existing = self.writer.read() or {} + if existing.get("state") == str(RunState.COMPLETED): + self._keep_existing = True + self.state = RunState.SKIPPED + self.writer.append( + {"event": str(RunState.SKIPPED), "attempt": self.attempt, "time": _now()} + ) + return + self.started = existing.get("started") + except Exception as exception: + logger.warning(f"Status record not written ({exception}).") + return + self._transition(RunState.SKIPPED) + + def failed(self, exception: BaseException): + """ + An exception left ``fit``: write its traceback beside the record, then + ``failed`` with its failure class. The caller re-raises. + + An exception raised before ``fit`` set its model on the paths (e.g. a model + with no free parameters) has no output path of its own and is not recorded. + """ + if self.writer is None and ( + getattr(self.paths, "model", None) is self._model_at_entry + ): + return + try: + self._begin() + except Exception as error: + logger.warning(f"Failure record not written ({error}).") + return + self._stop.set() + self.failure_class = failure_class_of(exception) + try: + self.writer.write_traceback( + "".join( + traceback.format_exception( + type(exception), exception, exception.__traceback__ + ) + ) + ) + except Exception as error: + logger.warning(f"Traceback not written ({error}).") + self._transition( + RunState.FAILED, + failure_class=str(self.failure_class), + error=f"{type(exception).__name__}: {exception}"[:200], + ) + + def flush_checkpoint(self): + """ + Save the search's resume state from the last internal state an update handed + over (``Checkpointer.save``; a native resume state, which the backend writes + itself, is a no-op). + """ + resume_state = getattr(self.search, "resume_state", None) + if resume_state is None or self.last_internal is None: + return + try: + resume_state.save(self.paths, self.last_internal) + except Exception as exception: + logger.warning(f"Checkpoint flush on signal failed ({exception}).") + + def timeout_pending(self, signal_name: str): + """ + SIGUSR1 / SIGTERM: flush the checkpoint, then write ``timeout_pending``. + """ + if self.state is None or self.state in (RunState.COMPLETED, RunState.FAILED): + return + self._lock_timeout = 1.0 + try: + self.flush_checkpoint() + self._transition(RunState.TIMEOUT_PENDING, signal=signal_name) + finally: + self._lock_timeout = None + + def close(self): + """ + End of ``fit`` (every exit): stop the cadence writes and restore the signal + handlers. A completed record is written once more, since the output is only + archived to its ``.zip`` after completion. + """ + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=5) + self._thread = None + if self._restore_signals is not None: + self._restore_signals() + self._restore_signals = None + if self._dirty or ( + self.state in (RunState.COMPLETED, RunState.SKIPPED) + and not self._keep_existing + ): + self._dirty = False + self.write() + + def _eta_s(self, fraction) -> Optional[int]: + """ + Seconds left, from this attempt's rate of progress (so a resumed run's earlier + progress does not inflate the rate). + """ + now = time.monotonic() + if self._progress_origin is None: + self._progress_origin = (now, fraction) + return None + t0, f0 = self._progress_origin + if fraction <= f0 or fraction >= 1: + return None + return int((now - t0) * (1 - fraction) / (fraction - f0)) + + def _outputs(self) -> Dict: + output_path = Path(self.paths.output_path) + zip_path = Path(f"{output_path}{Product.ZIP}") + outputs = {"zip": None, "zip_complete": None, "bytes": None} + if zip_path.exists(): + outputs["zip"] = str(zip_path) + outputs["bytes"] = zip_path.stat().st_size + try: + with zipfile.ZipFile(zip_path) as f: + outputs["zip_complete"] = str(Product.COMPLETED) in f.namelist() + except (OSError, zipfile.BadZipFile): + outputs["zip_complete"] = False + for candidate in ( + output_path / "files" / str(Product.DATASET_IDENTITY), + output_path / str(Product.DATASET_IDENTITY), + ): + if candidate.exists(): + outputs["dataset_identity"] = str(candidate) + break + if self.state == RunState.FAILED and self.writer.traceback_path.exists(): + outputs["traceback"] = str(self.writer.traceback_path) + return outputs + + def document(self) -> Dict: + """ + The status document (``schema`` 1) as it stands now. + """ + paths = self.paths + + def safe(attribute): + try: + value = getattr(paths, attribute) + except Exception: + return None + return None if value is None else str(value) + + progress = dict(self.progress) if self.progress else None + if progress is not None and progress.get("fraction") is not None: + progress.setdefault("eta_s", self._eta_s(progress["fraction"])) + + return { + "schema": SCHEMA, + "identifier": safe("identifier"), + "name": safe("name"), + "path_prefix": safe("path_prefix"), + "search": type(self.search).__name__, + "state": None if self.state is None else str(self.state), + "attempt": self.attempt, + "failure_class": None + if self.failure_class is None + else str(self.failure_class), + "job": self.job, + "progress": progress, + "resources": self._probe.resources() if self._probe else None, + "outputs": self._outputs(), + "updated": _now(), + "started": self.started, + } + + def write(self): + """ + Write the current document; best effort. Inside a signal handler the wait for + the writer's lock is bounded (the interrupted main thread may hold it); a write + that gives up is redone after the write it interrupted. + """ + if self.writer is None or self.state is None: + return + try: + written = self.writer.write(self.document(), timeout=self._lock_timeout) + except Exception as exception: + logger.warning(f"Status record not written ({exception}).") + return + if not written: + self._dirty = True + return + self.n_writes += 1 + if self._dirty and self._lock_timeout is None: + self._dirty = False + self.write() + + +@contextmanager +def status_recorded(search): + """ + Run one ``fit`` of ``search`` inside its status record: the record is kept on + ``search._fit_status`` for the fit's transitions (``status_notify``), an exception + leaving the fit is recorded (``failed`` + traceback) and re-raised unchanged, and + the record is closed on every exit. A fit nested in another fit of the same + search restores the outer record afterwards. + """ + outer = getattr(search, "_fit_status", None) + fit_status = search._fit_status = FitStatus.for_search(search) + try: + yield fit_status + except Exception as exception: + if fit_status is not None: + fit_status.failed(exception) + raise + finally: + if fit_status is not None: + fit_status.close() + search._fit_status = outer + + +def status_notify(search, transition: str, *args): + """ + Call ``transition`` (``start``, ``update``, ``completed``, ``skipped``) on the + status record of ``search``'s current fit, if it has one. + """ + fit_status = getattr(search, "_fit_status", None) + if fit_status is not None: + getattr(fit_status, transition)(*args) diff --git a/autofit/non_linear/status/signals.py b/autofit/non_linear/status/signals.py new file mode 100644 index 000000000..228e187ae --- /dev/null +++ b/autofit/non_linear/status/signals.py @@ -0,0 +1,103 @@ +""" +SIGUSR1 / SIGTERM handling for a fit's status record. + +A batch scheduler warns a job before its time limit with a signal (Slurm +``--signal=B:USR1@``) and stops it with SIGTERM. On either the handler flushes +the search's checkpoint (``FitStatus.flush_checkpoint``) and writes +``timeout_pending``; on SIGTERM it then exits with status 143 (128 + SIGTERM), raised as +``SystemExit`` so the fit's ``finally`` blocks (pool shutdown, handler restore) run. + +Handlers are installed only in the main thread of the top-level process (never in a +pool worker or a thread), and the previous handlers are restored when the fit ends. A +forked worker inherits the handlers; there they hand the signal to the previous +handler (or the default action) unchanged, so workers behave exactly as they did +without a status record. +""" +import logging +import multiprocessing +import os +import signal +import sys +import threading +from typing import Callable + +logger = logging.getLogger(__name__) + +EXIT_SIGTERM = 128 + 15 + + +def _signals(): + names = ("SIGUSR1", "SIGTERM") + return [getattr(signal, name) for name in names if hasattr(signal, name)] + + +def can_install() -> bool: + """ + Whether this is the main thread of a top-level process on a platform with the + signals (``SIGUSR1`` does not exist on Windows). + """ + return ( + hasattr(signal, "SIGUSR1") + and threading.current_thread() is threading.main_thread() + and multiprocessing.parent_process() is None + ) + + +def _pass_on(signum, frame, previous): + """ + Hand ``signum`` to ``previous``: call it if it is a handler, else restore it and + re-raise the signal so its default action (or ignore) applies. + """ + if callable(previous): + previous(signum, frame) + return + signal.signal(signum, previous if previous is not None else signal.SIG_DFL) + os.kill(os.getpid(), signum) + + +def install(status) -> Callable[[], None]: + """ + Install the SIGUSR1 / SIGTERM handlers of ``status`` (a ``FitStatus``) and return + the function that restores the previous ones; a no-op when ``can_install`` is + false. + """ + if not can_install(): + return lambda: None + + owner = os.getpid() + previous = {} + + def handler(signum, frame): + if os.getpid() != owner: + _pass_on(signum, frame, previous.get(signum)) + return + + name = signal.Signals(signum).name + logger.info(f"{name} received: flushing the checkpoint (timeout_pending).") + status.timeout_pending(name) + + if signum == signal.SIGTERM: + sys.exit(EXIT_SIGTERM) + + if callable(previous.get(signum)): + previous[signum](signum, frame) + + for signum in _signals(): + try: + previous[signum] = signal.signal(signum, handler) + except (ValueError, OSError): + pass + + def restore(): + if threading.current_thread() is not threading.main_thread(): + return + for signum, handler_before in previous.items(): + try: + signal.signal( + signum, + handler_before if handler_before is not None else signal.SIG_DFL, + ) + except (ValueError, OSError): + pass + + return restore diff --git a/autofit/non_linear/status/states.py b/autofit/non_linear/status/states.py new file mode 100644 index 000000000..154638311 --- /dev/null +++ b/autofit/non_linear/status/states.py @@ -0,0 +1,110 @@ +""" +The single vocabulary of run states, failure classes and products shared by the per-fit +status record (``autofit.non_linear.status``), the results library and any campaign +collector. + +A fit writes only the states in ``WRITTEN_STATES``. ``timeout``, ``oom`` and +``node_fail`` are assigned by a collector from the scheduler's accounting (e.g. +``sacct``), since a process killed by the scheduler cannot record its own death. + +All three enums subclass ``str``, so a member compares equal to (and serialises as) its +plain value: ``RunState.RUNNING == "running"``. +""" +from enum import Enum + + +class RunState(str, Enum): + PENDING = "pending" + """Written by a batch trap before Python starts; never by the fit itself.""" + RUNNING = "running" + """The fit is executing (written at start and at every cadence write).""" + COMPLETED = "completed" + """The fit finished and its output is marked complete.""" + FAILED = "failed" + """A Python exception left ``fit``; ``failure_class`` says which kind.""" + TIMEOUT_PENDING = "timeout_pending" + """SIGUSR1 / SIGTERM was received and the checkpoint flushed: the job is about to + hit its time limit (or be stopped) and should be resubmitted to resume.""" + SKIPPED = "skipped" + """The fit was already complete, so the search did not run again.""" + + # Collector-only states, assigned from scheduler accounting. + TIMEOUT = "timeout" + OOM = "oom" + NODE_FAIL = "node_fail" + + def __str__(self): + return self.value + + +WRITTEN_STATES = frozenset( + { + RunState.PENDING, + RunState.RUNNING, + RunState.COMPLETED, + RunState.FAILED, + RunState.TIMEOUT_PENDING, + RunState.SKIPPED, + } +) + +COLLECTOR_STATES = frozenset({RunState.TIMEOUT, RunState.OOM, RunState.NODE_FAIL}) + +TERMINAL_STATES = frozenset( + { + RunState.COMPLETED, + RunState.FAILED, + RunState.SKIPPED, + RunState.TIMEOUT, + RunState.OOM, + RunState.NODE_FAIL, + } +) + + +class FailureClass(str, Enum): + EXCEPTION = "exception" + """Any exception not classified more precisely (the default).""" + BAD_INPUT = "bad_input" + """The inputs of the fit are invalid (data, model or settings).""" + MISSING_SEED = "missing_seed" + """A result this fit is chained from (its seed) does not exist.""" + UPSTREAM_INCOMPLETE = "upstream_incomplete" + """A result this fit is chained from exists but is not complete.""" + + def __str__(self): + return self.value + + +class Product(str, Enum): + """ + The files a fit's status record names, relative to its ``output_path``. + """ + + STATUS = ".status.json" + """Suffix of ``.status.json``, the current status record.""" + HISTORY = ".status.jsonl" + """Suffix of ``.status.jsonl``, the append-only attempt history.""" + TRACEBACK = ".traceback.txt" + """Suffix of ``.traceback.txt``, the traceback of a failed attempt.""" + ZIP = ".zip" + """Suffix of ``.zip``, the fit's archived output.""" + COMPLETED = ".completed" + """The marker file inside ``output_path`` written when the fit completes.""" + DATASET_IDENTITY = "dataset_identity.json" + """The record of the fit's inputs (written by the results library).""" + + def __str__(self): + return self.value + + +def failure_class_of(exception: BaseException) -> FailureClass: + """ + The failure class of ``exception``: the ``failure_class`` attribute an exception may + carry (a ``FailureClass`` or its value), else ``FailureClass.EXCEPTION``. + """ + value = getattr(exception, "failure_class", None) + try: + return FailureClass(value) + except ValueError: + return FailureClass.EXCEPTION diff --git a/autofit/non_linear/status/writer.py b/autofit/non_linear/status/writer.py new file mode 100644 index 000000000..99e7645f6 --- /dev/null +++ b/autofit/non_linear/status/writer.py @@ -0,0 +1,148 @@ +""" +The files of a fit's status record, all siblings of its ``output_path`` so they survive +``output.remove_files`` (which deletes ``output_path`` and keeps only the ``.zip``): + +- ``.status.json``: the current status (``schema`` 1), at most + ``MAX_BYTES``, replaced atomically on every write; +- ``.status.jsonl``: the append-only history, one line per attempt start + and per state change (never per cadence write), from which the attempt number is + counted; +- ``.traceback.txt``: the traceback of the last failed attempt. +""" +import json +import threading +from contextlib import contextmanager +from pathlib import Path +from typing import Dict, List, Optional + +from autofit.non_linear.status.states import Product +from autofit.tools.util import open_atomic + +SCHEMA = 1 +MAX_BYTES = 2048 + +# Optional fields dropped, in this order, should a status document exceed `MAX_BYTES` +# (e.g. a very long `path_prefix`): the least useful to a collector goes first. +_TRIM_ORDER = ( + ("progress", ("n_eff", "f_live", "n_live", "log_z", "max_log_likelihood")), + ("resources", ("gpu", "cpu_model")), + ("outputs", ("dataset_identity", "traceback")), + ("job", ("node", "time_limit_s")), +) + + +def _dumps(document: Dict) -> str: + return json.dumps(document, separators=(",", ":"), default=str) + + +def compact(document: Dict, max_bytes: int = MAX_BYTES) -> str: + """ + ``document`` as compact JSON of at most ``max_bytes``, dropping optional fields + (``_TRIM_ORDER``) until it fits, then shortening ``path_prefix`` if it still + does not. + """ + text = _dumps(document) + for section, keys in _TRIM_ORDER: + if len(text.encode()) <= max_bytes: + break + value = document.get(section) + if isinstance(value, dict): + for key in keys: + value.pop(key, None) + text = _dumps(document) + # Still too large: only a pathological `path_prefix` gets here. Keep its tail (the + # most specific part), marked as cut. + excess = len(text.encode()) - max_bytes + if excess > 0 and isinstance(document.get("path_prefix"), str): + document["path_prefix"] = "..." + document["path_prefix"][excess + 3 :] + text = _dumps(document) + return text + + +class StatusWriter: + def __init__(self, output_path): + """ + Reads and writes the status record of the fit whose output folder is + ``output_path`` (the record's files sit beside that folder, not in it). + + Writes are serialised by a lock: the cadence thread and the main thread (an + update, a signal, completion) may both write. A signal handler runs in the main + thread, possibly while that thread holds the lock, so it passes a ``timeout`` + and a write that cannot take the lock in time is skipped (returns ``None``) + rather than deadlocking. + """ + output_path = str(output_path) + self.status_path = Path(f"{output_path}{Product.STATUS}") + self.history_path = Path(f"{output_path}{Product.HISTORY}") + self.traceback_path = Path(f"{output_path}{Product.TRACEBACK}") + self._lock = threading.Lock() + + @contextmanager + def _locked(self, timeout: Optional[float]): + acquired = self._lock.acquire(timeout=-1 if timeout is None else timeout) + try: + yield acquired + finally: + if acquired: + self._lock.release() + + def write(self, document: Dict, timeout: Optional[float] = None) -> Optional[str]: + """ + Atomically replace ``status.json`` with ``document`` (compacted to + ``MAX_BYTES``), returning the text written (``None`` if the lock timed out). + """ + text = compact(document) + with self._locked(timeout) as acquired: + if not acquired: + return None + with open_atomic(self.status_path, "w") as f: + f.write(text) + return text + + def read(self) -> Optional[Dict]: + try: + return json.loads(self.status_path.read_text()) + except (OSError, ValueError): + return None + + def append(self, entry: Dict, timeout: Optional[float] = None) -> bool: + """ + Append one line to the ``status.jsonl`` history (``False`` if the lock timed + out). + """ + with self._locked(timeout) as acquired: + if not acquired: + return False + self.history_path.parent.mkdir(parents=True, exist_ok=True) + with open(self.history_path, "a") as f: + f.write(_dumps(entry) + "\n") + return True + + def history(self) -> List[Dict]: + """ + The history's entries, skipping a line a killed process left truncated. + """ + entries = [] + try: + with open(self.history_path) as f: + for line in f: + try: + entries.append(json.loads(line)) + except ValueError: + pass + except OSError: + pass + return entries + + def next_attempt(self) -> int: + """ + The number of the attempt about to start: one more than the attempts the + history records. + """ + return 1 + sum(1 for entry in self.history() if entry.get("event") == "start") + + def write_traceback(self, text: str) -> Path: + with self._locked(None): + with open_atomic(self.traceback_path, "w") as f: + f.write(text) + return self.traceback_path diff --git a/test_autofit/non_linear/status/__init__.py b/test_autofit/non_linear/status/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test_autofit/non_linear/status/test_fit_status.py b/test_autofit/non_linear/status/test_fit_status.py new file mode 100644 index 000000000..470abf6d2 --- /dev/null +++ b/test_autofit/non_linear/status/test_fit_status.py @@ -0,0 +1,325 @@ +""" +The status record of real (cheap) fits: written beside the output on start and +completion, an attempt history across a kill and a resume, a failure record, cadence +writes, SIGUSR1 / SIGTERM in a fitting process, and nothing at all when +``general.hpc.status`` is off. +""" +import json +import os +import signal +import subprocess +import sys +import textwrap +import time +from pathlib import Path + +import numpy as np +import pytest + +import autofit as af +from autofit.non_linear.status import MAX_BYTES, FitStatus, StatusWriter, record +from autofit.non_linear.paths.directory import DirectoryPaths +from autonerves import conf + +CONFIG = Path(__file__).parents[2] / "config" + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + + +@pytest.fixture(name="output", autouse=True) +def make_output(tmp_path): + conf.instance.push(new_path=str(CONFIG), output_path=str(tmp_path / "output")) + return tmp_path / "output" + + +def _analysis(): + return af.ex.Analysis(data=np.ones(20), noise_map=np.ones(20) * 0.1, use_jax=False) + + +def _fit(name="status_fit", analysis=None): + search = af.Drawer(name=name, total_draws=5) + search.fit(model=af.Model(af.ex.Gaussian), analysis=analysis or _analysis()) + return search + + +def _writer(search): + return StatusWriter(search.paths.output_path) + + +def test_completed_fit_writes_a_small_record_beside_its_output(): + search = _fit() + writer = _writer(search) + status = writer.read() + + assert writer.status_path.parent == Path(search.paths.output_path).parent + assert len(writer.status_path.read_bytes()) <= MAX_BYTES + assert status["schema"] == 1 + assert status["state"] == "completed" + assert status["attempt"] == 1 + assert status["failure_class"] is None + assert status["identifier"] == search.paths.identifier + assert status["search"] == "Drawer" + assert status["resources"]["wall_s"] >= 0 + # `remove_files: true` in the test config: only the zip is left, and the record + # (which sits beside the folder) survives it. + assert not Path(search.paths.output_path).exists() + assert status["outputs"]["zip_complete"] is True + + events = [entry["event"] for entry in writer.history()] + assert events == ["start", "completed"] + + +def test_rerun_of_a_completed_fit_is_skipped_not_a_new_attempt(): + _fit() + search = _fit() + writer = _writer(search) + + assert writer.read()["state"] == "completed" + history = writer.history() + assert [entry["event"] for entry in history] == ["start", "completed", "skipped"] + assert history[-1]["attempt"] == 1 + + +class BadInput(ValueError): + failure_class = "bad_input" + + +class FailingAnalysis(af.ex.Analysis): + def modify_before_fit(self, paths, model): + raise BadInput("the data are all ones") + + +def test_exception_writes_a_failure_record_and_is_reraised(): + search = af.Drawer(name="status_failed", total_draws=5) + analysis = FailingAnalysis( + data=np.ones(20), noise_map=np.ones(20) * 0.1, use_jax=False + ) + + with pytest.raises(BadInput, match="all ones"): + search.fit(model=af.Model(af.ex.Gaussian), analysis=analysis) + + writer = _writer(search) + status = writer.read() + assert status["state"] == "failed" + assert status["failure_class"] == "bad_input" + assert status["outputs"]["traceback"] == str(writer.traceback_path) + assert "BadInput: the data are all ones" in writer.traceback_path.read_text() + assert writer.history()[-1]["event"] == "failed" + assert search._fit_status is None + + +def test_failure_before_the_model_is_set_is_not_recorded(output): + search = _fit(name="status_reused") + writer = _writer(search) + + class NoParameters(Exception): + pass + + def check_model(model): + raise NoParameters() + + search.check_model = check_model + with pytest.raises(NoParameters): + search.fit(model=af.Model(af.ex.Exponential), analysis=_analysis()) + + # The completed fit's record (the paths' model is still the first fit's) is kept. + assert writer.read()["state"] == "completed" + assert not writer.traceback_path.exists() + + +def test_disabled_writes_nothing_and_leaves_the_output_unchanged(output, monkeypatch): + on = _fit(name="status_on") + on_files = sorted( + str(p.relative_to(output / "status_on")) + for p in (output / "status_on").rglob("*") + ) + + monkeypatch.setattr(record, "status_enabled", lambda: False) + off = _fit(name="status_off") + off_files = sorted( + str(p.relative_to(output / "status_off")) + for p in (output / "status_off").rglob("*") + ) + + assert FitStatus.for_search(off) is None + status_files = {f for f in on_files if ".status." in f or ".traceback." in f} + assert len(status_files) == 2 + assert [f for f in on_files if f not in status_files] == off_files + assert not any(".status." in f for f in off_files) + + +def test_cadence_writes_while_running(tmp_path): + search = af.Drawer(name="status_cadence") + search.paths = DirectoryPaths(name="status_cadence") + search.paths.model = af.Model(af.ex.Gaussian) + status = FitStatus(search, interval_s=0.05) + + status.start() + time.sleep(0.5) + status.close() + + # One write at start, then roughly one per interval. + assert status.n_writes >= 5 + writes = status.n_writes + time.sleep(0.15) + assert status.n_writes == writes + assert _writer(search).read()["state"] == "running" + + +def test_update_takes_progress_and_a_signal_flushes_the_checkpoint(): + saved = [] + + class ResumeState: + def save(self, paths, internal): + saved.append(internal) + + search = af.Drawer(name="status_flush") + search.paths.model = af.Model(af.ex.Gaussian) + search.resume_state = ResumeState() + search.progress = lambda internal: {"n_like": internal, "fraction": internal / 10} + status = FitStatus(search, interval_s=0) + + status.start() + status.update(2) + status.update(4) + assert _writer(search).read()["progress"]["n_like"] == 4 + + status.timeout_pending("SIGUSR1") + status.close() + + assert saved == [4] + document = _writer(search).read() + assert document["state"] == "timeout_pending" + assert document["progress"]["eta_s"] is not None + assert _writer(search).history()[-1] == { + "event": "timeout_pending", + "attempt": 1, + "time": document["updated"], + "signal": "SIGUSR1", + } + + +SCRIPT = """ +import os, sys, time +import numpy as np +import autofit as af +from autonerves import conf + +conf.instance.push(new_path={config!r}, output_path={output!r}) + + +class Slow(af.ex.Analysis): + def log_likelihood_function(self, instance, shared=None, xp=np): + time.sleep(float(os.environ.get("STATUS_SLEEP", "0"))) + return super().log_likelihood_function(instance, shared, xp) + + +search = af.Drawer(name="status_subprocess", total_draws=40) +search.fit( + model=af.Model(af.ex.Gaussian), + analysis=Slow(data=np.ones(20), noise_map=np.ones(20) * 0.1, use_jax=False), +) +""" + + +def _script(tmp_path, output): + path = tmp_path / "fit_script.py" + path.write_text(SCRIPT.format(config=str(CONFIG), output=str(output))) + return path + + +def _status_path(output): + found = list((output / "status_subprocess").glob("*.status.json")) + return found[0] if found else None + + +def _launch(script, sleep): + # The subprocess imports the autofit under test, not whichever is installed. + source = str(Path(af.__file__).parents[1]) + pythonpath = os.pathsep.join( + filter(None, [source, os.environ.get("PYTHONPATH")]) + ) + env = dict(os.environ, STATUS_SLEEP=str(sleep), PYTHONPATH=pythonpath) + return subprocess.Popen( + [sys.executable, str(script)], + env=env, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + ) + + +def _wait_for_state(output, state, process, timeout=120): + start = time.time() + while time.time() - start < timeout: + path = _status_path(output) + if path is not None: + try: + if json.loads(path.read_text())["state"] == state: + return path + except (OSError, ValueError): + pass + if process.poll() is not None: + break + time.sleep(0.1) + raise AssertionError( + f"state {state} never written; stderr: {process.stderr.read()[-2000:]!r}" + if process.poll() is not None + else f"state {state} never written" + ) + + +needs_posix_signals = pytest.mark.skipif( + not hasattr(signal, "SIGUSR1"), reason="POSIX signals only" +) + + +@needs_posix_signals +def test_killed_and_resumed_fit_is_attempt_two(tmp_path, output): + script = _script(tmp_path, output) + + process = _launch(script, sleep=0.5) + _wait_for_state(output, "running", process) + process.kill() + process.wait(timeout=30) + + resumed = _launch(script, sleep=0) + assert resumed.wait(timeout=300) == 0, resumed.stderr.read()[-2000:] + + path = _status_path(output) + status = json.loads(path.read_text()) + assert status["state"] == "completed" + assert status["attempt"] == 2 + + history = [ + json.loads(line) + for line in Path(str(path) + "l").read_text().splitlines() + ] + assert [(e["event"], e["attempt"]) for e in history] == [ + ("start", 1), + ("start", 2), + ("completed", 2), + ] + + +@needs_posix_signals +def test_sigusr1_then_sigterm_in_a_fitting_process(tmp_path, output): + process = _launch(_script(tmp_path, output), sleep=0.5) + path = _wait_for_state(output, "running", process) + + process.send_signal(signal.SIGUSR1) + _wait_for_state(output, "timeout_pending", process) + # SIGUSR1 does not stop the fit. + assert process.poll() is None + + process.send_signal(signal.SIGTERM) + assert process.wait(timeout=60) == 143 + + history = [ + json.loads(line) + for line in Path(str(path) + "l").read_text().splitlines() + ] + assert [e.get("signal") for e in history if e["event"] == "timeout_pending"] == [ + "SIGUSR1", + "SIGTERM", + ] + assert json.loads(path.read_text())["state"] == "timeout_pending" diff --git a/test_autofit/non_linear/status/test_units.py b/test_autofit/non_linear/status/test_units.py new file mode 100644 index 000000000..eca1ab6d2 --- /dev/null +++ b/test_autofit/non_linear/status/test_units.py @@ -0,0 +1,230 @@ +""" +The status record's parts on their own: the state vocabulary, the size-capped atomic +writer and its history, the fail-soft probes and the signal handlers. +""" +import json +import os +import signal +import sys + +import pytest + +from autofit.non_linear.status import probe, signals +from autofit.non_linear.status.states import ( + COLLECTOR_STATES, + WRITTEN_STATES, + FailureClass, + Product, + RunState, + failure_class_of, +) +from autofit.non_linear.status.writer import MAX_BYTES, StatusWriter, compact + + +def test_vocabulary(): + assert [state.value for state in WRITTEN_STATES] and RunState.RUNNING == "running" + assert {str(state) for state in COLLECTOR_STATES} == {"timeout", "oom", "node_fail"} + assert not WRITTEN_STATES & COLLECTOR_STATES + assert {str(c) for c in FailureClass} == { + "exception", + "bad_input", + "missing_seed", + "upstream_incomplete", + } + assert json.dumps({"state": RunState.TIMEOUT_PENDING}) == '{"state": "timeout_pending"}' + assert str(Product.STATUS) == ".status.json" + + +def test_failure_class_of(): + assert failure_class_of(ValueError()) is FailureClass.EXCEPTION + + error = ValueError() + error.failure_class = "missing_seed" + assert failure_class_of(error) is FailureClass.MISSING_SEED + + error.failure_class = "not a class" + assert failure_class_of(error) is FailureClass.EXCEPTION + + +def test_compact_trims_optional_fields_to_the_cap(): + document = { + "schema": 1, + "path_prefix": "x" * 1900, + "progress": {"n_like": 1, "log_z": 1.0, "n_eff": 2.0, "f_live": 0.1}, + "resources": {"wall_s": 1.0, "gpu": {"model": "y" * 200}, "cpu_model": "z" * 200}, + } + text = compact(document) + + assert len(text.encode()) <= MAX_BYTES + loaded = json.loads(text) + assert loaded["progress"] == {"n_like": 1} + assert loaded["resources"] == {"wall_s": 1.0} + assert loaded["path_prefix"] == "x" * 1900 + + +def test_compact_cuts_a_pathological_path_prefix(): + text = compact({"schema": 1, "path_prefix": "a" * 5000 + "tail"}) + + assert len(text.encode()) <= MAX_BYTES + assert json.loads(text)["path_prefix"].endswith("tail") + + +def test_writer_files_are_siblings_and_history_counts_attempts(tmp_path): + output_path = tmp_path / "prefix" / "name" / "identifier" + writer = StatusWriter(output_path) + + assert writer.status_path == tmp_path / "prefix" / "name" / "identifier.status.json" + assert writer.next_attempt() == 1 + + writer.write({"state": "running"}) + assert writer.read() == {"state": "running"} + assert not list(writer.status_path.parent.glob("*.tmp")) + + writer.append({"event": "start", "attempt": 1}) + writer.append({"event": "failed", "attempt": 1}) + # A line a killed process left truncated is skipped. + with open(writer.history_path, "a") as f: + f.write('{"event": "sta') + assert writer.next_attempt() == 2 + + +def test_writer_gives_up_on_a_held_lock_with_a_timeout(tmp_path): + writer = StatusWriter(tmp_path / "out") + writer._lock.acquire() + try: + assert writer.write({"state": "running"}, timeout=0.01) is None + assert writer.append({"event": "x"}, timeout=0.01) is False + finally: + writer._lock.release() + assert not writer.status_path.exists() + + +SLURM_ENV = { + "SLURM_JOB_ID": "351102", + "SLURM_ARRAY_JOB_ID": "351085", + "SLURM_ARRAY_TASK_ID": "17", + "SLURMD_NODENAME": "ral-n12", + "SLURM_CPUS_PER_TASK": "4", + "SLURM_MEM_PER_CPU": "3072", + "SLURM_RESTART_COUNT": "1", + "SLURM_JOB_START_TIME": "1000", + "SLURM_JOB_END_TIME": "65800", +} + + +def test_job_from_slurm_env(): + assert probe.job_from_env(SLURM_ENV) == { + "scheduler": "slurm", + "job_id": "351085", + "task_id": 17, + "node": "ral-n12", + "cpus": 4, + "mem_mb": 12288, + "time_limit_s": 64800, + "restart_count": 1, + } + + +def test_job_outside_a_scheduler_is_none(): + assert probe.job_from_env({}) is None + assert probe.job_from_env({"SLURM_JOB_ID": ""}) is None + + +def test_gpu_allocated(): + assert not probe.gpu_allocated({}) + assert not probe.gpu_allocated({"CUDA_VISIBLE_DEVICES": "-1"}) + assert probe.gpu_allocated({"CUDA_VISIBLE_DEVICES": "0"}) + assert probe.gpu_allocated({"SLURM_JOB_GPUS": "0"}) + + +def test_probes_fail_soft(monkeypatch): + monkeypatch.setattr(probe, "_cgroup_memory_peak", lambda: None) + rss = probe.peak_rss() + assert rss["peak_rss_source"] in ("getrusage", None) + + monkeypatch.setattr(probe, "gpu_allocated", lambda env=None: False) + assert probe.gpu() is None + + resources = probe.ResourceProbe(cpus=2).resources() + assert set(resources) == { + "wall_s", + "cpu_s", + "cpu_efficiency", + "peak_rss_mb", + "peak_rss_source", + "gpu", + "cpu_model", + } + assert resources["wall_s"] >= 0 + + +def test_cgroup_used_only_inside_a_job(monkeypatch): + monkeypatch.setattr(probe, "_cgroup_memory_peak", lambda: 3 * 2**30) + assert probe.peak_rss(use_cgroup=True) == { + "peak_rss_mb": 3072, + "peak_rss_source": "cgroup", + } + assert probe.peak_rss(use_cgroup=False)["peak_rss_source"] != "cgroup" + + +class FakeStatus: + def __init__(self): + self.signals = [] + + def timeout_pending(self, name): + self.signals.append(name) + + +needs_usr1 = pytest.mark.skipif( + not hasattr(signal, "SIGUSR1"), reason="no SIGUSR1 on this platform" +) + + +@needs_usr1 +def test_sigusr1_records_timeout_pending_and_restores_the_previous_handler(): + before = signal.getsignal(signal.SIGUSR1) + status = FakeStatus() + + restore = signals.install(status) + try: + os.kill(os.getpid(), signal.SIGUSR1) + assert status.signals == ["SIGUSR1"] + finally: + restore() + + assert signal.getsignal(signal.SIGUSR1) == before + + +@needs_usr1 +def test_sigterm_records_then_exits_143(): + before = signal.getsignal(signal.SIGTERM) + status = FakeStatus() + + restore = signals.install(status) + try: + with pytest.raises(SystemExit) as exit_info: + os.kill(os.getpid(), signal.SIGTERM) + # The handler runs at the next bytecode boundary. + for _ in range(1000): + pass + assert exit_info.value.code == 143 + assert status.signals == ["SIGTERM"] + finally: + restore() + + assert signal.getsignal(signal.SIGTERM) == before + + +def test_no_handlers_outside_the_main_thread(): + import threading + + installed = [] + + def target(): + installed.append(signals.can_install()) + + thread = threading.Thread(target=target) + thread.start() + thread.join() + + assert installed == [False]