perf(reaction): stop generating 3D conformers just to count atoms (59x); allow execute_command to time out

[alongd]

Two changes

1. check_atom_balance() no longer builds 3D geometries

ARCReaction.check_atom_balance() built its reactant and product wells from
species.get_xyz(generate=True). For any species without coordinates that forces
get_cheap_conformer() → embed_rdkit() + rdkit_force_field() (MMFF94s) — a full 3D embedding
and force-field optimisation, purely to obtain a list of element symbols.

The new private _get_atom_balance_entry() returns the xyz string of an already available
geometry, and otherwise reads element symbols off the species' 2D graph with placeholder
coordinates. It mirrors the conditions of ARCSpecies.get_xyz(generate=True): coordinates first,
a TS never falls back to its graph, the mol_list fallback matches that method's
self.mol is not None or self.mol_list is not None guard, and a species with neither coordinates
nor a graph still empties the well and skips the check.

Building the 1460 reactions of a real 103-species RMG mechanism: 26.7 s → 1.2 s wall;
check_atom_balance itself 26.66 s → 0.34 s by cProfile; RDKit is never entered. Reproduced
independently at 228.3 ms/rxn → 3.9 ms/rxn (59×) on 15 reactions with fresh species.

2. execute_command() can take a timeout

Optional timeout: float | None = None. With None — the default — the call is byte-for-byte what
it was. When a timeout is given, the command runs in its own session and the timeout kills the whole
process group, SIGTERM then SIGKILL.

The verdict of check_atom_balance() is unchanged

A generated conformer is embedded from the very same 2D graph, so it always carries identical
element counts. Verified over ~51 species — deuterated and tritiated, cubane, prismane,
tetrahedrane, a 20-membered ring, trans-cyclooctene, cations and anions, Si, S8, noble gases,
adjlist-defined species, atomic C — with zero divergence. A 16-case oracle over every species
shape that can reach the function (including no-mol, no-graph-and-no-coordinates, resonance,
TS-in-well, duplicate species, label-only reactions) returns byte-identical verdicts and exception
types on base and head.

But this is NOT a behaviour-free change — four real changes

ARCReaction.__init__ used to write an MMFF cheap_conformer onto every polyatomic reaction
species. State diverges from main immediately after construction and re-converges exactly when
atom_map is first read, so the exposure window is __init__ → first atom_map access — and
Scheduler.__init__ sits inside it.

Monoatomics are the exception, and the table below was previously wrong about them.
ARCSpecies.__init__ calls _set_final_xyz_for_monoatomic() (arc/species/species.py:559, added
by 57a55de25), which populates final_xyz — and hence cheap_conformer — for every monoatomic
at construction, before any reaction exists. That commit landed after the d033cbbfd baseline
the original measurement was taken against, so monoatomics are unaffected by this PR.

Measured on base vs head, same species:

                    BASE                          HEAD
standalone H    cheap=T init=T final=T        cheap=T init=T final=T
in-reaction H   cheap=T init=T final=T        cheap=T init=T final=T   <-- unchanged
in-reaction CH4 cheap=T init=F final=F        cheap=F init=F final=F
in-reaction C2H6 cheap=T init=F final=F       cheap=F init=F final=F

(a) Polyatomic reaction species no longer carry an MMFF cheap_conformer after construction.
Monoatomics are unchanged: they still receive cheap_conformer/initial_xyz/final_xyz at
construction, from _set_final_xyz_for_monoatomic(). initial_xyz/final_xyz were never set for
polyatomics on either side, so the only thing this PR drops is cheap_conformer on polyatomics.
test_check_atom_balance_does_not_seed_coordinates documents the monoatomic exception.

(b) Under job_types: {opt: false, conf_opt: false} with no user xyz, reaction species now take
the conformer-generation path instead of running freq/sp on the MMFF geometry.
arc/scheduler.py:481 gates on species.get_xyz(generate=False) and not job_types['conf_opt'] and not job_types['opt'], and arc/scheduler.py:1309 gates conformer generation on
get_xyz(generate=False) is None. Both flip.

(c) Under consider_all_diastereomers: false with no user xyz, the diastereomer filter is no
longer pinned to whatever chirality RDKit's randomSeed=1 embedding happened to produce.
arc/species/species.py:1183 takes xyz = self.get_xyz(generate=False) and passes [xyz] as the
diastereomer to match; that xyz is now None.

(b) and (c) can change how many QM jobs a run spawns — the thing most worth a maintainer's
attention. Both make reaction species behave the way standalone species already behave on main,
so both are arguably bug fixes rather than regressions, but they are behaviour changes either way.

(d) as_dict() and restart YAML no longer carry cheap_conformer — for polyatomics. A
monoatomic's as_dict() still carries it at head, for the same reason as (a). test_as_dict
asserted on those exact RDKit-version-dependent strings; the eight dropped expectations are all
polyatomic entries. Restart files remain compatible in both directions: from_dict() guards the
read with if 'cheap_conformer' in species_dict.

atom_map is now a property with a side effect

ARCReaction.atom_map relied on the conformer the constructor left behind: its
get_xyz(generate=False) guard was only ever satisfied because generation had already happened.
It now uses generate=True, so reading the property generates conformers. Same predicate
(still None when a species has neither graph nor coordinates), and the cost moves to where the
geometry is actually used. Without this, six mapping/NMD tests fail.

Residual divergence — constructible, and it is the PR's second bug fix

Reclassified after review: this is a bug fix, not a risk. The case was reachable after all —
ARCSpecies(smiles='C12C3C4C5C1C6C4C3C25C6'), a strained C10 cage that RDKit sanitizes and embeds
but which neither MMFF94s nor UFF can optimise, so get_xyz(generate=True) returns None.

Measured on a genuinely imbalanced reaction built from it (C10H11 on the left, C10H10 on the
right):

BASE b361f605d:  constructed with NO error;  check_atom_balance(raise_error=False) -> True
HEAD b2359c22c:  ReactionError: The Reaction is not atom balanced.  {'C': 10, 'H': 11} vs {'C': 10, 'H': 10}

On base the balance check is skipped (get_xyz(generate=True) is None, the
if xyz is not None and xyz: guard fails, break) and the method returns True — ARC accepts an
unbalanced reaction and proceeds to compute a rate coefficient for it, silently. On head
_get_atom_balance_entry() falls back to the 2D graph, the check runs, and it raises at reaction
construction, before any QM job is submitted.

So the direction is the opposite of a risk: the abort only fires when the reaction really is
imbalanced, and head cannot raise for a genuinely balanced one (the element counts are identical
between the two paths). The class is narrow — highly strained polycyclics and atom types outside
MMFF/UFF coverage; a 19-species oracle including cubane, prismane, tetrahedrane, a 20-membered ring
and trans-cyclooctene did not reach it — but automated mechanism work enumerates thousands of
species with no human reading each one. TODO: add that cage species as a regression test. It is
the only known input where base and head give different verdicts, and it is currently untested.

No call site passes a timeout

All 24 callers pass none, and none is given one here — long QC runs must not acquire a deadline as
a side effect of this PR. The orphan bug is therefore not fixed by this PR, only made fixable.
_run_with_timeout, _kill_process_group and the SIGKILL escalation are infrastructure for a
follow-up that actually sets deadlines.

Two details worth knowing. The escalation is deliberately not conditioned on the direct child
having survived SIGTERM: with shell=True the direct child is a shell that dies on SIGTERM, while
a grandchild that ignores it is exactly the process that gets orphaned. And the child is
deliberately not reaped before the SIGKILL: it is the group leader, so leaving it unreaped keeps
its pid from being recycled and keeps the pgid pinned to this group. A timing-out call is bounded
by timeout + the 5 s grace period.

The except subprocess.CalledProcessError handler and the 30-attempt retry loop it drives are
dead code — subprocess.run() is called without check=True and so never raises it. Now
noted in a comment. The new TimeoutExpired handler is the only live one.

---

Review items — disposition

#	Item	Status
1	pgid recycling in _kill_process_group	Fixed, differently than proposed — see below
2	M11 (SIGTERM-only) survived	Test rebuilt; M11 now killed
3	M4 (mol_list[0] fallback) survived	Kept, now covered by a test; M4 killed
4	Rename to _get_atom_balance_entry	Done
5	_run_with_timeout docstring command type	See below — the annotation is correct as written

Item 1 — the proposed guard would have reintroduced item 2

The recycling defect is real: my code reaped the child inside the SIGTERM branch, freeing the
group-leader pid, and then SIGKILLed that pgid.

The proposed fix, if process.poll() is None:, does not work, because the direct child does die
on SIGTERM. Probed directly:

after SIGTERM: direct child poll()=-15   grandchild alive=True
 -> `if process.poll() is None` would send SIGKILL: False

So that guard skips the escalation in exactly the case item 2 asks me to pin. Run against the
rebuilt test, it fails:

######## coordinator's proposed guard: if process.poll() is None ########
FAILED arc/job/local_test.py::TestLocal::test_execute_command_timeout_kills_spawned_children

The same probe shows what actually closes the window: the unreaped child pins the pgid.

killpg(pgid, 0) with child unreaped: OK (group still exists)

The child is the group leader, so while it is unreaped its pid cannot be recycled and pgid stays
bound to this group. The fix is therefore to stop reaping before the SIGKILL —
time.sleep(grace_period) instead of process.wait(timeout=grace_period), then SIGKILL, then
reap. Recycling window closed, escalation preserved. Cost: a timing-out call now always takes the
full grace period.

Items 2 and 3 — mutation results

Rebuilt test: the command backgrounds a Python grandchild that installs signal.SIG_IGN for
SIGTERM and blocks, while the direct child is a plain sleep that dies on SIGTERM.

Mutant	Result
M11 — SIGKILL escalation removed, SIGTERM only	FAILED (killed) — 1 failed in 32.68 s
M-direct — process.kill() instead of the group	FAILED (killed) — 1 failed in 28.11 s
Proposed — if process.poll() is None: guard	FAILED (killed) — 1 failed in 27.45 s
HEAD implementation	passed in 12.63 s

Mutant	Result
M4 — mol = species.mol, mol_list[0] fallback dropped	FAILED (killed) — test_get_atom_balance_entry, 1 failed / 3 passed
HEAD implementation	4 passed

On item 3 I kept the fallback rather than removing it. Inside ARC the mol is None and mol_list is not None state is indeed unreachable — mol_list is only ever set by set_mol_list(),
which requires self.mol is not None (arc/species/species.py:1077); as_dict() writes a
mol_list key (:796) that from_dict() never reads back; and the only self.mol = None in ARC
is from_dict() at :936, which leaves mol_list at its __init__ value of None. But the state
is reachable through the public attribute, and ARCSpecies.get_xyz() itself contemplates it
(elif generate and (self.mol is not None or self.mol_list is not None), :1266), with
get_cheap_conformer() serving monoatomics and diatomics from mol_list[0] (:1213, :1217).
Dropping the fallback would therefore make the helper diverge from get_xyz in a state get_xyz
explicitly handles. The new test reaches it both directly and through a full reaction.

Item 5 — the annotation is right; I improved the docstring instead

_run_with_timeout cannot receive a str. execute_command() normalises before the retry loop:

if not isinstance(command, list):
    command = [command]
command = [' && '.join(command)]        # <- always a one-element list[str]
i, max_times_to_try = 1, 30
while i < max_times_to_try:
    ...
    _run_with_timeout(command=command, ...)

Widening the annotation to str | list[str] would document a call that cannot happen. I kept
list[str] and made the docstring say why it is a list ("already normalized by
execute_command() into the single-element list that subprocess expects"), which addresses the
readability concern without making the type less accurate. Say the word if you want it widened
anyway.

Also done

• git clean -fd — removed arc/testing/test_GaussianAdapter/, and also
  arc/testing/test_CFourAdapter/ and arc/testing/test_xTBAdapter_1/, which the suite creates
  the same way. Not committed, not gitignored. Working tree is clean.
• Comment added marking the CalledProcessError handler and its retry loop as dead code.

---

Tests

source ~/anaconda3/etc/profile.d/conda.sh && conda activate arc_env
cd /home/alon/Code/ARC-atom-balance
python -m pytest arc/ -ra -q -n auto -p no:randomly

Baseline re-established from scratch (all four touched files restored to d033cbbfd in place),
because an intermediate scratch file had been clobbered and I did not want to trust it.

	passed	failed	errors	skipped
base d033cbbfd	2718	26	72	7
head	2726	26	72	6

FAILED node-ID sets: identical, zero new and zero resolved.
ERROR node-ID sets: identical, 72 = 72.

Subsets:

• pytest arc/reaction/ -ra -q -p no:randomly → 33 passed, 0 failed
• pytest arc/job/local_test.py -ra -q -p no:randomly → 11 passed, 1 failed
  (test_determine_job_id, pre-existing on base — local settings.py has cluster_soft: local)
• pytest arc/job/ -ra -q -n auto -p no:randomly → 1105 passed, 15 failed, 37 errors — same set as
  base, zero new

Net new tests: 7 (3 in reaction_test.py, 4 in local_test.py), all passing.

Re-verified after the review edits: benchmark 1.22 s (from 26.7 s), and the 16-case verdict oracle
still byte-identical between base and head.

Concerns

1. (b) and (c) are the risk in this PR, not the speedup. They can change the number of QM jobs
   a run spawns under two specific configurations. Both align reaction species with standalone
   species, but no functional test was run and I have not exercised a real scheduler run.
2. The silent-pass → ReactionError direction of the residual divergence is the one failure
   mode that could abort a user's run at a place that never failed before. Unreached in ~51
   attempts; I could not construct it.
3. atom_map being side-effecting is a wart. It is the minimal fix, but a property that
   silently runs RDKit is a trap for the next reader; a follow-up might make it explicit.
4. A timing-out execute_command now always waits the full 5 s grace. Short-circuiting would
   require reaping the child, which is exactly what reopens the pid-recycling window. No caller
   passes a timeout, so nothing is affected today.
5. Downstream, T3 still pays ~3.4 s per build_t3() in PyYAML, ~100 times. The remaining fix
   is on the T3 side and is not in this PR.

Fifth behaviour change (undeclared until now): get_cheap_conformer()

arc/species/species.py::get_cheap_conformer() is public API and changed here, so it belongs on
this list. It is strictly more robust in three directions, with no regression:

• It resolves mol = self.mol if self.mol is not None else (self.mol_list[0] if self.mol_list else None) once and returns with a warning if that is None; base raised AttributeError in the
  polyatomic branch.
• The monoatomic and diatomic branches previously read self.mol_list[0] unconditionally and now
  prefer self.mol; for a species with mol set but mol_list still None, base raised
  TypeError.
• The polyatomic branch previously read self.mol unconditionally and now falls back to
  mol_list[0], matching get_xyz(), which contemplates that state at arc/species/species.py:1266.

---

Description corrected 2026-08-22 following review. Changes: (a) and (d) were measured against
d033cbbfd, before _set_final_xyz_for_monoatomic() (57a55de25) landed, and overstated the
monoatomic effect — corrected above; the residual divergence was reclassified from a risk to a bug
fix after the case was constructed; the fifth behaviour change was added. No code changed.

[Copilot]

Pull request overview

This PR optimizes reaction atom-balance checks and adds optional process-group timeouts to execute_command().

Changes:

• Avoids unnecessary 3D conformer generation during atom counting.
• Defers conformer generation for atom mapping.
• Adds SIGTERM/SIGKILL timeout handling and tests.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Summary
arc/reaction/reaction.py	Optimizes atom-balance processing and lazy atom mapping. Critical issues remain for timeout cleanup interaction and polyatomic mol_list-only species; downstream job-selection behavior also lacks regression coverage.
arc/reaction/reaction_test.py	Adds atom-balance and reaction behavior tests.
arc/job/local.py	Adds process-group timeout execution. Critical issue: context-manager cleanup may perform an unbounded wait after timeout.
arc/job/local_test.py	Adds timeout and backward-compatibility tests.

Suppressed comments (3)

arc/job/local.py:54

• This Raises entry says erroneous commands raise SettingsError, but both subprocess.run() and the new Popen path run without check=True; as the comment below acknowledges, a non-zero exit status never raises CalledProcessError and is returned as output. The documented error contract should describe the only live error path here (timeouts), or the implementation should actually check return codes.

    Raises:
        SettingsError: If the command is erroneous or timed out, and ``no_fail`` is ``False``.

arc/reaction/reaction.py:161

• atom_map is a property, but this read can now run RDKit embedding/force-field generation and mutate every reaction species lacking coordinates. The current one-line docstring does not expose that potentially expensive and failure-prone side effect, so callers can unexpectedly trigger conformer generation merely by inspecting the map. Please document this behavior (or make mapping/conformer generation explicit).

                and all(species.get_xyz(generate=True) is not None for species in self.r_species + self.p_species):

arc/reaction/reaction.py:1123

• The new graph fallback changes an observable failure path: for a non-TS species with a graph but no coordinates, get_xyz(generate=True) can return None after get_cheap_conformer() and generate_conformers(n_confs=1) produce no conformers (see arc/species/species.py:1266-1273). The old code then emptied the well and skipped the balance check, whereas this helper constructs a well and can raise ReactionError for an otherwise previously accepted reaction. Please preserve or explicitly handle this conformer-failure case rather than treating every graph as a successful replacement.

    mol = species.mol if species.mol is not None else (species.mol_list[0] if species.mol_list else None)
    if mol is None or not len(mol.atoms):
        return ''
    return '\n'.join(f'{atom.element.symbol} 0.0 0.0 0.0' for atom in mol.atoms)

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 64.60%. Comparing base (a08d314) to head (fc58df1).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1003      +/-   ##
==========================================
+ Coverage   64.55%   64.60%   +0.04%     
==========================================
  Files         119      119              
  Lines       39788    39843      +55     
  Branches    10307    10314       +7     
==========================================
+ Hits        25684    25739      +55     
- Misses      11123    11129       +6     
+ Partials     2981     2975       -6     

Flag	Coverage Δ
functionaltests	64.60% <ø> (+0.04%)	⬆️
unittests	64.60% <ø> (+0.04%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Suppressed comments (5)

arc/job/local.py:205

• After the unconditional grace_period sleep, this wait can block for another full grace_period when the direct child cannot be reaped. The timeout path can therefore take timeout + 2 * grace_period (up to 10 seconds beyond the deadline), rather than the timeout + 5 seconds bound stated in the PR description. Please avoid a second blocking grace period or document the actual upper bound.

    try:
        process.wait(timeout=grace_period)
    except subprocess.TimeoutExpired:

arc/job/local_test.py:139

• The PID-writing grandchild is started asynchronously, but the whole setup is given only three seconds before the timeout fires. On a loaded CI worker the process can be killed before it reaches the file write, making the assertion at line 140 fail even when process-group cleanup is correct. Synchronize on the PID file (or make the shell wait for it) before entering the long sleep.

        command = f'{sys.executable} {script_path} {pid_path} & sleep 300'
        try:
            with self.assertRaises(SettingsError):
                local.execute_command(command, timeout=3)

arc/reaction/reaction.py:707

• This change removes the constructor's geometry side effect, which flips existing scheduler/conformer paths for graph-only polyatomics: with both opt and conf_opt false, run_conformer_jobs() is now entered, and consider_all_diastereomers=False no longer receives the embedded xyz. Those paths can change conformer and QM-job counts, but the added tests only inspect species state; please add scheduler-level coverage that asserts the selected jobs and diastereomer inputs for these configurations.

            entry = _get_atom_balance_entry(species=reactant)
            if entry:
                r_well += (entry + '\n') * count

arc/reaction/reaction.py:168

• When get_xyz(generate=True) cannot produce a geometry, this all(...) stays false and _atom_map remains None. Every subsequent read of atom_map therefore retries get_cheap_conformer() and generate_conformers() for the same unembeddable species, repeatedly paying the failed RDKit path instead of failing cheaply. Cache the failed attempt or add a failure sentinel so the property does not repeat the generation.

                and all(species.get_xyz(generate=True) is not None for species in self.r_species + self.p_species):

arc/reaction/reaction.py:705

• By removing the constructor-time conformer, graph-only reaction species with force_field == 'cheap' now enter Scheduler.run_conformer_jobs() under opt=False, conf_opt=False. That branch sets initial_xyz from get_xyz() and then process_conformers() skips its job-spawn block because initial_xyz is already set, so requested freq/SP jobs are never submitted; the old reaction path did submit them from the seeded geometry. Please preserve the downstream scheduling or update the cheap path to continue with non-opt jobs.

            entry = _get_atom_balance_entry(species=reactant)