diff --git a/alphapulldown/prediction/prediction_batch.py b/alphapulldown/prediction/prediction_batch.py index 72e0fd1e..5867fd14 100644 --- a/alphapulldown/prediction/prediction_batch.py +++ b/alphapulldown/prediction/prediction_batch.py @@ -254,6 +254,46 @@ def _random_seed( return random.randrange(sys.maxsize) +def _setup_session( + backend: Any, fold_backend: str, configuration: Dict[str, Any] +) -> _AlphaPulldownSession: + """Validate the configuration, select the backend and set it up: the session every + job of a batch then reuses.""" + from alphapulldown.prediction.inference_flags import validate_model_configuration + + validate_model_configuration(configuration) + backend.change_backend(backend_name=fold_backend) + return _AlphaPulldownSession(dict(configuration), backend.setup(**configuration)) + + +def _predict_and_postprocess( + backend: Any, + session: _AlphaPulldownSession, + *, + fold_backend: str, + objects_to_model: list, + requested_seed: Optional[int], + postprocess_flags: Dict[str, Any], +) -> None: + """Predict the prepared objects with one seed and post-process each fold the backend + yields. A resumed fold it finds finished yields nothing, and still counts as done.""" + predicted_jobs = backend.predict( + **session.backend_values, + objects_to_model=objects_to_model, + random_seed=_random_seed(fold_backend, requested_seed, session), + # Resume compares them with those a finished fold was post-processed with. + postprocess_flags=postprocess_flags, + **session.configuration, + ) + for predicted_job in predicted_jobs: + backend.postprocess( + **postprocess_flags, + multimeric_object=predicted_job["object"], + prediction_results=predicted_job["prediction_results"], + output_dir=predicted_job["output_dir"], + ) + + class PreparedPredictionAdapter: """Adapter for the already-parsed invocation used by the legacy command.""" @@ -279,33 +319,18 @@ def configuration_for(self, job: PredictionJob) -> Dict[str, Any]: return self._model_flags def setup(self, configuration: Dict[str, Any]) -> _AlphaPulldownSession: - from alphapulldown.prediction.inference_flags import validate_model_configuration - - validate_model_configuration(configuration) - self._backend.change_backend(backend_name=self._fold_backend) - return _AlphaPulldownSession( - dict(configuration), self._backend.setup(**configuration) - ) + return _setup_session(self._backend, self._fold_backend, configuration) def predict(self, session: _AlphaPulldownSession, job: PredictionJob) -> None: del job - predicted_jobs = self._backend.predict( - **session.backend_values, + _predict_and_postprocess( + self._backend, + session, + fold_backend=self._fold_backend, objects_to_model=self._objects_to_model, - random_seed=_random_seed( - self._fold_backend, self._requested_random_seed, session - ), - # Resume compares them with those a finished fold was post-processed with. + requested_seed=self._requested_random_seed, postprocess_flags=self._postprocess_flags, - **session.configuration, ) - for predicted_job in predicted_jobs: - self._backend.postprocess( - **self._postprocess_flags, - multimeric_object=predicted_job["object"], - prediction_results=predicted_job["prediction_results"], - output_dir=predicted_job["output_dir"], - ) class AlphaPulldownPredictionAdapter: @@ -364,12 +389,7 @@ def configuration_for(self, job: PredictionJob) -> Dict[str, Any]: return model_flags def setup(self, configuration: Dict[str, Any]) -> _AlphaPulldownSession: - from alphapulldown.prediction.inference_flags import validate_model_configuration - - validate_model_configuration(configuration) - self._backend.change_backend(backend_name=self._flags.fold_backend) - backend_values = self._backend.setup(**configuration) - return _AlphaPulldownSession(dict(configuration), backend_values) + return _setup_session(self._backend, self._flags.fold_backend, configuration) def _prepare_protein_object(self, interactors: list, output_dir: str): from alphapulldown.prediction.fold_preparation import prepare_fold @@ -438,22 +458,11 @@ def _objects_to_model(self, job: PredictionJob) -> list: return objects_to_model def predict(self, session: _AlphaPulldownSession, job: PredictionJob) -> None: - objects_to_model = self._objects_to_model(job) - postprocess_flags = self._postprocess_flags() - predicted_jobs = self._backend.predict( - **session.backend_values, - objects_to_model=objects_to_model, - random_seed=_random_seed( - self._flags.fold_backend, self._flags.random_seed, session - ), - # Resume compares them with those a finished fold was post-processed with. - postprocess_flags=postprocess_flags, - **session.configuration, + _predict_and_postprocess( + self._backend, + session, + fold_backend=self._flags.fold_backend, + objects_to_model=self._objects_to_model(job), + requested_seed=self._flags.random_seed, + postprocess_flags=self._postprocess_flags(), ) - for predicted_job in predicted_jobs: - self._backend.postprocess( - **postprocess_flags, - multimeric_object=predicted_job["object"], - prediction_results=predicted_job["prediction_results"], - output_dir=predicted_job["output_dir"], - ) diff --git a/test/unit/test_prediction_batch.py b/test/unit/test_prediction_batch.py index a2178ff3..80232cf3 100644 --- a/test/unit/test_prediction_batch.py +++ b/test/unit/test_prediction_batch.py @@ -761,3 +761,114 @@ def __init__(self, interactors, **_kwargs): _, output_dir = fold_preparation.prepare_fold(interactors, str(tmp_path / "out"), flags) assert not (Path(output_dir) / fold_preparation.MULTIMERIC_FEATURES_FILENAME).exists() + + +class _RecordingBackend: + """Records every call an adapter makes and yields each prepared object back as a + predicted fold. ``fail`` = (stage, output directory) makes that stage raise for + that fold.""" + + def __init__(self, fail=None): + self.calls = [] + self.fail = fail + + def _maybe_fail(self, stage, output_dir): + if self.fail == (stage, output_dir): + raise RuntimeError(f"{stage} failed for {output_dir}") + + def change_backend(self, backend_name): + self.calls.append(("change_backend", backend_name)) + + def setup(self, **configuration): + self.calls.append(("setup", configuration)) + return {"model_runner": "resident-runner"} + + def predict(self, **kwargs): + self.calls.append(("predict", kwargs)) + for entry in kwargs["objects_to_model"]: + self._maybe_fail("predict", entry["output_dir"]) + yield { + "object": entry["object"], + "output_dir": entry["output_dir"], + "prediction_results": [1], + } + + def postprocess(self, **kwargs): + self.calls.append(("postprocess", kwargs)) + self._maybe_fail("postprocess", kwargs["output_dir"]) + + +def _write_af3_inputs(tmp_path, *names): + for name in names: + (tmp_path / name).write_text( + json.dumps({"name": name, "sequences": [], "modelSeeds": [1]}), encoding="utf-8" + ) + + +@pytest.mark.parametrize("requested_seed", [None, 13]) +def test_both_adapters_make_the_same_backend_calls_for_the_same_fold(tmp_path, requested_seed): + """The legacy command and the batch command share setup, seeding, prediction and + post-processing; given the same prepared fold they make the same calls.""" + from alphapulldown.prediction.inference_flags import model_flags, postprocess_flags + from alphapulldown.prediction.prediction_batch import ( + AlphaPulldownPredictionAdapter, + PredictionBatch, + PredictionJob, + PreparedPredictionAdapter, + ) + + _write_af3_inputs(tmp_path, "first.json", "second.json") + flags = _prediction_flags(tmp_path, "alphafold3") + flags.random_seed = requested_seed + batch_backend = _RecordingBackend() + PredictionBatch( + (PredictionJob("job", "first.json+second.json", tmp_path / "out"),) + ).run(AlphaPulldownPredictionAdapter(flags, backend=batch_backend)) + prepared = next(kwargs for call, kwargs in batch_backend.calls if call == "predict")[ + "objects_to_model" + ] + + legacy_backend = _RecordingBackend() + PredictionBatch((PredictionJob("job", "", tmp_path),)).run( + PreparedPredictionAdapter( + backend=legacy_backend, + fold_backend="alphafold3", + objects_to_model=prepared, + model_flags=model_flags(flags), + postprocess_flags=postprocess_flags(flags), + random_seed=requested_seed, + ) + ) + + assert legacy_backend.calls == batch_backend.calls + # The complex's two AF3 inputs stay one prediction; each yielded fold is post-processed. + assert [call for call, _ in batch_backend.calls] == [ + "change_backend", "setup", "predict", "postprocess", "postprocess", + ] + assert len(prepared) == 2 + + +@pytest.mark.parametrize("stage", ["predict", "postprocess"]) +def test_a_job_failing_in_prediction_or_postprocessing_does_not_stop_the_batch( + tmp_path, stage +): + from alphapulldown.prediction.prediction_batch import ( + AlphaPulldownPredictionAdapter, + PredictionBatch, + PredictionJob, + ) + + _write_af3_inputs(tmp_path, "first.json", "second.json") + backend = _RecordingBackend(fail=(stage, str(tmp_path / "out-first"))) + + summary = PredictionBatch( + ( + PredictionJob("first", "first.json", tmp_path / "out-first"), + PredictionJob("second", "second.json", tmp_path / "out-second"), + ) + ).run(AlphaPulldownPredictionAdapter(_prediction_flags(tmp_path, "alphafold3"), backend=backend)) + + assert summary.completed_job_ids == ("second",) + assert [failure.job_id for failure in summary.failures] == ["first"] + assert f"{stage} failed" in summary.failures[0].message + assert [call for call, _ in backend.calls].count("setup") == 1