fix: report mistyped backend options instead of dropping them

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, 

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,40 +181,61 @@ 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}")
 
         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...")
         prepared_inputs = AlphaFold3Backend.prepare_input(
             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)

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:

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."
         )
 
 

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))

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))