[feat] HSTUMatch: scalar item-tower export view

[tiankongdeguiji]

Summary

HSTUMatchItemTower gets a scalar export view (candidate.query per row over scalar video_id) so the recall item tower can be exported against a one-row-per-item parquet, decoupled from the training-shape sequence rows in train_input_path.

Sits on top of #506 (HSTUMatch encoder refactor), #519 (features/feature_groups as @Property end-to-end), and #520 (sampler bare-name attr_fields), all already on master.

Two-view model

view	item-tower key	data source
training	candidate.sequence (jagged)	train_input_path (sequence rows)
item-export	candidate.query (scalar)	--item_input_path (scalar parquet)

Key design decisions

• Lazy view-swap on HSTUMatchItemTower. features / feature_groups are @property-driven by _is_inference: training reads the parent's grouped sub-features; export reads a lazily-built scalar projection cached in _features_scalar / _feature_groups_scalar. set_is_inference(True) is a cheap flag toggle, no structural mutation; set_is_inference(False) reverts cleanly. forward() derives the candidate key inline ({group}.sequence vs {group}.query) from the flag — no _cand_key attribute.

• Materialized sequence defaults in project_grouped_sequence_feature_to_scalar. An id_feature projected from sequence mode (effective default_value="0" / value_dim=1) carries those defaults forward to the scalar proto; otherwise the export would silently flip to scalar defaults ("" / 0), diverging from the training sub-feature.

• Single set_is_inference flow. Hoisted to one site in main.py::export before every InferWrapper(...) call. recursive_setattr propagates to every sub-module before any wrapper builds its EmbeddingGroup off view-dependent features. Removed from the three post-wrap call sites in export_util.py.

• Item-tower input as a CLI flag. --item_input_path on tzrec/export.py plumbs through main.export() → export_model(data_input_path=...) → predict-mode dataloader. main.py::export applies the policy that only name == "item_tower" receives the override; the user tower keeps reading train_input_path. No proto change — EasyRecConfig is untouched.

Test plan

• python -m unittest tzrec.features.feature_test.ProjectGroupedSequenceFeatureToScalarTest tzrec.models.hstu_test — green on local A10. test_hstu_match now uses the production grouped-sequence pattern (sequence_feature wrapping video_id sub-feature, shared embedding across uih_seq / cand_seq via embedding_name) and asserts the scalar-view contract (bare names, non-grouped, scalar feature_groups.feature_names) as the final step.
• python -m unittest tzrec.tests.match_integration_test.MatchIntegrationTest.test_hstu_with_fg_train_eval — train + eval + AOT export + item/user predict, green on local A10.
• python -m unittest tzrec.tests.match_integration_test.MatchIntegrationTest.test_{dssm,dat,mind}_*_train_eval_export — non-HSTU match regression, green.
• pre-commit run — clean on all touched files.

Known limitations / follow-ups

• Scalar export only handles grouped sequence_feature sub-features. Top-level sequence_id_feature / sequence_raw_feature would raise ValueError from project_grouped_sequence_feature_to_scalar if reached; HSTUMatch doesn't use them today.

🤖 Generated with Claude Code

[github-actions]

Stale reference: item_input_path was moved out of pipeline_config in the precursor commit (1e885b53) and is now the --data_input_path CLI flag. Suggest:

Suggested change

	# `pipeline_config.item_input_path` instead of `train_input_path`.
	# so export_util.py can route item-tower export through the
	# `--data_input_path` CLI override instead of `train_input_path`.

[github-actions]

This EmbeddingGroup snapshots the current view at construction time, which only works because the caller in main.py::export already ran model.set_is_inference(True) before wrapping. That invariant is undocumented at this call site and silent if violated (EG built off training jagged features → predict path uses .query key → KeyError at inference). Consider an explicit assertion to lock the contract:

if hasattr(module, "_is_inference"):
    assert module._is_inference, (
        "TowerWoEGWrapper requires set_is_inference(True) to have run on "
        "the inner module before wrapping; otherwise EmbeddingGroup is "
        "built off the wrong view."
    )

[github-actions]

Review summary

Solid refactor with thoughtful comments and integration coverage. A few noteworthy points (left as inline comments):

Likely bug

• export_util.py:305 — hasattr(model, \"_feature_groups\") no longer fires on the post-PR ScriptWrapper/TowerWoEGWrapper, so the scalar-view groups are never written into the exported pipeline.config. The comment says they must be. Suggest dropping the gate and writing model.feature_groups unconditionally.

Doc

• model.py:351 — comment still references pipeline_config.item_input_path, now removed in favor of --data_input_path.

Tests / contracts

• match_model.py:514 — TowerWoEGWrapper relies on set_is_inference(True) having run before wrapping; consider an assert.
• hstu_test.py:207 — the scalar branch of HSTUMatchItemTower.forward is now only covered black-box via the AOT integration test; a focused unit test on set_is_inference(True) would lock the contract.
• export_util.py:744 — RTP data_input_path plumbing is untested; existing RTP tests only hit the fallback branch.

No concerns on perf (the hot-path lookup shrinks from dict-keyed to a single attribute load), the new copy.deepcopy(sampler_config) (sound, prevents cross-dataloader bleed), or the set_is_inference hoist (recursive_setattr is hasattr-gated and _is_inference is uniquely a BaseModule attribute).

[github-actions]

Worth a short comment here naming the reason this must flip on all ranks pre-wrap: export_rtp_model runs init_process_group() / dist.barrier() / DistributedModelParallel collectives, which require every rank to share the same model topology. If set_is_inference ever alters sharded-module shape (as it now does, via HSTUMatchItemTower.features's scalar projection), a rank-0-only flip would diverge ranks. A one-liner would prevent a future refactor from moving it back inside the rank-0 branch.

Also, asymmetry to watch for: the outer ScriptWrapper is constructed after the flip, so its own _is_inference is False while the inner self.model._is_inference is True. Functionally fine (nothing reads the outer flag), but surprising — same comment can call this out.

[github-actions]

Worth documenting the cache-validity invariant: set_is_inference(False) flips the flag back but does not clear _features_scalar / _feature_groups_scalar. Reuse on a subsequent flip is correct because _features is immutable post-__init__, but a future maintainer reading this will reasonably wonder whether the cache should be invalidated on revert. One line ("cache valid for tower lifetime; _features is immutable") would prevent a well-intentioned "fix".

[github-actions]

Two things:

1. The hasattr guards are meaningful only for OverlapFeature (its default_value is commented out in the proto); for the other 11 SeqFeatureConfig.feature variants both fields are always present, so the guard reads as defensive code with no purpose. Worth naming OverlapFeature in a comment.
2. Test gap: the "don't overwrite" branches (if not dst_msg.default_value / not dst_msg.HasField("value_dim")) are not exercised. The test in feature_test.py only covers the "both unset → materialize from source" path. Consider adding a case where the inner proto already sets e.g. default_value="-1" or value_dim=4 and asserting the materialization does not overwrite. Same for value_dim > 1 on raw_feature — currently only value_dim == 1 (default) is asserted.

[github-actions]

The scalar-view contract assertions are good, but the new forward() branch (hstu.py:240-241, the .query suffix path) is never invoked here — only via the AOT integration test in match_integration_test.py. If the wrapper's EmbeddingGroup ever emits the scalar group under a key that doesn't match self._group_name + ".query", this unit test won't catch it.

Cheap addition: build a fake grouped_features={"candidate.query": <tensor>} and call hstu.item_tower(grouped_features), asserting the output shape. Locks the suffix contract at unit-test latency rather than relying on the 30-min integration test to surface a regression.

[github-actions]

The user-tower fallback to train_input_path is asserted only implicitly by the integration test (which happens to predict the user tower over the eval parquet, not the train parquet). So if a refactor accidentally routed item_input_path to the user tower too, the integration test would still pass.

If you want to lock this in cheaply, a small unit test in main_test.py that monkey-patches export_model and asserts data_input_path per tower (item_tower → item_input_path, user_tower → None) would pin the policy.

[github-actions]

Review summary

Clean, well-scoped change. The two-view model (.sequence for training, .query for export) is implemented as a lazy property-driven flag flip on HSTUMatchItemTower with cached scalar projections — minimal blast radius, training path unchanged. Hoisting set_is_inference(True) to a single pre-wrap site in main.py::export (and removing it from three post-wrap sites in export_util.py) is the right consolidation: it ensures every rank shares the same model topology before any collective in the RTP path.

No perf, security, or correctness concerns surfaced. Findings are all low — mostly clarity/documentation and minor test gaps:

• main.py:946 — the all-rank pre-wrap flip is load-bearing for the RTP collective path; a one-line comment would prevent future regression.
• hstu.py:170-171 — cache survives set_is_inference(False); document the invariant.
• feature.py:1263-1266 — hasattr guards are only meaningful for OverlapFeature; the "don't overwrite" branches are untested.
• hstu_test.py:240-246 — forward() with _is_inference=True (the .query suffix branch) is only covered by the AOT integration test; a cheap unit test would tighten the contract.
• main.py:970-973 — user-tower fallback to train_input_path is asserted only implicitly; a unit test pinning the per-tower policy would lock it.

Nothing blocking. Nice work threading data_input_path through cleanly without touching the proto, and the test reshape to mirror the production grouped-sequence pattern (shared embedding via embedding_name) is a real improvement to test_hstu_match.