[feat] delta embedding dump upload to PAI Feature Store

requirements/feature_store.txt

@@ -0,0 +1 @@
+feature_store_py @ https://gecheng-rec.oss-accelerate-overseas.aliyuncs.com/software/feature_store_py-2.2.7-py3-none-any.whl

setup.py

@@ -80,5 +80,6 @@ def parse_require_file(fpath):
         "tests": parse_requirements("requirements/test.txt"),
         "gpu": parse_requirements("requirements/cu130.txt"),
         "cpu": parse_requirements("requirements/cpu.txt"),
+        "feature_store": parse_requirements("requirements/feature_store.txt"),
     },
 )

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],
@@ -452,11 +467,21 @@ 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:
+        # Timed dumps synchronize their per-step state, so uneven workers must
+        # fail together instead of silently truncating data 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}")
@@ -748,13 +773,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(
@@ -827,31 +861,65 @@ def train_and_evaluate(
     dist.barrier()
 
     if is_rank_zero:
-        config_util.save_message(
+        config_util.save_pipeline_config_artifact(
             pipeline_config, os.path.join(pipeline_config.model_dir, "pipeline.config")
         )
         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. The durable parquet
+            # outbox remains available for restart. 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
+            # Non-zero ranks have no uploader and can reach this rendezvous while
+            # rank zero drains its background queue. Propagate a late drain/SDK
+            # error uniformly instead of letting only rank zero fail 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,70 @@ message GradClipping {
     optional bool enable_global_grad_clip = 4 [default = false];
 }
 
+message FeatureStoreConfig {
+    // Runtime credentials are read only from ALIBABA_CLOUD_ACCESS_KEY_ID,
+    // ALIBABA_CLOUD_ACCESS_KEY_SECRET, FEATUREDB_USERNAME and FEATUREDB_PASSWORD.
+    reserved 2 to 5;
+    reserved "access_key_id", "access_key_secret", "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 = 6;
+    required string feature_view_name = 7;
+    // Existing FeatureStore entity associated with an automatically created
+    // DynamicEmbedding FeatureView.
+    required string feature_entity_name = 8;
+    // Explicit FeatureDB version for this incremental training run. It must be
+    // provisioned before upload starts; only positive-global-step dumps are
+    // written to this version with MERGE. One project/view/version may have
+    // only one active delta writer.
+    required string version = 9;
+
+
+    // Optional FeatureStore control-plane endpoint override.
+    optional string endpoint = 10;
+    optional string security_token = 11;
+    // Maximum records submitted before each SDK write_flush() completion gate.
+    optional uint32 upload_batch_size = 12 [default = 1000];
+    // Total attempts per process for one step. All retries reuse the version;
+    // each attempt durably reserves a newer monotonic ts range and fully replays
+    // the step so incremental readers cannot miss a partial earlier attempt.
+    optional uint32 max_retries = 13 [default = 3];
+    optional uint32 retry_backoff_secs = 14 [default = 5];
+    // Maximum time rank zero waits for every rank's atomic parquet shard.
+    // With FeatureStore enabled, output_dir must be one stable shared POSIX
+    // filesystem visible to rank zero across restarts and must support atomic
+    // rename, fsync and advisory file locking.
+    optional uint32 shard_wait_timeout_secs = 15 [default = 600];
+    // Maximum time normal training shutdown waits for the uploader to drain.
+    optional uint32 shutdown_timeout_secs = 16 [default = 600];
+    // Apply back-pressure when too many completed dumps await publication.
+    optional uint32 max_pending_steps = 17 [default = 32];
+    optional uint32 poll_interval_secs = 18 [default = 1];
+    // Creation settings used only when feature_view_name does not yet exist.
+    optional uint32 feature_view_ttl_secs = 19 [default = 1296000];
+    optional uint32 feature_view_shard_count = 20 [default = 20];
+    optional uint32 feature_view_replication_count = 21 [default = 1];
+}
+
 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 durable, rank-zero 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.