Fix crash when --omit_probs is combined with 3Di masking - #140
Open
gbouras13 wants to merge 1 commit into
Open
Conversation
--omit_probs sets output_probs=False, which makes the inference engine store
None in place of the per-residue probability array. write_predictions() then
indexed all_prob[0] unconditionally and died with
TypeError: 'NoneType' object is not subscriptable
at write time -- after the entire GPU prediction had already completed. On a
real workload that threw away 20 minutes of MI250X compute on 161,473
sequences.
The masking loop also ran even when masking was disabled: mask_threshold=0
makes the comparison all-False, so the loop did nothing but still required
the array. --mask_threshold defaults to 25 at the CLI (only the function
signature defaults to 0), so plain --omit_probs with no --mask_threshold was
a guaranteed crash too.
Changes:
* Add validate_mask_options() and call it in run, predict and
proteins-predict immediately after begin_phold -- before the database
check and before ProstT5 is loaded. Failing in a second beats failing
after the compute.
* write_predictions(): skip the masking loop entirely when mask_threshold
is 0, and warn-and-continue rather than crash when all_prob is None
(library callers of get_embeddings bypass the CLI check).
* subcommand_predict(): same guard on the amino acid masking path. Passing
None to mask_low_confidence_aa reaches np.asarray(None, dtype=np.float64)
-- a silent array(nan) on numpy 1.x, a TypeError on numpy 2.x that the
(KeyError, IndexError) handler does not catch. Proteins whose prediction
genuinely failed still become "X" * len as before; a merely absent
all_prob leaves the sequence unmasked instead of blanking a good
prediction.
predict_3di_12st.py is unaffected: it hardcodes output_probs=True and does
not pass mask_threshold. write_probs() is unaffected: all_probs is only
touched when output_path_all is non-None, itself gated on output_probs.
Adds tests/unit/test_predict_3di_masking.py -- 13 tests covering
mask_threshold=0, all_prob=None, the mixed case, proteins-mode headers, the
zero-length drop (issue #47), masking when probabilities are present, and
the CLI validator.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes the
--omit_probscrash that threw away 20 minutes of MI250X compute on 161,473 sequences.The defect
--omit_probssetsoutput_probs=False, which makes the inference engine storeNonein place of the per-residue probability array (inference.py:174-178).write_predictions()then indexedall_prob[0]unconditionally and died withat write time — after the entire GPU prediction had already completed.
The masking loop also ran when masking was disabled:
mask_threshold=0makes the comparison all-False, so the loop did nothing but still required the array.Worth noting:
--mask_thresholddefaults to 25 at the CLI (only the function signature defaults to 0), so plainphold predict --omit_probswith no--mask_thresholdat all was also a guaranteed crash — not just the--mask_threshold 0invocation from the report.The fix
Fail fast, plus defensive guards, since they cover different callers:
validate_mask_options()inphold/utils/validation.py, called inrun,predictandproteins-predictimmediately afterbegin_phold— before the database check and before ProstT5 is loaded. Verified end to end:proteins-predict --omit_probsnow exits in ~1s with a message naming both escape routes, and--omit_probs --mask_threshold 0proceeds normally.write_predictions()skips the masking loop entirely whenmask_thresholdis 0, and warns-and-continues rather than crashing whenall_prob is None(library callers ofget_embeddingsbypass the CLI check).subcommand_predict()gets the same guard on the amino-acid masking path. PassingNonetomask_low_confidence_aareachesnp.asarray(None, dtype=np.float64)— a silentarray(nan)on numpy 1.x, aTypeErroron numpy 2.x that the(KeyError, IndexError)handler doesn't catch. Proteins whose prediction genuinely failed still become"X" * lenas before; a merely absentall_probleaves the sequence unmasked rather than blanking a good prediction.Not affected
predict_3di_12st.pyhardcodesoutput_probs=Trueand doesn't passmask_threshold.write_probs()only touchesall_probswhenoutput_path_allis non-None, which is itself gated onoutput_probs.Testing
tests/unit/test_predict_3di_masking.py— 13 tests coveringmask_threshold=0,all_prob=None, the mixed case, proteins-mode headers, the zero-length drop (#47), masking when probabilities are present, and the CLI validator. Reverting onlypredict_3Di.pymakes 5 of them fail with the exactTypeErrorfrom the traceback.Full unit suite: 95 passed, 18 skipped. Ruff: new file clean,
src/phold/count unchanged.