Skip to content

Fix crash when --omit_probs is combined with 3Di masking - #140

Open
gbouras13 wants to merge 1 commit into
modernprost_12stfrom
fix-omit-probs-mask-threshold-12st
Open

Fix crash when --omit_probs is combined with 3Di masking#140
gbouras13 wants to merge 1 commit into
modernprost_12stfrom
fix-omit-probs-mask-threshold-12st

Conversation

@gbouras13

Copy link
Copy Markdown
Owner

Fixes the --omit_probs crash that threw away 20 minutes of MI250X compute on 161,473 sequences.

The defect

--omit_probs sets output_probs=False, which makes the inference engine store None in place of the per-residue probability array (inference.py:174-178). 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.

The masking loop also ran when masking was disabled: mask_threshold=0 makes the comparison all-False, so the loop did nothing but still required the array.

Worth noting: --mask_threshold defaults to 25 at the CLI (only the function signature defaults to 0), so plain phold predict --omit_probs with no --mask_threshold at all was also a guaranteed crash — not just the --mask_threshold 0 invocation from the report.

The fix

Fail fast, plus defensive guards, since they cover different callers:

  • validate_mask_options() in phold/utils/validation.py, called in run, predict and proteins-predict immediately after begin_phold — before the database check and before ProstT5 is loaded. Verified end to end: proteins-predict --omit_probs now exits in ~1s with a message naming both escape routes, and --omit_probs --mask_threshold 0 proceeds normally.
  • write_predictions() skips the masking loop entirely when mask_threshold is 0, and warns-and-continues rather than crashing when all_prob is None (library callers of get_embeddings bypass the CLI check).
  • subcommand_predict() gets the 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 doesn't catch. Proteins whose prediction genuinely failed still become "X" * len as before; a merely absent all_prob leaves the sequence unmasked rather than blanking a good prediction.

Not affected

predict_3di_12st.py hardcodes output_probs=True and doesn't pass mask_threshold. write_probs() only touches all_probs when output_path_all is non-None, which is itself gated on output_probs.

Testing

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 (#47), masking when probabilities are present, and the CLI validator. Reverting only predict_3Di.py makes 5 of them fail with the exact TypeError from the traceback.

Full unit suite: 95 passed, 18 skipped. Ruff: new file clean, src/phold/ count unchanged.

--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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant