From 7336cbc92c0efe5ff81af23c4c48fcf0c3ee3ca9 Mon Sep 17 00:00:00 2001 From: Clay Moore Date: Fri, 24 Jul 2026 20:36:07 -0500 Subject: [PATCH 1/2] fix(determinism): make --seed actually reproducible docs/prediction.md documents --seed as "Random seed for reproducible predictions", but two model inputs were drawn from RNGs that pl.seed_everything does not reach, so the same input could produce different atomistic features depending on --num_workers. 1. Ligand conformers. get_conformer embedded with ETKDGv3 and never set randomSeed (RDKit's default is -1, i.e. random). RDKit's global RNG starts from a fixed state per process, so a single-record run was already reproducible, but it advances with each embedding: in a batch the Nth ligand embedded inside a worker depends on how many preceded it, and therefore on how work is split across preprocess workers. Measured over 6 ligands, comparing --num_workers 1 vs 2 and 1 vs 4: 10 of 12 record/worker-count comparisons produced different geometry (up to 22.2 A); after the fix, 0 of 12. conformer_seed() now derives the seed from the base seed and the molecule's canonical SMILES, deliberately not from processing order. CCD and SDF ligands are unaffected: they already carry a conformer, so embedding never runs (verified: identical coordinates across seeds). 2. ref_pos augmentation. center_random_augmentation drew from the global torch RNG, which seed_everything(workers=True) seeds per DataLoader worker. With the conformer pinned so this was isolated, 0 of 6 records matched between --num_workers 0 and 1/2/4 (up to 14.5 A); after the fix, 6 of 6 match. The generator is now seeded from the existing per-record RandomState, so the augmentation distribution is unchanged. --seed is threaded to conformer generation so the flag stays meaningful, and base_seed=None keeps RDKit's previous non-deterministic behaviour. Also fixes `--num_workers 0`, which crashed with "max_workers must be greater than 0" because the value was reused for the preprocessing ProcessPoolExecutor. Note this is not output-preserving and cannot be: previous behaviour was random. The claim is identical output across runs and across --num_workers, not identical to any particular earlier run. Adds tests/test_determinism.py (6 tests). The two that rely only on existing APIs were confirmed to fail on main and pass here. --- nesso/data/featurizer.py | 12 ++- nesso/data/yaml_input.py | 95 ++++++++++++++--- nesso/main.py | 17 ++- nesso/model/modules/utils.py | 31 ++++-- tests/test_determinism.py | 193 +++++++++++++++++++++++++++++++++++ 5 files changed, 318 insertions(+), 30 deletions(-) create mode 100644 tests/test_determinism.py diff --git a/nesso/data/featurizer.py b/nesso/data/featurizer.py index 742304a..5efb2ba 100644 --- a/nesso/data/featurizer.py +++ b/nesso/data/featurizer.py @@ -469,12 +469,20 @@ def process_atom_features( center = center / resolved_mask.sum().clamp(min=1) coords = coords - center[:, None] - # Apply random roto-translation to the input conformers + # Apply random roto-translation to the input conformers. + # The generator is seeded from `random` (a per-record RandomState) rather than + # drawn from the global torch RNG, which is seeded per DataLoader worker and so + # would make `ref_pos` depend on `--num_workers`. + aug_generator = torch.Generator() + aug_generator.manual_seed(int(random.randint(0, 2**31 - 1))) for i in range(torch.max(ref_space_uid) + 1): included = ref_space_uid == i if torch.sum(included) > 0 and torch.any(resolved_mask[included]): ref_pos[included] = center_random_augmentation( - ref_pos[included][None], resolved_mask[included][None], centering=True + ref_pos[included][None], + resolved_mask[included][None], + centering=True, + generator=aug_generator, )[0] num_token_classes = max_tokens if max_tokens is not None else L diff --git a/nesso/data/yaml_input.py b/nesso/data/yaml_input.py index de9fe51..57989aa 100644 --- a/nesso/data/yaml_input.py +++ b/nesso/data/yaml_input.py @@ -107,10 +107,42 @@ def _load_mol( return pickle.load(f) # noqa: S301 -def get_conformer(mol: Chem.Mol) -> Chem.Conformer: +DEFAULT_CONFORMER_SEED = 42 + +# RDKit takes a signed 32-bit seed; -1 means "pick a random one". +_MAX_RDKIT_SEED = 2**31 - 1 + + +def conformer_seed(mol: Chem.Mol, base_seed: int) -> int: + """Derive a stable ETKDG seed for ``mol``. + + The seed is a function of ``base_seed`` and the molecule's canonical SMILES, + deliberately **not** of processing order: ``preprocess_yamls`` parses inputs in a + ``ProcessPoolExecutor``, so a counter-based scheme would vary with worker + scheduling and with ``--num_workers``. Keying on the molecule also means the same + ligand embeds identically wherever it appears. + """ + try: + key = Chem.MolToSmiles(mol) + except Exception: # unsanitized/exotic mol: fall back to a structural key + key = f"{mol.GetNumAtoms()}:{mol.GetNumBonds()}" + digest = hashlib.sha256(f"{base_seed}:{key}".encode("utf-8")).digest() + return int.from_bytes(digest[:4], "big") % _MAX_RDKIT_SEED + + +def get_conformer( + mol: Chem.Mol, base_seed: int | None = DEFAULT_CONFORMER_SEED +) -> Chem.Conformer: + """Return conformer 0, embedding one with ETKDG if the molecule has none. + + ``base_seed`` makes embedding reproducible; pass ``None`` for RDKit's default + non-deterministic behaviour. + """ if mol.GetNumConformers() == 0: opts = AllChem.ETKDGv3() opts.clearConfs = False + if base_seed is not None: + opts.randomSeed = conformer_seed(mol, base_seed) cid = AllChem.EmbedMolecule(mol, opts) if cid < 0: opts.useRandomCoords = True @@ -124,6 +156,7 @@ def _standard_residue( res_idx: int, *, ccd_dict: dict[str, Chem.Mol] | None = None, + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> dict[str, Any]: name = "MET" if res_name == "MSE" else res_name if name not in const.ref_atoms or not const.ref_atoms[name]: @@ -131,7 +164,7 @@ def _standard_residue( raise ValueError(msg) mol = _load_mol(mol_dir, name, ccd_dict=ccd_dict) mol = Chem.RemoveHs(mol, sanitize=False) - conf = get_conformer(mol) + conf = get_conformer(mol, base_seed) by_name = {a.GetProp("name"): a for a in mol.GetAtoms()} atoms: list[dict[str, Any]] = [] for atom_name in const.ref_atoms[name]: @@ -162,21 +195,27 @@ def _protein_residues( mol_dir: Path | None, *, ccd_dict: dict[str, Chem.Mol] | None = None, + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> tuple[int, list[dict[str, Any]]]: cmap = const.prot_letter_to_token unk = const.unk_token["PROTEIN"] mol_type = const.chain_type_ids["PROTEIN"] residues = [ - _standard_residue(cmap.get(c, unk), mol_dir, j, ccd_dict=ccd_dict) + _standard_residue( + cmap.get(c, unk), mol_dir, j, ccd_dict=ccd_dict, base_seed=base_seed + ) for j, c in enumerate(raw_seq) ] return mol_type, residues def _ligand_residue_from_mol( - mol: Chem.Mol, res_name: str, res_idx: int + mol: Chem.Mol, + res_name: str, + res_idx: int, + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> dict[str, Any]: - conf = get_conformer(mol) + conf = get_conformer(mol, base_seed) idx_map: dict[int, int] = {} atoms: list[dict[str, Any]] = [] for atom in mol.GetAtoms(): @@ -222,6 +261,7 @@ def _ligand_from_conformer_pkl( mol_dir: Path | None = None, rid: str = "", record_id: str = "", + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> tuple[int, list[dict[str, Any]]]: """Load an RDKit ``Mol`` (heavy-atom, with conformer) from ``path`` as a ligand residue.""" with Path(path).open("rb") as f: @@ -236,7 +276,7 @@ def _ligand_from_conformer_pkl( atom.SetProp("name", f"{sym}{int(rnk) + 1}"[:4]) _dump_ligand_mol(mol, mol_dir, rid, record_id=record_id) return const.chain_type_ids["NONPOLYMER"], [ - _ligand_residue_from_mol(mol, lig_tag, 0) + _ligand_residue_from_mol(mol, lig_tag, 0, base_seed) ] @@ -263,6 +303,7 @@ def _ligand_smiles_residue( rid: str = "", *, record_id: str = "", + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> tuple[int, list[dict[str, Any]]]: mol = Chem.MolFromSmiles(smi) if mol is None: @@ -270,12 +311,12 @@ def _ligand_smiles_residue( raise ValueError(msg) mol = Chem.AddHs(mol) Chem.AssignStereochemistry(mol, force=True, cleanIt=True) - get_conformer(mol) + get_conformer(mol, base_seed) mol_nh = Chem.RemoveHs(mol, sanitize=False) _assign_ligand_atom_names(mol_nh) _dump_ligand_mol(mol_nh, mol_dir, rid, record_id=record_id) return const.chain_type_ids["NONPOLYMER"], [ - _ligand_residue_from_mol(mol_nh, rid or "LIG", 0) + _ligand_residue_from_mol(mol_nh, rid or "LIG", 0, base_seed) ] @@ -286,14 +327,15 @@ def _ligand_ccd_residue( ccd_dict: dict[str, Chem.Mol] | None = None, rid: str = "", record_id: str = "", + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> tuple[int, list[dict[str, Any]]]: mol = _load_mol(mol_dir, ccd_code, ccd_dict=ccd_dict) mol = Chem.RemoveHs(mol, sanitize=False) _assign_ligand_atom_names(mol) - get_conformer(mol) + get_conformer(mol, base_seed) _dump_ligand_mol(mol, mol_dir, rid, record_id=record_id) return const.chain_type_ids["NONPOLYMER"], [ - _ligand_residue_from_mol(mol, ccd_code, 0) + _ligand_residue_from_mol(mol, ccd_code, 0, base_seed) ] @@ -303,6 +345,7 @@ def _ligand_sdf_residue( rid: str = "", *, record_id: str = "", + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> tuple[int, list[dict[str, Any]]]: path = Path(sdf_path) if not path.exists(): @@ -313,10 +356,10 @@ def _ligand_sdf_residue( msg = f"Invalid SDF: {path}" raise ValueError(msg) _assign_ligand_atom_names(mol) - get_conformer(mol) + get_conformer(mol, base_seed) _dump_ligand_mol(mol, mol_dir, rid, record_id=record_id) return const.chain_type_ids["NONPOLYMER"], [ - _ligand_residue_from_mol(mol, rid or "LIG", 0) + _ligand_residue_from_mol(mol, rid or "LIG", 0, base_seed) ] @@ -344,10 +387,11 @@ def _chain_data_for_entity( ccd_dict: dict[str, Chem.Mol] | None = None, rid: str = "", record_id: str = "", + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> _ChainData: if isinstance(entity, _EntityProtein): mol_type, residues = _protein_residues( - entity.sequence, mol_dir, ccd_dict=ccd_dict + entity.sequence, mol_dir, ccd_dict=ccd_dict, base_seed=base_seed ) elif isinstance(entity, _EntityLigandSmiles): smiles_ligand_idx[0] += 1 @@ -358,10 +402,15 @@ def _chain_data_for_entity( mol_dir=mol_dir, rid=rid, record_id=record_id, + base_seed=base_seed, ) else: mol_type, residues = _ligand_smiles_residue( - entity.smiles, mol_dir=mol_dir, rid=rid, record_id=record_id + entity.smiles, + mol_dir=mol_dir, + rid=rid, + record_id=record_id, + base_seed=base_seed, ) elif isinstance(entity, _EntityLigandCCD): mol_type, residues = _ligand_ccd_residue( @@ -370,10 +419,15 @@ def _chain_data_for_entity( ccd_dict=ccd_dict, rid=rid, record_id=record_id, + base_seed=base_seed, ) elif isinstance(entity, _EntityLigandSDF): mol_type, residues = _ligand_sdf_residue( - entity.sdf, mol_dir=mol_dir, rid=rid, record_id=record_id + entity.sdf, + mol_dir=mol_dir, + rid=rid, + record_id=record_id, + base_seed=base_seed, ) else: raise TypeError(entity) @@ -592,9 +646,13 @@ def parse_schema( *, ccd_dict: dict[str, Chem.Mol] | None = None, record_id: str = "record", + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> tuple[Structure, Record, dict[int, str], dict[int, str]]: """Parse a YAML schema dict into ``(Structure, Record, entity_to_seq, entity_to_esm_path)`` + + ``base_seed`` seeds ETKDG conformer embedding so parsing is reproducible; + pass ``None`` for RDKit's default non-deterministic behaviour. """ if not isinstance(schema, dict) or schema.get("version", 1) != 1: raise ValueError("Schema must be a dict with version: 1") @@ -630,6 +688,7 @@ def parse_schema( ccd_dict=ccd_dict, rid=lid, record_id=record_id, + base_seed=base_seed, ) data.entity_id = entity_id if isinstance(entity, _EntityProtein): @@ -671,11 +730,13 @@ def parse_yaml( ccd_pkl: Path | None = None, ccd_dict: dict[str, Chem.Mol] | None = None, record_id: str | None = None, + base_seed: int | None = DEFAULT_CONFORMER_SEED, ) -> tuple[Structure, Record, dict[int, str], dict[int, str]]: """Parse YAML into ``(Structure, Record, entity_to_seq, entity_to_esm_path)``. ``ccd_dict`` (in-memory) takes precedence over ``ccd_pkl`` (disk); ``_load_mol`` raises if a standard residue needs a CCD source that was not provided. + ``base_seed`` seeds ETKDG conformer embedding so parsing is reproducible. """ if ccd_dict is None and ccd_pkl is not None: ccd_dict = load_ccd_mol_dict(ccd_pkl) @@ -684,7 +745,9 @@ def parse_yaml( schema = yaml.safe_load(f) rid = record_id if record_id is not None else path.stem - return parse_schema(schema, mol_dir, ccd_dict=ccd_dict, record_id=rid) + return parse_schema( + schema, mol_dir, ccd_dict=ccd_dict, record_id=rid, base_seed=base_seed + ) def esm_keys(entity_to_seq: dict[int, str]) -> dict[str, str]: diff --git a/nesso/main.py b/nesso/main.py index c29a46e..9a7827b 100644 --- a/nesso/main.py +++ b/nesso/main.py @@ -24,7 +24,11 @@ from nesso.data.inference import NessoInferenceDataModule from nesso.data.types import Manifest, Record from nesso.data.writer import NessoWriter -from nesso.data.yaml_input import parse_yaml, validate_schema +from nesso.data.yaml_input import ( + DEFAULT_CONFORMER_SEED, + parse_yaml, + validate_schema, +) from nesso.model.models.nesso1 import Nesso1 from nesso.data.esm import ( @@ -164,8 +168,11 @@ def _process_single_yaml( mol_dir: Path, structures_dir: Path, records_dir: Path, + seed: int, ) -> Record: - struct, rec, _, _ = parse_yaml(yp, mol_dir, ccd_dict=_worker_ccd_dict) + struct, rec, _, _ = parse_yaml( + yp, mol_dir, ccd_dict=_worker_ccd_dict, base_seed=seed + ) struct.dump(structures_dir / f"{rec.id}.npz") rec.dump(records_dir / f"{rec.id}.json") return rec @@ -178,6 +185,7 @@ def preprocess_yamls( structures_dir: Path, records_dir: Path, num_workers: int = 2, + seed: int = DEFAULT_CONFORMER_SEED, ) -> tuple[Manifest, list[str]]: """Parse YAMLs into a Manifest, reporting which inputs failed. @@ -192,11 +200,11 @@ def preprocess_yamls( failed: list[str] = [] with ProcessPoolExecutor( - max_workers=num_workers, initializer=_init_worker, initargs=(ccd_pkl,) + max_workers=max(1, num_workers), initializer=_init_worker, initargs=(ccd_pkl,) ) as executor: futures = { executor.submit( - _process_single_yaml, yp, mol_dir, structures_dir, records_dir + _process_single_yaml, yp, mol_dir, structures_dir, records_dir, seed ): yp for yp in yaml_paths } @@ -547,6 +555,7 @@ def predict( paths.structures_dir, paths.records_dir, num_workers=num_workers, + seed=seed, ) manifest.dump(paths.manifest_path) if failed_preprocessing: diff --git a/nesso/model/modules/utils.py b/nesso/model/modules/utils.py index 5a9db9a..fd02e09 100644 --- a/nesso/model/modules/utils.py +++ b/nesso/model/modules/utils.py @@ -31,8 +31,10 @@ def forward( return F.silu(gates) * x -def randomly_rotate(coords, return_second_coords=False, second_coords=None): - R = random_rotations(len(coords), coords.dtype, coords.device) +def randomly_rotate( + coords, return_second_coords=False, second_coords=None, generator=None +): + R = random_rotations(len(coords), coords.dtype, coords.device, generator=generator) if return_second_coords: return torch.einsum("bmd,bds->bms", coords, R), torch.einsum( @@ -50,6 +52,7 @@ def center_random_augmentation( centering=True, return_second_coords=False, second_coords=None, + generator=None, ): """Algorithm 19""" if centering: @@ -64,9 +67,15 @@ def center_random_augmentation( if augmentation: atom_coords, second_coords = randomly_rotate( - atom_coords, return_second_coords=True, second_coords=second_coords + atom_coords, + return_second_coords=True, + second_coords=second_coords, + generator=generator, ) - random_trans = torch.randn_like(atom_coords[:, 0:1, :]) * s_trans + # randn_like() takes no generator; empty_like().normal_() is equivalent. + trans_shape = atom_coords[:, 0:1, :] + random_trans = torch.empty_like(trans_shape).normal_(generator=generator) + random_trans = random_trans * s_trans atom_coords = atom_coords + random_trans if second_coords is not None: @@ -132,7 +141,10 @@ def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor: def random_quaternions( - n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None + n: int, + dtype: Optional[torch.dtype] = None, + device: Optional[Device] = None, + generator: Optional[torch.Generator] = None, ) -> torch.Tensor: """ Generate random quaternions representing rotations, @@ -149,14 +161,17 @@ def random_quaternions( """ if isinstance(device, str): device = torch.device(device) - o = torch.randn((n, 4), dtype=dtype, device=device) + o = torch.randn((n, 4), dtype=dtype, device=device, generator=generator) s = (o * o).sum(1) o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None] return o def random_rotations( - n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None + n: int, + dtype: Optional[torch.dtype] = None, + device: Optional[Device] = None, + generator: Optional[torch.Generator] = None, ) -> torch.Tensor: """ Generate random rotations as 3x3 rotation matrices. @@ -170,5 +185,5 @@ def random_rotations( Returns: Rotation matrices as tensor of shape (n, 3, 3). """ - quaternions = random_quaternions(n, dtype=dtype, device=device) + quaternions = random_quaternions(n, dtype=dtype, device=device, generator=generator) return quaternion_to_matrix(quaternions) diff --git a/tests/test_determinism.py b/tests/test_determinism.py new file mode 100644 index 0000000..b6f6f96 --- /dev/null +++ b/tests/test_determinism.py @@ -0,0 +1,193 @@ +"""Predictions must be reproducible for a fixed seed. + +Two inputs to the model were previously drawn from unseeded RNGs, so the same +YAML gave different atomistic features on every run even though +``docs/prediction.md`` documents ``--seed`` as giving "reproducible predictions": + +* the ligand 3D conformer, embedded by RDKit ETKDG with no ``randomSeed`` (its + default is ``-1``, i.e. random); +* the per-residue roto-translation of ``ref_pos``, drawn from the *global* torch + RNG, which ``seed_everything(..., workers=True)`` seeds per DataLoader worker, + making the result depend on ``--num_workers``. + +These tests pin both. They are ligand-only so they need no CCD asset and run on +plain CI. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch +from rdkit import Chem +from torch.utils.data import DataLoader + +from nesso.data.featurizer import NessoFeaturizer +from nesso.data.inference import InferenceDataset, inference_collate +from nesso.data.types import Manifest +from nesso.data.yaml_input import ( + DEFAULT_CONFORMER_SEED, + conformer_seed, + get_conformer, + parse_yaml, +) + +# Flexible, drug-like: many rotatable bonds, so ETKDG genuinely varies. +_SMILES = "Cc1ccc(NC(=O)c2ccc(CN3CCN(C)CC3)cc2)cc1Nc1nccc(-c2cccnc2)n1" +_YAML = f'version: 1\nsequences:\n - ligand:\n id: B\n smiles: "{_SMILES}"\n' + + +def _parse_coords(tmp_path: Path, name: str, seed: int | None) -> np.ndarray: + work = tmp_path / name + mol_dir = work / "rdkit_conformers" + mol_dir.mkdir(parents=True, exist_ok=True) + yaml_path = work / "lig.yaml" + yaml_path.write_text(_YAML) + struct, _, _, _ = parse_yaml( + yaml_path, mol_dir, ccd_dict=None, record_id="lig", base_seed=seed + ) + return struct.coords["coords"].copy() + + +def _internal_distances(x: np.ndarray) -> np.ndarray: + """Rotation- and translation-invariant description of a conformer.""" + return np.linalg.norm(x[:, None, :] - x[None, :, :], axis=-1) + + +def test_same_seed_gives_identical_conformer(tmp_path: Path) -> None: + a = _parse_coords(tmp_path, "a", DEFAULT_CONFORMER_SEED) + b = _parse_coords(tmp_path, "b", DEFAULT_CONFORMER_SEED) + assert np.array_equal(a, b) + # Invariant under pose, so this also rules out "same shape, different frame". + assert np.array_equal(_internal_distances(a), _internal_distances(b)) + + +def test_different_seed_gives_different_conformer(tmp_path: Path) -> None: + """Guards against the seed being accepted but ignored.""" + a = _parse_coords(tmp_path, "a", DEFAULT_CONFORMER_SEED) + c = _parse_coords(tmp_path, "c", DEFAULT_CONFORMER_SEED + 1) + assert not np.array_equal(a, c) + + +def test_conformer_seed_depends_on_molecule_not_order() -> None: + """`preprocess_yamls` embeds in a ProcessPoolExecutor, so the seed must not + depend on processing order.""" + mol_a = Chem.AddHs(Chem.MolFromSmiles(_SMILES)) + mol_b = Chem.AddHs(Chem.MolFromSmiles(_SMILES)) + other = Chem.AddHs(Chem.MolFromSmiles("CC(=O)Oc1ccccc1C(=O)O")) + + assert conformer_seed(mol_a, 42) == conformer_seed(mol_b, 42) + assert conformer_seed(mol_a, 42) != conformer_seed(other, 42) + assert conformer_seed(mol_a, 42) != conformer_seed(mol_a, 43) + # RDKit takes a signed 32-bit seed. + assert 0 <= conformer_seed(mol_a, 42) < 2**31 - 1 + + +def test_conformer_seed_can_be_opted_out() -> None: + """`base_seed=None` keeps RDKit's default (non-deterministic) behaviour.""" + mol = Chem.AddHs(Chem.MolFromSmiles(_SMILES)) + get_conformer(mol, None) + assert mol.GetNumConformers() == 1 + + +def _ref_pos_by_record(tmp_path: Path, num_workers: int) -> dict[str, np.ndarray]: + """Featurize several records through a real DataLoader at a given worker count.""" + mol_dir = tmp_path / "rdkit_conformers" + struct_dir = tmp_path / "structures" + esm_dir = tmp_path / "esm" + for directory in (mol_dir, struct_dir, esm_dir): + directory.mkdir(parents=True, exist_ok=True) + + records = [] + for i in range(3): + yaml_path = tmp_path / f"lig{i}.yaml" + yaml_path.write_text(_YAML) + struct, record, _, _ = parse_yaml( + yaml_path, mol_dir, ccd_dict=None, record_id=f"lig{i}" + ) + struct.dump(struct_dir / f"{record.id}.npz") + records.append(record) + + dataset = InferenceDataset( + manifest=Manifest(records), + target_dir=tmp_path, + featurizer=NessoFeaturizer( + esm_emb_dir=esm_dir, esm_emb_dim=1280, esm_num_layers=33 + ), + ligand_dir=mol_dir, + ccd_pkl=None, + ) + loader = DataLoader( + dataset, + batch_size=1, + shuffle=False, + num_workers=num_workers, + collate_fn=inference_collate, + ) + return {b["record"][0].id: b["ref_pos"][0].numpy().copy() for b in loader} + + +def test_ref_pos_independent_of_num_workers(tmp_path: Path) -> None: + """ref_pos augmentation must not read the global (per-worker) torch RNG.""" + torch.manual_seed(0) + single = _ref_pos_by_record(tmp_path / "w0", 0) + torch.manual_seed(0) + multi = _ref_pos_by_record(tmp_path / "w2", 2) + + assert set(single) == set(multi) + for key in single: + assert np.array_equal(single[key], multi[key]), key + + +def test_ref_pos_unaffected_by_global_rng_state(tmp_path: Path) -> None: + """Consuming the global RNG beforehand must not change featurization.""" + torch.manual_seed(0) + baseline = _ref_pos_by_record(tmp_path / "base", 0) + torch.manual_seed(0) + _ = torch.randn(17) # perturb global RNG stream position + shifted = _ref_pos_by_record(tmp_path / "shift", 0) + + for key in baseline: + assert np.array_equal(baseline[key], shifted[key]), key + + +_OTHER_SMILES = [ + "CC(=O)Oc1ccccc1C(=O)O", + "CN1C=NC2=C1C(=O)N(C)C(=O)N2C", + "CC(C)Cc1ccc(cc1)C(C)C(=O)O", +] + + +def _screen(tmp_path: Path, name: str, smiles: list[str]) -> dict[str, np.ndarray]: + """Parse a batch of ligands, returning coordinates keyed by SMILES.""" + work = tmp_path / name + mol_dir = work / "rdkit_conformers" + mol_dir.mkdir(parents=True, exist_ok=True) + out: dict[str, np.ndarray] = {} + for i, smi in enumerate(smiles): + yaml_path = work / f"c{i}.yaml" + yaml_path.write_text( + f'version: 1\nsequences:\n - ligand:\n id: B\n smiles: "{smi}"\n' + ) + struct, _, _, _ = parse_yaml( + yaml_path, mol_dir, ccd_dict=None, record_id=f"c{i}" + ) + out[smi] = struct.coords["coords"].copy() + return out + + +def test_conformer_independent_of_batch_composition_and_order(tmp_path: Path) -> None: + """A ligand's geometry must depend only on that ligand. + + Previously RDKit's unseeded RNG advanced with every embedding, so a compound's + conformer depended on what else was in the batch and on the order it was + processed in. That silently changed already-computed results whenever a + screening library was extended or its inputs reordered. + """ + alone = _screen(tmp_path, "alone", [_SMILES]) + with_others = _screen(tmp_path, "with_others", [*_OTHER_SMILES, _SMILES]) + reordered = _screen(tmp_path, "reordered", [_SMILES, *reversed(_OTHER_SMILES)]) + + assert np.array_equal(alone[_SMILES], with_others[_SMILES]) + assert np.array_equal(alone[_SMILES], reordered[_SMILES]) From eaaf9c47cc79d0165441c65e5571a6a343d29090 Mon Sep 17 00:00:00 2001 From: Clay Moore Date: Tue, 4 Aug 2026 19:54:14 -0500 Subject: [PATCH 2/2] fix(determinism): key the per-record RNG on identity, not batch position Seeding the ref_pos augmentation was not sufficient on its own. Two further layers made a record's features depend on where it sat in the batch: 1. preprocess_yamls collected results with as_completed, which yields in completion order, so the manifest order followed worker scheduling. Measured over 40 records: order was stable at --num_workers 1 but differed between runs at 2 and 4. Results are now collected by submission index. 2. InferenceDataset seeded RandomState from the dataset index, so that scrambled order reached the augmentation. Even with order preserved, the index still shifts whenever inputs are added, removed, renamed or reordered. record_rng_seed() now derives the seed from the record id via sha256, matching how conformer_seed keys on the molecule. hash() is unusable here because it is salted per process. Diagnosis: hashing every feature tensor reaching predict_step across two --num_workers 4 CLI runs showed ref_pos differing in 39 of 40 records with all 26 other features identical, which pointed at the per-record RNG rather than the transfer path. An earlier suspicion that pin_memory plus non_blocking transfers were responsible was tested directly and disproved. End to end over 40 ligands against one target, same library run twice: comparison before after nw=1 twice 0/40 0/40 nw=4 twice 17/40 0/40 nw=1 vs nw=4 39/40 0/40 nw=1 vs nw=2 n/a 0/40 All max deltas are exactly 0.0000 and all Spearman correlations exactly 1.000000, so this is bit-identical rather than close. Adds two tests. test_ref_pos_independent_of_record_position featurizes the same records in forward and reversed manifest order and was confirmed to fail without the fix. --- nesso/data/inference.py | 18 ++++++++++++- nesso/main.py | 16 ++++++++---- tests/test_determinism.py | 53 +++++++++++++++++++++++++++++++++++---- 3 files changed, 76 insertions(+), 11 deletions(-) diff --git a/nesso/data/inference.py b/nesso/data/inference.py index 4a5e2ab..26c7626 100644 --- a/nesso/data/inference.py +++ b/nesso/data/inference.py @@ -3,6 +3,7 @@ Adapted from https://github.com/jwohlwend/boltz, MIT License, Copyright (c) 2024 Jeremy Wohlwend. """ +import hashlib import pickle from dataclasses import replace from pathlib import Path @@ -109,6 +110,21 @@ def collate( INFERENCE_EXCLUDED_KEYS = frozenset({"record", "exception", "record_id"}) +# numpy seeds must fit in uint32. +_MAX_NUMPY_SEED = 2**32 + + +def record_rng_seed(record_id: str) -> int: + """Per-record RNG seed derived from the record id. + + Keyed on identity rather than on the dataset index, so a record's features do + not change when the surrounding batch does. The index shifts whenever inputs + are added, removed, renamed, or reordered, and ``hash()`` is salted per + process, so neither is usable here. + """ + digest = hashlib.sha256(record_id.encode("utf-8")).digest() + return int.from_bytes(digest[:4], "big") % _MAX_NUMPY_SEED + def inference_collate(data: list[dict[str, Tensor]]) -> dict[str, Tensor]: # Must be module-level (picklable) for DataLoader multiprocessing on macOS/Windows. @@ -199,7 +215,7 @@ def __getitem__(self, idx: int) -> dict[str, Any]: structure=structure, record=record, ) - random = RandomState(idx) + random = RandomState(record_rng_seed(str(record_id))) molecules = self._setup_molecules(structure, str(record_id)) features = self.featurizer.process( tokenized, diff --git a/nesso/main.py b/nesso/main.py index 9a7827b..55af005 100644 --- a/nesso/main.py +++ b/nesso/main.py @@ -196,26 +196,32 @@ def preprocess_yamls( """ structures_dir.mkdir(parents=True, exist_ok=True) records_dir.mkdir(parents=True, exist_ok=True) - records: list[Record] = [] failed: list[str] = [] + # Results are collected by submission index, not completion order: + # `as_completed` yields whichever worker finishes first, so appending here + # would make the manifest order depend on scheduling and hence on + # `--num_workers`. Downstream that order is visible as the dataset index. + by_index: dict[int, Record] = {} + with ProcessPoolExecutor( max_workers=max(1, num_workers), initializer=_init_worker, initargs=(ccd_pkl,) ) as executor: futures = { executor.submit( _process_single_yaml, yp, mol_dir, structures_dir, records_dir, seed - ): yp - for yp in yaml_paths + ): (i, yp) + for i, yp in enumerate(yaml_paths) } for future in tqdm(as_completed(futures), total=len(futures), desc="YAML"): - yp = futures[future] + i, yp = futures[future] try: - records.append(future.result()) + by_index[i] = future.result() except Exception as e: failed.append(yp.stem) print(f"Error processing YAML {yp.name}: {e}", file=sys.stderr) + records = [by_index[i] for i in sorted(by_index)] return Manifest(records), failed diff --git a/tests/test_determinism.py b/tests/test_determinism.py index b6f6f96..90e4cda 100644 --- a/tests/test_determinism.py +++ b/tests/test_determinism.py @@ -10,8 +10,13 @@ RNG, which ``seed_everything(..., workers=True)`` seeds per DataLoader worker, making the result depend on ``--num_workers``. -These tests pin both. They are ligand-only so they need no CCD asset and run on -plain CI. +A third source sat between them: ``preprocess_yamls`` collected results via +``as_completed``, so the manifest order followed worker scheduling rather than +input order, and the per-record RNG was seeded from that position. Seeding the +augmentation was therefore not enough on its own. + +These tests pin all three. They are ligand-only so they need no CCD asset and run +on plain CI. """ from __future__ import annotations @@ -24,7 +29,11 @@ from torch.utils.data import DataLoader from nesso.data.featurizer import NessoFeaturizer -from nesso.data.inference import InferenceDataset, inference_collate +from nesso.data.inference import ( + InferenceDataset, + inference_collate, + record_rng_seed, +) from nesso.data.types import Manifest from nesso.data.yaml_input import ( DEFAULT_CONFORMER_SEED, @@ -91,7 +100,9 @@ def test_conformer_seed_can_be_opted_out() -> None: assert mol.GetNumConformers() == 1 -def _ref_pos_by_record(tmp_path: Path, num_workers: int) -> dict[str, np.ndarray]: +def _ref_pos_by_record( + tmp_path: Path, num_workers: int, n_records: int = 3, reverse: bool = False +) -> dict[str, np.ndarray]: """Featurize several records through a real DataLoader at a given worker count.""" mol_dir = tmp_path / "rdkit_conformers" struct_dir = tmp_path / "structures" @@ -100,7 +111,7 @@ def _ref_pos_by_record(tmp_path: Path, num_workers: int) -> dict[str, np.ndarray directory.mkdir(parents=True, exist_ok=True) records = [] - for i in range(3): + for i in range(n_records): yaml_path = tmp_path / f"lig{i}.yaml" yaml_path.write_text(_YAML) struct, record, _, _ = parse_yaml( @@ -109,6 +120,9 @@ def _ref_pos_by_record(tmp_path: Path, num_workers: int) -> dict[str, np.ndarray struct.dump(struct_dir / f"{record.id}.npz") records.append(record) + if reverse: + records = list(reversed(records)) + dataset = InferenceDataset( manifest=Manifest(records), target_dir=tmp_path, @@ -191,3 +205,32 @@ def test_conformer_independent_of_batch_composition_and_order(tmp_path: Path) -> assert np.array_equal(alone[_SMILES], with_others[_SMILES]) assert np.array_equal(alone[_SMILES], reordered[_SMILES]) + + +def test_record_rng_seed_depends_on_id_not_position() -> None: + """The per-record RNG must be keyed on identity, not on the dataset index. + + ``as_completed`` in ``preprocess_yamls`` made the manifest order follow worker + scheduling, so seeding from the index made ``ref_pos`` depend on + ``--num_workers``. Adding, removing or renaming inputs shifts it too. + """ + assert record_rng_seed("lig") == record_rng_seed("lig") + assert record_rng_seed("lig") != record_rng_seed("other") + # numpy requires a uint32 seed. + assert 0 <= record_rng_seed("lig") < 2**32 + + +def test_ref_pos_independent_of_record_position(tmp_path: Path) -> None: + """The same record must featurize identically wherever it sits in the batch. + + Seeding the per-record RNG from the dataset index made a record's features + depend on its position, so reordering the manifest (which `as_completed` did + on its own at ``--num_workers > 1``) changed ``ref_pos``. + """ + forward = _ref_pos_by_record(tmp_path / "fwd", 0, n_records=3) + # Identical records, reversed order, so every index changes. + backward = _ref_pos_by_record(tmp_path / "rev", 0, n_records=3, reverse=True) + + assert set(forward) == set(backward) + for key in forward: + assert np.array_equal(forward[key], backward[key]), key