diff --git a/benchmarks/bench_utils/__init__.py b/benchmarks/bench_utils/__init__.py index 4cabcc21..f238e5bb 100644 --- a/benchmarks/bench_utils/__init__.py +++ b/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", diff --git a/benchmarks/bench_utils/loaders.py b/benchmarks/bench_utils/loaders.py index d68f92d4..291eecd9 100644 --- a/benchmarks/bench_utils/loaders.py +++ b/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) diff --git a/benchmarks/bench_utils/molprep.py b/benchmarks/bench_utils/molprep.py index 21b88481..a37373b6 100644 --- a/benchmarks/bench_utils/molprep.py +++ b/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], *, diff --git a/benchmarks/conformer_rmsd_bench.py b/benchmarks/conformer_rmsd_bench.py index e9123a0b..b3314cb4 100644 --- a/benchmarks/conformer_rmsd_bench.py +++ b/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, diff --git a/benchmarks/ff_optimize_bench.py b/benchmarks/ff_optimize_bench.py index 3e1bc4fb..1147c62e 100644 --- a/benchmarks/ff_optimize_bench.py +++ b/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) diff --git a/benchmarks/tests/test_loaders.py b/benchmarks/tests/test_loaders.py new file mode 100644 index 00000000..5c4c25bd --- /dev/null +++ b/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 diff --git a/benchmarks/tfd_bench.py b/benchmarks/tfd_bench.py index b0b23979..97bae1c4 100644 --- a/benchmarks/tfd_bench.py +++ b/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, diff --git a/benchmarks/tfd_prepare_mols.py b/benchmarks/tfd_prepare_mols.py index 4f9eec39..fc3bbfdd 100644 --- a/benchmarks/tfd_prepare_mols.py +++ b/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()