From 7172c8311dee387b45742d6407c677c397d50832 Mon Sep 17 00:00:00 2001 From: "Josef M. Gallmetzer" <64498081+galjos@users.noreply.github.com> Date: Tue, 7 Jul 2026 11:42:43 +0200 Subject: [PATCH 1/2] Parallelize the screen driver over molecules Add a jobs parameter (and -j/--jobs CLI flag; default 1 = serial). With jobs > 1 the molecules run in a ProcessPoolExecutor over a module-level, picklable worker (_run_job) -- process-based because each job changes the working directory via os.chdir, so threads would race. Per-job error isolation, resume (skip completed), atomic incremental writes, and deterministic input-order output are all preserved. Tests use a synchronous fake executor to exercise the parallel branch without real subprocesses. Closes #96 --- ThermoScreening/cli/thermo.py | 6 + ThermoScreening/thermo/screening.py | 202 +++++++++++++++++++--------- docs/usage.rst | 3 + tests/thermo/test_main.py | 8 +- tests/thermo/test_screening.py | 70 +++++++++- 5 files changed, 221 insertions(+), 68 deletions(-) diff --git a/ThermoScreening/cli/thermo.py b/ThermoScreening/cli/thermo.py index 2988b26..55649b5 100644 --- a/ThermoScreening/cli/thermo.py +++ b/ThermoScreening/cli/thermo.py @@ -133,6 +133,11 @@ def _command_parser(): help="Reuse successful results from a prior .json and only (re)run " "the missing/failed molecules.", ) + screen_parser.add_argument( + "-j", "--jobs", type=int, default=1, + help="Number of molecules to run concurrently (default 1). Each job runs " + "in its own process and directory.", + ) conf_parser = subparsers.add_parser( "conformers", @@ -243,6 +248,7 @@ def run_screen(parser_args): engine=parser_args.engine, method=parser_args.method, resume=parser_args.resume, + jobs=parser_args.jobs, ) ranked = rank_by_gibbs(results) diff --git a/ThermoScreening/thermo/screening.py b/ThermoScreening/thermo/screening.py index 4747da9..d0dac29 100644 --- a/ThermoScreening/thermo/screening.py +++ b/ThermoScreening/thermo/screening.py @@ -8,9 +8,11 @@ """ import csv +import functools import json import logging import os +from concurrent.futures import ProcessPoolExecutor, as_completed from dataclasses import dataclass from pathlib import Path @@ -193,6 +195,80 @@ def write_csv(handle): return csv_path, json_path +def _run_job( + job, + *, + engine, + temperature, + pressure, + root, + method, + solvent, + dispersion, + quasi_rrho, + parameters, + spin_constants, +): + """ + Run one screening job and return its result record. + + A module-level, picklable worker (usable from a process pool). Failures are + isolated into the record (``status="error"``) rather than raised, so one bad + molecule never aborts the screen; the caller does the logging. + """ + record = { + "name": job.name, + "path": str(job.path), + "charge": job.charge, + "status": "ok", + "error": "", + } + try: + atoms = ase.io.read(str(job.path)) + if engine == "xtb-cli": + thermo = xtb_cli_thermo( + atoms, + temperature=temperature, + pressure=pressure, + charge=job.charge, + directory=str(root / job.name), + spin=job.spin, + method=method, + solvent=solvent, + quasi_rrho=quasi_rrho, + ) + elif engine == "xtb": + thermo = xtb_thermo( + atoms, + temperature=temperature, + pressure=pressure, + charge=job.charge, + directory=str(root / job.name), + spin=job.spin, + method=method, + quasi_rrho=quasi_rrho, + ) + else: + thermo = dftbplus_thermo( + atoms, + temperature=temperature, + pressure=pressure, + charge=job.charge, + directory=str(root / job.name), + spin=job.spin, + spin_constants=spin_constants, + solvent=solvent, + dispersion=dispersion, + quasi_rrho=quasi_rrho, + **parameters, + ) + record.update(_thermo_summary(thermo)) + except Exception as exc: # pylint: disable=broad-except + record["status"] = "error" + record["error"] = str(exc) + return record + + def screen( source, out="results", @@ -209,6 +285,7 @@ def screen( engine="dftb+", method="GFN2-xTB", resume=False, + jobs=1, ): """ Run a thermochemistry screen over a set of molecules. @@ -263,6 +340,11 @@ def screen( ``.json`` and skip those molecules; molecules that previously failed (or were never run) are (re-)run. Results are written incrementally after every molecule regardless, so an interrupted screen can be resumed. + jobs : int + Number of molecules to run concurrently. Default 1 (serial). With + ``jobs > 1`` the molecules run in a process pool (process-based because + each job changes the working directory). Results and output ordering are + unchanged; only the wall-clock time differs. Returns ------- @@ -273,82 +355,72 @@ def screen( raise TSValueError( f"Unknown engine {engine!r}; choose 'dftb+', 'xtb' or 'xtb-cli'." ) + if jobs < 1: + raise TSValueError(f"jobs must be >= 1, got {jobs}.") + spin_constants = None if engine == "dftb+": default_parameters, spin_constants = resolve_parameter_set(parameter_set) parameters = default_parameters if parameters is None else parameters - jobs = _load_jobs(source, charge, spin) + job_list = _load_jobs(source, charge, spin) root = Path(directory) - completed = _load_completed(out) if resume else {} - results = [] + run_one = functools.partial( + _run_job, + engine=engine, + temperature=temperature, + pressure=pressure, + root=root, + method=method, + solvent=solvent, + dispersion=dispersion, + quasi_rrho=quasi_rrho, + parameters=parameters, + spin_constants=spin_constants, + ) + + # records keyed by name; the returned list follows the input job order + records = {} csv_path = json_path = None - for job in jobs: - if job.name in completed: - logger.info(f"Skipping {job.name} (already completed)") - results.append(completed[job.name]) - csv_path, json_path = _write_results(results, out) - continue - record = { - "name": job.name, - "path": str(job.path), - "charge": job.charge, - "status": "ok", - "error": "", - } - logger.info(f"Screening {job.name} (charge {job.charge})") - try: - atoms = ase.io.read(str(job.path)) - if engine == "xtb-cli": - thermo = xtb_cli_thermo( - atoms, - temperature=temperature, - pressure=pressure, - charge=job.charge, - directory=str(root / job.name), - spin=job.spin, - method=method, - solvent=solvent, - quasi_rrho=quasi_rrho, - ) - elif engine == "xtb": - thermo = xtb_thermo( - atoms, - temperature=temperature, - pressure=pressure, - charge=job.charge, - directory=str(root / job.name), - spin=job.spin, - method=method, - quasi_rrho=quasi_rrho, - ) - else: - thermo = dftbplus_thermo( - atoms, - temperature=temperature, - pressure=pressure, - charge=job.charge, - directory=str(root / job.name), - spin=job.spin, - spin_constants=spin_constants, - solvent=solvent, - dispersion=dispersion, - quasi_rrho=quasi_rrho, - **parameters, - ) - record.update(_thermo_summary(thermo)) - except Exception as exc: # pylint: disable=broad-except - # isolate failures so one bad molecule does not abort the screen - record["status"] = "error" - record["error"] = str(exc) + + def _ordered(): + return [records[job.name] for job in job_list if job.name in records] + + def _store(job, record): + nonlocal csv_path, json_path + if record["status"] != "ok": # warning, not error: the active custom logger raises on error and # we must keep screening the remaining molecules - logger.warning(f"Screening failed for {job.name}: {exc}") - results.append(record) + logger.warning(f"Screening failed for {job.name}: {record['error']}") + records[job.name] = record # write after every molecule so an interrupted run stays resumable - csv_path, json_path = _write_results(results, out) + csv_path, json_path = _write_results(_ordered(), out) + + pending = [] + for job in job_list: + if job.name in completed: + logger.info(f"Skipping {job.name} (already completed)") + records[job.name] = completed[job.name] + else: + pending.append(job) + if records: + csv_path, json_path = _write_results(_ordered(), out) + + if jobs > 1 and len(pending) > 1: + with ProcessPoolExecutor(max_workers=jobs) as executor: + future_to_job = {} + for job in pending: + logger.info(f"Screening {job.name} (charge {job.charge})") + future_to_job[executor.submit(run_one, job)] = job + for future in as_completed(future_to_job): + job = future_to_job[future] + _store(job, future.result()) + else: + for job in pending: + logger.info(f"Screening {job.name} (charge {job.charge})") + _store(job, run_one(job)) logger.info(f"Wrote {csv_path} and {json_path}") - return results + return _ordered() diff --git a/docs/usage.rst b/docs/usage.rst index da0b037..0b0bcb2 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -38,6 +38,9 @@ The ``thermo`` command has four subcommands. # resume an interrupted screen thermo screen molecules/ -o results --resume + # run 4 molecules at a time (each in its own process/directory) + thermo screen molecules/ --jobs 4 + Python API ---------- diff --git a/tests/thermo/test_main.py b/tests/thermo/test_main.py index d63412f..fcf0355 100644 --- a/tests/thermo/test_main.py +++ b/tests/thermo/test_main.py @@ -180,7 +180,7 @@ def fake_screen(source, **kwargs): source="mols", out="results", charge=0.0, temperature=298.15, pressure=101325, directory="screening", parameter_set="3ob", solvent=None, dispersion="d3-bj", quasi_rrho=False, engine="dftb+", - method="GFN2-xTB", resume=False, + method="GFN2-xTB", resume=False, jobs=1, ) assert thermo.run_screen(args) == 0 assert captured["dispersion"] == "d3-bj" @@ -191,6 +191,12 @@ def test_parse_args_screen_dispersion_choice(): assert thermo.parse_args(["screen", "mols", "--dispersion", "d3-bj"]).dispersion == "d3-bj" +def test_parse_args_screen_jobs(): + assert thermo.parse_args(["screen", "mols"]).jobs == 1 + assert thermo.parse_args(["screen", "mols", "-j", "4"]).jobs == 4 + assert thermo.parse_args(["screen", "mols", "--jobs", "8"]).jobs == 8 + + def test_main_runs_conformers(monkeypatch): monkeypatch.setattr( thermo, "parse_args", lambda: argparse.Namespace(command="conformers") diff --git a/tests/thermo/test_screening.py b/tests/thermo/test_screening.py index 9b5bcb4..7a94f8f 100644 --- a/tests/thermo/test_screening.py +++ b/tests/thermo/test_screening.py @@ -1,6 +1,7 @@ import csv import json from argparse import Namespace +from concurrent.futures import Future from pathlib import Path import pytest @@ -9,6 +10,26 @@ from ThermoScreening.thermo import screening +class _SyncExecutor: + """A drop-in ProcessPoolExecutor that runs submit() synchronously in-process + (so monkeypatched engine functions apply and the parallel branch is testable + without real subprocesses).""" + + def __init__(self, max_workers=None): + self.max_workers = max_workers + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def submit(self, fn, *args, **kwargs): + future = Future() + future.set_result(fn(*args, **kwargs)) + return future + + class _FakeThermo: def electronic_energy(self): return -100.0 @@ -166,6 +187,51 @@ def fake_thermo(atoms, dispersion=None, **kwargs): assert captured["dispersion"] == "d3-bj" +def test_screen_rejects_non_positive_jobs(tmp_path): + with pytest.raises(TSValueError, match="jobs must be >= 1"): + screening.screen(str(tmp_path), out=str(tmp_path / "r"), jobs=0) + + +def test_screen_parallel_runs_all_and_preserves_order(monkeypatch, tmp_path): + for name in ("mol_a", "mol_b", "mol_c"): + _write_xyz(tmp_path / f"{name}.xyz") + + monkeypatch.setattr(screening, "dftbplus_thermo", lambda atoms, **kwargs: _FakeThermo()) + monkeypatch.setattr(screening, "ProcessPoolExecutor", _SyncExecutor) + + results = screening.screen( + str(tmp_path), out=str(tmp_path / "r"), directory=str(tmp_path / "runs"), jobs=3, + ) + + # results follow input (sorted) order regardless of completion order + assert [r["name"] for r in results] == ["mol_a", "mol_b", "mol_c"] + assert all(r["status"] == "ok" for r in results) + assert all(r["G_total_hartree"] == -111.0 for r in results) + + +def test_screen_parallel_isolates_failures(monkeypatch, tmp_path): + for name in ("good1", "bad", "good2"): + _write_xyz(tmp_path / f"{name}.xyz") + + def fake_thermo(atoms, directory=None, **kwargs): + if directory is not None and directory.endswith("bad"): + raise RuntimeError("boom") + return _FakeThermo() + + monkeypatch.setattr(screening, "dftbplus_thermo", fake_thermo) + monkeypatch.setattr(screening, "ProcessPoolExecutor", _SyncExecutor) + + results = screening.screen( + str(tmp_path), out=str(tmp_path / "r"), directory=str(tmp_path / "runs"), jobs=2, + ) + + by_name = {r["name"]: r for r in results} + assert by_name["bad"]["status"] == "error" + assert "boom" in by_name["bad"]["error"] + assert by_name["good1"]["status"] == "ok" + assert by_name["good2"]["status"] == "ok" + + def test_screen_dispatches_to_xtb_engine(monkeypatch, tmp_path): _write_xyz(tmp_path / "mol.xyz") @@ -461,7 +527,7 @@ def test_cli_run_screen_returns_failure_count(monkeypatch, capsys): source="x", out="res", charge=0.0, temperature=298.15, pressure=101325.0, directory="screening", parameter_set="3ob", solvent=None, dispersion=None, quasi_rrho=False, engine="dftb+", - method="GFN2-xTB", resume=False, + method="GFN2-xTB", resume=False, jobs=1, ) assert cli.run_screen(args) == 1 # one molecule failed @@ -484,7 +550,7 @@ def test_cli_run_screen_all_ok_returns_zero(monkeypatch, capsys): source="x", out="res", charge=0.0, temperature=298.15, pressure=101325.0, directory="screening", parameter_set="3ob", solvent=None, dispersion=None, quasi_rrho=False, engine="dftb+", - method="GFN2-xTB", resume=False, + method="GFN2-xTB", resume=False, jobs=1, ) assert cli.run_screen(args) == 0 # no failures From e0be0584e289e1f59d1c3e213d15d92a4de46663 Mon Sep 17 00:00:00 2001 From: "Josef M. Gallmetzer" <64498081+galjos@users.noreply.github.com> Date: Tue, 7 Jul 2026 11:51:55 +0200 Subject: [PATCH 2/2] Harden parallel screen: unique names, worker-death, worker cap Address review of the parallel screen: - Reject duplicate molecule names in _load_jobs (names key both the result record and the per-molecule directory, so a collision would drop a result and race two parallel jobs into the same directory). - Convert a dead worker process (BrokenProcessPool / OOM) into a per-job error record instead of aborting the whole screen. - Cap the pool at min(jobs, len(pending)). Add tests for duplicate-name rejection, worker-partial picklability (which the in-process fake executor can't catch), reversed completion order, and worker-death survival. --- ThermoScreening/thermo/screening.py | 43 +++++++++++++++++++----- tests/thermo/test_screening.py | 52 +++++++++++++++++++++++++++++ 2 files changed, 87 insertions(+), 8 deletions(-) diff --git a/ThermoScreening/thermo/screening.py b/ThermoScreening/thermo/screening.py index d0dac29..80ca872 100644 --- a/ThermoScreening/thermo/screening.py +++ b/ThermoScreening/thermo/screening.py @@ -94,12 +94,26 @@ def _jobs_from_manifest(manifest: Path, charge: float, spin=None): def _load_jobs(source, charge: float, spin=None): source = Path(source) if source.is_dir(): - return _jobs_from_directory(source, charge, spin) - if source.suffix.lower() == ".csv": - return _jobs_from_manifest(source, charge, spin) - raise TSValueError( - f"Screening input must be a directory or a .csv manifest, got '{source}'." - ) + jobs = _jobs_from_directory(source, charge, spin) + elif source.suffix.lower() == ".csv": + jobs = _jobs_from_manifest(source, charge, spin) + else: + raise TSValueError( + f"Screening input must be a directory or a .csv manifest, got '{source}'." + ) + + # Names key the result records and the per-molecule working directory, so a + # collision would silently drop a result and (in parallel) race two jobs into + # the same directory. Reject duplicates up front. (E.g. a directory holding + # both mol.xyz and mol.gen, or a manifest listing a name twice.) + names = [job.name for job in jobs] + duplicates = sorted({name for name in names if names.count(name) > 1}) + if duplicates: + raise TSValueError( + f"Duplicate molecule name(s) in the screen: {', '.join(duplicates)}. " + "Each molecule needs a unique name." + ) + return jobs def _thermo_summary(thermo): @@ -409,14 +423,27 @@ def _store(job, record): csv_path, json_path = _write_results(_ordered(), out) if jobs > 1 and len(pending) > 1: - with ProcessPoolExecutor(max_workers=jobs) as executor: + with ProcessPoolExecutor(max_workers=min(jobs, len(pending))) as executor: future_to_job = {} for job in pending: logger.info(f"Screening {job.name} (charge {job.charge})") future_to_job[executor.submit(run_one, job)] = job for future in as_completed(future_to_job): job = future_to_job[future] - _store(job, future.result()) + try: + record = future.result() + except Exception as exc: # pylint: disable=broad-except + # the worker process itself died (BrokenProcessPool, OOM, a + # segfault in a C extension); record this job as failed and + # keep going instead of aborting the whole screen + record = { + "name": job.name, + "path": str(job.path), + "charge": job.charge, + "status": "error", + "error": str(exc), + } + _store(job, record) else: for job in pending: logger.info(f"Screening {job.name} (charge {job.charge})") diff --git a/tests/thermo/test_screening.py b/tests/thermo/test_screening.py index 7a94f8f..31df9b2 100644 --- a/tests/thermo/test_screening.py +++ b/tests/thermo/test_screening.py @@ -30,6 +30,15 @@ def submit(self, fn, *args, **kwargs): return future +class _BrokenExecutor(_SyncExecutor): + """Simulates a dead worker: every future's result() raises.""" + + def submit(self, fn, *args, **kwargs): + future = Future() + future.set_exception(RuntimeError("worker died")) + return future + + class _FakeThermo: def electronic_energy(self): return -100.0 @@ -192,12 +201,39 @@ def test_screen_rejects_non_positive_jobs(tmp_path): screening.screen(str(tmp_path), out=str(tmp_path / "r"), jobs=0) +def test_load_jobs_rejects_duplicate_names(tmp_path): + # mol.xyz and mol.gen both have stem "mol" -> would collide on name/directory + _write_xyz(tmp_path / "mol.xyz") + (tmp_path / "mol.gen").write_text("dummy", encoding="utf-8") + with pytest.raises(TSValueError, match="Duplicate molecule name"): + screening._load_jobs(tmp_path, charge=0.0) + + +def test_run_job_worker_is_picklable(): + # spawn-based pools pickle the worker partial and each job; guard that here + # since the in-process fake executor never exercises pickling. + import pickle + import functools + + job = screening.ScreeningJob(name="m", path=Path("m.xyz"), charge=0.0, spin=None) + assert pickle.loads(pickle.dumps(job)).name == "m" + + partial = functools.partial( + screening._run_job, engine="dftb+", temperature=298.15, pressure=101325, + root=Path("runs"), method="GFN2-xTB", solvent=None, dispersion=None, + quasi_rrho=False, parameters={}, spin_constants=None, + ) + pickle.loads(pickle.dumps(partial)) # must round-trip for ProcessPoolExecutor + + def test_screen_parallel_runs_all_and_preserves_order(monkeypatch, tmp_path): for name in ("mol_a", "mol_b", "mol_c"): _write_xyz(tmp_path / f"{name}.xyz") monkeypatch.setattr(screening, "dftbplus_thermo", lambda atoms, **kwargs: _FakeThermo()) monkeypatch.setattr(screening, "ProcessPoolExecutor", _SyncExecutor) + # force completion order to be the reverse of submission order + monkeypatch.setattr(screening, "as_completed", lambda futures: list(futures)[::-1]) results = screening.screen( str(tmp_path), out=str(tmp_path / "r"), directory=str(tmp_path / "runs"), jobs=3, @@ -209,6 +245,22 @@ def test_screen_parallel_runs_all_and_preserves_order(monkeypatch, tmp_path): assert all(r["G_total_hartree"] == -111.0 for r in results) +def test_screen_parallel_survives_worker_death(monkeypatch, tmp_path): + for name in ("m1", "m2"): + _write_xyz(tmp_path / f"{name}.xyz") + + monkeypatch.setattr(screening, "dftbplus_thermo", lambda atoms, **kwargs: _FakeThermo()) + monkeypatch.setattr(screening, "ProcessPoolExecutor", _BrokenExecutor) + + results = screening.screen( + str(tmp_path), out=str(tmp_path / "r"), directory=str(tmp_path / "runs"), jobs=2, + ) + + # a dead worker becomes a per-job error, not an aborted screen + assert len(results) == 2 + assert all(r["status"] == "error" and "worker died" in r["error"] for r in results) + + def test_screen_parallel_isolates_failures(monkeypatch, tmp_path): for name in ("good1", "bad", "good2"): _write_xyz(tmp_path / f"{name}.xyz")