From dca9fc4b389d0353312783d9841474852c09992b Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Tue, 13 Oct 2026 20:31:59 +0200 Subject: [PATCH] fix: report mistyped backend options instead of dropping them The adapters give every backend's setup and predict the whole configuration, so each took the rest as **kwargs and read its own options with kwargs.get(name, default). A name mistyped on either side, in the configuration table or in a backend, fell back to the default without a word. setup and predict of AlphaFold 2, AlphaFold 3 and AlphaLink now name every option they use, and pass the options they do not name, meant for other backends, to validate_model_configuration, which rejects any no backend takes. AF3 hands its resume code an explicit dict of the settings it records instead of its kwargs, and no longer forwards af3_input_json, which nothing supplied. test_inference_flags reads the backends' signatures from source and checks both directions: every configuration key is named by some backend, every named option is a configuration key, and the unnamed options are only validated, never read. Co-Authored-By: Claude Opus 5.5 --- .../folding_backend/alphafold2_backend.py | 50 +++++++++---- .../folding_backend/alphafold3_backend.py | 50 ++++++++++--- .../folding_backend/alphalink_backend.py | 44 +++++++----- alphapulldown/prediction/inference_flags.py | 13 ++-- test/unit/test_alphafold2_backend_helpers.py | 10 +++ test/unit/test_alphafold3_backend_helpers.py | 10 +++ test/unit/test_inference_flags.py | 70 +++++++++++++++++++ 7 files changed, 198 insertions(+), 49 deletions(-) diff --git a/alphapulldown/folding_backend/alphafold2_backend.py b/alphapulldown/folding_backend/alphafold2_backend.py index 02453cb7..e4b6c498 100644 --- a/alphapulldown/folding_backend/alphafold2_backend.py +++ b/alphapulldown/folding_backend/alphafold2_backend.py @@ -25,6 +25,7 @@ #from alphapulldown.utils.calculate_rmsd import calculate_rmsd_and_superpose from alphapulldown.utils.modelling_setup import pad_input_features from alphapulldown.prediction import fast_kernels as fast_kernels_policy +from alphapulldown.prediction.inference_flags import validate_model_configuration from alphapulldown.prediction.jax_compilation_cache import enable_persistent_compilation_cache from alphapulldown.prediction.resume import ( finish_fold_record, @@ -617,7 +618,8 @@ def setup( msa_depth=None, dropout=False, fast_kernels="off", - **kwargs, + jax_compilation_cache_dir=None, + **other_backends_options, ) -> Dict: """ Initializes and configures multiple AlphaFold model runners. @@ -647,8 +649,11 @@ def setup( fast_kernels : str, optional "off" (default), "on" or "auto": ColabFold's fused Pallas kernels for the multimer models (``alphapulldown.prediction.fast_kernels``). - **kwargs : dict - Additional keyword arguments for model runner configuration. + jax_compilation_cache_dir : str, optional + A persistent on-disk JAX compile cache, shared by later processes. + **other_backends_options : dict + Options of other backends: the adapters give every backend the whole + configuration. One no backend takes is an error rather than dropped. Returns ------- @@ -661,13 +666,14 @@ def setup( If provided custom model names are not part of the available models. """ + validate_model_configuration(other_backends_options) _ensure_typing_dataclass_transform() from alphafold.model import config from alphafold.model import data, model # AlphaFold2 inference is JAX-compiled, and a fresh process recompiles every # model runner. Honour a persistent on-disk compile cache when one is given. - enable_persistent_compilation_cache(kwargs.get("jax_compilation_cache_dir")) + enable_persistent_compilation_cache(jax_compilation_cache_dir) num_ensemble = 1 model_runners = {} @@ -765,7 +771,10 @@ def predict_individual_job( skip_templates: bool, output_dir: Dict, random_seed: int, - **kwargs, + desired_num_res: int | None = None, + desired_num_msa: int | None = None, + debug_msas: bool = False, + debug_templates: bool = False, ) -> dict: """ Executes structure predictions using configured AlphaFold models on one individual job. @@ -785,8 +794,10 @@ def predict_individual_job( The directory to save prediction results and PDB files. random_seed : int, optional A seed for random number generation to ensure reproducibility, default is 42. - **kwargs : dict - Additional keyword arguments for prediction. + desired_num_res, desired_num_msa : int, optional + Pad the features to this many residues and MSA rows (``--desired_num_res``). + debug_msas, debug_templates : bool, optional + Write the processed MSA and templates next to the outputs, for inspection. Returns ------- @@ -845,8 +856,6 @@ def predict_individual_job( f"All predictions for {multimeric_object.description} are already completed.") return prediction_results # first check whether the desired num_res and num_msa are specified for padding - desired_num_res, desired_num_msa = kwargs.get( - "desired_num_res", None), kwargs.get("desired_num_msa", None) if (desired_num_res is not None and type(multimeric_object) == MultimericObject): # Padding every fold in a batch to one shape lets a resident session compile @@ -889,14 +898,14 @@ def predict_individual_job( raise ValueError( "No template_all_atom_positions key found in processed_feature_dict." ) - if kwargs.get("debug_msas", False): + if debug_msas: _write_processed_msa_debug_artifact( processed_feature_dict=processed_feature_dict, multimeric_object=multimeric_object, output_dir=output_dir, model_name=model_name, ) - if kwargs.get("debug_templates", False): + if debug_templates: _write_processed_template_debug_artifacts( processed_feature_dict=processed_feature_dict, output_dir=output_dir, @@ -962,7 +971,19 @@ def predict(model_runners: Dict, skip_templates: bool, random_seed: int = 42, postprocess_flags: Dict = None, - **kwargs): + *, + desired_num_res: int | None = None, + desired_num_msa: int | None = None, + debug_msas: bool = False, + debug_templates: bool = False, + **other_backends_options): + """Predict each object with every model runner, yielding one job per object. + + Options of other backends arrive in ``other_backends_options`` (the adapters + give every backend the whole configuration); one no backend takes is an error + rather than dropped. + """ + validate_model_configuration(other_backends_options) for entry in objects_to_model: object_to_model = entry['object'] output_dir = entry['output_dir'] @@ -981,7 +1002,10 @@ def predict(model_runners: Dict, skip_templates=skip_templates, output_dir=output_dir, random_seed=random_seed, - **kwargs + desired_num_res=desired_num_res, + desired_num_msa=desired_num_msa, + debug_msas=debug_msas, + debug_templates=debug_templates, ) yield {'object': object_to_model, diff --git a/alphapulldown/folding_backend/alphafold3_backend.py b/alphapulldown/folding_backend/alphafold3_backend.py index af4f38de..f315bd3b 100644 --- a/alphapulldown/folding_backend/alphafold3_backend.py +++ b/alphapulldown/folding_backend/alphafold3_backend.py @@ -42,6 +42,7 @@ process_fold_input, ) from alphapulldown.prediction import af3_fused_triangles +from alphapulldown.prediction.inference_flags import validate_model_configuration from alphapulldown.prediction.jax_compilation_cache import enable_persistent_compilation_cache from alphapulldown.prediction.resume import finish_fold_record from alphapulldown.objects import MultimericObject, MonomericObject, ChoppedObject @@ -63,9 +64,15 @@ def setup( return_embeddings: bool = False, return_distogram: bool = False, fast_kernels: str = "off", - **kwargs, + **other_backends_options, ) -> Dict: - """Sets up the ModelRunner with the given configurations.""" + """Sets up the ModelRunner with the given configurations. + + The adapters give every backend the whole configuration; the options of other + backends arrive in ``other_backends_options``, and an option no backend takes is an + error rather than dropped. + """ + validate_model_configuration(other_backends_options) enable_persistent_compilation_cache(jax_compilation_cache_dir) _initialise_tokamax_trace_context() @@ -174,14 +181,30 @@ def predict( objects_to_model: List[Dict[str, Union[MultimericObject, MonomericObject, ChoppedObject, 'folding_input.Input', str]]], random_seed: int, buckets: int | Sequence[int], - **kwargs, + *, + allow_resume: bool = False, + postprocess_flags: Dict | None = None, + features_directory=None, + debug_templates: bool = False, + debug_msas: bool = False, + num_seeds: int | None = None, + num_diffusion_samples: int | None = None, + num_recycles: int | None = None, + return_embeddings: bool = False, + return_distogram: bool = False, + **other_backends_options, ): """Predicts structures for a list of objects using AlphaFold 3. Supports merging AlphaPulldown protein objects and input.json objects into a single job. ``random_seed=None`` reuses an earlier run's seed when resuming. With ``allow_resume``, a fold already complete is neither predicted nor yielded. + ``num_diffusion_samples``, ``num_recycles`` and the embedding and distogram + switches, set up in ``setup``, are recorded with each fold for resume. Options of + other backends arrive in ``other_backends_options``; one no backend takes is an + error rather than dropped. """ + validate_model_configuration(other_backends_options) logging.info(f"predict called with {len(objects_to_model)} objects to model") for i, obj in enumerate(objects_to_model): logging.info(f"Object {i}: {obj}") @@ -189,8 +212,14 @@ def predict( if isinstance(buckets, int): buckets = [buckets] buckets = tuple(int(b) for b in buckets) - allow_resume = bool(kwargs.get("allow_resume", False)) - postprocess_flags = kwargs.get("postprocess_flags") or {} + postprocess_flags = postprocess_flags or {} + model_flags = { + 'num_seeds': num_seeds, + 'num_diffusion_samples': num_diffusion_samples, + 'num_recycles': num_recycles, + 'return_embeddings': return_embeddings, + 'return_distogram': return_distogram, + } # Prepare inputs logging.info("Calling prepare_input...") @@ -198,16 +227,15 @@ def predict( objects_to_model=objects_to_model, # A placeholder: _start_or_resume_fold settles each fold's seeds. random_seed=0 if random_seed is None else random_seed, - af3_input_json=kwargs.get("af3_input_json"), - features_directory=kwargs.get("features_directory"), - debug_templates=kwargs.get("debug_templates", False), - debug_msas=kwargs.get("debug_msas", False), + features_directory=features_directory, + debug_templates=debug_templates, + debug_msas=debug_msas, ) # Run predictions for prepared in prepared_inputs: output_dir = prepared.output_dir started = _start_or_resume_fold( - prepared.fold_input, output_dir, random_seed, allow_resume, kwargs, postprocess_flags + prepared.fold_input, output_dir, random_seed, allow_resume, model_flags, postprocess_flags ) if started is None: continue @@ -223,7 +251,7 @@ def predict( output_dir=output_dir, buckets=buckets, resolve_msa_overlaps=prepared.resolve_msa_overlaps, - debug_msas=kwargs.get('debug_msas', False), + debug_msas=debug_msas, seeds_done=seeds_done, modelcif_metadata=( _modelcif_metadata(fold_input_obj, output_dir) diff --git a/alphapulldown/folding_backend/alphalink_backend.py b/alphapulldown/folding_backend/alphalink_backend.py index 4358fab0..876158d9 100644 --- a/alphapulldown/folding_backend/alphalink_backend.py +++ b/alphapulldown/folding_backend/alphalink_backend.py @@ -15,6 +15,7 @@ import re, json from alphapulldown.folding_backend.alphafold2_backend import _save_pae_json_file from alphapulldown.objects import MultimericObject, MonomericObject, ChoppedObject +from alphapulldown.prediction.inference_flags import validate_model_configuration from alphapulldown.utils.plotting import plot_pae_from_matrix from .folding_backend import FoldingBackend import torch @@ -49,7 +50,7 @@ class AlphaLinkBackend(FoldingBackend): def setup( model_dir: str, model_name: str = "multimer_af2_crop", - **kwargs, + **other_backends_options, ) -> Dict: """ Initializes and configures an AlphaLink model runner including crosslinking data. @@ -63,10 +64,9 @@ def setup( 1. A directory containing AlphaLink weights files (e.g., AlphaLink-Multimer_SDA_v3.pt) 2. A specific AlphaLink weights file (e.g., /path/to/AlphaLink-Multimer_SDA_v3.pt) Expected file names: AlphaLink-Multimer_SDA_v2.pt or AlphaLink-Multimer_SDA_v3.pt - crosslinks : str - The path to the file containing crosslinking data. - **kwargs : dict - Additional keyword arguments for model configuration. + **other_backends_options : dict + Options of other backends: the adapters give every backend the whole + configuration. One no backend takes is an error rather than dropped. Returns ------- @@ -82,6 +82,7 @@ def setup( ValueError If the file does not have a .pt extension or is not a recognized AlphaLink weights file. """ + validate_model_configuration(other_backends_options) # Check if model_dir is a file or directory if os.path.isfile(model_dir): # Direct file path provided @@ -472,7 +473,13 @@ def predict_iterations(feature_dict, output_dir='', param_path='', input_seqs=[] @staticmethod def predict( objects_to_model: List[Dict[str, Union[MultimericObject, MonomericObject, ChoppedObject, str]]], - **kwargs, + configs=None, + param_path=None, + crosslinks=None, + num_predictions_per_model: int = 1, + random_seed=None, + postprocess_flags=None, + **other_backends_options, ): """ Predicts the structure of proteins using configured AlphaLink models. @@ -483,17 +490,19 @@ def predict( A list of dictionaries. Each dictionary has keys 'object' and 'output_dir'. The 'object' key contains an instance of MultimericObject, MonomericObject, or ChoppedObject. The 'output_dir' key contains the corresponding output directory to save the modelling results. - **kwargs : dict - Additional keyword arguments including: - - configs: Configuration dictionary for the AlphaLink model - - param_path: Path to the AlphaLink model parameters - - crosslinks: Path to crosslink information pickle for AlphaLink + configs, param_path + The AlphaLink model configuration and weights, from ``setup``. + crosslinks : str, optional + Path to the crosslink information pickle. + num_predictions_per_model : int + Predictions per object, default 1. + random_seed, postprocess_flags + Given to every backend's ``predict``; AlphaLink uses neither. + **other_backends_options : dict + Options of other backends (the adapters give every backend the whole + configuration); one no backend takes is an error rather than dropped. """ - # Extract required parameters from kwargs - configs = kwargs.get('configs') - param_path = kwargs.get('param_path') - crosslinks = kwargs.get('crosslinks') - + validate_model_configuration(other_backends_options) if not all([configs, param_path]): raise ValueError("Missing required parameters: configs or param_path") @@ -508,9 +517,6 @@ def predict( output_dir = entry['output_dir'] os.makedirs(output_dir, exist_ok=True) - # Get num_predictions_per_model from kwargs, default to 1 - num_predictions_per_model = kwargs.get('num_predictions_per_model', 1) - # Get chain_id_map if available, otherwise create a default one chain_id_map = getattr(object_to_model, 'chain_id_map', None) if chain_id_map is None: diff --git a/alphapulldown/prediction/inference_flags.py b/alphapulldown/prediction/inference_flags.py index f36ecf9f..587d5df5 100644 --- a/alphapulldown/prediction/inference_flags.py +++ b/alphapulldown/prediction/inference_flags.py @@ -116,8 +116,9 @@ def unsupported_flags(backend_name: str, present: Iterable[str]) -> list[str]: "convert_to_modelcif": "convert_to_modelcif", } -# Keys a backend's setup() may legitimately be given. Backends ignore what they do not -# use, so this set is the only thing standing between a mistyped key and silence. +# Keys a backend's setup() and predict() may be given. Every backend is given them all, +# names the ones it uses and checks the rest against this set; test_inference_flags ties +# the set to the backends' parameter names, so a key no backend names cannot slip in. MODEL_CONFIGURATION_KEYS = frozenset(_MODEL_FLAG_SOURCES) | { "model_name", # added for AlphaFold2 multimer batches @@ -166,15 +167,15 @@ def postprocess_flags(flags: Any) -> Dict[str, Any]: def validate_model_configuration(configuration: Mapping[str, Any]) -> None: """Reject configuration keys no backend consumes. - ``setup(**configuration)`` lets every backend absorb keys meant for another one in - ``**kwargs``, so without this a mistyped key is dropped without a word and the - setting it was meant to carry simply never takes effect. + Every backend's ``setup`` and ``predict`` take the keys meant for other backends as + ``**other_backends_options`` and pass them here, so a mistyped key is reported + instead of never taking effect. """ unknown = sorted(set(configuration) - MODEL_CONFIGURATION_KEYS) if unknown: raise ValueError( f"Unknown model configuration key(s): {unknown}. " - "Backends ignore keys they do not use, so this would otherwise be silent." + "No backend takes them, so they would have no effect." ) diff --git a/test/unit/test_alphafold2_backend_helpers.py b/test/unit/test_alphafold2_backend_helpers.py index be39a41c..af1e8f5a 100644 --- a/test/unit/test_alphafold2_backend_helpers.py +++ b/test/unit/test_alphafold2_backend_helpers.py @@ -2095,3 +2095,13 @@ def test_msa_depths_follow_the_fixed_depth_or_the_logarithmic_scan(af2_backend_m ] with pytest.raises(ValueError): af2_backend_module._msa_depths(evoformer, 2) + + +def test_af2_predict_reports_an_option_no_backend_takes(af2_backend_module): + predict = af2_backend_module.AlphaFold2Backend.predict + common = dict(model_runners={}, objects_to_model=[], allow_resume=False, skip_templates=False) + + # Other backends' options are expected: every backend is given the whole configuration. + assert list(predict(**common, num_diffusion_samples=5, crosslinks=None)) == [] + with pytest.raises(ValueError, match="debug_msa"): + list(predict(**common, debug_msa=True)) diff --git a/test/unit/test_alphafold3_backend_helpers.py b/test/unit/test_alphafold3_backend_helpers.py index c102e9ff..feb999db 100644 --- a/test/unit/test_alphafold3_backend_helpers.py +++ b/test/unit/test_alphafold3_backend_helpers.py @@ -2774,3 +2774,13 @@ def test_prepare_input_writes_template_error_artifacts_on_mmcif_generation_failu assert "Chain ID: A" in error_report assert "Domain name:" in error_report assert "Template sequence:" in error_report + + +def test_af3_predict_reports_an_option_no_backend_takes(af3_backend_module): + predict = af3_backend_module.AlphaFold3Backend.predict + common = dict(model_runner=object(), objects_to_model=[], random_seed=1, buckets=256) + + # Other backends' options are expected: every backend is given the whole configuration. + assert list(predict(**common, num_cycle=3, skip_templates=False)) == [] + with pytest.raises(ValueError, match="debug_msa"): + list(predict(**common, debug_msa=True)) diff --git a/test/unit/test_inference_flags.py b/test/unit/test_inference_flags.py index a645edf5..64c59244 100644 --- a/test/unit/test_inference_flags.py +++ b/test/unit/test_inference_flags.py @@ -1,5 +1,8 @@ """The inference flag set has one home; these pin what that home guarantees.""" +import ast +from pathlib import Path + import pytest from alphapulldown.prediction import inference_flags @@ -114,3 +117,70 @@ def test_unifold_validation_rejects_the_backend_before_checking_flags(): def test_both_alphafold_backends_accept_allow_resume(): for backend in ("alphafold2", "alphafold3"): assert inference_flags.unsupported_flags(backend, ["allow_resume"]) == [], backend + + +# The backends are read from source: the AlphaFold 2 and 3 backends cannot be imported +# into one environment. +BACKENDS = Path(__file__).resolve().parents[2] / "alphapulldown" / "folding_backend" +# What each backend's setup returns, handed by the adapters to its own predict. +SETUP_VALUES = { + "alphafold2_backend.py": {"model_runners"}, + "alphafold3_backend.py": {"model_runner"}, + "alphalink_backend.py": {"configs", "param_path"}, +} +# What the adapters give every predict besides the configuration. +PREDICT_ARGUMENTS = {"objects_to_model", "random_seed", "postprocess_flags"} + + +def _setup_and_predict(filename): + """The backend class's setup and predict: each one's definition and named parameters.""" + tree = ast.parse((BACKENDS / filename).read_text()) + (backend,) = [ + node for node in tree.body + if isinstance(node, ast.ClassDef) + and {"setup", "predict"} <= {item.name for item in node.body if isinstance(item, ast.FunctionDef)} + ] + methods = {} + for node in backend.body: + if isinstance(node, ast.FunctionDef) and node.name in ("setup", "predict"): + arguments = node.args + named = {arg.arg for arg in arguments.posonlyargs + arguments.args + arguments.kwonlyargs} + methods[node.name] = (node, named) + return methods + + +def test_every_configuration_key_is_named_by_a_backend(): + # A key no backend names is a setting that never takes effect, such as a + # mistyped key in _MODEL_FLAG_SOURCES. + named = set() + for filename in SETUP_VALUES: + for _, parameters in _setup_and_predict(filename).values(): + named |= parameters + assert sorted(inference_flags.MODEL_CONFIGURATION_KEYS - named) == [] + + +@pytest.mark.parametrize("filename", sorted(SETUP_VALUES)) +def test_every_backend_option_is_a_configuration_key(filename): + # An option the configuration never supplies keeps its default whatever the flags + # say, such as a mistyped parameter name. + for name, (_, parameters) in _setup_and_predict(filename).items(): + options = parameters - PREDICT_ARGUMENTS - SETUP_VALUES[filename] + assert sorted(options - inference_flags.MODEL_CONFIGURATION_KEYS) == [], name + + +@pytest.mark.parametrize("filename", sorted(SETUP_VALUES)) +def test_backends_only_validate_the_options_they_do_not_name(filename): + # Reading an unnamed option with kwargs.get(...) would bring back the silent default. + for name, (method, _) in _setup_and_predict(filename).items(): + rest = method.args.kwarg.arg + uses = [ + node for node in ast.walk(method) + if isinstance(node, ast.Name) and node.id == rest + ] + checks = [ + node for node in ast.walk(method) + if isinstance(node, ast.Call) + and getattr(node.func, "id", None) == "validate_model_configuration" + and [getattr(arg, "id", None) for arg in node.args] == [rest] + ] + assert len(uses) == len(checks) == 1, name