Add PySCF as a local in-core ESS adapter (sp + opt + freq) - #924
Conversation
|
The issue I had with pyscf or rather with any locally run ESS is the serial run of jobs and since we spawn so many conformers we would be waiting weeks to months. I had tried to implement PySCF with utilising parallel runs by "submitting" to the system and then holding onto their PID and checking it every so often, so emulating HPC systems but the issue then came down to the hardware of the local PC https://github.com/ReactionMechanismGenerator/ARC/tree/pyscf This was written before claude code etc so there may be nuances or other things I could have done |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #924 +/- ##
==========================================
+ Coverage 63.43% 63.47% +0.04%
==========================================
Files 114 115 +1
Lines 38325 38484 +159
Branches 10030 10054 +24
==========================================
+ Hits 24312 24429 +117
- Misses 11096 11126 +30
- Partials 2917 2929 +12
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
End-to-end validation: real ARC runs through the PySCF adapter, CPU vs GPURan this branch on real species to check the adapter works end-to-end and to measure whether the GPU Setup
Geometry optimization
The GPU only pays off above roughly 100 basis functions, and the advantage grows with system size.
Reproducibility. GPU timings are near-identical across three independent runs (propanol ARC's second
|
PySCF grids.level |
radial × angular (C/N/O) | Gaussian equivalent |
|---|---|---|
| 3 (PySCF default) | (75, 302) | FineGrid — exact match |
4 ← used for fine |
(90, 590) | UltraFine (99,590) |
| 5 | (105, 770) | overshoots UltraFine by ~1.5× in grid points |
Level 4 is therefore the honest analogue. The NLC (VV10) grid is left at PySCF's default level 3, since
the nonlocal kernel is smooth and conventionally integrated more coarsely than the XC term.
Not yet measured
- Frequencies on either device. The opt → fine-opt → freq pipeline now runs end-to-end (verified —
freq jobs spawn and execute), but no freq job has been carried to completion yet. - n-dodecane on CPU. The 1-core and 10-core runs were interrupted. At the point of interruption
10 cores had already passed the GPU's 1449 s without converging, so the GPU is >1.2× faster at
this size — a lower bound, not a measurement.
GPU limitations worth knowing before routing production jobs
- Closed-shell analytic Hessians are forced back to the CPU.
gpu4pyscf1.8.0 fails with
assert a.flags.c_contiguous(lib/cupy_helper.py:527, viahessian/rhf.py:753). This reproduces
with B3LYP too, so it is not NLC-specific — it looks like a general gpu4pyscf bug worth reporting
upstream. Consequence: singletfreqgets no GPU acceleration. - Open-shell frequencies use a numerical Hessian, because PySCF raises
NotImplementedError: UKS Hessian for NLC functional(hessian/uks.py:44). RKS NLC Hessians are
implemented. This is expensive —[CH2]CN's numerical Hessian was still running after 25 min. - Scaling is sub-linear. A 10-thread job sustained only ~650% CPU, not ~1000%. For throughput this
argues for several concurrent modest-core jobs (what the local pipe backend does) over one wide job. - Use the cuda12x wheel stack, not cu13 — the
nvidia-*-cu13pip wheels are unpopulated stubs
(nvidia-nvjitlink-cu13is 0.0.1) and cupy fails to findlibnvJitLink.so.
Bugs this shook out
Running a local in-core ESS end-to-end had evidently never been done before; five distinct failures had
to be fixed to get a single species through opt → freq. Four are pre-existing ARC issues, not
PySCF-specific.
| # | Site | Symptom |
|---|---|---|
| 1 | arc/common.py:156 |
check_ess_settings rejects server local unless a local entry exists in servers. Any in-core ESS needs one. |
| 2 | arc/scheduler.py:1235 |
conf_opt: false is ignored for a species defined only by SMILES — the gate is job_types['conf_opt'] ... or species.get_xyz(generate=False) is None, so conformers are generated anyway and dispatched to the conformer-level ESS. |
| 3 | arc/job/local.py:186 |
ValueError: Server cluster software local is not supported — a machine with no queueing system has no queue job IDs to report. |
| 4 | arc/job/adapter.py:799 |
ValueError: Unrecognized cluster software: local — tries to collect scheduler stdout/stderr files that cannot exist without a scheduler. |
| 5 | arc/job/trsh.py |
The significant one. determine_ess_status() only understands the text logs of gaussian/qchem/orca/molpro/terachem; anything else falls through to return '', .... Status '' is not 'done', so ARC marked every successfully completed in-core job as failed, tried to troubleshoot it, found no methods, and declared the species unconverged — while output.yml said success: True the whole time. This is not PySCF-specific: ase, torchani, gcn and mockter all write output.yml and hit the identical fallthrough. Fixed generally via a new determine_ess_status_from_yaml() that treats the adapter's own success/error keys as authoritative. |
Fixes for 3, 4 and 5 are in this branch, with tests. check_running_jobs_ids() had no test coverage at
all previously. Item 1 is a settings/documentation gap — the shipped servers['local'] template now
documents 'local' as a valid cluster_soft for a machine with no queue. Item 2 is untouched and
still open; it is orthogonal to this PR and probably deserves its own issue.
There was a problem hiding this comment.
Pull request overview
Adds a new in-core electronic-structure adapter for running PySCF jobs locally via ARC’s YAML-handoff contract, and extends ARC’s “pipe” execution mode to optionally run as a local worker pool (no scheduler). This fits ARC’s adapter architecture by introducing a new ESS backend (pyscf) and ensuring failures are surfaced consistently through YAML parsing / status detection.
Changes:
- Added a new PySCF adapter + standalone engine script to support
sp,opt,freq, andoptfreqlocally viainput.yml/output.yml. - Added a
localpipe backend (worker pool) plus settings to cap local concurrency by cores/memory. - Improved YAML-based failure surfacing by having YAML parsing recognize
success: falseand returning the associated error message.
Reviewed changes
Copilot reviewed 15 out of 15 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| arc/settings/submit.py | Adds a local pipe submit template that runs a background worker pool instead of submitting to a scheduler. |
| arc/settings/settings.py | Registers pyscf, adds PYSCF_PYTHON/PYSCF_DEVICE, and introduces local pipe settings (run_locally, local_max_workers). |
| arc/parser/adapters/yaml.py | Makes YAMLParser surface in-core failures via success: false + error. |
| arc/main_test.py | Updates settings serialization test expectations to include pyscf. |
| arc/job/trsh.py | Adds YAML-aware ESS status detection (determine_ess_status_from_yaml) and routes .yml outputs through it. |
| arc/job/trsh_test.py | Adds tests for YAML-based ESS status determination. |
| arc/job/pipe/pipe_run.py | Adds local submission path, local worker limiting, and local-aware TMPDIR derivation. |
| arc/job/pipe/pipe_run_test.py | Adds tests for local pipe submit content and worker limit capping behavior. |
| arc/job/local.py | Treats cluster_soft == 'local' as “no queue” and returns no running job IDs. |
| arc/job/local_test.py | Adds tests for cluster_soft == 'local' behavior and unsupported scheduler rejection. |
| arc/job/adapters/scripts/pyscf_script.py | New standalone PySCF engine script producing schema-v1 output.yml. |
| arc/job/adapters/pyscf_test.py | New unit tests for PySCF adapter IO, YAML round-trip, and error surfacing. |
| arc/job/adapters/pyscf_adapter.py | New PySCFAdapter implementing YAML in/out and isolated execution in pyscf_env. |
| arc/job/adapters/init.py | Registers/imports the new adapter module. |
| arc/job/adapter.py | Adds JobEnum.pyscf and skips scheduler stdout/stderr collection for cluster_soft == 'local'. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Run DFT on the workstation with no PBS/zeus round-trip, using the YAML-handoff in-core pattern (mirroring the ASE adapter): the engine script runs PySCF in the dedicated pyscf_env conda env, writes an output.yml in ARC's internal schema, and YAMLParser re-parses it. Level: wb97m-v/def2-tzvp (native VV10 nonlocal dispersion). Closed-shell (RKS) frequencies use PySCF's analytic Hessian; open-shell (UKS) frequencies use a numerical Hessian (finite difference of analytic gradients), since PySCF has no analytic Hessian for NLC functionals in the unrestricted case. geomeTRIC drives geometry optimization, with tightened trust radii for TS searches. IRC and scans/rotors are out of scope. PySCF is opt-in per level rather than auto-routed: adding a levels_ess phrase such as 'wb97m-v' would silently move levels that currently run on QChem/Gaussian onto the local adapter, so it is deliberately omitted. The engine honors ARC's `fine` flag (denser XC/NLC integration grids) and pins PySCF's thread count to the job's cpu_cores, so concurrent local jobs cannot each grab every core. SCF convergence and grid levels are exposed as optional settings and otherwise left at PySCF's defaults. Also surface failed in-core jobs: YAMLParser.logfile_contains_errors now returns the error when an output.yml sets success: false (unaffected for producers that don't emit success). Registers pyscf in supported_ess, global_ess_settings, input/output_filenames, PYSCF_PYTHON, the JobEnum, and the adapter registry.
Setting `device: gpu` (globally via PYSCF_DEVICE, or per job via args['keyword']['device']) runs the SCF, the nuclear gradients and the geomeTRIC geometry optimization on an NVIDIA GPU. Energies agree with the CPU path to ~3e-5 kJ/mol, and the speedup grows with system size: measured 2.5x for ethanol (129 basis functions) and 3.0x for toluene (265) at wb97m-v/def2-TZVP against 28 CPU cores, with no benefit below ~100 basis functions. Closed-shell analytic Hessians always fall back to the CPU: gpu4pyscf 1.8.0 raises an assertion error in transpose_sum() for every functional tried, not just NLC ones. Open-shell numerical Hessians stay on the GPU, since they only need SCF energies and gradients, which is the expensive path for radicals. If gpu4pyscf is not installed the engine logs a notice and runs on the CPU, so the device setting is safe to leave on in a CPU-only environment.
ARC's pipe mechanism already batches independent jobs (conformers, sp, freq, rotor scans) into an array of workers that claim tasks from a shared pipe root, so the workers need no scheduler to coordinate them. This adds a 'local' backend that runs that array as a pool of background processes on the current machine, which is what lets locally-executed engines such as PySCF run in parallel instead of one job at a time. Enabled with pipe_settings['run_locally']. The pool size is capped by the number of cores and by the available memory, rather than by pipe_settings['max_workers'] alone, since nothing else throttles it without a queue; pipe_settings ['local_max_workers'] overrides the derived limit. Each worker is pinned to its task's core count so the pool cannot oversubscribe the machine. Completion is detected by reconciling the per-task state files, exactly as for a queued pipe run, so the returned job id is only used for bookkeeping.
|
@alongd Missing the installation of PySCF scripts - both with and without gpu support |
|
I think also pin it to 1.7.6 for GPU4PySCF since there is a regression: |
|
Also, please write up an update to the documentation on the utilisation of PySCF so a user can look on the website to understand what is required, what they need to change and what each change does |
| # =========================================================================== | ||
|
|
||
|
|
||
| def local_worker_limit(cpus_per_worker: int, memory_mb_per_worker: int) -> int: |
There was a problem hiding this comment.
The problem that I see here is that this function is stateless, so to speak. It only takes how big one worker, and not how many workers are already running. So every batch asks it the same questions and gets the same answers because it is recomputing form the full machine as if it were one batch only.
so we could have a sp batch and a freq batch and a ts batch, and each time we ask 'how many fit?', it would do
os.cpu_count() // 5 and say 4 for each cause it still believes 20 cores are free.
| status handling (e.g., ``determine_job_status()``) does not query a queueing system | ||
| for a job that never reached one. | ||
| """ | ||
| logger.warning('Queue execution is not supported for PySCF in ARC; running incore instead.') |
There was a problem hiding this comment.
This gets spammed quite a bit, I don't know if we should have this warning?
| xyz=xyz, | ||
| ) | ||
|
|
||
| def determine_settings(self) -> dict: |
There was a problem hiding this comment.
settings_dict is seeded from args['keyword'] before ARC setdefaults its own keys, so a user arg named method, basis, memory_mb, cpu_cores or device silently wins.
args: {keyword: {method: 'pbe'}} would run PBE while the level, logs and restart file all still say b3lyp.
There was a problem hiding this comment.
┌────────────────────────────────────────────────────────────┬───────────────────────────────────┐
│ step │ settings_dict['method'] │
├────────────────────────────────────────────────────────────┼───────────────────────────────────┤
│ user args seed the dict │ 'pbe' │
├────────────────────────────────────────────────────────────┼───────────────────────────────────┤
│ setdefault('method', 'b3lyp') │ 'pbe' ← no-op, key already exists │
├────────────────────────────────────────────────────────────┼───────────────────────────────────┤
│ what PySCF actually runs │ PBE │
├────────────────────────────────────────────────────────────┼───────────────────────────────────┤
│ what ARC's logs, restart.yml, and the Arkane level all say │ B3LYP │
└────────────────────────────────────────────────────────────┴───────────────────────────────────┘
| import unittest | ||
| from unittest import mock | ||
|
|
||
| import arc.job.pipe.pipe_run as pipe_run_module |
local_worker_limit() sized the local pipe pool off os.cpu_count(), ignoring the 'cpus' ceiling that the in-core ESS path already applies per job. As a result there was no single knob to cap total local CPU use. Add local_cpu_budget(), which reads the 'cpus' of the server whose cluster_soft is 'local' (falling back to the physical core count), and have local_worker_limit() derive from it. One number -- the local server's 'cpus' -- now bounds both the in-core jobs and the local worker pool; pipe_settings['local_max_workers'] remains an optional override. Documented in settings.py.
Summary
Adds PySCF as a local (in-core) electronic-structure adapter so ARC can run DFT on the
workstation with no PBS/zeus round-trip. Uses the YAML-handoff in-core pattern (mirroring
ase_adapter.py): the engine script runs PySCF in a dedicatedpyscf_envconda env, writes anoutput.ymlin ARC's internal schema, and ARC'sYAMLParserre-parses it (the on-diskoutput.ymlis the entire contract).Scope: exactly
sp+opt+freq(+optfreq), closed- and open-shell (RKS/UKS).IRC (no native PySCF path — needs Sella) and scans/rotors are out of scope, deferred.
Level of theory & frequency method
(ωB97X-D3, the originally-requested level, is not available in PySCF —
wb97x_d3isblack-listed in
pyscf/scf/dispersion.pyand raisesNotImplementedError. ωB97M-V was chosenas the replacement.)
because PySCF raises
NotImplementedError: UKS Hessian for NLC functional— there is no analyticHessian for NLC functionals in the unrestricted case.
Validation (raw PySCF 2.14.0, ωB97M-V/def2-TZVP, before wiring the adapter)
End-to-end round-trip (through the full ARC adapter path, local, no zeus)
A real
optfreqon H₂O viaPySCFAdapter.execute_incore()→output.yml→YAMLParser:e_elect=−200665.877 kJ/mol, geometry, freqs[1617.5, 3823.7, 3927.5], ZPE56.038 kJ/mol,dipole
2.0657 D, normal modes(3,3,3)— all ingest cleanly.Schema decisions (output.yml, schema v1)
spandzpein kJ/mol (matching whatYAMLParserconsumes); scanenergieswould be Hartree — documented, but scans are out of scope here.opt_xyz(opt jobs) /xyz(sp/freq),never a string — addressing the geometry key/format drift.
freqs(cm⁻¹, imaginary as negative reals) +modes+zpe, i.e.exactly the keys
YAMLParseralready ingests.hessian/reduced_masses/force_constantsaredeliberately NOT emitted since
YAMLParserdrops them (no written-but-dropped fields). Ingestingthe analytic Hessian into Arkane is a follow-up (see below).
success: false+error, andYAMLParser.logfile_contains_errors()now returns that error (previously alwaysNone), so acrashed local job isn't silently treated as done. Producers that don't emit
successare unaffected.Files
arc/job/adapters/scripts/pyscf_script.py— engine script (RKS/UKS sp/opt/freq via geomeTRIC +analytic/numerical Hessian; emits schema-v1
output.yml). Runs inpyscf_env.arc/job/adapters/pyscf_adapter.py— theJobAdapter(YAML in/out; launches the script viarun_in_conda_envto isolatepyscf_env).arc/job/adapters/pyscf_test.py— unit tests (IO/logic mocked for CI; engine test skipped unlessPySCF is importable). 7 passed, 1 skipped locally.
arc/parser/adapters/yaml.py— error-surfacing fix (above).supported_ess,global_ess_settings,levels_ess(wb97m-v),input_filenames,output_filenames,PYSCF_PYTHON(settings.py);JobEnum.pyscf(adapter.py);adapter import (
adapters/__init__.py);register_job_adapter('pyscf', ...).Reviewer notes / follow-ups
levels_ess['pyscf'] = ['wb97m-v']routes wb97m-v jobs to PySCF. A local~/.arc/settings.py(if present) overrideslevels_ess/supported_ess, so addpyscfthere tooto exercise routing on a machine that uses a local settings override.
pyscf_envconda env withpyscf,geometric,pyyaml.(
YAMLParsercurrently reads freqs+modes+zpe, not the full Hessian).is_ts) attempts a geomeTRIC saddle search (transition=True); validated here onlyfor minimization.