[niklashoelter]

Fixes the slowdown introduced by #40 while keeping its fix.

Problem

#40 stopped capping max_atoms_to_try at the job size, so the probe ceiling can no longer land below the job's own largest structure. As a side effect, every optimize_structure_batch call now runs five probe sequences, each growing replicated copies of one system towards 100k atoms or OOM: two in _chunked_apply's BinningAutoBatcher and three in InFlightAutoBatcher._get_first_batch.

Profiled on a Quadro RTX 6000 with UMA-s on 36 conformers (927 atoms): 838 s total, of which 754 s is probing and ~20 s is the optimization itself.

Fix

Before optimizing, run one forward pass over the whole job:

• Fits within max_memory_padding of the GPU budget: set max_memory_scaler to the job's total n_edges metric. Everything runs as one batch and neither batcher probes. The ceiling can't be below any structure in the job, so optimizer: size the autobatcher by the GPU, not by the input #40's fix holds.
• OOM, too little headroom, or job larger than max_atoms_to_try: leave the scaler unset, and torch-sim probes on the OptimState exactly as after optimizer: size the autobatcher by the GPU, not by the input #40.

The "Memory estimation" timing phase is timed again and now covers this check.

Results

examples/full_workflow_benchmark.py, FIRE, fconv 5e-2, dedicated GPU (optimization time):

dataset	model	0.7.0	0.8.0	this PR
small	uma-s-1p2	123 s	1659 s	81 s
small	orb-v3-direct	65 s	66 s	14 s
small	7net-0	29 s	128 s	18 s
medium	uma-s-1p2	165 s*	1514 s	373 s
medium	orb-v3-direct	13 s*	64 s	11 s
medium	7net-0	32 s*	149 s	22 s

*0.7.0 medium had 189 structures; 0.8.0 and this PR have 243/244 (#33 keeps more conformers).

• Every structure succeeded in every run.
• Energies are identical to 0.8.0 on the same inputs.
• The large set exceeds max_atoms_to_try, so it takes the unchanged probe path.

UMA medium vs 0.7.0: UMA gets slower per structure as batches grow. With the scaler forced to 1/2, 1/4 and 1/8 of the job, medium drops from 431 s to 316, 253 and 217 s. The cause is in fairchem: the MoLE layer (uma/nn/mole.py) loops over systems once per activation-checkpoint chunk, and the backward pass takes 0.2 s at 30 systems vs 11 s at 240. 0.7.0's smaller batches were a side effect of its capped probe. This is left for an upstream fairchem fix rather than capped here.

Tests

5 new GPU tests in tests/optimization/test_batch.py::TestAutobatcherMemoryProbe:

• A job that fits never calls determine_max_batch_size. This fails on master with 5 probe calls.
• A single structure and a small-then-large job still work (the optimizer: size the autobatcher by the GPU, not by the input #40 regression cases).
• An OOM on the whole-job pass falls back to the probe.
• A whole-job pass without enough headroom falls back to the probe.

Full suite with HF_TOKEN set: 299 passed, 3 skipped, in 6 min 38 s (22 min 47 s on master, mostly probing).