diff --git a/autofit/campaign/scheduler/slurm.py b/autofit/campaign/scheduler/slurm.py index 40b789cca..b4f99f19f 100644 --- a/autofit/campaign/scheduler/slurm.py +++ b/autofit/campaign/scheduler/slurm.py @@ -107,9 +107,12 @@ def sbatch_script(request: SubmitRequest) -> str: "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=$?", + # A trapped signal ends `wait` early (128 + signal); wait again until the fit + # itself has exited and keep only its own status, so a fit that checkpoints on + # USR1 and then finishes cleanly exits 0, not 138. + "while true; do", + " status=0", + ' wait "$child" || status=$?', ' kill -0 "$child" 2>/dev/null || break', "done", 'exit "$status"', diff --git a/autofit/results/index.py b/autofit/results/index.py index b96a6e381..ff1e2fa39 100644 --- a/autofit/results/index.py +++ b/autofit/results/index.py @@ -43,7 +43,13 @@ reruns of that fit with another model or search setting (sibling hashes). """ -_SKIP_DIRS = {"files", "image", INDEX_DIRNAME} +#: Batch-side task records (``//.status.json``) a campaign's +#: stage entry writes before its fit starts, so a task that fails before Python runs is +#: still counted. They are the campaign collector's (``autofit.campaign.collector`` +#: walks for them itself), never results. +CAMPAIGN_TASKS_DIRNAME = ".campaign_tasks" + +_SKIP_DIRS = {"files", "image", INDEX_DIRNAME, CAMPAIGN_TASKS_DIRNAME} _STATUS_SUFFIX = str(Product.STATUS) _ZIP_SUFFIX = str(Product.ZIP) diff --git a/test_autofit/campaign/test_slurm.py b/test_autofit/campaign/test_slurm.py index 7e59aed4e..5afb35838 100644 --- a/test_autofit/campaign/test_slurm.py +++ b/test_autofit/campaign/test_slurm.py @@ -78,6 +78,69 @@ def test_sbatch_script(): assert "--gres" not in script +def _run_generated_script(tmp_path, fit_body, signal_usr1): + """Run the generated batch script under bash, with a stand-in fit.""" + import os + import shutil + import signal + import time + + if shutil.which("bash") is None: + pytest.skip("needs bash") + ready = tmp_path / "ready" + fit = tmp_path / "fit.py" + fit.write_text( + "import os, signal, sys, time\n" + "got = []\n" + "signal.signal(signal.SIGUSR1, lambda *_: got.append(1))\n" + f"open({str(ready)!r}, 'w').close()\n" + fit_body + ) + request = _request(cpus=1) + request.script = str(fit) + request.project_root = str(tmp_path) + request.env = {} + script = tmp_path / "job.sh" + script.write_text(sbatch_script(request)) + # The script runs `python