[akwarii]

Summary

With autobatcher=True, a call probed the GPU memory 2 (ts.static, ts.integrate) or 5 (ts.optimize) times.

This PR introduces the following changes to reduce the number of probes:

• Probe identical systems once. estimate_max_memory_scaler probes the smallest and the largest system; for a single system or identical ones they are the same, so the second probe is skipped.
• Reuse the init's estimate in ts.optimize. _chunked_apply returns the limit its BinningAutoBatcher estimated, and the in-flight autobatcher takes it instead of probing again. The autobatcher's cutoff is now passed to _chunked_apply too, so that with memory_scales_with="n_edges" the reused limit counts edges at the in-flight autobatcher's cutoff.
• No bootstrap probe when all states are known. When InFlightAutoBatcher.load_states gets a list or a batched SimState, it estimates the limit over all of them, with the metrics of the batched state as BinningAutoBatcher.load_states does. Lazy iterators keep the bootstrap probe of the first state.

ts.static and ts.integrate now probe once for identical systems (twice otherwise), and ts.optimize once or twice instead of 5 times. The limits are estimated the same way, except that the in-flight autobatcher's is now estimated over all the states instead of its first batch.

Closes #640

Checklist

Before a pull request can be merged, the following items must be checked:

• Doc strings have been added in the Google docstring format.
• Run ruff on your code.
• Tests have been added for any new functionality or bug fixes.

[lkkmpn]

Somewhat unrelated to the original issue, but I noticed this PR and just tested it on my own workload (molecules with sizes ranging from 9 to 19 atoms). In addition to resolving the multiple memory probe calls, it also resolved an out-of-memory error I got on the original behavior. My guess is that it's related to your mention of the in-flight autobatcher's memory estimate now being based on all states instead of the first batch, which could have caused underestimation of the maximum memory usage (although I have not investigated this in more detail).

[lil-lon]

Thanks for this PR! The fix basically looks good to me.

The tests could be simpler by patching determine_max_batch_size instead of measure_model_memory_forward, as the probe itself is tested separately by the test_determine_max_batch_size_* tests. With one call per probe, the tests can assert which systems are probed (e.g. [8, 256]) rather than counting the batches of one system, and the OOM threshold is no longer needed.
A shared fixture would also replace the five copies of mock_measure:

@pytest.fixture
def probed_n_atoms(monkeypatch: pytest.MonkeyPatch) -> list[int]:
    """Record the number of atoms of each system whose batch size is probed."""
    probed: list[int] = []

    def recording_probe(state: ts.SimState, *_args: Any, **_kwargs: Any) -> int:
        probed.append(state.n_atoms)
        return 200

    monkeypatch.setattr(ts.autobatching, "determine_max_batch_size", recording_probe)
    return probed

With it, the first two tests can be merged by parametrizing identical, with expected being [8] or [8, 256].

[akwarii]

Thank you for the review!

• The tests now patch determine_max_batch_size with your fixture and assert which systems are probed instead of counting the batches. I also merged 2 tests into one parametrized one.

• memory_scaling_factor is now correctly passed to _chunked_apply within ts.runners.optimize.

• After running Claude on the changes, it also noticed a bug within autobatching.py that I fixed. I quote:

on the lazy-iterator path, autobatching.py:1054 computes the first state’s metric without self.cutoff, so memory_scales_with="n_edges" with a custom cutoff counts that first state at 6.0 Å.

[lil-lon]

LGTM, thanks for addressing my comments and catching an issue!
I think this is a really nice performance improvement!