[bugfix] fix flaky TRT predict test by making the mock model confident

[tiankongdeguiji]

Problem

test_multi_tower_with_fg_train_eval_export_trt (and the _zch variant) flake: the strict TRT-vs-baseline comparison in _test_rank_with_fg_trt occasionally fails by ~1e-3 prob on a few rows.

Root cause

The mock fixture trains on random labels, so the model can't learn and every logit ≈ 0 (prob ≈ 0.5). That is a maximally ill-conditioned regime (κ ≈ 1e5): TRT's tiny per-process GEMM-reduction nondeterminism (~1e-7 at the fp32 accumulator) gets amplified into ~1e-3 prob swings on these decision-boundary rows. It is a degenerate-fixture artifact, not a TRT correctness bug — eager and TRT compute the same function (they agree to ~1e-7 once logits are confident).

Fix

Make the model confident rather than loosening the tolerance. test_train_eval gains two optional, default-off params:

• learnable_label — encodes the label bimodally into a raw feature (0 → ~0.05, 1 → ~0.95), which lands in distinct bucketize bins with a clear gap, so the signal is cleanly separable.
• num_epochs — lets a test train enough to fit it.

The TRT tests encode into raw_1 and train 3 epochs. The model then drives |logit| large, eager and TRT agree to the fp32 floor (~1e-7), and the strict 2e-5 comparison is stable. Default behavior (random labels, original epochs) is unchanged for every other test.

Verification

• 240 separate-process launches per config: worst atol 1.8e-7 (custom_1 AutoDis) / 4.8e-7 (bucketized) ≪ 2e-5 — vs 3.6e-4 / 1.2e-3 unfixed.
• The TRT per-process nondeterminism still occurs (some launches show a nonzero process-vs-process diff) but is no longer κ-amplified, since no eval row sits at logit ≈ 0.
• Both test_multi_tower_with_fg_train_eval_export_trt and test_multi_tower_zch_with_fg_train_eval_export_trt pass locally.

🤖 Generated with Claude Code

[github-actions]

The fix is well-reasoned, but it can silently self-disable: the encoding in create_mock_join_data is guarded by learnable_label in t.column_names (utils.py:558), so if raw_1 is ever renamed/dropped from these configs the encoding is skipped, the model trains on random labels again, and this test quietly reverts to the exact flaky state it was meant to fix — with no failure signal, since nothing asserts the model became confident (only 2e-5 closeness + file existence are checked).

Consider adding a cheap confidence assertion to lock in the premise the whole fix rests on, e.g. after eval, assert the predicted probs are bimodal (((probs < 0.05) | (probs > 0.95)).mean() > 0.9) or assert the AUC in train/eval_result.txt is high. That turns a silent regression into a loud one and also answers "is num_epochs=3 always enough?" on every run.

[github-actions]

Minor/optional: this encodes inside the label loop keyed to the loop's last label, whereas create_mock_join_data keys to label_fields[0] (line 558-559). Harmless for the current single-label (clk) configs, but the two helpers would diverge for a multi-label test. Suggest hoisting it out of the loop and keying to a single explicit label for consistency:

    for label_field in label_fields:
        input_data[label_field] = pa.array(np.random.randint(2, size=(num_rows,)))
    if learnable_label and label_fields and learnable_label in input_data:
        label = input_data[label_fields[0]].to_numpy(zero_copy_only=False)
        input_data[learnable_label] = _learnable_label_value(label)

[github-actions]

Review summary — test-only change; LGTM in substance.

Nicely diagnosed and documented: fixing the ill-conditioning at the source (confident model) instead of loosening the 2e-5 tolerance is the right call, the inline kappa-amplification comment explains why well, and the 240-launch verification is convincing. I verified the mechanics: raw_1 is an item-side feature (expression: "item:raw_1") present in both TRT configs, and with boundaries: [0.1,0.2,0.3,0.4] the bimodal encoding lands label-0 in bin 0 (≈0.05±0.03) and label-1 in bin 4 (≈0.95±0.03) with a clean gap — and the ±0.03 noise never crosses a boundary, so bucketization is deterministic. Both train and eval data get the encoding, defaults stay None/unchanged for all other tests, and the docstrings are accurate. No perf or safety concerns.

Two non-blocking suggestions left inline:

1. Guard against silent regression — the encoding silently no-ops if raw_1 is ever renamed (and nothing asserts the model became confident), so a future config change could quietly revert the test to flaky. A cheap bimodality/AUC assertion would make that fail loudly.
2. Minor — create_mock_data keys the encoding to the loop's last label while create_mock_join_data uses label_fields[0]; consistent only for single-label configs.