Python bindings and benchmark for substructlib

benchmarks/tests/test_substruct_library_bench.py

@@ -0,0 +1,122 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+import pytest
+from bench_utils import Deadline
+from rdkit import Chem
+from rdkit.Chem import rdSubstructLibrary
+from substruct_library_bench import (
+    Measurement,
+    benchmark_nvmolkit,
+    benchmark_rdkit,
+    first_mismatch,
+    load_queries,
+    result_row,
+    run_rdkit_queries,
+)
+
+TARGETS = ["CCO", "c1ccccc1O", "CC(=O)O", "CCN", "c1ccncc1", "OCCO"]
+QUERIES = ["[OX2H]", "c1ccccc1", "[#7]", "C=O", "[Si]"]
+
+
+@pytest.fixture
+def mols():
+    return [Chem.MolFromSmiles(smiles) for smiles in TARGETS]
+
+
+@pytest.fixture
+def queries():
+    return [Chem.MolFromSmarts(smarts) for smarts in QUERIES]
+
+
+def expected_matches(mols, queries):
+    return [[index for index, mol in enumerate(mols) if mol.HasSubstructMatch(query)] for query in queries]
+
+
+@pytest.mark.parametrize("holder", ["mol", "cached-pattern"])
+def test_rdkit_backend_matches_direct_substructure_search(mols, queries, holder):
+    measurements = benchmark_rdkit(
+        mols,
+        queries,
+        operations=["has", "count", "get"],
+        holder=holder,
+        num_threads=1,
+        max_results=-1,
+        runs=1,
+        max_seconds=0,
+    )
+    expected = expected_matches(mols, queries)
+    assert measurements["get"].results == expected
+    assert measurements["count"].results == [len(matches) for matches in expected]
+    assert measurements["has"].results == [bool(matches) for matches in expected]
+    assert all(measurement.completed_queries == len(queries) for measurement in measurements.values())
+
+
+def test_rdkit_queries_stop_at_the_deadline(mols, queries):
+    library = rdSubstructLibrary.SubstructLibrary(rdSubstructLibrary.MolHolder())
+    for mol in mols:
+        library.AddMol(mol)
+
+    class ExpiresAfter(Deadline):
+        def __init__(self, checks):
+            super().__init__(0)
+            self.checks = checks
+
+        def expired(self):
+            self.checks -= 1
+            return self.checks < 0
+
+    assert run_rdkit_queries(library, queries, "get", -1, 1, ExpiresAfter(2)) == expected_matches(mols, queries)[:2]
+
+
+def test_validation_compares_the_completed_prefix():
+    assert first_mismatch([[0, 2], [1], [3]], [[0, 2], [1]], "get") is None
+    assert "query 1" in first_mismatch([[0, 2], [1]], [[0, 2], [1, 4]], "get")
+
+
+def test_result_row_reports_throughput_over_completed_queries():
+    measurement = Measurement(
+        staging_ms=10, finalize_ms=20, steady_ms=50, steady_std_ms=1, results=[[1], [], [2, 3]], completed_queries=3
+    )
+    row = result_row(measurement, operation="get", num_mols=100, num_queries=4, backend="nvmolkit")
+
+    assert row["num_queries"] == 4
+    assert row["positive_queries"] == 2
+    assert row["amortized_ms"] == 80
+    assert row["steady_queries_per_s"] == pytest.approx(60)
+    assert row["steady_pairs_per_s"] == pytest.approx(6000)
+    assert row["amortized_queries_per_s"] == pytest.approx(37.5)
+
+
+def test_smiles_queries_ignore_stereochemistry(tmp_path):
+    path = tmp_path / "queries.smi"
+    path.write_text("C[C@H](O)CC\n")
+    args = SimpleNamespace(smarts=None, query_smiles=str(path), num_queries=0, sanitize=True, seed=0)
+
+    (query,) = load_queries(args)
+    assert Chem.MolToSmiles(query) == "CCC(C)O"
+
+
+def test_nvmolkit_backend_matches_rdkit_in_every_query_mode(mols, queries):
+    torch = pytest.importorskip("torch")
+    if not torch.cuda.is_available():
+        pytest.skip("requires a CUDA device")
+    from nvmolkit.substructure import SubstructSearchConfig
+
+    measurements = benchmark_nvmolkit(
+        mols,
+        queries,
+        operations=["has", "count", "get"],
+        query_modes=["serial", "concurrent"],
+        config=SubstructSearchConfig(),
+        max_results=-1,
+        runs=1,
+        warmups=0,
+    )
+    expected = expected_matches(mols, queries)
+    for mode in ("serial", "concurrent"):
+        assert measurements[("get", mode)].results == expected
+        assert measurements[("count", mode)].results == [len(matches) for matches in expected]
+        assert measurements[("has", mode)].results == [bool(matches) for matches in expected]

docs/api/nvmolkit.rst

@@ -146,6 +146,7 @@ Substructure Search
 
    substructure.SubstructSearchConfig
    substructure.SubstructMatchResults
+   substruct_library.SubstructLibrary
 
 Conformer RMSD
 --------------

docs/conf.py

@@ -108,6 +108,7 @@ def __getattr__(self, name):
     "_Fingerprints",
     "_mcs",
     "_mmffOptimization",
+    "_substructLibrary",
     "_substructure",
     "_TFD",
     "_types",

docs/index.rst

@@ -186,6 +186,7 @@ nvMolKit currently supports the following features:
 
 * **Substructure Search**: GPU-accelerated substructure matching against batches of molecules
     * Supports SMILES and recursive SMARTS-based query molecules via RDKit
+    * ``SubstructLibrary`` keeps target molecules on one or more GPUs for repeated queries, as a GPU counterpart to RDKit's ``SubstructLibrary``
     * Does not yet support chirality-aware matching, enhanced stereochemistry, or other advanced RDKit ``SubstructMatchParameters`` options
 
 * **Maximum Common Substructure (MCS)**: GPU-accelerated MCS search across batches of molecule pairs

nvmolkit/CMakeLists.txt

@@ -138,6 +138,14 @@ target_include_directories(_substructure PUBLIC ${Boost_INCLUDE_DIRS}
                                                 ${Python_INCLUDE_DIRS})
 installpythontarget(_substructure ./)
 
+add_library(_substructLibrary MODULE substructLibrary.cpp)
+target_link_libraries(_substructLibrary PUBLIC ${Boost_LIBRARIES}
+                                               ${PYTHON_LIBRARIES})
+target_link_libraries(_substructLibrary PRIVATE ${RDKit_LIBS} substruct_library)
+target_include_directories(_substructLibrary PUBLIC ${Boost_INCLUDE_DIRS}
+                                                    ${Python_INCLUDE_DIRS})
+installpythontarget(_substructLibrary ./)
+
 add_library(_mcs MODULE mcs.cpp)
 target_link_libraries(_mcs PUBLIC ${Boost_LIBRARIES} ${PYTHON_LIBRARIES})
 target_link_libraries(_mcs PRIVATE ${RDKit_LIBS} mcs_search CUDA::cudart)

nvmolkit/substructLibrary.cpp

@@ -0,0 +1,109 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+#include <GraphMol/ROMol.h>
+
+#include <boost/noncopyable.hpp>
+#include <boost/python.hpp>
+#include <boost/python/stl_iterator.hpp>
+#include <cstddef>
+#include <vector>
+
+#include "src/substruct/substruct_library.h"
+
+namespace {
+
+using namespace boost::python;
+
+class ScopedGilRelease {
+ public:
+  ScopedGilRelease() : state_(PyEval_SaveThread()) {}
+  ScopedGilRelease(const ScopedGilRelease&)            = delete;
+  ScopedGilRelease& operator=(const ScopedGilRelease&) = delete;
+  ~ScopedGilRelease() { PyEval_RestoreThread(state_); }
+
+ private:
+  PyThreadState* state_;
+};
+
+unsigned int addMolecule(nvMolKit::SubstructLibrary& library, const RDKit::ROMol& molecule) {
+  const ScopedGilRelease release;
+  return library.addMol(molecule);
+}
+
+list addMolecules(nvMolKit::SubstructLibrary& library, const object& molecules) {
+  std::vector<object>              owners;
+  std::vector<const RDKit::ROMol*> pointers;
+  stl_input_iterator<object>       iterator(molecules), end;
+  for (; iterator != end; ++iterator) {
+    owners.push_back(*iterator);
+    pointers.push_back(extract<const RDKit::ROMol*>(owners.back()));
+  }
+
+  std::vector<unsigned int> ids;
+  {
+    const ScopedGilRelease release;
+    ids = library.addMols(pointers);
+  }
+
+  list result;
+  for (const unsigned int id : ids) {
+    result.append(id);
+  }
+  return result;
+}
+
+void finalize(nvMolKit::SubstructLibrary& library) {
+  const ScopedGilRelease release;
+  library.finalize();
+}
+
+list getMatches(nvMolKit::SubstructLibrary& library, const RDKit::ROMol& query, int maxResults) {
+  std::vector<unsigned int> ids;
+  {
+    const ScopedGilRelease release;
+    ids = library.getMatches(query, maxResults);
+  }
+  list result;
+  for (const unsigned int id : ids) {
+    result.append(id);
+  }
+  return result;
+}
+
+std::size_t countMatches(nvMolKit::SubstructLibrary& library, const RDKit::ROMol& query) {
+  const ScopedGilRelease release;
+  return library.countMatches(query);
+}
+
+bool hasMatch(nvMolKit::SubstructLibrary& library, const RDKit::ROMol& query) {
+  const ScopedGilRelease release;
+  return library.hasMatch(query);
+}
+
+// Accessors wait on the library's lock, so they release the GIL like the operations do.
+template <std::size_t (nvMolKit::SubstructLibrary::*Accessor)() const>
+std::size_t withoutGil(const nvMolKit::SubstructLibrary& library) {
+  const ScopedGilRelease release;
+  return (library.*Accessor)();
+}
+
+}  // namespace
+
+BOOST_PYTHON_MODULE(_substructLibrary) {
+  import("rdkit.Chem.rdchem");
+  import("nvmolkit._substructure");
+
+  class_<nvMolKit::SubstructLibrary, boost::noncopyable>(
+    "SubstructLibrary",
+    init<nvMolKit::SubstructSearchConfig>((arg("config") = nvMolKit::SubstructSearchConfig())))
+    .def("addMol", &addMolecule)
+    .def("addMols", &addMolecules)
+    .def("finalize", &finalize)
+    .def("getMatches", &getMatches, (arg("query"), arg("maxResults") = -1))
+    .def("countMatches", &countMatches)
+    .def("hasMatch", &hasMatch)
+    .def("__len__", &withoutGil<&nvMolKit::SubstructLibrary::size>)
+    .add_property("pendingSize", &withoutGil<&nvMolKit::SubstructLibrary::pendingSize>)
+    .add_property("maxConcurrentQueries", &withoutGil<&nvMolKit::SubstructLibrary::maxConcurrentQueries>);
+}