[feat]online export of dense models#604
Open
eric-gecheng wants to merge 85 commits into
Open
Conversation
…ew instead of mutating the checkpoint dir
…asyRec into feature/export_split
The full export special-cases MatchModel (per-tower user/item dirs) and TDM (separate embedding/ and model/ dirs), but the online dense export unconditionally wrapped the whole model and published one scripted_model.pt at the version root, so a hot swap for such a deployment would load an artifact whose layout doesn't match what the processor deployed from the full export. Reject these model families with a clear error right after model build instead of silently producing an incompatible monolithic export. Co-Authored-By: Claude <noreply@anthropic.com>
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Contributor
|
Follow-up static review on the latest head: I left six inline comments covering meta-parameter warm-up failure, Kafka consumer-group interference, retention deleting the published target, two exporter shutdown/lifecycle races, and unbounded export logs. The meta warm-up, Kafka input reuse, publication retention, and shutdown issues are blockers because they can prevent exports, disrupt live training, or publish after cleanup. Static review only; tests and builds were not run as requested. |
export_dense_model_cpu's pre-trace warm-up runs a forward to materialize lazy dense submodules, but it ran before embedding params were materialized: models default to meta device and restore_model happens only later, so the warm-up's sparse lookup executed on meta tensors and crashed. The distributed path avoids this by restoring a materialized DMP first; the E2E test hid the bug by calling init_parameters(..., cpu) on the model before export, so real online dense exports (which never init_parameters) failed at warm-up. Call init_parameters(model, device) right before the warm-up to allocate meta embeddings -- the warm-up only needs shapes to flow through lazy modules, and real weights still come from restore_model(..., error_on_missing_keys=True) below. Rework the E2E test to export a meta model (no init_parameters) so it guards the fix instead of masking it. Co-Authored-By: Claude <noreply@anthropic.com>
export_dense_model_cpu falls back to pipeline_config.train_input_path when data_input_path is unset. For Kafka training that path carries the live training group.id, so the export's predict dataloader subscribed with it, joined the training consumer group, and held partitions through the slow tracing/restoring/scripting -- rebalancing training on every checkpoint. Rewrite group.id to an isolated <orig>__dense_export group when the fallback input is a kafka URI, so the export consumer never joins the training group; non-kafka inputs and an explicit data_input_path are left untouched. Co-Authored-By: Claude <noreply@anthropic.com>
_prune_old_dense_versions keeps the lexicographically-newest K versions and deletes the rest, but the just-published version (what current.json points at) can sort outside the newest K: an explicit --version with an older timestamp, or clock rollback after a restart -- the worker's _last_version starts empty and is not seeded from disk, so _make_monotonic_version degrades to the raw clock. Deleting that version left current.json referencing a missing directory. Read the version current.json points at and always spare it from deletion, even when it sorts outside the newest K; a missing or corrupt current.json degrades to pure newest-K retention. Co-Authored-By: Claude <noreply@anthropic.com>
… backstop close() detached the finalizer before joining the worker, and _close_timeout covered only a single task timeout. At close the worker can have an in-flight export plus one pending (latest-wins coalescing caps it at one), so draining could need two task timeouts while _close_timeout allowed one: the join timed out with the worker alive, and the already-detached finalizer left no atexit backstop, so close() returned with a live subprocess publisher after ckpt_manager.close(). Make _close_timeout cover an in-flight plus one pending task, and detach the finalizer only after the worker actually stops so a worker that outlives the close timeout keeps the drain backstop. Co-Authored-By: Claude <noreply@anthropic.com>
…eption _train_and_evaluate only called online_dense_exporter.close() and ckpt_manager.close() on the happy path, after the final save and eval. Any raise in the training loop, on_train_end, final save, or eval skipped them; and the bound self._worker_loop thread target keeps the exporter reachable, so its weakref.finalize backstop does not fire when the function unwinds -- a caller that catches the error leaks the daemon worker and the protected checkpoint and may see a late current.json publish. Wrap the training body (the epoch loop through the final run_eval) in try/finally so both close() calls always run. Co-Authored-By: Claude <noreply@anthropic.com>
Each export attempt writes one <version>.log under dense_hot_export/logs/, but version retention only prunes version dirs (and only on successful exports), never these logs -- a long-running online-learning job accumulates one file/inode per checkpoint including failed exports, and failed attempts never reach the version-retention call at all. Prune logs in the worker's _run_task finally (runs for success, timeout, and failure), keeping the newest K (ONLINE_DENSE_EXPORT_KEEP_LOGS, default 3); version names are timestamps, so name order is chronological. Co-Authored-By: Claude <noreply@anthropic.com>
eric-gecheng
changed the title
dynamicemb features are FG-hashed into a 2**63-1 bucket, so warm-up batches carry ~1e18 ids. On the GPU/DMP path the dynamicemb backend remaps them into [0, max_capacity), but export_dense_model_cpu runs the raw model on CPU where that backend is not active, so ebc_user degrades to F.embedding_bag and its strict range-check rejects the id. The warm-up only materializes lazy-module shapes for FX tracing -- its lookup values are discarded (real weights come from restore_model; the trace uses symbolic proxies) -- so zero sparse_feature values in-place before the warm-up forward; 0 is unconditionally in range. The GPU distributed export path is unaffected. Co-Authored-By: Claude <noreply@anthropic.com>
…dense_hot_export ONLINE_DENSE_EXPORT_DIR was returned verbatim as the export root, so the publish tree landed directly at the serving root without the dense_hot_export leaf the inference processor reads from. resolve_dense_export_root now always appends dense_hot_export, treating the env var as the serving root in both the override and default branches. The manager passes the pre-suffix serving root to the export subprocess so it re-resolves to the same leaf instead of double-nesting dense_hot_export. Co-Authored-By: Claude <noreply@anthropic.com>
tiankongdeguiji
requested changes
Jul 23, 2026
| return flag | ||
|
|
||
|
|
||
| def use_distributed_embedding() -> bool: |
Collaborator
There was a problem hiding this comment.
use_distributed_embedding already in acc/utils.py
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
2 participants
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.
No description provided.