From c49f4558c828fed8c23801e8723bc36b4c233fc5 Mon Sep 17 00:00:00 2001 From: gecheng Date: Thu, 16 Jul 2026 11:22:11 +0800 Subject: [PATCH 01/13] basic impl of delta dump upload to feature store --- requirements/feature_store.txt | 1 + setup.py | 1 + tzrec/main.py | 108 +- tzrec/protos/train.proto | 55 + tzrec/utils/config_util.py | 68 +- tzrec/utils/config_util_test.py | 61 + tzrec/utils/delta_embedding_dump.py | 483 ++++- tzrec/utils/delta_embedding_dump_test.py | 224 +- tzrec/utils/export_util.py | 2 +- tzrec/utils/feature_store_delta_uploader.py | 1872 +++++++++++++++++ .../feature_store_delta_uploader_test.py | 1573 ++++++++++++++ tzrec/utils/sparse_embedding_contract.py | 105 + tzrec/utils/sparse_embedding_contract_test.py | 66 + 13 files changed, 4546 insertions(+), 73 deletions(-) create mode 100644 requirements/feature_store.txt create mode 100644 tzrec/utils/feature_store_delta_uploader.py create mode 100644 tzrec/utils/feature_store_delta_uploader_test.py create mode 100644 tzrec/utils/sparse_embedding_contract.py create mode 100644 tzrec/utils/sparse_embedding_contract_test.py diff --git a/requirements/feature_store.txt b/requirements/feature_store.txt new file mode 100644 index 000000000..d4390de4a --- /dev/null +++ b/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 diff --git a/setup.py b/setup.py index 88f82b111..9ca775554 100644 --- a/setup.py +++ b/setup.py @@ -80,5 +80,6 @@ def parse_require_file(fpath): "tests": parse_requirements("requirements/test.txt"), "gpu": parse_requirements("requirements/cu129.txt"), "cpu": parse_requirements("requirements/cpu.txt"), + "feature_store": parse_requirements("requirements/feature_store.txt"), }, ) diff --git a/tzrec/main.py b/tzrec/main.py index 7ebf34983..a5a901a8e 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -126,6 +126,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], @@ -745,13 +760,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( @@ -824,31 +848,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.") diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index 052103032..170fad703 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -29,6 +29,59 @@ message GradClipping { optional bool enable_global_grad_clip = 4 [default = false]; } +message FeatureStoreConfig { + // FeatureStore control-plane region. An explicitly empty value falls back + // to ALIBABA_CLOUD_REGION at runtime. + required string region = 1; + // Runtime credentials. Explicit values take precedence over the standard + // ALIBABA_CLOUD_* environment variables. Persisted model artifacts replace + // these required secret values with empty strings. + required string access_key_id = 2; + required string access_key_secret = 3; + // FeatureDB data-plane credential pair. Explicit values take + // precedence over FEATUREDB_USERNAME and FEATUREDB_PASSWORD. + required string featuredb_username = 4; + required string featuredb_password = 5; + // 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 @@ -39,6 +92,8 @@ message DeltaEmbeddingDumpConfig { 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; } message TrainConfig { diff --git a/tzrec/utils/config_util.py b/tzrec/utils/config_util.py index cb7f2ff4c..d62d72b76 100644 --- a/tzrec/utils/config_util.py +++ b/tzrec/utils/config_util.py @@ -17,7 +17,7 @@ from google.protobuf import json_format, text_format from google.protobuf.message import Message -from tzrec.protos import data_pb2, pipeline_pb2 +from tzrec.protos import data_pb2, pipeline_pb2, train_pb2 from tzrec.protos.data_pb2 import FgMode from tzrec.utils.logging_util import logger @@ -63,6 +63,72 @@ def save_message(message: Message, filepath: str) -> None: f.write(pbtxt) +_FEATURE_STORE_ARTIFACT_FIELDS = ( + "region", + "endpoint", + "project_name", + "feature_entity_name", + "feature_view_name", + "feature_view_ttl_secs", + "feature_view_shard_count", + "feature_view_replication_count", + "version", + "upload_batch_size", + "max_retries", + "retry_backoff_secs", + "shard_wait_timeout_secs", + "shutdown_timeout_secs", + "max_pending_steps", + "poll_interval_secs", +) + +_FEATURE_STORE_ARTIFACT_REQUIRED_SECRETS = ( + "access_key_id", + "access_key_secret", + "featuredb_username", + "featuredb_password", +) + + +def sanitize_pipeline_config_for_artifact( + pipeline_config: pipeline_pb2.EasyRecConfig, +) -> pipeline_pb2.EasyRecConfig: + """Return an artifact-safe config without FeatureStore credentials. + + The runtime config remains unchanged. FeatureStore public routing and + version fields are copied through an allowlist, so newly introduced fields + do not silently become persistent secrets. Required credential fields are + set to explicit empty strings to keep the protobuf initialized while still + allowing the runtime environment fallbacks. + """ + sanitized = pipeline_pb2.EasyRecConfig() + sanitized.CopyFrom(pipeline_config) + train_config = sanitized.train_config + if not train_config.HasField("delta_embedding_dump_config"): + return sanitized + dump_config = train_config.delta_embedding_dump_config + if not dump_config.HasField("feature_store_config"): + return sanitized + + runtime_config = dump_config.feature_store_config + public_config = train_pb2.FeatureStoreConfig() + for field_name in _FEATURE_STORE_ARTIFACT_FIELDS: + if runtime_config.HasField(field_name): + setattr(public_config, field_name, getattr(runtime_config, field_name)) + for field_name in _FEATURE_STORE_ARTIFACT_REQUIRED_SECRETS: + setattr(public_config, field_name, "") + dump_config.ClearField("feature_store_config") + dump_config.feature_store_config.CopyFrom(public_config) + return sanitized + + +def save_pipeline_config_artifact( + pipeline_config: pipeline_pb2.EasyRecConfig, filepath: str +) -> None: + """Persist a pipeline config after removing runtime-only credentials.""" + save_message(sanitize_pipeline_config_for_artifact(pipeline_config), filepath) + + def config_to_kwargs(config: Message) -> Dict[str, Any]: """Convert a message to a config dict.""" return json_format.MessageToDict( diff --git a/tzrec/utils/config_util_test.py b/tzrec/utils/config_util_test.py index ab29f9d5d..be6182049 100644 --- a/tzrec/utils/config_util_test.py +++ b/tzrec/utils/config_util_test.py @@ -9,6 +9,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +import tempfile import unittest from tzrec.utils import config_util @@ -40,6 +42,65 @@ def test_edit_config(self): ) self.assertEqual(pipeline_config.feature_configs[4].id_feature.num_buckets, 3) + def test_pipeline_artifact_redacts_feature_store_credentials(self): + pipeline_config = config_util.load_pipeline_config( + "examples/multi_tower_taobao.config" + ) + dump_config = pipeline_config.train_config.delta_embedding_dump_config + feature_store_config = dump_config.feature_store_config + feature_store_config.region = "cn-test" + feature_store_config.endpoint = "feature-store.example" + feature_store_config.project_name = "project_a" + feature_store_config.feature_entity_name = "embedding_entity" + feature_store_config.feature_view_name = "shared_embeddings" + feature_store_config.feature_view_ttl_secs = 86400 + feature_store_config.feature_view_shard_count = 4 + feature_store_config.feature_view_replication_count = 2 + feature_store_config.version = "model_a@export_1" + feature_store_config.access_key_id = "SECRET_AK_ID" + feature_store_config.access_key_secret = "SECRET_AK_VALUE" + feature_store_config.security_token = "SECRET_STS" + feature_store_config.featuredb_username = "SECRET_FDB_USER" + feature_store_config.featuredb_password = "SECRET_FDB_PASSWORD" + + sanitized = config_util.sanitize_pipeline_config_for_artifact(pipeline_config) + sanitized_fs = ( + sanitized.train_config.delta_embedding_dump_config.feature_store_config + ) + self.assertEqual(sanitized_fs.project_name, "project_a") + self.assertEqual(sanitized_fs.feature_entity_name, "embedding_entity") + self.assertEqual(sanitized_fs.feature_view_name, "shared_embeddings") + self.assertEqual(sanitized_fs.feature_view_ttl_secs, 86400) + self.assertEqual(sanitized_fs.feature_view_shard_count, 4) + self.assertEqual(sanitized_fs.feature_view_replication_count, 2) + self.assertEqual(sanitized_fs.version, "model_a@export_1") + self.assertTrue(sanitized_fs.IsInitialized()) + for field_name in ( + "access_key_id", + "access_key_secret", + "featuredb_username", + "featuredb_password", + ): + self.assertTrue(sanitized_fs.HasField(field_name)) + self.assertEqual(getattr(sanitized_fs, field_name), "") + self.assertTrue(feature_store_config.HasField(field_name)) + self.assertFalse(sanitized_fs.HasField("security_token")) + self.assertTrue(feature_store_config.HasField("security_token")) + + with tempfile.TemporaryDirectory() as tmp_dir: + path = os.path.join(tmp_dir, "pipeline.config") + config_util.save_pipeline_config_artifact(pipeline_config, path) + with open(path) as source: + artifact_text = source.read() + for secret in ( + "SECRET_AK_ID", + "SECRET_AK_VALUE", + "SECRET_STS", + "SECRET_FDB_USER", + "SECRET_FDB_PASSWORD", + ): + self.assertNotIn(secret, artifact_text) + if __name__ == "__main__": unittest.main() diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 3bb8f3d88..5fd1faa24 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -9,8 +9,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import hashlib import os import re +import time +import uuid from contextlib import contextmanager from dataclasses import dataclass from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Set, Tuple @@ -20,26 +23,64 @@ import torch from torch import nn from torch.distributed._shard.sharded_tensor import ShardedTensor +from torchrec.distributed.embedding import ShardedEmbeddingCollection +from torchrec.distributed.embeddingbag import ShardedEmbeddingBagCollection from torchrec.distributed.model_tracker.model_delta_tracker import ( ModelDeltaTrackerTrec, ) from torchrec.distributed.model_tracker.types import TrackingMode from tzrec.protos.train_pb2 import DeltaEmbeddingDumpConfig +from tzrec.utils.feature_store_delta_uploader import ( + DELTA_DUMP_GENERATION_METADATA_KEY, + DELTA_DUMP_SCHEMA_VERSION, + DELTA_OPERATION_UPSERT, + FeatureStoreDeltaUploader, + FeatureStoreUploadError, + _durable_makedirs, + feature_store_delta_file_prefix, + feature_store_upload_error_marker_path, + validate_feature_store_config, +) from tzrec.utils.logging_util import logger +from tzrec.utils.sparse_embedding_contract import ( + SPARSE_EBC_ROLE, + SPARSE_EC_ROLE, + SPARSE_EMBEDDING_INVALID_KEY, + SPARSE_EMBEDDING_ROLES, + SparseEmbeddingIdentity, + build_sparse_embedding_name_map, + resolve_sparse_embedding_name, + sparse_embedding_role_from_state_key, +) _CONSUMER = "delta_embedding_dump" +_DURABILITY_CONSUMER = "delta_embedding_dump_durable_ack" _DELTA_DUMP_SCHEMA = pa.schema( [ ("global_step", pa.int64()), ("rank", pa.int32()), ("world_size", pa.int32()), + ("embedding_name", pa.string()), + ("embedding_role", pa.string()), ("feature_name", pa.string()), ("table_fqn", pa.string()), ("key_id", pa.int64()), ("embedding", pa.list_(pa.float32())), + ("operation", pa.string()), ("source", pa.string()), - ] + ], + metadata={ + b"tzrec.delta_embedding.schema_version": DELTA_DUMP_SCHEMA_VERSION.encode( + "ascii" + ), + b"tzrec.delta_embedding.dynamic_key_encoding": ( + b"UINT64_BIT_PATTERN_IN_SIGNED_INT64" + ), + b"tzrec.delta_embedding.invalid_key": str(SPARSE_EMBEDDING_INVALID_KEY).encode( + "ascii" + ), + }, ) @@ -87,6 +128,8 @@ def validate_delta_embedding_dump_config( ) if config.dump_interval_steps <= 0: raise ValueError("delta_embedding_dump_config.dump_interval_steps must be > 0.") + if config.HasField("feature_store_config"): + validate_feature_store_config(config.feature_store_config) def _has_proto_field(config: Any, field_name: str) -> bool: @@ -334,16 +377,40 @@ def __init__( self._output_dir = config.output_dir or os.path.join( model_dir, "delta_embedding_dump" ) - self._file_prefix = config.file_prefix or "delta_embedding" + file_prefix = config.file_prefix or "delta_embedding" + self._file_prefix = file_prefix self._rank, self._world_size = _distributed_rank_world_size() + self._device = device self._tracking_pause_depth = 0 - os.makedirs(self._output_dir, exist_ok=True) + self._feature_store_enabled = config.HasField("feature_store_config") + self._run_generation: Optional[str] = None + self._feature_store_error_marker_path: Optional[str] = None + self._next_feature_store_error_check = 0.0 + self._feature_store_error_check_interval_secs = 1 + if self._feature_store_enabled: + self._file_prefix = feature_store_delta_file_prefix( + config.feature_store_config, self._file_prefix + ) + self._feature_store_error_check_interval_secs = max( + int(config.feature_store_config.poll_interval_secs), 1 + ) + self._feature_store_error_marker_path = ( + feature_store_upload_error_marker_path( + config.feature_store_config, self._output_dir + ) + ) + _durable_makedirs(self._output_dir) self._table_shard_infos = self._collect_table_shard_infos() self._validate_supported_table_sharding(self._table_shard_infos) + tracker_consumers = [_CONSUMER] + if self._feature_store_enabled: + # Keep tracked rows resident until their parquet outbox is durable. + # The second consumer advances only after atomic os.replace(). + tracker_consumers.append(_DURABILITY_CONSUMER) self._tracker = ModelDeltaTrackerTrec( model, - consumers=[_CONSUMER], + consumers=tracker_consumers, delete_on_read=True, auto_compact=True, mode=TrackingMode.ID_ONLY, @@ -356,20 +423,236 @@ def __init__( } self._fqn_to_feature_names: Dict[str, List[str]] = {} self._fqn_to_feature_names.update(self._tracker.fqn_to_feature_names()) + self._fqn_to_identity, embedding_dimensions = ( + self._build_sparse_embedding_contract() + ) + self._uploader: Optional[FeatureStoreDeltaUploader] = None + if self._feature_store_enabled and self._rank == 0: + self._uploader = FeatureStoreDeltaUploader( + config.feature_store_config, + output_dir=self._output_dir, + file_prefix=file_prefix, + world_size=self._world_size, + embedding_dimensions=embedding_dimensions, + ) logger.info( "Delta embedding dump enabled: interval=%s output_dir=%s " - "rank=%s/%s tables=%s", + "rank=%s/%s tables=%s feature_store_upload=%s", self._interval, self._output_dir, self._rank, self._world_size, sorted(self._table_to_fqn.keys()), + self._feature_store_enabled, ) def clear(self) -> None: """Clear tracked sparse ids, usually after restore-time dummy steps.""" - self._tracker.clear(_CONSUMER) + self._tracker.clear() + + def start(self) -> None: + """Start rank-zero background publication after training initialization.""" + if self._feature_store_enabled: + self._initialize_run_generation() + if self._uploader is not None: + self._uploader.start() + + def close(self, raise_on_error: bool = True, drain: bool = True) -> None: + """Close the rank-zero uploader; abnormal shutdown can skip draining.""" + if self._uploader is not None: + self._uploader.close(raise_on_error=raise_on_error, drain=drain) + + def _feature_store_upload_error( + self, force: bool = False + ) -> Optional[BaseException]: + """Collect a local/shared uploader error without changing rank control flow.""" + if not getattr(self, "_feature_store_enabled", False): + return None + + local_error: Optional[BaseException] = None + uploader = getattr(self, "_uploader", None) + if uploader is not None: + try: + uploader.check_error() + except BaseException as exc: + local_error = exc + + if local_error is not None: + return local_error + marker_path = getattr(self, "_feature_store_error_marker_path", None) + now = time.monotonic() + next_check = getattr(self, "_next_feature_store_error_check", 0.0) + if not force and now < next_check: + return None + self._next_feature_store_error_check = now + getattr( + self, "_feature_store_error_check_interval_secs", 1 + ) + if marker_path and os.path.isfile(marker_path): + return FeatureStoreUploadError( + "FeatureStore delta upload failed on rank zero; all training " + "workers are stopping and parquet outbox files were retained" + ) + return None + + def _check_feature_store_upload_error(self, force: bool = False) -> None: + """Surface the rank-zero background failure through the shared outbox.""" + error = self._feature_store_upload_error(force=force) + if error is not None: + raise error.with_traceback(error.__traceback__) + + def _initialize_run_generation(self) -> None: + """Broadcast one run fence after every rank constructed successfully.""" + if self._run_generation is not None: + return + generation = uuid.uuid4().bytes if self._rank == 0 else bytes(16) + if self._world_size > 1: + if not ( + torch.distributed.is_available() and torch.distributed.is_initialized() + ): + raise RuntimeError( + "distributed FeatureStore delta dump requires an initialized " + "process group" + ) + token = torch.tensor( + list(generation), dtype=torch.uint8, device=self._device + ) + torch.distributed.broadcast(token, src=0) + generation = bytes(token.cpu().tolist()) + self._run_generation = generation.hex() + + def _next_dump_generation(self, global_step: int) -> Optional[str]: + """Derive a stable per-step token from the process-run generation fence.""" + if not getattr(self, "_feature_store_enabled", False): + return None + if self._run_generation is None: + raise RuntimeError("FeatureStore delta dumper must be started before use") + value = f"{self._run_generation}:{global_step}".encode("ascii") + return hashlib.sha256(value).hexdigest()[:32] + + def _build_sparse_embedding_contract( + self, + ) -> Tuple[Dict[str, SparseEmbeddingIdentity], Dict[str, int]]: + """Build the same physical-table identity consumed by sparse export.""" + metadata_by_identity: Dict[Tuple[str, str], Tuple[int, Tuple[str, ...]]] = {} + owner_by_identity: Dict[Tuple[str, str], str] = {} + roles_by_table: Dict[str, Set[str]] = {} + for module_fqn, module in self._tracker.get_tracked_modules().items(): + if isinstance(module, ShardedEmbeddingCollection): + role = SPARSE_EC_ROLE + elif isinstance(module, ShardedEmbeddingBagCollection): + role = SPARSE_EBC_ROLE + else: + continue + table_name_to_config = getattr(module, "_table_name_to_config", {}) + for table_name, table_config in table_name_to_config.items(): + dimension = _int_attr(table_config, "embedding_dim") + if dimension <= 0: + dimension = self._table_shard_infos.get( + table_name, _TableShardInfo() + ).global_cols + feature_names = tuple(getattr(table_config, "feature_names", ())) + identity_key = (role, table_name) + previous_owner = owner_by_identity.get(identity_key) + if previous_owner is not None and previous_owner != module_fqn: + raise ValueError( + "delta embedding dump cannot distinguish duplicate physical " + f"table identity {identity_key}: {previous_owner!r} vs " + f"{module_fqn!r}" + ) + owner_by_identity[identity_key] = module_fqn + previous = metadata_by_identity.get(identity_key) + current = (dimension, feature_names) + if previous is not None and previous != current: + raise ValueError( + "inconsistent sparse embedding metadata for " + f"role={role} table={table_name}: {previous} vs {current}" + ) + metadata_by_identity[identity_key] = current + roles_by_table.setdefault(table_name, set()).add(role) + + ambiguous_tables = sorted( + table_name for table_name, roles in roles_by_table.items() if len(roles) > 1 + ) + if ambiguous_tables: + # ModelDeltaTrackerTrec currently stores table_name -> FQN and + # overwrites one collection when EC/EBC reuse a raw name. Refuse to + # publish incomplete data instead of silently assigning a wrong PK. + raise ValueError( + "delta embedding dump cannot safely track table names reused by " + "both EmbeddingCollection and EmbeddingBagCollection: " + f"{ambiguous_tables}. TorchRec tracker needs role-aware identity." + ) + + name_by_identity = build_sparse_embedding_name_map(metadata_by_identity) + identity_by_fqn: Dict[str, SparseEmbeddingIdentity] = {} + embedding_dimensions: Dict[str, int] = {} + for table_name, fqn in self._table_to_fqn.items(): + role = sparse_embedding_role_from_state_key(fqn) + if role is None: + roles = roles_by_table.get(table_name, set()) + if len(roles) == 1: + role = next(iter(roles)) + if role is None or (role, table_name) not in metadata_by_identity: + raise ValueError( + "cannot resolve sparse embedding collection role for " + f"table={table_name!r}, fqn={fqn!r}" + ) + dimension, feature_names = metadata_by_identity[(role, table_name)] + if dimension <= 0: + raise ValueError( + f"invalid embedding dimension for table {table_name!r}: {dimension}" + ) + embedding_name = resolve_sparse_embedding_name( + name_by_identity, table_name, role + ) + identity = SparseEmbeddingIdentity( + role=role, + table_name=table_name, + embedding_name=embedding_name, + dimension=dimension, + feature_names=feature_names, + ) + identity_by_fqn[fqn] = identity + previous_dimension = embedding_dimensions.get(embedding_name) + if previous_dimension is not None and previous_dimension != dimension: + raise ValueError( + f"canonical embedding {embedding_name!r} has inconsistent " + f"dimensions: {previous_dimension} vs {dimension}" + ) + embedding_dimensions[embedding_name] = dimension + return identity_by_fqn, embedding_dimensions + + def _tracker_cursor_before_read(self) -> Optional[int]: + if not getattr(self, "_feature_store_enabled", False): + return None + return int(self._tracker.per_consumer_batch_idx[_CONSUMER]) + + def _rollback_tracker_read(self, cursor: Optional[int]) -> None: + if cursor is not None: + self._tracker.per_consumer_batch_idx[_CONSUMER] = cursor + + def _ack_durable_tracker_read(self) -> None: + if getattr(self, "_feature_store_enabled", False): + # Do not call get_unique() for the guard: that would repeat the + # expensive GPU cat/unique already performed by the real consumer. + # The parquet shard is durable now, so advance the guard cursor to + # the exact snapshot consumed above, then best-effort reclaim rows. + durable_cursor = int(self._tracker.per_consumer_batch_idx[_CONSUMER]) + self._tracker.per_consumer_batch_idx[_DURABILITY_CONSUMER] = durable_cursor + if getattr(self._tracker, "_delete_on_read", False): + try: + self._tracker.store.delete( + up_to_idx=min(self._tracker.per_consumer_batch_idx.values()) + ) + except BaseException as exc: + # Cursor advancement is the durability acknowledgement; + # deletion is only memory reclamation and can be retried by + # a later dump without republishing rows. + logger.warning( + "Failed to reclaim durable delta tracker rows (%s).", + type(exc).__name__, + ) @contextmanager def pause_tracking(self) -> Iterator[None]: @@ -386,6 +669,9 @@ def maybe_dump(self, global_step: int) -> None: Args: global_step: Current training step. """ + # This is a throttled local shared-filesystem check, not a distributed + # collective, so all ranks can surface an async rank-zero failure quickly. + self._check_feature_store_upload_error() if global_step > 0 and global_step % self._interval == 0: self.dump(global_step) self._tracker.step() @@ -404,6 +690,10 @@ def final_dump(self, global_step: int) -> Optional[str]: Path to the dumped parquet file, or None if skipped. """ global_step = self._sync_final_step(global_step) + if global_step <= 0: + # Step zero is excluded from the delta publication contract. + logger.info("Skipping delta embedding dump at step %s.", global_step) + return None if global_step > 0 and global_step % self._interval == 0: # Boundary steps were already written (with full delta) by # ``maybe_dump``. Re-dumping here has no new delta to flush -- every @@ -413,10 +703,20 @@ def final_dump(self, global_step: int) -> Optional[str]: # consumer window. Re-dumping would also overwrite the already-written # boundary shards (with an empty file under multi-GPU), so skip. return None - return self.dump(global_step) + # The error bit was included in the mandatory final collective. Do not + # perform another rank-local marker check here: a marker appearing between + # ranks could otherwise make only part of the shard set return early. + output_path: Optional[str] = None + local_error: Optional[BaseException] = None + try: + output_path = self.dump(global_step, check_upload_error=False) + except BaseException as exc: + local_error = exc + self._raise_if_any_final_dump_failed(local_error) + return output_path def _sync_final_step(self, global_step: int) -> int: - """Align the final step across ranks before the trailing flush. + """Align the final step and uploader failure state before the trailing flush. ``maybe_dump`` runs in lockstep so every rank shares ``global_step``, but ``final_dump`` is reached with each rank's own last step. With @@ -426,51 +726,112 @@ def _sync_final_step(self, global_step: int) -> int: empty-shard logic prevents. Reduce with MAX so the furthest-progressed rank's trailing delta is never swallowed by the boundary-step skip, and so every rank takes the same skip/dump decision into the same dir. + + The same collective carries a local/shared uploader-failure bit. No rank + raises before all ranks have entered it, preventing an asynchronous marker + from stranding another worker inside the final all-reduce. """ - if self._world_size <= 1: - return global_step - if not ( + local_error = self._feature_store_upload_error(force=True) + any_failed = local_error is not None + synced_step = global_step + if self._world_size > 1 and ( torch.distributed.is_available() and torch.distributed.is_initialized() ): - return global_step - device = torch.device(f"cuda:{torch.cuda.current_device()}") - step_tensor = torch.tensor(global_step, dtype=torch.long, device=device) - torch.distributed.all_reduce(step_tensor, op=torch.distributed.ReduceOp.MAX) - return int(step_tensor.item()) + device = torch.device(f"cuda:{torch.cuda.current_device()}") + final_state = torch.tensor( + [global_step, int(any_failed)], dtype=torch.long, device=device + ) + torch.distributed.all_reduce(final_state, op=torch.distributed.ReduceOp.MAX) + synced_step = int(final_state[0].item()) + any_failed = bool(final_state[1].item()) + + if local_error is not None: + raise local_error.with_traceback(local_error.__traceback__) + if any_failed: + raise FeatureStoreUploadError( + "FeatureStore delta upload failed on another distributed worker; " + "all workers are stopping and parquet outbox files were retained" + ) + return synced_step - def dump(self, global_step: int) -> Optional[str]: + def _raise_if_any_final_dump_failed( + self, local_error: Optional[BaseException] + ) -> None: + """Make rank-local final shard/submit failures visible before checkpointing.""" + any_failed = local_error is not None + if self._world_size > 1 and ( + torch.distributed.is_available() and torch.distributed.is_initialized() + ): + device = torch.device(f"cuda:{torch.cuda.current_device()}") + failed = torch.tensor([int(any_failed)], dtype=torch.int32, device=device) + torch.distributed.all_reduce(failed, op=torch.distributed.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( + "final delta embedding dump failed on another distributed worker; " + "no worker may enter the final checkpoint" + ) + + def dump(self, global_step: int, check_upload_error: bool = True) -> Optional[str]: """Dump currently tracked sparse ids and embeddings to a parquet file. Args: global_step: Current training step. + check_upload_error: Whether to perform a rank-local async error check. + ``final_dump`` disables it after synchronizing the error bit. Returns: Path to the dumped parquet file, or None if no data to dump. """ - table_weights = self._collect_table_weights() - dynamic_modules = self._collect_dynamic_modules() - table_chunks: List[pa.Table] = [] - num_rows = self._append_model_delta_rows( - table_chunks, - global_step=global_step, - table_weights=table_weights, - dynamic_modules=dynamic_modules, - ) - if num_rows == 0: - if self._world_size == 1: + global_step = int(global_step) + if global_step <= 0: + raise ValueError("delta embedding dump global_step must be > 0") + if check_upload_error: + self._check_feature_store_upload_error(force=True) + uploader = getattr(self, "_uploader", None) + dump_generation = self._next_dump_generation(global_step) + tracker_cursor = self._tracker_cursor_before_read() + try: + table_weights = self._collect_table_weights() + dynamic_modules = self._collect_dynamic_modules() + table_chunks: List[pa.Table] = [] + num_rows = self._append_model_delta_rows( + table_chunks, + global_step=global_step, + table_weights=table_weights, + dynamic_modules=dynamic_modules, + ) + if ( + num_rows == 0 + and self._world_size == 1 + and not getattr(self, "_feature_store_enabled", False) + ): logger.info("No delta embedding rows to dump at step %s.", global_step) return None output_path = self._output_path(global_step) - self._write_table_chunks(table_chunks, output_path) + self._write_table_chunks( + table_chunks, output_path, dump_generation=dump_generation + ) + except BaseException: + # The durability guard still owns the rows, so rewinding this + # consumer makes a caller retry observe the same snapshot. + self._rollback_tracker_read(tracker_cursor) + raise + + self._ack_durable_tracker_read() + if uploader is not None: + uploader.submit(global_step) + if num_rows == 0: logger.info( "Dumped empty delta embedding shard to %s at step %s.", output_path, global_step, ) - return output_path - output_path = self._output_path(global_step) - self._write_table_chunks(table_chunks, output_path) - logger.info("Dumped %s delta embedding rows to %s.", num_rows, output_path) + else: + logger.info("Dumped %s delta embedding rows to %s.", num_rows, output_path) return output_path def _output_path(self, global_step: int) -> str: @@ -479,7 +840,7 @@ def _output_path(self, global_step: int) -> str: self._output_dir, f"{self._file_prefix}_step_{global_step}.parquet" ) step_dir = os.path.join(self._output_dir, f"step_{global_step}") - os.makedirs(step_dir, exist_ok=True) + _durable_makedirs(step_dir) return os.path.join( step_dir, ( @@ -535,6 +896,11 @@ def _append_model_delta_rows( if table_name is None: logger.warning("Skip delta rows for unknown table fqn: %s", fqn) continue + identity = self._fqn_to_identity.get(fqn) + if identity is None: + raise ValueError( + f"Missing sparse embedding contract for table fqn {fqn!r}" + ) ids = ids.unique(sorted=True) embeddings, key_ids = self._lookup_embeddings( table_name, @@ -547,6 +913,9 @@ def _append_model_delta_rows( num_rows += self._append_table_chunk( table_chunks, global_step=global_step, + embedding_name=identity.embedding_name, + embedding_role=identity.role, + expected_dimension=identity.dimension, feature_name=feature_name, table_fqn=fqn, key_ids=key_ids, @@ -737,6 +1106,9 @@ def _append_table_chunk( self, table_chunks: List[pa.Table], global_step: int, + embedding_name: str, + embedding_role: str, + expected_dimension: int, feature_name: str, table_fqn: str, key_ids: torch.Tensor, @@ -745,6 +1117,12 @@ def _append_table_chunk( ) -> int: key_ids_cpu = key_ids.detach().cpu().to(torch.int64).contiguous() embeddings_cpu = embeddings.detach().cpu().to(torch.float32).contiguous() + if not embedding_name: + raise ValueError("delta embedding dump embedding_name must not be empty") + if embedding_role not in SPARSE_EMBEDDING_ROLES: + raise ValueError( + f"delta embedding dump has invalid embedding_role={embedding_role!r}" + ) if embeddings_cpu.dim() != 2: raise ValueError( "delta embedding dump expects a 2-D embedding tensor, " @@ -758,17 +1136,31 @@ def _append_table_chunk( "delta embedding dump key ids and embeddings row count mismatch: " f"key_ids={num_rows}, embeddings={embeddings_cpu.size(0)}." ) - + if embeddings_cpu.size(1) != expected_dimension: + raise ValueError( + f"delta embedding dimension mismatch for {embedding_name!r}: " + f"expected={expected_dimension}, actual={embeddings_cpu.size(1)}" + ) + if not bool(torch.isfinite(embeddings_cpu).all().item()): + raise ValueError(f"delta embedding {embedding_name!r} contains NaN or Inf") + if bool((key_ids_cpu == SPARSE_EMBEDDING_INVALID_KEY).any().item()): + raise ValueError( + "delta embedding key_id=-1 is reserved as the Processor/NvEmbeddings " + "invalid-key sentinel" + ) table_chunks.append( pa.Table.from_arrays( [ pa.repeat(pa.scalar(global_step, pa.int64()), num_rows), pa.repeat(pa.scalar(self._rank, pa.int32()), num_rows), pa.repeat(pa.scalar(self._world_size, pa.int32()), num_rows), + pa.repeat(pa.scalar(embedding_name, pa.string()), num_rows), + pa.repeat(pa.scalar(embedding_role, pa.string()), num_rows), pa.repeat(pa.scalar(feature_name, pa.string()), num_rows), pa.repeat(pa.scalar(table_fqn, pa.string()), num_rows), pa.array(key_ids_cpu.numpy(), type=pa.int64()), self._embedding_array(embeddings_cpu), + pa.repeat(pa.scalar(DELTA_OPERATION_UPSERT, pa.string()), num_rows), pa.repeat(pa.scalar(source, pa.string()), num_rows), ], schema=_DELTA_DUMP_SCHEMA, @@ -792,7 +1184,10 @@ def _embedding_array(self, embeddings: torch.Tensor) -> pa.ListArray: return pa.ListArray.from_arrays(pa.array(offsets, type=pa.int32()), values) def _write_table_chunks( - self, table_chunks: List[pa.Table], output_path: str + self, + table_chunks: List[pa.Table], + output_path: str, + dump_generation: Optional[str] = None, ) -> None: # Write to a sibling temp file and atomically os.replace() it into place # only after the writer closes cleanly. A kill or exception mid-write @@ -800,11 +1195,25 @@ def _write_table_chunks( # ignores), never a truncated shard at the canonical path. tmp_path = f"{output_path}.rank{self._rank}.tmp" try: - with pq.ParquetWriter(tmp_path, _DELTA_DUMP_SCHEMA) as writer: + writer_schema = _DELTA_DUMP_SCHEMA + if dump_generation is not None: + metadata = dict(writer_schema.metadata or {}) + metadata[DELTA_DUMP_GENERATION_METADATA_KEY] = dump_generation.encode( + "ascii" + ) + writer_schema = writer_schema.with_metadata(metadata) + with pq.ParquetWriter(tmp_path, writer_schema) as writer: chunks = table_chunks or [_DELTA_DUMP_SCHEMA.empty_table()] for table_chunk in chunks: writer.write_table(table_chunk) + with open(tmp_path, "rb") as source: + os.fsync(source.fileno()) os.replace(tmp_path, output_path) + directory_fd = os.open(os.path.dirname(output_path) or ".", os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) except BaseException: if os.path.exists(tmp_path): os.remove(tmp_path) diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 381953f70..fccea83b9 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -44,7 +44,9 @@ from tzrec.tests import utils as test_utils from tzrec.utils import config_util from tzrec.utils.delta_embedding_dump import ( + _CONSUMER, _DELTA_DUMP_SCHEMA, + _DURABILITY_CONSUMER, DeltaEmbeddingDumper, _table_shard_info_from_config, _TableShardInfo, @@ -54,6 +56,10 @@ validate_delta_embedding_dump_no_zch_features, ) from tzrec.utils.dynamicemb_util import has_dynamicemb +from tzrec.utils.feature_store_delta_uploader import ( + DELTA_DUMP_GENERATION_METADATA_KEY, + FeatureStoreUploadError, +) from tzrec.utils.test_util import gpu_unavailable, make_test_dir, mark_ci_scope _SHARDED_TABLE_NAME = "table_1" @@ -156,6 +162,11 @@ def _assert_sharded_dump_file(rank: int, output_path: str, dumper) -> None: testcase.assertEqual( set(table["feature_name"].to_pylist()), {_SHARDED_FEATURE_NAME} ) + testcase.assertEqual( + set(table["embedding_name"].to_pylist()), {_SHARDED_TABLE_NAME} + ) + testcase.assertEqual(set(table["embedding_role"].to_pylist()), {"ebc"}) + testcase.assertEqual(set(table["operation"].to_pylist()), {"UPSERT"}) testcase.assertEqual(set(table["source"].to_pylist()), {"model_delta_tracker"}) table_weight = dumper._collect_table_weights()[_SHARDED_TABLE_NAME] @@ -329,9 +340,12 @@ def test_dump_rows_include_rank_metadata(self): num_rows = dumper._append_table_chunk( table_chunks, global_step=10, + embedding_name="user_emb", + embedding_role="ebc", + expected_dimension=2, feature_name="user_id", table_fqn="model.ebc.user_emb", - key_ids=torch.tensor([42]), + key_ids=torch.tensor([-42]), embeddings=torch.tensor([[1.0, 2.0]]), source="model_delta_tracker", ) @@ -341,8 +355,29 @@ def test_dump_rows_include_rank_metadata(self): self.assertEqual(table.schema, _DELTA_DUMP_SCHEMA) self.assertEqual(table["rank"].to_pylist(), [1]) self.assertEqual(table["world_size"].to_pylist(), [4]) - self.assertEqual(table["key_id"].to_pylist(), [42]) + self.assertEqual(table["embedding_name"].to_pylist(), ["user_emb"]) + self.assertEqual(table["embedding_role"].to_pylist(), ["ebc"]) + self.assertEqual(table["key_id"].to_pylist(), [-42]) self.assertEqual(table["embedding"].to_pylist(), [[1.0, 2.0]]) + self.assertEqual(table["operation"].to_pylist(), ["UPSERT"]) + + def test_dump_rows_reject_processor_invalid_key_sentinel(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._rank = 0 + dumper._world_size = 1 + with self.assertRaisesRegex(ValueError, "invalid-key sentinel"): + dumper._append_table_chunk( + [], + global_step=10, + embedding_name="user_emb", + embedding_role="ebc", + expected_dimension=2, + feature_name="user_id", + table_fqn="model.ebc.user_emb", + key_ids=torch.tensor([-1]), + embeddings=torch.tensor([[1.0, 2.0]]), + source="model_delta_tracker", + ) def test_write_table_chunks_preserves_parquet_schema(self): dumper = object.__new__(DeltaEmbeddingDumper) @@ -352,6 +387,9 @@ def test_write_table_chunks_preserves_parquet_schema(self): dumper._append_table_chunk( table_chunks, global_step=5, + embedding_name="user_emb", + embedding_role="ebc", + expected_dimension=2, feature_name="user_id", table_fqn="model.ebc.user_emb", key_ids=torch.tensor([7, 8]), @@ -379,6 +417,19 @@ def test_write_empty_table_chunks_preserves_parquet_schema(self): self.assertEqual(table.schema, _DELTA_DUMP_SCHEMA) self.assertEqual(table.num_rows, 0) + def test_write_table_chunks_persists_dump_generation_metadata(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._rank = 0 + generation = "00112233445566778899aabbccddeeff" + with tempfile.TemporaryDirectory() as tmp_dir: + output_path = os.path.join(tmp_dir, "delta.parquet") + dumper._write_table_chunks([], output_path, dump_generation=generation) + metadata = pq.read_schema(output_path).metadata + + self.assertEqual( + metadata[DELTA_DUMP_GENERATION_METADATA_KEY], generation.encode("ascii") + ) + def test_write_table_chunks_leaves_no_partial_shard_on_error(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._rank = 0 @@ -408,12 +459,19 @@ def test_final_dump_skips_boundary_step_to_avoid_overwrite(self): self.assertIsNone(dumper.final_dump(100)) dump_mock.assert_not_called() - # Trailing partial interval (and step 0) must still be flushed. - dumper.final_dump(0) + # Step 0 is not publishable. A positive trailing partial interval + # must still be flushed. + self.assertIsNone(dumper.final_dump(0)) dumper.final_dump(73) self.assertEqual( [call.args[0] for call in dump_mock.call_args_list], - [0, 73], + [73], + ) + self.assertTrue( + all( + call.kwargs == {"check_upload_error": False} + for call in dump_mock.call_args_list + ) ) def test_final_dump_syncs_step_across_ranks_before_flush(self): @@ -427,7 +485,11 @@ def test_final_dump_syncs_step_across_ranks_before_flush(self): def fake_all_reduce(tensor, op=None): self.assertIs(op, torch.distributed.ReduceOp.MAX) - tensor.fill_(73) + if tensor.numel() == 2: + tensor[0] = 73 + tensor[1] = 0 + else: + tensor[0] = 0 with ( mock.patch.object(dumper, "dump") as dump_mock, @@ -436,18 +498,140 @@ def fake_all_reduce(tensor, op=None): mock.patch("torch.cuda.current_device", return_value=0), mock.patch( "torch.tensor", - side_effect=lambda *a, **k: torch.zeros(1, dtype=torch.long), + side_effect=lambda value, *a, **k: torch.zeros( + len(value), dtype=torch.long + ), ), mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), ): dumper.final_dump(50) - dump_mock.assert_called_once_with(73) + dump_mock.assert_called_once_with(73, check_upload_error=False) + + def test_final_dump_syncs_ranks_before_skipping_step_zero(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval = 50 + dumper._world_size = 2 + collective_count = 0 + + def fake_all_reduce(tensor, op=None): + nonlocal collective_count + self.assertIs(op, torch.distributed.ReduceOp.MAX) + tensor[0] = 0 + tensor[1] = 0 + collective_count += 1 + + with ( + mock.patch.object(dumper, "dump") as dump_mock, + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.cuda.current_device", return_value=0), + mock.patch( + "torch.tensor", + side_effect=lambda value, *a, **k: torch.zeros( + len(value), dtype=torch.long + ), + ), + mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), + ): + self.assertIsNone(dumper.final_dump(0)) + + self.assertEqual(collective_count, 1) + dump_mock.assert_not_called() + + def test_final_dump_reduces_upload_error_before_raising(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval = 50 + dumper._world_size = 2 + + def fake_all_reduce(tensor, op=None): + self.assertIs(op, torch.distributed.ReduceOp.MAX) + tensor[0] = 73 + tensor[1] = 1 + + with ( + mock.patch.object(dumper, "_feature_store_upload_error", return_value=None), + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.cuda.current_device", return_value=0), + mock.patch( + "torch.tensor", + side_effect=lambda value, *a, **k: torch.zeros( + len(value), dtype=torch.long + ), + ), + mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), + ): + with self.assertRaisesRegex( + FeatureStoreUploadError, "another distributed worker" + ): + dumper.final_dump(50) + + def test_final_dump_reduces_rank_local_dump_failure_before_checkpoint(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval = 50 + dumper._world_size = 2 + collective_index = 0 + + def fake_all_reduce(tensor, op=None): + nonlocal collective_index + self.assertIs(op, torch.distributed.ReduceOp.MAX) + if collective_index == 0: + tensor[0] = 73 + tensor[1] = 0 + else: + tensor[0] = 1 + collective_index += 1 + + with ( + mock.patch.object(dumper, "_feature_store_upload_error", return_value=None), + mock.patch.object(dumper, "dump", return_value="rank-local.parquet"), + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.cuda.current_device", return_value=0), + mock.patch( + "torch.tensor", + side_effect=lambda value, *a, **k: torch.zeros( + len(value), dtype=torch.long + ), + ), + mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), + ): + with self.assertRaisesRegex(RuntimeError, "final delta embedding dump"): + dumper.final_dump(50) + self.assertEqual(collective_index, 2) + + def test_rank_zero_upload_failure_stops_non_uploader_rank(self): + with tempfile.TemporaryDirectory() as tmp_dir: + marker_path = os.path.join(tmp_dir, "last_error.json") + with open(marker_path, "w") as output: + output.write("{}") + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._feature_store_enabled = True + dumper._feature_store_error_marker_path = marker_path + dumper._uploader = None + + with self.assertRaisesRegex(FeatureStoreUploadError, "rank zero"): + dumper._check_feature_store_upload_error() + + def test_dump_generation_is_stable_per_step_and_unique_per_run(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._feature_store_enabled = True + dumper._rank = 0 + dumper._world_size = 1 + dumper._run_generation = None + dumper._initialize_run_generation() + + step_10_generation = dumper._next_dump_generation(10) + self.assertEqual(step_10_generation, dumper._next_dump_generation(10)) + self.assertNotEqual(step_10_generation, dumper._next_dump_generation(11)) def test_maybe_dump_uses_checkpoint_aligned_global_step(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._interval = 50 dumper._tracker = mock.MagicMock() with mock.patch.object(dumper, "dump") as dump_mock: + dumper.maybe_dump(0) + dump_mock.assert_not_called() dumper.maybe_dump(49) dump_mock.assert_not_called() dumper.maybe_dump(50) @@ -459,7 +643,12 @@ def test_maybe_dump_uses_checkpoint_aligned_global_step(self): [call.args[0] for call in dump_mock.call_args_list], [50, 100], ) - self.assertEqual(dumper._tracker.step.call_count, 4) + self.assertEqual(dumper._tracker.step.call_count, 5) + + def test_direct_dump_rejects_step_zero(self): + dumper = object.__new__(DeltaEmbeddingDumper) + with self.assertRaisesRegex(ValueError, "global_step must be > 0"): + dumper.dump(0) def test_tracker_uses_auto_compact(self): tracker = mock.MagicMock() @@ -483,6 +672,23 @@ def test_tracker_uses_auto_compact(self): self.assertTrue(tracker_cls.call_args.kwargs["auto_compact"]) + def test_durable_ack_advances_guard_without_recomputing_unique_rows(self): + tracker = mock.MagicMock() + tracker.per_consumer_batch_idx = { + _CONSUMER: 12, + _DURABILITY_CONSUMER: 5, + } + tracker._delete_on_read = True + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._feature_store_enabled = True + dumper._tracker = tracker + + dumper._ack_durable_tracker_read() + + self.assertEqual(tracker.per_consumer_batch_idx[_DURABILITY_CONSUMER], 12) + tracker.store.delete.assert_called_once_with(up_to_idx=12) + tracker.get_unique.assert_not_called() + def test_multi_gpu_output_path_uses_step_underscore_dir(self): with tempfile.TemporaryDirectory() as tmp_dir: dumper = object.__new__(DeltaEmbeddingDumper) diff --git a/tzrec/utils/export_util.py b/tzrec/utils/export_util.py index cec13353b..89eca78c8 100644 --- a/tzrec/utils/export_util.py +++ b/tzrec/utils/export_util.py @@ -296,7 +296,7 @@ def export_model_normal( pipeline_config = copy.copy(pipeline_config) pipeline_config.ClearField("feature_configs") pipeline_config.feature_configs.extend(feature_configs) - config_util.save_message( + config_util.save_pipeline_config_artifact( pipeline_config, os.path.join(save_dir, "pipeline.config") ) logger.info("saving fg json...") diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py new file mode 100644 index 000000000..69d5fddce --- /dev/null +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -0,0 +1,1872 @@ +# 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. + +"""Durable rank-zero uploader for delta-embedding parquet outboxes.""" + +import fcntl +import hashlib +import json +import os +import re +import shutil +import threading +import time +import uuid +from contextlib import ExitStack +from dataclasses import dataclass, field +from typing import ( + Any, + BinaryIO, + Callable, + Dict, + Iterable, + List, + Mapping, + Optional, + Tuple, +) +from urllib.parse import urlsplit + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +from tzrec.protos.train_pb2 import FeatureStoreConfig +from tzrec.utils.logging_util import logger +from tzrec.utils.sparse_embedding_contract import ( + SPARSE_EMBEDDING_INVALID_KEY, + SPARSE_EMBEDDING_ROLES, +) + +FEATURE_STORE_PK_FIELD = "embedding_name" +FEATURE_STORE_SK_FIELD = "key_id" +FEATURE_STORE_VALUE_FIELD = "embedding" +FEATURE_STORE_WRITE_MODE = "MERGE" +FEATURE_STORE_SDK_BATCH_SIZE = 1000 +DELTA_OPERATION_UPSERT = "UPSERT" +DELTA_DUMP_SCHEMA_VERSION = "2" +DELTA_DUMP_GENERATION_METADATA_KEY = b"tzrec.delta_embedding.dump_generation" + +_POISONED_WRITER_LOCKS: List[BinaryIO] = [] + +_SCHEMA_VERSION_METADATA_KEY = b"tzrec.delta_embedding.schema_version" +_REQUIRED_PARQUET_FIELDS = { + "global_step": pa.int64(), + "rank": pa.int32(), + "world_size": pa.int32(), + "embedding_name": pa.string(), + "embedding_role": pa.string(), + "feature_name": pa.string(), + "table_fqn": pa.string(), + "key_id": pa.int64(), + "embedding": pa.list_(pa.float32()), + "operation": pa.string(), + "source": pa.string(), +} + + +class FeatureStoreUploadError(RuntimeError): + """Safe, credential-free error propagated from the uploader thread.""" + + +class _UploadAborted(RuntimeError): + """Internal control flow for abnormal, non-draining shutdown.""" + + +class _ShardSetNotReady(RuntimeError): + """A multi-rank canonical shard set is still being atomically replaced.""" + + +def _config_or_env(config: FeatureStoreConfig, field_name: str, env_name: str) -> str: + value = getattr(config, field_name, "") + return value or os.environ.get(env_name, "") + + +def _validate_pair( + first: str, + second: str, + first_name: str, + second_name: str, + required: bool, +) -> None: + if bool(first) != bool(second): + raise ValueError(f"{first_name} and {second_name} must be configured together") + if required and not first: + raise ValueError(f"{first_name} and {second_name} are required") + + +@dataclass(frozen=True) +class FeatureStoreUploadSettings: + """Validated immutable settings copied from the runtime protobuf.""" + + region: str + endpoint: str + access_key_id: str = field(repr=False) + access_key_secret: str = field(repr=False) + security_token: str = field(repr=False) + featuredb_username: str = field(repr=False) + featuredb_password: str = field(repr=False) + project_name: str + feature_entity_name: str + feature_view_name: str + feature_view_ttl_secs: int + feature_view_shard_count: int + feature_view_replication_count: int + version: str + upload_batch_size: int + max_retries: int + retry_backoff_secs: int + shard_wait_timeout_secs: int + shutdown_timeout_secs: int + max_pending_steps: int + poll_interval_secs: int + + @classmethod + def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": + """Resolve standard environment fallbacks without logging credentials.""" + initialization_errors = config.FindInitializationErrors() + if initialization_errors: + raise ValueError( + "feature_store_config is missing required fields: " + + ", ".join(initialization_errors) + ) + region = _config_or_env(config, "region", "ALIBABA_CLOUD_REGION") + endpoint = config.endpoint + access_key_id = _config_or_env( + config, "access_key_id", "ALIBABA_CLOUD_ACCESS_KEY_ID" + ) + access_key_secret = _config_or_env( + config, "access_key_secret", "ALIBABA_CLOUD_ACCESS_KEY_SECRET" + ) + security_token = _config_or_env( + config, "security_token", "ALIBABA_CLOUD_SECURITY_TOKEN" + ) + featuredb_username = _config_or_env( + config, "featuredb_username", "FEATUREDB_USERNAME" + ) + featuredb_password = _config_or_env( + config, "featuredb_password", "FEATUREDB_PASSWORD" + ) + + if not region: + raise ValueError( + "feature_store_config.region must not be empty " + "(it may come from ALIBABA_CLOUD_REGION)" + ) + endpoint_url = endpoint if "://" in endpoint else f"//{endpoint}" + parsed_endpoint = urlsplit(endpoint_url) + if parsed_endpoint.username is not None or parsed_endpoint.password is not None: + raise ValueError( + "feature_store_config.endpoint must not contain URI userinfo" + ) + _validate_pair( + access_key_id, + access_key_secret, + "access_key_id", + "access_key_secret", + required=True, + ) + _validate_pair( + featuredb_username, + featuredb_password, + "featuredb_username", + "featuredb_password", + required=True, + ) + + project_name = config.project_name.strip() + feature_entity_name = config.feature_entity_name.strip() + feature_view_name = config.feature_view_name.strip() + version = config.version.strip() + if not project_name: + raise ValueError("feature_store_config.project_name must not be empty") + if not feature_entity_name: + raise ValueError( + "feature_store_config.feature_entity_name must not be empty" + ) + if not feature_view_name: + raise ValueError("feature_store_config.feature_view_name must not be empty") + if not version or version == "default": + raise ValueError( + "feature_store_config.version must be an explicit non-default version" + ) + + positive_values = { + "feature_view_ttl_secs": int(config.feature_view_ttl_secs), + "upload_batch_size": int(config.upload_batch_size), + "max_retries": int(config.max_retries), + "shard_wait_timeout_secs": int(config.shard_wait_timeout_secs), + "shutdown_timeout_secs": int(config.shutdown_timeout_secs), + "max_pending_steps": int(config.max_pending_steps), + "poll_interval_secs": int(config.poll_interval_secs), + } + for name, value in positive_values.items(): + if value <= 0: + raise ValueError(f"feature_store_config.{name} must be > 0") + feature_view_shard_count = int(config.feature_view_shard_count) + if not 1 <= feature_view_shard_count <= 20: + raise ValueError( + "feature_store_config.feature_view_shard_count must be in [1, 20]" + ) + feature_view_replication_count = int(config.feature_view_replication_count) + if not 1 <= feature_view_replication_count <= 3: + raise ValueError( + "feature_store_config.feature_view_replication_count must be in [1, 3]" + ) + if positive_values["upload_batch_size"] > FEATURE_STORE_SDK_BATCH_SIZE: + raise ValueError( + "feature_store_config.upload_batch_size must be <= " + f"{FEATURE_STORE_SDK_BATCH_SIZE} so one publish timestamp maps to " + "exactly one FeatureStore SDK HTTP batch" + ) + + return cls( + region=region, + endpoint=endpoint, + access_key_id=access_key_id, + access_key_secret=access_key_secret, + security_token=security_token, + featuredb_username=featuredb_username, + featuredb_password=featuredb_password, + project_name=project_name, + feature_entity_name=feature_entity_name, + feature_view_name=feature_view_name, + feature_view_ttl_secs=positive_values["feature_view_ttl_secs"], + feature_view_shard_count=feature_view_shard_count, + feature_view_replication_count=feature_view_replication_count, + version=version, + upload_batch_size=positive_values["upload_batch_size"], + max_retries=positive_values["max_retries"], + retry_backoff_secs=int(config.retry_backoff_secs), + shard_wait_timeout_secs=positive_values["shard_wait_timeout_secs"], + shutdown_timeout_secs=positive_values["shutdown_timeout_secs"], + max_pending_steps=positive_values["max_pending_steps"], + poll_interval_secs=positive_values["poll_interval_secs"], + ) + + +def validate_feature_store_config(config: FeatureStoreConfig) -> None: + """Validate upload configuration and its environment-resolved credentials.""" + FeatureStoreUploadSettings.from_proto(config) + + +def _feature_store_target_hash(settings: FeatureStoreUploadSettings) -> str: + target_identity = { + "region": settings.region, + "endpoint": settings.endpoint, + "project_name": settings.project_name, + "feature_view_name": settings.feature_view_name, + "version": settings.version, + } + return _json_digest(target_identity) + + +def feature_store_delta_file_prefix( + config: FeatureStoreConfig, file_prefix: str +) -> str: + """Scope canonical parquet names to one immutable FeatureStore target.""" + settings = FeatureStoreUploadSettings.from_proto(config) + return _scoped_feature_store_file_prefix(settings, file_prefix) + + +def _scoped_feature_store_file_prefix( + settings: FeatureStoreUploadSettings, file_prefix: str +) -> str: + target_hash = _feature_store_target_hash(settings) + return f"{file_prefix}__fs_{target_hash[:16]}" + + +def _feature_store_state_dir( + settings: FeatureStoreUploadSettings, output_dir: str +) -> str: + target_hash = _feature_store_target_hash(settings) + return os.path.join( + os.path.abspath(output_dir), ".feature_store_upload", target_hash[:16] + ) + + +def feature_store_upload_error_marker_path( + config: FeatureStoreConfig, output_dir: str +) -> str: + """Return the credential-free shared failure marker path for all ranks.""" + settings = FeatureStoreUploadSettings.from_proto(config) + return os.path.join( + _feature_store_state_dir(settings, output_dir), "last_error.json" + ) + + +def _fsync_parent_directory(path: str) -> None: + directory_fd = os.open(os.path.dirname(path) or ".", os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + + +def _durable_makedirs(path: str) -> None: + """Create missing directories and durably publish every new directory entry.""" + path = os.path.abspath(path) + missing = [] + current = path + while not os.path.exists(current): + missing.append(current) + parent = os.path.dirname(current) + if parent == current: + break + current = parent + if os.path.exists(current) and not os.path.isdir(current): + raise NotADirectoryError(current) + + for directory in reversed(missing): + try: + os.mkdir(directory) + except FileExistsError: + if not os.path.isdir(directory): + raise + # Another creator may have won the race but not fsynced its parent. + _fsync_parent_directory(directory) + else: + _fsync_parent_directory(directory) + if not os.path.isdir(path): + raise NotADirectoryError(path) + # Close the race where another rank created the final directory after our + # initial exists() check but before it fsynced the parent directory entry. + _fsync_parent_directory(path) + + +def _atomic_write_json(path: str, value: Mapping[str, Any]) -> None: + _durable_makedirs(os.path.dirname(path) or ".") + tmp_path = f"{path}.tmp-{os.getpid()}-{threading.get_ident()}-{uuid.uuid4().hex}" + try: + with open(tmp_path, "w") as output: + json.dump(value, output, indent=2, sort_keys=True) + output.write("\n") + output.flush() + os.fsync(output.fileno()) + os.replace(tmp_path, path) + _fsync_parent_directory(path) + except BaseException: + if os.path.exists(tmp_path): + os.remove(tmp_path) + raise + + +def _read_json(path: str) -> Dict[str, Any]: + with open(path) as source: + value = json.load(source) + if not isinstance(value, dict): + raise ValueError(f"expected a JSON object in {path}") + return value + + +def _json_digest(value: Mapping[str, Any]) -> str: + encoded = json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _file_sha256(source: BinaryIO) -> str: + digest = hashlib.sha256() + source.seek(0) + while True: + chunk = source.read(8 * 1024 * 1024) + if not chunk: + break + digest.update(chunk) + source.seek(0) + return digest.hexdigest() + + +class FeatureStoreDeltaUploader: + """Publish complete delta-dump steps from a persistent local outbox. + + The training thread only writes/parquet-renames and enqueues a step. This + object's single worker waits for the exact rank shard set and persists a + monotonic timestamp range before each full-step MERGE attempt. It never + invokes a torch.distributed collective. + """ + + def __init__( + self, + config: FeatureStoreConfig, + output_dir: str, + file_prefix: str, + world_size: int, + embedding_dimensions: Mapping[str, int], + client_factory: Optional[Callable[..., Any]] = None, + clock_ms: Optional[Callable[[], int]] = None, + ) -> None: + self._settings = FeatureStoreUploadSettings.from_proto(config) + self._output_dir = os.path.abspath(output_dir) + self._file_prefix = _scoped_feature_store_file_prefix( + self._settings, file_prefix + ) + self._world_size = int(world_size) + if self._world_size <= 0: + raise ValueError("world_size must be > 0") + self._embedding_dimensions = { + str(name): int(dimension) + for name, dimension in embedding_dimensions.items() + } + invalid_dimensions = { + name: dimension + for name, dimension in self._embedding_dimensions.items() + if not name or dimension <= 0 + } + if invalid_dimensions: + raise ValueError( + "invalid sparse embedding dimensions in FeatureStore contract: " + f"{invalid_dimensions}" + ) + + self._client_factory = client_factory + self._clock_ms = clock_ms or (lambda: time.time_ns() // 1_000_000) + self._view = None + self._condition = threading.Condition() + self._pending: Dict[int, float] = {} + self._started = False + self._closing = False + self._aborting = False + self._closed = False + self._worker: Optional[threading.Thread] = None + self._error: Optional[FeatureStoreUploadError] = None + self._writer_quiescence_failed = False + + contract = { + "schema_version": 5, + "delta_dump_schema_version": DELTA_DUMP_SCHEMA_VERSION, + "region": self._settings.region, + "endpoint": self._settings.endpoint, + "project_name": self._settings.project_name, + "feature_entity_name": self._settings.feature_entity_name, + "feature_view_name": self._settings.feature_view_name, + "feature_view_ttl_secs": self._settings.feature_view_ttl_secs, + "feature_view_shard_count": self._settings.feature_view_shard_count, + "feature_view_replication_count": ( + self._settings.feature_view_replication_count + ), + "version": self._settings.version, + "version_initialization": "PREPROVISIONED_FOR_DELTA_MERGE", + "feature_view_provisioning": "CHECK_OR_CREATE_DYNAMIC_EMBEDDING", + "writer_ownership": "SINGLE_WRITER_PER_VERSION_REQUIRED", + "publish_semantics": "MONOTONIC_TS_RANGE_PER_ATTEMPT_FULL_REPLAY", + "consistency": "ROW_LEVEL_EVENTUAL", + "step_atomicity": False, + "outbox_snapshot": "PRIVATE_READ_ONLY_CONTENT_VERIFIED", + "canonical_outbox_scope": "FEATURE_STORE_TARGET_HASH", + "dump_generation": "ONE_SHARED_TOKEN_PER_GLOBAL_STEP_SHARD_SET", + "minimum_global_step": 1, + "world_size": self._world_size, + "file_prefix": self._file_prefix, + "upload_batch_size": self._settings.upload_batch_size, + "pk_field": FEATURE_STORE_PK_FIELD, + "sk_field": FEATURE_STORE_SK_FIELD, + "value_field": FEATURE_STORE_VALUE_FIELD, + "key_dtype": "INT64", + "dynamic_key_encoding": "UINT64_BIT_PATTERN_IN_SIGNED_INT64", + "value_dtype": "FLOAT32", + "operation": DELTA_OPERATION_UPSERT, + "invalid_key": SPARSE_EMBEDDING_INVALID_KEY, + "embedding_dimensions": dict(sorted(self._embedding_dimensions.items())), + } + contract_bytes = json.dumps( + contract, sort_keys=True, separators=(",", ":") + ).encode("utf-8") + self._contract = contract + self._contract_hash = hashlib.sha256(contract_bytes).hexdigest() + self._state_dir = _feature_store_state_dir(self._settings, self._output_dir) + # The lock file needs its directory to exist. Directory creation itself + # contains no remote-write state, but every new parent entry is fsynced so + # a later durable attempt journal cannot disappear after a host crash. + _durable_makedirs(self._state_dir) + self._snapshot_root = os.path.join(self._state_dir, "snapshots") + self._error_marker_path = os.path.join(self._state_dir, "last_error.json") + self._writer_lock: Optional[BinaryIO] = None + self._contract_path = os.path.join(self._state_dir, "contract.json") + self._journal_initialized = False + self._committed_global_step = -1 + self._last_publish_ts = 0 + + @property + def state_dir(self) -> str: + """Directory containing credential-free manifests and success markers.""" + return self._state_dir + + def start(self) -> None: + """Start the single background worker and discover restart work.""" + with self._condition: + if self._started: + return + if self._closed: + raise RuntimeError("FeatureStoreDeltaUploader is already closed") + self._raise_if_failed_locked() + self._acquire_writer_lock() + try: + self._initialize_journal() + self._cleanup_stale_snapshot_staging() + self._cleanup_committed_snapshots() + self._add_discovered_steps_locked() + # Validate or provision the remote DynamicEmbedding FeatureView + # synchronously. Training must not start while the target is + # missing, has the wrong schema, or lacks the configured version. + self._get_view() + self._clear_error_marker() + self._started = True + self._worker = threading.Thread( + target=self._run, + name="tzrec-feature-store-delta-uploader", + daemon=True, + ) + self._worker.start() + except BaseException: + self._started = False + # A later retry must reconcile again after reacquiring the lock; + # another process may have advanced the journal in between. + self._journal_initialized = False + self._reset_view(suppress_errors=True) + self._release_writer_lock() + raise + logger.info( + "FeatureStore delta uploader started: project=%s feature_view=%s " + "version=%s", + self._settings.project_name, + self._settings.feature_view_name, + self._settings.version, + ) + + def submit(self, global_step: int) -> None: + """Enqueue a durably written rank-zero shard with bounded back-pressure.""" + global_step = int(global_step) + if global_step <= 0: + raise ValueError("FeatureStore delta global_step must be > 0") + with self._condition: + self._raise_if_failed_locked() + if not self._started: + raise RuntimeError( + "FeatureStoreDeltaUploader.start() must be called before submit()" + ) + if self._closing or self._closed: + raise RuntimeError("cannot submit to a closing FeatureStore uploader") + if self._is_committed(global_step): + return + if global_step <= self._committed_global_step: + raise FeatureStoreUploadError( + "refusing an uncommitted delta step older than the local " + "FeatureStore committed watermark" + ) + while ( + global_step not in self._pending + and len(self._pending) >= self._settings.max_pending_steps + ): + self._condition.wait(self._settings.poll_interval_secs) + self._raise_if_failed_locked() + self._pending.setdefault(global_step, time.monotonic()) + self._condition.notify_all() + + def check_error(self) -> None: + """Surface a background failure at a safe training-thread boundary.""" + with self._condition: + self._raise_if_failed_locked() + + def close(self, raise_on_error: bool = True, drain: bool = True) -> None: + """Close the worker, draining only during a normal training shutdown.""" + with self._condition: + if self._closed: + if raise_on_error: + self._raise_if_failed_locked() + return + if not self._started: + self._closed = True + return + self._closing = True + self._aborting = not drain + self._condition.notify_all() + worker = self._worker + + if drain and worker is not None: + worker.join(timeout=self._settings.shutdown_timeout_secs) + if worker.is_alive(): + timeout_error = FeatureStoreUploadError( + "FeatureStore uploader did not drain before shutdown timeout; " + "parquet outbox files were retained for restart" + ) + with self._condition: + if self._error is None: + self._error = timeout_error + self._aborting = True + self._condition.notify_all() + + with self._condition: + self._closed = True + if raise_on_error: + self._raise_if_failed_locked() + + def _run(self) -> None: + current_step: Optional[int] = None + try: + while True: + with self._condition: + if self._aborting: + return + self._add_discovered_steps_locked() + if not self._pending: + if self._closing: + return + self._condition.wait(self._settings.poll_interval_secs) + continue + current_step = min(self._pending) + pending_since = self._pending[current_step] + + canonical_paths = self._expected_shard_paths(current_step) + manifest_exists = os.path.isfile(self._manifest_path(current_step)) + snapshot_paths = self._snapshot_paths(current_step) + snapshot_dir_exists = os.path.isdir(self._snapshot_dir(current_step)) + if manifest_exists or snapshot_dir_exists: + if not all(os.path.isfile(path) for path in snapshot_paths): + raise FeatureStoreUploadError( + "an uncommitted FeatureStore upload journal is missing " + "its durable shard snapshot; recovery cannot continue" + ) + else: + try: + snapshot_paths = self._snapshot_canonical_shards( + current_step, canonical_paths + ) + except _ShardSetNotReady: + snapshot_paths = None + + if snapshot_paths is None: + elapsed = time.monotonic() - pending_since + if elapsed >= self._settings.shard_wait_timeout_secs: + raise TimeoutError( + "timed out waiting for a complete same-generation " + "delta shard set" + ) + with self._condition: + self._condition.wait(self._settings.poll_interval_secs) + continue + + with self._condition: + if self._aborting: + return + # Hash, parse and upload only uploader-owned read-only snapshots. + # The canonical dump files may be atomically replaced or cleaned + # after this point without changing restart/replay semantics. + with ExitStack() as stack: + shard_sources = [ + stack.enter_context(open(path, "rb")) for path in snapshot_paths + ] + shard_descriptions = self._describe_shards( + snapshot_paths, shard_sources + ) + self._validate_dump_generation(shard_descriptions) + records = self._load_records( + current_step, snapshot_paths, shard_sources + ) + manifest = self._load_or_create_manifest( + current_step, + snapshot_paths, + len(records), + shard_descriptions=shard_descriptions, + ) + summary = self._upload_with_retries(current_step, records, manifest) + # Rehashing can scan large shards. Keep it outside the + # condition so close(drain=False) can publish the abort bit + # immediately; the small durable commit remains serialized. + self._validate_shard_identities( + snapshot_paths, manifest["shards"], shard_sources + ) + with self._condition: + if self._aborting: + return + # Keep the condition locked through the local commit so + # aborting close cannot return before success-state writes. + self._commit_success(current_step, manifest, summary) + self._pending.pop(current_step, None) + self._condition.notify_all() + self._reclaim_snapshot(current_step) + current_step = None + except _UploadAborted: + return + except BaseException as exc: + step_context = ( + f" at global_step={current_step}" if current_step is not None else "" + ) + if isinstance(exc, FeatureStoreUploadError): + safe_error = FeatureStoreUploadError( + f"{exc}{step_context}; parquet outbox files were retained" + ) + else: + safe_error = FeatureStoreUploadError( + f"FeatureStore delta upload failed{step_context} " + f"({type(exc).__name__}); parquet outbox files were retained" + ) + try: + _atomic_write_json( + self._error_marker_path, + { + "schema_version": 1, + "project_name": self._settings.project_name, + "feature_view_name": self._settings.feature_view_name, + "version": self._settings.version, + "contract_hash": self._contract_hash, + "error": str(safe_error), + }, + ) + except BaseException as marker_error: + logger.error( + "Failed to persist FeatureStore upload failure marker (%s).", + type(marker_error).__name__, + ) + with self._condition: + self._error = safe_error + self._condition.notify_all() + logger.error("%s", safe_error) + finally: + self._reset_view(suppress_errors=True) + self._release_writer_lock() + + def _acquire_writer_lock(self) -> None: + if self._writer_lock is not None: + return + lock_path = os.path.join(self._state_dir, "writer.lock") + lock_file = open(lock_path, "a+b") + try: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as exc: + lock_file.close() + raise FeatureStoreUploadError( + "another process is already publishing this FeatureStore target " + "from the configured output_dir" + ) from exc + self._writer_lock = lock_file + + def _release_writer_lock(self) -> None: + lock_file = self._writer_lock + self._writer_lock = None + if lock_file is None: + return + if self._writer_quiescence_failed: + # close(wait=True) is the only SDK proof that no asynchronous write + # remains in flight. If it fails, retain the flock until process exit + # so another local writer cannot overlap an indeterminate old request. + _POISONED_WRITER_LOCKS.append(lock_file) + logger.error( + "FeatureStore writer lock retained until process exit because " + "SDK writer quiescence could not be confirmed." + ) + return + try: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) + finally: + lock_file.close() + + def _initialize_journal(self) -> None: + """Validate and reload all timestamp fences while holding the writer lock.""" + if self._writer_lock is None: + raise RuntimeError( + "FeatureStore upload journal must be initialized under writer lock" + ) + if os.path.isfile(self._contract_path): + if _read_json(self._contract_path) != self._contract: + raise ValueError( + "FeatureStore upload contract changed for an existing remote " + "target; provision and configure a new immutable version" + ) + else: + _atomic_write_json(self._contract_path, self._contract) + + # A previously constructed uploader may have waited behind another + # process. Always reconcile and reload after flock acquisition so it + # cannot allocate timestamps from stale constructor-time state. + self._reconcile_committed_state() + self._committed_global_step = self._load_committed_global_step() + self._last_publish_ts = self._load_latest_publish_ts() + self._journal_initialized = True + + def _cleanup_stale_snapshot_staging(self) -> None: + if not os.path.isdir(self._snapshot_root): + return + pattern = re.compile(r"^step_\d+\.tmp-") + for name in os.listdir(self._snapshot_root): + path = os.path.join(self._snapshot_root, name) + if pattern.match(name) and os.path.isdir(path): + shutil.rmtree(path) + _fsync_parent_directory(path) + + def _cleanup_committed_snapshots(self) -> None: + """Retry best-effort reclamation left behind after a committed upload.""" + if not os.path.isdir(self._snapshot_root): + return + pattern = re.compile(r"^step_(\d+)$") + for name in os.listdir(self._snapshot_root): + match = pattern.match(name) + if match is None: + continue + global_step = int(match.group(1)) + if self._is_committed(global_step): + self._reclaim_snapshot(global_step) + + def _clear_error_marker(self) -> None: + if not os.path.exists(self._error_marker_path): + return + os.remove(self._error_marker_path) + _fsync_parent_directory(self._error_marker_path) + + def _raise_if_failed_locked(self) -> None: + if self._error is not None: + raise self._error + + def _raise_if_aborting(self) -> None: + with self._condition: + if self._aborting: + raise _UploadAborted() + + def _add_discovered_steps_locked(self) -> None: + remaining_capacity = self._settings.max_pending_steps - len(self._pending) + if remaining_capacity <= 0: + return + for step in sorted(self._discover_steps()): + if remaining_capacity <= 0: + break + if step <= 0: + raise ValueError( + "found invalid FeatureStore delta outbox global_step; " + "global_step must be > 0" + ) + if self._is_committed(step): + continue + if step <= self._committed_global_step: + raise ValueError( + "found an uncommitted delta step older than the local " + "FeatureStore committed watermark" + ) + if step not in self._pending: + self._pending[step] = time.monotonic() + remaining_capacity -= 1 + + def _discover_steps(self) -> Iterable[int]: + if not os.path.isdir(self._output_dir): + return [] + steps = set() + manifest_pattern = re.compile(r"^step_(\d+)\.manifest\.json$") + for name in os.listdir(self._state_dir): + match = manifest_pattern.match(name) + if match: + steps.add(int(match.group(1))) + if os.path.isdir(self._snapshot_root): + snapshot_pattern = re.compile(r"^step_(\d+)$") + for name in os.listdir(self._snapshot_root): + match = snapshot_pattern.match(name) + if match and os.path.isdir(os.path.join(self._snapshot_root, name)): + steps.add(int(match.group(1))) + if self._world_size == 1: + pattern = re.compile( + rf"^{re.escape(self._file_prefix)}_step_(\d+)\.parquet$" + ) + for name in os.listdir(self._output_dir): + match = pattern.match(name) + if match: + steps.add(int(match.group(1))) + else: + pattern = re.compile(r"^step_(\d+)$") + for name in os.listdir(self._output_dir): + match = pattern.match(name) + if match is None or not os.path.isdir( + os.path.join(self._output_dir, name) + ): + continue + step = int(match.group(1)) + # A partial shard set for this exact contract is real pending + # work and must time out fail-safe. Empty dirs and shards from + # another prefix/world-size contract are ignored. + if any( + os.path.isfile(path) for path in self._expected_shard_paths(step) + ): + steps.add(step) + return steps + + def _snapshot_dir(self, global_step: int) -> str: + return os.path.join(self._snapshot_root, f"step_{global_step}") + + def _snapshot_paths(self, global_step: int) -> List[str]: + snapshot_dir = self._snapshot_dir(global_step) + return [ + os.path.join(snapshot_dir, f"rank_{rank}.parquet") + for rank in range(self._world_size) + ] + + def _expected_shard_paths(self, global_step: int) -> List[str]: + if self._world_size == 1: + return [ + os.path.join( + self._output_dir, + f"{self._file_prefix}_step_{global_step}.parquet", + ) + ] + step_dir = os.path.join(self._output_dir, f"step_{global_step}") + return [ + os.path.join( + step_dir, + f"{self._file_prefix}_step_{global_step}_rank_{rank}" + f"_of_{self._world_size}.parquet", + ) + for rank in range(self._world_size) + ] + + def _manifest_path(self, global_step: int) -> str: + return os.path.join(self._state_dir, f"step_{global_step}.manifest.json") + + def _success_path(self, global_step: int) -> str: + return os.path.join(self._state_dir, f"step_{global_step}._FS_SUCCESS.json") + + def _is_committed(self, global_step: int) -> bool: + if global_step <= 0: + raise ValueError("FeatureStore delta global_step must be > 0") + path = self._success_path(global_step) + if not os.path.isfile(path): + return False + success = _read_json(path) + expected = { + "global_step": global_step, + "version": self._settings.version, + "contract_hash": self._contract_hash, + } + for name, value in expected.items(): + if success.get(name) != value: + raise ValueError(f"FeatureStore success marker mismatch for {name}") + publish_ts = success.get("publish_ts") + if type(publish_ts) is not int or publish_ts <= 0: + raise ValueError("FeatureStore success marker has invalid publish_ts") + manifest_path = self._manifest_path(global_step) + if not os.path.isfile(manifest_path): + raise ValueError("FeatureStore success marker is missing its manifest") + manifest = _read_json(manifest_path) + if success.get("manifest_digest") != _json_digest(manifest): + raise ValueError("FeatureStore success marker manifest digest mismatch") + shards = success.get("shards") + if not isinstance(shards, list) or shards != manifest.get("shards"): + raise ValueError("FeatureStore success marker has invalid shards") + return True + + def _reconcile_committed_state(self) -> None: + """Repair a crash between atomic success and committed-state writes.""" + pattern = re.compile(r"^step_(\d+)\._FS_SUCCESS\.json$") + latest_success: Optional[Dict[str, Any]] = None + for name in os.listdir(self._state_dir): + match = pattern.match(name) + if match is None: + continue + step = int(match.group(1)) + if not self._is_committed(step): + continue + success = _read_json(self._success_path(step)) + if latest_success is None or step > int(latest_success["global_step"]): + latest_success = success + if latest_success is None: + return + + committed_path = os.path.join(self._state_dir, "committed.json") + committed_step = -1 + if os.path.isfile(committed_path): + committed_step = int( + _read_json(committed_path).get("committed_global_step", -1) + ) + latest_step = int(latest_success["global_step"]) + if committed_step >= latest_step: + return + committed = { + "schema_version": 1, + "project_name": self._settings.project_name, + "feature_view_name": self._settings.feature_view_name, + "version": self._settings.version, + "committed_global_step": latest_step, + "publish_ts": int(latest_success["publish_ts"]), + "contract_hash": self._contract_hash, + } + _atomic_write_json(committed_path, committed) + + def _load_latest_publish_ts(self) -> int: + latest = 0 + committed_path = os.path.join(self._state_dir, "committed.json") + if os.path.isfile(committed_path): + latest = max(latest, int(_read_json(committed_path).get("publish_ts", 0))) + if os.path.isdir(self._state_dir): + for name in os.listdir(self._state_dir): + if not name.endswith(".manifest.json"): + continue + manifest = _read_json(os.path.join(self._state_dir, name)) + for attempt in manifest.get("attempts", []): + latest = max(latest, int(attempt.get("range_end", 0))) + return latest + + def _load_committed_global_step(self) -> int: + path = os.path.join(self._state_dir, "committed.json") + if not os.path.isfile(path): + return -1 + global_step = int(_read_json(path).get("committed_global_step", -1)) + if global_step == 0 or global_step < -1: + raise ValueError( + "FeatureStore committed global_step must be -1 or greater than 0" + ) + return global_step + + @staticmethod + def _file_identity(stat_result: os.stat_result) -> Dict[str, int]: + return { + "device": int(stat_result.st_dev), + "inode": int(stat_result.st_ino), + "size_bytes": int(stat_result.st_size), + "mtime_ns": int(stat_result.st_mtime_ns), + "ctime_ns": int(stat_result.st_ctime_ns), + } + + def _snapshot_canonical_shards( + self, global_step: int, canonical_paths: List[str] + ) -> Optional[List[str]]: + """Copy one stable, same-generation shard set into the recovery journal.""" + if not all(os.path.isfile(path) for path in canonical_paths): + return None + _durable_makedirs(self._snapshot_root) + snapshot_dir = self._snapshot_dir(global_step) + if os.path.isdir(snapshot_dir): + snapshot_paths = self._snapshot_paths(global_step) + if not all(os.path.isfile(path) for path in snapshot_paths): + raise FeatureStoreUploadError( + "durable FeatureStore shard snapshot is incomplete" + ) + return snapshot_paths + + staging_dir = f"{snapshot_dir}.tmp-{os.getpid()}-{uuid.uuid4().hex}" + os.mkdir(staging_dir) + _fsync_parent_directory(staging_dir) + staging_paths = [ + os.path.join(staging_dir, f"rank_{rank}.parquet") + for rank in range(self._world_size) + ] + try: + for source_path, snapshot_path in zip(canonical_paths, staging_paths): + with open(source_path, "rb") as source: + before_identity = self._file_identity(os.fstat(source.fileno())) + with open(snapshot_path, "xb") as output: + source.seek(0) + shutil.copyfileobj(source, output, length=8 * 1024 * 1024) + output.flush() + os.fsync(output.fileno()) + after_identity = self._file_identity(os.fstat(source.fileno())) + if before_identity != after_identity: + raise _ShardSetNotReady( + "canonical shard changed while creating snapshot" + ) + os.chmod(snapshot_path, 0o400) + + # Validate the generation before atomically publishing the directory. + descriptions = self._describe_shards(staging_paths) + self._validate_dump_generation(descriptions) + _fsync_parent_directory(staging_paths[0]) + os.replace(staging_dir, snapshot_dir) + _fsync_parent_directory(snapshot_dir) + return self._snapshot_paths(global_step) + except BaseException: + shutil.rmtree(staging_dir, ignore_errors=True) + raise + + def _describe_shards( + self, + shard_paths: List[str], + shard_sources: Optional[List[BinaryIO]] = None, + ) -> List[Dict[str, Any]]: + if shard_sources is None: + with ExitStack() as stack: + opened_sources = [ + stack.enter_context(open(path, "rb")) for path in shard_paths + ] + return self._describe_shards(shard_paths, opened_sources) + if len(shard_paths) != len(shard_sources): + raise ValueError("delta shard paths and sources must have equal length") + + descriptions = [] + for path, source in zip(shard_paths, shard_sources): + before_identity = self._file_identity(os.fstat(source.fileno())) + source.seek(0) + parquet_file = pq.ParquetFile(source) + num_rows = int(parquet_file.metadata.num_rows) + metadata = parquet_file.schema_arrow.metadata or {} + schema_version = metadata.get(_SCHEMA_VERSION_METADATA_KEY) + dump_generation = metadata.get(DELTA_DUMP_GENERATION_METADATA_KEY) + if schema_version != DELTA_DUMP_SCHEMA_VERSION.encode("ascii"): + raise ValueError(f"unsupported delta dump schema version in {path}") + if not dump_generation: + raise ValueError(f"delta dump generation is missing in {path}") + try: + dump_generation_text = dump_generation.decode("ascii") + except UnicodeDecodeError as exc: + raise ValueError( + f"delta dump generation is not ASCII in {path}" + ) from exc + sha256 = _file_sha256(source) + after_identity = self._file_identity(os.fstat(source.fileno())) + try: + path_identity = self._file_identity(os.stat(path)) + except OSError as exc: + raise RuntimeError( + "delta shard snapshot path changed while being described" + ) from exc + if before_identity != after_identity or path_identity != after_identity: + raise RuntimeError("delta shard changed while building upload manifest") + descriptions.append( + { + "path": os.path.relpath(path, self._output_dir), + "size_bytes": after_identity["size_bytes"], + "num_rows": num_rows, + "sha256": sha256, + "delta_dump_schema_version": DELTA_DUMP_SCHEMA_VERSION, + "dump_generation": dump_generation_text, + } + ) + return descriptions + + @staticmethod + def _validate_dump_generation(descriptions: List[Dict[str, Any]]) -> str: + generations = { + description.get("dump_generation") for description in descriptions + } + if None in generations or "" in generations: + raise ValueError("delta shard set has invalid dump generation metadata") + if len(generations) != 1: + raise _ShardSetNotReady( + "delta shards from different dump generations are present" + ) + return str(next(iter(generations))) + + def _validate_shard_identities( + self, + shard_paths: List[str], + descriptions: List[Dict[str, Any]], + shard_sources: Optional[List[BinaryIO]] = None, + ) -> None: + if shard_sources is None: + with ExitStack() as stack: + opened_sources = [ + stack.enter_context(open(path, "rb")) for path in shard_paths + ] + self._validate_shard_identities( + shard_paths, descriptions, opened_sources + ) + return + if len(shard_paths) != len(descriptions) or ( + len(shard_paths) != len(shard_sources) + ): + raise ValueError("FeatureStore shard identity count mismatch") + for index, (path, description) in enumerate(zip(shard_paths, descriptions)): + expected_path = os.path.relpath(path, self._output_dir) + if description.get("path") != expected_path: + raise ValueError("FeatureStore shard identity metadata is invalid") + try: + source = shard_sources[index] + before_identity = self._file_identity(os.fstat(source.fileno())) + path_identity = self._file_identity(os.stat(path)) + sha256 = _file_sha256(source) + after_identity = self._file_identity(os.fstat(source.fileno())) + except OSError as exc: + raise RuntimeError( + "delta shard changed after upload snapshot was claimed" + ) from exc + if ( + before_identity != after_identity + or path_identity != after_identity + or after_identity["size_bytes"] != description.get("size_bytes") + or sha256 != description.get("sha256") + ): + raise RuntimeError( + "delta shard changed after upload snapshot was claimed" + ) + + def _load_or_create_manifest( + self, + global_step: int, + shard_paths: List[str], + record_count: int, + shard_descriptions: Optional[List[Dict[str, Any]]] = None, + ) -> Dict[str, Any]: + if not self._journal_initialized or self._writer_lock is None: + raise RuntimeError( + "FeatureStore upload manifests may only be changed under the " + "initialized writer journal lock" + ) + path = self._manifest_path(global_step) + expected_shards = ( + self._describe_shards(shard_paths) + if shard_descriptions is None + else shard_descriptions + ) + dump_generation = self._validate_dump_generation(expected_shards) + if os.path.isfile(path): + manifest = _read_json(path) + expected = { + "schema_version": 3, + "global_step": global_step, + "world_size": self._world_size, + "version": self._settings.version, + "contract_hash": self._contract_hash, + "record_count": record_count, + "dump_generation": dump_generation, + "shards": expected_shards, + } + for name, value in expected.items(): + if manifest.get(name) != value: + raise ValueError( + f"FeatureStore upload manifest mismatch for {name}" + ) + attempts = manifest.get("attempts") + if not isinstance(attempts, list): + raise ValueError("FeatureStore upload manifest has invalid attempts") + for attempt in attempts: + if ( + not isinstance(attempt, dict) + or type(attempt.get("range_start")) is not int + or type(attempt.get("range_end")) is not int + or attempt["range_start"] <= 0 + or attempt["range_end"] < attempt["range_start"] + ): + raise ValueError( + "FeatureStore upload manifest has an invalid ts range" + ) + self._last_publish_ts = max( + self._last_publish_ts, int(attempt["range_end"]) + ) + return manifest + + manifest = { + "schema_version": 3, + "global_step": global_step, + "world_size": self._world_size, + "project_name": self._settings.project_name, + "feature_view_name": self._settings.feature_view_name, + "version": self._settings.version, + "write_mode": FEATURE_STORE_WRITE_MODE, + "contract_hash": self._contract_hash, + "record_count": record_count, + "dump_generation": dump_generation, + "shards": expected_shards, + "attempts": [], + } + _atomic_write_json(path, manifest) + return manifest + + def _start_attempt( + self, + global_step: int, + record_count: int, + manifest: Dict[str, Any], + ) -> Dict[str, int]: + if not self._journal_initialized or self._writer_lock is None: + raise RuntimeError( + "FeatureStore timestamp ranges may only be reserved under the " + "initialized writer journal lock" + ) + batch_count = ( + record_count + self._settings.upload_batch_size - 1 + ) // self._settings.upload_batch_size + reserved_count = max(batch_count, 1) + range_start = max(int(self._clock_ms()), self._last_publish_ts + 1, 1) + attempt = { + "attempt_id": len(manifest["attempts"]) + 1, + "record_count": record_count, + "batch_count": batch_count, + "range_start": range_start, + "range_end": range_start + reserved_count - 1, + } + # Persist the whole range before any remote request. After a partial + # write or process crash, restart allocates a newer range and replays + # the entire step so a Processor Next-Ts cursor cannot miss late rows. + manifest["attempts"].append(attempt) + _atomic_write_json(self._manifest_path(global_step), manifest) + self._last_publish_ts = int(attempt["range_end"]) + return attempt + + def _upload_with_retries( + self, + global_step: int, + records: List[Tuple[str, int, np.ndarray]], + manifest: Dict[str, Any], + ) -> Dict[str, int]: + for local_attempt in range(1, self._settings.max_retries + 1): + self._raise_if_aborting() + attempt = self._start_attempt(global_step, len(records), manifest) + try: + if records: + summary = self._upload_records(records, attempt) + else: + # An empty step still validates the FeatureView schema and + # pre-provisioned target version before it can be committed. + self._get_view() + summary = { + "total_batches": 0, + "failed_batches": 0, + "total_records": 0, + "success_records": 0, + "failed_records": 0, + } + summary.update(attempt) + return summary + except _UploadAborted: + raise + except BaseException as exc: + self._reset_view() + if local_attempt >= self._settings.max_retries: + raise + logger.warning( + "FeatureStore delta upload attempt %s/%s failed at step %s " + "(%s); replaying the full step with a newer persisted ts range.", + local_attempt, + self._settings.max_retries, + global_step, + type(exc).__name__, + ) + if self._settings.retry_backoff_secs > 0: + with self._condition: + if self._aborting: + raise _UploadAborted() from None + self._condition.wait( + self._settings.retry_backoff_secs * local_attempt + ) + if self._aborting: + raise _UploadAborted() from None + raise AssertionError("unreachable FeatureStore retry state") + + def _wait_for_dynamic_embedding_view(self, project: Any) -> Any: + """Bounded re-get after a concurrent or partially completed create.""" + last_error: Optional[Exception] = None + for attempt in range(1, self._settings.max_retries + 1): + try: + view = project.get_dynamic_embedding_feature_view( + self._settings.feature_view_name + ) + except Exception as exc: + last_error = exc + else: + if view is not None: + return view + if ( + attempt < self._settings.max_retries + and self._settings.retry_backoff_secs > 0 + ): + time.sleep(self._settings.retry_backoff_secs * attempt) + if last_error is not None: + raise RuntimeError( + "DynamicEmbedding FeatureView did not become ready after creation" + ) from last_error + return None + + def _wait_for_feature_view_metadata(self, project: Any) -> Any: + """Bounded control-plane lookup used to validate the created resource.""" + last_error: Optional[Exception] = None + for attempt in range(1, self._settings.max_retries + 1): + try: + feature_view = project.get_feature_view( + self._settings.feature_view_name + ) + except Exception as exc: + last_error = exc + else: + if feature_view is not None: + return feature_view + if ( + attempt < self._settings.max_retries + and self._settings.retry_backoff_secs > 0 + ): + time.sleep(self._settings.retry_backoff_secs * attempt) + if last_error is not None: + raise RuntimeError( + "FeatureView control-plane metadata did not become ready" + ) from last_error + return None + + def _validate_feature_view_metadata(self, feature_view: Any) -> None: + """Validate immutable control-plane schema and provisioning settings.""" + actual_type = getattr(feature_view, "type", None) + if actual_type != "DynamicEmbedding": + raise RuntimeError( + "configured FeatureView exists with an incompatible type: " + f"expected='DynamicEmbedding', actual={actual_type!r}" + ) + actual_entity = getattr(feature_view, "feature_entity_name", None) + if actual_entity != self._settings.feature_entity_name: + raise RuntimeError( + "DynamicEmbedding FeatureView entity mismatch: " + f"expected={self._settings.feature_entity_name!r}, " + f"actual={actual_entity!r}" + ) + + expected_fields = { + FEATURE_STORE_PK_FIELD: ("STRING", {"PrimaryKey"}), + FEATURE_STORE_SK_FIELD: ("INT64", {"SubKey"}), + FEATURE_STORE_VALUE_FIELD: ("ARRAY", set()), + } + fields = getattr(feature_view, "fields_dict", None) + if not isinstance(fields, dict) or set(fields) != set(expected_fields): + actual_names = sorted(fields) if isinstance(fields, dict) else None + raise RuntimeError( + "DynamicEmbedding FeatureView field set mismatch: " + f"expected={sorted(expected_fields)}, actual={actual_names}" + ) + for name, (expected_type, expected_attributes) in expected_fields.items(): + field_info = fields[name] + if not isinstance(field_info, dict): + raise RuntimeError( + "DynamicEmbedding FeatureView has invalid field metadata for " + f"{name!r}" + ) + actual_field_type = field_info.get("Type") + attributes = field_info.get("Attributes", []) + if not isinstance(attributes, (list, tuple, set)): + raise RuntimeError( + "DynamicEmbedding FeatureView has invalid field attributes for " + f"{name!r}" + ) + actual_attributes = set(attributes) + if ( + actual_field_type != expected_type + or actual_attributes != expected_attributes + ): + raise RuntimeError( + "DynamicEmbedding FeatureView field contract mismatch for " + f"{name!r}: expected_type={expected_type!r}, " + f"actual_type={actual_field_type!r}, " + f"expected_attributes={sorted(expected_attributes)}, " + f"actual_attributes={sorted(actual_attributes)}" + ) + + summary = getattr(feature_view, "summary", None) + config_value = summary.get("Config") if isinstance(summary, dict) else None + try: + provisioning = ( + json.loads(config_value) + if isinstance(config_value, str) + else dict(config_value) + ) + except (TypeError, ValueError) as exc: + raise RuntimeError( + "DynamicEmbedding FeatureView has invalid provisioning config" + ) from exc + expected_provisioning = { + "ttl": self._settings.feature_view_ttl_secs, + "shard_count": self._settings.feature_view_shard_count, + "replication_count": self._settings.feature_view_replication_count, + } + actual_provisioning = { + name: provisioning.get(name) for name in expected_provisioning + } + if actual_provisioning != expected_provisioning: + raise RuntimeError( + "DynamicEmbedding FeatureView provisioning mismatch: " + f"expected={expected_provisioning}, actual={actual_provisioning}" + ) + + def _get_or_create_view(self, project: Any) -> Any: + """Return the configured DynamicEmbedding view, creating it if absent.""" + provisioned = False + view = project.get_dynamic_embedding_feature_view( + self._settings.feature_view_name + ) + if view is not None: + self._view = view + metadata = self._wait_for_feature_view_metadata(project) + if view is None and metadata is not None: + self._validate_feature_view_metadata(metadata) + view = self._wait_for_dynamic_embedding_view(project) + if view is None: + raise RuntimeError( + "configured DynamicEmbedding FeatureView exists but did not " + "become ready" + ) + self._view = view + elif view is None: + create_error: Optional[Exception] = None + try: + view = project.create_dynamic_embedding_feature_view( + name=self._settings.feature_view_name, + entity=self._settings.feature_entity_name, + pk_field_name=FEATURE_STORE_PK_FIELD, + sk_field_name=FEATURE_STORE_SK_FIELD, + embedding_field_name=FEATURE_STORE_VALUE_FIELD, + pk_field_type="STRING", + sk_field_type="INT64", + ttl=self._settings.feature_view_ttl_secs, + shard_count=self._settings.feature_view_shard_count, + replication_count=self._settings.feature_view_replication_count, + ) + provisioned = True + except Exception as exc: + # Another writer may have won the create race, or control-plane + # creation may have succeeded before SDK data-plane initialization + # failed. Re-get before declaring the operation failed. + create_error = exc + view = self._wait_for_dynamic_embedding_view(project) + if view is None: + error = RuntimeError( + "failed to create configured DynamicEmbedding FeatureView; " + "verify that feature_entity_name already exists" + ) + if create_error is not None: + raise error from create_error + raise error + self._view = view + metadata = self._wait_for_feature_view_metadata(project) + + if metadata is None: + # The specialized view owns the FeatureDB writer. Retain it before + # reporting a control-plane validation failure so start() can close + # it through the normal reset path. + self._view = view + raise RuntimeError( + "DynamicEmbedding FeatureView control-plane metadata was not found" + ) + self._view = view + self._validate_feature_view_metadata(metadata) + if provisioned: + logger.info( + "Created DynamicEmbedding FeatureView: project=%s entity=%s view=%s", + self._settings.project_name, + self._settings.feature_entity_name, + self._settings.feature_view_name, + ) + return view + + def _get_view(self) -> Any: + if self._view is not None: + return self._view + if self._client_factory is None: + try: + from feature_store_py import FeatureStoreClient + except ImportError as exc: + raise RuntimeError( + "feature_store_py is required when feature_store_config is set" + ) from exc + client_factory = FeatureStoreClient + else: + client_factory = self._client_factory + + kwargs = { + "access_key_id": self._settings.access_key_id, + "access_key_secret": self._settings.access_key_secret, + "region": self._settings.region or None, + "endpoint": self._settings.endpoint or None, + "security_token": self._settings.security_token or None, + "featuredb_username": self._settings.featuredb_username or None, + "featuredb_password": self._settings.featuredb_password or None, + } + client = client_factory(**kwargs) + project = client.get_project(self._settings.project_name) + if project is None: + raise RuntimeError("configured FeatureStore project was not found") + view = self._get_or_create_view(project) + # Own the SDK writer before validating the remote contract so every + # failure path is closed by the retry/reset logic. + self._view = view + actual_fields = (view.pk_field, view.sk_field, view.embedding_field) + expected_fields = ( + FEATURE_STORE_PK_FIELD, + FEATURE_STORE_SK_FIELD, + FEATURE_STORE_VALUE_FIELD, + ) + if actual_fields != expected_fields: + raise RuntimeError( + "DynamicEmbedding FeatureView schema mismatch: " + f"expected={expected_fields}, actual={actual_fields}" + ) + sdk_batch_size = getattr(view, "_batch_size", FEATURE_STORE_SDK_BATCH_SIZE) + if ( + type(sdk_batch_size) is not int + or sdk_batch_size < self._settings.upload_batch_size + ): + raise RuntimeError( + "FeatureStore SDK batch_size is smaller than the configured outer " + "batch; one publish timestamp could span multiple HTTP requests" + ) + version_info = view.list_versions() + versions = version_info.get("versions", []) if version_info else [] + if not any( + isinstance(item, dict) and item.get("version") == self._settings.version + for item in versions + ): + raise FeatureStoreUploadError( + "configured FeatureStore version does not exist; provision it " + "before enabling delta MERGE upload" + ) + return view + + def _reset_view(self, suppress_errors: bool = False) -> None: + view = self._view + self._view = None + if view is not None: + try: + view.close(wait=True) + except BaseException as exc: + self._writer_quiescence_failed = True + close_error = FeatureStoreUploadError( + "FeatureStore SDK writer close failed; retry was aborted to " + "avoid overlapping requests" + ) + if not suppress_errors: + raise close_error from None + + with self._condition: + new_error = self._error is None + if new_error: + self._error = close_error + self._condition.notify_all() + if new_error: + try: + _atomic_write_json( + self._error_marker_path, + { + "schema_version": 1, + "project_name": self._settings.project_name, + "feature_view_name": self._settings.feature_view_name, + "version": self._settings.version, + "contract_hash": self._contract_hash, + "error": str(close_error), + }, + ) + except BaseException as marker_error: + logger.error( + "Failed to persist FeatureStore upload failure marker " + "after SDK close failure (%s).", + type(marker_error).__name__, + ) + logger.error( + "Failed to close FeatureStore SDK writer cleanly (%s); the " + "local writer lock will be retained until process exit.", + type(exc).__name__, + ) + + def _upload_records( + self, + records: List[Tuple[str, int, np.ndarray]], + attempt: Mapping[str, int], + ) -> Dict[str, int]: + view = self._get_view() + aggregate = { + "total_batches": 0, + "failed_batches": 0, + "total_records": 0, + "success_records": 0, + "failed_records": 0, + } + try: + for batch_index, offset in enumerate( + range(0, len(records), self._settings.upload_batch_size) + ): + self._raise_if_aborting() + chunk = records[offset : offset + self._settings.upload_batch_size] + payload = [ + { + FEATURE_STORE_PK_FIELD: embedding_name, + FEATURE_STORE_SK_FIELD: key_id, + FEATURE_STORE_VALUE_FIELD: embedding, + } + for embedding_name, key_id, embedding in chunk + ] + view.write_features( + data=payload, + version=self._settings.version, + write_mode=FEATURE_STORE_WRITE_MODE, + ts=int(attempt["range_start"]) + batch_index, + ) + summary = view.write_flush() + self._validate_flush_summary(summary, len(chunk)) + for name in aggregate: + aggregate[name] += int(summary[name]) + self._raise_if_aborting() + except BaseException: + # A write_features() call can enqueue part of its work before raising. + # Drain it so a retry never mixes futures from two attempts. + try: + view.write_flush() + except BaseException: + pass + raise + return aggregate + + @staticmethod + def _validate_flush_summary(summary: Any, expected_records: int) -> None: + required = { + "total_batches", + "failed_batches", + "total_records", + "success_records", + "failed_records", + } + if not isinstance(summary, dict) or not required.issubset(summary): + raise RuntimeError("FeatureStore write_flush returned an invalid summary") + if ( + int(summary["total_batches"]) != 1 + or int(summary["failed_batches"]) != 0 + or int(summary["failed_records"]) != 0 + or int(summary["success_records"]) != int(summary["total_records"]) + or int(summary["total_records"]) != expected_records + ): + raise RuntimeError("FeatureStore write_flush reported incomplete writes") + + def _load_records( + self, + global_step: int, + shard_paths: List[str], + shard_sources: Optional[List[BinaryIO]] = None, + ) -> List[Tuple[str, int, np.ndarray]]: + if shard_sources is None: + with ExitStack() as stack: + opened_sources = [ + stack.enter_context(open(path, "rb")) for path in shard_paths + ] + return self._load_records(global_step, shard_paths, opened_sources) + if len(shard_paths) != len(shard_sources): + raise ValueError("delta shard paths and sources must have equal length") + + # Exact duplicates are harmless; conflicting duplicates are corruption. + # The final sort makes retry/restart batch boundaries deterministic. + records: Dict[Tuple[str, int], np.ndarray] = {} + for expected_rank, (path, source) in enumerate(zip(shard_paths, shard_sources)): + source.seek(0) + table = pq.read_table(source) + source.seek(0) + self._validate_parquet_schema(table.schema, path) + for row in table.to_pylist(): + if int(row["global_step"]) != global_step: + raise ValueError("delta shard global_step mismatch") + if int(row["rank"]) != expected_rank: + raise ValueError("delta shard rank mismatch") + if int(row["world_size"]) != self._world_size: + raise ValueError("delta shard world_size mismatch") + if row["operation"] != DELTA_OPERATION_UPSERT: + raise ValueError("only UPSERT delta operations are supported") + if row["embedding_role"] not in SPARSE_EMBEDDING_ROLES: + raise ValueError("delta shard has an invalid embedding role") + + embedding_name = row["embedding_name"] + if not isinstance(embedding_name, str) or not embedding_name: + raise ValueError("delta shard embedding_name must not be empty") + if embedding_name not in self._embedding_dimensions: + raise ValueError( + "delta shard embedding_name is absent from model contract: " + f"{embedding_name!r}" + ) + key_id = row["key_id"] + if isinstance(key_id, bool) or not isinstance(key_id, int): + raise ValueError("delta shard key_id must be a signed int64") + if key_id == SPARSE_EMBEDDING_INVALID_KEY: + raise ValueError( + "delta shard key_id=-1 is reserved as the Processor/" + "NvEmbeddings invalid-key sentinel" + ) + embedding = np.asarray(row["embedding"], dtype=np.float32) + expected_dimension = self._embedding_dimensions[embedding_name] + if embedding.ndim != 1 or embedding.size != expected_dimension: + raise ValueError( + f"delta embedding dimension mismatch for {embedding_name!r}: " + f"expected={expected_dimension}, actual_shape={embedding.shape}" + ) + if not np.isfinite(embedding).all(): + raise ValueError("delta embedding contains NaN or Inf") + record_key = (embedding_name, key_id) + previous = records.get(record_key) + if previous is not None: + if not np.array_equal(previous, embedding): + raise ValueError( + "conflicting duplicate delta row for " + f"embedding_name={embedding_name!r}, key_id={key_id}" + ) + continue + records[record_key] = embedding + + return [ + (embedding_name, key_id, records[(embedding_name, key_id)]) + for embedding_name, key_id in sorted(records) + ] + + @staticmethod + def _validate_parquet_schema(schema: pa.Schema, path: str) -> None: + metadata = schema.metadata or {} + if metadata.get( + _SCHEMA_VERSION_METADATA_KEY + ) != DELTA_DUMP_SCHEMA_VERSION.encode("ascii"): + raise ValueError(f"unsupported delta dump schema version in {path}") + for field_name, field_type in _REQUIRED_PARQUET_FIELDS.items(): + index = schema.get_field_index(field_name) + if index < 0 or schema.field(index).type != field_type: + raise ValueError( + f"delta dump schema mismatch for field {field_name!r} in {path}" + ) + + def _commit_success( + self, + global_step: int, + manifest: Mapping[str, Any], + summary: Mapping[str, int], + ) -> None: + manifest_digest = _json_digest(manifest) + publish_ts = int(summary["range_end"]) + success = { + "schema_version": 2, + "global_step": global_step, + "project_name": self._settings.project_name, + "feature_view_name": self._settings.feature_view_name, + "version": self._settings.version, + "attempt_id": int(summary["attempt_id"]), + "range_start": int(summary["range_start"]), + "range_end": publish_ts, + "publish_ts": publish_ts, + "write_mode": FEATURE_STORE_WRITE_MODE, + "contract_hash": self._contract_hash, + "manifest_digest": manifest_digest, + "shards": manifest["shards"], + "total_records": int(summary["total_records"]), + "success_records": int(summary["success_records"]), + } + _atomic_write_json(self._success_path(global_step), success) + committed = { + "schema_version": 1, + "project_name": self._settings.project_name, + "feature_view_name": self._settings.feature_view_name, + "version": self._settings.version, + "committed_global_step": global_step, + "publish_ts": publish_ts, + "contract_hash": self._contract_hash, + "manifest_digest": manifest_digest, + } + _atomic_write_json(os.path.join(self._state_dir, "committed.json"), committed) + self._committed_global_step = global_step + + logger.info( + "FeatureStore delta upload committed: step=%s version=%s records=%s ts=%s", + global_step, + self._settings.version, + summary["success_records"], + publish_ts, + ) + + def _reclaim_snapshot(self, global_step: int) -> None: + """Best-effort snapshot reclamation after durable success publication.""" + snapshot_dir = self._snapshot_dir(global_step) + if not os.path.isdir(snapshot_dir): + return + try: + shutil.rmtree(snapshot_dir) + _fsync_parent_directory(snapshot_dir) + except OSError as exc: + # Success marker + manifest are already durable. Snapshot cleanup is + # storage reclamation only and must not turn a committed write into + # a reported failure. + logger.warning( + "Failed to reclaim committed FeatureStore shard snapshot (%s).", + type(exc).__name__, + ) diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py new file mode 100644 index 000000000..1c03e9dac --- /dev/null +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -0,0 +1,1573 @@ +# Copyright (c) 2025, 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. + +import json +import os +import tempfile +import threading +import unittest +from contextlib import contextmanager +from unittest import mock + +import pyarrow as pa +import pyarrow.parquet as pq +from google.protobuf.descriptor import FieldDescriptor + +from tzrec.protos.train_pb2 import FeatureStoreConfig +from tzrec.utils.delta_embedding_dump import _DELTA_DUMP_SCHEMA +from tzrec.utils.feature_store_delta_uploader import ( + DELTA_DUMP_GENERATION_METADATA_KEY, + FeatureStoreDeltaUploader, + FeatureStoreUploadError, + FeatureStoreUploadSettings, + feature_store_delta_file_prefix, +) + +_TEST_DUMP_GENERATION = "00112233445566778899aabbccddeeff" + + +@contextmanager +def _initialized_journal(uploader: FeatureStoreDeltaUploader): + uploader._acquire_writer_lock() + try: + uploader._initialize_journal() + yield + finally: + uploader._release_writer_lock() + + +def _schema_with_generation(generation: str = _TEST_DUMP_GENERATION) -> pa.Schema: + metadata = dict(_DELTA_DUMP_SCHEMA.metadata or {}) + metadata[DELTA_DUMP_GENERATION_METADATA_KEY] = generation.encode("ascii") + return _DELTA_DUMP_SCHEMA.with_metadata(metadata) + + +def _feature_store_config(**overrides) -> FeatureStoreConfig: + config = FeatureStoreConfig( + region="cn-test", + access_key_id="ak-id-secret", + access_key_secret="ak-value-secret", + security_token="sts-secret", + featuredb_username="fdb-user-secret", + featuredb_password="fdb-password-secret", + project_name="project_a", + feature_entity_name="embedding_entity", + feature_view_name="shared_embeddings", + version="model_a@export_1", + upload_batch_size=2, + max_retries=1, + retry_backoff_secs=0, + shard_wait_timeout_secs=2, + shutdown_timeout_secs=5, + max_pending_steps=8, + poll_interval_secs=1, + ) + for name, value in overrides.items(): + setattr(config, name, value) + return config + + +def _row( + step: int, + rank: int, + key_id: int, + values, + name: str = "user_emb", + world_size: int = 1, +): + return { + "global_step": step, + "rank": rank, + "world_size": world_size, + "embedding_name": name, + "embedding_role": "ebc", + "feature_name": "user_id", + "table_fqn": f"model.ebc.embedding_bags.{name}.weight", + "key_id": key_id, + "embedding": values, + "operation": "UPSERT", + "source": "model_delta_tracker", + } + + +def _write_single_shard( + output_dir: str, + step: int, + rows, + generation: str = _TEST_DUMP_GENERATION, + file_prefix=None, +) -> str: + if file_prefix is None: + file_prefix = feature_store_delta_file_prefix(_feature_store_config(), "delta") + path = os.path.join(output_dir, f"{file_prefix}_step_{step}.parquet") + schema = _schema_with_generation(generation) + table = pa.Table.from_pylist(rows, schema=schema) if rows else schema.empty_table() + pq.write_table(table, path) + return path + + +def _write_rank_shard( + output_dir: str, + step: int, + rank: int, + world_size: int, + rows, + generation: str = _TEST_DUMP_GENERATION, + file_prefix=None, +) -> str: + if file_prefix is None: + file_prefix = feature_store_delta_file_prefix(_feature_store_config(), "delta") + step_dir = os.path.join(output_dir, f"step_{step}") + os.makedirs(step_dir, exist_ok=True) + path = os.path.join( + step_dir, + f"{file_prefix}_step_{step}_rank_{rank}_of_{world_size}.parquet", + ) + table = pa.Table.from_pylist(rows, schema=_schema_with_generation(generation)) + pq.write_table(table, path) + return path + + +class _FakeView: + pk_field = "embedding_name" + sk_field = "key_id" + embedding_field = "embedding" + + def __init__(self, summaries=None, versions=None, close_error=None): + self.calls = [] + self.closed = [] + self._summaries = list(summaries or []) + self._versions = ["model_a@export_1"] if versions is None else list(versions) + self._close_error = close_error + self._batch_size = 1000 + self._last_size = 0 + + def write_features(self, **kwargs): + self.calls.append(kwargs) + self._last_size = len(kwargs["data"]) + + def write_flush(self): + if self._summaries: + return self._summaries.pop(0) + return { + "total_batches": 1, + "failed_batches": 0, + "total_records": self._last_size, + "success_records": self._last_size, + "failed_records": 0, + "errors": [], + } + + def list_versions(self): + return { + "default_version": "", + "versions": [{"version": version} for version in self._versions], + } + + def close(self, wait=True): + self.closed.append(wait) + if self._close_error is not None: + raise self._close_error + + +class _BlockingView(_FakeView): + def __init__(self): + super().__init__() + self.flush_started = threading.Event() + self.release_flush = threading.Event() + self.close_finished = threading.Event() + + def write_flush(self): + self.flush_started.set() + self.release_flush.wait(timeout=5) + return super().write_flush() + + def close(self, wait=True): + super().close(wait=wait) + self.close_finished.set() + + +class _FakeGenericFeatureView: + def __init__( + self, + *, + feature_view_type="DynamicEmbedding", + entity="embedding_entity", + fields=None, + provisioning=None, + ): + self.type = feature_view_type + self.feature_entity_name = entity + self.fields_dict = ( + { + "embedding_name": { + "Name": "embedding_name", + "Type": "STRING", + "Attributes": ["PrimaryKey"], + }, + "key_id": { + "Name": "key_id", + "Type": "INT64", + "Attributes": ["SubKey"], + }, + "embedding": { + "Name": "embedding", + "Type": "ARRAY", + "Attributes": [], + }, + } + if fields is None + else fields + ) + if provisioning is None: + provisioning = { + "ttl": 1296000, + "shard_count": 20, + "replication_count": 1, + } + self.summary = {"Config": json.dumps(provisioning)} + + +class _FakeProject: + def __init__( + self, + view, + *, + created_view=None, + generic_view=None, + create_error=None, + view_after_create_error=None, + ): + self._view = view + self._created_view = created_view + self._generic_view = generic_view + self._create_error = create_error + self._view_after_create_error = view_after_create_error + self.dynamic_get_calls = [] + self.generic_get_calls = [] + self.create_calls = [] + + def get_dynamic_embedding_feature_view(self, name): + self.dynamic_get_calls.append(name) + return self._view + + def get_feature_view(self, name): + self.generic_get_calls.append(name) + if self._generic_view is None and self._view is not None: + self._generic_view = _FakeGenericFeatureView() + return self._generic_view + + def create_dynamic_embedding_feature_view(self, **kwargs): + self.create_calls.append(kwargs) + if self._create_error is not None: + self._view = self._view_after_create_error + raise self._create_error + self._view = self._created_view or _FakeView() + self._generic_view = _FakeGenericFeatureView() + return self._view + + +class _FakeClient: + def __init__(self, project, kwargs): + self._project = project + self.kwargs = kwargs + + def get_project(self, name): + return self._project + + +class _FakeClientFactory: + def __init__(self, view, **project_kwargs): + self.view = view + self.calls = [] + self.project = _FakeProject(view, **project_kwargs) + + def __call__(self, **kwargs): + self.calls.append(kwargs) + return _FakeClient(self.project, kwargs) + + +class _SequencedClientFactory: + def __init__(self, views): + self._projects = [_FakeProject(view) for view in views] + self.calls = [] + + def __call__(self, **kwargs): + self.calls.append(kwargs) + return _FakeClient(self._projects.pop(0), kwargs) + + +class FeatureStoreDeltaUploaderTest(unittest.TestCase): + def test_proto_groups_required_fields_before_optional_fields(self): + required_fields = [ + "region", + "access_key_id", + "access_key_secret", + "featuredb_username", + "featuredb_password", + "project_name", + "feature_view_name", + "version", + "feature_entity_name", + ] + optional_fields = [ + "endpoint", + "security_token", + "upload_batch_size", + "max_retries", + "retry_backoff_secs", + "shard_wait_timeout_secs", + "shutdown_timeout_secs", + "max_pending_steps", + "poll_interval_secs", + "feature_view_ttl_secs", + "feature_view_shard_count", + "feature_view_replication_count", + ] + fields = list(FeatureStoreConfig.DESCRIPTOR.fields) + + self.assertEqual( + [field.name for field in fields], required_fields + optional_fields + ) + self.assertTrue( + all( + field.label == FieldDescriptor.LABEL_REQUIRED + for field in fields[: len(required_fields)] + ) + ) + self.assertTrue( + all( + field.label == FieldDescriptor.LABEL_OPTIONAL + for field in fields[len(required_fields) :] + ) + ) + self.assertEqual([field.number for field in fields], list(range(1, 22))) + for field_name in required_fields: + with self.subTest(field_name=field_name): + config = _feature_store_config() + config.ClearField(field_name) + self.assertFalse(config.IsInitialized()) + self.assertIn(field_name, config.FindInitializationErrors()) + with self.assertRaisesRegex(ValueError, field_name): + FeatureStoreUploadSettings.from_proto(config) + + def test_feature_entity_is_required_for_view_creation(self): + config = _feature_store_config() + config.ClearField("feature_entity_name") + + self.assertFalse(config.IsInitialized()) + self.assertIn("feature_entity_name", config.FindInitializationErrors()) + with self.assertRaisesRegex(ValueError, "feature_entity_name"): + FeatureStoreUploadSettings.from_proto(config) + + def test_version_is_required_and_must_be_explicit(self): + config = _feature_store_config() + config.ClearField("version") + + self.assertFalse(config.IsInitialized()) + self.assertIn("version", config.FindInitializationErrors()) + with self.assertRaisesRegex(ValueError, "required fields.*version"): + FeatureStoreUploadSettings.from_proto(config) + + with self.assertRaisesRegex(ValueError, "explicit non-default version"): + FeatureStoreUploadSettings.from_proto( + _feature_store_config(version="default") + ) + + def test_environment_fallback_and_credential_pairs(self): + config = _feature_store_config() + for name in ( + "region", + "access_key_id", + "access_key_secret", + "featuredb_username", + "featuredb_password", + ): + setattr(config, name, "") + config.ClearField("security_token") + self.assertTrue(config.IsInitialized()) + environment = { + "ALIBABA_CLOUD_REGION": "cn-env", + "ALIBABA_CLOUD_ACCESS_KEY_ID": "env-ak", + "ALIBABA_CLOUD_ACCESS_KEY_SECRET": "env-sk", + "ALIBABA_CLOUD_SECURITY_TOKEN": "env-sts", + "FEATUREDB_USERNAME": "env-user", + "FEATUREDB_PASSWORD": "env-password", + } + with mock.patch.dict(os.environ, environment, clear=False): + settings = FeatureStoreUploadSettings.from_proto(config) + self.assertEqual(settings.region, "cn-env") + self.assertEqual(settings.access_key_id, "env-ak") + + config = _feature_store_config() + config.access_key_secret = "" + with mock.patch.dict( + os.environ, {"ALIBABA_CLOUD_ACCESS_KEY_SECRET": ""}, clear=False + ): + with self.assertRaisesRegex(ValueError, "configured together"): + FeatureStoreUploadSettings.from_proto(config) + + config = _feature_store_config() + config.featuredb_username = "" + config.featuredb_password = "" + with mock.patch.dict( + os.environ, + {"FEATUREDB_USERNAME": "", "FEATUREDB_PASSWORD": ""}, + clear=False, + ): + with self.assertRaisesRegex(ValueError, "are required"): + FeatureStoreUploadSettings.from_proto(config) + + config.featuredb_username = "only-one-side" + with mock.patch.dict( + os.environ, + {"FEATUREDB_USERNAME": "", "FEATUREDB_PASSWORD": ""}, + clear=False, + ): + with self.assertRaisesRegex(ValueError, "configured together"): + FeatureStoreUploadSettings.from_proto(config) + + with self.assertRaisesRegex(ValueError, "URI userinfo"): + FeatureStoreUploadSettings.from_proto( + _feature_store_config(endpoint="https://user:secret@example.com") + ) + with self.assertRaisesRegex(ValueError, "must be <= 1000"): + FeatureStoreUploadSettings.from_proto( + _feature_store_config(upload_batch_size=1001) + ) + with self.assertRaisesRegex(ValueError, "shard_count must be in"): + FeatureStoreUploadSettings.from_proto( + _feature_store_config(feature_view_shard_count=21) + ) + with self.assertRaisesRegex(ValueError, "replication_count must be in"): + FeatureStoreUploadSettings.from_proto( + _feature_store_config(feature_view_replication_count=4) + ) + + def test_start_reuses_existing_dynamic_embedding_feature_view(self): + with tempfile.TemporaryDirectory() as output_dir: + view = _FakeView() + factory = _FakeClientFactory(view) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + + uploader.start() + uploader.close() + + self.assertEqual(factory.project.dynamic_get_calls, ["shared_embeddings"]) + self.assertEqual(factory.project.generic_get_calls, ["shared_embeddings"]) + self.assertEqual(factory.project.create_calls, []) + self.assertEqual(view.closed, [True]) + + def test_start_creates_missing_dynamic_embedding_feature_view(self): + with tempfile.TemporaryDirectory() as output_dir: + created_view = _FakeView() + factory = _FakeClientFactory(None, created_view=created_view) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + + uploader.start() + uploader.close() + + self.assertEqual( + factory.project.generic_get_calls, + ["shared_embeddings", "shared_embeddings"], + ) + self.assertEqual( + factory.project.create_calls, + [ + { + "name": "shared_embeddings", + "entity": "embedding_entity", + "pk_field_name": "embedding_name", + "sk_field_name": "key_id", + "embedding_field_name": "embedding", + "pk_field_type": "STRING", + "sk_field_type": "INT64", + "ttl": 1296000, + "shard_count": 20, + "replication_count": 1, + } + ], + ) + self.assertNotIn("embedding_dim", factory.project.create_calls[0]) + self.assertEqual(created_view.closed, [True]) + + def test_start_rejects_same_name_non_dynamic_feature_view(self): + with tempfile.TemporaryDirectory() as output_dir: + generic_view = mock.Mock(type="Batch") + factory = _FakeClientFactory(None, generic_view=generic_view) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + + with self.assertRaisesRegex(RuntimeError, "incompatible type"): + uploader.start() + + self.assertEqual(factory.project.create_calls, []) + + def test_start_rejects_existing_feature_view_with_wrong_entity(self): + with tempfile.TemporaryDirectory() as output_dir: + view = _FakeView() + generic_view = _FakeGenericFeatureView(entity="another_entity") + factory = _FakeClientFactory(view, generic_view=generic_view) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + + with self.assertRaisesRegex(RuntimeError, "entity mismatch"): + uploader.start() + + self.assertEqual(factory.project.create_calls, []) + self.assertEqual(view.closed, [True]) + + def test_start_rejects_existing_feature_view_with_wrong_field_contract(self): + with tempfile.TemporaryDirectory() as output_dir: + view = _FakeView() + generic_view = _FakeGenericFeatureView() + generic_view.fields_dict["embedding"]["Type"] = "ARRAY" + factory = _FakeClientFactory(view, generic_view=generic_view) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + + with self.assertRaisesRegex(RuntimeError, "field contract mismatch"): + uploader.start() + + self.assertEqual(factory.project.create_calls, []) + self.assertEqual(view.closed, [True]) + + def test_start_rejects_existing_feature_view_with_wrong_provisioning(self): + with tempfile.TemporaryDirectory() as output_dir: + view = _FakeView() + generic_view = _FakeGenericFeatureView( + provisioning={ + "ttl": 60, + "shard_count": 20, + "replication_count": 1, + } + ) + factory = _FakeClientFactory(view, generic_view=generic_view) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + + with self.assertRaisesRegex(RuntimeError, "provisioning mismatch"): + uploader.start() + + self.assertEqual(factory.project.create_calls, []) + self.assertEqual(view.closed, [True]) + + def test_start_recovers_from_concurrent_feature_view_creation(self): + with tempfile.TemporaryDirectory() as output_dir: + concurrent_view = _FakeView() + factory = _FakeClientFactory( + None, + create_error=RuntimeError("already exists"), + view_after_create_error=concurrent_view, + ) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + + uploader.start() + uploader.close() + + self.assertEqual(len(factory.project.create_calls), 1) + self.assertEqual( + factory.project.dynamic_get_calls, + ["shared_embeddings", "shared_embeddings"], + ) + self.assertEqual(concurrent_view.closed, [True]) + + def test_start_closes_new_feature_view_with_incompatible_schema(self): + with tempfile.TemporaryDirectory() as output_dir: + created_view = _FakeView() + created_view.pk_field = "wrong_pk" + factory = _FakeClientFactory(None, created_view=created_view) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + + with self.assertRaisesRegex(RuntimeError, "schema mismatch"): + uploader.start() + + self.assertEqual(len(factory.project.create_calls), 1) + self.assertEqual(created_view.closed, [True]) + + def test_start_creates_missing_view_but_requires_preexisting_version(self): + with tempfile.TemporaryDirectory() as output_dir: + created_view = _FakeView(versions=[]) + factory = _FakeClientFactory(None, created_view=created_view) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + + with self.assertRaisesRegex( + FeatureStoreUploadError, "version does not exist" + ): + uploader.start() + + self.assertEqual(len(factory.project.create_calls), 1) + self.assertEqual(created_view.closed, [True]) + + def test_start_reports_missing_entity_when_view_creation_fails(self): + with tempfile.TemporaryDirectory() as output_dir: + factory = _FakeClientFactory( + None, create_error=ValueError("Entity not found") + ) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + + with self.assertRaisesRegex(RuntimeError, "feature_entity_name"): + uploader.start() + + self.assertEqual(len(factory.project.create_calls), 1) + + def test_existing_remote_target_rejects_contract_change(self): + with tempfile.TemporaryDirectory() as output_dir: + first = FeatureStoreDeltaUploader( + _feature_store_config(upload_batch_size=2), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + first.start() + first.close() + second = FeatureStoreDeltaUploader( + _feature_store_config(upload_batch_size=1), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with self.assertRaisesRegex(ValueError, "upload contract changed"): + second.start() + second.close() + + def test_submit_requires_started_uploader(self): + with tempfile.TemporaryDirectory() as output_dir: + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with self.assertRaisesRegex(RuntimeError, "start.*before submit"): + uploader.submit(10) + uploader.close() + + def test_complete_step_uploads_merge_with_stable_version_and_ts(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard( + output_dir, + 10, + [ + _row(10, 0, 1, [1.0, 2.0]), + _row(10, 0, 2, [3.0, 4.0]), + _row(10, 0, 3, [5.0, 6.0]), + ], + ) + view = _FakeView() + factory = _FakeClientFactory(view) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir=output_dir, + file_prefix="delta", + world_size=1, + embedding_dimensions={"user_emb": 2}, + client_factory=factory, + clock_ms=lambda: 123456, + ) + uploader.start() + uploader.submit(10) + uploader.close() + + self.assertEqual(len(view.calls), 2) + self.assertEqual([len(call["data"]) for call in view.calls], [2, 1]) + self.assertEqual( + {call["version"] for call in view.calls}, {"model_a@export_1"} + ) + self.assertEqual({call["write_mode"] for call in view.calls}, {"MERGE"}) + self.assertEqual([call["ts"] for call in view.calls], [123456, 123457]) + self.assertEqual(view.closed, [True]) + + success_path = os.path.join(uploader.state_dir, "step_10._FS_SUCCESS.json") + self.assertTrue(os.path.isfile(success_path)) + with open(os.path.join(uploader.state_dir, "committed.json")) as source: + committed = json.load(source) + self.assertEqual(committed["committed_global_step"], 10) + self.assertEqual(committed["publish_ts"], 123457) + + generated_text = "" + for root, _, names in os.walk(uploader.state_dir): + for name in names: + if name.endswith(".json"): + with open(os.path.join(root, name)) as source: + generated_text += source.read() + for secret in ( + "ak-id-secret", + "ak-value-secret", + "sts-secret", + "fdb-user-secret", + "fdb-password-secret", + ): + self.assertNotIn(secret, generated_text) + + def test_first_positive_dump_step_is_not_filtered_by_an_implicit_base(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard(output_dir, 1, [_row(1, 0, 1, [1.0, 2.0])]) + view = _FakeView() + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(view), + clock_ms=lambda: 100, + ) + + uploader.start() + uploader.close() + + self.assertEqual(len(view.calls), 1) + self.assertEqual(view.calls[0]["ts"], 100) + self.assertEqual(view.calls[0]["version"], "model_a@export_1") + self.assertTrue( + os.path.isfile( + os.path.join(uploader.state_dir, "step_1._FS_SUCCESS.json") + ) + ) + + def test_step_zero_outbox_fails_loud_without_upload(self): + with tempfile.TemporaryDirectory() as output_dir: + step_zero_path = _write_single_shard( + output_dir, 0, [_row(0, 0, 1, [1.0, 2.0])] + ) + view = _FakeView() + factory = _FakeClientFactory(view) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + + with self.assertRaisesRegex(ValueError, "global_step must be > 0"): + uploader.start() + + self.assertEqual(factory.calls, []) + self.assertEqual(view.calls, []) + + # Retrying the same uploader after quarantining the invalid file + # reacquires and reconciles the journal instead of using stale state. + os.remove(step_zero_path) + uploader.start() + uploader.close() + + def test_step_zero_success_marker_fails_before_reconcile(self): + with tempfile.TemporaryDirectory() as output_dir: + view = _FakeView() + factory = _FakeClientFactory(view) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + with open( + os.path.join(uploader.state_dir, "step_0._FS_SUCCESS.json"), "w" + ) as output: + json.dump({}, output) + + with self.assertRaisesRegex(ValueError, "global_step must be > 0"): + uploader.start() + + self.assertEqual(factory.calls, []) + self.assertEqual(view.calls, []) + + def test_step_zero_committed_watermark_is_rejected(self): + with tempfile.TemporaryDirectory() as output_dir: + view = _FakeView() + factory = _FakeClientFactory(view) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + with open(os.path.join(uploader.state_dir, "committed.json"), "w") as out: + json.dump({"committed_global_step": 0}, out) + + with self.assertRaisesRegex(ValueError, "must be -1 or greater than 0"): + uploader.start() + + self.assertEqual(factory.calls, []) + self.assertEqual(view.calls, []) + + def test_submit_rejects_step_zero(self): + with tempfile.TemporaryDirectory() as output_dir: + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + + uploader.start() + try: + with self.assertRaisesRegex(ValueError, "global_step must be > 0"): + uploader.submit(0) + finally: + uploader.close() + + def test_target_scoped_prefix_prevents_cross_version_parquet_replay(self): + with tempfile.TemporaryDirectory() as output_dir: + version_a = _feature_store_config(version="model_a@run_1") + version_b = _feature_store_config(version="model_a@run_2") + prefix_a = feature_store_delta_file_prefix(version_a, "delta") + prefix_b = feature_store_delta_file_prefix(version_b, "delta") + self.assertNotEqual(prefix_a, prefix_b) + self.assertEqual( + prefix_a, feature_store_delta_file_prefix(version_a, "delta") + ) + + _write_single_shard( + output_dir, + 10, + [_row(10, 0, 1, [1.0, 2.0])], + file_prefix=prefix_a, + ) + view = _FakeView(versions=["model_a@run_2"]) + factory = _FakeClientFactory(view) + uploader = FeatureStoreDeltaUploader( + version_b, + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + self.assertEqual(uploader._file_prefix, prefix_b) + + uploader.start() + uploader.close() + + self.assertEqual(len(factory.calls), 1) + self.assertEqual(view.calls, []) + + def test_preconstructed_uploader_refreshes_timestamp_fence_after_lock(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + first_view = _FakeView() + first = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(first_view), + clock_ms=lambda: 777, + ) + stale_view = _FakeView() + stale = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(stale_view), + clock_ms=lambda: 777, + ) + + first.start() + first.close() + _write_single_shard(output_dir, 11, [_row(11, 0, 2, [3.0, 4.0])]) + stale.start() + stale.close() + + self.assertEqual([call["ts"] for call in first_view.calls], [777]) + self.assertEqual([call["ts"] for call in stale_view.calls], [778]) + + def test_exact_duplicate_key_is_deduplicated_before_async_sdk_batches(self): + with tempfile.TemporaryDirectory() as output_dir: + shard = _write_single_shard( + output_dir, + 10, + [ + _row(10, 0, 1, [1.0, 2.0]), + _row(10, 0, 1, [1.0, 2.0]), + ], + ) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + records = uploader._load_records(10, [shard]) + self.assertEqual(len(records), 1) + self.assertEqual(records[0][2].tolist(), [1.0, 2.0]) + + def test_conflicting_duplicate_key_is_rejected(self): + with tempfile.TemporaryDirectory() as output_dir: + shard = _write_single_shard( + output_dir, + 10, + [ + _row(10, 0, 1, [1.0, 2.0]), + _row(10, 0, 1, [7.0, 8.0]), + ], + ) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with self.assertRaisesRegex(ValueError, "conflicting duplicate"): + uploader._load_records(10, [shard]) + + def test_flush_failure_does_not_publish_success_marker(self): + failed_summary = { + "total_batches": 1, + "failed_batches": 1, + "total_records": 0, + "success_records": 0, + "failed_records": 0, + "errors": ["failed future"], + } + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + view = _FakeView([failed_summary, failed_summary]) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(max_retries=1), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(view), + ) + uploader.start() + uploader.submit(10) + with self.assertRaises(FeatureStoreUploadError): + uploader.close() + self.assertFalse( + os.path.exists( + os.path.join(uploader.state_dir, "step_10._FS_SUCCESS.json") + ) + ) + + def test_retry_uses_fresh_view_and_newer_persisted_timestamp_range(self): + failed_summary = { + "total_batches": 1, + "failed_batches": 1, + "total_records": 1, + "success_records": 0, + "failed_records": 1, + "errors": ["failed future"], + } + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + first_view = _FakeView([failed_summary]) + second_view = _FakeView() + factory = _SequencedClientFactory([first_view, second_view]) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(max_retries=2), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + clock_ms=lambda: 777, + ) + uploader.start() + uploader.submit(10) + uploader.close() + + self.assertEqual(len(factory.calls), 2) + self.assertEqual(first_view.closed, [True]) + self.assertEqual(second_view.closed, [True]) + all_calls = first_view.calls + second_view.calls + self.assertEqual( + {call["version"] for call in all_calls}, {"model_a@export_1"} + ) + self.assertEqual([call["ts"] for call in all_calls], [777, 778]) + + def test_restart_skips_abandoned_attempt_timestamp_range(self): + with tempfile.TemporaryDirectory() as output_dir: + shard = _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + interrupted = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + clock_ms=lambda: 777, + ) + with _initialized_journal(interrupted): + snapshot_paths = interrupted._snapshot_canonical_shards(10, [shard]) + self.assertIsNotNone(snapshot_paths) + manifest = interrupted._load_or_create_manifest(10, snapshot_paths, 1) + interrupted._start_attempt(10, 1, manifest) + os.remove(shard) + + view = _FakeView() + restarted = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(view), + clock_ms=lambda: 777, + ) + restarted.start() + restarted.close() + + self.assertEqual([call["ts"] for call in view.calls], [778]) + + def test_uncommitted_manifest_without_snapshot_fails_recovery(self): + with tempfile.TemporaryDirectory() as output_dir: + shard = _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + interrupted = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with _initialized_journal(interrupted): + snapshot_paths = interrupted._snapshot_canonical_shards(10, [shard]) + interrupted._load_or_create_manifest(10, snapshot_paths, 1) + os.chmod(snapshot_paths[0], 0o600) + os.remove(snapshot_paths[0]) + os.rmdir(interrupted._snapshot_dir(10)) + os.remove(shard) + + factory = _FakeClientFactory(_FakeView()) + restarted = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + restarted.start() + with self.assertRaisesRegex( + FeatureStoreUploadError, "missing its durable shard snapshot" + ): + restarted.close() + self.assertEqual(len(factory.calls), 1) + + def test_different_multi_rank_dump_generations_are_not_snapshotted(self): + with tempfile.TemporaryDirectory() as output_dir: + shard_0 = _write_rank_shard( + output_dir, + 10, + 0, + 2, + [_row(10, 0, 1, [1.0, 2.0], world_size=2)], + generation="generation-a", + ) + shard_1 = _write_rank_shard( + output_dir, + 10, + 1, + 2, + [_row(10, 1, 2, [3.0, 4.0], world_size=2)], + generation="generation-b", + ) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 2, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with self.assertRaisesRegex(RuntimeError, "different dump generations"): + uploader._snapshot_canonical_shards(10, [shard_0, shard_1]) + self.assertFalse(os.path.exists(uploader._snapshot_dir(10))) + + def test_same_output_target_rejects_a_second_active_writer(self): + with tempfile.TemporaryDirectory() as output_dir: + first = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + second = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + first.start() + with self.assertRaisesRegex(FeatureStoreUploadError, "another process"): + second.start() + first.close() + second.close() + + def test_start_clears_a_stale_failure_marker_for_restart(self): + with tempfile.TemporaryDirectory() as output_dir: + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with open(uploader._error_marker_path, "w") as output: + output.write("{}") + uploader.start() + self.assertFalse(os.path.exists(uploader._error_marker_path)) + uploader.close() + + def test_missing_preprovisioned_version_blocks_merge(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + view = _FakeView(versions=[]) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(max_retries=1), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(view), + ) + with self.assertRaisesRegex( + FeatureStoreUploadError, "version does not exist" + ): + uploader.start() + uploader.close() + self.assertEqual(view.calls, []) + self.assertEqual(view.closed, [True]) + self.assertFalse( + os.path.exists( + os.path.join(uploader.state_dir, "step_10._FS_SUCCESS.json") + ) + ) + + def test_close_failure_aborts_retry(self): + failed_summary = { + "total_batches": 1, + "failed_batches": 1, + "total_records": 1, + "success_records": 0, + "failed_records": 1, + "errors": ["failed future"], + } + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + first_view = _FakeView( + [failed_summary], close_error=RuntimeError("unsafe close") + ) + second_view = _FakeView() + factory = _SequencedClientFactory([first_view, second_view]) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(max_retries=2), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + uploader.start() + uploader.submit(10) + with self.assertRaisesRegex( + FeatureStoreUploadError, "close failed; retry was aborted" + ): + uploader.close() + self.assertEqual(len(factory.calls), 1) + self.assertEqual(second_view.calls, []) + contender = FeatureStoreDeltaUploader( + _feature_store_config(max_retries=2), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with self.assertRaisesRegex(FeatureStoreUploadError, "another process"): + contender.start() + contender.close() + + def test_restart_reclaims_snapshot_left_after_committed_upload(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with mock.patch( + "tzrec.utils.feature_store_delta_uploader.shutil.rmtree", + side_effect=OSError("cleanup interrupted"), + ): + uploader.start() + uploader.close() + snapshot_dir = uploader._snapshot_dir(10) + self.assertTrue(os.path.isdir(snapshot_dir)) + + restarted = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + restarted.start() + restarted.close() + self.assertFalse(os.path.exists(snapshot_dir)) + + def test_non_draining_close_stops_after_inflight_batch_without_commit(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard( + output_dir, + 10, + [ + _row(10, 0, 1, [1.0, 2.0]), + _row(10, 0, 2, [3.0, 4.0]), + _row(10, 0, 3, [5.0, 6.0]), + ], + ) + view = _BlockingView() + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(view), + ) + uploader.start() + uploader.submit(10) + self.assertTrue(view.flush_started.wait(timeout=2)) + uploader.close(raise_on_error=False, drain=False) + self.assertFalse( + os.path.exists( + os.path.join(uploader.state_dir, "step_10._FS_SUCCESS.json") + ) + ) + view.release_flush.set() + self.assertTrue(view.close_finished.wait(timeout=2)) + self.assertEqual(len(view.calls), 1) + self.assertFalse( + os.path.exists( + os.path.join(uploader.state_dir, "step_10._FS_SUCCESS.json") + ) + ) + + def test_signed_int64_key_is_preserved(self): + with tempfile.TemporaryDirectory() as output_dir: + shard = _write_single_shard(output_dir, 10, [_row(10, 0, -2, [1.0, 2.0])]) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + records = uploader._load_records(10, [shard]) + self.assertEqual(records[0][1], -2) + + def test_reserved_invalid_key_is_rejected(self): + with tempfile.TemporaryDirectory() as output_dir: + shard = _write_single_shard(output_dir, 10, [_row(10, 0, -1, [1.0, 2.0])]) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with self.assertRaisesRegex(ValueError, "invalid-key sentinel"): + uploader._load_records(10, [shard]) + + def test_empty_step_validates_remote_contract_before_commit(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard(output_dir, 10, []) + factory = _FakeClientFactory(_FakeView()) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + uploader.start() + uploader.submit(10) + uploader.close() + self.assertEqual(len(factory.calls), 1) + self.assertTrue( + os.path.isfile( + os.path.join(uploader.state_dir, "step_10._FS_SUCCESS.json") + ) + ) + + def test_multi_rank_discovery_uses_exact_contract_shards(self): + with tempfile.TemporaryDirectory() as output_dir: + os.makedirs(os.path.join(output_dir, "step_9")) + wrong_path = os.path.join( + output_dir, + "step_9", + "other_step_9_rank_0_of_4.parquet", + ) + with open(wrong_path, "w"): + pass + view = _FakeView() + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 2, + {"user_emb": 2}, + client_factory=_FakeClientFactory(view), + ) + self.assertEqual(set(uploader._discover_steps()), set()) + + _write_rank_shard( + output_dir, + 10, + 0, + 2, + [_row(10, 0, 1, [1.0, 2.0], world_size=2)], + ) + self.assertEqual(set(uploader._discover_steps()), {10}) + _write_rank_shard( + output_dir, + 10, + 1, + 2, + [_row(10, 1, 2, [3.0, 4.0], world_size=2)], + ) + uploader.start() + uploader.close() + + self.assertEqual(len(view.calls), 1) + self.assertEqual([item["key_id"] for item in view.calls[0]["data"]], [1, 2]) + + def test_dimension_and_finite_value_validation(self): + with tempfile.TemporaryDirectory() as output_dir: + bad_dimension = _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0])]) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with self.assertRaisesRegex(ValueError, "dimension mismatch"): + uploader._load_records(10, [bad_dimension]) + + bad_finite = _write_single_shard( + output_dir, 11, [_row(11, 0, 1, [1.0, float("nan")])] + ) + with self.assertRaisesRegex(ValueError, "NaN or Inf"): + uploader._load_records(11, [bad_finite]) + + def test_persisted_manifest_rejects_changed_shard_content(self): + with tempfile.TemporaryDirectory() as output_dir: + shard = _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with _initialized_journal(uploader): + uploader._load_or_create_manifest(10, [shard], 1) + _write_single_shard(output_dir, 10, [_row(10, 0, 1, [7.0, 8.0])]) + with self.assertRaisesRegex(ValueError, "manifest mismatch for shards"): + uploader._load_or_create_manifest(10, [shard], 1) + + def test_opened_shard_snapshot_rejects_atomic_path_replacement(self): + with tempfile.TemporaryDirectory() as output_dir: + shard = _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with open(shard, "rb") as source: + descriptions = uploader._describe_shards([shard], [source]) + replacement = f"{shard}.replacement" + pq.write_table( + pa.Table.from_pylist( + [_row(10, 0, 1, [7.0, 8.0])], + schema=_schema_with_generation(), + ), + replacement, + ) + os.replace(replacement, shard) + + records = uploader._load_records(10, [shard], [source]) + self.assertEqual(records[0][2].tolist(), [1.0, 2.0]) + with self.assertRaisesRegex( + RuntimeError, "changed after upload snapshot" + ): + uploader._validate_shard_identities([shard], descriptions, [source]) + + def test_opened_shard_snapshot_rejects_same_inode_content_rewrite(self): + with tempfile.TemporaryDirectory() as output_dir: + shard = _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + original_stat = os.stat(shard) + with open(shard, "rb") as source: + descriptions = uploader._describe_shards([shard], [source]) + pq.write_table( + pa.Table.from_pylist( + [_row(10, 0, 1, [7.0, 8.0])], + schema=_schema_with_generation(), + ), + shard, + ) + os.utime( + shard, + ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns), + ) + with self.assertRaisesRegex( + RuntimeError, "changed after upload snapshot" + ): + uploader._validate_shard_identities([shard], descriptions, [source]) + + def test_restart_skips_successfully_committed_step(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard(output_dir, 10, []) + first = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + first.start() + first.submit(10) + first.close() + + factory = _FakeClientFactory(_FakeView()) + restarted = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=factory, + ) + restarted.start() + restarted.close() + self.assertEqual(len(factory.calls), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/utils/sparse_embedding_contract.py b/tzrec/utils/sparse_embedding_contract.py new file mode 100644 index 000000000..f64f93733 --- /dev/null +++ b/tzrec/utils/sparse_embedding_contract.py @@ -0,0 +1,105 @@ +# 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. + +"""Shared sparse-embedding identity used by export, delta dump and serving.""" + +from collections import defaultdict +from dataclasses import dataclass +from typing import Dict, Iterable, Optional, Tuple + +SPARSE_EC_ROLE = "ec" +SPARSE_EBC_ROLE = "ebc" +SPARSE_EMBEDDING_ROLES = frozenset((SPARSE_EC_ROLE, SPARSE_EBC_ROLE)) +# NvEmbeddings and the future Processor consumer reserve this key as invalid. +# Other negative int64 values remain valid bit patterns for dynamic uint64 IDs. +SPARSE_EMBEDDING_INVALID_KEY = -1 + + +@dataclass(frozen=True) +class SparseEmbeddingIdentity: + """Physical sparse table identity and its cross-system canonical name.""" + + role: str + table_name: str + embedding_name: str + dimension: int + feature_names: Tuple[str, ...] = () + + +def build_sparse_embedding_name_map( + table_identities: Iterable[Tuple[str, str]], +) -> Dict[Tuple[str, str], str]: + """Allocate canonical names for ``(collection role, table name)`` pairs. + + EC and EBC are separate physical collections. A table name used in only one + collection keeps its historical name. If it appears in both collections, + each physical table receives a role suffix. Candidate collisions with other + raw table names are resolved deterministically with a numeric suffix. + """ + roles_by_name = defaultdict(set) + for role, table_name in table_identities: + if role not in SPARSE_EMBEDDING_ROLES: + raise ValueError(f"unsupported sparse embedding role: {role!r}") + if not table_name: + raise ValueError("sparse embedding table_name must not be empty") + roles_by_name[table_name].add(role) + + used_names = set(roles_by_name.keys()) + name_by_identity: Dict[Tuple[str, str], str] = {} + for table_name, roles in roles_by_name.items(): + if len(roles) == 1: + role = next(iter(roles)) + name_by_identity[(role, table_name)] = table_name + continue + + for role in sorted(roles): + base_candidate = f"{table_name}__{role}" + candidate = base_candidate + suffix = 1 + while candidate in used_names: + candidate = f"{base_candidate}_{suffix}" + suffix += 1 + used_names.add(candidate) + name_by_identity[(role, table_name)] = candidate + return name_by_identity + + +def resolve_sparse_embedding_name( + name_by_identity: Dict[Tuple[str, str], str], + table_name: str, + role: Optional[str], +) -> str: + """Resolve a table to its canonical name, rejecting ambiguous identities.""" + if role is not None and (role, table_name) in name_by_identity: + return name_by_identity[(role, table_name)] + + candidates = [ + embedding_name + for (_role, name), embedding_name in name_by_identity.items() + if name == table_name + ] + if len(candidates) == 1: + return candidates[0] + if not candidates: + raise KeyError(f"sparse embedding {table_name!r} is not in model metadata") + raise ValueError( + f"sparse embedding {table_name!r} appears in multiple collection kinds; " + f"cannot resolve canonical name without role, got role={role!r}" + ) + + +def sparse_embedding_role_from_state_key(state_key: str) -> Optional[str]: + """Infer the sparse collection role from a TorchRec state/table FQN.""" + if ".embedding_bags." in state_key: + return SPARSE_EBC_ROLE + if ".embeddings." in state_key: + return SPARSE_EC_ROLE + return None diff --git a/tzrec/utils/sparse_embedding_contract_test.py b/tzrec/utils/sparse_embedding_contract_test.py new file mode 100644 index 000000000..2a9f54648 --- /dev/null +++ b/tzrec/utils/sparse_embedding_contract_test.py @@ -0,0 +1,66 @@ +# Copyright (c) 2025, 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. + +import unittest + +from tzrec.utils.sparse_embedding_contract import ( + build_sparse_embedding_name_map, + resolve_sparse_embedding_name, + sparse_embedding_role_from_state_key, +) + + +class SparseEmbeddingContractTest(unittest.TestCase): + def test_single_collection_keeps_table_name(self): + names = build_sparse_embedding_name_map( + [("ec", "sequence_emb"), ("ebc", "user_emb")] + ) + self.assertEqual(names[("ec", "sequence_emb")], "sequence_emb") + self.assertEqual(names[("ebc", "user_emb")], "user_emb") + + def test_cross_collection_collision_uses_role_and_numeric_suffix(self): + names = build_sparse_embedding_name_map( + [ + ("ec", "shared"), + ("ebc", "shared"), + ("ec", "shared__ec"), + ] + ) + self.assertEqual(names[("ec", "shared")], "shared__ec_1") + self.assertEqual(names[("ebc", "shared")], "shared__ebc") + self.assertEqual(names[("ec", "shared__ec")], "shared__ec") + + def test_resolver_requires_role_for_ambiguous_table(self): + names = build_sparse_embedding_name_map([("ec", "shared"), ("ebc", "shared")]) + with self.assertRaisesRegex(ValueError, "multiple collection"): + resolve_sparse_embedding_name(names, "shared", None) + self.assertEqual( + resolve_sparse_embedding_name(names, "shared", "ec"), "shared__ec" + ) + + def test_state_key_role(self): + self.assertEqual( + sparse_embedding_role_from_state_key( + "model.ebc.embedding_bags.user_emb.weight" + ), + "ebc", + ) + self.assertEqual( + sparse_embedding_role_from_state_key( + "model.ec.embeddings.sequence_emb.weight" + ), + "ec", + ) + self.assertIsNone(sparse_embedding_role_from_state_key("model.unknown.weight")) + + +if __name__ == "__main__": + unittest.main() From 01481c1504109ce9e8bfcf69e8cbf7bd9e22ff67 Mon Sep 17 00:00:00 2001 From: gecheng Date: Thu, 16 Jul 2026 17:55:09 +0800 Subject: [PATCH 02/13] parallelize FeatureStore delta uploads --- tzrec/utils/feature_store_delta_uploader.py | 141 ++++++++++++------ .../feature_store_delta_uploader_test.py | 140 ++++++++++++----- 2 files changed, 203 insertions(+), 78 deletions(-) diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index 69d5fddce..2a9b559c2 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -55,6 +55,10 @@ DELTA_DUMP_SCHEMA_VERSION = "2" DELTA_DUMP_GENERATION_METADATA_KEY = b"tzrec.delta_embedding.dump_generation" +_FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES = 100 +_VERSION_INITIALIZATION = "AUTO_CREATE_ON_FIRST_DELTA_MERGE" +_LEGACY_VERSION_INITIALIZATION = "PREPROVISIONED_FOR_DELTA_MERGE" + _POISONED_WRITER_LOCKS: List[BinaryIO] = [] _SCHEMA_VERSION_METADATA_KEY = b"tzrec.delta_embedding.schema_version" @@ -453,7 +457,7 @@ def __init__( self._settings.feature_view_replication_count ), "version": self._settings.version, - "version_initialization": "PREPROVISIONED_FOR_DELTA_MERGE", + "version_initialization": _VERSION_INITIALIZATION, "feature_view_provisioning": "CHECK_OR_CREATE_DYNAMIC_EMBEDDING", "writer_ownership": "SINGLE_WRITER_PER_VERSION_REQUIRED", "publish_semantics": "MONOTONIC_TS_RANGE_PER_ATTEMPT_FULL_REPLAY", @@ -514,8 +518,7 @@ def start(self) -> None: self._cleanup_committed_snapshots() self._add_discovered_steps_locked() # Validate or provision the remote DynamicEmbedding FeatureView - # synchronously. Training must not start while the target is - # missing, has the wrong schema, or lacks the configured version. + # synchronously. Training must not start with an incompatible target. self._get_view() self._clear_error_marker() self._started = True @@ -775,11 +778,22 @@ def _initialize_journal(self) -> None: "FeatureStore upload journal must be initialized under writer lock" ) if os.path.isfile(self._contract_path): - if _read_json(self._contract_path) != self._contract: - raise ValueError( - "FeatureStore upload contract changed for an existing remote " - "target; provision and configure a new immutable version" + stored_contract = _read_json(self._contract_path) + if stored_contract != self._contract: + legacy_contract = dict(self._contract) + legacy_contract["version_initialization"] = ( + _LEGACY_VERSION_INITIALIZATION ) + if stored_contract == legacy_contract: + # Preserve the old hash so existing manifests and success + # markers remain valid after removing the version precheck. + self._contract = stored_contract + self._contract_hash = _json_digest(stored_contract) + else: + raise ValueError( + "FeatureStore upload contract changed for an existing remote " + "target; provision and configure a new immutable version" + ) else: _atomic_write_json(self._contract_path, self._contract) @@ -1303,10 +1317,10 @@ def _upload_with_retries( attempt = self._start_attempt(global_step, len(records), manifest) try: if records: - summary = self._upload_records(records, attempt) + summary = self._upload_records(global_step, records, attempt) else: - # An empty step still validates the FeatureView schema and - # pre-provisioned target version before it can be committed. + # An empty step still validates the FeatureView schema before + # it can be committed. self._get_view() summary = { "total_batches": 0, @@ -1563,6 +1577,7 @@ def _get_view(self) -> Any: "security_token": self._settings.security_token or None, "featuredb_username": self._settings.featuredb_username or None, "featuredb_password": self._settings.featuredb_password or None, + "test_mode": True, } client = client_factory(**kwargs) project = client.get_project(self._settings.project_name) @@ -1592,16 +1607,9 @@ def _get_view(self) -> Any: "FeatureStore SDK batch_size is smaller than the configured outer " "batch; one publish timestamp could span multiple HTTP requests" ) - version_info = view.list_versions() - versions = version_info.get("versions", []) if version_info else [] - if not any( - isinstance(item, dict) and item.get("version") == self._settings.version - for item in versions - ): - raise FeatureStoreUploadError( - "configured FeatureStore version does not exist; provision it " - "before enabling delta MERGE upload" - ) + sdk_max_workers = getattr(view, "_max_workers", 1) + if type(sdk_max_workers) is not int or sdk_max_workers <= 0: + raise RuntimeError("FeatureStore SDK max_workers must be a positive int") return view def _reset_view(self, suppress_errors: bool = False) -> None: @@ -1651,10 +1659,15 @@ def _reset_view(self, suppress_errors: bool = False) -> None: def _upload_records( self, + global_step: int, records: List[Tuple[str, int, np.ndarray]], attempt: Mapping[str, int], ) -> Dict[str, int]: view = self._get_view() + max_in_flight = int(getattr(view, "_max_workers", 1)) + batch_count = ( + len(records) + self._settings.upload_batch_size - 1 + ) // self._settings.upload_batch_size aggregate = { "total_batches": 0, "failed_batches": 0, @@ -1662,30 +1675,72 @@ def _upload_records( "success_records": 0, "failed_records": 0, } + started_at = time.monotonic() + next_progress_batch = _FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES + logger.info( + "FeatureStore delta upload started: step=%s attempt=%s version=%s " + "records=%s batches=%s max_in_flight=%s ts_range=%s-%s", + global_step, + attempt["attempt_id"], + self._settings.version, + len(records), + batch_count, + max_in_flight, + attempt["range_start"], + attempt["range_end"], + ) try: - for batch_index, offset in enumerate( - range(0, len(records), self._settings.upload_batch_size) - ): - self._raise_if_aborting() - chunk = records[offset : offset + self._settings.upload_batch_size] - payload = [ - { - FEATURE_STORE_PK_FIELD: embedding_name, - FEATURE_STORE_SK_FIELD: key_id, - FEATURE_STORE_VALUE_FIELD: embedding, - } - for embedding_name, key_id, embedding in chunk - ] - view.write_features( - data=payload, - version=self._settings.version, - write_mode=FEATURE_STORE_WRITE_MODE, - ts=int(attempt["range_start"]) + batch_index, - ) + for window_start in range(0, batch_count, max_in_flight): + window_end = min(window_start + max_in_flight, batch_count) + window_records = 0 + for batch_index in range(window_start, window_end): + self._raise_if_aborting() + offset = batch_index * self._settings.upload_batch_size + chunk = records[offset : offset + self._settings.upload_batch_size] + payload = [ + { + FEATURE_STORE_PK_FIELD: embedding_name, + FEATURE_STORE_SK_FIELD: key_id, + FEATURE_STORE_VALUE_FIELD: embedding, + } + for embedding_name, key_id, embedding in chunk + ] + view.write_features( + data=payload, + version=self._settings.version, + write_mode=FEATURE_STORE_WRITE_MODE, + ts=int(attempt["range_start"]) + batch_index, + ) + window_records += len(chunk) summary = view.write_flush() - self._validate_flush_summary(summary, len(chunk)) + self._validate_flush_summary( + summary, + expected_records=window_records, + expected_batches=window_end - window_start, + ) for name in aggregate: aggregate[name] += int(summary[name]) + completed_batches = window_end + if ( + window_start == 0 + or completed_batches >= next_progress_batch + or completed_batches == batch_count + ): + logger.info( + "FeatureStore delta upload progress: step=%s attempt=%s " + "batches=%s/%s records=%s/%s elapsed_secs=%.1f", + global_step, + attempt["attempt_id"], + completed_batches, + batch_count, + aggregate["success_records"], + len(records), + time.monotonic() - started_at, + ) + while next_progress_batch <= completed_batches: + next_progress_batch += ( + _FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES + ) self._raise_if_aborting() except BaseException: # A write_features() call can enqueue part of its work before raising. @@ -1698,7 +1753,9 @@ def _upload_records( return aggregate @staticmethod - def _validate_flush_summary(summary: Any, expected_records: int) -> None: + def _validate_flush_summary( + summary: Any, expected_records: int, expected_batches: int + ) -> None: required = { "total_batches", "failed_batches", @@ -1709,7 +1766,7 @@ def _validate_flush_summary(summary: Any, expected_records: int) -> None: if not isinstance(summary, dict) or not required.issubset(summary): raise RuntimeError("FeatureStore write_flush returned an invalid summary") if ( - int(summary["total_batches"]) != 1 + int(summary["total_batches"]) != expected_batches or int(summary["failed_batches"]) != 0 or int(summary["failed_records"]) != 0 or int(summary["success_records"]) != int(summary["total_records"]) diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index 1c03e9dac..c239bb01b 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -28,6 +28,7 @@ FeatureStoreDeltaUploader, FeatureStoreUploadError, FeatureStoreUploadSettings, + _json_digest, feature_store_delta_file_prefix, ) @@ -141,37 +142,36 @@ class _FakeView: sk_field = "key_id" embedding_field = "embedding" - def __init__(self, summaries=None, versions=None, close_error=None): + def __init__(self, summaries=None, close_error=None, max_workers=4): self.calls = [] self.closed = [] + self.flush_calls = [] self._summaries = list(summaries or []) - self._versions = ["model_a@export_1"] if versions is None else list(versions) self._close_error = close_error self._batch_size = 1000 - self._last_size = 0 + self._max_workers = max_workers + self._pending_sizes = [] def write_features(self, **kwargs): self.calls.append(kwargs) - self._last_size = len(kwargs["data"]) + self._pending_sizes.append(len(kwargs["data"])) def write_flush(self): + pending_sizes = self._pending_sizes + self._pending_sizes = [] + self.flush_calls.append(pending_sizes) if self._summaries: return self._summaries.pop(0) + total_records = sum(pending_sizes) return { - "total_batches": 1, + "total_batches": len(pending_sizes), "failed_batches": 0, - "total_records": self._last_size, - "success_records": self._last_size, + "total_records": total_records, + "success_records": total_records, "failed_records": 0, "errors": [], } - def list_versions(self): - return { - "default_version": "", - "versions": [{"version": version} for version in self._versions], - } - def close(self, wait=True): self.closed.append(wait) if self._close_error is not None: @@ -645,9 +645,9 @@ def test_start_closes_new_feature_view_with_incompatible_schema(self): self.assertEqual(len(factory.project.create_calls), 1) self.assertEqual(created_view.closed, [True]) - def test_start_creates_missing_view_but_requires_preexisting_version(self): + def test_start_creates_missing_view_without_version_precheck(self): with tempfile.TemporaryDirectory() as output_dir: - created_view = _FakeView(versions=[]) + created_view = _FakeView() factory = _FakeClientFactory(None, created_view=created_view) uploader = FeatureStoreDeltaUploader( _feature_store_config(), @@ -658,10 +658,8 @@ def test_start_creates_missing_view_but_requires_preexisting_version(self): client_factory=factory, ) - with self.assertRaisesRegex( - FeatureStoreUploadError, "version does not exist" - ): - uploader.start() + uploader.start() + uploader.close() self.assertEqual(len(factory.project.create_calls), 1) self.assertEqual(created_view.closed, [True]) @@ -709,6 +707,40 @@ def test_existing_remote_target_rejects_contract_change(self): second.start() second.close() + def test_existing_legacy_version_initialization_contract_is_reused(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + legacy = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + legacy._contract["version_initialization"] = ( + "PREPROVISIONED_FOR_DELTA_MERGE" + ) + legacy._contract_hash = _json_digest(legacy._contract) + legacy.start() + legacy.close() + + restarted = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + self.assertNotEqual(restarted._contract_hash, legacy._contract_hash) + + restarted.start() + restarted.close() + + self.assertEqual(restarted._contract_hash, legacy._contract_hash) + self.assertEqual(restarted._committed_global_step, 10) + def test_submit_requires_started_uploader(self): with tempfile.TemporaryDirectory() as output_dir: uploader = FeatureStoreDeltaUploader( @@ -751,6 +783,7 @@ def test_complete_step_uploads_merge_with_stable_version_and_ts(self): self.assertEqual(len(view.calls), 2) self.assertEqual([len(call["data"]) for call in view.calls], [2, 1]) + self.assertEqual(view.flush_calls, [[2, 1]]) self.assertEqual( {call["version"] for call in view.calls}, {"model_a@export_1"} ) @@ -780,6 +813,32 @@ def test_complete_step_uploads_merge_with_stable_version_and_ts(self): ): self.assertNotIn(secret, generated_text) + def test_upload_uses_bounded_sdk_worker_windows(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard( + output_dir, + 10, + [_row(10, 0, key, [1.0, 2.0]) for key in range(1, 6)], + ) + view = _FakeView(max_workers=2) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(upload_batch_size=1), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(view), + clock_ms=lambda: 100, + ) + + uploader.start() + uploader.close() + + self.assertEqual( + [call["ts"] for call in view.calls], [100, 101, 102, 103, 104] + ) + self.assertEqual(view.flush_calls, [[1, 1], [1, 1], [1]]) + def test_first_positive_dump_step_is_not_filtered_by_an_implicit_base(self): with tempfile.TemporaryDirectory() as output_dir: _write_single_shard(output_dir, 1, [_row(1, 0, 1, [1.0, 2.0])]) @@ -913,7 +972,7 @@ def test_target_scoped_prefix_prevents_cross_version_parquet_replay(self): [_row(10, 0, 1, [1.0, 2.0])], file_prefix=prefix_a, ) - view = _FakeView(versions=["model_a@run_2"]) + view = _FakeView() factory = _FakeClientFactory(view) uploader = FeatureStoreDeltaUploader( version_b, @@ -1009,16 +1068,24 @@ def test_conflicting_duplicate_key_is_rejected(self): def test_flush_failure_does_not_publish_success_marker(self): failed_summary = { - "total_batches": 1, + "total_batches": 2, "failed_batches": 1, - "total_records": 0, - "success_records": 0, - "failed_records": 0, + "total_records": 3, + "success_records": 2, + "failed_records": 1, "errors": ["failed future"], } with tempfile.TemporaryDirectory() as output_dir: - _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) - view = _FakeView([failed_summary, failed_summary]) + _write_single_shard( + output_dir, + 10, + [ + _row(10, 0, 1, [1.0, 2.0]), + _row(10, 0, 2, [3.0, 4.0]), + _row(10, 0, 3, [5.0, 6.0]), + ], + ) + view = _FakeView([failed_summary]) uploader = FeatureStoreDeltaUploader( _feature_store_config(max_retries=1), output_dir, @@ -1031,6 +1098,7 @@ def test_flush_failure_does_not_publish_success_marker(self): uploader.submit(10) with self.assertRaises(FeatureStoreUploadError): uploader.close() + self.assertEqual(view.flush_calls, [[2, 1], []]) self.assertFalse( os.path.exists( os.path.join(uploader.state_dir, "step_10._FS_SUCCESS.json") @@ -1212,10 +1280,10 @@ def test_start_clears_a_stale_failure_marker_for_restart(self): self.assertFalse(os.path.exists(uploader._error_marker_path)) uploader.close() - def test_missing_preprovisioned_version_blocks_merge(self): + def test_merge_does_not_require_preprovisioned_version(self): with tempfile.TemporaryDirectory() as output_dir: _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) - view = _FakeView(versions=[]) + view = _FakeView() uploader = FeatureStoreDeltaUploader( _feature_store_config(max_retries=1), output_dir, @@ -1224,14 +1292,14 @@ def test_missing_preprovisioned_version_blocks_merge(self): {"user_emb": 2}, client_factory=_FakeClientFactory(view), ) - with self.assertRaisesRegex( - FeatureStoreUploadError, "version does not exist" - ): - uploader.start() + uploader.start() uploader.close() - self.assertEqual(view.calls, []) + + self.assertEqual(len(view.calls), 1) + self.assertEqual(view.calls[0]["version"], "model_a@export_1") + self.assertEqual(view.calls[0]["write_mode"], "MERGE") self.assertEqual(view.closed, [True]) - self.assertFalse( + self.assertTrue( os.path.exists( os.path.join(uploader.state_dir, "step_10._FS_SUCCESS.json") ) @@ -1313,7 +1381,7 @@ def test_restart_reclaims_snapshot_left_after_committed_upload(self): restarted.close() self.assertFalse(os.path.exists(snapshot_dir)) - def test_non_draining_close_stops_after_inflight_batch_without_commit(self): + def test_non_draining_close_stops_after_inflight_window_without_commit(self): with tempfile.TemporaryDirectory() as output_dir: _write_single_shard( output_dir, @@ -1344,7 +1412,7 @@ def test_non_draining_close_stops_after_inflight_batch_without_commit(self): ) view.release_flush.set() self.assertTrue(view.close_finished.wait(timeout=2)) - self.assertEqual(len(view.calls), 1) + self.assertEqual(len(view.calls), 2) self.assertFalse( os.path.exists( os.path.join(uploader.state_dir, "step_10._FS_SUCCESS.json") From e4cffe73508b483732cbe2c90d651a2f8b7a5390 Mon Sep 17 00:00:00 2001 From: gecheng Date: Thu, 16 Jul 2026 20:06:58 +0800 Subject: [PATCH 03/13] support minute-based delta embedding dumps --- tzrec/protos/train.proto | 8 +- tzrec/utils/delta_embedding_dump.py | 78 ++++++++++++-- tzrec/utils/delta_embedding_dump_test.py | 130 +++++++++++++++++++++-- 3 files changed, 200 insertions(+), 16 deletions(-) diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index 170fad703..b0dc81adc 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -84,8 +84,9 @@ message FeatureStoreConfig { 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 @@ -94,6 +95,9 @@ message DeltaEmbeddingDumpConfig { 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 { diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 5fd1faa24..a69c869e4 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -126,7 +126,17 @@ def validate_delta_embedding_dump_config( "delta_embedding_dump_config only supports CUDA training, " f"but got device={device}." ) - if config.dump_interval_steps <= 0: + if config.HasField("dump_interval_minutes"): + if config.HasField("dump_interval_steps"): + raise ValueError( + "delta_embedding_dump_config must configure only one of " + "dump_interval_steps and dump_interval_minutes." + ) + if config.dump_interval_minutes <= 0: + raise ValueError( + "delta_embedding_dump_config.dump_interval_minutes must be > 0." + ) + elif config.dump_interval_steps <= 0: raise ValueError("delta_embedding_dump_config.dump_interval_steps must be > 0.") if config.HasField("feature_store_config"): validate_feature_store_config(config.feature_store_config) @@ -373,7 +383,14 @@ def __init__( validate_delta_embedding_dump_no_zch_features(feature_configs) self._model = model self._config = config - self._interval = config.dump_interval_steps + self._interval_steps: Optional[int] = None + self._interval_secs: Optional[float] = None + if config.HasField("dump_interval_minutes"): + self._interval_secs = float(config.dump_interval_minutes * 60) + else: + self._interval_steps = int(config.dump_interval_steps) + self._next_dump_time: Optional[float] = None + self._last_dump_step: Optional[int] = None self._output_dir = config.output_dir or os.path.join( model_dir, "delta_embedding_dump" ) @@ -436,10 +453,17 @@ def __init__( embedding_dimensions=embedding_dimensions, ) + interval_name = "minutes" if self._interval_secs is not None else "steps" + interval_value = ( + config.dump_interval_minutes + if self._interval_secs is not None + else self._interval_steps + ) logger.info( - "Delta embedding dump enabled: interval=%s output_dir=%s " + "Delta embedding dump enabled: interval_%s=%s output_dir=%s " "rank=%s/%s tables=%s feature_store_upload=%s", - self._interval, + interval_name, + interval_value, self._output_dir, self._rank, self._world_size, @@ -452,11 +476,13 @@ def clear(self) -> None: self._tracker.clear() def start(self) -> None: - """Start rank-zero background publication after training initialization.""" + """Start timed cadence and rank-zero publication after initialization.""" if self._feature_store_enabled: self._initialize_run_generation() if self._uploader is not None: self._uploader.start() + if self._interval_secs is not None: + self._next_dump_time = time.monotonic() + self._interval_secs def close(self, raise_on_error: bool = True, drain: bool = True) -> None: """Close the rank-zero uploader; abnormal shutdown can skip draining.""" @@ -664,7 +690,7 @@ def pause_tracking(self) -> Iterator[None]: self._tracking_pause_depth -= 1 def maybe_dump(self, global_step: int) -> None: - """Dump on the configured global-step interval and advance tracker state. + """Dump on the configured step or time interval and advance tracker state. Args: global_step: Current training step. @@ -672,10 +698,42 @@ def maybe_dump(self, global_step: int) -> None: # This is a throttled local shared-filesystem check, not a distributed # collective, so all ranks can surface an async rank-zero failure quickly. self._check_feature_store_upload_error() - if global_step > 0 and global_step % self._interval == 0: + if self._should_dump(global_step): self.dump(global_step) + self._last_dump_step = global_step + if self._interval_secs is not None and self._rank == 0: + # Schedule from dump completion to avoid immediate catch-up dumps + # when writing the previous interval took longer than expected. + self._next_dump_time = time.monotonic() + self._interval_secs self._tracker.step() + def _should_dump(self, global_step: int) -> bool: + """Return one rank-consistent dump decision for the current step.""" + if global_step <= 0: + return False + if self._interval_steps is not None: + return global_step % self._interval_steps == 0 + if self._next_dump_time is None: + raise RuntimeError( + "time-based delta embedding dumper must be started before training" + ) + + should_dump = self._rank == 0 and time.monotonic() >= self._next_dump_time + if self._world_size > 1: + if not ( + torch.distributed.is_available() and torch.distributed.is_initialized() + ): + raise RuntimeError( + "distributed time-based delta embedding dump requires an " + "initialized process group" + ) + decision = torch.tensor( + [int(should_dump)], dtype=torch.int32, device=self._device + ) + torch.distributed.broadcast(decision, src=0) + should_dump = bool(decision.item()) + return should_dump + def final_dump(self, global_step: int) -> Optional[str]: """Flush the trailing partial interval at the end of training. @@ -694,7 +752,7 @@ def final_dump(self, global_step: int) -> Optional[str]: # Step zero is excluded from the delta publication contract. logger.info("Skipping delta embedding dump at step %s.", global_step) return None - if global_step > 0 and global_step % self._interval == 0: + if self._interval_steps is not None and global_step % self._interval_steps == 0: # Boundary steps were already written (with full delta) by # ``maybe_dump``. Re-dumping here has no new delta to flush -- every # rank's consumer cursor has already advanced past the boundary's @@ -703,6 +761,10 @@ def final_dump(self, global_step: int) -> Optional[str]: # consumer window. Re-dumping would also overwrite the already-written # boundary shards (with an empty file under multi-GPU), so skip. return None + if self._interval_secs is not None and global_step == self._last_dump_step: + # A timed dump can land on any step. Avoid replacing that step's full + # delta with an empty final shard when training ends immediately after. + return None # The error bit was included in the mandatory final collective. Do not # perform another rank-local marker check here: a marker appearing between # ranks could otherwise make only part of the shard set return early. diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index fccea83b9..a66d4f2c6 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -240,6 +240,22 @@ def test_present_config_requires_positive_interval(self): with self.assertRaisesRegex(ValueError, "dump_interval_steps"): validate_delta_embedding_dump_config(config, torch.device("cuda:0")) + def test_present_config_accepts_minutes_interval(self): + config = DeltaEmbeddingDumpConfig(dump_interval_minutes=5) + validate_delta_embedding_dump_config(config, torch.device("cuda:0")) + + def test_present_config_requires_positive_minutes_interval(self): + config = DeltaEmbeddingDumpConfig(dump_interval_minutes=0) + with self.assertRaisesRegex(ValueError, "dump_interval_minutes"): + validate_delta_embedding_dump_config(config, torch.device("cuda:0")) + + def test_present_config_rejects_both_intervals(self): + config = DeltaEmbeddingDumpConfig( + dump_interval_steps=10, dump_interval_minutes=5 + ) + with self.assertRaisesRegex(ValueError, "only one"): + validate_delta_embedding_dump_config(config, torch.device("cuda:0")) + def test_zch_feature_fails_fast(self): feature_configs = [ feature_pb2.FeatureConfig( @@ -450,7 +466,8 @@ def test_write_table_chunks_leaves_no_partial_shard_on_error(self): def test_final_dump_skips_boundary_step_to_avoid_overwrite(self): dumper = object.__new__(DeltaEmbeddingDumper) - dumper._interval = 50 + dumper._interval_steps = 50 + dumper._interval_secs = None dumper._world_size = 1 with mock.patch.object(dumper, "dump") as dump_mock: # Boundary steps were already written by maybe_dump; skip them so a @@ -480,7 +497,8 @@ def test_final_dump_syncs_step_across_ranks_before_flush(self): # skip and write no shard, leaving step_73/ ragged. The MAX all_reduce # lifts every rank to 73 so all take the same dump-into-step_73 path. dumper = object.__new__(DeltaEmbeddingDumper) - dumper._interval = 50 + dumper._interval_steps = 50 + dumper._interval_secs = None dumper._world_size = 2 def fake_all_reduce(tensor, op=None): @@ -509,7 +527,8 @@ def fake_all_reduce(tensor, op=None): def test_final_dump_syncs_ranks_before_skipping_step_zero(self): dumper = object.__new__(DeltaEmbeddingDumper) - dumper._interval = 50 + dumper._interval_steps = 50 + dumper._interval_secs = None dumper._world_size = 2 collective_count = 0 @@ -540,7 +559,8 @@ def fake_all_reduce(tensor, op=None): def test_final_dump_reduces_upload_error_before_raising(self): dumper = object.__new__(DeltaEmbeddingDumper) - dumper._interval = 50 + dumper._interval_steps = 50 + dumper._interval_secs = None dumper._world_size = 2 def fake_all_reduce(tensor, op=None): @@ -568,7 +588,8 @@ def fake_all_reduce(tensor, op=None): def test_final_dump_reduces_rank_local_dump_failure_before_checkpoint(self): dumper = object.__new__(DeltaEmbeddingDumper) - dumper._interval = 50 + dumper._interval_steps = 50 + dumper._interval_secs = None dumper._world_size = 2 collective_index = 0 @@ -627,7 +648,8 @@ def test_dump_generation_is_stable_per_step_and_unique_per_run(self): def test_maybe_dump_uses_checkpoint_aligned_global_step(self): dumper = object.__new__(DeltaEmbeddingDumper) - dumper._interval = 50 + dumper._interval_steps = 50 + dumper._interval_secs = None dumper._tracker = mock.MagicMock() with mock.patch.object(dumper, "dump") as dump_mock: dumper.maybe_dump(0) @@ -645,6 +667,79 @@ def test_maybe_dump_uses_checkpoint_aligned_global_step(self): ) self.assertEqual(dumper._tracker.step.call_count, 5) + def test_maybe_dump_uses_elapsed_time_and_resets_from_completion(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_steps = None + dumper._interval_secs = 60.0 + dumper._next_dump_time = 160.0 + dumper._last_dump_step = None + dumper._rank = 0 + dumper._world_size = 1 + dumper._device = torch.device("cpu") + dumper._tracker = mock.MagicMock() + with ( + mock.patch.object(dumper, "_check_feature_store_upload_error"), + mock.patch.object(dumper, "dump") as dump_mock, + mock.patch( + "tzrec.utils.delta_embedding_dump.time.monotonic", + side_effect=[159.0, 160.0, 162.0, 221.0, 222.0, 223.0], + ), + ): + dumper.maybe_dump(10) + dumper.maybe_dump(11) + dumper.maybe_dump(12) + dumper.maybe_dump(13) + + self.assertEqual( + [call.args[0] for call in dump_mock.call_args_list], + [11, 13], + ) + self.assertEqual(dumper._next_dump_time, 283.0) + self.assertEqual(dumper._last_dump_step, 13) + self.assertEqual(dumper._tracker.step.call_count, 4) + + def test_time_dump_decision_is_broadcast_from_rank_zero(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_steps = None + dumper._interval_secs = 60.0 + dumper._next_dump_time = 160.0 + dumper._rank = 1 + dumper._world_size = 2 + dumper._device = torch.device("cpu") + + def fake_broadcast(decision, src): + self.assertEqual(src, 0) + decision.fill_(1) + + with ( + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.broadcast", side_effect=fake_broadcast), + ): + self.assertTrue(dumper._should_dump(10)) + + def test_final_dump_skips_step_already_dumped_by_time_interval(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_steps = None + dumper._interval_secs = 60.0 + dumper._last_dump_step = 73 + dumper._world_size = 1 + with mock.patch.object(dumper, "dump") as dump_mock: + self.assertIsNone(dumper.final_dump(73)) + dump_mock.assert_not_called() + + def test_start_initializes_minutes_interval_from_training_start(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._feature_store_enabled = False + dumper._uploader = None + dumper._interval_secs = 120.0 + dumper._next_dump_time = None + with mock.patch( + "tzrec.utils.delta_embedding_dump.time.monotonic", return_value=100.0 + ): + dumper.start() + self.assertEqual(dumper._next_dump_time, 220.0) + def test_direct_dump_rejects_step_zero(self): dumper = object.__new__(DeltaEmbeddingDumper) with self.assertRaisesRegex(ValueError, "global_step must be > 0"): @@ -672,6 +767,29 @@ def test_tracker_uses_auto_compact(self): self.assertTrue(tracker_cls.call_args.kwargs["auto_compact"]) + def test_minutes_interval_is_converted_to_seconds(self): + tracker = mock.MagicMock() + tracker.table_to_fqn = {} + tracker.fqn_to_feature_names.return_value = {} + tracker.get_tracked_modules.return_value = {} + with ( + tempfile.TemporaryDirectory() as tmp_dir, + mock.patch( + "tzrec.utils.delta_embedding_dump.ModelDeltaTrackerTrec", + return_value=tracker, + ), + ): + dumper = DeltaEmbeddingDumper( + torch.nn.Module(), + DeltaEmbeddingDumpConfig(dump_interval_minutes=2), + tmp_dir, + torch.device("cuda"), + [], + ) + + self.assertIsNone(dumper._interval_steps) + self.assertEqual(dumper._interval_secs, 120.0) + def test_durable_ack_advances_guard_without_recomputing_unique_rows(self): tracker = mock.MagicMock() tracker.per_consumer_batch_idx = { From 953a70cf9ab14196deee8b2d085ee3025bd36fd7 Mon Sep 17 00:00:00 2001 From: gecheng Date: Thu, 16 Jul 2026 20:38:31 +0800 Subject: [PATCH 04/13] add FeatureStore delta readback checker --- tzrec/tools/feature_store/__init__.py | 10 + .../check_feature_store_delta.py | 590 ++++++++++++++++++ .../check_feature_store_delta_test.py | 173 +++++ 3 files changed, 773 insertions(+) create mode 100644 tzrec/tools/feature_store/__init__.py create mode 100644 tzrec/tools/feature_store/check_feature_store_delta.py create mode 100644 tzrec/tools/feature_store/check_feature_store_delta_test.py diff --git a/tzrec/tools/feature_store/__init__.py b/tzrec/tools/feature_store/__init__.py new file mode 100644 index 000000000..eedc773bc --- /dev/null +++ b/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. diff --git a/tzrec/tools/feature_store/check_feature_store_delta.py b/tzrec/tools/feature_store/check_feature_store_delta.py new file mode 100644 index 000000000..c7c5cdd48 --- /dev/null +++ b/tzrec/tools/feature_store/check_feature_store_delta.py @@ -0,0 +1,590 @@ +# 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. + +r"""Read back sampled delta embeddings from FeatureStore. + +The tool uses the latest locally committed FeatureStore upload by default, samples +keys from its canonical parquet shard set, and queries the configured explicit +FeatureDB version through ``DynamicEmbeddingFeatureView.get_online_features``. + +Example:: + + python -m tzrec.tools.feature_store.check_feature_store_delta \ + --pipeline_config path/to/pipeline.config \ + --ak YOUR_ACCESS_KEY_ID \ + --sk YOUR_ACCESS_KEY_SECRET \ + --featuredb_username YOUR_FEATUREDB_USERNAME \ + --featuredb_password YOUR_FEATUREDB_PASSWORD \ + --output_dir path/to/delta_embedding_dump \ + --sample_count 10 +""" + +import argparse +import glob +import inspect +import json +import os +import sys +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Any, DefaultDict, Dict, List, Mapping, Optional, Sequence, Tuple + +import numpy as np +import numpy.typing as npt +import pyarrow.parquet as pq + +from tzrec.utils import config_util +from tzrec.utils.feature_store_delta_uploader import ( + DELTA_OPERATION_UPSERT, + FEATURE_STORE_PK_FIELD, + FEATURE_STORE_SK_FIELD, + FEATURE_STORE_VALUE_FIELD, + FeatureStoreUploadSettings, + feature_store_delta_file_prefix, +) + + +@dataclass(frozen=True) +class LocalSample: + """One local parquet record selected for remote readback.""" + + embedding_name: str + key_id: int + embedding: npt.NDArray[np.float32] = field(repr=False) + source_path: str + + +def resolve_output_dir( + pipeline_config_path: str, + model_dir: str, + configured_output_dir: str, + output_dir_override: Optional[str], +) -> str: + """Resolve the local delta outbox, including relocated pipeline configs. + + Args: + pipeline_config_path: Source pipeline config path. + model_dir: Model directory from the pipeline config. + configured_output_dir: Explicit delta dump output directory, if any. + output_dir_override: Command-line output directory override, if any. + + Returns: + Absolute delta outbox directory path. + """ + if output_dir_override: + return os.path.abspath(output_dir_override) + if configured_output_dir: + return os.path.abspath(configured_output_dir) + + configured_path = os.path.abspath(os.path.join(model_dir, "delta_embedding_dump")) + if os.path.isdir(configured_path): + return configured_path + + colocated_path = os.path.join( + os.path.dirname(os.path.abspath(pipeline_config_path)), + "delta_embedding_dump", + ) + if os.path.isdir(colocated_path): + return colocated_path + return configured_path + + +def _read_json_object(path: str) -> Dict[str, Any]: + """Read one JSON object from disk.""" + with open(path) as source: + value = json.load(source) + if not isinstance(value, dict): + raise ValueError(f"expected a JSON object in {path}") + return value + + +def _target_matches( + value: Mapping[str, Any], settings: FeatureStoreUploadSettings +) -> bool: + """Return whether a credential-free marker identifies this remote target.""" + return ( + value.get("project_name") == settings.project_name + and value.get("feature_view_name") == settings.feature_view_name + and value.get("version") == settings.version + ) + + +def load_committed_upload( + output_dir: str, + settings: FeatureStoreUploadSettings, + global_step: Optional[int] = None, +) -> Tuple[str, Dict[str, Any], Dict[str, Any]]: + """Load the local committed watermark and its success marker. + + Args: + output_dir: Delta parquet outbox directory. + settings: Resolved FeatureStore target settings. + global_step: Specific committed step to inspect, or None for the latest. + + Returns: + Tuple of state directory, committed watermark, and success marker. + + Raises: + FileNotFoundError: If no committed marker exists for the configured target. + ValueError: If markers are ambiguous, inconsistent, or the step is invalid. + """ + state_root = os.path.join(output_dir, ".feature_store_upload") + committed_candidates: List[Tuple[str, Dict[str, Any]]] = [] + for path in sorted(glob.glob(os.path.join(state_root, "*", "committed.json"))): + committed = _read_json_object(path) + if _target_matches(committed, settings): + committed_candidates.append((os.path.dirname(path), committed)) + + if not committed_candidates: + raise FileNotFoundError( + "no committed FeatureStore upload marker was found for " + f"project={settings.project_name!r}, " + f"view={settings.feature_view_name!r}, version={settings.version!r} " + f"under {state_root}" + ) + if len(committed_candidates) > 1: + raise ValueError( + "multiple local FeatureStore state directories match the configured " + f"target under {state_root}; use a target-specific output_dir" + ) + + state_dir, committed = committed_candidates[0] + committed_step = int(committed.get("committed_global_step", -1)) + selected_step = committed_step if global_step is None else int(global_step) + if selected_step <= 0: + raise ValueError("global_step must be > 0") + if selected_step > committed_step: + raise ValueError( + f"step {selected_step} is newer than the locally committed " + f"FeatureStore watermark {committed_step}" + ) + + success_path = os.path.join(state_dir, f"step_{selected_step}._FS_SUCCESS.json") + if not os.path.isfile(success_path): + raise FileNotFoundError( + f"committed FeatureStore success marker was not found: {success_path}" + ) + success = _read_json_object(success_path) + if not _target_matches(success, settings): + raise ValueError(f"FeatureStore success marker target mismatch: {success_path}") + if int(success.get("global_step", -1)) != selected_step: + raise ValueError(f"FeatureStore success marker step mismatch: {success_path}") + if int(success.get("success_records", -1)) != int(success.get("total_records", -2)): + raise ValueError( + f"FeatureStore success marker is not fully successful: {success_path}" + ) + return state_dir, committed, success + + +def committed_parquet_paths( + output_dir: str, + file_prefix: str, + global_step: int, + expected_shards: int, +) -> List[str]: + """Resolve the canonical parquet shards for one committed upload step.""" + single_rank_path = os.path.join( + output_dir, f"{file_prefix}_step_{global_step}.parquet" + ) + if os.path.isfile(single_rank_path): + paths = [single_rank_path] + else: + paths = sorted( + glob.glob( + os.path.join( + output_dir, + f"step_{global_step}", + f"{file_prefix}_step_{global_step}_rank_*_of_*.parquet", + ) + ) + ) + if not paths: + raise FileNotFoundError( + f"no canonical delta parquet shards found for committed step {global_step}" + ) + if expected_shards > 0 and len(paths) != expected_shards: + raise ValueError( + f"committed step {global_step} has {len(paths)} local parquet shards, " + f"but its success marker records {expected_shards}" + ) + return paths + + +def sample_local_records( + parquet_paths: Sequence[str], + sample_count: int, + embedding_name: Optional[str] = None, +) -> List[LocalSample]: + """Read a bounded set of unique UPSERT records from canonical parquet shards.""" + if sample_count <= 0: + raise ValueError("sample_count must be > 0") + + columns = [ + FEATURE_STORE_PK_FIELD, + FEATURE_STORE_SK_FIELD, + FEATURE_STORE_VALUE_FIELD, + "operation", + ] + samples: List[LocalSample] = [] + seen: set[Tuple[str, int]] = set() + for path in parquet_paths: + parquet_file = pq.ParquetFile(path) + missing_columns = [ + name for name in columns if name not in parquet_file.schema_arrow.names + ] + if missing_columns: + raise ValueError( + f"delta parquet {path} is missing columns {missing_columns}" + ) + for batch in parquet_file.iter_batches(batch_size=1024, columns=columns): + values = batch.to_pydict() + for name, key_id, vector, operation in zip( + values[FEATURE_STORE_PK_FIELD], + values[FEATURE_STORE_SK_FIELD], + values[FEATURE_STORE_VALUE_FIELD], + values["operation"], + ): + if operation != DELTA_OPERATION_UPSERT: + continue + name = str(name) + if embedding_name is not None and name != embedding_name: + continue + identity = (name, int(key_id)) + if identity in seen: + continue + if vector is None or len(vector) == 0: + raise ValueError( + f"delta parquet {path} contains an empty embedding" + ) + seen.add(identity) + samples.append( + LocalSample( + embedding_name=name, + key_id=int(key_id), + embedding=np.asarray(vector, dtype=np.float32), + source_path=path, + ) + ) + if len(samples) >= sample_count: + return samples + if not samples: + suffix = f" for embedding_name={embedding_name!r}" if embedding_name else "" + raise ValueError(f"no sampleable UPSERT delta records were found{suffix}") + return samples + + +def _normalize_remote_key(value: Any) -> int: + """Normalize the SDK's string/bytes/integer SK representation.""" + if isinstance(value, bytes): + value = value.decode("utf-8") + return int(value) + + +def verify_samples( + view: Any, + version: str, + samples: Sequence[LocalSample], +) -> Tuple[List[Dict[str, Any]], Dict[str, int]]: + """Query sampled keys and classify presence and value equality.""" + grouped: DefaultDict[str, List[LocalSample]] = defaultdict(list) + for sample in samples: + grouped[sample.embedding_name].append(sample) + + results: List[Dict[str, Any]] = [] + for embedding_name, group in grouped.items(): + remote_rows = view.get_online_features( + feature_name=embedding_name, + keys=[sample.key_id for sample in group], + version=version, + ) + remote_by_key: Dict[int, npt.NDArray[np.float32]] = {} + for row in remote_rows: + raw_key = row.get("sk", row.get(FEATURE_STORE_SK_FIELD)) + if raw_key is None: + raise ValueError("FeatureStore readback row is missing its SK") + key_id = _normalize_remote_key(raw_key) + if key_id in remote_by_key: + raise ValueError( + "FeatureStore returned duplicate rows for " + f"{embedding_name}/{key_id}" + ) + vector = row.get(FEATURE_STORE_VALUE_FIELD) + if vector is None: + raise ValueError( + f"FeatureStore readback row is missing embedding for " + f"{embedding_name}/{key_id}" + ) + remote_by_key[key_id] = np.asarray(vector, dtype=np.float32) + + for sample in group: + remote = remote_by_key.get(sample.key_id) + result: Dict[str, Any] = { + "embedding_name": sample.embedding_name, + "key_id": sample.key_id, + "local_dimension": int(sample.embedding.size), + "source_path": sample.source_path, + } + if remote is None: + result.update( + { + "status": "MISSING", + "remote_dimension": None, + "remote_embedding": None, + } + ) + else: + same_shape = remote.shape == sample.embedding.shape + matches = same_shape and bool( + np.allclose(remote, sample.embedding, rtol=1e-5, atol=1e-6) + ) + max_abs_diff = ( + float(np.max(np.abs(remote - sample.embedding))) + if same_shape and remote.size > 0 + else None + ) + result.update( + { + "status": "MATCH" if matches else "PRESENT_DIFFERENT", + "remote_dimension": int(remote.size), + "remote_embedding": remote.tolist(), + "max_abs_diff": max_abs_diff, + } + ) + results.append(result) + + summary = { + "requested": len(results), + "found": sum(result["status"] != "MISSING" for result in results), + "matching": sum(result["status"] == "MATCH" for result in results), + "present_different": sum( + result["status"] == "PRESENT_DIFFERENT" for result in results + ), + "missing": sum(result["status"] == "MISSING" for result in results), + } + return results, summary + + +def create_feature_store_view(settings: FeatureStoreUploadSettings) -> Any: + """Create the SDK client and return the existing DynamicEmbedding view.""" + try: + from feature_store_py import FeatureStoreClient + except ImportError as exc: + raise RuntimeError( + "feature_store_py is required; install requirements/feature_store.txt" + ) from exc + + kwargs = { + "access_key_id": settings.access_key_id, + "access_key_secret": settings.access_key_secret, + "region": settings.region or None, + "endpoint": settings.endpoint or None, + "security_token": settings.security_token or None, + "featuredb_username": settings.featuredb_username or None, + "featuredb_password": settings.featuredb_password or None, + } + try: + parameters = inspect.signature(FeatureStoreClient).parameters + except (TypeError, ValueError): + parameters = {} + if "test_mode" in parameters: + # Match the uploader's FeatureDB routing when supported by the SDK wheel. + kwargs["test_mode"] = True + + client = FeatureStoreClient(**kwargs) + project = client.get_project(settings.project_name) + if project is None: + raise RuntimeError( + f"FeatureStore project {settings.project_name!r} was not found" + ) + view = project.get_dynamic_embedding_feature_view(settings.feature_view_name) + if view is None: + raise RuntimeError( + f"DynamicEmbedding FeatureView {settings.feature_view_name!r} was not found" + ) + actual_fields = (view.pk_field, view.sk_field, view.embedding_field) + expected_fields = ( + FEATURE_STORE_PK_FIELD, + FEATURE_STORE_SK_FIELD, + FEATURE_STORE_VALUE_FIELD, + ) + if actual_fields != expected_fields: + raise RuntimeError( + "DynamicEmbedding FeatureView schema mismatch: " + f"expected={expected_fields}, actual={actual_fields}" + ) + return view + + +def run_check(args: argparse.Namespace) -> int: + """Run one local-marker plus remote-readback verification.""" + pipeline_config = config_util.load_pipeline_config(args.pipeline_config) + train_config = pipeline_config.train_config + if not train_config.HasField("delta_embedding_dump_config"): + raise ValueError("pipeline config has no delta_embedding_dump_config") + dump_config = train_config.delta_embedding_dump_config + if not dump_config.HasField("feature_store_config"): + raise ValueError( + "pipeline config delta_embedding_dump_config has no feature_store_config" + ) + feature_store_config = dump_config.feature_store_config + feature_store_config.access_key_id = args.access_key_id + feature_store_config.access_key_secret = args.access_key_secret + feature_store_config.featuredb_username = args.featuredb_username + feature_store_config.featuredb_password = args.featuredb_password + settings = FeatureStoreUploadSettings.from_proto(feature_store_config) + output_dir = resolve_output_dir( + args.pipeline_config, + pipeline_config.model_dir, + dump_config.output_dir, + args.output_dir, + ) + if not os.path.isdir(output_dir): + raise FileNotFoundError( + f"delta embedding output directory not found: {output_dir}" + ) + + _, committed, success = load_committed_upload( + output_dir, settings, args.global_step + ) + global_step = int(success["global_step"]) + base_prefix = dump_config.file_prefix or "delta_embedding" + scoped_prefix = feature_store_delta_file_prefix(feature_store_config, base_prefix) + shard_count = len(success.get("shards", [])) + parquet_paths = committed_parquet_paths( + output_dir, + scoped_prefix, + global_step, + expected_shards=shard_count, + ) + samples = sample_local_records( + parquet_paths, + args.sample_count, + embedding_name=args.embedding_name, + ) + + view = create_feature_store_view(settings) + try: + results, summary = verify_samples(view, settings.version, samples) + finally: + close = getattr(view, "close", None) + if callable(close): + close(wait=True) + + report: Dict[str, Any] = { + "target": { + "project_name": settings.project_name, + "feature_view_name": settings.feature_view_name, + "version": settings.version, + }, + "local_commit": { + "global_step": global_step, + "committed_watermark": int(committed["committed_global_step"]), + "success_records": int(success["success_records"]), + "total_records": int(success["total_records"]), + "parquet_paths": [ + os.path.relpath(path, output_dir) for path in parquet_paths + ], + }, + "summary": summary, + "presence_verified": summary["missing"] == 0, + "value_match_verified": summary["matching"] == summary["requested"], + "samples": results, + } + if summary["present_different"]: + report["value_match_note"] = ( + "PRESENT_DIFFERENT confirms the key exists but its value differs from " + "the sampled committed parquet; a later in-flight step may have updated " + "the same key in this version." + ) + print(json.dumps(report, indent=2, sort_keys=True)) + + if summary["missing"]: + return 1 + if args.require_value_match and summary["present_different"]: + return 1 + return 0 + + +def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser( + description=( + "Sample a committed delta parquet and read the same keys from " + "FeatureStore using the configured explicit version." + ) + ) + parser.add_argument("--pipeline_config", required=True) + parser.add_argument( + "--access_key_id", + "--ak", + dest="access_key_id", + required=True, + help="Alibaba Cloud AccessKey ID.", + ) + parser.add_argument( + "--access_key_secret", + "--sk", + dest="access_key_secret", + required=True, + help="Alibaba Cloud AccessKey Secret.", + ) + parser.add_argument( + "--featuredb_username", + required=True, + help="FeatureDB username.", + ) + parser.add_argument( + "--featuredb_password", + required=True, + help="FeatureDB password.", + ) + parser.add_argument( + "--output_dir", + default=None, + help="Override delta_embedding_dump output_dir (useful across mounts).", + ) + parser.add_argument( + "--global_step", + type=int, + default=None, + help="Committed step to inspect; defaults to the latest committed watermark.", + ) + parser.add_argument( + "--sample_count", + type=int, + default=10, + help="Maximum number of local keys to read back (default: 10).", + ) + parser.add_argument( + "--embedding_name", + default=None, + help="Only sample rows for this canonical embedding name.", + ) + parser.add_argument( + "--require_value_match", + action="store_true", + help="Exit nonzero when a key exists but differs from the sampled local value.", + ) + return parser.parse_args(argv) + + +def main() -> None: + """CLI entry point.""" + try: + exit_code = run_check(parse_args()) + except Exception as exc: + print(f"ERROR: {exc}", file=sys.stderr) + exit_code = 2 + raise SystemExit(exit_code) + + +if __name__ == "__main__": + main() diff --git a/tzrec/tools/feature_store/check_feature_store_delta_test.py b/tzrec/tools/feature_store/check_feature_store_delta_test.py new file mode 100644 index 000000000..f8f0e53b2 --- /dev/null +++ b/tzrec/tools/feature_store/check_feature_store_delta_test.py @@ -0,0 +1,173 @@ +# 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. + +import json +import os +import tempfile +import unittest +from types import SimpleNamespace + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +from tzrec.tools.feature_store.check_feature_store_delta import ( + LocalSample, + committed_parquet_paths, + load_committed_upload, + parse_args, + resolve_output_dir, + sample_local_records, + verify_samples, +) + + +class _FakeView: + def get_online_features(self, feature_name, keys, version): + self.feature_name = feature_name + self.keys = keys + self.version = version + return [ + {"sk": str(keys[0]), "embedding": [1.0, 2.0]}, + {"sk": str(keys[1]), "embedding": [9.0, 9.0]}, + ] + + +class CheckFeatureStoreDeltaTest(unittest.TestCase): + def test_parse_args_requires_command_line_access_key_pair(self): + args = parse_args( + [ + "--pipeline_config", + "pipeline.config", + "--ak", + "access-key-id", + "--sk", + "access-key-secret", + "--featuredb_username", + "featuredb-user", + "--featuredb_password", + "featuredb-password", + ] + ) + self.assertEqual(args.access_key_id, "access-key-id") + self.assertEqual(args.access_key_secret, "access-key-secret") + self.assertEqual(args.featuredb_username, "featuredb-user") + self.assertEqual(args.featuredb_password, "featuredb-password") + + def test_resolve_output_dir_uses_colocated_relocated_outbox(self): + with tempfile.TemporaryDirectory() as tmp_dir: + config_path = os.path.join(tmp_dir, "pipeline.config") + colocated = os.path.join(tmp_dir, "delta_embedding_dump") + os.mkdir(colocated) + self.assertEqual( + resolve_output_dir(config_path, "/missing/model", "", None), + colocated, + ) + + def test_load_committed_upload_checks_target_and_success(self): + settings = SimpleNamespace( + project_name="project", feature_view_name="view", version="v1" + ) + with tempfile.TemporaryDirectory() as output_dir: + state_dir = os.path.join(output_dir, ".feature_store_upload", "target") + os.makedirs(state_dir) + target = { + "project_name": "project", + "feature_view_name": "view", + "version": "v1", + } + with open(os.path.join(state_dir, "committed.json"), "w") as output: + json.dump({**target, "committed_global_step": 20}, output) + with open( + os.path.join(state_dir, "step_20._FS_SUCCESS.json"), "w" + ) as output: + json.dump( + { + **target, + "global_step": 20, + "success_records": 3, + "total_records": 3, + "shards": [{}], + }, + output, + ) + + actual_state_dir, committed, success = load_committed_upload( + output_dir, settings + ) + + self.assertEqual(actual_state_dir, state_dir) + self.assertEqual(committed["committed_global_step"], 20) + self.assertEqual(success["global_step"], 20) + + def test_committed_parquet_paths_and_sampling(self): + with tempfile.TemporaryDirectory() as output_dir: + path = os.path.join(output_dir, "delta__fs_target_step_20.parquet") + pq.write_table( + pa.table( + { + "embedding_name": ["table_a", "table_a", "table_b"], + "key_id": pa.array([1, 2, 3], type=pa.int64()), + "embedding": pa.array( + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], + type=pa.list_(pa.float32()), + ), + "operation": ["UPSERT", "UPSERT", "UPSERT"], + } + ), + path, + ) + paths = committed_parquet_paths( + output_dir, "delta__fs_target", 20, expected_shards=1 + ) + samples = sample_local_records(paths, 2) + + self.assertEqual(paths, [path]) + self.assertEqual( + [(sample.embedding_name, sample.key_id) for sample in samples], + [("table_a", 1), ("table_a", 2)], + ) + np.testing.assert_array_equal(samples[0].embedding, [1.0, 2.0]) + + def test_verify_samples_separates_presence_from_value_match(self): + samples = [ + LocalSample("table_a", 1, np.array([1.0, 2.0], np.float32), "a"), + LocalSample("table_a", 2, np.array([3.0, 4.0], np.float32), "a"), + LocalSample("table_a", 3, np.array([5.0, 6.0], np.float32), "a"), + ] + view = _FakeView() + + results, summary = verify_samples(view, "v1", samples) + + self.assertEqual(view.feature_name, "table_a") + self.assertEqual(view.keys, [1, 2, 3]) + self.assertEqual(view.version, "v1") + self.assertEqual( + [result["status"] for result in results], + ["MATCH", "PRESENT_DIFFERENT", "MISSING"], + ) + self.assertEqual(results[0]["remote_embedding"], [1.0, 2.0]) + self.assertEqual(results[1]["remote_embedding"], [9.0, 9.0]) + self.assertIsNone(results[2]["remote_embedding"]) + self.assertEqual( + summary, + { + "requested": 3, + "found": 2, + "matching": 1, + "present_different": 1, + "missing": 1, + }, + ) + + +if __name__ == "__main__": + unittest.main() From bb4f1d39a3d0b41fbe4e77d6f775e10a7213a6b2 Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 17 Jul 2026 10:56:57 +0800 Subject: [PATCH 05/13] put feature store credentials into envs --- tzrec/protos/train.proto | 13 +-- .../check_feature_store_delta.py | 36 +------- .../check_feature_store_delta_test.py | 27 ++---- tzrec/utils/config_util.py | 15 +--- tzrec/utils/config_util_test.py | 24 +----- tzrec/utils/feature_store_delta_uploader.py | 78 +++++++----------- .../feature_store_delta_uploader_test.py | 82 ++++++++----------- 7 files changed, 83 insertions(+), 192 deletions(-) diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index b0dc81adc..7e2785ec7 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -30,18 +30,13 @@ message GradClipping { } 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; - // Runtime credentials. Explicit values take precedence over the standard - // ALIBABA_CLOUD_* environment variables. Persisted model artifacts replace - // these required secret values with empty strings. - required string access_key_id = 2; - required string access_key_secret = 3; - // FeatureDB data-plane credential pair. Explicit values take - // precedence over FEATUREDB_USERNAME and FEATUREDB_PASSWORD. - required string featuredb_username = 4; - required string featuredb_password = 5; // 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; diff --git a/tzrec/tools/feature_store/check_feature_store_delta.py b/tzrec/tools/feature_store/check_feature_store_delta.py index c7c5cdd48..b2d9e08a0 100644 --- a/tzrec/tools/feature_store/check_feature_store_delta.py +++ b/tzrec/tools/feature_store/check_feature_store_delta.py @@ -17,12 +17,12 @@ Example:: + export ALIBABA_CLOUD_ACCESS_KEY_ID=YOUR_ACCESS_KEY_ID + export ALIBABA_CLOUD_ACCESS_KEY_SECRET=YOUR_ACCESS_KEY_SECRET + export FEATUREDB_USERNAME=YOUR_FEATUREDB_USERNAME + export FEATUREDB_PASSWORD=YOUR_FEATUREDB_PASSWORD python -m tzrec.tools.feature_store.check_feature_store_delta \ --pipeline_config path/to/pipeline.config \ - --ak YOUR_ACCESS_KEY_ID \ - --sk YOUR_ACCESS_KEY_SECRET \ - --featuredb_username YOUR_FEATUREDB_USERNAME \ - --featuredb_password YOUR_FEATUREDB_PASSWORD \ --output_dir path/to/delta_embedding_dump \ --sample_count 10 """ @@ -435,10 +435,6 @@ def run_check(args: argparse.Namespace) -> int: "pipeline config delta_embedding_dump_config has no feature_store_config" ) feature_store_config = dump_config.feature_store_config - feature_store_config.access_key_id = args.access_key_id - feature_store_config.access_key_secret = args.access_key_secret - feature_store_config.featuredb_username = args.featuredb_username - feature_store_config.featuredb_password = args.featuredb_password settings = FeatureStoreUploadSettings.from_proto(feature_store_config) output_dir = resolve_output_dir( args.pipeline_config, @@ -522,30 +518,6 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: ) ) parser.add_argument("--pipeline_config", required=True) - parser.add_argument( - "--access_key_id", - "--ak", - dest="access_key_id", - required=True, - help="Alibaba Cloud AccessKey ID.", - ) - parser.add_argument( - "--access_key_secret", - "--sk", - dest="access_key_secret", - required=True, - help="Alibaba Cloud AccessKey Secret.", - ) - parser.add_argument( - "--featuredb_username", - required=True, - help="FeatureDB username.", - ) - parser.add_argument( - "--featuredb_password", - required=True, - help="FeatureDB password.", - ) parser.add_argument( "--output_dir", default=None, diff --git a/tzrec/tools/feature_store/check_feature_store_delta_test.py b/tzrec/tools/feature_store/check_feature_store_delta_test.py index f8f0e53b2..97c1e6bf3 100644 --- a/tzrec/tools/feature_store/check_feature_store_delta_test.py +++ b/tzrec/tools/feature_store/check_feature_store_delta_test.py @@ -42,25 +42,14 @@ def get_online_features(self, feature_name, keys, version): class CheckFeatureStoreDeltaTest(unittest.TestCase): - def test_parse_args_requires_command_line_access_key_pair(self): - args = parse_args( - [ - "--pipeline_config", - "pipeline.config", - "--ak", - "access-key-id", - "--sk", - "access-key-secret", - "--featuredb_username", - "featuredb-user", - "--featuredb_password", - "featuredb-password", - ] - ) - self.assertEqual(args.access_key_id, "access-key-id") - self.assertEqual(args.access_key_secret, "access-key-secret") - self.assertEqual(args.featuredb_username, "featuredb-user") - self.assertEqual(args.featuredb_password, "featuredb-password") + def test_parse_args_does_not_accept_credentials(self): + args = parse_args(["--pipeline_config", "pipeline.config"]) + + self.assertEqual(args.pipeline_config, "pipeline.config") + self.assertFalse(hasattr(args, "access_key_id")) + self.assertFalse(hasattr(args, "access_key_secret")) + self.assertFalse(hasattr(args, "featuredb_username")) + self.assertFalse(hasattr(args, "featuredb_password")) def test_resolve_output_dir_uses_colocated_relocated_outbox(self): with tempfile.TemporaryDirectory() as tmp_dir: diff --git a/tzrec/utils/config_util.py b/tzrec/utils/config_util.py index d62d72b76..cd757b002 100644 --- a/tzrec/utils/config_util.py +++ b/tzrec/utils/config_util.py @@ -82,24 +82,15 @@ def save_message(message: Message, filepath: str) -> None: "poll_interval_secs", ) -_FEATURE_STORE_ARTIFACT_REQUIRED_SECRETS = ( - "access_key_id", - "access_key_secret", - "featuredb_username", - "featuredb_password", -) - def sanitize_pipeline_config_for_artifact( pipeline_config: pipeline_pb2.EasyRecConfig, ) -> pipeline_pb2.EasyRecConfig: - """Return an artifact-safe config without FeatureStore credentials. + """Return an artifact-safe config without runtime FeatureStore secrets. The runtime config remains unchanged. FeatureStore public routing and version fields are copied through an allowlist, so newly introduced fields - do not silently become persistent secrets. Required credential fields are - set to explicit empty strings to keep the protobuf initialized while still - allowing the runtime environment fallbacks. + do not silently become persistent secrets. """ sanitized = pipeline_pb2.EasyRecConfig() sanitized.CopyFrom(pipeline_config) @@ -115,8 +106,6 @@ def sanitize_pipeline_config_for_artifact( for field_name in _FEATURE_STORE_ARTIFACT_FIELDS: if runtime_config.HasField(field_name): setattr(public_config, field_name, getattr(runtime_config, field_name)) - for field_name in _FEATURE_STORE_ARTIFACT_REQUIRED_SECRETS: - setattr(public_config, field_name, "") dump_config.ClearField("feature_store_config") dump_config.feature_store_config.CopyFrom(public_config) return sanitized diff --git a/tzrec/utils/config_util_test.py b/tzrec/utils/config_util_test.py index be6182049..c9721f01f 100644 --- a/tzrec/utils/config_util_test.py +++ b/tzrec/utils/config_util_test.py @@ -42,7 +42,7 @@ def test_edit_config(self): ) self.assertEqual(pipeline_config.feature_configs[4].id_feature.num_buckets, 3) - def test_pipeline_artifact_redacts_feature_store_credentials(self): + def test_pipeline_artifact_redacts_feature_store_security_token(self): pipeline_config = config_util.load_pipeline_config( "examples/multi_tower_taobao.config" ) @@ -57,11 +57,7 @@ def test_pipeline_artifact_redacts_feature_store_credentials(self): feature_store_config.feature_view_shard_count = 4 feature_store_config.feature_view_replication_count = 2 feature_store_config.version = "model_a@export_1" - feature_store_config.access_key_id = "SECRET_AK_ID" - feature_store_config.access_key_secret = "SECRET_AK_VALUE" feature_store_config.security_token = "SECRET_STS" - feature_store_config.featuredb_username = "SECRET_FDB_USER" - feature_store_config.featuredb_password = "SECRET_FDB_PASSWORD" sanitized = config_util.sanitize_pipeline_config_for_artifact(pipeline_config) sanitized_fs = ( @@ -75,15 +71,6 @@ def test_pipeline_artifact_redacts_feature_store_credentials(self): self.assertEqual(sanitized_fs.feature_view_replication_count, 2) self.assertEqual(sanitized_fs.version, "model_a@export_1") self.assertTrue(sanitized_fs.IsInitialized()) - for field_name in ( - "access_key_id", - "access_key_secret", - "featuredb_username", - "featuredb_password", - ): - self.assertTrue(sanitized_fs.HasField(field_name)) - self.assertEqual(getattr(sanitized_fs, field_name), "") - self.assertTrue(feature_store_config.HasField(field_name)) self.assertFalse(sanitized_fs.HasField("security_token")) self.assertTrue(feature_store_config.HasField("security_token")) @@ -92,14 +79,7 @@ def test_pipeline_artifact_redacts_feature_store_credentials(self): config_util.save_pipeline_config_artifact(pipeline_config, path) with open(path) as source: artifact_text = source.read() - for secret in ( - "SECRET_AK_ID", - "SECRET_AK_VALUE", - "SECRET_STS", - "SECRET_FDB_USER", - "SECRET_FDB_PASSWORD", - ): - self.assertNotIn(secret, artifact_text) + self.assertNotIn("SECRET_STS", artifact_text) if __name__ == "__main__": diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index 2a9b559c2..54d6a9d6f 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -89,24 +89,6 @@ class _ShardSetNotReady(RuntimeError): """A multi-rank canonical shard set is still being atomically replaced.""" -def _config_or_env(config: FeatureStoreConfig, field_name: str, env_name: str) -> str: - value = getattr(config, field_name, "") - return value or os.environ.get(env_name, "") - - -def _validate_pair( - first: str, - second: str, - first_name: str, - second_name: str, - required: bool, -) -> None: - if bool(first) != bool(second): - raise ValueError(f"{first_name} and {second_name} must be configured together") - if required and not first: - raise ValueError(f"{first_name} and {second_name} are required") - - @dataclass(frozen=True) class FeatureStoreUploadSettings: """Validated immutable settings copied from the runtime protobuf.""" @@ -135,30 +117,43 @@ class FeatureStoreUploadSettings: @classmethod def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": - """Resolve standard environment fallbacks without logging credentials.""" + """Resolve environment-backed credentials without logging them.""" initialization_errors = config.FindInitializationErrors() if initialization_errors: raise ValueError( "feature_store_config is missing required fields: " + ", ".join(initialization_errors) ) - region = _config_or_env(config, "region", "ALIBABA_CLOUD_REGION") + region = config.region or os.environ.get("ALIBABA_CLOUD_REGION", "") endpoint = config.endpoint - access_key_id = _config_or_env( - config, "access_key_id", "ALIBABA_CLOUD_ACCESS_KEY_ID" - ) - access_key_secret = _config_or_env( - config, "access_key_secret", "ALIBABA_CLOUD_ACCESS_KEY_SECRET" - ) - security_token = _config_or_env( - config, "security_token", "ALIBABA_CLOUD_SECURITY_TOKEN" - ) - featuredb_username = _config_or_env( - config, "featuredb_username", "FEATUREDB_USERNAME" - ) - featuredb_password = _config_or_env( - config, "featuredb_password", "FEATUREDB_PASSWORD" + security_token = config.security_token or os.environ.get( + "ALIBABA_CLOUD_SECURITY_TOKEN", "" ) + credential_env_names = { + "access_key_id": "ALIBABA_CLOUD_ACCESS_KEY_ID", + "access_key_secret": "ALIBABA_CLOUD_ACCESS_KEY_SECRET", + "featuredb_username": "FEATUREDB_USERNAME", + "featuredb_password": "FEATUREDB_PASSWORD", + } + credentials = { + field_name: os.environ.get(env_name, "") + for field_name, env_name in credential_env_names.items() + } + missing_env_names = [ + credential_env_names[field_name] + for field_name, value in credentials.items() + if not value + ] + if missing_env_names: + raise ValueError( + "feature_store_config requires non-empty environment variables: " + + ", ".join(missing_env_names) + + "; set them in the training process environment" + ) + access_key_id = credentials["access_key_id"] + access_key_secret = credentials["access_key_secret"] + featuredb_username = credentials["featuredb_username"] + featuredb_password = credentials["featuredb_password"] if not region: raise ValueError( @@ -171,21 +166,6 @@ def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": raise ValueError( "feature_store_config.endpoint must not contain URI userinfo" ) - _validate_pair( - access_key_id, - access_key_secret, - "access_key_id", - "access_key_secret", - required=True, - ) - _validate_pair( - featuredb_username, - featuredb_password, - "featuredb_username", - "featuredb_password", - required=True, - ) - project_name = config.project_name.strip() feature_entity_name = config.feature_entity_name.strip() feature_view_name = config.feature_view_name.strip() diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index c239bb01b..62c919211 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -54,11 +54,7 @@ def _schema_with_generation(generation: str = _TEST_DUMP_GENERATION) -> pa.Schem def _feature_store_config(**overrides) -> FeatureStoreConfig: config = FeatureStoreConfig( region="cn-test", - access_key_id="ak-id-secret", - access_key_secret="ak-value-secret", security_token="sts-secret", - featuredb_username="fdb-user-secret", - featuredb_password="fdb-password-secret", project_name="project_a", feature_entity_name="embedding_entity", feature_view_name="shared_embeddings", @@ -306,17 +302,27 @@ def __call__(self, **kwargs): class FeatureStoreDeltaUploaderTest(unittest.TestCase): + def setUp(self): + self._credential_env = mock.patch.dict( + os.environ, + { + "ALIBABA_CLOUD_ACCESS_KEY_ID": "ak-id-secret", + "ALIBABA_CLOUD_ACCESS_KEY_SECRET": "ak-value-secret", + "FEATUREDB_USERNAME": "fdb-user-secret", + "FEATUREDB_PASSWORD": "fdb-password-secret", + }, + clear=False, + ) + self._credential_env.start() + self.addCleanup(self._credential_env.stop) + def test_proto_groups_required_fields_before_optional_fields(self): required_fields = [ "region", - "access_key_id", - "access_key_secret", - "featuredb_username", - "featuredb_password", "project_name", "feature_view_name", - "version", "feature_entity_name", + "version", ] optional_fields = [ "endpoint", @@ -349,7 +355,7 @@ def test_proto_groups_required_fields_before_optional_fields(self): for field in fields[len(required_fields) :] ) ) - self.assertEqual([field.number for field in fields], list(range(1, 22))) + self.assertEqual([field.number for field in fields], [1, *list(range(6, 22))]) for field_name in required_fields: with self.subTest(field_name=field_name): config = _feature_store_config() @@ -382,18 +388,9 @@ def test_version_is_required_and_must_be_explicit(self): _feature_store_config(version="default") ) - def test_environment_fallback_and_credential_pairs(self): - config = _feature_store_config() - for name in ( - "region", - "access_key_id", - "access_key_secret", - "featuredb_username", - "featuredb_password", - ): - setattr(config, name, "") + def test_environment_resolution_and_required_credentials(self): + config = _feature_store_config(region="") config.ClearField("security_token") - self.assertTrue(config.IsInitialized()) environment = { "ALIBABA_CLOUD_REGION": "cn-env", "ALIBABA_CLOUD_ACCESS_KEY_ID": "env-ak", @@ -406,34 +403,23 @@ def test_environment_fallback_and_credential_pairs(self): settings = FeatureStoreUploadSettings.from_proto(config) self.assertEqual(settings.region, "cn-env") self.assertEqual(settings.access_key_id, "env-ak") - - config = _feature_store_config() - config.access_key_secret = "" - with mock.patch.dict( - os.environ, {"ALIBABA_CLOUD_ACCESS_KEY_SECRET": ""}, clear=False - ): - with self.assertRaisesRegex(ValueError, "configured together"): - FeatureStoreUploadSettings.from_proto(config) - - config = _feature_store_config() - config.featuredb_username = "" - config.featuredb_password = "" - with mock.patch.dict( - os.environ, - {"FEATUREDB_USERNAME": "", "FEATUREDB_PASSWORD": ""}, - clear=False, - ): - with self.assertRaisesRegex(ValueError, "are required"): - FeatureStoreUploadSettings.from_proto(config) - - config.featuredb_username = "only-one-side" - with mock.patch.dict( - os.environ, - {"FEATUREDB_USERNAME": "", "FEATUREDB_PASSWORD": ""}, - clear=False, + self.assertEqual(settings.access_key_secret, "env-sk") + self.assertEqual(settings.security_token, "env-sts") + self.assertEqual(settings.featuredb_username, "env-user") + self.assertEqual(settings.featuredb_password, "env-password") + + for missing_env_name in ( + "ALIBABA_CLOUD_ACCESS_KEY_ID", + "ALIBABA_CLOUD_ACCESS_KEY_SECRET", + "FEATUREDB_USERNAME", + "FEATUREDB_PASSWORD", ): - with self.assertRaisesRegex(ValueError, "configured together"): - FeatureStoreUploadSettings.from_proto(config) + with self.subTest(missing_env_name=missing_env_name): + incomplete_environment = dict(environment) + incomplete_environment.pop(missing_env_name) + with mock.patch.dict(os.environ, incomplete_environment, clear=True): + with self.assertRaisesRegex(ValueError, missing_env_name): + FeatureStoreUploadSettings.from_proto(config) with self.assertRaisesRegex(ValueError, "URI userinfo"): FeatureStoreUploadSettings.from_proto( From 035fc1a582a0398333a8ffb363dc070a8ebc9f9c Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 17 Jul 2026 14:35:52 +0800 Subject: [PATCH 06/13] remove redundant delta dump columns --- .../check_feature_store_delta.py | 11 +++------ .../check_feature_store_delta_test.py | 1 - tzrec/utils/delta_embedding_dump.py | 7 ------ tzrec/utils/delta_embedding_dump_test.py | 23 +++++++++++-------- tzrec/utils/feature_store_delta_uploader.py | 4 ---- .../feature_store_delta_uploader_test.py | 2 -- 6 files changed, 17 insertions(+), 31 deletions(-) diff --git a/tzrec/tools/feature_store/check_feature_store_delta.py b/tzrec/tools/feature_store/check_feature_store_delta.py index b2d9e08a0..8942629d7 100644 --- a/tzrec/tools/feature_store/check_feature_store_delta.py +++ b/tzrec/tools/feature_store/check_feature_store_delta.py @@ -43,7 +43,6 @@ from tzrec.utils import config_util from tzrec.utils.feature_store_delta_uploader import ( - DELTA_OPERATION_UPSERT, FEATURE_STORE_PK_FIELD, FEATURE_STORE_SK_FIELD, FEATURE_STORE_VALUE_FIELD, @@ -223,7 +222,7 @@ def sample_local_records( sample_count: int, embedding_name: Optional[str] = None, ) -> List[LocalSample]: - """Read a bounded set of unique UPSERT records from canonical parquet shards.""" + """Read a bounded set of unique records from canonical parquet shards.""" if sample_count <= 0: raise ValueError("sample_count must be > 0") @@ -231,7 +230,6 @@ def sample_local_records( FEATURE_STORE_PK_FIELD, FEATURE_STORE_SK_FIELD, FEATURE_STORE_VALUE_FIELD, - "operation", ] samples: List[LocalSample] = [] seen: set[Tuple[str, int]] = set() @@ -246,14 +244,11 @@ def sample_local_records( ) for batch in parquet_file.iter_batches(batch_size=1024, columns=columns): values = batch.to_pydict() - for name, key_id, vector, operation in zip( + for name, key_id, vector in zip( values[FEATURE_STORE_PK_FIELD], values[FEATURE_STORE_SK_FIELD], values[FEATURE_STORE_VALUE_FIELD], - values["operation"], ): - if operation != DELTA_OPERATION_UPSERT: - continue name = str(name) if embedding_name is not None and name != embedding_name: continue @@ -277,7 +272,7 @@ def sample_local_records( return samples if not samples: suffix = f" for embedding_name={embedding_name!r}" if embedding_name else "" - raise ValueError(f"no sampleable UPSERT delta records were found{suffix}") + raise ValueError(f"no sampleable delta records were found{suffix}") return samples diff --git a/tzrec/tools/feature_store/check_feature_store_delta_test.py b/tzrec/tools/feature_store/check_feature_store_delta_test.py index 97c1e6bf3..58405ad04 100644 --- a/tzrec/tools/feature_store/check_feature_store_delta_test.py +++ b/tzrec/tools/feature_store/check_feature_store_delta_test.py @@ -109,7 +109,6 @@ def test_committed_parquet_paths_and_sampling(self): [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], type=pa.list_(pa.float32()), ), - "operation": ["UPSERT", "UPSERT", "UPSERT"], } ), path, diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index a69c869e4..1c14679dc 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -34,7 +34,6 @@ from tzrec.utils.feature_store_delta_uploader import ( DELTA_DUMP_GENERATION_METADATA_KEY, DELTA_DUMP_SCHEMA_VERSION, - DELTA_OPERATION_UPSERT, FeatureStoreDeltaUploader, FeatureStoreUploadError, _durable_makedirs, @@ -67,8 +66,6 @@ ("table_fqn", pa.string()), ("key_id", pa.int64()), ("embedding", pa.list_(pa.float32())), - ("operation", pa.string()), - ("source", pa.string()), ], metadata={ b"tzrec.delta_embedding.schema_version": DELTA_DUMP_SCHEMA_VERSION.encode( @@ -982,7 +979,6 @@ def _append_model_delta_rows( table_fqn=fqn, key_ids=key_ids, embeddings=embeddings, - source="model_delta_tracker", ) return num_rows @@ -1175,7 +1171,6 @@ def _append_table_chunk( table_fqn: str, key_ids: torch.Tensor, embeddings: torch.Tensor, - source: str, ) -> int: key_ids_cpu = key_ids.detach().cpu().to(torch.int64).contiguous() embeddings_cpu = embeddings.detach().cpu().to(torch.float32).contiguous() @@ -1222,8 +1217,6 @@ def _append_table_chunk( pa.repeat(pa.scalar(table_fqn, pa.string()), num_rows), pa.array(key_ids_cpu.numpy(), type=pa.int64()), self._embedding_array(embeddings_cpu), - pa.repeat(pa.scalar(DELTA_OPERATION_UPSERT, pa.string()), num_rows), - pa.repeat(pa.scalar(source, pa.string()), num_rows), ], schema=_DELTA_DUMP_SCHEMA, ) diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index a66d4f2c6..f1ef779c1 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -166,8 +166,6 @@ def _assert_sharded_dump_file(rank: int, output_path: str, dumper) -> None: set(table["embedding_name"].to_pylist()), {_SHARDED_TABLE_NAME} ) testcase.assertEqual(set(table["embedding_role"].to_pylist()), {"ebc"}) - testcase.assertEqual(set(table["operation"].to_pylist()), {"UPSERT"}) - testcase.assertEqual(set(table["source"].to_pylist()), {"model_delta_tracker"}) table_weight = dumper._collect_table_weights()[_SHARDED_TABLE_NAME] expected_key_ids = [ @@ -363,7 +361,6 @@ def test_dump_rows_include_rank_metadata(self): table_fqn="model.ebc.user_emb", key_ids=torch.tensor([-42]), embeddings=torch.tensor([[1.0, 2.0]]), - source="model_delta_tracker", ) self.assertEqual(num_rows, 1) self.assertEqual(len(table_chunks), 1) @@ -375,7 +372,20 @@ def test_dump_rows_include_rank_metadata(self): self.assertEqual(table["embedding_role"].to_pylist(), ["ebc"]) self.assertEqual(table["key_id"].to_pylist(), [-42]) self.assertEqual(table["embedding"].to_pylist(), [[1.0, 2.0]]) - self.assertEqual(table["operation"].to_pylist(), ["UPSERT"]) + self.assertEqual( + table.column_names, + [ + "global_step", + "rank", + "world_size", + "embedding_name", + "embedding_role", + "feature_name", + "table_fqn", + "key_id", + "embedding", + ], + ) def test_dump_rows_reject_processor_invalid_key_sentinel(self): dumper = object.__new__(DeltaEmbeddingDumper) @@ -392,7 +402,6 @@ def test_dump_rows_reject_processor_invalid_key_sentinel(self): table_fqn="model.ebc.user_emb", key_ids=torch.tensor([-1]), embeddings=torch.tensor([[1.0, 2.0]]), - source="model_delta_tracker", ) def test_write_table_chunks_preserves_parquet_schema(self): @@ -410,7 +419,6 @@ def test_write_table_chunks_preserves_parquet_schema(self): table_fqn="model.ebc.user_emb", key_ids=torch.tensor([7, 8]), embeddings=torch.tensor([[1.0, 2.0], [3.0, 4.0]]), - source="model_delta_tracker", ) with tempfile.TemporaryDirectory() as tmp_dir: output_path = os.path.join(tmp_dir, "delta.parquet") @@ -1244,9 +1252,6 @@ def test_dynamicemb_multi_gpu_delta_dump_writes_uniform_shards(self): continue dumped_real_rows = True self.assertEqual(set(table["world_size"].to_pylist()), {world_size}) - self.assertEqual( - set(table["source"].to_pylist()), {"model_delta_tracker"} - ) # dynamic lookup must return a real embedding vector per id. self.assertTrue( all(len(emb) > 0 for emb in table["embedding"].to_pylist()) diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index 54d6a9d6f..b0c00a1c3 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -72,8 +72,6 @@ "table_fqn": pa.string(), "key_id": pa.int64(), "embedding": pa.list_(pa.float32()), - "operation": pa.string(), - "source": pa.string(), } @@ -1784,8 +1782,6 @@ def _load_records( raise ValueError("delta shard rank mismatch") if int(row["world_size"]) != self._world_size: raise ValueError("delta shard world_size mismatch") - if row["operation"] != DELTA_OPERATION_UPSERT: - raise ValueError("only UPSERT delta operations are supported") if row["embedding_role"] not in SPARSE_EMBEDDING_ROLES: raise ValueError("delta shard has an invalid embedding role") diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index 62c919211..b931a4cb9 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -90,8 +90,6 @@ def _row( "table_fqn": f"model.ebc.embedding_bags.{name}.weight", "key_id": key_id, "embedding": values, - "operation": "UPSERT", - "source": "model_delta_tracker", } From 2faabccef435eb475331e03ad1a31f1dc5a7df33 Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 17 Jul 2026 14:45:45 +0800 Subject: [PATCH 07/13] export zeros for missing dynamic embedding keys --- tzrec/utils/delta_embedding_dump.py | 7 ++- tzrec/utils/delta_embedding_dump_test.py | 48 ++++++++++++++----- .../feature_store_delta_uploader_test.py | 3 +- 3 files changed, 44 insertions(+), 14 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 1c14679dc..ffacc2a8c 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -1049,13 +1049,16 @@ def _lookup_dynamic_embeddings( ) emb_dim = dynamic_module._dynamicemb_options[table_id].dim founds = founds.to(dtype=torch.bool) + embeddings = values[:, :emb_dim].detach() if not bool(founds.all().item()): + embeddings = embeddings.clone() + embeddings[~founds] = 0 logger.warning( - "Skip %s missing dynamic embedding ids for table %s.", + "Use zero embeddings for %s missing dynamic embedding ids in table %s.", int((~founds).sum().item()), table_name, ) - return values[founds, :emb_dim].detach(), ids[founds] + return embeddings, ids def _collect_table_shard_infos(self) -> Dict[str, _TableShardInfo]: table_shard_infos: Dict[str, _TableShardInfo] = {} diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index f1ef779c1..8770b6aa7 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -90,16 +90,17 @@ def forward(self, features: KeyedJaggedTensor) -> torch.Tensor: class _FakeDynamicTables: - def __init__(self) -> None: + def __init__(self, founds=None) -> None: self.ids = None self.table_ids = None self.copy_mode = None + self._founds = founds if founds is not None else [True, False, True] def find(self, ids, table_ids, copy_mode): self.ids = ids.detach().clone() self.table_ids = table_ids.detach().clone() self.copy_mode = copy_mode - founds = torch.tensor([True, False, True], device=ids.device) + founds = torch.tensor(self._founds, device=ids.device) values = torch.tensor( [ [1.0, 2.0, 20.0], @@ -1083,12 +1084,9 @@ def test_row_wise_lookup_requires_shard_metadata(self): ) @unittest.skipUnless(has_dynamicemb, "dynamicemb is not installed; skipping.") - @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for dynamicemb.") - @mark_ci_scope("gpu") - def test_lookup_dynamic_embeddings_filters_missing_ids(self): + def test_lookup_dynamic_embeddings_zero_fills_missing_ids(self): from dynamicemb.types import CopyMode - torch.cuda.set_device(0) dumper = object.__new__(DeltaEmbeddingDumper) fake_tables = _FakeDynamicTables() dynamic_module = SimpleNamespace( @@ -1098,19 +1096,47 @@ def test_lookup_dynamic_embeddings_filters_missing_ids(self): _dynamicemb_options=[SimpleNamespace(dim=2)], ) - embeddings, key_ids = dumper._lookup_dynamic_embeddings( - dynamic_module, "dyn_table", torch.tensor([101, 102, 103]) - ) + cpu_device = torch.device("cpu") + with ( + mock.patch.object(torch.cuda, "current_device", return_value=0), + mock.patch.object(torch, "device", return_value=cpu_device), + ): + embeddings, key_ids = dumper._lookup_dynamic_embeddings( + dynamic_module, "dyn_table", torch.tensor([101, 102, 103]) + ) dynamic_module.flush.assert_called_once_with() self.assertIs(fake_tables.copy_mode, CopyMode.EMBEDDING) torch.testing.assert_close(fake_tables.ids.cpu(), torch.tensor([101, 102, 103])) torch.testing.assert_close(fake_tables.table_ids.cpu(), torch.tensor([0, 0, 0])) - torch.testing.assert_close(key_ids.cpu(), torch.tensor([101, 103])) + torch.testing.assert_close(key_ids.cpu(), torch.tensor([101, 102, 103])) torch.testing.assert_close( - embeddings.cpu(), torch.tensor([[1.0, 2.0], [5.0, 6.0]]) + embeddings.cpu(), + torch.tensor([[1.0, 2.0], [0.0, 0.0], [5.0, 6.0]]), ) + @unittest.skipUnless(has_dynamicemb, "dynamicemb is not installed; skipping.") + def test_lookup_dynamic_embeddings_zero_fills_all_missing_ids(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dynamic_module = SimpleNamespace( + table_names=["dyn_table"], + tables=_FakeDynamicTables(founds=[False, False, False]), + flush=mock.MagicMock(), + _dynamicemb_options=[SimpleNamespace(dim=2)], + ) + + cpu_device = torch.device("cpu") + with ( + mock.patch.object(torch.cuda, "current_device", return_value=0), + mock.patch.object(torch, "device", return_value=cpu_device), + ): + embeddings, key_ids = dumper._lookup_dynamic_embeddings( + dynamic_module, "dyn_table", torch.tensor([101, 102, 103]) + ) + + torch.testing.assert_close(key_ids.cpu(), torch.tensor([101, 102, 103])) + torch.testing.assert_close(embeddings.cpu(), torch.zeros(3, 2)) + @unittest.skipUnless(has_dynamicemb, "dynamicemb is not installed; skipping.") @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for dynamicemb.") @mark_ci_scope("gpu") diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index b931a4cb9..c30729aac 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -747,7 +747,7 @@ def test_complete_step_uploads_merge_with_stable_version_and_ts(self): [ _row(10, 0, 1, [1.0, 2.0]), _row(10, 0, 2, [3.0, 4.0]), - _row(10, 0, 3, [5.0, 6.0]), + _row(10, 0, 3, [0.0, 0.0]), ], ) view = _FakeView() @@ -773,6 +773,7 @@ def test_complete_step_uploads_merge_with_stable_version_and_ts(self): ) self.assertEqual({call["write_mode"] for call in view.calls}, {"MERGE"}) self.assertEqual([call["ts"] for call in view.calls], [123456, 123457]) + self.assertEqual(view.calls[1]["data"][0]["embedding"].tolist(), [0.0, 0.0]) self.assertEqual(view.closed, [True]) success_path = os.path.join(uploader.state_dir, "step_10._FS_SUCCESS.json") From 14ba85d59325b459515b4556a214e1ac353bb037 Mon Sep 17 00:00:00 2001 From: gecheng Date: Sun, 19 Jul 2026 11:04:42 +0800 Subject: [PATCH 08/13] [bugfix] replay restored FeatureStore deltas by generation --- .../check_feature_store_delta.py | 14 +- .../check_feature_store_delta_test.py | 34 +++ tzrec/utils/delta_embedding_dump.py | 4 +- tzrec/utils/delta_embedding_dump_test.py | 27 ++ tzrec/utils/feature_store_delta_uploader.py | 260 ++++++++++++++---- .../feature_store_delta_uploader_test.py | 172 +++++++++++- 6 files changed, 444 insertions(+), 67 deletions(-) diff --git a/tzrec/tools/feature_store/check_feature_store_delta.py b/tzrec/tools/feature_store/check_feature_store_delta.py index 8942629d7..a2ad744d4 100644 --- a/tzrec/tools/feature_store/check_feature_store_delta.py +++ b/tzrec/tools/feature_store/check_feature_store_delta.py @@ -121,7 +121,7 @@ def load_committed_upload( settings: FeatureStoreUploadSettings, global_step: Optional[int] = None, ) -> Tuple[str, Dict[str, Any], Dict[str, Any]]: - """Load the local committed watermark and its success marker. + """Load the latest local commit or a requested committed step. Args: output_dir: Delta parquet outbox directory. @@ -129,7 +129,7 @@ def load_committed_upload( global_step: Specific committed step to inspect, or None for the latest. Returns: - Tuple of state directory, committed watermark, and success marker. + Tuple of state directory, latest commit, and selected success marker. Raises: FileNotFoundError: If no committed marker exists for the configured target. @@ -160,12 +160,6 @@ def load_committed_upload( selected_step = committed_step if global_step is None else int(global_step) if selected_step <= 0: raise ValueError("global_step must be > 0") - if selected_step > committed_step: - raise ValueError( - f"step {selected_step} is newer than the locally committed " - f"FeatureStore watermark {committed_step}" - ) - success_path = os.path.join(state_dir, f"step_{selected_step}._FS_SUCCESS.json") if not os.path.isfile(success_path): raise FileNotFoundError( @@ -477,7 +471,7 @@ def run_check(args: argparse.Namespace) -> int: }, "local_commit": { "global_step": global_step, - "committed_watermark": int(committed["committed_global_step"]), + "latest_committed_step": int(committed["committed_global_step"]), "success_records": int(success["success_records"]), "total_records": int(success["total_records"]), "parquet_paths": [ @@ -522,7 +516,7 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: "--global_step", type=int, default=None, - help="Committed step to inspect; defaults to the latest committed watermark.", + help="Committed step to inspect; defaults to the latest publication.", ) parser.add_argument( "--sample_count", diff --git a/tzrec/tools/feature_store/check_feature_store_delta_test.py b/tzrec/tools/feature_store/check_feature_store_delta_test.py index 58405ad04..38dff3999 100644 --- a/tzrec/tools/feature_store/check_feature_store_delta_test.py +++ b/tzrec/tools/feature_store/check_feature_store_delta_test.py @@ -97,6 +97,40 @@ def test_load_committed_upload_checks_target_and_success(self): self.assertEqual(committed["committed_global_step"], 20) self.assertEqual(success["global_step"], 20) + def test_load_committed_upload_accepts_step_before_latest_replay(self): + settings = SimpleNamespace( + project_name="project", feature_view_name="view", version="v1" + ) + with tempfile.TemporaryDirectory() as output_dir: + state_dir = os.path.join(output_dir, ".feature_store_upload", "target") + os.makedirs(state_dir) + target = { + "project_name": "project", + "feature_view_name": "view", + "version": "v1", + } + with open(os.path.join(state_dir, "committed.json"), "w") as output: + json.dump({**target, "committed_global_step": 15}, output) + with open( + os.path.join(state_dir, "step_20._FS_SUCCESS.json"), "w" + ) as output: + json.dump( + { + **target, + "global_step": 20, + "success_records": 1, + "total_records": 1, + }, + output, + ) + + _, committed, success = load_committed_upload( + output_dir, settings, global_step=20 + ) + + self.assertEqual(committed["committed_global_step"], 15) + self.assertEqual(success["global_step"], 20) + def test_committed_parquet_paths_and_sampling(self): with tempfile.TemporaryDirectory() as output_dir: path = os.path.join(output_dir, "delta__fs_target_step_20.parquet") diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index ffacc2a8c..b97ae50b1 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -874,6 +874,8 @@ def dump(self, global_step: int, check_upload_error: bool = True) -> Optional[st self._write_table_chunks( table_chunks, output_path, dump_generation=dump_generation ) + if uploader is not None: + uploader.submit(global_step, dump_generation) except BaseException: # The durability guard still owns the rows, so rewinding this # consumer makes a caller retry observe the same snapshot. @@ -881,8 +883,6 @@ def dump(self, global_step: int, check_upload_error: bool = True) -> Optional[st raise self._ack_durable_tracker_read() - if uploader is not None: - uploader.submit(global_step) if num_rows == 0: logger.info( "Dumped empty delta embedding shard to %s at step %s.", diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 8770b6aa7..56b722403 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -816,6 +816,33 @@ def test_durable_ack_advances_guard_without_recomputing_unique_rows(self): tracker.store.delete.assert_called_once_with(up_to_idx=12) tracker.get_unique.assert_not_called() + def test_dump_does_not_ack_tracker_when_submission_is_rejected(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._feature_store_enabled = True + dumper._world_size = 1 + dumper._tracker = mock.MagicMock() + dumper._tracker.per_consumer_batch_idx = {_CONSUMER: 7} + dumper._uploader = mock.MagicMock() + dumper._uploader.submit.side_effect = RuntimeError("submission rejected") + + with ( + mock.patch.object(dumper, "_check_feature_store_upload_error"), + mock.patch.object(dumper, "_next_dump_generation", return_value="run-b"), + mock.patch.object(dumper, "_collect_table_weights", return_value={}), + mock.patch.object(dumper, "_collect_dynamic_modules", return_value={}), + mock.patch.object(dumper, "_append_model_delta_rows", return_value=1), + mock.patch.object(dumper, "_output_path", return_value="delta.parquet"), + mock.patch.object(dumper, "_write_table_chunks"), + mock.patch.object(dumper, "_rollback_tracker_read") as rollback, + mock.patch.object(dumper, "_ack_durable_tracker_read") as durable_ack, + ): + with self.assertRaisesRegex(RuntimeError, "submission rejected"): + dumper.dump(10) + + dumper._uploader.submit.assert_called_once_with(10, "run-b") + rollback.assert_called_once_with(7) + durable_ack.assert_not_called() + def test_multi_gpu_output_path_uses_step_underscore_dir(self): with tempfile.TemporaryDirectory() as tmp_dir: dumper = object.__new__(DeltaEmbeddingDumper) diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index b0c00a1c3..3dc360de3 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -372,7 +372,8 @@ class FeatureStoreDeltaUploader: The training thread only writes/parquet-renames and enqueues a step. This object's single worker waits for the exact rank shard set and persists a monotonic timestamp range before each full-step MERGE attempt. It never - invokes a torch.distributed collective. + invokes a torch.distributed collective. Restored training may replay a + repeated or lower step; dump generation identifies the new publication. """ def __init__( @@ -522,11 +523,13 @@ def start(self) -> None: self._settings.version, ) - def submit(self, global_step: int) -> None: + def submit(self, global_step: int, dump_generation: Optional[str] = None) -> None: """Enqueue a durably written rank-zero shard with bounded back-pressure.""" global_step = int(global_step) if global_step <= 0: raise ValueError("FeatureStore delta global_step must be > 0") + if dump_generation is not None and not dump_generation: + raise ValueError("FeatureStore delta dump_generation must not be empty") with self._condition: self._raise_if_failed_locked() if not self._started: @@ -535,13 +538,8 @@ def submit(self, global_step: int) -> None: ) if self._closing or self._closed: raise RuntimeError("cannot submit to a closing FeatureStore uploader") - if self._is_committed(global_step): + if self._is_committed(global_step, dump_generation): return - if global_step <= self._committed_global_step: - raise FeatureStoreUploadError( - "refusing an uncommitted delta step older than the local " - "FeatureStore committed watermark" - ) while ( global_step not in self._pending and len(self._pending) >= self._settings.max_pending_steps @@ -609,19 +607,50 @@ def _run(self) -> None: manifest_exists = os.path.isfile(self._manifest_path(current_step)) snapshot_paths = self._snapshot_paths(current_step) snapshot_dir_exists = os.path.isdir(self._snapshot_dir(current_step)) - if manifest_exists or snapshot_dir_exists: - if not all(os.path.isfile(path) for path in snapshot_paths): - raise FeatureStoreUploadError( - "an uncommitted FeatureStore upload journal is missing " - "its durable shard snapshot; recovery cannot continue" - ) - else: + if snapshot_dir_exists and self._is_committed(current_step): try: - snapshot_paths = self._snapshot_canonical_shards( - current_step, canonical_paths + canonical_generation = self._canonical_dump_generation( + current_step ) except _ShardSetNotReady: + canonical_generation = None snapshot_paths = None + manifest_generation = _read_json( + self._manifest_path(current_step) + ).get("dump_generation") + if ( + canonical_generation is not None + and canonical_generation != manifest_generation + ): + self._reclaim_snapshot(current_step) + snapshot_dir_exists = os.path.isdir( + self._snapshot_dir(current_step) + ) + if snapshot_dir_exists: + raise FeatureStoreUploadError( + "cannot replace a committed FeatureStore shard " + "snapshot with a newer dump generation" + ) + if snapshot_paths is not None: + if snapshot_dir_exists: + if not all(os.path.isfile(path) for path in snapshot_paths): + raise FeatureStoreUploadError( + "an uncommitted FeatureStore upload journal is " + "missing its durable shard snapshot; recovery " + "cannot continue" + ) + elif manifest_exists and not self._is_committed(current_step): + raise FeatureStoreUploadError( + "an uncommitted FeatureStore upload journal is missing " + "its durable shard snapshot; recovery cannot continue" + ) + else: + try: + snapshot_paths = self._snapshot_canonical_shards( + current_step, canonical_paths + ) + except _ShardSetNotReady: + snapshot_paths = None if snapshot_paths is None: elapsed = time.monotonic() - pending_since @@ -834,12 +863,17 @@ def _add_discovered_steps_locked(self) -> None: "global_step must be > 0" ) if self._is_committed(step): - continue - if step <= self._committed_global_step: - raise ValueError( - "found an uncommitted delta step older than the local " - "FeatureStore committed watermark" + try: + canonical_generation = self._canonical_dump_generation(step) + except _ShardSetNotReady: + canonical_generation = "" + if canonical_generation is None: + continue + manifest_generation = _read_json(self._manifest_path(step)).get( + "dump_generation" ) + if canonical_generation == manifest_generation: + continue if step not in self._pending: self._pending[step] = time.monotonic() remaining_capacity -= 1 @@ -913,18 +947,52 @@ def _expected_shard_paths(self, global_step: int) -> List[str]: for rank in range(self._world_size) ] + def _canonical_dump_generation(self, global_step: int) -> Optional[str]: + """Read one complete canonical shard set's generation without hashing it.""" + paths = self._expected_shard_paths(global_step) + existing = [os.path.isfile(path) for path in paths] + if not any(existing): + return None + if not all(existing): + raise _ShardSetNotReady("canonical delta shard set is incomplete") + + generations = set() + for path in paths: + metadata = pq.read_schema(path).metadata or {} + if metadata.get( + _SCHEMA_VERSION_METADATA_KEY + ) != DELTA_DUMP_SCHEMA_VERSION.encode("ascii"): + raise ValueError(f"unsupported delta dump schema version in {path}") + generation = metadata.get(DELTA_DUMP_GENERATION_METADATA_KEY) + if not generation: + raise ValueError(f"delta dump generation is missing in {path}") + try: + generations.add(generation.decode("ascii")) + except UnicodeDecodeError as exc: + raise ValueError( + f"delta dump generation is not ASCII in {path}" + ) from exc + if len(generations) != 1: + raise _ShardSetNotReady( + "delta shards from different dump generations are present" + ) + return next(iter(generations)) + def _manifest_path(self, global_step: int) -> str: return os.path.join(self._state_dir, f"step_{global_step}.manifest.json") def _success_path(self, global_step: int) -> str: return os.path.join(self._state_dir, f"step_{global_step}._FS_SUCCESS.json") - def _is_committed(self, global_step: int) -> bool: + def _load_success_state( + self, global_step: int + ) -> Optional[Tuple[Dict[str, Any], Dict[str, Any], bool]]: + """Validate one success marker and report whether its manifest is active.""" if global_step <= 0: raise ValueError("FeatureStore delta global_step must be > 0") path = self._success_path(global_step) if not os.path.isfile(path): - return False + return None success = _read_json(path) expected = { "global_step": global_step, @@ -941,12 +1009,49 @@ def _is_committed(self, global_step: int) -> bool: if not os.path.isfile(manifest_path): raise ValueError("FeatureStore success marker is missing its manifest") manifest = _read_json(manifest_path) - if success.get("manifest_digest") != _json_digest(manifest): - raise ValueError("FeatureStore success marker manifest digest mismatch") shards = success.get("shards") - if not isinstance(shards, list) or shards != manifest.get("shards"): + if not isinstance(shards, list): raise ValueError("FeatureStore success marker has invalid shards") - return True + success_manifest_digest = success.get("manifest_digest") + if not isinstance(success_manifest_digest, str) or not success_manifest_digest: + raise ValueError("FeatureStore success marker has invalid manifest_digest") + success_dump_generation = self._validate_dump_generation(shards) + if success.get("dump_generation", success_dump_generation) != ( + success_dump_generation + ): + raise ValueError("FeatureStore success marker has invalid dump_generation") + if success_manifest_digest == _json_digest(manifest): + if ( + shards != manifest.get("shards") + or manifest.get("dump_generation") != success_dump_generation + ): + raise ValueError("FeatureStore success marker has invalid shards") + return success, manifest, True + + if ( + manifest.get("supersedes_manifest_digest") != success_manifest_digest + or manifest.get("supersedes_publish_ts") != publish_ts + or manifest.get("supersedes_dump_generation") != success_dump_generation + ): + raise ValueError("FeatureStore success marker manifest digest mismatch") + if manifest.get("dump_generation") == manifest.get( + "supersedes_dump_generation" + ): + raise ValueError("FeatureStore superseding manifest generation is invalid") + return success, manifest, False + + def _is_committed( + self, global_step: int, dump_generation: Optional[str] = None + ) -> bool: + state = self._load_success_state(global_step) + if state is None: + return False + _, manifest, active = state + if not active: + return False + return dump_generation is None or ( + manifest.get("dump_generation") == dump_generation + ) def _reconcile_committed_state(self) -> None: """Repair a crash between atomic success and committed-state writes.""" @@ -957,22 +1062,23 @@ def _reconcile_committed_state(self) -> None: if match is None: continue step = int(match.group(1)) - if not self._is_committed(step): + state = self._load_success_state(step) + if state is None: continue - success = _read_json(self._success_path(step)) - if latest_success is None or step > int(latest_success["global_step"]): + success, _, _ = state + if latest_success is None or int(success["publish_ts"]) > int( + latest_success["publish_ts"] + ): latest_success = success if latest_success is None: return committed_path = os.path.join(self._state_dir, "committed.json") - committed_step = -1 + committed_publish_ts = 0 if os.path.isfile(committed_path): - committed_step = int( - _read_json(committed_path).get("committed_global_step", -1) - ) + committed_publish_ts = int(_read_json(committed_path).get("publish_ts", 0)) latest_step = int(latest_success["global_step"]) - if committed_step >= latest_step: + if committed_publish_ts >= int(latest_success["publish_ts"]): return committed = { "schema_version": 1, @@ -982,6 +1088,9 @@ def _reconcile_committed_state(self) -> None: "committed_global_step": latest_step, "publish_ts": int(latest_success["publish_ts"]), "contract_hash": self._contract_hash, + "dump_generation": latest_success.get("dump_generation") + or self._validate_dump_generation(latest_success["shards"]), + "manifest_digest": latest_success["manifest_digest"], } _atomic_write_json(committed_path, committed) @@ -995,6 +1104,7 @@ def _load_latest_publish_ts(self) -> int: if not name.endswith(".manifest.json"): continue manifest = _read_json(os.path.join(self._state_dir, name)) + latest = max(latest, int(manifest.get("supersedes_publish_ts", 0))) for attempt in manifest.get("attempts", []): latest = max(latest, int(attempt.get("range_end", 0))) return latest @@ -1212,28 +1322,68 @@ def _load_or_create_manifest( "dump_generation": dump_generation, "shards": expected_shards, } - for name, value in expected.items(): - if manifest.get(name) != value: + mismatch = next( + ( + name + for name, value in expected.items() + if manifest.get(name) != value + ), + None, + ) + if mismatch is None: + attempts = manifest.get("attempts") + if not isinstance(attempts, list): raise ValueError( - f"FeatureStore upload manifest mismatch for {name}" + "FeatureStore upload manifest has invalid attempts" ) - attempts = manifest.get("attempts") - if not isinstance(attempts, list): - raise ValueError("FeatureStore upload manifest has invalid attempts") - for attempt in attempts: - if ( - not isinstance(attempt, dict) - or type(attempt.get("range_start")) is not int - or type(attempt.get("range_end")) is not int - or attempt["range_start"] <= 0 - or attempt["range_end"] < attempt["range_start"] - ): - raise ValueError( - "FeatureStore upload manifest has an invalid ts range" + for attempt in attempts: + if ( + not isinstance(attempt, dict) + or type(attempt.get("range_start")) is not int + or type(attempt.get("range_end")) is not int + or attempt["range_start"] <= 0 + or attempt["range_end"] < attempt["range_start"] + ): + raise ValueError( + "FeatureStore upload manifest has an invalid ts range" + ) + self._last_publish_ts = max( + self._last_publish_ts, int(attempt["range_end"]) ) - self._last_publish_ts = max( - self._last_publish_ts, int(attempt["range_end"]) + return manifest + + if manifest.get( + "dump_generation" + ) == dump_generation or not self._is_committed(global_step): + raise ValueError( + f"FeatureStore upload manifest mismatch for {mismatch}" ) + + success_state = self._load_success_state(global_step) + if success_state is None or not success_state[2]: + raise ValueError( + "FeatureStore committed manifest cannot be superseded safely" + ) + # Keep the old success verifiable across a crash during replacement. + success, previous_manifest, _ = success_state + manifest = { + "schema_version": 3, + "global_step": global_step, + "world_size": self._world_size, + "project_name": self._settings.project_name, + "feature_view_name": self._settings.feature_view_name, + "version": self._settings.version, + "write_mode": FEATURE_STORE_WRITE_MODE, + "contract_hash": self._contract_hash, + "record_count": record_count, + "dump_generation": dump_generation, + "shards": expected_shards, + "attempts": [], + "supersedes_manifest_digest": success["manifest_digest"], + "supersedes_publish_ts": int(success["publish_ts"]), + "supersedes_dump_generation": previous_manifest["dump_generation"], + } + _atomic_write_json(path, manifest) return manifest manifest = { @@ -1860,6 +2010,7 @@ def _commit_success( "publish_ts": publish_ts, "write_mode": FEATURE_STORE_WRITE_MODE, "contract_hash": self._contract_hash, + "dump_generation": manifest["dump_generation"], "manifest_digest": manifest_digest, "shards": manifest["shards"], "total_records": int(summary["total_records"]), @@ -1874,6 +2025,7 @@ def _commit_success( "committed_global_step": global_step, "publish_ts": publish_ts, "contract_hash": self._contract_hash, + "dump_generation": manifest["dump_generation"], "manifest_digest": manifest_digest, } _atomic_write_json(os.path.join(self._state_dir, "committed.json"), committed) diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index c30729aac..1a4aa25b1 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -29,6 +29,7 @@ FeatureStoreUploadError, FeatureStoreUploadSettings, _json_digest, + _read_json, feature_store_delta_file_prefix, ) @@ -901,7 +902,7 @@ def test_step_zero_success_marker_fails_before_reconcile(self): self.assertEqual(factory.calls, []) self.assertEqual(view.calls, []) - def test_step_zero_committed_watermark_is_rejected(self): + def test_step_zero_committed_state_is_rejected(self): with tempfile.TemporaryDirectory() as output_dir: view = _FakeView() factory = _FakeClientFactory(view) @@ -1621,6 +1622,175 @@ def test_restart_skips_successfully_committed_step(self): restarted.close() self.assertEqual(len(factory.calls), 1) + def test_restart_republishes_same_step_from_new_dump_generation(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard( + output_dir, + 10, + [_row(10, 0, 1, [1.0, 2.0])], + generation="run-a", + ) + first = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + clock_ms=lambda: 100, + ) + first.start() + first.close() + + _write_single_shard( + output_dir, + 10, + [_row(10, 0, 1, [7.0, 8.0])], + generation="run-b", + ) + replay_view = _FakeView() + replayed = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(replay_view), + clock_ms=lambda: 100, + ) + replayed.start() + replayed.close() + + self.assertEqual(len(replay_view.calls), 1) + self.assertEqual( + replay_view.calls[0]["data"][0]["embedding"].tolist(), [7.0, 8.0] + ) + manifest = _read_json(replayed._manifest_path(10)) + success = _read_json(replayed._success_path(10)) + self.assertEqual(manifest["dump_generation"], "run-b") + self.assertEqual(manifest["supersedes_dump_generation"], "run-a") + self.assertEqual(manifest["supersedes_publish_ts"], 100) + self.assertEqual(success["publish_ts"], 101) + self.assertEqual(success["manifest_digest"], _json_digest(manifest)) + + def test_restart_recovers_interrupted_same_step_supersession(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard( + output_dir, + 10, + [_row(10, 0, 1, [1.0, 2.0])], + generation="run-a", + ) + first = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + clock_ms=lambda: 100, + ) + first.start() + first.close() + + shard = _write_single_shard( + output_dir, + 10, + [_row(10, 0, 1, [7.0, 8.0])], + generation="run-b", + ) + interrupted = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + clock_ms=lambda: 100, + ) + with _initialized_journal(interrupted): + snapshot_paths = interrupted._snapshot_canonical_shards(10, [shard]) + self.assertIsNotNone(snapshot_paths) + manifest = interrupted._load_or_create_manifest(10, snapshot_paths, 1) + self.assertEqual(manifest["dump_generation"], "run-b") + self.assertFalse(interrupted._is_committed(10)) + + replay_view = _FakeView() + replayed = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(replay_view), + clock_ms=lambda: 100, + ) + replayed.start() + replayed.close() + + self.assertEqual(len(replay_view.calls), 1) + self.assertEqual(replay_view.calls[0]["ts"], 101) + self.assertTrue(replayed._is_committed(10, "run-b")) + + def test_restart_republishes_new_generation_below_previous_step(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard( + output_dir, + 20, + [_row(20, 0, 1, [1.0, 2.0])], + generation="run-a", + ) + first = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + clock_ms=lambda: 100, + ) + first.start() + first.close() + + _write_single_shard( + output_dir, + 15, + [_row(15, 0, 1, [7.0, 8.0])], + generation="run-b", + ) + replay_view = _FakeView() + replayed = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(replay_view), + clock_ms=lambda: 100, + ) + replayed.start() + replayed.close() + + self.assertEqual(len(replay_view.calls), 1) + self.assertEqual(replay_view.calls[0]["ts"], 101) + with open(os.path.join(replayed.state_dir, "committed.json")) as source: + committed = json.load(source) + self.assertEqual(committed["committed_global_step"], 15) + self.assertEqual(committed["publish_ts"], 101) + + verified = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + clock_ms=lambda: 100, + ) + verified.start() + verified.close() + self.assertEqual(verified._committed_global_step, 15) + if __name__ == "__main__": unittest.main() From 8eb37f2bebff6bbace2e9c3294fb76ffc22e26be Mon Sep 17 00:00:00 2001 From: gecheng Date: Sun, 19 Jul 2026 11:38:27 +0800 Subject: [PATCH 09/13] [bugfix] synchronize timed delta dump workers --- tzrec/main.py | 8 +++++++- tzrec/utils/delta_embedding_dump.py | 5 +++++ tzrec/utils/delta_embedding_dump_test.py | 21 +++++++++++++++++++++ 3 files changed, 33 insertions(+), 1 deletion(-) diff --git a/tzrec/main.py b/tzrec/main.py index 1886db5f9..e0b47e8c9 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -467,11 +467,17 @@ 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 + sync_train_data_exhaustion = check_all_workers_data_status or ( + delta_embedding_dumper is not None + and delta_embedding_dumper.requires_synced_dataloader_exhaustion + ) for i_epoch in epoch_iter: + # Timed dumps broadcast their per-step decision, so all ranks must stop + # consuming data together before any worker can leave that collective. 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, ) if plogger is not None: plogger.set_description(f"Training Epoch {i_epoch}") diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index b97ae50b1..3499a72ce 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -472,6 +472,11 @@ def clear(self) -> None: """Clear tracked sparse ids, usually after restore-time dummy steps.""" self._tracker.clear() + @property + def requires_synced_dataloader_exhaustion(self) -> bool: + """Return whether input exhaustion must stay aligned across ranks.""" + return self._interval_secs is not None and self._world_size > 1 + def start(self) -> None: """Start timed cadence and rank-zero publication after initialization.""" if self._feature_store_enabled: diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 56b722403..2c573eb8c 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -799,6 +799,27 @@ def test_minutes_interval_is_converted_to_seconds(self): self.assertIsNone(dumper._interval_steps) self.assertEqual(dumper._interval_secs, 120.0) + def test_multi_rank_minutes_requires_synced_dataloader_exhaustion(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_secs = 60.0 + dumper._world_size = 2 + + self.assertTrue(dumper.requires_synced_dataloader_exhaustion) + + def test_step_interval_does_not_require_synced_dataloader_exhaustion(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_secs = None + dumper._world_size = 2 + + self.assertFalse(dumper.requires_synced_dataloader_exhaustion) + + def test_single_rank_minutes_does_not_require_synced_dataloader_exhaustion(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_secs = 60.0 + dumper._world_size = 1 + + self.assertFalse(dumper.requires_synced_dataloader_exhaustion) + def test_durable_ack_advances_guard_without_recomputing_unique_rows(self): tracker = mock.MagicMock() tracker.per_consumer_batch_idx = { From 0183ec0281959bb05f9ce1c0b2a33d8d8e263907 Mon Sep 17 00:00:00 2001 From: gecheng Date: Sun, 19 Jul 2026 11:53:06 +0800 Subject: [PATCH 10/13] [bugfix] use VPC routing for FeatureStore uploads --- .../check_feature_store_delta.py | 2 +- .../check_feature_store_delta_test.py | 41 +++++++++++++++++++ tzrec/utils/feature_store_delta_uploader.py | 1 - .../feature_store_delta_uploader_test.py | 1 + 4 files changed, 43 insertions(+), 2 deletions(-) diff --git a/tzrec/tools/feature_store/check_feature_store_delta.py b/tzrec/tools/feature_store/check_feature_store_delta.py index a2ad744d4..4e5f53eec 100644 --- a/tzrec/tools/feature_store/check_feature_store_delta.py +++ b/tzrec/tools/feature_store/check_feature_store_delta.py @@ -384,7 +384,7 @@ def create_feature_store_view(settings: FeatureStoreUploadSettings) -> Any: except (TypeError, ValueError): parameters = {} if "test_mode" in parameters: - # Match the uploader's FeatureDB routing when supported by the SDK wheel. + # This manual readback tool normally runs outside the production VPC. kwargs["test_mode"] = True client = FeatureStoreClient(**kwargs) diff --git a/tzrec/tools/feature_store/check_feature_store_delta_test.py b/tzrec/tools/feature_store/check_feature_store_delta_test.py index 38dff3999..92fadc35b 100644 --- a/tzrec/tools/feature_store/check_feature_store_delta_test.py +++ b/tzrec/tools/feature_store/check_feature_store_delta_test.py @@ -11,9 +11,11 @@ import json import os +import sys import tempfile import unittest from types import SimpleNamespace +from unittest import mock import numpy as np import pyarrow as pa @@ -22,6 +24,7 @@ from tzrec.tools.feature_store.check_feature_store_delta import ( LocalSample, committed_parquet_paths, + create_feature_store_view, load_committed_upload, parse_args, resolve_output_dir, @@ -42,6 +45,44 @@ def get_online_features(self, feature_name, keys, version): class CheckFeatureStoreDeltaTest(unittest.TestCase): + def test_create_feature_store_view_uses_public_endpoint_when_supported(self): + captured_kwargs = {} + view = SimpleNamespace( + pk_field="embedding_name", + sk_field="key_id", + embedding_field="embedding", + ) + project = SimpleNamespace(get_dynamic_embedding_feature_view=lambda name: view) + + class FakeFeatureStoreClient: + def __init__(self, test_mode=False, **kwargs): + captured_kwargs.update(kwargs) + captured_kwargs["test_mode"] = test_mode + + def get_project(self, name): + return project + + settings = SimpleNamespace( + access_key_id="ak-id", + access_key_secret="ak-secret", + region="cn-test", + endpoint="", + security_token="", + featuredb_username="featuredb-user", + featuredb_password="featuredb-password", + project_name="project", + feature_view_name="view", + ) + feature_store_module = SimpleNamespace( + FeatureStoreClient=FakeFeatureStoreClient + ) + + with mock.patch.dict(sys.modules, {"feature_store_py": feature_store_module}): + actual = create_feature_store_view(settings) + + self.assertIs(actual, view) + self.assertTrue(captured_kwargs["test_mode"]) + def test_parse_args_does_not_accept_credentials(self): args = parse_args(["--pipeline_config", "pipeline.config"]) diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index 3dc360de3..ac6b8d18c 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -1705,7 +1705,6 @@ def _get_view(self) -> Any: "security_token": self._settings.security_token or None, "featuredb_username": self._settings.featuredb_username or None, "featuredb_password": self._settings.featuredb_password or None, - "test_mode": True, } client = client_factory(**kwargs) project = client.get_project(self._settings.project_name) diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index 1a4aa25b1..8403e0539 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -456,6 +456,7 @@ def test_start_reuses_existing_dynamic_embedding_feature_view(self): self.assertEqual(factory.project.dynamic_get_calls, ["shared_embeddings"]) self.assertEqual(factory.project.generic_get_calls, ["shared_embeddings"]) self.assertEqual(factory.project.create_calls, []) + self.assertNotIn("test_mode", factory.calls[0]) self.assertEqual(view.closed, [True]) def test_start_creates_missing_dynamic_embedding_feature_view(self): From 70ee3b848c3b4d671fd213aad5608b9e6578da16 Mon Sep 17 00:00:00 2001 From: gecheng Date: Sun, 19 Jul 2026 11:58:41 +0800 Subject: [PATCH 11/13] [bugfix] sanitize distributed export configs --- tzrec/utils/export_util.py | 2 +- tzrec/utils/export_util_test.py | 25 +++++++++++++++++++++++-- 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/tzrec/utils/export_util.py b/tzrec/utils/export_util.py index d07002737..0a76fbff7 100644 --- a/tzrec/utils/export_util.py +++ b/tzrec/utils/export_util.py @@ -1597,7 +1597,7 @@ def export_distributed_embedding( pipeline_config = copy.copy(pipeline_config) pipeline_config.ClearField("feature_configs") pipeline_config.feature_configs.extend(feature_configs) - config_util.save_message( + config_util.save_pipeline_config_artifact( pipeline_config, os.path.join(save_dir, "pipeline.config") ) diff --git a/tzrec/utils/export_util_test.py b/tzrec/utils/export_util_test.py index 6e3c53d7a..c253efb83 100644 --- a/tzrec/utils/export_util_test.py +++ b/tzrec/utils/export_util_test.py @@ -30,6 +30,7 @@ MLPDenseEmbeddingConfig, ) from tzrec.protos.pipeline_pb2 import EasyRecConfig +from tzrec.utils import config_util from tzrec.utils.export_util import ( _dedup_key_files_by_realpath, _get_dense_embedding_leaf_module_names, @@ -174,7 +175,9 @@ def test_distributed_embedding_export_skips_nonzero_rank_before_pg_init( else: os.environ[key] = value - def test_distributed_embedding_export_uses_export_overrides(self) -> None: + def test_distributed_embedding_export_uses_overrides_and_sanitizes_config( + self, + ) -> None: class FakeBatch: def to(self, device): # type: ignore[no-untyped-def] return self @@ -222,6 +225,14 @@ def forward(self, data, device=None): # type: ignore[no-untyped-def] eval_input_path="eval_input", model_dir="model_dir", ) + dump_config = pipeline_config.train_config.delta_embedding_dump_config + feature_store_config = dump_config.feature_store_config + feature_store_config.region = "cn-test" + feature_store_config.project_name = "project_a" + feature_store_config.feature_entity_name = "embedding_entity" + feature_store_config.feature_view_name = "shared_embeddings" + feature_store_config.version = "model_a@export_1" + feature_store_config.security_token = "SECRET_STS" model_acc = {"SPARSE_INT64": "1", "cand_seq_pk": "cand_seq"} fake_scripted = mock.Mock() @@ -255,7 +266,6 @@ def forward(self, data, device=None): # type: ignore[no-untyped-def] "tzrec.utils.export_util._get_sparse_embedding_tensor", return_value=({}, {}, {}, {}), ), - mock.patch("tzrec.utils.export_util.config_util.save_message"), mock.patch( "tzrec.utils.export_util.create_fg_json", return_value={"features": []}, @@ -289,6 +299,17 @@ def forward(self, data, device=None): # type: ignore[no-untyped-def] ) with open(os.path.join(tmp, "model_acc.json")) as f: self.assertEqual(json.load(f), model_acc) + pipeline_config_path = os.path.join(tmp, "pipeline.config") + with open(pipeline_config_path) as f: + self.assertNotIn("SECRET_STS", f.read()) + exported_config = config_util.load_pipeline_config(pipeline_config_path) + exported_dump_config = ( + exported_config.train_config.delta_embedding_dump_config + ) + exported_feature_store_config = exported_dump_config.feature_store_config + self.assertEqual(exported_feature_store_config.project_name, "project_a") + self.assertFalse(exported_feature_store_config.HasField("security_token")) + self.assertTrue(feature_store_config.HasField("security_token")) finally: _restore_env(old_env) shutil.rmtree(tmp, ignore_errors=True) From a6112431a0f7e8f0c8ecdecdb6abaf3608aeee96 Mon Sep 17 00:00:00 2001 From: gecheng Date: Sun, 19 Jul 2026 15:17:02 +0800 Subject: [PATCH 12/13] [bugfix] synchronize timed delta dump failures across ranks --- tzrec/utils/delta_embedding_dump.py | 92 +++++++--- tzrec/utils/delta_embedding_dump_test.py | 209 ++++++++++++++++++++++- 2 files changed, 275 insertions(+), 26 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 3499a72ce..99a93502e 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -697,11 +697,20 @@ def maybe_dump(self, global_step: int) -> None: Args: global_step: Current training step. """ - # This is a throttled local shared-filesystem check, not a distributed - # collective, so all ranks can surface an async rank-zero failure quickly. - self._check_feature_store_upload_error() - if self._should_dump(global_step): - self.dump(global_step) + upload_error = self._feature_store_upload_error() + if self._should_dump(global_step, upload_error): + dump_error: Optional[BaseException] = None + try: + if ( + getattr(self, "_world_size", 1) > 1 + and self._interval_secs is not None + ): + self.dump(global_step, check_upload_error=False) + else: + self.dump(global_step) + except BaseException as exc: + dump_error = exc + self._raise_if_any_timed_dump_failed(dump_error) self._last_dump_step = global_step if self._interval_secs is not None and self._rank == 0: # Schedule from dump completion to avoid immediate catch-up dumps @@ -709,19 +718,23 @@ def maybe_dump(self, global_step: int) -> None: self._next_dump_time = time.monotonic() + self._interval_secs self._tracker.step() - def _should_dump(self, global_step: int) -> bool: - """Return one rank-consistent dump decision for the current step.""" + def _should_dump( + self, global_step: int, local_error: Optional[BaseException] = None + ) -> bool: + """Synchronize upload failure state and return one dump decision.""" if global_step <= 0: - return False - if self._interval_steps is not None: - return global_step % self._interval_steps == 0 - if self._next_dump_time is None: - raise RuntimeError( - "time-based delta embedding dumper must be started before training" - ) + should_dump = False + elif self._interval_steps is not None: + should_dump = global_step % self._interval_steps == 0 + else: + if self._next_dump_time is None: + raise RuntimeError( + "time-based delta embedding dumper must be started before training" + ) + should_dump = self._rank == 0 and time.monotonic() >= self._next_dump_time - should_dump = self._rank == 0 and time.monotonic() >= self._next_dump_time - if self._world_size > 1: + any_failed = local_error is not None + if getattr(self, "_world_size", 1) > 1 and self._interval_secs is not None: if not ( torch.distributed.is_available() and torch.distributed.is_initialized() ): @@ -729,13 +742,52 @@ def _should_dump(self, global_step: int) -> bool: "distributed time-based delta embedding dump requires an " "initialized process group" ) - decision = torch.tensor( - [int(should_dump)], dtype=torch.int32, device=self._device + state = torch.tensor( + [int(should_dump), int(any_failed)], + dtype=torch.int32, + device=self._device, + ) + torch.distributed.all_reduce(state, op=torch.distributed.ReduceOp.MAX) + should_dump = bool(state[0].item()) + any_failed = bool(state[1].item()) + + if local_error is not None: + raise local_error.with_traceback(local_error.__traceback__) + if any_failed: + raise FeatureStoreUploadError( + "FeatureStore delta upload failed on another distributed worker; " + "all training workers are stopping and parquet outbox files were " + "retained" ) - torch.distributed.broadcast(decision, src=0) - should_dump = bool(decision.item()) return should_dump + def _raise_if_any_timed_dump_failed( + self, local_error: Optional[BaseException] + ) -> None: + """Propagate one timed dump failure before the next training step.""" + any_failed = local_error is not None + if getattr(self, "_world_size", 1) > 1 and self._interval_secs is not None: + if not ( + torch.distributed.is_available() and torch.distributed.is_initialized() + ): + raise RuntimeError( + "distributed time-based delta embedding dump requires an " + "initialized process group" + ) + failed = torch.tensor( + [int(any_failed)], dtype=torch.int32, device=self._device + ) + torch.distributed.all_reduce(failed, op=torch.distributed.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( + "timed delta embedding dump failed on another distributed worker; " + "all workers are stopping" + ) + def final_dump(self, global_step: int) -> Optional[str]: """Flush the trailing partial interval at the end of training. diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 2c573eb8c..ffb973628 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -687,7 +687,7 @@ def test_maybe_dump_uses_elapsed_time_and_resets_from_completion(self): dumper._device = torch.device("cpu") dumper._tracker = mock.MagicMock() with ( - mock.patch.object(dumper, "_check_feature_store_upload_error"), + mock.patch.object(dumper, "_feature_store_upload_error", return_value=None), mock.patch.object(dumper, "dump") as dump_mock, mock.patch( "tzrec.utils.delta_embedding_dump.time.monotonic", @@ -707,7 +707,7 @@ def test_maybe_dump_uses_elapsed_time_and_resets_from_completion(self): self.assertEqual(dumper._last_dump_step, 13) self.assertEqual(dumper._tracker.step.call_count, 4) - def test_time_dump_decision_is_broadcast_from_rank_zero(self): + def test_time_dump_decision_is_reduced_from_rank_zero(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._interval_steps = None dumper._interval_secs = 60.0 @@ -716,17 +716,214 @@ def test_time_dump_decision_is_broadcast_from_rank_zero(self): dumper._world_size = 2 dumper._device = torch.device("cpu") - def fake_broadcast(decision, src): - self.assertEqual(src, 0) - decision.fill_(1) + def fake_all_reduce(state, op=None): + self.assertIs(op, torch.distributed.ReduceOp.MAX) + self.assertEqual(state.tolist(), [0, 0]) + state[0] = 1 with ( mock.patch("torch.distributed.is_available", return_value=True), mock.patch("torch.distributed.is_initialized", return_value=True), - mock.patch("torch.distributed.broadcast", side_effect=fake_broadcast), + mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), ): self.assertTrue(dumper._should_dump(10)) + def test_timed_maybe_dump_reduces_local_upload_error_before_raising(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_steps = None + dumper._interval_secs = 60.0 + dumper._next_dump_time = 0.0 + dumper._last_dump_step = None + dumper._rank = 0 + dumper._world_size = 2 + dumper._device = torch.device("cpu") + dumper._tracker = mock.MagicMock() + local_error = FeatureStoreUploadError("local upload failed") + collective_called = False + + def fake_all_reduce(state, op=None): + nonlocal collective_called + self.assertIs(op, torch.distributed.ReduceOp.MAX) + self.assertEqual(state.tolist(), [1, 1]) + collective_called = True + + with ( + mock.patch.object( + dumper, "_feature_store_upload_error", return_value=local_error + ), + mock.patch.object(dumper, "dump") as dump_mock, + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), + ): + with self.assertRaises(FeatureStoreUploadError) as context: + dumper.maybe_dump(10) + + self.assertTrue(collective_called) + self.assertIs(context.exception, local_error) + dump_mock.assert_not_called() + dumper._tracker.step.assert_not_called() + + def test_timed_maybe_dump_surfaces_remote_upload_error_when_marker_is_throttled( + self, + ): + with tempfile.TemporaryDirectory() as tmp_dir: + marker_path = os.path.join(tmp_dir, "last_error.json") + with open(marker_path, "w") as output: + output.write("{}") + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_steps = None + dumper._interval_secs = 60.0 + dumper._next_dump_time = 0.0 + dumper._last_dump_step = None + dumper._rank = 1 + dumper._world_size = 2 + dumper._device = torch.device("cpu") + dumper._tracker = mock.MagicMock() + dumper._feature_store_enabled = True + dumper._feature_store_error_marker_path = marker_path + dumper._next_feature_store_error_check = 100.0 + dumper._feature_store_error_check_interval_secs = 1 + dumper._uploader = None + + def fake_all_reduce(state, op=None): + self.assertIs(op, torch.distributed.ReduceOp.MAX) + self.assertEqual(state.tolist(), [0, 0]) + state[0] = 1 + state[1] = 1 + + with ( + mock.patch.object(dumper, "dump") as dump_mock, + mock.patch( + "tzrec.utils.delta_embedding_dump.time.monotonic", + return_value=50.0, + ), + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), + ): + with self.assertRaisesRegex( + FeatureStoreUploadError, "another distributed worker" + ): + dumper.maybe_dump(10) + + dump_mock.assert_not_called() + dumper._tracker.step.assert_not_called() + + def test_timed_maybe_dump_advances_after_all_workers_succeed(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_steps = None + dumper._interval_secs = 60.0 + dumper._next_dump_time = 0.0 + dumper._last_dump_step = None + dumper._rank = 0 + dumper._world_size = 2 + dumper._device = torch.device("cpu") + dumper._tracker = mock.MagicMock() + collective_states = [] + + def fake_all_reduce(state, op=None): + self.assertIs(op, torch.distributed.ReduceOp.MAX) + collective_states.append(state.tolist()) + + with ( + mock.patch.object(dumper, "_feature_store_upload_error", return_value=None), + mock.patch.object( + dumper, "dump", return_value="delta.parquet" + ) as dump_mock, + mock.patch( + "tzrec.utils.delta_embedding_dump.time.monotonic", + side_effect=[1.0, 2.0], + ), + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), + ): + dumper.maybe_dump(10) + + self.assertEqual(collective_states, [[1, 0], [0]]) + dump_mock.assert_called_once_with(10, check_upload_error=False) + self.assertEqual(dumper._last_dump_step, 10) + self.assertEqual(dumper._next_dump_time, 62.0) + dumper._tracker.step.assert_called_once_with() + + def test_timed_maybe_dump_reduces_local_dump_failure_before_raising(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_steps = None + dumper._interval_secs = 60.0 + dumper._next_dump_time = 0.0 + dumper._last_dump_step = None + dumper._rank = 0 + dumper._world_size = 2 + dumper._device = torch.device("cpu") + dumper._tracker = mock.MagicMock() + dump_error = RuntimeError("local dump failed") + collective_index = 0 + + def fake_all_reduce(state, op=None): + nonlocal collective_index + self.assertIs(op, torch.distributed.ReduceOp.MAX) + if collective_index == 0: + self.assertEqual(state.tolist(), [1, 0]) + else: + self.assertEqual(state.tolist(), [1]) + collective_index += 1 + + with ( + mock.patch.object(dumper, "_feature_store_upload_error", return_value=None), + mock.patch.object(dumper, "dump", side_effect=dump_error) as dump_mock, + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), + ): + with self.assertRaises(RuntimeError) as context: + dumper.maybe_dump(10) + + self.assertEqual(collective_index, 2) + self.assertIs(context.exception, dump_error) + dump_mock.assert_called_once_with(10, check_upload_error=False) + self.assertIsNone(dumper._last_dump_step) + dumper._tracker.step.assert_not_called() + + def test_timed_maybe_dump_surfaces_remote_dump_failure(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_steps = None + dumper._interval_secs = 60.0 + dumper._next_dump_time = 0.0 + dumper._last_dump_step = None + dumper._rank = 0 + dumper._world_size = 2 + dumper._device = torch.device("cpu") + dumper._tracker = mock.MagicMock() + collective_index = 0 + + def fake_all_reduce(state, op=None): + nonlocal collective_index + self.assertIs(op, torch.distributed.ReduceOp.MAX) + if collective_index == 0: + self.assertEqual(state.tolist(), [1, 0]) + else: + self.assertEqual(state.tolist(), [0]) + state[0] = 1 + collective_index += 1 + + with ( + mock.patch.object(dumper, "_feature_store_upload_error", return_value=None), + mock.patch.object( + dumper, "dump", return_value="delta.parquet" + ) as dump_mock, + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), + ): + with self.assertRaisesRegex(RuntimeError, "another distributed worker"): + dumper.maybe_dump(10) + + self.assertEqual(collective_index, 2) + dump_mock.assert_called_once_with(10, check_upload_error=False) + self.assertIsNone(dumper._last_dump_step) + dumper._tracker.step.assert_not_called() + def test_final_dump_skips_step_already_dumped_by_time_interval(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._interval_steps = None From cc03f25922645a27b54dadc52417decca04773fd Mon Sep 17 00:00:00 2001 From: gecheng Date: Sun, 19 Jul 2026 15:46:41 +0800 Subject: [PATCH 13/13] [bugfix] reject uneven batches during timed delta dump --- tzrec/main.py | 10 +++++++--- tzrec/utils/dist_util.py | 16 +++++++++++++++- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index e0b47e8c9..cc19ddd4e 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -467,17 +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 - sync_train_data_exhaustion = check_all_workers_data_status or ( + 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 broadcast their per-step decision, so all ranks must stop - # consuming data together before any worker can leave that collective. + # 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=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}") diff --git a/tzrec/utils/dist_util.py b/tzrec/utils/dist_util.py index bc916a7f1..1ab6859ec 100644 --- a/tzrec/utils/dist_util.py +++ b/tzrec/utils/dist_util.py @@ -237,6 +237,7 @@ def __init__( dmp_collection_sync_interval_batches: Optional[int] = 1, enqueue_batch_after_forward: bool = False, check_all_workers_data_status: bool = False, + fail_on_uneven_data: bool = False, ) -> None: super().__init__( model, @@ -251,6 +252,7 @@ def __init__( enqueue_batch_after_forward, ) self._check_all_workers_data_status = check_all_workers_data_status + self._fail_on_uneven_data = fail_on_uneven_data self._sync_at_progress_entry = ( device.type == "cuda" and dist.is_initialized() @@ -288,7 +290,15 @@ def _next_batch(self, dataloader_iter: Iterator[In]) -> Optional[In]: 0 if batch is None else 1, dtype=torch.float, device=self._device ) dist.all_reduce(has_batch, dist.ReduceOp.AVG) - if has_batch.item() < 1: + available_fraction = has_batch.item() + if self._fail_on_uneven_data and 0 < available_fraction < 1: + self._dataloader_exhausted = True + raise RuntimeError( + "training dataloader exhausted unevenly across workers; " + "refusing to drop remainder batches; ensure every worker " + "yields the same number of batches" + ) + if available_fraction < 1: # We drop remainder batches on all workers, # if one worker does not have a batch self._dataloader_exhausted = True @@ -336,6 +346,7 @@ def create_train_pipeline( model: nn.Module, optimizer: Optional[torch.optim.Optimizer] = None, check_all_workers_data_status: bool = False, + fail_on_uneven_data: bool = False, ) -> TrainPipeline: """Create TrainPipeline. @@ -344,6 +355,8 @@ def create_train_pipeline( optimizer (torch.optim.Optimizer): a KeyedOptimizer. check_all_workers_data_status (bool): check data on all workers is available or not. + fail_on_uneven_data (bool): raise instead of dropping remainder batches + when workers exhaust their dataloaders at different times. Return: a TrainPipeline. @@ -373,4 +386,5 @@ def create_train_pipeline( model.device, execute_all_batches=True, check_all_workers_data_status=check_all_workers_data_status, + fail_on_uneven_data=fail_on_uneven_data, )