[feat] delta embedding dump upload to PAI Feature Store (alibaba#603)

Co-authored-by: Claude <noreply@anthropic.com>

Author: eric-gecheng

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

@@ -606,11 +606,9 @@ 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
+            # (all ranks run the same step count, so every rank participated in
+            # those dumps and reaches the same final step).
             delta_embedding_dumper.final_dump(i_step)
 
         _log_train(
@@ -900,6 +898,8 @@ def train_and_evaluate(
         with open(os.path.join(pipeline_config.model_dir, "version"), "w") as f:
             f.write(tzrec_version + "\n")
 
+    if delta_embedding_dumper is not None:
+        delta_embedding_dumper.start()
     # 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(
@@ -922,6 +922,12 @@ def train_and_evaluate(
         dense_ema=dense_ema,
         export_config=pipeline_config.export_config,
     )
+    # Drain background uploads only after training succeeds. A training failure
+    # terminates the whole job (torchrun tears down every rank) and pending
+    # in-memory deltas are intentionally abandoned: the restarted run re-dumps
+    # from the latest checkpoint, so there is nothing to roll back or undo.
+    if delta_embedding_dumper is not None:
+        delta_embedding_dumper.close()
     if is_local_rank_zero:
         logger.info("Train and Evaluate Finished.")
 

tzrec/protos/train.proto

@@ -29,16 +29,71 @@ 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,
+    // together with a default FeatureStore entity that is auto-created on demand.
+    required string project_name = 2;
+    required string feature_view_name = 3;
+    // FeatureDB version for this incremental training run.
+    required string version = 5;
+
+
+    // Optional FeatureStore control-plane endpoint override, passed directly to
+    // the FeatureStore SDK, which owns endpoint handling and connection errors.
+    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];
+    reserved 4, 16;
+    reserved "feature_entity_name", "allow_custom_endpoint";
+    // 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];
+    // Wire format for delta upload. "ARROW" (default) streams a columnar Arrow
+    // IPC batch through write_features_arrow(), avoiding the JSON path's per-row
+    // dict construction and embedding deep-copy; "JSON" keeps the legacy
+    // write_features() per-row payload. Both paths use MERGE write_mode.
+    optional string upload_format = 18 [default = "ARROW"];
+}
+
 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.