[feat] delta embedding dump upload to PAI Feature Store

requirements/runtime.txt

@@ -3,6 +3,7 @@ anytree
 common_io @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/common_io-0.4.1%2Btunnel-py2.py3-none-any.whl
 confluent-kafka
 fbgemm-gpu==1.7.0
+feature_store_py @ https://feature-store-py.oss-cn-beijing.aliyuncs.com/package/feature_store_py-2.2.7-py3-none-any.whl
 fsspec
 graphlearn @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/graphlearn/graphlearn-1.3.8-cp312-cp312-linux_x86_64.whl ; python_version=="3.12"
 graphlearn @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/graphlearn/graphlearn-1.3.8-cp311-cp311-linux_x86_64.whl ; python_version=="3.11"

tzrec/main.py

@@ -129,6 +129,21 @@ def _get_sampler_type(data_config: DataConfig) -> Optional[str]:
     return sampler_type
 
 
+def _raise_if_any_worker_failed(
+    local_error: Optional[BaseException], device: torch.device, action: str
+) -> None:
+    """Make rank-local initialization failures visible to every worker."""
+    any_failed = local_error is not None
+    if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
+        failed = torch.tensor([int(any_failed)], dtype=torch.int32, device=device)
+        dist.all_reduce(failed, op=dist.ReduceOp.MAX)
+        any_failed = bool(failed.item())
+    if local_error is not None:
+        raise local_error.with_traceback(local_error.__traceback__)
+    if any_failed:
+        raise RuntimeError(f"{action} failed on another distributed worker")
+
+
 def _create_model(
     model_config: ModelConfig,
     features: List[BaseFeature],
@@ -454,11 +469,22 @@ def run_eval(step: int, epoch: int) -> None:
 
     # this rank's last consumed event-time, reused by the epoch / final saves
     data_timestamp = -1.0
+    require_equal_train_batches = (
+        delta_embedding_dumper is not None
+        and delta_embedding_dumper.requires_synced_dataloader_exhaustion
+    )
+    sync_train_data_exhaustion = (
+        check_all_workers_data_status or require_equal_train_batches
+    )
     for i_epoch in epoch_iter:
+        # Multi-rank dumps synchronize their per-step state, so uneven workers
+        # must fail together instead of silently dropping a lagging rank's
+        # trailing delta rows or leaving collectives.
         pipeline = create_train_pipeline(
             model,
             optimizer,
-            check_all_workers_data_status=check_all_workers_data_status,
+            check_all_workers_data_status=sync_train_data_exhaustion,
+            fail_on_uneven_data=require_equal_train_batches,
         )
         if plogger is not None:
             plogger.set_description(f"Training Epoch {i_epoch}")
@@ -560,11 +586,10 @@ def run_eval(step: int, epoch: int) -> None:
     _model.on_train_end()
     if delta_embedding_dumper is not None:
         # Flush the trailing partial interval before the final checkpoint.
-        # final_dump skips dump-boundary steps already written by maybe_dump,
-        # so it never overwrites their shards with an empty file. Ranks can
-        # reach here at different i_step (independent dataloader exhaustion with
-        # check_all_workers_data_status=False), so final_dump all-reduces the
-        # step across ranks to keep one complete shard set per step dir.
+        # final_dump skips dump-boundary steps already written by maybe_dump
+        # (multi-rank dumps force synced exhaustion, so every rank participated
+        # in those dumps); the MAX all-reduce in final_dump keeps one complete
+        # shard set per step dir as a defensive guard.
         delta_embedding_dumper.final_dump(i_step)
 
     _log_train(
@@ -750,13 +775,22 @@ def train_and_evaluate(
         plan=plan,
     )
     delta_embedding_dumper = None
+    delta_embedding_dumper_error: Optional[BaseException] = None
     if enable_delta_embedding_dump:
-        delta_embedding_dumper = DeltaEmbeddingDumper(
-            model,
-            train_config.delta_embedding_dump_config,
-            pipeline_config.model_dir,
+        try:
+            delta_embedding_dumper = DeltaEmbeddingDumper(
+                model,
+                train_config.delta_embedding_dump_config,
+                pipeline_config.model_dir,
+                device,
+                pipeline_config.feature_configs,
+            )
+        except BaseException as exc:
+            delta_embedding_dumper_error = exc
+        _raise_if_any_worker_failed(
+            delta_embedding_dumper_error,
             device,
-            pipeline_config.feature_configs,
+            "delta embedding dumper initialization",
         )
 
     dense_optim_cls, dense_optim_kwargs = optimizer_builder.create_dense_optimizer(
@@ -835,25 +869,60 @@ def train_and_evaluate(
         with open(os.path.join(pipeline_config.model_dir, "version"), "w") as f:
             f.write(tzrec_version + "\n")
 
-    # when slice batch by sample cost, data on all workers may not be balanced
-    check_all_workers_data_status = data_config.HasField("batch_cost_size")
-    _train_and_evaluate(
-        model,
-        optimizer,
-        train_dataloader,
-        eval_dataloader,
-        [sparse_lr, dense_lr, *part_lrs],
-        pipeline_config.model_dir,
-        train_config=train_config,
-        eval_config=pipeline_config.eval_config,
-        ckpt_manager=ckpt_manager,
-        skip_steps=skip_steps,
-        ckpt_path=ckpt_path,
-        check_all_workers_data_status=check_all_workers_data_status,
-        ignore_restore_optimizer=ignore_restore_optimizer,
-        dataloader_state=dataloader_state,
-        delta_embedding_dumper=delta_embedding_dumper,
-    )
+    try:
+        if delta_embedding_dumper is not None:
+            delta_embedding_dumper_start_error: Optional[BaseException] = None
+            try:
+                delta_embedding_dumper.start()
+            except BaseException as exc:
+                delta_embedding_dumper_start_error = exc
+            _raise_if_any_worker_failed(
+                delta_embedding_dumper_start_error,
+                device,
+                "FeatureStore delta uploader startup",
+            )
+        # when slice batch by sample cost, data on all workers may not be balanced
+        check_all_workers_data_status = data_config.HasField("batch_cost_size")
+        _train_and_evaluate(
+            model,
+            optimizer,
+            train_dataloader,
+            eval_dataloader,
+            [sparse_lr, dense_lr, *part_lrs],
+            pipeline_config.model_dir,
+            train_config=train_config,
+            eval_config=pipeline_config.eval_config,
+            ckpt_manager=ckpt_manager,
+            skip_steps=skip_steps,
+            ckpt_path=ckpt_path,
+            check_all_workers_data_status=check_all_workers_data_status,
+            ignore_restore_optimizer=ignore_restore_optimizer,
+            dataloader_state=dataloader_state,
+            delta_embedding_dumper=delta_embedding_dumper,
+        )
+    except BaseException:
+        if delta_embedding_dumper is not None:
+            # Keep the original training error primary. Pending in-memory
+            # deltas are abandoned; the restarted run re-dumps from the latest
+            # checkpoint. Do not wait up to the normal upload-drain timeout
+            # while unwinding a training failure.
+            delta_embedding_dumper.close(raise_on_error=False, drain=False)
+        raise
+    else:
+        if delta_embedding_dumper is not None:
+            delta_embedding_dumper_close_error: Optional[BaseException] = None
+            try:
+                delta_embedding_dumper.close()
+            except BaseException as exc:
+                delta_embedding_dumper_close_error = exc
+            # Each rank drains its own background queue at a different pace.
+            # Propagate a late drain/SDK error uniformly instead of letting a
+            # single rank fail alone at shutdown.
+            _raise_if_any_worker_failed(
+                delta_embedding_dumper_close_error,
+                device,
+                "FeatureStore delta uploader shutdown",
+            )
     if is_local_rank_zero:
         logger.info("Train and Evaluate Finished.")
 

tzrec/protos/train.proto

@@ -29,16 +29,74 @@ message GradClipping {
     optional bool enable_global_grad_clip = 4 [default = false];
 }
 
+message FeatureStoreConfig {
+    // Cloud credentials (AK/SK/STS) are resolved at runtime through the
+    // Alibaba Cloud default credential provider chain (alibabacloud_credentials).
+    // FeatureDB credentials are read from FEATUREDB_USERNAME/FEATUREDB_PASSWORD.
+    // FeatureStore control-plane region. An explicitly empty value falls back
+    // to ALIBABA_CLOUD_REGION at runtime.
+    required string region = 1;
+    // Existing FeatureStore project name and target DynamicEmbedding FeatureView
+    // name. The view is validated at startup and created when it does not exist.
+    required string project_name = 2;
+    required string feature_view_name = 3;
+    // Existing FeatureStore entity associated with an automatically created
+    // DynamicEmbedding FeatureView.
+    required string feature_entity_name = 4;
+    // FeatureDB version for this incremental training run.
+    required string version = 5;
+
+
+    // Optional FeatureStore control-plane endpoint override. Cloud and
+    // FeatureDB credentials are sent to this host, so it must be a trusted
+    // *.aliyuncs.com FeatureStore endpoint (an explicit https:// scheme is
+    // required when a scheme is given) without userinfo, port, path, query or
+    // fragment components, unless allow_custom_endpoint is set.
+    optional string endpoint = 6;
+    // Maximum records submitted before each SDK write_flush() completion gate.
+    optional uint32 upload_batch_size = 7 [default = 1000];
+    // Total attempts per rank for one step. All retries reuse the version;
+    // each attempt reserves a newer monotonic ts range and fully replays the
+    // rank's delta so incremental readers cannot miss a partial earlier attempt.
+    optional uint32 max_retries = 8 [default = 3];
+    optional uint32 retry_backoff_secs = 9 [default = 5];
+    // Maximum time normal training shutdown waits for the uploader to drain.
+    optional uint32 shutdown_timeout_secs = 10 [default = 600];
+    // Apply back-pressure when too many completed dumps await upload; bounds
+    // the per-rank memory retained by pending in-memory deltas.
+    optional uint32 max_pending_steps = 11 [default = 32];
+    optional uint32 poll_interval_secs = 12 [default = 1];
+    // Creation settings used only when feature_view_name does not yet exist.
+    optional uint32 feature_view_ttl_secs = 13 [default = 1296000];
+    optional uint32 feature_view_shard_count = 14 [default = 20];
+    optional uint32 feature_view_replication_count = 15 [default = 1];
+    // Explicit opt-in allowing endpoint to name a non-*.aliyuncs.com host
+    // (e.g. a vetted private FeatureStore deployment with a custom port).
+    // Credentials are sent to that host; structural endpoint checks (HTTPS
+    // scheme, no userinfo/path/query/fragment) still apply.
+    optional bool allow_custom_endpoint = 16 [default = false];
+    // Also write the local per-rank delta parquet files while uploading to
+    // FeatureStore (e.g. for the offline readback checker); by default the
+    // delta is handed to the uploader in memory and never touches disk.
+    optional bool retain_local_dump = 17 [default = false];
+}
+
 message DeltaEmbeddingDumpConfig {
     // MC/ZCH features are not supported; use dynamicemb for delta dump.
-    // dump touched ids and their latest embedding every N training steps. Larger
-    // intervals retain a longer id window in memory; auto compaction reduces
+    // Dump touched ids and their latest embedding every N training steps. Do not
+    // set this together with dump_interval_minutes.
+    // Larger intervals retain a longer id window in memory; auto compaction reduces
     // per-batch tensor buildup but unique ids still scale with the interval.
     optional uint32 dump_interval_steps = 1 [default = 1000];
     // output directory. default is ${model_dir}/delta_embedding_dump
     optional string output_dir = 2;
     // parquet file prefix
     optional string file_prefix = 3 [default = "delta_embedding"];
+    // Presence enables best-effort per-rank background upload to FeatureStore.
+    optional FeatureStoreConfig feature_store_config = 4;
+    // Dump after this many elapsed minutes. The timer starts when training starts.
+    // Do not set this together with dump_interval_steps.
+    optional uint32 dump_interval_minutes = 5;
 }
 
 message TrainConfig {

tzrec/tools/feature_store/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) 2026, Alibaba Group;
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#    http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.