[feat] delta embedding dump upload to PAI Feature Store by eric-gecheng · Pull Request #603 · alibaba/TorchEasyRec

eric-gecheng · 2026-07-16T12:21:34Z

No description provided.

github-actions · 2026-07-19T08:06:11Z

Static review summary (tests/builds intentionally not run):

Distributed correctness blockers: step-cadence final flushing can drop trailing rows from a shorter rank, and rank-local dump/uploader failures can leave peers hanging.
Scaling risks: minute cadence adds redundant blocking collectives to every batch, whole steps are materialized in Python memory before upload, and the worker repeatedly rescans retained history.
Safety and verification: tighten SDK provenance and endpoint validation, and bind readback samples to the committed shard hashes/generation.
Documentation: the version-provisioning comment contradicts the implementation. Please also add user-facing docs and a runnable config covering installation, credential environment variables (including the security token), shared-POSIX outbox requirements, CUDA/sharding and equal-batch constraints, restart/single-writer behavior, and the readback command.

Detailed actionable findings are inline.

… dumps With dump_interval_steps and check_all_workers_data_status=False, ranks could exhaust their dataloaders at different steps (e.g. 49 and 50). final_dump all-reduces the final step with MAX, so the rank that stopped at 49 adopted 50 and hit the boundary-step skip even though it never ran maybe_dump(50); its trailing tracked rows were never written and were silently lost. requires_synced_dataloader_exhaustion now returns True for any multi-rank cadence, not only timed dumps, so the train pipeline checks all workers' data status and fails loudly on uneven exhaustion instead of letting a lagging rank skip its final delta. This keeps every rank participating in every boundary dump, the invariant the boundary-step skip in final_dump relies on. Adds a 49/50 regression test asserting the uneven stop is rejected, plus a contract unit test. Co-Authored-By: Claude <noreply@anthropic.com>

…p cadence The uploader-failure rendezvous in _should_dump only ran for timed (minute) dumps. In the default step cadence with FeatureStore upload, rank zero observes an async upload failure immediately through uploader.check_error() while peers only see the shared error marker through a throttled poll, so rank zero raised alone and its peers marched on into the next training collective and hung. Boundary dump failures had the same hole: _raise_if_any_timed_dump_failed also gated its failure-bit all-reduce on the timed cadence, so any rank-local boundary dump failure (e.g. a shard write error on one node) raised on one rank only. Add _requires_dump_state_rendezvous() and use it for the _should_dump rendezvous and for selecting check_upload_error=False in maybe_dump, so distributed FeatureStore dumps of either cadence synchronize the failure bit before any rank raises. Also widen the interval dump failure rendezvous (renamed to _raise_if_any_interval_dump_failed) to every multi-rank cadence, mirroring _raise_if_any_final_dump_failed. Co-Authored-By: Claude <noreply@anthropic.com>

…path Multi-rank dumps all-reduced the timer vote and uploader-failure bit, with host-synchronizing .item() reads, on every training step, while the synced-exhaustion pipeline already pays one blocking availability collective per batch: two blocking collectives per step for dumps that fire minutes apart. The rendezvous is now launched asynchronously and consumed on the next maybe_dump call. NCCL orders it ahead of the next step's communication on the same process group, so it has completed by the time the next call reads it back and the hot path pays only an async launch plus a non-blocking read; the decision lands at most one step late, which is immaterial for minute-scale intervals. Timed votes arm an in-flight guard so the still-elapsed deadline cannot fire a second launch before the dump reschedules from its completion time. Step-boundary decisions stay deterministic and on their exact boundary step; only the failure bit is pipelined there. Coalescing the bits into the pipeline's availability collective would also remove the per-step launch, but couples the dumper into pipeline prefetch internals and moves its raise/dump control flow across components; the overlap keeps the failure domain inside the dumper while taking the blocking synchronization off the hot path. Co-Authored-By: Claude <noreply@anthropic.com>

…led dedup _load_records read each shard with pq.read_table() and to_pylist(), materializing every embedding element as Python objects, then kept the deduplicated rows plus a sorted copy in memory for the whole upload. Peak memory therefore scaled with a single dump's size, so one large delta could exhaust rank-zero training memory despite max_pending_steps bounding only the number of pending steps, not their size. Shards are now streamed with ParquetFile.iter_batches and validated with vectorized per-batch checks, and dedup state lives in a throwaway SQLite database (bounded 64MB page cache, journal off: the store is disposable because a crashed step replays from its durable shard snapshot). Uploads stream from an ORDER BY cursor in upload_batch_size windows, and retries re-iterate the same deterministic ordering, so peak memory depends only on one Parquet batch, the page cache, and one in-flight upload batch. Co-Authored-By: Claude <noreply@anthropic.com>

…ta dump torchrun terminates all workers when any rank exits with an uncaught exception, so the cross-rank failure-bit all_reduce propagation is unnecessary. Remove the error-bit from the timed-dump rendezvous (keep only the dump-vote OR-reduce), remove _raise_if_any_interval_dump_failed and _raise_if_any_final_dump_failed, and let dump exceptions propagate directly. Also replace getattr(..., default) with direct attribute access for fields unconditionally initialized in __init__, and remove the requires_synced_dataloader_exhaustion property and its fail_on_uneven_data usage in main.py.

…dentials Remove security_token from FeatureStoreConfig proto (now reserved) and remove all credential fields from FeatureStoreUploadSettings. Cloud credentials (AK/SK/STS) are now resolved at runtime through the alibabacloud_credentials default provider chain, enabling automatic STS token refresh for long-running training jobs. FeatureDB credentials remain env-var based (FEATUREDB_USERNAME/FEATUREDB_PASSWORD). Since the config no longer contains secrets, revert the save_pipeline_config_artifact sanitization and use save_message directly.

…heel URL

Rank shards no longer funnel through parquet files on a shared filesystem to a rank-zero uploader: each rank streams its own in-memory delta table to the SDK, coordinating only view creation and startup via a barrier. Per-rank monotonic timestamps suffice because sharding is fixed within a process, and data-parallel duplicate MERGE writes are idempotent. Local parquet files are now opt-in via feature_store_config.retain_local_dump, which the offline readback checker requires.

The TorchRec delta tracker keys its bookkeeping by raw table name, so a table name legitimately reused by several modules (an EC sequence path and an EBC deep path, or two data-group impls) collapsed into one entry and the dump either raised or published one owner's weights under the other's identity. Identities, weights and dynamicemb lookups are now keyed by (module_fqn, table_name), and the merged id set is fanned out to every owner so each publishes its own true rows under its own role-suffixed serving name. Same-role duplicates still raise for FeatureStore upload because the export naming contract cannot address them distinctly.

Minute-cadence dumps all-reduced a rank-zero clock vote on every training step so all ranks dumped together. With per-rank FeatureStore uploads and rank-local tracker windows, ranks no longer need to dump on the same step, so each rank now checks its own clock against a shared fixed-rate deadline sequence armed at start(), removing the hot-path collective entirely. Multi-rank timed local parquet sets are best-effort step-aligned.

… test fb53143 removed the credential fields from FeatureStoreConfig (3062112 renumbered the proto) and reverted the export-time sanitization to plain save_message, but this test -- added earlier by 70ee3b8 when sanitization existed -- still set feature_store_config.security_token and asserted it was scrubbed from the exported pipeline.config, so it fails on the CPU CI lane with AttributeError: Protocol message FeatureStoreConfig has no "security_token" field. Drop the dead credential assertions, keep the override and FeatureStoreConfig round-trip coverage the test still provides, and rename it to describe what it now verifies. Co-Authored-By: Claude <noreply@anthropic.com>

eeaace1 removed requires_synced_dataloader_exhaustion and the main.py fail_on_uneven_data wiring as a side item of the error-bit rendezvous removal, but final_dump still MAX-all-reduces the final step across ranks and skips steps that fall on a dump boundary. Without the enforcement, a rank whose dataloader stops before the synced step (e.g. 49 batches vs 50) never ran maybe_dump on that boundary, adopts the boundary skip, and silently drops every trailing tracked row -- the exact data loss the regression test left behind by that commit guards against, which now fails with "RuntimeError not raised" on the GPU CI lane. Restore the property (True whenever world_size > 1, either cadence) and wire check_all_workers_data_status/fail_on_uneven_data from it in _train_and_evaluate, so uneven exhaustion fails loudly on every rank instead of dropping a lagging rank's delta; also restore the property's contract tests and the pipeline kwargs in the regression worker, and reword the final_dump call-site comment the removal left false. The error-bit removal itself stays. Co-Authored-By: Claude <noreply@anthropic.com>

877ece7's per-owner fan-out calls _lookup_embeddings for every owner of a tracker-merged table name, but _collect_table_weights and _collect_dynamic_modules only gather this rank's local shards. An owner whose shard lives on another rank -- e.g. the table_wise EBC copy of item_id_emb, placed on rank 0 only, queried on rank 1 in the dynamicemb multi-GPU integration test -- is absent from both maps and raised KeyError, crashing the whole dump. Pre-filter owners this rank cannot resolve and let the hosting rank dump their rows, restoring the pre-877ece7 robustness (the model-wide collection then only ever found local shards) while keeping per-owner fan-out and identities for locally resolvable owners. The KeyError in _lookup_embeddings stays as a defensive net for genuine internal inconsistency. Co-Authored-By: Claude <noreply@anthropic.com>

The dumper probed torchrec internals and proto fields with getattr-defaults and an _int_attr helper, hiding which attributes are type-guaranteed versus optional. Use direct access for fields the torchrec dataclasses always expose (ParameterSharding.ranks/sharding_spec, Shard.placement, shards[0].metadata, EmbeddingBagConfig.num_embeddings/embedding_dim/feature_names, the isinstance-narrowed _table_name_to_config) and hasattr guards for genuinely optional attributes on arbitrary modules (_lookups, config/_config, embedding_tables, module_sharding_plan, post_*_tracker_fn). _int_attr and _feature_config_name are inlined, _has_proto_field simplified. Semantics are unchanged: hasattr plus direct access is equivalent to the prior getattr defaults, and tests are updated to use faithful, real-type-shaped configs. The only remaining getattr accesses a proto oneof field by dynamic name (feature_config.WhichOneof -> getattr), which has no getattr-free equivalent. Co-Authored-By: Claude <noreply@anthropic.com>

`identity` was overloaded: it denoted both a `(role, table_name)` dict key (the `table_identities` param / `name_by_identity` local in the contract and delta dump) and a `SparseEmbeddingIdentity` instance, confusing readers. Rename the tuple-keyed identifiers to `role_table_pairs` / `name_by_role_table` so the key shape is explicit, and note that `identity_by_owner`'s key is `(module_fqn, table_name)` -- distinct from the contract's `(role, table_name)`. The `SparseEmbeddingIdentity` dataclass and its instance locals stay (the legitimate singular). Pure rename; all callers are positional, so behavior and the exported embedding_name are unchanged. Co-Authored-By: Claude <noreply@anthropic.com>

export_util duplicated the EC/EBC role-suffix naming algorithm that sparse_embedding_contract already owns, risking drift between the exported feat_meta["embedding_name"] (read by serving) and the delta-dump name. Drop the private role constants and the _build_sparse_export_name_map / _resolve_sparse_export_name copies; import the shared build_sparse_embedding_name_map / resolve_sparse_embedding_name instead. _build_sparse_export_name_map stays as a thin adapter (two config lists -> (role, name) pairs) and the four resolve call sites call the shared function directly. Output is byte-identical for every realistic config (differential fuzz over 610 cases incl. shared-table and raw-candidate collisions); the only behavior change is raising on an empty table name, which real configs never produce (feature.py `embedding_name or f"{name}_emb"` fallback). This addresses the instinct behind review alibaba#2 without the unsafe FQN switch: embedding_name is a cross-system serving key (export + delta dump + serving share it) and the export path has no module FQN, so a delta-side FQN would diverge and orphan every delta row. Co-Authored-By: Claude <noreply@anthropic.com>

Remove the endpoint allowlist and control-plane FeatureView metadata validation that review flagged as over-engineered: endpoint selection is an explicit operator action owned by the FeatureStore SDK, and view schema compatibility is already enforced on the data-plane writer. Also drop the redundant invalid_dimensions check and validate_feature_store_config helper, which duplicate validation performed when the uploader settings are built from the proto. Co-Authored-By: Claude <noreply@anthropic.com>

Submissions arrive in training order, so the pending queue is a plain in-process FIFO: replace the step-keyed dict (with min() and a setdefault dedup guard) with a collections.deque. Also demote the per-rank, per-dump log lines to DEBUG, since their volume scaled with world size and dump frequency; periodic upload progress and the final completion stay at INFO. Co-Authored-By: Claude <noreply@anthropic.com>

Each dump traversed every tracked module to rebuild the owner-to-weight and owner-to-dynamic-module maps, even though those maps never change during a run: the owning modules and their local shard tensors keep a stable identity and training updates the storage in place. Discover both once on the first dump and reuse the cached references, so later dumps read the current rows through the same handles without re-traversing. Co-Authored-By: Claude <noreply@anthropic.com>

…ta dump Multi-rank dumps already force synced dataloader exhaustion, so every rank reaches the same final step before final_dump(); the MAX all-reduce that re-aligned it was redundant and is removed in favor of the local step. The tracker read rollback is also removed: get_unique() deletes the consumed rows, so restoring only the consumer cursor could never recover them, and under the best-effort contract a dump failure simply propagates and torchrun tears down the job. Co-Authored-By: Claude <noreply@anthropic.com>

Stop catching training failures to close the dumper and stop rendezvousing rank-local startup/shutdown errors across workers via all-reduce. torchrun terminates the whole job when any rank exits with an uncaught exception, and pending in-memory deltas are intentionally discarded on restart from the latest checkpoint, so the failure propagation added nothing. The dumper is now started directly and closed only after training succeeds, to drain background uploads; _raise_if_any_worker_failed is removed entirely. Co-Authored-By: Claude <noreply@anthropic.com>

The barrier sequencing rank-zero view creation ahead of non-primary opens lived in DeltaEmbeddingDumper.start, leaking the uploader's manage_remote_view startup protocol to the caller. Move the create-barrier-open sequence into FeatureStoreDeltaUploader.start so the dumper only delegates and arms the timed cadence; the failure rendezvous (rank-zero create error still joins the barrier before re-raising) is preserved verbatim, with world_size plumbed through the uploader ctor. Co-Authored-By: Claude <noreply@anthropic.com>

…lient FeatureStoreClient construction and credential resolution were inlined in _get_view and reached from __init__ via client_factory/credentials_client ctor params, leaking an SDK-specific concern into the constructor. Move both into a private _create_client() seam and drop the two ctor params; tests now inject a fake by patching that one method. The test_mode kwarg hygiene check moves to a focused unit test asserting _create_client forwards exactly the fixed credential allowlist. Co-Authored-By: Claude <noreply@anthropic.com>

Reconciles master alibaba#604 (online dense export), alibaba#607 (track delta by owner-qualified FQN), alibaba#608 (export sparse by table FQN), and alibaba#609 (drop INPUT_TILE mapping) with the branch's FeatureStore delta-dump work. Keeps master alibaba#607's FQN-keyed delta tracking (local ModelDeltaTracker subclass, source/table_fqn schema) and alibaba#608's FQN sparse-export naming as the canonical base, and re-grafts the branch's FeatureStore upload, minute-cadence, and synced-dataloader-exhaustion onto alibaba#607's dumper. The uploader reads table_fqn and publishes remap_input_tile_user_key(table_fqn) as the FeatureStore embedding_name, aligning delta upload with alibaba#608's export naming and serving. Synced-exhaustion is kept so alibaba#607's _sync_final_step boundary-skip does not drop a shorter-exhausted rank's trailing delta (SEANQ). alibaba#604's online dense export and alibaba#609's INPUT_TILE removal are adopted as-is; main.py keeps the synced-exhaustion gating and FeatureStore start/close hooks. Co-Authored-By: Claude <noreply@anthropic.com>

Drop feature_entity_name from FeatureStoreConfig; the rank-zero uploader now provisions a default entity (default_dynemb_entity) on demand when it creates the DynamicEmbedding feature view. Co-Authored-By: Claude <noreply@anthropic.com>

Add FeatureStoreConfig.upload_format (default ARROW) to stream delta batches through the SDK's write_features_arrow() columnar IPC path, avoiding the JSON path's per-row dict construction and embedding deep-copy. The legacy JSON write_features() path stays behind upload_format="JSON"; both reuse a shared _validate_delta_batch and the existing MERGE/ts/window/flush/retry semantics. Two pre-existing defects surfaced and are fixed in passing: streaming now materializes the actual to_batches() list (ceil(total/batch_size) undercounted multi-FQN multi-chunk tables, letting a stuck clock reuse a prior step's ts so Next-Ts readers miss updates), and the NaN/Inf check is bounded to each batch's value range instead of re-scanning the whole chunk buffer per batch. Co-Authored-By: Claude <noreply@anthropic.com>

eric-gecheng added 14 commits July 16, 2026 11:22

basic impl of delta dump upload to feature store

c49f455

Merge branch 'master' into feat/delta_export_optimize

b697ee5

parallelize FeatureStore delta uploads

01481c1

support minute-based delta embedding dumps

e4cffe7

add FeatureStore delta readback checker

953a70c

put feature store credentials into envs

bb4f1d3

remove redundant delta dump columns

035fc1a

export zeros for missing dynamic embedding keys

2faabcc

[bugfix] replay restored FeatureStore deltas by generation

14ba85d

[bugfix] synchronize timed delta dump workers

8eb37f2

[bugfix] use VPC routing for FeatureStore uploads

0183ec0

[bugfix] sanitize distributed export configs

70ee3b8

[bugfix] synchronize timed delta dump failures across ranks

a611243

[bugfix] reject uneven batches during timed delta dump

cc03f25

eric-gecheng added the codex-review Let Codex Review label Jul 19, 2026

github-actions Bot removed the codex-review Let Codex Review label Jul 19, 2026