Consolidate more benchmark code, make ff optimize more robust

benchmarks/bench_utils/__init__.py

@@ -20,7 +20,13 @@
 """
 
 from bench_utils.loaders import load_pickle, load_sdf, load_smarts, load_smiles
-from bench_utils.molprep import clone_mols_with_conformers, embed_and_jitter, perturb_conformer, prep_mols
+from bench_utils.molprep import (
+    available_cpu_count,
+    clone_mols_with_conformers,
+    embed_and_jitter,
+    perturb_conformer,
+    prep_mols,
+)
 from bench_utils.timing import (
     Deadline,
     TimingResult,
@@ -34,6 +40,7 @@
     "Deadline",
     "TimingResult",
     "add_rdkit_max_seconds_arg",
+    "available_cpu_count",
     "clone_mols_with_conformers",
     "embed_and_jitter",
     "load_pickle",

benchmarks/bench_utils/loaders.py

@@ -23,6 +23,7 @@
 import pickle
 import random
 from functools import partial
+from math import ceil
 from typing import Iterator
 
 from rdkit import Chem, RDLogger
@@ -33,7 +34,17 @@ def _mol_from_binary(binary_mol: bytes) -> Chem.Mol:
     return Chem.Mol(binary_mol)
 
 
-def load_pickle(filepath: str, max_count: int = 0, seed: int | None = None) -> list[Chem.Mol]:
+def _buffered_count(max_count: int) -> int:
+    """Return a 10% candidate reserve for a positive requested count."""
+    return max_count + max(1, ceil(max_count * 0.1)) if max_count > 0 else 0
+
+
+def load_pickle(
+    filepath: str,
+    max_count: int = 0,
+    seed: int | None = None,
+    keep_buffer: bool = False,
+) -> list[Chem.Mol]:
     """Load molecules from a pickle file containing a list of RDKit binary molecules.
 
     Args:
@@ -42,6 +53,7 @@ def load_pickle(filepath: str, max_count: int = 0, seed: int | None = None) -> l
         max_count: When positive and the source has more entries, draw a
             uniform random sample of this size before unpickling.
         seed: Optional seed for the sampling RNG.
+        keep_buffer: Return the 10% candidate reserve instead of trimming it.
 
     Returns:
         List of parsed RDKit molecules. The list is always shuffled
@@ -51,8 +63,9 @@ def load_pickle(filepath: str, max_count: int = 0, seed: int | None = None) -> l
     with open(filepath, "rb") as fh:
         binary_mols = pickle.load(fh)
     rng = random.Random(seed)
-    if max_count > 0 and len(binary_mols) > max_count:
-        binary_mols = rng.sample(binary_mols, max_count)
+    sample_count = _buffered_count(max_count) if keep_buffer else max_count
+    if sample_count > 0 and len(binary_mols) > sample_count:
+        binary_mols = rng.sample(binary_mols, sample_count)
     else:
         binary_mols = list(binary_mols)
         rng.shuffle(binary_mols)
@@ -94,6 +107,7 @@ def load_smiles(
     max_count: int = 0,
     sanitize: bool = True,
     seed: int | None = None,
+    keep_buffer: bool = False,
 ) -> list[Chem.Mol]:
     """Load and parse molecules from a SMILES file.
 
@@ -105,8 +119,15 @@ def load_smiles(
     The returned list is always shuffled (deterministic with ``seed``) so
     benches that consume a head slice get a representative cross-section
     rather than file-order bias (some upstream files are sorted by size).
+
+    Args:
+        filepath: Path to the SMILES file.
+        max_count: Maximum number of requested molecules; zero loads all.
+        sanitize: Sanitize molecules while parsing.
+        seed: Optional seed for reservoir sampling and shuffling.
+        keep_buffer: Return the 10% candidate reserve instead of trimming it.
     """
-    read_limit = int(max_count * 1.1) if max_count > 0 else 0
+    read_limit = _buffered_count(max_count)
     rng = random.Random(seed)
 
     if read_limit > 0:
@@ -135,7 +156,7 @@ def load_smiles(
         if parse_failures > 0:
             print(f"    ({parse_failures} parse failures)")
 
-    if max_count > 0 and len(mols) > max_count:
+    if not keep_buffer and max_count > 0 and len(mols) > max_count:
         mols = rng.sample(mols, max_count)
     else:
         rng.shuffle(mols)
@@ -179,15 +200,24 @@ def load_sdf(
     seed: int | None = None,
     removeHs: bool = False,
     sanitize: bool = True,
+    keep_buffer: bool = False,
 ) -> list[Chem.Mol]:
     """Load molecules from an SDF file with optional reservoir sampling.
 
     The returned list is always shuffled (deterministic with ``seed``) so
     benches that consume a head slice get a representative cross-section
     rather than file-order bias (some upstream files are sorted by size).
+
+    Args:
+        filepath: Path to the SDF file.
+        max_count: Maximum number of requested molecules; zero loads all.
+        seed: Optional seed for reservoir sampling and shuffling.
+        removeHs: Remove hydrogens while reading the SDF.
+        sanitize: Sanitize molecules while reading the SDF.
+        keep_buffer: Return the 10% candidate reserve instead of trimming it.
     """
     supplier = Chem.SDMolSupplier(filepath, removeHs=removeHs, sanitize=sanitize)
-    read_limit = int(max_count * 1.1) if max_count > 0 else 0
+    read_limit = _buffered_count(max_count)
     rng = random.Random(seed)
 
     parse_failures = 0
@@ -214,7 +244,7 @@ def load_sdf(
                 continue
             mols.append(mol)
 
-    if max_count > 0 and len(mols) > max_count:
+    if not keep_buffer and max_count > 0 and len(mols) > max_count:
         mols = rng.sample(mols, max_count)
     else:
         rng.shuffle(mols)

benchmarks/bench_utils/molprep.py

@@ -15,6 +15,7 @@
 
 """Molecule preparation helpers shared across nvMolKit benchmarks."""
 
+import os
 import random
 from functools import partial
 
@@ -28,6 +29,17 @@
 JITTER_SPREAD = 0.6
 
 
+def available_cpu_count() -> int:
+    """Return CPUs available to this process across supported Python versions."""
+    process_cpu_count = getattr(os, "process_cpu_count", None)
+    if process_cpu_count is not None:
+        return process_cpu_count() or 1
+    try:
+        return len(os.sched_getaffinity(0))
+    except AttributeError:
+        return os.cpu_count() or 1
+
+
 def prep_mols(
     mols: list[Chem.Mol],
     *,

benchmarks/conformer_rmsd_bench.py

@@ -21,13 +21,12 @@
 
 import argparse
 import csv
-import multiprocessing as mp
 import statistics
 import time
 from pathlib import Path
 
 import torch
-from bench_utils import Deadline, add_rdkit_max_seconds_arg, embed_and_jitter, load_smiles
+from bench_utils import Deadline, add_rdkit_max_seconds_arg, available_cpu_count, embed_and_jitter, load_smiles
 from benchmark_timing import time_it
 from rdkit import Chem
 from rdkit.Chem import AllChem
@@ -42,7 +41,7 @@ def prepare_mols(
     num_workers: int,
 ) -> list[Chem.Mol]:
     """Embed one base conformer per mol, then perturb to ``confs_per_mol``."""
-    workers = num_workers if num_workers > 0 else max(1, mp.cpu_count() // 2)
+    workers = num_workers if num_workers > 0 else max(1, available_cpu_count() // 2)
     return embed_and_jitter(
         raw_mols,
         confs_per_mol=confs_per_mol,

benchmarks/ff_optimize_bench.py

@@ -41,6 +41,7 @@
 from bench_utils import (
     Deadline,
     add_rdkit_max_seconds_arg,
+    available_cpu_count,
     clone_mols_with_conformers,
     embed_and_jitter,
     load_pickle,
@@ -50,10 +51,11 @@
     throughput_per_s,
     time_it,
 )
+from rdkit import Chem
+from rdkit.Chem import AllChem
+
 from nvmolkit import autotune as nv_autotune
 from nvmolkit.types import HardwareOptions
-from rdkit import Chem
-from rdkit.Chem import AllChem, rdDistGeom
 
 OPTUNA_AVAILABLE = nv_autotune.is_available()
 
@@ -273,6 +275,12 @@ def main() -> None:
         default=1,
         help="Threads passed to RDKit FF optimizer via numThreads (default: 1)",
     )
+    parser.add_argument(
+        "--conf_gen_workers",
+        type=int,
+        default=0,
+        help="Processes used to generate base conformers (default: 0 = all available CPUs)",
+    )
     add_rdkit_max_seconds_arg(
         parser,
         extra_help="The RDKit FF optimizer loop stops at the next molecule boundary once the budget is hit.",
@@ -391,6 +399,8 @@ def main() -> None:
     print(f"  Validate (energy diffs): {args.validate}")
     print(f"  Run nvmolkit: {not args.no_nvmolkit}")
     print(f"  Run RDKit: {not args.no_rdkit}")
+    conf_gen_workers = args.conf_gen_workers if args.conf_gen_workers > 0 else available_cpu_count()
+    print(f"  Conformer generation workers: {conf_gen_workers}")
     if not args.no_rdkit:
         print(f"  RDKit threads: {args.rdkit_threads}")
     if not args.no_nvmolkit:
@@ -402,24 +412,33 @@ def main() -> None:
 
     print("\nLoading molecules...")
     if args.smiles:
-        raw_mols = load_smiles(args.smiles, args.num_mols, args.sanitize, seed=args.seed)
+        raw_mols = load_smiles(args.smiles, args.num_mols, args.sanitize, seed=args.seed, keep_buffer=True)
     elif args.sdf:
-        raw_mols = load_sdf(args.sdf, args.num_mols, seed=args.seed, sanitize=args.sanitize)
+        raw_mols = load_sdf(args.sdf, args.num_mols, seed=args.seed, sanitize=args.sanitize, keep_buffer=True)
     else:
-        raw_mols = load_pickle(args.pickle, args.num_mols, seed=args.seed)
+        raw_mols = load_pickle(args.pickle, args.num_mols, seed=args.seed, keep_buffer=True)
     if not raw_mols:
         print("Error: No valid molecules loaded")
         sys.exit(1)
 
     print("\nPreparing molecules (AddHs / sanitize / clear conformers)...")
     mols = prep_mols(raw_mols)
+    if args.ff == "mmff":
+        mols = [mol for mol in mols if AllChem.MMFFHasAllMoleculeParams(mol)]
+    else:
+        mols = [mol for mol in mols if AllChem.UFFHasAllMoleculeParams(mol)]
+    if args.num_mols > 0:
+        mols = mols[: args.num_mols]
     if not mols:
-        print("Error: No molecules survived preparation")
+        print(f"Error: No molecules survived preparation with {args.ff.upper()} parameters")
         sys.exit(1)
     print(f"  {len(mols)} molecules ready")
 
-    print(f"\nEmbedding {args.confs_per_mol} conformer(s) per molecule with RDKit ETKDGv3...")
-    mols = embed_and_jitter(mols, args.confs_per_mol, seed=args.seed, num_workers=args.rdkit_threads)
+    print(
+        f"\nEmbedding one RDKit ETKDGv3 base conformer per molecule "
+        f"with {conf_gen_workers} processes, then jittering to {args.confs_per_mol}..."
+    )
+    mols = embed_and_jitter(mols, args.confs_per_mol, seed=args.seed, num_workers=conf_gen_workers)
     if not mols:
         print("Error: No molecules retained after embedding")
         sys.exit(1)

benchmarks/tests/test_loaders.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Unit tests for benchmark molecule-loader candidate reserves."""
+
+import os
+
+from bench_utils import loaders, molprep
+
+
+def test_load_smiles_can_return_candidate_buffer(tmp_path, monkeypatch):
+    smiles_path = tmp_path / "mols.smi"
+    smiles_path.write_text("\n".join(["C", "CC", "CCC", "CCCC", "CCCCC", "CCCCCC"]))
+    monkeypatch.setattr(loaders, "process_map", lambda fn, values, **_kwargs: [fn(value) for value in values])
+
+    trimmed = loaders.load_smiles(str(smiles_path), max_count=3, seed=42)
+    buffered = loaders.load_smiles(str(smiles_path), max_count=3, seed=42, keep_buffer=True)
+
+    assert len(trimmed) == 3
+    assert len(buffered) == 4
+
+
+def test_buffered_count_reserves_at_least_one_candidate():
+    assert loaders._buffered_count(0) == 0
+    assert loaders._buffered_count(1) == 2
+    assert loaders._buffered_count(1000) == 1100
+
+
+def test_available_cpu_count_supports_python_312(monkeypatch):
+    monkeypatch.delattr(os, "process_cpu_count", raising=False)
+    monkeypatch.setattr(os, "sched_getaffinity", lambda _pid: {2, 4, 6}, raising=False)
+
+    assert molprep.available_cpu_count() == 3

benchmarks/tfd_bench.py

@@ -32,7 +32,6 @@
 """
 
 import argparse
-import multiprocessing
 import os
 import pickle
 import sys
@@ -41,7 +40,7 @@
 
 import pandas as pd
 import torch
-from bench_utils import embed_and_jitter, load_smiles
+from bench_utils import available_cpu_count, embed_and_jitter, load_smiles
 from rdkit import Chem
 from rdkit.Chem import TorsionFingerprints
 
@@ -89,7 +88,7 @@ def generate_conformers_batch(
     """
     if num_confs < 2:
         raise ValueError(f"num_confs must be >= 2 for TFD, got {num_confs}")
-    workers = num_workers if num_workers > 0 else max(1, multiprocessing.cpu_count() // 2)
+    workers = num_workers if num_workers > 0 else max(1, available_cpu_count() // 2)
     return embed_and_jitter(
         mols,
         confs_per_mol=num_confs,

benchmarks/tfd_prepare_mols.py

@@ -42,6 +42,7 @@
 from functools import partial
 
 import pandas as pd
+from bench_utils import available_cpu_count
 from rdkit import Chem, RDLogger
 from rdkit.Chem import AllChem
 
@@ -98,11 +99,12 @@ def main():
         default=0,
         help="Maximum number of molecules to prepare (default: 0 = all valid SMILES)",
     )
+    default_workers = max(1, available_cpu_count() // 2)
     parser.add_argument(
         "--workers",
         type=int,
-        default=max(1, multiprocessing.cpu_count() // 2),
-        help=f"Number of parallel workers (default: {max(1, multiprocessing.cpu_count() // 2)})",
+        default=default_workers,
+        help=f"Number of parallel workers (default: {default_workers})",
     )
     args = parser.parse_args()