From c06a9fd1f2a1a43e1118d1ccbded198432f4d574 Mon Sep 17 00:00:00 2001 From: Timur Rvachov Date: Sat, 18 Jul 2026 01:02:28 -0400 Subject: [PATCH] perf: accelerate fingerprint index queries with verified parity Run fingerprint cdist on the model-assigned GPU and accelerate invalid-SMILES Levenshtein lookup with RapidFuzz. Thread explicit multi-GPU device routing, canonicalize CUDA cache keys, restore legacy CPU topk tie identity, pin RapidFuzz in both manifests, and use workers=1. Add tests/test_fpindex.py for GPU/CPU parity, device routing, cache behavior, fallback paths, and the real index. Reviewed-and-tested: adversarially verified on the 211k fingerprint index with an A100, on CPU-only paths, and against both model checkpoints. --- env.yml | 1 + pyproject.toml | 1 + reasyn/chem/fpindex.py | 54 +++++++-- reasyn/sampler/sampler.py | 16 ++- reasyn/utils/sample_utils.py | 7 +- tests/test_fpindex.py | 224 +++++++++++++++++++++++++++++++++++ 6 files changed, 291 insertions(+), 12 deletions(-) create mode 100644 tests/test_fpindex.py diff --git a/env.yml b/env.yml index dc2cd04..a0fd531 100644 --- a/env.yml +++ b/env.yml @@ -35,5 +35,6 @@ dependencies: - rdkit>=2023.09 - scikit-learn==1.2.2 - editdistance + - rapidfuzz==3.14.5 - torchtyping - datasets diff --git a/pyproject.toml b/pyproject.toml index 68b41a5..21c0365 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ rdkit = "2025.3.2" numpy = "1.26.4" scikit-learn = "1.2.2" editdistance = "0.8.1" +rapidfuzz = "3.14.5" torchtyping = "0.1.5" [build-system] diff --git a/reasyn/chem/fpindex.py b/reasyn/chem/fpindex.py index 55c5e83..5f22443 100644 --- a/reasyn/chem/fpindex.py +++ b/reasyn/chem/fpindex.py @@ -131,22 +131,62 @@ def query(self, q: np.ndarray, k: int) -> list[list[_QueryResult]]: results.append(res) return results - @functools.cache def fp_cuda(self, device: torch.device) -> torch.Tensor: + device = torch.device(device) + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + return self._fp_cuda(device) + + @functools.cache + def _fp_cuda(self, device: torch.device) -> torch.Tensor: return torch.tensor(self._fp, dtype=torch.float, device=device) @torch.no_grad() - def query_cuda(self, q: torch.Tensor | str, k: int) -> list[list[_QueryResult]]: + def query_cuda( + self, + q: torch.Tensor | str, + k: int, + device: torch.device | None = None, + ) -> list[list[_QueryResult]]: if isinstance(q, str): # if mol is invalid, use edit distance instead bsz = 1 - pwdist = torch.Tensor([editdistance.eval(q, s) for s in self._smiles]) + # RapidFuzz computes identical Levenshtein distances in one C++/SIMD + # call. Falls back to editdistance if rapidfuzz is absent. + # REASYN_RAPIDFUZZ=0 restores the original loop for A/B measurement. + if os.environ.get('REASYN_RAPIDFUZZ', '1') != '0': + try: + from rapidfuzz import process as _rf_process + from rapidfuzz.distance import Levenshtein as _rf_lev + dists = _rf_process.cdist([q], self._smiles, + scorer=_rf_lev.distance, workers=1)[0] + pwdist = torch.Tensor(dists.astype('float32')) + except ImportError: + pwdist = torch.Tensor([editdistance.eval(q, s) for s in self._smiles]) + else: + pwdist = torch.Tensor([editdistance.eval(q, s) for s in self._smiles]) else: bsz = q.size(0) q = q.reshape([-1, self._fp_option.dim]) - pwdist = torch.cdist(self.fp_cuda(q.device), q, p=1) # (n_mols, n_queries) - dist_t, idx_t = torch.topk(pwdist, k=k, dim=0, largest=False) # (k, n_queries) - dist = dist_t.t().reshape([bsz, -1]).cpu().numpy() - idx = idx_t.t().reshape([bsz, -1]).cpu().numpy() + # Use the caller's model device; an existing CUDA query takes precedence. + # REASYN_GPU_QUERY=0 restores the original path for A/B measurement. + if os.environ.get('REASYN_GPU_QUERY', '1') != '0' and torch.cuda.is_available(): + if q.is_cuda: + dev = q.device + elif device is not None: + dev = device + else: + dev = torch.device("cuda", torch.cuda.current_device()) + dev = torch.device(dev) + if dev.type == "cuda" and dev.index is None: + dev = torch.device("cuda", torch.cuda.current_device()) + q = q.to(dev) + else: + dev = q.device + pwdist = torch.cdist(self.fp_cuda(dev), q, p=1) # (n_mols, n_queries) + # Keep legacy CPU topk tie-breaking while calculating distances on the GPU. + dist_t, idx_t = torch.topk(pwdist.cpu(), k=k, dim=0, largest=False) # (k, n_queries) + dist = dist_t.t().reshape([bsz, -1]).numpy() + idx = idx_t.t().reshape([bsz, -1]).numpy() results: list[list[_QueryResult]] = [] for i in range(dist.shape[0]): diff --git a/reasyn/sampler/sampler.py b/reasyn/sampler/sampler.py index 6617c9e..acc3a46 100644 --- a/reasyn/sampler/sampler.py +++ b/reasyn/sampler/sampler.py @@ -203,7 +203,13 @@ def _sample_token(tokens, enforce_td=False): token_sampled, _ = _sample_token(tokens) token_sampled_bb.append(token_sampled) smiles = decode_smiles(token_sampled_bb) - sampled_item = get_reactants(smiles, fpindex=self._fpindex, topk=100, use_edit_distance=use_edit_distance) + sampled_item = get_reactants( + smiles, + fpindex=self._fpindex, + topk=100, + use_edit_distance=use_edit_distance, + device=self.device, + ) if sampled_item is None: sampled_type = 'ABORTED' elif token_sampled >= TokenType.RXN_MIN: @@ -406,7 +412,13 @@ def _predict_editflow( success = state.stack.push_rxn(rxn, rxn_idx) if not success: break else: - sample = get_reactants(sample, fpindex=self._fpindex, topk=1, use_edit_distance=use_edit_distance) + sample = get_reactants( + sample, + fpindex=self._fpindex, + topk=1, + use_edit_distance=use_edit_distance, + device=self.device, + ) if sample is None: break sample = sample[0] mol, mol_idx = sample.reactant, sample.index diff --git a/reasyn/utils/sample_utils.py b/reasyn/utils/sample_utils.py index 7ebd7ea..f227317 100644 --- a/reasyn/utils/sample_utils.py +++ b/reasyn/utils/sample_utils.py @@ -86,17 +86,18 @@ def get_reactants( mol: str | Molecule, fpindex: FingerprintIndex, topk=1, - use_edit_distance: bool = False) -> list[list[_ReactantItem]]: + use_edit_distance: bool = False, + device: torch.device | None = None) -> list[list[_ReactantItem]]: if isinstance(mol, str): mol = Molecule(mol) if mol._rdmol is None: if not use_edit_distance: return None - query_res = fpindex.query_cuda(q=mol.smiles, k=topk)[0] + query_res = fpindex.query_cuda(q=mol.smiles, k=topk, device=device)[0] else: fp = torch.Tensor(mol.get_fingerprint(option=fpindex._fp_option)) - query_res = fpindex.query_cuda(q=fp[None, :], k=topk)[0] + query_res = fpindex.query_cuda(q=fp[None, :], k=topk, device=device)[0] mols = np.array([q.molecule for q in query_res]) mol_idxs = np.array([q.index for q in query_res]) distances = np.array([q.distance for q in query_res]) diff --git a/tests/test_fpindex.py b/tests/test_fpindex.py new file mode 100644 index 0000000..7e5fb0a --- /dev/null +++ b/tests/test_fpindex.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import builtins +import os +import pickle +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from reasyn.chem.fpindex import FingerprintIndex +from reasyn.chem.mol import FingerprintOption, Molecule +from reasyn.utils.sample_utils import get_reactants + + +def _index(fingerprints=None, smiles=None): + if fingerprints is None: + smiles = smiles or ["C", "CC", "CCC", "COC", "N#N", "C1CC1"] + fingerprints = np.zeros((len(smiles), 8), dtype=np.uint8) + elif smiles is None: + smiles = [str(index) for index in range(len(fingerprints))] + index = object.__new__(FingerprintIndex) + index._molecules = tuple(smiles) + index._smiles = smiles + index._fp = np.asarray(fingerprints, dtype=np.uint8) + index._fp_option = SimpleNamespace(dim=index._fp.shape[1]) + return index + + +def _pairs(results): + return [ + [(result.index, float(result.distance)) for result in row] + for row in results + ] + + +@pytest.fixture(autouse=True) +def _clear_fp_cache(): + FingerprintIndex._fp_cuda.cache_clear() + yield + FingerprintIndex._fp_cuda.cache_clear() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_gpu_query_matches_legacy_cpu_topk_with_ties(monkeypatch): + index = _index(np.zeros((100, 8), dtype=np.uint8)) + queries = torch.stack((torch.zeros((3, 8)), torch.ones((3, 8)))) + + monkeypatch.setenv("REASYN_GPU_QUERY", "0") + expected = _pairs(index.query_cuda(queries, k=3)) + monkeypatch.setenv("REASYN_GPU_QUERY", "1") + actual = _pairs( + index.query_cuda( + queries, + k=3, + device=torch.device("cuda", torch.cuda.current_device()), + ) + ) + + assert actual == expected + assert len(actual) == 2 + assert all(len(row) == 9 for row in actual) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_existing_cuda_query_keeps_its_device(monkeypatch): + index = _index(np.zeros((4, 8), dtype=np.uint8)) + query = torch.zeros((1, 8), device="cuda") + seen = [] + fp_cuda = index.fp_cuda + + def record_device(device): + seen.append(torch.device(device)) + return fp_cuda(device) + + monkeypatch.setattr(index, "fp_cuda", record_device) + monkeypatch.setenv("REASYN_GPU_QUERY", "1") + index.query_cuda(query, k=1, device=torch.device("cpu")) + + assert seen == [query.device] + + +def test_query_uses_explicit_indexed_device(monkeypatch): + index = _index(np.zeros((4, 8), dtype=np.uint8)) + query = torch.zeros((1, 8)) + moves = [] + fp_devices = [] + + def fake_to(tensor, device): + moves.append(torch.device(device)) + return tensor + + def fake_fp_cuda(device): + fp_devices.append(torch.device(device)) + return torch.tensor(index._fp, dtype=torch.float) + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.Tensor, "to", fake_to) + monkeypatch.setattr(index, "fp_cuda", fake_fp_cuda) + monkeypatch.setenv("REASYN_GPU_QUERY", "1") + index.query_cuda(query, k=1, device=torch.device("cuda:7")) + + assert moves == [torch.device("cuda:7")] + assert fp_devices == [torch.device("cuda:7")] + + +def test_fp_cache_canonicalizes_bare_cuda(monkeypatch): + index = _index() + allocations = [] + sentinel = object() + + def fake_tensor(*args, **kwargs): + allocations.append(kwargs["device"]) + return sentinel + + monkeypatch.setattr(torch.cuda, "current_device", lambda: 7) + monkeypatch.setattr(torch, "tensor", fake_tensor) + + assert index.fp_cuda(torch.device("cuda")) is sentinel + assert index.fp_cuda(torch.device("cuda:7")) is sentinel + assert allocations == [torch.device("cuda:7")] + assert FingerprintIndex._fp_cuda.cache_info().hits == 1 + + +def test_cpu_only_host_ignores_requested_cuda_device(monkeypatch): + index = _index(np.zeros((4, 8), dtype=np.uint8)) + query = torch.ones((1, 8)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setenv("REASYN_GPU_QUERY", "1") + + actual = _pairs(index.query_cuda(query, k=2, device=torch.device("cuda:7"))) + monkeypatch.setenv("REASYN_GPU_QUERY", "0") + expected = _pairs(index.query_cuda(query, k=2)) + + assert actual == expected + + +def test_get_reactants_forwards_model_device(): + class SpyIndex: + _fp_option = FingerprintOption(morgan_n_bits=8) + + def query_cuda(self, q, k, device=None): + self.device = device + result = SimpleNamespace(molecule=Molecule("C"), index=0, distance=0.0) + return [[result]] + + index = SpyIndex() + device = torch.device("cuda:7") + get_reactants("C", fpindex=index, device=device) + + assert index.device == device + + +def test_rapidfuzz_matches_editdistance_and_uses_one_worker(monkeypatch): + from rapidfuzz import process as rf_process + + index = _index() + real_cdist = rf_process.cdist + workers = [] + + def record_workers(*args, **kwargs): + workers.append(kwargs["workers"]) + return real_cdist(*args, **kwargs) + + monkeypatch.setattr(rf_process, "cdist", record_workers) + for query in ("", "not-a-smiles(", "C", "😀"): + monkeypatch.setenv("REASYN_RAPIDFUZZ", "0") + expected = _pairs(index.query_cuda(query, k=len(index._smiles))) + monkeypatch.setenv("REASYN_RAPIDFUZZ", "1") + actual = _pairs(index.query_cuda(query, k=len(index._smiles))) + assert actual == expected + + assert workers == [1, 1, 1, 1] + + +def test_missing_rapidfuzz_falls_back_to_editdistance(monkeypatch): + index = _index() + query = "not-a-smiles(" + monkeypatch.setenv("REASYN_RAPIDFUZZ", "0") + expected = _pairs(index.query_cuda(query, k=len(index._smiles))) + real_import = builtins.__import__ + attempts = [] + + def without_rapidfuzz(name, *args, **kwargs): + if name == "rapidfuzz" or name.startswith("rapidfuzz."): + attempts.append(name) + raise ImportError("simulated missing RapidFuzz") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", without_rapidfuzz) + monkeypatch.setenv("REASYN_RAPIDFUZZ", "1") + actual = _pairs(index.query_cuda(query, k=len(index._smiles))) + + assert actual == expected + assert attempts + + +def test_real_index_gpu_query_matches_legacy_cpu_topk(monkeypatch): + path = os.environ.get("REASYN_REAL_FPINDEX") + if not path: + pytest.skip("set REASYN_REAL_FPINDEX to exercise the real index") + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + with Path(path).open("rb") as handle: + index = pickle.load(handle) + query_ids = [0, len(index._fp) // 4, len(index._fp) // 2, len(index._fp) - 1] + queries = torch.from_numpy(index._fp[query_ids].astype(np.float32, copy=True)) + + monkeypatch.setenv("REASYN_GPU_QUERY", "0") + expected = _pairs(index.query_cuda(queries, k=100)) + monkeypatch.setenv("REASYN_GPU_QUERY", "1") + actual = _pairs( + index.query_cuda( + queries, + k=100, + device=torch.device("cuda", torch.cuda.current_device()), + ) + ) + + assert actual == expected