From c49f4558c828fed8c23801e8723bc36b4c233fc5 Mon Sep 17 00:00:00 2001 From: gecheng Date: Thu, 16 Jul 2026 11:22:11 +0800 Subject: [PATCH 01/50] 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/50] 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/50] 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/50] 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/50] 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/50] 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/50] 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/50] [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/50] [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/50] [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/50] [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/50] [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/50] [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, ) From 0af5f74fc2188761a686ac6510abfeea5931b033 Mon Sep 17 00:00:00 2001 From: gecheng Date: Tue, 21 Jul 2026 15:58:07 +0800 Subject: [PATCH 14/50] [bugfix] reject uneven dataloader exhaustion for all multi-rank delta dumps With dump_interval_steps and check_all_workers_data_status=False, ranks could exhaust their dataloaders at different steps (e.g. 49 and 50). final_dump all-reduces the final step with MAX, so the rank that stopped at 49 adopted 50 and hit the boundary-step skip even though it never ran maybe_dump(50); its trailing tracked rows were never written and were silently lost. requires_synced_dataloader_exhaustion now returns True for any multi-rank cadence, not only timed dumps, so the train pipeline checks all workers' data status and fails loudly on uneven exhaustion instead of letting a lagging rank skip its final delta. This keeps every rank participating in every boundary dump, the invariant the boundary-step skip in final_dump relies on. Adds a 49/50 regression test asserting the uneven stop is rejected, plus a contract unit test. Co-Authored-By: Claude --- tzrec/main.py | 5 +- tzrec/utils/delta_embedding_dump.py | 36 +++++++++----- tzrec/utils/delta_embedding_dump_test.py | 63 +++++++++++++++++++++++- 3 files changed, 87 insertions(+), 17 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index cc19ddd4e..cefc308fe 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -475,8 +475,9 @@ def run_eval(step: int, epoch: int) -> None: check_all_workers_data_status or require_equal_train_batches ) for i_epoch in epoch_iter: - # Timed dumps synchronize their per-step state, so uneven workers must - # fail together instead of silently truncating data or leaving collectives. + # Multi-rank dumps synchronize their per-step state, so uneven workers + # must fail together instead of silently dropping a lagging rank's + # trailing delta rows or leaving collectives. pipeline = create_train_pipeline( model, optimizer, diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 99a93502e..48d104203 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -474,8 +474,15 @@ def clear(self) -> None: @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 + """Return whether input exhaustion must stay aligned across ranks. + + Multi-rank dumps of either cadence need every rank to reach the same + final step. ``final_dump`` skips steps already written on their dump + boundary, and a rank that stops before the synced MAX step never ran + that boundary's ``maybe_dump``, so an unsynced stop would make it + adopt the boundary skip and silently drop its trailing tracked rows. + """ + return self._world_size > 1 def start(self) -> None: """Start timed cadence and rank-zero publication after initialization.""" @@ -810,10 +817,12 @@ def final_dump(self, global_step: int) -> Optional[str]: # 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 - # delta -- and torchrec's ``get_unique`` raises - # ``torch.cat(): expected a non-empty list of Tensors`` on the empty - # consumer window. Re-dumping would also overwrite the already-written - # boundary shards (with an empty file under multi-GPU), so skip. + # delta (multi-rank dumps force synced exhaustion, so every rank + # participated in the boundary dump) -- and torchrec's ``get_unique`` + # raises ``torch.cat(): expected a non-empty list of Tensors`` on the + # empty 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 @@ -835,13 +844,14 @@ def _sync_final_step(self, global_step: int) -> int: """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 - ``check_all_workers_data_status=False`` ranks can exhaust the dataloader - at different steps, so without syncing each would write a lone shard to - its own ``step_/`` dir -- exactly the ragged shard set the per-rank - 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. + and the training loop keeps ``final_dump`` aligned too: multi-rank + dumps require synced dataloader exhaustion, because a rank stopping at + an earlier step would adopt the synced MAX step, hit the boundary-step + skip without ever running that boundary's ``maybe_dump``, and silently + drop its trailing tracked rows. The MAX all-reduce remains as a + defensive guard so every rank still takes the same skip/dump decision + into the same ``step_/`` dir -- the ragged shard set the per-rank + empty-shard logic prevents. The same collective carries a local/shared uploader-failure bit. No rank raises before all ranks have entered it, preventing an asynchronous marker diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index ffb973628..29dccda21 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -55,6 +55,7 @@ validate_delta_embedding_dump_config, validate_delta_embedding_dump_no_zch_features, ) +from tzrec.utils.dist_util import create_train_pipeline from tzrec.utils.dynamicemb_util import has_dynamicemb from tzrec.utils.feature_store_delta_uploader import ( DELTA_DUMP_GENERATION_METADATA_KEY, @@ -217,6 +218,38 @@ def _run_sharded_delta_embedding_dump(rank: int, world_size: int, output_dir: st torch.distributed.barrier() +def _run_uneven_exhaustion_rejection(rank: int, world_size: int, output_dir: str): + # Regression: with rank 0 stopping at step 49 and rank 1 at step 50, + # final_dump used to sync both ranks to the boundary 50, both took the + # boundary-step skip, and rank 0's trailing tracked rows were silently + # lost. Multi-rank dumps now force aligned dataloader exhaustion, so the + # uneven stop must fail loudly on every rank instead. + with MultiProcessContext(rank=rank, world_size=world_size, backend="nccl") as ctx: + model = _build_sharded_delta_dump_model(rank, world_size, ctx) + dumper = DeltaEmbeddingDumper( + model, + DeltaEmbeddingDumpConfig(dump_interval_steps=50, output_dir=output_dir), + output_dir, + torch.device(f"cuda:{rank}"), + [], + ) + testcase = unittest.TestCase() + testcase.assertTrue(dumper.requires_synced_dataloader_exhaustion) + pipeline = create_train_pipeline( + model, + check_all_workers_data_status=dumper.requires_synced_dataloader_exhaustion, + fail_on_uneven_data=dumper.requires_synced_dataloader_exhaustion, + ) + # Rank 0's 50th fetch returns None while rank 1 still has batch 50; the + # pipeline's collective data-status check must raise on both ranks. + num_batches = 49 if rank == 0 else 50 + batches = iter([_sharded_features(rank) for _ in range(num_batches)]) + with testcase.assertRaisesRegex(RuntimeError, "exhausted unevenly"): + for _ in range(num_batches + 1): + pipeline._next_batch(batches) + torch.distributed.barrier() + + class DeltaEmbeddingDumpValidationTest(unittest.TestCase): def test_missing_config_skips_runtime_validation(self): with mock.patch.dict(os.environ, {"WORLD_SIZE": "2"}): @@ -1003,12 +1036,18 @@ def test_multi_rank_minutes_requires_synced_dataloader_exhaustion(self): self.assertTrue(dumper.requires_synced_dataloader_exhaustion) - def test_step_interval_does_not_require_synced_dataloader_exhaustion(self): + def test_multi_rank_step_interval_requires_synced_dataloader_exhaustion(self): + # Regression: with ranks finishing at steps 49 and 50, final_dump synced + # both to the boundary 50, both took the boundary-step skip, and the + # rank that never ran maybe_dump(50) silently dropped its trailing + # tracked rows. Every multi-rank cadence must keep exhaustion aligned + # so all ranks participate in every boundary dump. dumper = object.__new__(DeltaEmbeddingDumper) dumper._interval_secs = None + dumper._interval_steps = 50 dumper._world_size = 2 - self.assertFalse(dumper.requires_synced_dataloader_exhaustion) + self.assertTrue(dumper.requires_synced_dataloader_exhaustion) def test_single_rank_minutes_does_not_require_synced_dataloader_exhaustion(self): dumper = object.__new__(DeltaEmbeddingDumper) @@ -1445,6 +1484,26 @@ def test_row_wise_sharded_dump_writes_global_key_ids(self): ) ) + @unittest.skipIf(torch.cuda.device_count() < 2, "test requires 2+ GPUs") + @mark_ci_scope("gpu") + def test_uneven_exhaustion_is_rejected_for_step_interval_dump(self): + with ( + tempfile.TemporaryDirectory() as tmp_dir, + mock.patch.dict( + os.environ, + { + "NCCL_DEBUG": "WARN", + "FORCED_NCCL_DEBUG": "WARN", + "NCCL_DEBUG_SUBSYS": "", + }, + ), + ): + self._run_multi_process_test( + callable=_run_uneven_exhaustion_rejection, + world_size=self.world_size, + output_dir=tmp_dir, + ) + class DeltaEmbeddingDumpDynamicembIntegrationTest(unittest.TestCase): """End-to-end multi-process delta dump over a sharded dynamicemb model. From 00a874b399c023736780f5b222be0f4e93a33924 Mon Sep 17 00:00:00 2001 From: gecheng Date: Tue, 21 Jul 2026 16:25:03 +0800 Subject: [PATCH 15/50] [bugfix] rendezvous FeatureStore failure bit on every distributed dump cadence The uploader-failure rendezvous in _should_dump only ran for timed (minute) dumps. In the default step cadence with FeatureStore upload, rank zero observes an async upload failure immediately through uploader.check_error() while peers only see the shared error marker through a throttled poll, so rank zero raised alone and its peers marched on into the next training collective and hung. Boundary dump failures had the same hole: _raise_if_any_timed_dump_failed also gated its failure-bit all-reduce on the timed cadence, so any rank-local boundary dump failure (e.g. a shard write error on one node) raised on one rank only. Add _requires_dump_state_rendezvous() and use it for the _should_dump rendezvous and for selecting check_upload_error=False in maybe_dump, so distributed FeatureStore dumps of either cadence synchronize the failure bit before any rank raises. Also widen the interval dump failure rendezvous (renamed to _raise_if_any_interval_dump_failed) to every multi-rank cadence, mirroring _raise_if_any_final_dump_failed. Co-Authored-By: Claude --- tzrec/utils/delta_embedding_dump.py | 53 +++++-- tzrec/utils/delta_embedding_dump_test.py | 183 +++++++++++++++++++++++ 2 files changed, 221 insertions(+), 15 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 48d104203..0b38ed6b2 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -708,16 +708,15 @@ def maybe_dump(self, global_step: int) -> None: 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 - ): + if self._requires_dump_state_rendezvous(): + # The uploader-failure bit was already rendezvoused in + # _should_dump; a rank-local recheck here could diverge. 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._raise_if_any_interval_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 @@ -725,6 +724,21 @@ def maybe_dump(self, global_step: int) -> None: self._next_dump_time = time.monotonic() + self._interval_secs self._tracker.step() + def _requires_dump_state_rendezvous(self) -> bool: + """Return whether maybe_dump must all-reduce its dump and failure state. + + Timed dumps OR-reduce the rank-zero clock decision so every rank dumps + together. Distributed FeatureStore dumps of either cadence rendezvous + the uploader-failure bit: rank zero observes an async upload failure + immediately while peers only see the shared error marker through a + throttled poll, so without the collective rank zero could raise alone + and strand its peers in the next training collective. + """ + return getattr(self, "_world_size", 1) > 1 and ( + self._interval_secs is not None + or getattr(self, "_feature_store_enabled", False) + ) + def _should_dump( self, global_step: int, local_error: Optional[BaseException] = None ) -> bool: @@ -741,13 +755,13 @@ def _should_dump( should_dump = self._rank == 0 and time.monotonic() >= self._next_dump_time any_failed = local_error is not None - if getattr(self, "_world_size", 1) > 1 and self._interval_secs is not None: + if self._requires_dump_state_rendezvous(): if not ( torch.distributed.is_available() and torch.distributed.is_initialized() ): raise RuntimeError( - "distributed time-based delta embedding dump requires an " - "initialized process group" + "distributed delta embedding dump requires an initialized " + "process group" ) state = torch.tensor( [int(should_dump), int(any_failed)], @@ -768,18 +782,26 @@ def _should_dump( ) return should_dump - def _raise_if_any_timed_dump_failed( + def _raise_if_any_interval_dump_failed( self, local_error: Optional[BaseException] ) -> None: - """Propagate one timed dump failure before the next training step.""" + """Propagate one interval dump failure before the next training step. + + Boundary and timed dumps run on every rank, but a failure can still be + rank-local (a shard write error on one node, or one rank observing the + shared uploader-error marker ahead of the others). Reduce the failure + bit on every multi-rank cadence so all workers stop together instead + of the failing rank raising alone and stranding its peers in the next + training collective. + """ any_failed = local_error is not None - if getattr(self, "_world_size", 1) > 1 and self._interval_secs is not None: + if getattr(self, "_world_size", 1) > 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" + "distributed delta embedding dump requires an initialized " + "process group" ) failed = torch.tensor( [int(any_failed)], dtype=torch.int32, device=self._device @@ -791,7 +813,7 @@ def _raise_if_any_timed_dump_failed( raise local_error.with_traceback(local_error.__traceback__) if any_failed: raise RuntimeError( - "timed delta embedding dump failed on another distributed worker; " + "delta embedding dump failed on another distributed worker; " "all workers are stopping" ) @@ -907,7 +929,8 @@ def dump(self, global_step: int, check_upload_error: bool = True) -> Optional[st 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. + ``maybe_dump`` and ``final_dump`` disable it after synchronizing + the error bit across ranks. Returns: Path to the dumped parquet file, or None if no data to dump. diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 29dccda21..e466209a2 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -957,6 +957,189 @@ def fake_all_reduce(state, op=None): self.assertIsNone(dumper._last_dump_step) dumper._tracker.step.assert_not_called() + def test_step_maybe_dump_reduces_local_upload_error_before_raising(self): + # Rank zero observes an async uploader failure immediately; without a + # step-cadence rendezvous it would raise alone and strand peers that + # skipped the throttled marker poll in the next training collective. + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_steps = 50 + dumper._interval_secs = None + dumper._last_dump_step = None + dumper._rank = 0 + dumper._world_size = 2 + dumper._device = torch.device("cpu") + dumper._tracker = mock.MagicMock() + dumper._feature_store_enabled = True + 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(), [0, 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_step_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 = 50 + dumper._interval_secs = None + 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 + # The throttled poll hides the marker on this rank; only the + # rendezvoused failure bit from rank zero surfaces the error. + 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[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_step_feature_store_boundary_dump_skips_rank_local_upload_recheck(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_steps = 50 + dumper._interval_secs = None + 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 + 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("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(50) + + self.assertEqual(collective_states, [[1, 0], [0]]) + dump_mock.assert_called_once_with(50, check_upload_error=False) + self.assertEqual(dumper._last_dump_step, 50) + dumper._tracker.step.assert_called_once_with() + + def test_step_boundary_dump_failure_is_reduced_before_raising(self): + # Even without FeatureStore, a rank-local boundary dump failure (e.g. a + # shard write error on one node) must be reduced so the failing rank + # does not raise alone and strand peers in the next collective. + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_steps = 50 + dumper._interval_secs = None + dumper._last_dump_step = None + dumper._rank = 1 + dumper._world_size = 2 + dumper._device = torch.device("cpu") + dumper._tracker = mock.MagicMock() + dump_error = RuntimeError("local dump 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]) + collective_called = True + + with ( + 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(50) + + self.assertTrue(collective_called) + self.assertIs(context.exception, dump_error) + dump_mock.assert_called_once_with(50) + self.assertIsNone(dumper._last_dump_step) + dumper._tracker.step.assert_not_called() + + def test_step_boundary_dump_surfaces_remote_failure(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_steps = 50 + dumper._interval_secs = None + dumper._last_dump_step = None + dumper._rank = 1 + dumper._world_size = 2 + dumper._device = torch.device("cpu") + dumper._tracker = mock.MagicMock() + + def fake_all_reduce(state, op=None): + self.assertIs(op, torch.distributed.ReduceOp.MAX) + self.assertEqual(state.tolist(), [0]) + state[0] = 1 + + with ( + 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(50) + + dump_mock.assert_called_once_with(50) + 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 736a4712ead9450ec7bf637b9b20569363754771 Mon Sep 17 00:00:00 2001 From: gecheng Date: Tue, 21 Jul 2026 16:52:22 +0800 Subject: [PATCH 16/50] [perf] pipeline the delta dump state rendezvous off the training hot path Multi-rank dumps all-reduced the timer vote and uploader-failure bit, with host-synchronizing .item() reads, on every training step, while the synced-exhaustion pipeline already pays one blocking availability collective per batch: two blocking collectives per step for dumps that fire minutes apart. The rendezvous is now launched asynchronously and consumed on the next maybe_dump call. NCCL orders it ahead of the next step's communication on the same process group, so it has completed by the time the next call reads it back and the hot path pays only an async launch plus a non-blocking read; the decision lands at most one step late, which is immaterial for minute-scale intervals. Timed votes arm an in-flight guard so the still-elapsed deadline cannot fire a second launch before the dump reschedules from its completion time. Step-boundary decisions stay deterministic and on their exact boundary step; only the failure bit is pipelined there. Coalescing the bits into the pipeline's availability collective would also remove the per-step launch, but couples the dumper into pipeline prefetch internals and moves its raise/dump control flow across components; the overlap keeps the failure domain inside the dumper while taking the blocking synchronization off the hot path. Co-Authored-By: Claude --- tzrec/utils/delta_embedding_dump.py | 153 ++++++++++++----- tzrec/utils/delta_embedding_dump_test.py | 205 ++++++++++++++--------- 2 files changed, 236 insertions(+), 122 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 0b38ed6b2..3a6b4af91 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -388,6 +388,9 @@ def __init__( self._interval_steps = int(config.dump_interval_steps) self._next_dump_time: Optional[float] = None self._last_dump_step: Optional[int] = None + self._pending_rendezvous: Optional[torch.Tensor] = None + self._pending_upload_error: Optional[BaseException] = None + self._timed_dump_in_flight = False self._output_dir = config.output_dir or os.path.join( model_dir, "delta_embedding_dump" ) @@ -704,16 +707,31 @@ def maybe_dump(self, global_step: int) -> None: Args: global_step: Current training step. """ - upload_error = self._feature_store_upload_error() - if self._should_dump(global_step, upload_error): + if self._requires_dump_state_rendezvous(): + should_dump, any_failed = self._consume_dump_state_rendezvous() + if self._interval_steps is not None: + # Step-boundary decisions are deterministic across ranks and + # stay on their exact boundary step; only the failure bit + # travels through the pipelined rendezvous. + should_dump = ( + global_step > 0 and global_step % self._interval_steps == 0 + ) + if any_failed: + self._raise_rendezvoused_upload_failure() + self._launch_dump_state_rendezvous(global_step) + check_upload_error = False + else: + should_dump = self._local_dump_decision(global_step) + check_upload_error = True + if should_dump: dump_error: Optional[BaseException] = None try: - if self._requires_dump_state_rendezvous(): - # The uploader-failure bit was already rendezvoused in - # _should_dump; a rank-local recheck here could diverge. - self.dump(global_step, check_upload_error=False) - else: + if check_upload_error: self.dump(global_step) + else: + # The failure bit was already rendezvoused; a rank-local + # recheck inside dump() could diverge across ranks. + self.dump(global_step, check_upload_error=False) except BaseException as exc: dump_error = exc self._raise_if_any_interval_dump_failed(dump_error) @@ -722,6 +740,7 @@ def maybe_dump(self, global_step: int) -> None: # 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._timed_dump_in_flight = False self._tracker.step() def _requires_dump_state_rendezvous(self) -> bool: @@ -732,55 +751,101 @@ def _requires_dump_state_rendezvous(self) -> bool: the uploader-failure bit: rank zero observes an async upload failure immediately while peers only see the shared error marker through a throttled poll, so without the collective rank zero could raise alone - and strand its peers in the next training collective. + and strand its peers in the next training collective. The rendezvous + is pipelined across steps so it never blocks the training hot path. """ return getattr(self, "_world_size", 1) > 1 and ( self._interval_secs is not None or getattr(self, "_feature_store_enabled", False) ) - def _should_dump( - self, global_step: int, local_error: Optional[BaseException] = None - ) -> bool: - """Synchronize upload failure state and return one dump decision.""" + def _local_dump_decision(self, global_step: int) -> bool: + """Return the dump decision when no cross-rank rendezvous is needed.""" + upload_error = self._feature_store_upload_error() + if upload_error is not None: + raise upload_error.with_traceback(upload_error.__traceback__) if global_step <= 0: - 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 + 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" + ) + return self._rank == 0 and time.monotonic() >= self._next_dump_time - any_failed = local_error is not None - if self._requires_dump_state_rendezvous(): - if not ( - torch.distributed.is_available() and torch.distributed.is_initialized() - ): - raise RuntimeError( - "distributed delta embedding dump requires an initialized " - "process group" - ) - state = torch.tensor( - [int(should_dump), int(any_failed)], - dtype=torch.int32, - device=self._device, + def _launch_dump_state_rendezvous(self, global_step: int) -> None: + """Asynchronously all-reduce this step's dump vote and failure bit. + + The reduced tensor is consumed by the next ``maybe_dump`` call. NCCL + orders this collective ahead of the next training step's communication + on the same process group, so it has already completed by the time + the next call reads it back: the hot path pays an asynchronous launch + instead of a host-blocking all_reduce plus ``.item()`` pair on every + training step, and the decision lands at most one step late. + """ + if not ( + torch.distributed.is_available() and torch.distributed.is_initialized() + ): + raise RuntimeError( + "distributed delta embedding dump requires an initialized process group" + ) + local_error = self._feature_store_upload_error() + if local_error is not None: + self._pending_upload_error = local_error + state = torch.tensor( + [ + int(self._timed_dump_vote(global_step)), + int(self._pending_upload_error is not None), + ], + dtype=torch.int32, + device=self._device, + ) + torch.distributed.all_reduce(state, op=torch.distributed.ReduceOp.MAX) + self._pending_rendezvous = state + + def _consume_dump_state_rendezvous(self) -> Tuple[bool, bool]: + """Read back the previous step's reduced dump vote and failure bit.""" + state = self._pending_rendezvous + self._pending_rendezvous = None + if state is None: + return False, False + # Launched one training step ago and ordered ahead of this step's + # communication, so this read does not block the host. + should_dump, any_failed = state.tolist() + return bool(should_dump), bool(any_failed) + + def _timed_dump_vote(self, global_step: int) -> bool: + """Return rank zero's local vote for whether a timed dump is due. + + A fired vote arms ``_timed_dump_in_flight`` so the following step's + vote stays low until the dump executes and reschedules the timer from + its completion time; without the guard the still-elapsed deadline + would fire on the next launch and dump twice. + """ + if self._interval_secs is None or global_step <= 0: + return False + if self._next_dump_time is None: + raise RuntimeError( + "time-based delta embedding dumper must be started before training" ) - torch.distributed.all_reduce(state, op=torch.distributed.ReduceOp.MAX) - should_dump = bool(state[0].item()) - any_failed = bool(state[1].item()) + if self._timed_dump_in_flight: + return False + if self._rank == 0 and time.monotonic() >= self._next_dump_time: + self._timed_dump_in_flight = True + return True + return False + def _raise_rendezvoused_upload_failure(self) -> None: + """Raise the rendezvoused upload failure identically on every rank.""" + local_error = self._pending_upload_error 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" - ) - return should_dump + raise FeatureStoreUploadError( + "FeatureStore delta upload failed on another distributed worker; " + "all training workers are stopping and parquet outbox files were " + "retained" + ) def _raise_if_any_interval_dump_failed( self, local_error: Optional[BaseException] diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index e466209a2..07aaea7cf 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -740,37 +740,63 @@ 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_reduced_from_rank_zero(self): + def _new_rendezvous_dumper(self) -> DeltaEmbeddingDumper: dumper = object.__new__(DeltaEmbeddingDumper) + dumper._last_dump_step = None + dumper._device = torch.device("cpu") + dumper._tracker = mock.MagicMock() + dumper._pending_rendezvous = None + dumper._pending_upload_error = None + dumper._timed_dump_in_flight = False + return dumper + + def test_time_dump_decision_is_reduced_from_rank_zero(self): + # The rank-zero timer vote is launched on one step and consumed on the + # next: the pipelined rendezvous lands a timed dump at most one step + # after the timer fires without blocking any step on the collective. + dumper = self._new_rendezvous_dumper() 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") + launched_states = [] def fake_all_reduce(state, op=None): self.assertIs(op, torch.distributed.ReduceOp.MAX) - self.assertEqual(state.tolist(), [0, 0]) - state[0] = 1 + launched_states.append(state.tolist()) + if state.numel() == 2: + state[0] = 1 with ( + 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), ): - self.assertTrue(dumper._should_dump(10)) + dumper.maybe_dump(10) + dump_mock.assert_not_called() + dumper.maybe_dump(11) + + dump_mock.assert_called_once_with(11, check_upload_error=False) + self.assertEqual(dumper._last_dump_step, 11) + # Both steps launched [0, 0] locally; rank zero's vote arrived through + # the MAX reduce. The trailing [0] is the dump-failure rendezvous. + self.assertEqual(launched_states, [[0, 0], [0, 0], [0]]) + self.assertEqual(dumper._tracker.step.call_count, 2) def test_timed_maybe_dump_reduces_local_upload_error_before_raising(self): - dumper = object.__new__(DeltaEmbeddingDumper) + # Rank zero observes an async uploader failure immediately; the bit is + # launched on this step and consumed on the next so every rank raises + # together instead of rank zero raising alone. + dumper = self._new_rendezvous_dumper() 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 @@ -789,13 +815,14 @@ def fake_all_reduce(state, op=None): mock.patch("torch.distributed.is_initialized", return_value=True), mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), ): + dumper.maybe_dump(10) with self.assertRaises(FeatureStoreUploadError) as context: - dumper.maybe_dump(10) + dumper.maybe_dump(11) self.assertTrue(collective_called) self.assertIs(context.exception, local_error) dump_mock.assert_not_called() - dumper._tracker.step.assert_not_called() + self.assertEqual(dumper._tracker.step.call_count, 1) def test_timed_maybe_dump_surfaces_remote_upload_error_when_marker_is_throttled( self, @@ -804,17 +831,16 @@ def test_timed_maybe_dump_surfaces_remote_upload_error_when_marker_is_throttled( marker_path = os.path.join(tmp_dir, "last_error.json") with open(marker_path, "w") as output: output.write("{}") - dumper = object.__new__(DeltaEmbeddingDumper) + dumper = self._new_rendezvous_dumper() 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 + # The throttled poll hides the marker locally; only rank zero's + # rendezvoused failure bit surfaces the error, one step later. dumper._next_feature_store_error_check = 100.0 dumper._feature_store_error_check_interval_secs = 1 dumper._uploader = None @@ -822,7 +848,6 @@ def test_timed_maybe_dump_surfaces_remote_upload_error_when_marker_is_throttled( 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 ( @@ -835,24 +860,22 @@ def fake_all_reduce(state, op=None): mock.patch("torch.distributed.is_initialized", return_value=True), mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), ): + dumper.maybe_dump(10) with self.assertRaisesRegex( FeatureStoreUploadError, "another distributed worker" ): - dumper.maybe_dump(10) + dumper.maybe_dump(11) dump_mock.assert_not_called() - dumper._tracker.step.assert_not_called() + self.assertEqual(dumper._tracker.step.call_count, 1) def test_timed_maybe_dump_advances_after_all_workers_succeed(self): - dumper = object.__new__(DeltaEmbeddingDumper) + dumper = self._new_rendezvous_dumper() 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): @@ -866,41 +889,36 @@ def fake_all_reduce(state, op=None): ) as dump_mock, mock.patch( "tzrec.utils.delta_embedding_dump.time.monotonic", - side_effect=[1.0, 2.0], + side_effect=[1.0, 2.0, 3.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) + dump_mock.assert_not_called() + dumper.maybe_dump(11) - 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(collective_states, [[1, 0], [0, 0], [0]]) + dump_mock.assert_called_once_with(11, check_upload_error=False) + self.assertEqual(dumper._last_dump_step, 11) self.assertEqual(dumper._next_dump_time, 62.0) - dumper._tracker.step.assert_called_once_with() + self.assertFalse(dumper._timed_dump_in_flight) + self.assertEqual(dumper._tracker.step.call_count, 2) def test_timed_maybe_dump_reduces_local_dump_failure_before_raising(self): - dumper = object.__new__(DeltaEmbeddingDumper) + dumper = self._new_rendezvous_dumper() 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 + collective_states = [] 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 + collective_states.append(state.tolist()) with ( mock.patch.object(dumper, "_feature_store_upload_error", return_value=None), @@ -909,36 +927,30 @@ def fake_all_reduce(state, op=None): mock.patch("torch.distributed.is_initialized", return_value=True), mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), ): + dumper.maybe_dump(10) with self.assertRaises(RuntimeError) as context: - dumper.maybe_dump(10) + dumper.maybe_dump(11) - self.assertEqual(collective_index, 2) + self.assertEqual(collective_states, [[1, 0], [0, 0], [1]]) self.assertIs(context.exception, dump_error) - dump_mock.assert_called_once_with(10, check_upload_error=False) + dump_mock.assert_called_once_with(11, check_upload_error=False) self.assertIsNone(dumper._last_dump_step) - dumper._tracker.step.assert_not_called() + self.assertEqual(dumper._tracker.step.call_count, 1) def test_timed_maybe_dump_surfaces_remote_dump_failure(self): - dumper = object.__new__(DeltaEmbeddingDumper) + dumper = self._new_rendezvous_dumper() 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 + collective_states = [] 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]) + collective_states.append(state.tolist()) + if state.numel() == 1: state[0] = 1 - collective_index += 1 with ( mock.patch.object(dumper, "_feature_store_upload_error", return_value=None), @@ -949,26 +961,63 @@ def fake_all_reduce(state, op=None): mock.patch("torch.distributed.is_initialized", return_value=True), mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), ): + dumper.maybe_dump(10) with self.assertRaisesRegex(RuntimeError, "another distributed worker"): - dumper.maybe_dump(10) + dumper.maybe_dump(11) - self.assertEqual(collective_index, 2) - dump_mock.assert_called_once_with(10, check_upload_error=False) + self.assertEqual(collective_states, [[1, 0], [0, 0], [0]]) + dump_mock.assert_called_once_with(11, check_upload_error=False) self.assertIsNone(dumper._last_dump_step) - dumper._tracker.step.assert_not_called() + self.assertEqual(dumper._tracker.step.call_count, 1) + + def test_timed_dump_vote_does_not_fire_twice_before_completion(self): + # Once rank zero's timer vote fires, the next step's vote stays low + # until the consumed dump reschedules from its completion time; the + # still-elapsed deadline must not launch a second dump. + dumper = self._new_rendezvous_dumper() + dumper._interval_steps = None + dumper._interval_secs = 60.0 + dumper._next_dump_time = 0.0 + dumper._rank = 0 + dumper._world_size = 2 + 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=[100.0, 100.0, 100.0, 100.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) + dumper.maybe_dump(11) + dumper.maybe_dump(12) + + dump_mock.assert_called_once_with(11, check_upload_error=False) + self.assertEqual(dumper._next_dump_time, 160.0) + self.assertFalse(dumper._timed_dump_in_flight) + self.assertEqual(collective_states, [[1, 0], [0, 0], [0], [0, 0]]) + self.assertEqual(dumper._tracker.step.call_count, 3) def test_step_maybe_dump_reduces_local_upload_error_before_raising(self): - # Rank zero observes an async uploader failure immediately; without a - # step-cadence rendezvous it would raise alone and strand peers that - # skipped the throttled marker poll in the next training collective. - dumper = object.__new__(DeltaEmbeddingDumper) + # Rank zero observes an async uploader failure immediately; the + # pipelined rendezvous launches the bit on this step and raises on the + # next so every rank stops together instead of rank zero raising alone. + dumper = self._new_rendezvous_dumper() dumper._interval_steps = 50 dumper._interval_secs = None - dumper._last_dump_step = None dumper._rank = 0 dumper._world_size = 2 - dumper._device = torch.device("cpu") - dumper._tracker = mock.MagicMock() dumper._feature_store_enabled = True local_error = FeatureStoreUploadError("local upload failed") collective_called = False @@ -988,13 +1037,14 @@ def fake_all_reduce(state, op=None): mock.patch("torch.distributed.is_initialized", return_value=True), mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), ): + dumper.maybe_dump(10) with self.assertRaises(FeatureStoreUploadError) as context: - dumper.maybe_dump(10) + dumper.maybe_dump(11) self.assertTrue(collective_called) self.assertIs(context.exception, local_error) dump_mock.assert_not_called() - dumper._tracker.step.assert_not_called() + self.assertEqual(dumper._tracker.step.call_count, 1) def test_step_maybe_dump_surfaces_remote_upload_error_when_marker_is_throttled( self, @@ -1003,14 +1053,11 @@ def test_step_maybe_dump_surfaces_remote_upload_error_when_marker_is_throttled( marker_path = os.path.join(tmp_dir, "last_error.json") with open(marker_path, "w") as output: output.write("{}") - dumper = object.__new__(DeltaEmbeddingDumper) + dumper = self._new_rendezvous_dumper() dumper._interval_steps = 50 dumper._interval_secs = None - 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 # The throttled poll hides the marker on this rank; only the @@ -1034,23 +1081,23 @@ def fake_all_reduce(state, op=None): mock.patch("torch.distributed.is_initialized", return_value=True), mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), ): + dumper.maybe_dump(10) with self.assertRaisesRegex( FeatureStoreUploadError, "another distributed worker" ): - dumper.maybe_dump(10) + dumper.maybe_dump(11) dump_mock.assert_not_called() - dumper._tracker.step.assert_not_called() + self.assertEqual(dumper._tracker.step.call_count, 1) def test_step_feature_store_boundary_dump_skips_rank_local_upload_recheck(self): - dumper = object.__new__(DeltaEmbeddingDumper) + # Only the failure bit travels through the pipelined rendezvous; the + # boundary decision stays deterministic and on its exact boundary step. + dumper = self._new_rendezvous_dumper() dumper._interval_steps = 50 dumper._interval_secs = None - 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 collective_states = [] @@ -1067,12 +1114,14 @@ def fake_all_reduce(state, op=None): mock.patch("torch.distributed.is_initialized", return_value=True), mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), ): + dumper.maybe_dump(49) + dump_mock.assert_not_called() dumper.maybe_dump(50) - self.assertEqual(collective_states, [[1, 0], [0]]) + self.assertEqual(collective_states, [[0, 0], [0, 0], [0]]) dump_mock.assert_called_once_with(50, check_upload_error=False) self.assertEqual(dumper._last_dump_step, 50) - dumper._tracker.step.assert_called_once_with() + self.assertEqual(dumper._tracker.step.call_count, 2) def test_step_boundary_dump_failure_is_reduced_before_raising(self): # Even without FeatureStore, a rank-local boundary dump failure (e.g. a From c22051d116a9ba3e63a1256c97dc6e9624a01922 Mon Sep 17 00:00:00 2001 From: gecheng Date: Tue, 21 Jul 2026 17:31:42 +0800 Subject: [PATCH 17/50] [perf] cap rank-zero delta upload memory with streamed reads and spilled dedup _load_records read each shard with pq.read_table() and to_pylist(), materializing every embedding element as Python objects, then kept the deduplicated rows plus a sorted copy in memory for the whole upload. Peak memory therefore scaled with a single dump's size, so one large delta could exhaust rank-zero training memory despite max_pending_steps bounding only the number of pending steps, not their size. Shards are now streamed with ParquetFile.iter_batches and validated with vectorized per-batch checks, and dedup state lives in a throwaway SQLite database (bounded 64MB page cache, journal off: the store is disposable because a crashed step replays from its durable shard snapshot). Uploads stream from an ORDER BY cursor in upload_batch_size windows, and retries re-iterate the same deterministic ordering, so peak memory depends only on one Parquet batch, the page cache, and one in-flight upload batch. Co-Authored-By: Claude --- tzrec/utils/feature_store_delta_uploader.py | 365 +++++++++++++----- .../feature_store_delta_uploader_test.py | 118 +++++- 2 files changed, 379 insertions(+), 104 deletions(-) diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index ac6b8d18c..eba866f42 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -17,6 +17,7 @@ import os import re import shutil +import sqlite3 import threading import time import uuid @@ -28,10 +29,12 @@ Callable, Dict, Iterable, + Iterator, List, Mapping, Optional, Tuple, + cast, ) from urllib.parse import urlsplit @@ -56,6 +59,8 @@ DELTA_DUMP_GENERATION_METADATA_KEY = b"tzrec.delta_embedding.dump_generation" _FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES = 100 +_RECORD_STORE_DB_FILENAME = "upload_records.sqlite" +_RECORD_STORE_CACHE_SIZE_KB = 65536 _VERSION_INITIALIZATION = "AUTO_CREATE_ON_FIRST_DELTA_MERGE" _LEGACY_VERSION_INITIALIZATION = "PREPROVISIONED_FOR_DELTA_MERGE" @@ -366,6 +371,137 @@ def _file_sha256(source: BinaryIO) -> str: return digest.hexdigest() +class _RecordStore: + """Deduplicated delta records spilled to a throwaway SQLite database. + + One delta dump can contain an arbitrary number of unique keys, so keeping + the dedup state and the upload ordering in memory would make rank-zero + peak memory scale with a single dump's size. The store keeps both on disk + behind a bounded page cache, so memory depends only on the streaming + Parquet read batch and one in-flight upload batch. The database is + disposable: after a crash the whole step replays from its durable shard + snapshot. + """ + + def __init__(self, db_path: str) -> None: + self._db_path = db_path + if os.path.exists(db_path): + os.remove(db_path) + self._connection = sqlite3.connect(db_path) + self._connection.execute("PRAGMA journal_mode = OFF") + self._connection.execute("PRAGMA synchronous = OFF") + self._connection.execute(f"PRAGMA cache_size = -{_RECORD_STORE_CACHE_SIZE_KB}") + self._connection.execute( + "CREATE TABLE dedup (" + "embedding_name TEXT NOT NULL, " + "key_id INTEGER NOT NULL, " + "embedding BLOB NOT NULL, " + "PRIMARY KEY (embedding_name, key_id))" + ) + self._record_count: Optional[int] = None + + def __enter__(self) -> "_RecordStore": + return self + + def __exit__(self, *exc_info: Any) -> None: + self.close() + + @property + def record_count(self) -> int: + """Return the number of unique deduplicated records.""" + if self._record_count is None: + cursor = self._connection.execute("SELECT COUNT(*) FROM dedup") + self._record_count = int(cursor.fetchone()[0]) + return self._record_count + + def add_batch( + self, + embedding_names: List[str], + key_ids: np.ndarray, + flat_embeddings: np.ndarray, + offsets: np.ndarray, + ) -> None: + """Spill one streamed shard batch, deduplicating on (name, key_id).""" + rows = [ + ( + embedding_names[index], + int(key_ids[index]), + flat_embeddings[ + int(offsets[index]) : int(offsets[index + 1]) + ].tobytes(), + ) + for index in range(len(embedding_names)) + ] + try: + self._connection.executemany("INSERT INTO dedup VALUES (?, ?, ?)", rows) + except sqlite3.IntegrityError: + # A duplicate can be exact (harmless) or conflicting (corruption); + # replay the batch row by row to tell the two apart. + for row in rows: + try: + self._connection.execute("INSERT INTO dedup VALUES (?, ?, ?)", row) + except sqlite3.IntegrityError: + existing = self._connection.execute( + "SELECT embedding FROM dedup " + "WHERE embedding_name = ? AND key_id = ?", + (row[0], row[1]), + ).fetchone() + if existing[0] != row[2]: + raise ValueError( + "conflicting duplicate delta row for " + f"embedding_name={row[0]!r}, key_id={row[1]}" + ) from None + self._connection.commit() + self._record_count = None + + def iter_upload_batches( + self, batch_size: int + ) -> Iterator[Tuple[List[Dict[str, Any]], int]]: + """Yield (payload, record count) upload batches in a deterministic order.""" + cursor = self._connection.execute( + "SELECT embedding_name, key_id, embedding FROM dedup " + "ORDER BY embedding_name, key_id" + ) + while True: + rows = cursor.fetchmany(batch_size) + if not rows: + return + payload = [ + { + FEATURE_STORE_PK_FIELD: embedding_name, + FEATURE_STORE_SK_FIELD: key_id, + FEATURE_STORE_VALUE_FIELD: np.frombuffer( + embedding, dtype=np.float32 + ).copy(), + } + for embedding_name, key_id, embedding in rows + ] + yield payload, len(payload) + + def to_records(self) -> List[Tuple[str, int, np.ndarray]]: + """Materialize every deduplicated record in upload order.""" + cursor = self._connection.execute( + "SELECT embedding_name, key_id, embedding FROM dedup " + "ORDER BY embedding_name, key_id" + ) + return [ + ( + embedding_name, + int(key_id), + np.frombuffer(embedding, dtype=np.float32).copy(), + ) + for embedding_name, key_id, embedding in cursor + ] + + def close(self) -> None: + """Close the connection and remove the throwaway database file.""" + try: + self._connection.close() + finally: + if os.path.exists(self._db_path): + os.remove(self._db_path) + + class FeatureStoreDeltaUploader: """Publish complete delta-dump steps from a persistent local outbox. @@ -677,16 +813,20 @@ def _run(self) -> None: snapshot_paths, shard_sources ) self._validate_dump_generation(shard_descriptions) - records = self._load_records( - current_step, snapshot_paths, shard_sources + record_store = stack.enter_context( + self._build_record_store( + current_step, snapshot_paths, shard_sources + ) ) manifest = self._load_or_create_manifest( current_step, snapshot_paths, - len(records), + record_store.record_count, shard_descriptions=shard_descriptions, ) - summary = self._upload_with_retries(current_step, records, manifest) + summary = self._upload_with_retries( + current_step, record_store, 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. @@ -1437,15 +1577,15 @@ def _start_attempt( def _upload_with_retries( self, global_step: int, - records: List[Tuple[str, int, np.ndarray]], + store: _RecordStore, 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) + attempt = self._start_attempt(global_step, store.record_count, manifest) try: - if records: - summary = self._upload_records(global_step, records, attempt) + if store.record_count: + summary = self._upload_records(global_step, store, attempt) else: # An empty step still validates the FeatureView schema before # it can be committed. @@ -1787,13 +1927,14 @@ def _reset_view(self, suppress_errors: bool = False) -> None: def _upload_records( self, global_step: int, - records: List[Tuple[str, int, np.ndarray]], + store: _RecordStore, attempt: Mapping[str, int], ) -> Dict[str, int]: view = self._get_view() max_in_flight = int(getattr(view, "_max_workers", 1)) + total_records = store.record_count batch_count = ( - len(records) + self._settings.upload_batch_size - 1 + total_records + self._settings.upload_batch_size - 1 ) // self._settings.upload_batch_size aggregate = { "total_batches": 0, @@ -1810,46 +1951,42 @@ def _upload_records( global_step, attempt["attempt_id"], self._settings.version, - len(records), + total_records, batch_count, max_in_flight, attempt["range_start"], attempt["range_end"], ) try: - 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) + completed_batches = 0 + window_batches = 0 + window_records = 0 + logged_first_window = False + for payload, payload_records in store.iter_upload_batches( + self._settings.upload_batch_size + ): + self._raise_if_aborting() + view.write_features( + data=payload, + version=self._settings.version, + write_mode=FEATURE_STORE_WRITE_MODE, + ts=int(attempt["range_start"]) + completed_batches, + ) + completed_batches += 1 + window_batches += 1 + window_records += payload_records + if window_batches < max_in_flight and completed_batches < batch_count: + continue summary = view.write_flush() self._validate_flush_summary( summary, expected_records=window_records, - expected_batches=window_end - window_start, + expected_batches=window_batches, ) for name in aggregate: aggregate[name] += int(summary[name]) - completed_batches = window_end if ( - window_start == 0 + not logged_first_window or completed_batches >= next_progress_batch or completed_batches == batch_count ): @@ -1861,13 +1998,16 @@ def _upload_records( completed_batches, batch_count, aggregate["success_records"], - len(records), + total_records, time.monotonic() - started_at, ) + logged_first_window = True while next_progress_batch <= completed_batches: next_progress_batch += ( _FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES ) + window_batches = 0 + window_records = 0 self._raise_if_aborting() except BaseException: # A write_features() call can enqueue part of its work before raising. @@ -1901,79 +2041,110 @@ def _validate_flush_summary( ): raise RuntimeError("FeatureStore write_flush reported incomplete writes") - def _load_records( + def _build_record_store( self, global_step: int, shard_paths: List[str], shard_sources: Optional[List[BinaryIO]] = None, - ) -> List[Tuple[str, int, np.ndarray]]: + ) -> _RecordStore: + """Stream the step's shards into a disk-spilled deduplicated store. + + Shards are read batch by batch and validated with vectorized checks, + so peak memory stays bounded by one Parquet batch and the store's + page cache, independently of one dump's size. + """ 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) + return self._build_record_store( + 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["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 + store = _RecordStore(os.path.join(self._state_dir, _RECORD_STORE_DB_FILENAME)) + try: + for expected_rank, (path, source) in enumerate( + zip(shard_paths, shard_sources) + ): + source.seek(0) + parquet_file = pq.ParquetFile(source) + self._validate_parquet_schema(parquet_file.schema_arrow, path) + for batch in parquet_file.iter_batches( + batch_size=self._settings.upload_batch_size + ): + self._ingest_record_batch(store, global_step, expected_rank, batch) + source.seek(0) + except BaseException: + store.close() + raise + return store - return [ - (embedding_name, key_id, records[(embedding_name, key_id)]) - for embedding_name, key_id in sorted(records) - ] + def _ingest_record_batch( + self, + store: _RecordStore, + global_step: int, + expected_rank: int, + batch: pa.RecordBatch, + ) -> None: + """Validate one streamed shard batch and spill it into the store.""" + num_rows = batch.num_rows + if num_rows == 0: + return + # The parquet schema check already enforces every column type, so the + # value checks below stay vectorized instead of per-row Python loops. + global_steps = batch.column("global_step").to_numpy(zero_copy_only=False) + if not bool((global_steps == global_step).all()): + raise ValueError("delta shard global_step mismatch") + ranks = batch.column("rank").to_numpy(zero_copy_only=False) + if not bool((ranks == expected_rank).all()): + raise ValueError("delta shard rank mismatch") + world_sizes = batch.column("world_size").to_numpy(zero_copy_only=False) + if not bool((world_sizes == self._world_size).all()): + raise ValueError("delta shard world_size mismatch") + + batch_roles = set(batch.column("embedding_role").to_pylist()) + if not batch_roles <= set(SPARSE_EMBEDDING_ROLES): + raise ValueError("delta shard has an invalid embedding role") + embedding_names = batch.column("embedding_name").to_pylist() + for embedding_name in set(embedding_names): + if 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_ids = batch.column("key_id").to_numpy(zero_copy_only=False) + if bool((key_ids == SPARSE_EMBEDDING_INVALID_KEY).any()): + raise ValueError( + "delta shard key_id=-1 is reserved as the Processor/" + "NvEmbeddings invalid-key sentinel" + ) + + embedding_column = cast(pa.ListArray, batch.column("embedding")) + flat_embeddings = embedding_column.values.to_numpy(zero_copy_only=False) + if not bool(np.isfinite(flat_embeddings).all()): + raise ValueError("delta embedding contains NaN or Inf") + offsets = embedding_column.offsets.to_numpy() + lengths = np.diff(offsets) + expected_dims = np.array( + [self._embedding_dimensions[name] for name in embedding_names], + dtype=lengths.dtype, + ) + bad_rows = np.flatnonzero(lengths != expected_dims) + if bad_rows.size > 0: + row = int(bad_rows[0]) + raise ValueError( + f"delta embedding dimension mismatch for {embedding_names[row]!r}: " + f"expected={int(expected_dims[row])}, " + f"actual_shape=({int(lengths[row])},)" + ) + + store.add_batch(embedding_names, key_ids, flat_embeddings, offsets) @staticmethod def _validate_parquet_schema(schema: pa.Schema, path: str) -> None: diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index 8403e0539..2dabd4c1b 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -24,6 +24,7 @@ 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 ( + _RECORD_STORE_DB_FILENAME, DELTA_DUMP_GENERATION_METADATA_KEY, FeatureStoreDeltaUploader, FeatureStoreUploadError, @@ -1028,7 +1029,13 @@ def test_exact_duplicate_key_is_deduplicated_before_async_sdk_batches(self): {"user_emb": 2}, client_factory=_FakeClientFactory(_FakeView()), ) - records = uploader._load_records(10, [shard]) + with uploader._build_record_store(10, [shard]) as store: + records = store.to_records() + self.assertFalse( + os.path.exists( + os.path.join(uploader.state_dir, _RECORD_STORE_DB_FILENAME) + ) + ) self.assertEqual(len(records), 1) self.assertEqual(records[0][2].tolist(), [1.0, 2.0]) @@ -1051,7 +1058,102 @@ def test_conflicting_duplicate_key_is_rejected(self): client_factory=_FakeClientFactory(_FakeView()), ) with self.assertRaisesRegex(ValueError, "conflicting duplicate"): - uploader._load_records(10, [shard]) + uploader._build_record_store(10, [shard]) + # The throwaway spill database must not outlive a failed ingest. + self.assertFalse( + os.path.exists( + os.path.join(uploader.state_dir, _RECORD_STORE_DB_FILENAME) + ) + ) + + def test_duplicate_key_across_streamed_batches_is_deduplicated(self): + # Shards stream in upload_batch_size=2 record batches; the duplicate + # arrives in a later ingest batch than its first occurrence, so dedup + # must hold across spilled batches, not just within one. + with tempfile.TemporaryDirectory() as output_dir: + shard = _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, 1, [1.0, 2.0]), + _row(10, 0, 3, [5.0, 6.0]), + ], + ) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with uploader._build_record_store(10, [shard]) as store: + records = store.to_records() + + self.assertEqual([record[1] for record in records], [1, 2, 3]) + self.assertEqual(records[0][2].tolist(), [1.0, 2.0]) + + def test_conflicting_duplicate_across_streamed_batches_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, 2, [3.0, 4.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._build_record_store(10, [shard]) + + def test_record_store_reiteration_is_deterministic_across_attempts(self): + # A retried attempt re-streams the whole store; the spill ordering + # must be identical every time so replayed ts ranges carry identical + # batch boundaries. + with tempfile.TemporaryDirectory() as output_dir: + shard = _write_single_shard( + output_dir, + 10, + [ + _row(10, 0, 3, [5.0, 6.0]), + _row(10, 0, 1, [1.0, 2.0]), + _row(10, 0, 2, [3.0, 4.0]), + ], + ) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with uploader._build_record_store(10, [shard]) as store: + first = list(store.iter_upload_batches(2)) + second = list(store.iter_upload_batches(2)) + + self.assertEqual([count for _, count in first], [2, 1]) + + def keys_and_values(batches): + return [ + (item["key_id"], item["embedding"].tolist()) + for payload, _ in batches + for item in payload + ] + + self.assertEqual(keys_and_values(first), keys_and_values(second)) + self.assertEqual([key for key, _ in keys_and_values(first)], [1, 2, 3]) def test_flush_failure_does_not_publish_success_marker(self): failed_summary = { @@ -1417,7 +1519,8 @@ def test_signed_int64_key_is_preserved(self): {"user_emb": 2}, client_factory=_FakeClientFactory(_FakeView()), ) - records = uploader._load_records(10, [shard]) + with uploader._build_record_store(10, [shard]) as store: + records = store.to_records() self.assertEqual(records[0][1], -2) def test_reserved_invalid_key_is_rejected(self): @@ -1432,7 +1535,7 @@ def test_reserved_invalid_key_is_rejected(self): client_factory=_FakeClientFactory(_FakeView()), ) with self.assertRaisesRegex(ValueError, "invalid-key sentinel"): - uploader._load_records(10, [shard]) + uploader._build_record_store(10, [shard]) def test_empty_step_validates_remote_contract_before_commit(self): with tempfile.TemporaryDirectory() as output_dir: @@ -1510,13 +1613,13 @@ def test_dimension_and_finite_value_validation(self): client_factory=_FakeClientFactory(_FakeView()), ) with self.assertRaisesRegex(ValueError, "dimension mismatch"): - uploader._load_records(10, [bad_dimension]) + uploader._build_record_store(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]) + uploader._build_record_store(11, [bad_finite]) def test_persisted_manifest_rejects_changed_shard_content(self): with tempfile.TemporaryDirectory() as output_dir: @@ -1558,7 +1661,8 @@ def test_opened_shard_snapshot_rejects_atomic_path_replacement(self): ) os.replace(replacement, shard) - records = uploader._load_records(10, [shard], [source]) + with uploader._build_record_store(10, [shard], [source]) as store: + records = store.to_records() self.assertEqual(records[0][2].tolist(), [1.0, 2.0]) with self.assertRaisesRegex( RuntimeError, "changed after upload snapshot" From 36cad98647a621397f41605e36749186e518936a Mon Sep 17 00:00:00 2001 From: gecheng Date: Tue, 21 Jul 2026 17:53:49 +0800 Subject: [PATCH 18/50] [perf] reconcile committed delta steps once instead of on every poll The uploader worker called _add_discovered_steps_locked() on every poll while holding the condition lock, and each call rescanned all historical steps: because committed manifests and success markers are retained, every poll re-read the success/manifest JSON of every committed step and re-opened its canonical shard Parquet schemas (world_size footer reads per step), plus stat storms over the retained step directories. Lock hold time therefore grew with the job's committed history, and submit()/close()/check_error() on the training thread stalled behind it. Committed steps are now reconciled at most once per process: a reconciled set records steps verified committed with a matching generation (and steps committed by this run, marked at commit time), and discovery skips their JSON/Parquet re-reads and per-entry stat checks entirely. Startup still reconciles from scratch with an empty set, so crash recovery and the committed-generation replacement alarm are unchanged; new dumps enter the queue through submit(). Each poll now costs a few directory listings regardless of history. Co-Authored-By: Claude --- tzrec/utils/feature_store_delta_uploader.py | 45 +++++- .../feature_store_delta_uploader_test.py | 151 ++++++++++++++++++ 2 files changed, 191 insertions(+), 5 deletions(-) diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index eba866f42..067943ca4 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -33,6 +33,7 @@ List, Mapping, Optional, + Set, Tuple, cast, ) @@ -550,6 +551,10 @@ def __init__( self._view = None self._condition = threading.Condition() self._pending: Dict[int, float] = {} + # Steps already reconciled as committed with a matching generation; + # discovery skips their JSON/Parquet re-reads so poll lock hold time + # does not grow with the job's committed history. + self._reconciled_steps: Set[int] = set() self._started = False self._closing = False self._aborting = False @@ -648,6 +653,7 @@ def start(self) -> None: # A later retry must reconcile again after reacquiring the lock; # another process may have advanced the journal in between. self._journal_initialized = False + self._reconciled_steps = set() self._reset_view(suppress_errors=True) self._release_writer_lock() raise @@ -991,6 +997,16 @@ def _raise_if_aborting(self) -> None: raise _UploadAborted() def _add_discovered_steps_locked(self) -> None: + """Incrementally enqueue new outbox steps without rescanning history. + + Full reconciliation happens once in ``start()`` with an empty + reconciled set. Afterwards, steps already pending or already verified + committed with a matching generation are skipped without any JSON or + Parquet reads, so the condition lock hold time stays independent of + how many steps this job has committed. A committed step whose + canonical generation no longer matches its manifest is deliberately + never marked reconciled, keeping the replacement alarm alive. + """ remaining_capacity = self._settings.max_pending_steps - len(self._pending) if remaining_capacity <= 0: return @@ -1002,17 +1018,21 @@ def _add_discovered_steps_locked(self) -> None: "found invalid FeatureStore delta outbox global_step; " "global_step must be > 0" ) + if step in self._reconciled_steps or step in self._pending: + continue if self._is_committed(step): try: canonical_generation = self._canonical_dump_generation(step) except _ShardSetNotReady: canonical_generation = "" if canonical_generation is None: + self._reconciled_steps.add(step) continue manifest_generation = _read_json(self._manifest_path(step)).get( "dump_generation" ) if canonical_generation == manifest_generation: + self._reconciled_steps.add(step) continue if step not in self._pending: self._pending[step] = time.monotonic() @@ -1022,6 +1042,10 @@ def _discover_steps(self) -> Iterable[int]: if not os.path.isdir(self._output_dir): return [] steps = set() + # Already pending or reconciled steps are known to be real entries; + # re-statting their directories/shards on every poll would make the + # scan cost grow with the retained outbox history. + known_steps = self._reconciled_steps | set(self._pending) manifest_pattern = re.compile(r"^step_(\d+)\.manifest\.json$") for name in os.listdir(self._state_dir): match = manifest_pattern.match(name) @@ -1031,8 +1055,13 @@ def _discover_steps(self) -> Iterable[int]: 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 match is None: + continue + step = int(match.group(1)) + if step in known_steps or os.path.isdir( + os.path.join(self._snapshot_root, name) + ): + steps.add(step) if self._world_size == 1: pattern = re.compile( rf"^{re.escape(self._file_prefix)}_step_(\d+)\.parquet$" @@ -1045,11 +1074,14 @@ def _discover_steps(self) -> Iterable[int]: 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) - ): + if match is None: continue step = int(match.group(1)) + if step in known_steps: + steps.add(step) + continue + if not os.path.isdir(os.path.join(self._output_dir, name)): + continue # 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. @@ -2200,6 +2232,9 @@ def _commit_success( } _atomic_write_json(os.path.join(self._state_dir, "committed.json"), committed) self._committed_global_step = global_step + # Committed under the condition lock: later polls never re-read this + # step's success marker, manifest, or canonical shard schemas again. + self._reconciled_steps.add(global_step) logger.info( "FeatureStore delta upload committed: step=%s version=%s records=%s ts=%s", diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index 2dabd4c1b..2727f0f95 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -1601,6 +1601,157 @@ def test_multi_rank_discovery_uses_exact_contract_shards(self): self.assertEqual(len(view.calls), 1) self.assertEqual([item["key_id"] for item in view.calls[0]["data"]], [1, 2]) + def test_commit_marks_step_reconciled_for_later_polls(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()), + ) + uploader.start() + uploader.submit(10) + uploader.close() + + self.assertIn(10, uploader._reconciled_steps) + load_calls = 0 + original_load = uploader._load_success_state + + def counting_load(step): + nonlocal load_calls + load_calls += 1 + return original_load(step) + + with mock.patch.object( + uploader, "_load_success_state", side_effect=counting_load + ): + for _ in range(3): + with uploader._condition: + uploader._add_discovered_steps_locked() + + # The committed step was reconciled at commit time, so repeated + # polls never re-read its success marker or manifest again. + self.assertEqual(load_calls, 0) + self.assertEqual(uploader._pending, {}) + + def test_inherited_committed_step_is_reconciled_exactly_once(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + first = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + first.start() + first.submit(10) + first.close() + + # A fresh process starts with no reconciled history. + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + load_calls = 0 + generation_calls = 0 + original_load = uploader._load_success_state + original_generation = uploader._canonical_dump_generation + + def counting_load(step): + nonlocal load_calls + load_calls += 1 + return original_load(step) + + def counting_generation(step): + nonlocal generation_calls + generation_calls += 1 + return original_generation(step) + + with ( + mock.patch.object( + uploader, "_load_success_state", side_effect=counting_load + ), + mock.patch.object( + uploader, + "_canonical_dump_generation", + side_effect=counting_generation, + ), + ): + for _ in range(3): + with uploader._condition: + uploader._add_discovered_steps_locked() + + self.assertEqual(load_calls, 1) + self.assertEqual(generation_calls, 1) + self.assertIn(10, uploader._reconciled_steps) + self.assertEqual(uploader._pending, {}) + + def test_committed_step_with_replaced_generation_stays_unreconciled(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + first = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + first.start() + first.submit(10) + first.close() + + # Replace the committed step's canonical shard with a dump from a + # newer generation; discovery must keep checking it (the worker's + # replacement alarm relies on that) instead of reconciling it away. + _write_single_shard( + output_dir, + 10, + [_row(10, 0, 1, [7.0, 8.0])], + generation="ffffffffffffffffffffffffffffffff", + ) + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with uploader._condition: + uploader._add_discovered_steps_locked() + + self.assertNotIn(10, uploader._reconciled_steps) + self.assertIn(10, uploader._pending) + + def test_discovery_finds_new_outbox_entries_incrementally(self): + with tempfile.TemporaryDirectory() as output_dir: + uploader = FeatureStoreDeltaUploader( + _feature_store_config(), + output_dir, + "delta", + 1, + {"user_emb": 2}, + client_factory=_FakeClientFactory(_FakeView()), + ) + with uploader._condition: + uploader._add_discovered_steps_locked() + self.assertEqual(uploader._pending, {}) + + _write_single_shard(output_dir, 5, [_row(5, 0, 1, [1.0, 2.0])]) + with uploader._condition: + uploader._add_discovered_steps_locked() + self.assertIn(5, uploader._pending) + 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])]) From a2d89673bfae7a23d26b1e47f6f1f6a969be50ce Mon Sep 17 00:00:00 2001 From: gecheng Date: Tue, 21 Jul 2026 18:29:58 +0800 Subject: [PATCH 19/50] [bugfix] restrict feature_store_config.endpoint to trusted HTTPS hosts The endpoint override was passed to the FeatureStore SDK together with the access keys, STS token and FeatureDB credentials, but validation only rejected URI userinfo: any plaintext or attacker-controlled host in a pipeline config (configs are routinely checked into repos and shipped as templates) would receive all of those credentials, and an explicit http:// scheme would send them in cleartext. Endpoint validation now requires an https:// scheme when a scheme is given, rejects userinfo, path, query and fragment components, and only accepts trusted *.aliyuncs.com hosts without an explicit port. A new optional allow_custom_endpoint proto field is an explicit opt-in for vetted private deployments; it waives the trusted-host and port rules only, while every structural credential-safety check still applies. Co-Authored-By: Claude --- tzrec/protos/train.proto | 11 +++- tzrec/utils/feature_store_delta_uploader.py | 55 +++++++++++++++-- .../feature_store_delta_uploader_test.py | 61 ++++++++++++++++++- 3 files changed, 119 insertions(+), 8 deletions(-) diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index 7e2785ec7..957eec751 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -51,7 +51,11 @@ message FeatureStoreConfig { required string version = 9; - // Optional FeatureStore control-plane endpoint override. + // Optional FeatureStore control-plane endpoint override. Cloud and + // FeatureDB credentials are sent to this host, so it must be a trusted + // *.aliyuncs.com FeatureStore endpoint (an explicit https:// scheme is + // required when a scheme is given) without userinfo, port, path, query or + // fragment components, unless allow_custom_endpoint is set. optional string endpoint = 10; optional string security_token = 11; // Maximum records submitted before each SDK write_flush() completion gate. @@ -75,6 +79,11 @@ message FeatureStoreConfig { 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]; + // Explicit opt-in allowing endpoint to name a non-*.aliyuncs.com host + // (e.g. a vetted private FeatureStore deployment with a custom port). + // Credentials are sent to that host; structural endpoint checks (HTTPS + // scheme, no userinfo/path/query/fragment) still apply. + optional bool allow_custom_endpoint = 22 [default = false]; } message DeltaEmbeddingDumpConfig { diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index 067943ca4..01a88e687 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -118,6 +118,7 @@ class FeatureStoreUploadSettings: shutdown_timeout_secs: int max_pending_steps: int poll_interval_secs: int + allow_custom_endpoint: bool @classmethod def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": @@ -164,12 +165,8 @@ def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": "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" - ) + allow_custom_endpoint = bool(config.allow_custom_endpoint) + _validate_feature_store_endpoint(endpoint, allow_custom_endpoint) project_name = config.project_name.strip() feature_entity_name = config.feature_entity_name.strip() feature_view_name = config.feature_view_name.strip() @@ -238,6 +235,7 @@ def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": shutdown_timeout_secs=positive_values["shutdown_timeout_secs"], max_pending_steps=positive_values["max_pending_steps"], poll_interval_secs=positive_values["poll_interval_secs"], + allow_custom_endpoint=allow_custom_endpoint, ) @@ -246,6 +244,51 @@ def validate_feature_store_config(config: FeatureStoreConfig) -> None: FeatureStoreUploadSettings.from_proto(config) +def _validate_feature_store_endpoint( + endpoint: str, allow_custom_endpoint: bool +) -> None: + """Reject endpoints that could divert cloud or FeatureDB credentials. + + The endpoint receives the access keys, STS token and FeatureDB + credentials, so only trusted Alibaba Cloud hosts are accepted unless the + operator explicitly opts in to a vetted custom deployment. + """ + if not endpoint: + return + parsed_endpoint = urlsplit(endpoint if "://" in endpoint else f"//{endpoint}") + if parsed_endpoint.scheme and parsed_endpoint.scheme != "https": + raise ValueError( + "feature_store_config.endpoint must use HTTPS when a scheme is given" + ) + 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") + if ( + parsed_endpoint.path not in ("", "/") + or parsed_endpoint.query + or parsed_endpoint.fragment + ): + raise ValueError( + "feature_store_config.endpoint must not contain path, query, or " + "fragment components" + ) + hostname = parsed_endpoint.hostname + if not hostname: + raise ValueError("feature_store_config.endpoint must name a host") + if allow_custom_endpoint: + return + if parsed_endpoint.port is not None: + raise ValueError( + "feature_store_config.endpoint must not contain a port unless " + "allow_custom_endpoint is set" + ) + if not hostname.endswith(".aliyuncs.com"): + raise ValueError( + "feature_store_config.endpoint must be a trusted *.aliyuncs.com " + "FeatureStore endpoint; set allow_custom_endpoint only for a " + "vetted private deployment" + ) + + def _feature_store_target_hash(settings: FeatureStoreUploadSettings) -> str: target_identity = { "region": settings.region, diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index 2727f0f95..c76a8f0ce 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -337,6 +337,7 @@ def test_proto_groups_required_fields_before_optional_fields(self): "feature_view_ttl_secs", "feature_view_shard_count", "feature_view_replication_count", + "allow_custom_endpoint", ] fields = list(FeatureStoreConfig.DESCRIPTOR.fields) @@ -355,7 +356,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], [1, *list(range(6, 22))]) + self.assertEqual([field.number for field in fields], [1, *list(range(6, 23))]) for field_name in required_fields: with self.subTest(field_name=field_name): config = _feature_store_config() @@ -438,6 +439,64 @@ def test_environment_resolution_and_required_credentials(self): _feature_store_config(feature_view_replication_count=4) ) + def test_endpoint_must_be_a_trusted_https_host(self): + # Cloud and FeatureDB credentials are sent to this host, so trusted + # Alibaba Cloud endpoints are accepted bare or under explicit https. + for endpoint in ( + "featurestore.cn-hangzhou.aliyuncs.com", + "https://featurestore.cn-hangzhou.aliyuncs.com", + "https://featurestore-vpc.cn-beijing.aliyuncs.com/", + ): + with self.subTest(endpoint=endpoint): + settings = FeatureStoreUploadSettings.from_proto( + _feature_store_config(endpoint=endpoint) + ) + self.assertEqual(settings.endpoint, endpoint) + self.assertFalse(settings.allow_custom_endpoint) + + rejected_endpoints = { + "http://featurestore.cn-hangzhou.aliyuncs.com": "HTTPS", + "ftp://featurestore.cn-hangzhou.aliyuncs.com": "HTTPS", + "evil.example.com": "trusted", + "https://aliyuncs.com": "trusted", + "https://featurestore.cn-hangzhou.aliyuncs.com.evil.com": "trusted", + "featurestore.cn-hangzhou.aliyuncs.com:8443": "port", + "https://featurestore.cn-hangzhou.aliyuncs.com/v1": "path, query", + "https://featurestore.cn-hangzhou.aliyuncs.com?x=1": "path, query", + "https://featurestore.cn-hangzhou.aliyuncs.com#frag": "fragment", + "https://": "host", + } + for endpoint, message in rejected_endpoints.items(): + with self.subTest(endpoint=endpoint): + with self.assertRaisesRegex(ValueError, message): + FeatureStoreUploadSettings.from_proto( + _feature_store_config(endpoint=endpoint) + ) + + def test_allow_custom_endpoint_opts_into_vetted_hosts(self): + settings = FeatureStoreUploadSettings.from_proto( + _feature_store_config( + endpoint="https://featurestore.internal.example.com:8443", + allow_custom_endpoint=True, + ) + ) + self.assertTrue(settings.allow_custom_endpoint) + + # The opt-in waives the trusted-host/port rules only; the structural + # credential-safety checks still apply. + for endpoint in ( + "http://featurestore.internal.example.com", + "https://user:secret@featurestore.internal.example.com", + "https://featurestore.internal.example.com/v1", + ): + with self.subTest(endpoint=endpoint): + with self.assertRaisesRegex(ValueError, "endpoint"): + FeatureStoreUploadSettings.from_proto( + _feature_store_config( + endpoint=endpoint, allow_custom_endpoint=True + ) + ) + def test_start_reuses_existing_dynamic_embedding_feature_view(self): with tempfile.TemporaryDirectory() as output_dir: view = _FakeView() From 42ee91253de67c7f2af2cdcbc651537992507c92 Mon Sep 17 00:00:00 2001 From: gecheng Date: Tue, 21 Jul 2026 18:58:06 +0800 Subject: [PATCH 20/50] fix comment --- tzrec/protos/train.proto | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index 957eec751..712174130 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -44,10 +44,7 @@ message FeatureStoreConfig { // 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. + // FeatureDB version for this incremental training run. required string version = 9; From b169a9862388d5c8fc748f1e20aca84a9e61c876 Mon Sep 17 00:00:00 2001 From: gecheng Date: Wed, 22 Jul 2026 10:35:09 +0800 Subject: [PATCH 21/50] [bugfix] let parquet delta dump proceed for shared EC/EBC table names The ambiguous-table guard added to _build_sparse_embedding_contract raises when EmbeddingCollection and EmbeddingBagCollection reuse a table name, but that is TZRec's standard weight-sharing pattern (sequence and non-sequence features on one table) and the pre-existing dynamicemb multi-GPU integration test uses exactly such a config, so the guard broke a test that passes on master. The guard exists to keep a wrong primary key out of the shared FeatureStore: the torchrec delta tracker keys bookkeeping by raw table name and collapses shared-name EC/EBC to one physical table. Keep the hard refusal on the FeatureStore path, where a wrong PK corrupts the store irreversibly, and degrade the re-derivable local parquet dump to a warning, restoring master's behavior. No test asserted the raise. Co-Authored-By: Claude --- tzrec/utils/delta_embedding_dump.py | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 3a6b4af91..5e91e38bc 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -613,13 +613,25 @@ def _build_sparse_embedding_contract( 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." + # The TorchRec delta tracker keys bookkeeping by raw table name, so + # EC/EBC sharing a name collapse to one physical table and a wrong + # primary key could be published. Refuse for the shared FeatureStore + # (corrupting); the local parquet dump is re-derivable, so it warns + # and publishes the surviving table until the tracker is role-aware. + if self._feature_store_enabled: + raise ValueError( + "delta embedding dump cannot safely upload to FeatureStore " + "while table names are reused by both EmbeddingCollection " + f"and EmbeddingBagCollection: {ambiguous_tables}. The " + "TorchRec tracker would publish a wrong primary key; " + "role-aware tracker identity is required." + ) + logger.warning( + "Delta embedding dump cannot distinguish tables reused by both " + "EmbeddingCollection and EmbeddingBagCollection (%s); the TorchRec " + "tracker collapses them to one physical table, so only that " + "table's deltas are published.", + ambiguous_tables, ) name_by_identity = build_sparse_embedding_name_map(metadata_by_identity) From 2f7dc483a3908ef1132b18f3703c38adfc5bcc34 Mon Sep 17 00:00:00 2001 From: gecheng Date: Wed, 22 Jul 2026 10:35:17 +0800 Subject: [PATCH 22/50] [ci] scope-mark the dynamicemb zero-fill delta dump tests test_lookup_dynamic_embeddings_zero_fills_missing_ids and ..._all_missing_ids skip on the CPU lane (no dynamicemb) and, lacking @mark_ci_scope("gpu"), were filtered out of the GPU lane too, so they ran on no per-PR lane and tripped ci_scope_coverage_test. Add the marker, matching the sibling flushes_module_once_per_dump test. Co-Authored-By: Claude --- tzrec/utils/delta_embedding_dump_test.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 07aaea7cf..ea990ee36 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -1600,6 +1600,7 @@ def test_row_wise_lookup_requires_shard_metadata(self): ) @unittest.skipUnless(has_dynamicemb, "dynamicemb is not installed; skipping.") + @mark_ci_scope("gpu") def test_lookup_dynamic_embeddings_zero_fills_missing_ids(self): from dynamicemb.types import CopyMode @@ -1632,6 +1633,7 @@ def test_lookup_dynamic_embeddings_zero_fills_missing_ids(self): ) @unittest.skipUnless(has_dynamicemb, "dynamicemb is not installed; skipping.") + @mark_ci_scope("gpu") def test_lookup_dynamic_embeddings_zero_fills_all_missing_ids(self): dumper = object.__new__(DeltaEmbeddingDumper) dynamic_module = SimpleNamespace( From c0c965cbe8c13e4ecc73ba9765fb5de54a9b7d1b Mon Sep 17 00:00:00 2001 From: gecheng Date: Thu, 23 Jul 2026 21:34:14 +0800 Subject: [PATCH 23/50] [refactor] replace durable outbox uploader with ephemeral in-process pipeline The FeatureStore delta uploader implemented cross-restart durability (journal, snapshots, manifests, fsync, reconciliation, replay, writer lock, error markers) that is unnecessary under the product contract: delta upload is best-effort for the current live process only, and a crash means restart from checkpoint with pending deltas discarded. Replace the 2300-line durable outbox with a ~750-line ephemeral uploader that streams parquet batches directly to the SDK without intermediate SQLite dedup or snapshot copies. Remove fsync, durable makedirs, the durability tracker consumer, and the file-based error marker polling from the dump side. --- .../check_feature_store_delta.py | 206 +- .../check_feature_store_delta_test.py | 138 +- tzrec/utils/delta_embedding_dump.py | 84 +- tzrec/utils/delta_embedding_dump_test.py | 35 +- tzrec/utils/feature_store_delta_uploader.py | 2083 ++++------------- .../feature_store_delta_uploader_test.py | 1194 ++-------- 6 files changed, 765 insertions(+), 2975 deletions(-) diff --git a/tzrec/tools/feature_store/check_feature_store_delta.py b/tzrec/tools/feature_store/check_feature_store_delta.py index 4e5f53eec..9bc4b960c 100644 --- a/tzrec/tools/feature_store/check_feature_store_delta.py +++ b/tzrec/tools/feature_store/check_feature_store_delta.py @@ -11,8 +11,8 @@ 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 +The tool scans the local delta parquet outbox for the latest (or specified) +step, samples keys from its shard set, and queries the configured explicit FeatureDB version through ``DynamicEmbeddingFeatureView.get_online_features``. Example:: @@ -35,7 +35,7 @@ import sys from collections import defaultdict from dataclasses import dataclass, field -from typing import Any, DefaultDict, Dict, List, Mapping, Optional, Sequence, Tuple +from typing import Any, DefaultDict, Dict, List, Optional, Sequence, Tuple import numpy as np import numpy.typing as npt @@ -96,119 +96,94 @@ def resolve_output_dir( 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( +def resolve_upload_step( output_dir: str, - settings: FeatureStoreUploadSettings, + file_prefix: str, + world_size: int, global_step: Optional[int] = None, -) -> Tuple[str, Dict[str, Any], Dict[str, Any]]: - """Load the latest local commit or a requested committed step. +) -> Tuple[int, List[str]]: + """Find the latest (or specified) step with complete parquet shards. Args: output_dir: Delta parquet outbox directory. - settings: Resolved FeatureStore target settings. - global_step: Specific committed step to inspect, or None for the latest. + file_prefix: Scoped file prefix for parquet filenames. + world_size: Expected number of rank shards per step. + global_step: Specific step to inspect, or None for the latest. Returns: - Tuple of state directory, latest commit, and selected success marker. + Tuple of (global_step, shard_paths). Raises: - FileNotFoundError: If no committed marker exists for the configured target. - ValueError: If markers are ambiguous, inconsistent, or the step is invalid. + FileNotFoundError: If no complete shard set is found. """ - 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") - success_path = os.path.join(state_dir, f"step_{selected_step}._FS_SUCCESS.json") - if not os.path.isfile(success_path): + if global_step is not None: + if global_step <= 0: + raise ValueError("global_step must be > 0") + paths = _shard_paths_for_step(output_dir, file_prefix, global_step, world_size) + if not paths: + raise FileNotFoundError( + f"no delta parquet shards found for step {global_step} " + f"under {output_dir}" + ) + return global_step, paths + + best_step = -1 + best_paths: List[str] = [] + if world_size == 1: + pattern = os.path.join(output_dir, f"{file_prefix}_step_*.parquet") + for path in glob.glob(pattern): + basename = os.path.basename(path) + step_str = basename.replace(f"{file_prefix}_step_", "").replace( + ".parquet", "" + ) + try: + step = int(step_str) + except ValueError: + continue + if step > best_step: + best_step = step + best_paths = [path] + else: + pattern = os.path.join(output_dir, "step_*") + for step_dir in sorted(glob.glob(pattern)): + if not os.path.isdir(step_dir): + continue + dir_name = os.path.basename(step_dir) + step_str = dir_name.replace("step_", "") + try: + step = int(step_str) + except ValueError: + continue + paths = _shard_paths_for_step(output_dir, file_prefix, step, world_size) + if paths and step > best_step: + best_step = step + best_paths = paths + + if best_step <= 0: 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}" + f"no complete delta parquet shard set found under {output_dir}" ) - return state_dir, committed, success + return best_step, best_paths -def committed_parquet_paths( - output_dir: str, - file_prefix: str, - global_step: int, - expected_shards: int, +def _shard_paths_for_step( + output_dir: str, file_prefix: str, global_step: int, world_size: 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}" + """Resolve expected shard paths for one step, returning them if complete.""" + if world_size == 1: + path = os.path.join(output_dir, f"{file_prefix}_step_{global_step}.parquet") + return [path] if os.path.isfile(path) else [] + step_dir = os.path.join(output_dir, f"step_{global_step}") + paths = [ + os.path.join( + step_dir, + f"{file_prefix}_step_{global_step}_rank_{rank}_of_{world_size}.parquet", ) - 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 + for rank in range(world_size) + ] + if all(os.path.isfile(path) for path in paths): + return paths + return [] def sample_local_records( @@ -413,7 +388,7 @@ def create_feature_store_view(settings: FeatureStoreUploadSettings) -> Any: def run_check(args: argparse.Namespace) -> int: - """Run one local-marker plus remote-readback verification.""" + """Run one local-parquet 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"): @@ -436,18 +411,11 @@ def run_check(args: argparse.Namespace) -> int: 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, + world_size = args.world_size + global_step, parquet_paths = resolve_upload_step( + output_dir, scoped_prefix, world_size, args.global_step ) samples = sample_local_records( parquet_paths, @@ -469,11 +437,8 @@ def run_check(args: argparse.Namespace) -> int: "feature_view_name": settings.feature_view_name, "version": settings.version, }, - "local_commit": { + "parquet_source": { "global_step": global_step, - "latest_committed_step": 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 ], @@ -486,8 +451,7 @@ def run_check(args: argparse.Namespace) -> int: 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." + "the sampled parquet; a later upload may have updated the same key." ) print(json.dumps(report, indent=2, sort_keys=True)) @@ -516,7 +480,13 @@ 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 publication.", + help="Step to inspect; defaults to the latest available shard set.", + ) + parser.add_argument( + "--world_size", + type=int, + default=1, + help="Number of rank shards per step (default: 1 for single-rank).", ) 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 92fadc35b..7aa24960d 100644 --- a/tzrec/tools/feature_store/check_feature_store_delta_test.py +++ b/tzrec/tools/feature_store/check_feature_store_delta_test.py @@ -9,7 +9,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json import os import sys import tempfile @@ -23,11 +22,10 @@ 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, + resolve_upload_step, sample_local_records, verify_samples, ) @@ -102,79 +100,82 @@ def test_resolve_output_dir_uses_colocated_relocated_outbox(self): colocated, ) - def test_load_committed_upload_checks_target_and_success(self): - settings = SimpleNamespace( - project_name="project", feature_view_name="view", version="v1" - ) + def test_resolve_upload_step_finds_latest_single_rank(self): 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, + prefix = "delta__fs_target" + for step in (10, 20, 5): + path = os.path.join(output_dir, f"{prefix}_step_{step}.parquet") + pq.write_table( + pa.table( + { + "embedding_name": ["a"], + "key_id": pa.array([1], type=pa.int64()), + "embedding": pa.array([[1.0]], type=pa.list_(pa.float32())), + } + ), + path, ) + step, paths = resolve_upload_step(output_dir, prefix, world_size=1) + self.assertEqual(step, 20) + self.assertEqual(len(paths), 1) + self.assertIn("step_20", paths[0]) - 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_resolve_upload_step_finds_latest_multi_rank(self): + with tempfile.TemporaryDirectory() as output_dir: + prefix = "delta__fs_target" + for step in (10, 20): + step_dir = os.path.join(output_dir, f"step_{step}") + os.makedirs(step_dir) + for rank in range(2): + path = os.path.join( + step_dir, + f"{prefix}_step_{step}_rank_{rank}_of_2.parquet", + ) + pq.write_table( + pa.table( + { + "embedding_name": ["a"], + "key_id": pa.array([rank + 1], type=pa.int64()), + "embedding": pa.array( + [[1.0]], type=pa.list_(pa.float32()) + ), + } + ), + path, + ) + step, paths = resolve_upload_step(output_dir, prefix, world_size=2) + self.assertEqual(step, 20) + self.assertEqual(len(paths), 2) - def test_load_committed_upload_accepts_step_before_latest_replay(self): - settings = SimpleNamespace( - project_name="project", feature_view_name="view", version="v1" - ) + def test_resolve_upload_step_explicit_step(self): 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( + prefix = "delta__fs_target" + path = os.path.join(output_dir, f"{prefix}_step_15.parquet") + pq.write_table( + pa.table( { - **target, - "global_step": 20, - "success_records": 1, - "total_records": 1, - }, - output, - ) - - _, committed, success = load_committed_upload( - output_dir, settings, global_step=20 + "embedding_name": ["a"], + "key_id": pa.array([1], type=pa.int64()), + "embedding": pa.array([[1.0]], type=pa.list_(pa.float32())), + } + ), + path, ) + step, paths = resolve_upload_step( + output_dir, prefix, world_size=1, global_step=15 + ) + self.assertEqual(step, 15) + self.assertEqual(paths, [path]) - self.assertEqual(committed["committed_global_step"], 15) - self.assertEqual(success["global_step"], 20) + def test_resolve_upload_step_raises_when_not_found(self): + with tempfile.TemporaryDirectory() as output_dir: + with self.assertRaises(FileNotFoundError): + resolve_upload_step(output_dir, "delta__fs_target", world_size=1) - def test_committed_parquet_paths_and_sampling(self): + def test_parquet_paths_and_sampling(self): with tempfile.TemporaryDirectory() as output_dir: - path = os.path.join(output_dir, "delta__fs_target_step_20.parquet") + prefix = "delta__fs_target" + path = os.path.join(output_dir, f"{prefix}_step_20.parquet") pq.write_table( pa.table( { @@ -188,11 +189,12 @@ def test_committed_parquet_paths_and_sampling(self): ), path, ) - paths = committed_parquet_paths( - output_dir, "delta__fs_target", 20, expected_shards=1 + step, paths = resolve_upload_step( + output_dir, prefix, world_size=1, global_step=20 ) samples = sample_local_records(paths, 2) + self.assertEqual(step, 20) self.assertEqual(paths, [path]) self.assertEqual( [(sample.embedding_name, sample.key_id) for sample in samples], diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 5e91e38bc..fec286714 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -36,9 +36,7 @@ DELTA_DUMP_SCHEMA_VERSION, 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 @@ -54,7 +52,6 @@ ) _CONSUMER = "delta_embedding_dump" -_DURABILITY_CONSUMER = "delta_embedding_dump_durable_ack" _DELTA_DUMP_SCHEMA = pa.schema( [ ("global_step", pa.int64()), @@ -401,33 +398,17 @@ def __init__( self._tracking_pause_depth = 0 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) + os.makedirs(self._output_dir, exist_ok=True) 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=tracker_consumers, + consumers=[_CONSUMER], delete_on_read=True, auto_compact=True, mode=TrackingMode.ID_ONLY, @@ -504,33 +485,16 @@ def close(self, raise_on_error: bool = True, drain: bool = True) -> None: def _feature_store_upload_error( self, force: bool = False ) -> Optional[BaseException]: - """Collect a local/shared uploader error without changing rank control flow.""" + """Collect a local 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 exc return None def _check_feature_store_upload_error(self, force: bool = False) -> None: @@ -682,28 +646,6 @@ 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]: """Temporarily skip delta tracking for non-training forward passes.""" @@ -1042,14 +984,11 @@ def dump(self, global_step: int, check_upload_error: bool = True) -> Optional[st table_chunks, output_path, dump_generation=dump_generation ) if uploader is not None: - uploader.submit(global_step, dump_generation) + uploader.submit(global_step) 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 num_rows == 0: logger.info( "Dumped empty delta embedding shard to %s at step %s.", @@ -1066,7 +1005,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}") - _durable_makedirs(step_dir) + os.makedirs(step_dir, exist_ok=True) return os.path.join( step_dir, ( @@ -1414,10 +1353,6 @@ def _write_table_chunks( 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 - # then leaves at most an orphan .tmp (which the step_*/*.parquet glob - # ignores), never a truncated shard at the canonical path. tmp_path = f"{output_path}.rank{self._rank}.tmp" try: writer_schema = _DELTA_DUMP_SCHEMA @@ -1431,14 +1366,7 @@ def _write_table_chunks( 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 ea990ee36..1b4d68542 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -46,7 +46,6 @@ from tzrec.utils.delta_embedding_dump import ( _CONSUMER, _DELTA_DUMP_SCHEMA, - _DURABILITY_CONSUMER, DeltaEmbeddingDumper, _table_shard_info_from_config, _TableShardInfo, @@ -663,19 +662,6 @@ def fake_all_reduce(tensor, op=None): 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 @@ -1288,23 +1274,6 @@ def test_single_rank_minutes_does_not_require_synced_dataloader_exhaustion(self) 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 = { - _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_dump_does_not_ack_tracker_when_submission_is_rejected(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._feature_store_enabled = True @@ -1323,14 +1292,12 @@ def test_dump_does_not_ack_tracker_when_submission_is_rejected(self): 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") + dumper._uploader.submit.assert_called_once_with(10) 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: diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index 01a88e687..2b2a36e3d 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -9,31 +9,26 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Durable rank-zero uploader for delta-embedding parquet outboxes.""" +"""Ephemeral rank-zero uploader for delta-embedding parquet shards. + +Best-effort upload for the current live training process only. No cross-restart +recovery, no durable state, no replay. A process crash means restart from the +latest checkpoint and pending deltas are discarded. +""" -import fcntl import hashlib import json import os -import re -import shutil -import sqlite3 import threading import time -import uuid -from contextlib import ExitStack from dataclasses import dataclass, field from typing import ( Any, - BinaryIO, Callable, Dict, - Iterable, - Iterator, List, Mapping, Optional, - Set, Tuple, cast, ) @@ -60,12 +55,7 @@ DELTA_DUMP_GENERATION_METADATA_KEY = b"tzrec.delta_embedding.dump_generation" _FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES = 100 -_RECORD_STORE_DB_FILENAME = "upload_records.sqlite" -_RECORD_STORE_CACHE_SIZE_KB = 65536 _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" _REQUIRED_PARQUET_FIELDS = { @@ -258,35 +248,45 @@ def _validate_feature_store_endpoint( parsed_endpoint = urlsplit(endpoint if "://" in endpoint else f"//{endpoint}") if parsed_endpoint.scheme and parsed_endpoint.scheme != "https": raise ValueError( - "feature_store_config.endpoint must use HTTPS when a scheme is given" + "feature_store_config.endpoint must use HTTPS when set; " + f"got scheme={parsed_endpoint.scheme!r}" ) - 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") - if ( - parsed_endpoint.path not in ("", "/") - or parsed_endpoint.query - or parsed_endpoint.fragment - ): + if parsed_endpoint.username or parsed_endpoint.password: raise ValueError( - "feature_store_config.endpoint must not contain path, query, or " - "fragment components" + "feature_store_config.endpoint must not contain userinfo credentials" ) - hostname = parsed_endpoint.hostname - if not hostname: - raise ValueError("feature_store_config.endpoint must name a host") - if allow_custom_endpoint: - return - if parsed_endpoint.port is not None: + if parsed_endpoint.path not in ("", "/"): raise ValueError( - "feature_store_config.endpoint must not contain a port unless " - "allow_custom_endpoint is set" + "feature_store_config.endpoint must not contain a path, query, or fragment" ) - if not hostname.endswith(".aliyuncs.com"): + if parsed_endpoint.query or parsed_endpoint.fragment: raise ValueError( - "feature_store_config.endpoint must be a trusted *.aliyuncs.com " - "FeatureStore endpoint; set allow_custom_endpoint only for a " - "vetted private deployment" + "feature_store_config.endpoint must not contain a path, query, or fragment" ) + if parsed_endpoint.port is not None and not allow_custom_endpoint: + raise ValueError( + "feature_store_config.endpoint must not contain an explicit port" + ) + hostname = (parsed_endpoint.hostname or "").lower() + if not hostname: + raise ValueError("feature_store_config.endpoint must contain a hostname") + trusted_suffixes = ( + ".aliyuncs.com", + ".alibabacloud.com", + ".aliyun.com", + ) + if not hostname.endswith(trusted_suffixes): + if not allow_custom_endpoint: + raise ValueError( + "feature_store_config.endpoint must be a trusted Alibaba Cloud " + f"host (got {hostname!r}); set allow_custom_endpoint only for a " + "vetted private deployment" + ) + + +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 _feature_store_target_hash(settings: FeatureStoreUploadSettings) -> str: @@ -315,245 +315,17 @@ def _scoped_feature_store_file_prefix( 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 _RecordStore: - """Deduplicated delta records spilled to a throwaway SQLite database. - - One delta dump can contain an arbitrary number of unique keys, so keeping - the dedup state and the upload ordering in memory would make rank-zero - peak memory scale with a single dump's size. The store keeps both on disk - behind a bounded page cache, so memory depends only on the streaming - Parquet read batch and one in-flight upload batch. The database is - disposable: after a crash the whole step replays from its durable shard - snapshot. - """ - - def __init__(self, db_path: str) -> None: - self._db_path = db_path - if os.path.exists(db_path): - os.remove(db_path) - self._connection = sqlite3.connect(db_path) - self._connection.execute("PRAGMA journal_mode = OFF") - self._connection.execute("PRAGMA synchronous = OFF") - self._connection.execute(f"PRAGMA cache_size = -{_RECORD_STORE_CACHE_SIZE_KB}") - self._connection.execute( - "CREATE TABLE dedup (" - "embedding_name TEXT NOT NULL, " - "key_id INTEGER NOT NULL, " - "embedding BLOB NOT NULL, " - "PRIMARY KEY (embedding_name, key_id))" - ) - self._record_count: Optional[int] = None - - def __enter__(self) -> "_RecordStore": - return self - - def __exit__(self, *exc_info: Any) -> None: - self.close() - - @property - def record_count(self) -> int: - """Return the number of unique deduplicated records.""" - if self._record_count is None: - cursor = self._connection.execute("SELECT COUNT(*) FROM dedup") - self._record_count = int(cursor.fetchone()[0]) - return self._record_count - - def add_batch( - self, - embedding_names: List[str], - key_ids: np.ndarray, - flat_embeddings: np.ndarray, - offsets: np.ndarray, - ) -> None: - """Spill one streamed shard batch, deduplicating on (name, key_id).""" - rows = [ - ( - embedding_names[index], - int(key_ids[index]), - flat_embeddings[ - int(offsets[index]) : int(offsets[index + 1]) - ].tobytes(), - ) - for index in range(len(embedding_names)) - ] - try: - self._connection.executemany("INSERT INTO dedup VALUES (?, ?, ?)", rows) - except sqlite3.IntegrityError: - # A duplicate can be exact (harmless) or conflicting (corruption); - # replay the batch row by row to tell the two apart. - for row in rows: - try: - self._connection.execute("INSERT INTO dedup VALUES (?, ?, ?)", row) - except sqlite3.IntegrityError: - existing = self._connection.execute( - "SELECT embedding FROM dedup " - "WHERE embedding_name = ? AND key_id = ?", - (row[0], row[1]), - ).fetchone() - if existing[0] != row[2]: - raise ValueError( - "conflicting duplicate delta row for " - f"embedding_name={row[0]!r}, key_id={row[1]}" - ) from None - self._connection.commit() - self._record_count = None - - def iter_upload_batches( - self, batch_size: int - ) -> Iterator[Tuple[List[Dict[str, Any]], int]]: - """Yield (payload, record count) upload batches in a deterministic order.""" - cursor = self._connection.execute( - "SELECT embedding_name, key_id, embedding FROM dedup " - "ORDER BY embedding_name, key_id" - ) - while True: - rows = cursor.fetchmany(batch_size) - if not rows: - return - payload = [ - { - FEATURE_STORE_PK_FIELD: embedding_name, - FEATURE_STORE_SK_FIELD: key_id, - FEATURE_STORE_VALUE_FIELD: np.frombuffer( - embedding, dtype=np.float32 - ).copy(), - } - for embedding_name, key_id, embedding in rows - ] - yield payload, len(payload) - - def to_records(self) -> List[Tuple[str, int, np.ndarray]]: - """Materialize every deduplicated record in upload order.""" - cursor = self._connection.execute( - "SELECT embedding_name, key_id, embedding FROM dedup " - "ORDER BY embedding_name, key_id" - ) - return [ - ( - embedding_name, - int(key_id), - np.frombuffer(embedding, dtype=np.float32).copy(), - ) - for embedding_name, key_id, embedding in cursor - ] - - def close(self) -> None: - """Close the connection and remove the throwaway database file.""" - try: - self._connection.close() - finally: - if os.path.exists(self._db_path): - os.remove(self._db_path) - - class FeatureStoreDeltaUploader: - """Publish complete delta-dump steps from a persistent local outbox. + """Ephemeral in-process uploader for delta-embedding parquet shards. - 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. Restored training may replay a - repeated or lower step; dump generation identifies the new publication. + Best-effort upload for the current live process only. No cross-restart + recovery, no durable state, no replay. A process crash means restart from + the latest checkpoint and pending deltas are discarded. + + The training thread writes parquet shards and enqueues a step. This + object's single background worker waits for the complete shard set, + streams parquet batches directly to the FeatureStore SDK, and cleans up + local files on success. """ def __init__( @@ -566,6 +338,7 @@ def __init__( client_factory: Optional[Callable[..., Any]] = None, clock_ms: Optional[Callable[[], int]] = None, ) -> None: + """Initialize the uploader with validated settings and in-memory state.""" self._settings = FeatureStoreUploadSettings.from_proto(config) self._output_dir = os.path.abspath(output_dir) self._file_prefix = _scoped_feature_store_file_prefix( @@ -594,127 +367,47 @@ def __init__( self._view = None self._condition = threading.Condition() self._pending: Dict[int, float] = {} - # Steps already reconciled as committed with a matching generation; - # discovery skips their JSON/Parquet re-reads so poll lock hold time - # does not grow with the job's committed history. - self._reconciled_steps: Set[int] = set() 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": _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", - "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 + self._last_publish_ts: int = 0 def start(self) -> None: - """Start the single background worker and discover restart work.""" + """Start the background upload worker after validating the remote view.""" 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 with an incompatible target. 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._reconciled_steps = set() 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, + self._started = True + self._worker = threading.Thread( + target=self._run, + name="tzrec-feature-store-delta-uploader", + daemon=True, ) + self._worker.start() + 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, dump_generation: Optional[str] = None) -> None: - """Enqueue a durably written rank-zero shard with bounded back-pressure.""" + def submit(self, global_step: int) -> None: + """Enqueue a completed step for background upload with 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: @@ -723,8 +416,6 @@ def submit(self, global_step: int, dump_generation: Optional[str] = None) -> Non ) if self._closing or self._closed: raise RuntimeError("cannot submit to a closing FeatureStore uploader") - if self._is_committed(global_step, dump_generation): - return while ( global_step not in self._pending and len(self._pending) >= self._settings.max_pending_steps @@ -758,8 +449,7 @@ def close(self, raise_on_error: bool = True, drain: bool = True) -> 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" + "FeatureStore uploader did not drain before shutdown timeout" ) with self._condition: if self._error is None: @@ -779,7 +469,6 @@ def _run(self) -> None: with self._condition: if self._aborting: return - self._add_discovered_steps_locked() if not self._pending: if self._closing: return @@ -788,109 +477,30 @@ def _run(self) -> None: 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 snapshot_dir_exists and self._is_committed(current_step): - try: - 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: + shard_paths = self._expected_shard_paths(current_step) + if not all(os.path.isfile(path) for path in shard_paths): 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" + "timed out waiting for a complete delta shard set " + f"at step {current_step}" ) with self._condition: self._condition.wait(self._settings.poll_interval_secs) continue + self._validate_shard_generation(shard_paths) + 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) - record_store = stack.enter_context( - self._build_record_store( - current_step, snapshot_paths, shard_sources - ) - ) - manifest = self._load_or_create_manifest( - current_step, - snapshot_paths, - record_store.record_count, - shard_descriptions=shard_descriptions, - ) - summary = self._upload_with_retries( - current_step, record_store, 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) + + self._upload_with_retries(current_step, shard_paths) + self._cleanup_shard_files(shard_paths) + + with self._condition: + self._pending.pop(current_step, None) + self._condition.notify_all() current_step = None except _UploadAborted: return @@ -898,137 +508,21 @@ def _run(self) -> None: 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__, - ) + error = FeatureStoreUploadError( + f"FeatureStore delta upload failed{step_context}: {exc}" + ) with self._condition: - self._error = safe_error + if self._error is None: + self._error = 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." + "FeatureStore delta upload failed%s: %s", + step_context, + exc, + exc_info=True, ) - 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): - 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) - - # 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) + self._reset_view(suppress_errors=True) def _raise_if_failed_locked(self) -> None: if self._error is not None: @@ -1039,111 +533,6 @@ def _raise_if_aborting(self) -> None: if self._aborting: raise _UploadAborted() - def _add_discovered_steps_locked(self) -> None: - """Incrementally enqueue new outbox steps without rescanning history. - - Full reconciliation happens once in ``start()`` with an empty - reconciled set. Afterwards, steps already pending or already verified - committed with a matching generation are skipped without any JSON or - Parquet reads, so the condition lock hold time stays independent of - how many steps this job has committed. A committed step whose - canonical generation no longer matches its manifest is deliberately - never marked reconciled, keeping the replacement alarm alive. - """ - 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 step in self._reconciled_steps or step in self._pending: - continue - if self._is_committed(step): - try: - canonical_generation = self._canonical_dump_generation(step) - except _ShardSetNotReady: - canonical_generation = "" - if canonical_generation is None: - self._reconciled_steps.add(step) - continue - manifest_generation = _read_json(self._manifest_path(step)).get( - "dump_generation" - ) - if canonical_generation == manifest_generation: - self._reconciled_steps.add(step) - continue - 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() - # Already pending or reconciled steps are known to be real entries; - # re-statting their directories/shards on every poll would make the - # scan cost grow with the retained outbox history. - known_steps = self._reconciled_steps | set(self._pending) - 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 is None: - continue - step = int(match.group(1)) - if step in known_steps or os.path.isdir( - os.path.join(self._snapshot_root, name) - ): - steps.add(step) - 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: - continue - step = int(match.group(1)) - if step in known_steps: - steps.add(step) - continue - if not os.path.isdir(os.path.join(self._output_dir, name)): - continue - # 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 [ @@ -1162,542 +551,388 @@ 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") - + def _validate_shard_generation(self, shard_paths: List[str]) -> None: + """Verify all shards share the same dump generation.""" 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}") + for path in shard_paths: + parquet_file = pq.ParquetFile(path) + metadata = parquet_file.schema_arrow.metadata or {} 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: + generations.add(generation) + if len(generations) > 1: raise _ShardSetNotReady( - "delta shards from different dump generations are present" + "delta shard set has inconsistent dump generations; " + "a concurrent dump may be replacing files" ) - 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 _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 None - 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) - shards = success.get("shards") - if not isinstance(shards, list): - raise ValueError("FeatureStore success marker has invalid shards") - 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 + def _upload_with_retries(self, global_step: int, shard_paths: List[str]) -> None: + for attempt in range(1, self._settings.max_retries + 1): + self._raise_if_aborting() + try: + self._stream_upload(global_step, shard_paths) + return + except _UploadAborted: + raise + except BaseException as exc: + self._reset_view(suppress_errors=True) + if attempt >= self._settings.max_retries: + raise + logger.warning( + "FeatureStore delta upload attempt %s/%s failed for step %s " + "(%s); retrying after backoff", + attempt, + self._settings.max_retries, + global_step, + exc, + ) + if self._settings.retry_backoff_secs > 0: + time.sleep(self._settings.retry_backoff_secs * attempt) + raise AssertionError("unreachable FeatureStore retry state") - 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 _allocate_timestamp_range(self, batch_count: int) -> Tuple[int, int]: + """Allocate a monotonically increasing timestamp range (in-memory only).""" + reserved = max(batch_count, 1) + range_start = max(int(self._clock_ms()), self._last_publish_ts + 1, 1) + range_end = range_start + reserved - 1 + self._last_publish_ts = range_end + return range_start, range_end - 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)) - state = self._load_success_state(step) - if state is None: - continue - 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 + def _stream_upload(self, global_step: int, shard_paths: List[str]) -> None: + """Stream parquet batches directly to the FeatureStore SDK.""" + view = self._get_view() + max_in_flight = int(getattr(view, "_max_workers", 1)) - committed_path = os.path.join(self._state_dir, "committed.json") - committed_publish_ts = 0 - if os.path.isfile(committed_path): - committed_publish_ts = int(_read_json(committed_path).get("publish_ts", 0)) - latest_step = int(latest_success["global_step"]) - if committed_publish_ts >= int(latest_success["publish_ts"]): - 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, - "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) - - 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)) - 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 - - 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 + total_batches = self._count_total_batches(shard_paths) + ts_range = self._allocate_timestamp_range(total_batches) + range_start = ts_range[0] - @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), - } + completed_batches = 0 + window_batches = 0 + window_records = 0 + started_at = time.monotonic() + next_progress_batch = _FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES + logged_first_window = False - 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 + logger.info( + "FeatureStore delta upload started: step=%s version=%s " + "batches=%s ts_range=%s-%s", + global_step, + self._settings.version, + total_batches, + ts_range[0], + ts_range[1], + ) - 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" + for expected_rank, shard_path in enumerate(shard_paths): + parquet_file = pq.ParquetFile(shard_path) + self._validate_parquet_schema(parquet_file.schema_arrow, shard_path) + for batch in parquet_file.iter_batches( + batch_size=self._settings.upload_batch_size + ): + self._raise_if_aborting() + payload = self._validate_and_build_payload( + batch, global_step, expected_rank + ) + if not payload: + continue + view.write_features( + data=payload, + version=self._settings.version, + write_mode=FEATURE_STORE_WRITE_MODE, + ts=range_start + completed_batches, + ) + completed_batches += 1 + window_batches += 1 + window_records += len(payload) + + if ( + window_batches < max_in_flight + and completed_batches < total_batches + ): + continue + summary = view.write_flush() + self._validate_flush_summary( + summary, + expected_records=window_records, + expected_batches=window_batches, + ) + if ( + not logged_first_window + or completed_batches >= next_progress_batch + or completed_batches == total_batches + ): + logger.info( + "FeatureStore delta upload progress: step=%s " + "batches=%s/%s elapsed_secs=%.1f", + global_step, + completed_batches, + total_batches, + time.monotonic() - started_at, ) - 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 + logged_first_window = True + while next_progress_batch <= completed_batches: + next_progress_batch += ( + _FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES + ) + window_batches = 0 + window_records = 0 - 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())) + if window_batches > 0: + summary = view.write_flush() + self._validate_flush_summary( + summary, + expected_records=window_records, + expected_batches=window_batches, + ) + except BaseException: 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 + view.write_flush() + except BaseException: + pass + raise - @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))) + logger.info( + "FeatureStore delta upload completed: step=%s batches=%s elapsed_secs=%.1f", + global_step, + completed_batches, + time.monotonic() - started_at, + ) - 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 _count_total_batches(self, shard_paths: List[str]) -> int: + """Count the total number of upload batches across all shards.""" + total_rows = 0 + for path in shard_paths: + parquet_file = pq.ParquetFile(path) + total_rows += int(parquet_file.metadata.num_rows) + if total_rows == 0: + return 1 + return ( + total_rows + self._settings.upload_batch_size - 1 + ) // self._settings.upload_batch_size - def _load_or_create_manifest( + def _validate_and_build_payload( self, + batch: pa.RecordBatch, 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, - } - 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( - "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 + expected_rank: int, + ) -> List[Dict[str, Any]]: + """Validate one streamed shard batch and build the SDK payload.""" + num_rows = batch.num_rows + if num_rows == 0: + return [] + global_steps = batch.column("global_step").to_numpy(zero_copy_only=False) + if not bool((global_steps == global_step).all()): + raise ValueError("delta shard global_step mismatch") + ranks = batch.column("rank").to_numpy(zero_copy_only=False) + if not bool((ranks == expected_rank).all()): + raise ValueError("delta shard rank mismatch") + world_sizes = batch.column("world_size").to_numpy(zero_copy_only=False) + if not bool((world_sizes == self._world_size).all()): + raise ValueError("delta shard world_size mismatch") - if manifest.get( - "dump_generation" - ) == dump_generation or not self._is_committed(global_step): + batch_roles = set(batch.column("embedding_role").to_pylist()) + if not batch_roles <= set(SPARSE_EMBEDDING_ROLES): + raise ValueError("delta shard has an invalid embedding role") + embedding_names = batch.column("embedding_name").to_pylist() + for embedding_name in set(embedding_names): + if not embedding_name: + raise ValueError("delta shard embedding_name must not be empty") + if embedding_name not in self._embedding_dimensions: raise ValueError( - f"FeatureStore upload manifest mismatch for {mismatch}" + "delta shard embedding_name is absent from model contract: " + f"{embedding_name!r}" ) - 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"], + key_ids = batch.column("key_id").to_numpy(zero_copy_only=False) + if bool((key_ids == SPARSE_EMBEDDING_INVALID_KEY).any()): + raise ValueError( + "delta shard key_id=-1 is reserved as the Processor/" + "NvEmbeddings invalid-key sentinel" + ) + + embedding_column = cast(pa.ListArray, batch.column("embedding")) + flat_embeddings = embedding_column.values.to_numpy(zero_copy_only=False) + if not bool(np.isfinite(flat_embeddings).all()): + raise ValueError("delta embedding contains NaN or Inf") + offsets = embedding_column.offsets.to_numpy() + lengths = np.diff(offsets) + expected_dims = np.array( + [self._embedding_dimensions[name] for name in embedding_names], + dtype=lengths.dtype, + ) + bad_rows = np.flatnonzero(lengths != expected_dims) + if bad_rows.size > 0: + row = int(bad_rows[0]) + raise ValueError( + f"delta embedding dimension mismatch for {embedding_names[row]!r}: " + f"expected={int(expected_dims[row])}, " + f"actual={int(lengths[row])}" + ) + + return [ + { + FEATURE_STORE_PK_FIELD: embedding_names[i], + FEATURE_STORE_SK_FIELD: int(key_ids[i]), + FEATURE_STORE_VALUE_FIELD: flat_embeddings[ + int(offsets[i]) : int(offsets[i + 1]) + ].copy(), } - _atomic_write_json(path, manifest) - 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 + for i in range(num_rows) + ] - 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: + def _cleanup_shard_files(self, shard_paths: List[str]) -> None: + """Remove uploaded shard files (best-effort).""" + for path in shard_paths: + try: + os.remove(path) + except OSError: + logger.warning("Failed to remove uploaded shard file: %s", path) + if self._world_size > 1 and shard_paths: + step_dir = os.path.dirname(shard_paths[0]) + try: + os.rmdir(step_dir) + except OSError: + pass + + 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) + 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( - "FeatureStore timestamp ranges may only be reserved under the " - "initialized writer journal lock" + "DynamicEmbedding FeatureView schema mismatch: " + f"expected={expected_fields}, actual={actual_fields}" ) - 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, - store: _RecordStore, - 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, store.record_count, manifest) + 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" + ) + 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: + view = self._view + self._view = None + if view is not None: try: - if store.record_count: - summary = self._upload_records(global_step, store, attempt) - else: - # An empty step still validates the FeatureView schema 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 + view.close(wait=True) 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, + close_error = FeatureStoreUploadError( + "FeatureStore SDK writer close failed" + ) + if not suppress_errors: + raise close_error from exc + with self._condition: + if self._error is None: + self._error = close_error + self._condition.notify_all() + logger.error( + "Failed to close FeatureStore SDK writer cleanly (%s)", 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 _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: + 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: + 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 _wait_for_dynamic_embedding_view(self, project: Any) -> Any: """Bounded re-get after a concurrent or partially completed create.""" @@ -1828,272 +1063,6 @@ def _validate_feature_view_metadata(self, feature_view: Any) -> None: 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" - ) - 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: - 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, - global_step: int, - store: _RecordStore, - attempt: Mapping[str, int], - ) -> Dict[str, int]: - view = self._get_view() - max_in_flight = int(getattr(view, "_max_workers", 1)) - total_records = store.record_count - batch_count = ( - total_records + self._settings.upload_batch_size - 1 - ) // self._settings.upload_batch_size - aggregate = { - "total_batches": 0, - "failed_batches": 0, - "total_records": 0, - "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, - total_records, - batch_count, - max_in_flight, - attempt["range_start"], - attempt["range_end"], - ) - try: - completed_batches = 0 - window_batches = 0 - window_records = 0 - logged_first_window = False - for payload, payload_records in store.iter_upload_batches( - self._settings.upload_batch_size - ): - self._raise_if_aborting() - view.write_features( - data=payload, - version=self._settings.version, - write_mode=FEATURE_STORE_WRITE_MODE, - ts=int(attempt["range_start"]) + completed_batches, - ) - completed_batches += 1 - window_batches += 1 - window_records += payload_records - if window_batches < max_in_flight and completed_batches < batch_count: - continue - summary = view.write_flush() - self._validate_flush_summary( - summary, - expected_records=window_records, - expected_batches=window_batches, - ) - for name in aggregate: - aggregate[name] += int(summary[name]) - if ( - not logged_first_window - 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"], - total_records, - time.monotonic() - started_at, - ) - logged_first_window = True - while next_progress_batch <= completed_batches: - next_progress_batch += ( - _FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES - ) - window_batches = 0 - window_records = 0 - 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, expected_batches: int @@ -2116,111 +1085,6 @@ def _validate_flush_summary( ): raise RuntimeError("FeatureStore write_flush reported incomplete writes") - def _build_record_store( - self, - global_step: int, - shard_paths: List[str], - shard_sources: Optional[List[BinaryIO]] = None, - ) -> _RecordStore: - """Stream the step's shards into a disk-spilled deduplicated store. - - Shards are read batch by batch and validated with vectorized checks, - so peak memory stays bounded by one Parquet batch and the store's - page cache, independently of one dump's size. - """ - if shard_sources is None: - with ExitStack() as stack: - opened_sources = [ - stack.enter_context(open(path, "rb")) for path in shard_paths - ] - return self._build_record_store( - global_step, shard_paths, opened_sources - ) - if len(shard_paths) != len(shard_sources): - raise ValueError("delta shard paths and sources must have equal length") - - store = _RecordStore(os.path.join(self._state_dir, _RECORD_STORE_DB_FILENAME)) - try: - for expected_rank, (path, source) in enumerate( - zip(shard_paths, shard_sources) - ): - source.seek(0) - parquet_file = pq.ParquetFile(source) - self._validate_parquet_schema(parquet_file.schema_arrow, path) - for batch in parquet_file.iter_batches( - batch_size=self._settings.upload_batch_size - ): - self._ingest_record_batch(store, global_step, expected_rank, batch) - source.seek(0) - except BaseException: - store.close() - raise - return store - - def _ingest_record_batch( - self, - store: _RecordStore, - global_step: int, - expected_rank: int, - batch: pa.RecordBatch, - ) -> None: - """Validate one streamed shard batch and spill it into the store.""" - num_rows = batch.num_rows - if num_rows == 0: - return - # The parquet schema check already enforces every column type, so the - # value checks below stay vectorized instead of per-row Python loops. - global_steps = batch.column("global_step").to_numpy(zero_copy_only=False) - if not bool((global_steps == global_step).all()): - raise ValueError("delta shard global_step mismatch") - ranks = batch.column("rank").to_numpy(zero_copy_only=False) - if not bool((ranks == expected_rank).all()): - raise ValueError("delta shard rank mismatch") - world_sizes = batch.column("world_size").to_numpy(zero_copy_only=False) - if not bool((world_sizes == self._world_size).all()): - raise ValueError("delta shard world_size mismatch") - - batch_roles = set(batch.column("embedding_role").to_pylist()) - if not batch_roles <= set(SPARSE_EMBEDDING_ROLES): - raise ValueError("delta shard has an invalid embedding role") - embedding_names = batch.column("embedding_name").to_pylist() - for embedding_name in set(embedding_names): - if 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_ids = batch.column("key_id").to_numpy(zero_copy_only=False) - if bool((key_ids == SPARSE_EMBEDDING_INVALID_KEY).any()): - raise ValueError( - "delta shard key_id=-1 is reserved as the Processor/" - "NvEmbeddings invalid-key sentinel" - ) - - embedding_column = cast(pa.ListArray, batch.column("embedding")) - flat_embeddings = embedding_column.values.to_numpy(zero_copy_only=False) - if not bool(np.isfinite(flat_embeddings).all()): - raise ValueError("delta embedding contains NaN or Inf") - offsets = embedding_column.offsets.to_numpy() - lengths = np.diff(offsets) - expected_dims = np.array( - [self._embedding_dimensions[name] for name in embedding_names], - dtype=lengths.dtype, - ) - bad_rows = np.flatnonzero(lengths != expected_dims) - if bad_rows.size > 0: - row = int(bad_rows[0]) - raise ValueError( - f"delta embedding dimension mismatch for {embedding_names[row]!r}: " - f"expected={int(expected_dims[row])}, " - f"actual_shape=({int(lengths[row])},)" - ) - - store.add_batch(embedding_names, key_ids, flat_embeddings, offsets) - @staticmethod def _validate_parquet_schema(schema: pa.Schema, path: str) -> None: metadata = schema.metadata or {} @@ -2234,72 +1098,3 @@ def _validate_parquet_schema(schema: pa.Schema, path: str) -> None: 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, - "dump_generation": manifest["dump_generation"], - "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, - "dump_generation": manifest["dump_generation"], - "manifest_digest": manifest_digest, - } - _atomic_write_json(os.path.join(self._state_dir, "committed.json"), committed) - self._committed_global_step = global_step - # Committed under the condition lock: later polls never re-read this - # step's success marker, manifest, or canonical shard schemas again. - self._reconciled_steps.add(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 index c76a8f0ce..77fd053a5 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -14,7 +14,6 @@ import tempfile import threading import unittest -from contextlib import contextmanager from unittest import mock import pyarrow as pa @@ -24,29 +23,16 @@ 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 ( - _RECORD_STORE_DB_FILENAME, DELTA_DUMP_GENERATION_METADATA_KEY, FeatureStoreDeltaUploader, FeatureStoreUploadError, FeatureStoreUploadSettings, - _json_digest, - _read_json, 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") @@ -338,6 +324,7 @@ def test_proto_groups_required_fields_before_optional_fields(self): "feature_view_shard_count", "feature_view_replication_count", "allow_custom_endpoint", + "upload_mode", ] fields = list(FeatureStoreConfig.DESCRIPTOR.fields) @@ -356,7 +343,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], [1, *list(range(6, 23))]) + self.assertEqual([field.number for field in fields], [1, *list(range(6, 24))]) for field_name in required_fields: with self.subTest(field_name=field_name): config = _feature_store_config() @@ -422,7 +409,7 @@ def test_environment_resolution_and_required_credentials(self): with self.assertRaisesRegex(ValueError, missing_env_name): FeatureStoreUploadSettings.from_proto(config) - with self.assertRaisesRegex(ValueError, "URI userinfo"): + with self.assertRaisesRegex(ValueError, "userinfo"): FeatureStoreUploadSettings.from_proto( _feature_store_config(endpoint="https://user:secret@example.com") ) @@ -440,8 +427,6 @@ def test_environment_resolution_and_required_credentials(self): ) def test_endpoint_must_be_a_trusted_https_host(self): - # Cloud and FeatureDB credentials are sent to this host, so trusted - # Alibaba Cloud endpoints are accepted bare or under explicit https. for endpoint in ( "featurestore.cn-hangzhou.aliyuncs.com", "https://featurestore.cn-hangzhou.aliyuncs.com", @@ -482,8 +467,6 @@ def test_allow_custom_endpoint_opts_into_vetted_hosts(self): ) self.assertTrue(settings.allow_custom_endpoint) - # The opt-in waives the trusted-host/port rules only; the structural - # credential-safety checks still apply. for endpoint in ( "http://featurestore.internal.example.com", "https://user:secret@featurestore.internal.example.com", @@ -556,7 +539,6 @@ def test_start_creates_missing_dynamic_embedding_feature_view(self): } ], ) - 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): @@ -729,64 +711,6 @@ def test_start_reports_missing_entity_when_view_creation_fails(self): 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_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( @@ -838,28 +762,6 @@ def test_complete_step_uploads_merge_with_stable_version_and_ts(self): 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") - 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_upload_uses_bounded_sdk_worker_windows(self): with tempfile.TemporaryDirectory() as output_dir: _write_single_shard( @@ -879,6 +781,7 @@ def test_upload_uses_bounded_sdk_worker_windows(self): ) uploader.start() + uploader.submit(10) uploader.close() self.assertEqual( @@ -886,7 +789,7 @@ def test_upload_uses_bounded_sdk_worker_windows(self): ) self.assertEqual(view.flush_calls, [[1, 1], [1, 1], [1]]) - def test_first_positive_dump_step_is_not_filtered_by_an_implicit_base(self): + def test_first_positive_dump_step_is_not_filtered(self): with tempfile.TemporaryDirectory() as output_dir: _write_single_shard(output_dir, 1, [_row(1, 0, 1, [1.0, 2.0])]) view = _FakeView() @@ -901,88 +804,12 @@ def test_first_positive_dump_step_is_not_filtered_by_an_implicit_base(self): ) uploader.start() + uploader.submit(1) 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_state_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: @@ -1037,184 +864,7 @@ def test_target_scoped_prefix_prevents_cross_version_parquet_replay(self): 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()), - ) - with uploader._build_record_store(10, [shard]) as store: - records = store.to_records() - self.assertFalse( - os.path.exists( - os.path.join(uploader.state_dir, _RECORD_STORE_DB_FILENAME) - ) - ) - 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._build_record_store(10, [shard]) - # The throwaway spill database must not outlive a failed ingest. - self.assertFalse( - os.path.exists( - os.path.join(uploader.state_dir, _RECORD_STORE_DB_FILENAME) - ) - ) - - def test_duplicate_key_across_streamed_batches_is_deduplicated(self): - # Shards stream in upload_batch_size=2 record batches; the duplicate - # arrives in a later ingest batch than its first occurrence, so dedup - # must hold across spilled batches, not just within one. - with tempfile.TemporaryDirectory() as output_dir: - shard = _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, 1, [1.0, 2.0]), - _row(10, 0, 3, [5.0, 6.0]), - ], - ) - uploader = FeatureStoreDeltaUploader( - _feature_store_config(), - output_dir, - "delta", - 1, - {"user_emb": 2}, - client_factory=_FakeClientFactory(_FakeView()), - ) - with uploader._build_record_store(10, [shard]) as store: - records = store.to_records() - - self.assertEqual([record[1] for record in records], [1, 2, 3]) - self.assertEqual(records[0][2].tolist(), [1.0, 2.0]) - - def test_conflicting_duplicate_across_streamed_batches_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, 2, [3.0, 4.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._build_record_store(10, [shard]) - - def test_record_store_reiteration_is_deterministic_across_attempts(self): - # A retried attempt re-streams the whole store; the spill ordering - # must be identical every time so replayed ts ranges carry identical - # batch boundaries. - with tempfile.TemporaryDirectory() as output_dir: - shard = _write_single_shard( - output_dir, - 10, - [ - _row(10, 0, 3, [5.0, 6.0]), - _row(10, 0, 1, [1.0, 2.0]), - _row(10, 0, 2, [3.0, 4.0]), - ], - ) - uploader = FeatureStoreDeltaUploader( - _feature_store_config(), - output_dir, - "delta", - 1, - {"user_emb": 2}, - client_factory=_FakeClientFactory(_FakeView()), - ) - with uploader._build_record_store(10, [shard]) as store: - first = list(store.iter_upload_batches(2)) - second = list(store.iter_upload_batches(2)) - - self.assertEqual([count for _, count in first], [2, 1]) - - def keys_and_values(batches): - return [ - (item["key_id"], item["embedding"].tolist()) - for payload, _ in batches - for item in payload - ] - - self.assertEqual(keys_and_values(first), keys_and_values(second)) - self.assertEqual([key for key, _ in keys_and_values(first)], [1, 2, 3]) - - def test_flush_failure_does_not_publish_success_marker(self): + def test_flush_failure_raises_error(self): failed_summary = { "total_batches": 2, "failed_batches": 1, @@ -1247,13 +897,8 @@ def test_flush_failure_does_not_publish_success_marker(self): 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") - ) - ) - def test_retry_uses_fresh_view_and_newer_persisted_timestamp_range(self): + def test_retry_uses_fresh_view_and_newer_timestamp_range(self): failed_summary = { "total_batches": 1, "failed_batches": 1, @@ -1289,258 +934,172 @@ def test_retry_uses_fresh_view_and_newer_persisted_timestamp_range(self): ) self.assertEqual([call["ts"] for call in all_calls], [777, 778]) - def test_restart_skips_abandoned_attempt_timestamp_range(self): + def test_merge_does_not_require_preprovisioned_version(self): with tempfile.TemporaryDirectory() as output_dir: - shard = _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) - interrupted = FeatureStoreDeltaUploader( + _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + view = _FakeView() + uploader = FeatureStoreDeltaUploader( _feature_store_config(), output_dir, "delta", 1, {"user_emb": 2}, - client_factory=_FakeClientFactory(_FakeView()), - clock_ms=lambda: 777, + client_factory=_FakeClientFactory(view), ) - 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) + uploader.start() + uploader.submit(10) + uploader.close() - view = _FakeView() - restarted = FeatureStoreDeltaUploader( - _feature_store_config(), + self.assertEqual(len(view.calls), 1) + self.assertEqual(view.calls[0]["write_mode"], "MERGE") + + def test_non_draining_close_stops_without_commit(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, 10)], + ) + view = _BlockingView() + uploader = FeatureStoreDeltaUploader( + _feature_store_config(upload_batch_size=1), 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]) + uploader.start() + uploader.submit(10) + self.assertTrue(view.flush_started.wait(timeout=5)) + uploader.close(raise_on_error=False, drain=False) + view.release_flush.set() + self.assertTrue(view.close_finished.wait(timeout=5)) + self.assertTrue(len(view.calls) < 9) - def test_uncommitted_manifest_without_snapshot_fails_recovery(self): + def test_signed_int64_key_is_preserved(self): with tempfile.TemporaryDirectory() as output_dir: - shard = _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) - interrupted = FeatureStoreDeltaUploader( + large_key = (1 << 63) - 1 + _write_single_shard(output_dir, 10, [_row(10, 0, large_key, [1.0, 2.0])]) + view = _FakeView() + uploader = FeatureStoreDeltaUploader( _feature_store_config(), output_dir, "delta", 1, {"user_emb": 2}, - client_factory=_FakeClientFactory(_FakeView()), + client_factory=_FakeClientFactory(view), ) - 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(), + uploader.start() + uploader.submit(10) + uploader.close() + + self.assertEqual(view.calls[0]["data"][0]["key_id"], large_key) + + def test_reserved_invalid_key_is_rejected(self): + with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard(output_dir, 10, [_row(10, 0, -1, [1.0, 2.0])]) + view = _FakeView() + uploader = FeatureStoreDeltaUploader( + _feature_store_config(max_retries=1), 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()), + client_factory=_FakeClientFactory(view), ) - first.start() - with self.assertRaisesRegex(FeatureStoreUploadError, "another process"): - second.start() - first.close() - second.close() + uploader.start() + uploader.submit(10) + with self.assertRaises(FeatureStoreUploadError): + uploader.close() + self.assertEqual(view.calls, []) - def test_start_clears_a_stale_failure_marker_for_restart(self): + def test_empty_step_validates_remote_contract(self): with tempfile.TemporaryDirectory() as output_dir: + _write_single_shard(output_dir, 10, []) + view = _FakeView() uploader = FeatureStoreDeltaUploader( _feature_store_config(), output_dir, "delta", 1, {"user_emb": 2}, - client_factory=_FakeClientFactory(_FakeView()), + client_factory=_FakeClientFactory(view), ) - with open(uploader._error_marker_path, "w") as output: - output.write("{}") uploader.start() - self.assertFalse(os.path.exists(uploader._error_marker_path)) + uploader.submit(10) uploader.close() - def test_merge_does_not_require_preprovisioned_version(self): + self.assertEqual(view.calls, []) + + def test_multi_rank_waits_for_all_shards(self): with tempfile.TemporaryDirectory() as output_dir: - _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + _write_rank_shard( + output_dir, 10, 0, 2, [_row(10, 0, 1, [1.0, 2.0], world_size=2)] + ) + _write_rank_shard( + output_dir, 10, 1, 2, [_row(10, 1, 2, [3.0, 4.0], world_size=2)] + ) view = _FakeView() uploader = FeatureStoreDeltaUploader( - _feature_store_config(max_retries=1), + _feature_store_config(), output_dir, "delta", - 1, + 2, {"user_emb": 2}, client_factory=_FakeClientFactory(view), + clock_ms=lambda: 100, ) uploader.start() + uploader.submit(10) uploader.close() - 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.assertTrue( - os.path.exists( - os.path.join(uploader.state_dir, "step_10._FS_SUCCESS.json") - ) - ) + self.assertEqual(len(view.calls), 2) + self.assertEqual(view.calls[0]["data"][0]["key_id"], 1) + self.assertEqual(view.calls[1]["data"][0]["key_id"], 2) - 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"], - } + def test_dimension_and_finite_value_validation(self): 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]) + _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0, 3.0])]) + view = _FakeView() uploader = FeatureStoreDeltaUploader( - _feature_store_config(max_retries=2), + _feature_store_config(max_retries=1), output_dir, "delta", 1, {"user_emb": 2}, - client_factory=factory, + client_factory=_FakeClientFactory(view), ) uploader.start() uploader.submit(10) - with self.assertRaisesRegex( - FeatureStoreUploadError, "close failed; retry was aborted" - ): + with self.assertRaises(FeatureStoreUploadError): 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() + self.assertEqual(view.calls, []) - 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])]) + _write_single_shard(output_dir, 10, [_row(10, 0, 1, [float("nan"), 2.0])]) + view = _FakeView() uploader = FeatureStoreDeltaUploader( - _feature_store_config(), + _feature_store_config(max_retries=1), output_dir, "delta", 1, {"user_emb": 2}, - client_factory=_FakeClientFactory(_FakeView()), + client_factory=_FakeClientFactory(view), ) - with mock.patch( - "tzrec.utils.feature_store_delta_uploader.shutil.rmtree", - side_effect=OSError("cleanup interrupted"), - ): - uploader.start() + uploader.start() + uploader.submit(10) + with self.assertRaises(FeatureStoreUploadError): 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)) + self.assertEqual(view.calls, []) - def test_non_draining_close_stops_after_inflight_window_without_commit(self): + def test_successful_upload_deletes_shard_files(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]), - ], + shard_path = _write_single_shard( + output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])] ) - view = _BlockingView() + self.assertTrue(os.path.isfile(shard_path)) + view = _FakeView() uploader = FeatureStoreDeltaUploader( _feature_store_config(), output_dir, @@ -1551,94 +1110,80 @@ def test_non_draining_close_stops_after_inflight_window_without_commit(self): ) 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), 2) - self.assertFalse( - os.path.exists( - os.path.join(uploader.state_dir, "step_10._FS_SUCCESS.json") - ) - ) + uploader.close() - 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()), - ) - with uploader._build_record_store(10, [shard]) as store: - records = store.to_records() - self.assertEqual(records[0][1], -2) + self.assertFalse(os.path.isfile(shard_path)) - def test_reserved_invalid_key_is_rejected(self): + def test_in_memory_timestamp_monotonicity_across_steps(self): with tempfile.TemporaryDirectory() as output_dir: - shard = _write_single_shard(output_dir, 10, [_row(10, 0, -1, [1.0, 2.0])]) + _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + _write_single_shard(output_dir, 20, [_row(20, 0, 2, [3.0, 4.0])]) + view = _FakeView() uploader = FeatureStoreDeltaUploader( _feature_store_config(), output_dir, "delta", 1, {"user_emb": 2}, - client_factory=_FakeClientFactory(_FakeView()), + client_factory=_FakeClientFactory(view), + clock_ms=lambda: 100, ) - with self.assertRaisesRegex(ValueError, "invalid-key sentinel"): - uploader._build_record_store(10, [shard]) + uploader.start() + uploader.submit(10) + uploader.submit(20) + uploader.close() - def test_empty_step_validates_remote_contract_before_commit(self): + ts_values = [call["ts"] for call in view.calls] + self.assertEqual(ts_values, [100, 101]) + + def test_in_memory_timestamp_monotonicity_across_retries(self): + failed_summary = { + "total_batches": 1, + "failed_batches": 1, + "total_records": 1, + "success_records": 0, + "failed_records": 1, + } with tempfile.TemporaryDirectory() as output_dir: - _write_single_shard(output_dir, 10, []) - factory = _FakeClientFactory(_FakeView()) + _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(), + _feature_store_config(max_retries=2), output_dir, "delta", 1, {"user_emb": 2}, client_factory=factory, + clock_ms=lambda: 500, ) 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): + self.assertEqual([call["ts"] for call in first_view.calls], [500]) + self.assertEqual([call["ts"] for call in second_view.calls], [501]) + + def test_shard_wait_timeout_raises(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(), + _feature_store_config(shard_wait_timeout_secs=1, poll_interval_secs=1), output_dir, "delta", - 2, + 1, {"user_emb": 2}, client_factory=_FakeClientFactory(view), ) - self.assertEqual(set(uploader._discover_steps()), set()) + uploader.start() + uploader.submit(10) + with self.assertRaises(FeatureStoreUploadError): + uploader.close() + self.assertEqual(view.calls, []) + def test_data_parallel_duplicate_keys_uploaded_without_dedup(self): + with tempfile.TemporaryDirectory() as output_dir: _write_rank_shard( output_dir, 10, @@ -1646,465 +1191,48 @@ def test_multi_rank_discovery_uses_exact_contract_shards(self): 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)], + [_row(10, 1, 1, [1.0, 2.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_commit_marks_step_reconciled_for_later_polls(self): - with tempfile.TemporaryDirectory() as output_dir: - _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) + view = _FakeView() uploader = FeatureStoreDeltaUploader( _feature_store_config(), output_dir, "delta", - 1, + 2, {"user_emb": 2}, - client_factory=_FakeClientFactory(_FakeView()), + client_factory=_FakeClientFactory(view), + clock_ms=lambda: 100, ) uploader.start() uploader.submit(10) uploader.close() - self.assertIn(10, uploader._reconciled_steps) - load_calls = 0 - original_load = uploader._load_success_state - - def counting_load(step): - nonlocal load_calls - load_calls += 1 - return original_load(step) - - with mock.patch.object( - uploader, "_load_success_state", side_effect=counting_load - ): - for _ in range(3): - with uploader._condition: - uploader._add_discovered_steps_locked() - - # The committed step was reconciled at commit time, so repeated - # polls never re-read its success marker or manifest again. - self.assertEqual(load_calls, 0) - self.assertEqual(uploader._pending, {}) + all_keys = [item["key_id"] for call in view.calls for item in call["data"]] + self.assertEqual(all_keys, [1, 1]) - def test_inherited_committed_step_is_reconciled_exactly_once(self): + def test_close_error_surfaces_via_check_error(self): with tempfile.TemporaryDirectory() as output_dir: - _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) - first = FeatureStoreDeltaUploader( - _feature_store_config(), - output_dir, - "delta", - 1, - {"user_emb": 2}, - client_factory=_FakeClientFactory(_FakeView()), - ) - first.start() - first.submit(10) - first.close() - - # A fresh process starts with no reconciled history. - uploader = FeatureStoreDeltaUploader( - _feature_store_config(), - output_dir, - "delta", - 1, - {"user_emb": 2}, - client_factory=_FakeClientFactory(_FakeView()), - ) - load_calls = 0 - generation_calls = 0 - original_load = uploader._load_success_state - original_generation = uploader._canonical_dump_generation - - def counting_load(step): - nonlocal load_calls - load_calls += 1 - return original_load(step) - - def counting_generation(step): - nonlocal generation_calls - generation_calls += 1 - return original_generation(step) - - with ( - mock.patch.object( - uploader, "_load_success_state", side_effect=counting_load - ), - mock.patch.object( - uploader, - "_canonical_dump_generation", - side_effect=counting_generation, - ), - ): - for _ in range(3): - with uploader._condition: - uploader._add_discovered_steps_locked() - - self.assertEqual(load_calls, 1) - self.assertEqual(generation_calls, 1) - self.assertIn(10, uploader._reconciled_steps) - self.assertEqual(uploader._pending, {}) - - def test_committed_step_with_replaced_generation_stays_unreconciled(self): - with tempfile.TemporaryDirectory() as output_dir: - _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) - first = FeatureStoreDeltaUploader( - _feature_store_config(), - output_dir, - "delta", - 1, - {"user_emb": 2}, - client_factory=_FakeClientFactory(_FakeView()), - ) - first.start() - first.submit(10) - first.close() - - # Replace the committed step's canonical shard with a dump from a - # newer generation; discovery must keep checking it (the worker's - # replacement alarm relies on that) instead of reconciling it away. - _write_single_shard( - output_dir, - 10, - [_row(10, 0, 1, [7.0, 8.0])], - generation="ffffffffffffffffffffffffffffffff", - ) - uploader = FeatureStoreDeltaUploader( - _feature_store_config(), - output_dir, - "delta", - 1, - {"user_emb": 2}, - client_factory=_FakeClientFactory(_FakeView()), - ) - with uploader._condition: - uploader._add_discovered_steps_locked() - - self.assertNotIn(10, uploader._reconciled_steps) - self.assertIn(10, uploader._pending) - - def test_discovery_finds_new_outbox_entries_incrementally(self): - with tempfile.TemporaryDirectory() as output_dir: - uploader = FeatureStoreDeltaUploader( - _feature_store_config(), - output_dir, - "delta", - 1, - {"user_emb": 2}, - client_factory=_FakeClientFactory(_FakeView()), - ) - with uploader._condition: - uploader._add_discovered_steps_locked() - self.assertEqual(uploader._pending, {}) - - _write_single_shard(output_dir, 5, [_row(5, 0, 1, [1.0, 2.0])]) - with uploader._condition: - uploader._add_discovered_steps_locked() - self.assertIn(5, uploader._pending) - - 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._build_record_store(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._build_record_store(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) - - with uploader._build_record_store(10, [shard], [source]) as store: - records = store.to_records() - 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])]) + _write_single_shard(output_dir, 10, [_row(10, 0, -1, [1.0, 2.0])]) + view = _FakeView() 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) - - 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(), + _feature_store_config(max_retries=1), output_dir, "delta", 1, {"user_emb": 2}, - client_factory=_FakeClientFactory(_FakeView()), - clock_ms=lambda: 100, + client_factory=_FakeClientFactory(view), ) - verified.start() - verified.close() - self.assertEqual(verified._committed_global_step, 15) + uploader.start() + uploader.submit(10) + with self.assertRaises(FeatureStoreUploadError): + uploader.close() + with self.assertRaises(FeatureStoreUploadError): + uploader.check_error() if __name__ == "__main__": From eeaace12c961b255b85f3b1a3af76e266d114913 Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 24 Jul 2026 08:02:48 +0800 Subject: [PATCH 24/50] [refactor] remove error-bit rendezvous and getattr fallbacks from delta dump torchrun terminates all workers when any rank exits with an uncaught exception, so the cross-rank failure-bit all_reduce propagation is unnecessary. Remove the error-bit from the timed-dump rendezvous (keep only the dump-vote OR-reduce), remove _raise_if_any_interval_dump_failed and _raise_if_any_final_dump_failed, and let dump exceptions propagate directly. Also replace getattr(..., default) with direct attribute access for fields unconditionally initialized in __init__, and remove the requires_synced_dataloader_exhaustion property and its fail_on_uneven_data usage in main.py. --- tzrec/main.py | 13 +- tzrec/utils/delta_embedding_dump.py | 217 ++--------- tzrec/utils/delta_embedding_dump_test.py | 447 ++--------------------- 3 files changed, 57 insertions(+), 620 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index 6ab7c6164..eeae060c1 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -469,22 +469,11 @@ def run_eval(step: int, epoch: int) -> None: # this rank's last consumed event-time, reused by the epoch / final saves data_timestamp = -1.0 - require_equal_train_batches = ( - delta_embedding_dumper is not None - and delta_embedding_dumper.requires_synced_dataloader_exhaustion - ) - sync_train_data_exhaustion = ( - check_all_workers_data_status or require_equal_train_batches - ) for i_epoch in epoch_iter: - # Multi-rank dumps synchronize their per-step state, so uneven workers - # must fail together instead of silently dropping a lagging rank's - # trailing delta rows or leaving collectives. pipeline = create_train_pipeline( model, optimizer, - check_all_workers_data_status=sync_train_data_exhaustion, - fail_on_uneven_data=require_equal_train_batches, + check_all_workers_data_status=check_all_workers_data_status, ) 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 fec286714..0ea7feeef 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -35,7 +35,6 @@ DELTA_DUMP_GENERATION_METADATA_KEY, DELTA_DUMP_SCHEMA_VERSION, FeatureStoreDeltaUploader, - FeatureStoreUploadError, feature_store_delta_file_prefix, validate_feature_store_config, ) @@ -386,7 +385,6 @@ def __init__( self._next_dump_time: Optional[float] = None self._last_dump_step: Optional[int] = None self._pending_rendezvous: Optional[torch.Tensor] = None - self._pending_upload_error: Optional[BaseException] = None self._timed_dump_in_flight = False self._output_dir = config.output_dir or os.path.join( model_dir, "delta_embedding_dump" @@ -456,18 +454,6 @@ 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. - - Multi-rank dumps of either cadence need every rank to reach the same - final step. ``final_dump`` skips steps already written on their dump - boundary, and a rank that stops before the synced MAX step never ran - that boundary's ``maybe_dump``, so an unsynced stop would make it - adopt the boundary skip and silently drop its trailing tracked rows. - """ - return self._world_size > 1 - def start(self) -> None: """Start timed cadence and rank-zero publication after initialization.""" if self._feature_store_enabled: @@ -486,13 +472,12 @@ def _feature_store_upload_error( self, force: bool = False ) -> Optional[BaseException]: """Collect a local uploader error without changing rank control flow.""" - if not getattr(self, "_feature_store_enabled", False): + if not self._feature_store_enabled: return None - uploader = getattr(self, "_uploader", None) - if uploader is not None: + if self._uploader is not None: try: - uploader.check_error() + self._uploader.check_error() except BaseException as exc: return exc return None @@ -525,7 +510,7 @@ def _initialize_run_generation(self) -> None: 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): + if not self._feature_store_enabled: return None if self._run_generation is None: raise RuntimeError("FeatureStore delta dumper must be started before use") @@ -638,7 +623,7 @@ def _build_sparse_embedding_contract( return identity_by_fqn, embedding_dimensions def _tracker_cursor_before_read(self) -> Optional[int]: - if not getattr(self, "_feature_store_enabled", False): + if not self._feature_store_enabled: return None return int(self._tracker.per_consumer_batch_idx[_CONSUMER]) @@ -661,57 +646,28 @@ def maybe_dump(self, global_step: int) -> None: Args: global_step: Current training step. """ + self._check_feature_store_upload_error(force=True) if self._requires_dump_state_rendezvous(): - should_dump, any_failed = self._consume_dump_state_rendezvous() - if self._interval_steps is not None: - # Step-boundary decisions are deterministic across ranks and - # stay on their exact boundary step; only the failure bit - # travels through the pipelined rendezvous. - should_dump = ( - global_step > 0 and global_step % self._interval_steps == 0 - ) - if any_failed: - self._raise_rendezvoused_upload_failure() + should_dump = self._consume_dump_state_rendezvous() self._launch_dump_state_rendezvous(global_step) - check_upload_error = False else: should_dump = self._local_dump_decision(global_step) - check_upload_error = True if should_dump: - dump_error: Optional[BaseException] = None - try: - if check_upload_error: - self.dump(global_step) - else: - # The failure bit was already rendezvoused; a rank-local - # recheck inside dump() could diverge across ranks. - self.dump(global_step, check_upload_error=False) - except BaseException as exc: - dump_error = exc - self._raise_if_any_interval_dump_failed(dump_error) + 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._timed_dump_in_flight = False self._tracker.step() def _requires_dump_state_rendezvous(self) -> bool: - """Return whether maybe_dump must all-reduce its dump and failure state. + """Return whether maybe_dump must all-reduce its dump vote. Timed dumps OR-reduce the rank-zero clock decision so every rank dumps - together. Distributed FeatureStore dumps of either cadence rendezvous - the uploader-failure bit: rank zero observes an async upload failure - immediately while peers only see the shared error marker through a - throttled poll, so without the collective rank zero could raise alone - and strand its peers in the next training collective. The rendezvous - is pipelined across steps so it never blocks the training hot path. + together. Step-based dumps are deterministic across ranks and need no + collective. """ - return getattr(self, "_world_size", 1) > 1 and ( - self._interval_secs is not None - or getattr(self, "_feature_store_enabled", False) - ) + return self._world_size > 1 and self._interval_secs is not None def _local_dump_decision(self, global_step: int) -> bool: """Return the dump decision when no cross-rank rendezvous is needed.""" @@ -729,14 +685,12 @@ def _local_dump_decision(self, global_step: int) -> bool: return self._rank == 0 and time.monotonic() >= self._next_dump_time def _launch_dump_state_rendezvous(self, global_step: int) -> None: - """Asynchronously all-reduce this step's dump vote and failure bit. + """Asynchronously all-reduce this step's timed dump vote. The reduced tensor is consumed by the next ``maybe_dump`` call. NCCL orders this collective ahead of the next training step's communication on the same process group, so it has already completed by the time - the next call reads it back: the hot path pays an asynchronous launch - instead of a host-blocking all_reduce plus ``.item()`` pair on every - training step, and the decision lands at most one step late. + the next call reads it back. """ if not ( torch.distributed.is_available() and torch.distributed.is_initialized() @@ -744,30 +698,21 @@ def _launch_dump_state_rendezvous(self, global_step: int) -> None: raise RuntimeError( "distributed delta embedding dump requires an initialized process group" ) - local_error = self._feature_store_upload_error() - if local_error is not None: - self._pending_upload_error = local_error state = torch.tensor( - [ - int(self._timed_dump_vote(global_step)), - int(self._pending_upload_error is not None), - ], + [int(self._timed_dump_vote(global_step))], dtype=torch.int32, device=self._device, ) torch.distributed.all_reduce(state, op=torch.distributed.ReduceOp.MAX) self._pending_rendezvous = state - def _consume_dump_state_rendezvous(self) -> Tuple[bool, bool]: - """Read back the previous step's reduced dump vote and failure bit.""" + def _consume_dump_state_rendezvous(self) -> bool: + """Read back the previous step's reduced timed dump vote.""" state = self._pending_rendezvous self._pending_rendezvous = None if state is None: - return False, False - # Launched one training step ago and ordered ahead of this step's - # communication, so this read does not block the host. - should_dump, any_failed = state.tolist() - return bool(should_dump), bool(any_failed) + return False + return bool(state.tolist()[0]) def _timed_dump_vote(self, global_step: int) -> bool: """Return rank zero's local vote for whether a timed dump is due. @@ -790,52 +735,6 @@ def _timed_dump_vote(self, global_step: int) -> bool: return True return False - def _raise_rendezvoused_upload_failure(self) -> None: - """Raise the rendezvoused upload failure identically on every rank.""" - local_error = self._pending_upload_error - if local_error is not None: - raise local_error.with_traceback(local_error.__traceback__) - raise FeatureStoreUploadError( - "FeatureStore delta upload failed on another distributed worker; " - "all training workers are stopping and parquet outbox files were " - "retained" - ) - - def _raise_if_any_interval_dump_failed( - self, local_error: Optional[BaseException] - ) -> None: - """Propagate one interval dump failure before the next training step. - - Boundary and timed dumps run on every rank, but a failure can still be - rank-local (a shard write error on one node, or one rank observing the - shared uploader-error marker ahead of the others). Reduce the failure - bit on every multi-rank cadence so all workers stop together instead - of the failing rank raising alone and stranding its peers in the next - training collective. - """ - any_failed = local_error is not None - if getattr(self, "_world_size", 1) > 1: - if not ( - torch.distributed.is_available() and torch.distributed.is_initialized() - ): - raise RuntimeError( - "distributed 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( - "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. @@ -869,87 +768,29 @@ def final_dump(self, global_step: int) -> Optional[str]: # 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. - 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 + return self.dump(global_step) def _sync_final_step(self, global_step: int) -> int: - """Align the final step and uploader failure state before the trailing flush. - - ``maybe_dump`` runs in lockstep so every rank shares ``global_step``, - and the training loop keeps ``final_dump`` aligned too: multi-rank - dumps require synced dataloader exhaustion, because a rank stopping at - an earlier step would adopt the synced MAX step, hit the boundary-step - skip without ever running that boundary's ``maybe_dump``, and silently - drop its trailing tracked rows. The MAX all-reduce remains as a - defensive guard so every rank still takes the same skip/dump decision - into the same ``step_/`` dir -- the ragged shard set the per-rank - empty-shard logic prevents. - - 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. + """Align the final step across ranks before the trailing flush. + + The MAX all-reduce ensures every rank takes the same skip/dump + decision into the same ``step_/`` directory. """ - 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() ): device = torch.device(f"cuda:{torch.cuda.current_device()}") - final_state = torch.tensor( - [global_step, int(any_failed)], dtype=torch.long, device=device - ) + final_state = torch.tensor([global_step], 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 _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]: + def dump(self, global_step: int) -> 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. - ``maybe_dump`` and ``final_dump`` disable it after synchronizing - the error bit across ranks. Returns: Path to the dumped parquet file, or None if no data to dump. @@ -957,9 +798,7 @@ def dump(self, global_step: int, check_upload_error: bool = True) -> Optional[st 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) + uploader = self._uploader dump_generation = self._next_dump_generation(global_step) tracker_cursor = self._tracker_cursor_before_read() try: @@ -975,7 +814,7 @@ def dump(self, global_step: int, check_upload_error: bool = True) -> Optional[st if ( num_rows == 0 and self._world_size == 1 - and not getattr(self, "_feature_store_enabled", False) + and not self._feature_store_enabled ): logger.info("No delta embedding rows to dump at step %s.", global_step) return None diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 1b4d68542..a2a80c798 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -58,7 +58,6 @@ 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 @@ -225,7 +224,7 @@ def _run_uneven_exhaustion_rejection(rank: int, world_size: int, output_dir: str # uneven stop must fail loudly on every rank instead. with MultiProcessContext(rank=rank, world_size=world_size, backend="nccl") as ctx: model = _build_sharded_delta_dump_model(rank, world_size, ctx) - dumper = DeltaEmbeddingDumper( + DeltaEmbeddingDumper( model, DeltaEmbeddingDumpConfig(dump_interval_steps=50, output_dir=output_dir), output_dir, @@ -233,11 +232,9 @@ def _run_uneven_exhaustion_rejection(rank: int, world_size: int, output_dir: str [], ) testcase = unittest.TestCase() - testcase.assertTrue(dumper.requires_synced_dataloader_exhaustion) pipeline = create_train_pipeline( model, - check_all_workers_data_status=dumper.requires_synced_dataloader_exhaustion, - fail_on_uneven_data=dumper.requires_synced_dataloader_exhaustion, + check_all_workers_data_status=True, ) # Rank 0's 50th fetch returns None while rank 1 still has batch 50; the # pipeline's collective data-status check must raise on both ranks. @@ -525,12 +522,6 @@ def test_final_dump_skips_boundary_step_to_avoid_overwrite(self): [call.args[0] for call in dump_mock.call_args_list], [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): # A lagging rank reaches final_dump at a boundary step (50) while the @@ -544,11 +535,7 @@ 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) - if tensor.numel() == 2: - tensor[0] = 73 - tensor[1] = 0 - else: - tensor[0] = 0 + tensor[0] = 73 with ( mock.patch.object(dumper, "dump") as dump_mock, @@ -564,7 +551,7 @@ def fake_all_reduce(tensor, op=None): mock.patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), ): dumper.final_dump(50) - dump_mock.assert_called_once_with(73, check_upload_error=False) + dump_mock.assert_called_once_with(73) def test_final_dump_syncs_ranks_before_skipping_step_zero(self): dumper = object.__new__(DeltaEmbeddingDumper) @@ -577,7 +564,6 @@ 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 ( @@ -598,70 +584,6 @@ def fake_all_reduce(tensor, op=None): 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_steps = 50 - dumper._interval_secs = None - 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_steps = 50 - dumper._interval_secs = None - 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_dump_generation_is_stable_per_step_and_unique_per_run(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._feature_store_enabled = True @@ -678,6 +600,8 @@ def test_maybe_dump_uses_checkpoint_aligned_global_step(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._interval_steps = 50 dumper._interval_secs = None + dumper._world_size = 1 + dumper._feature_store_enabled = False dumper._tracker = mock.MagicMock() with mock.patch.object(dumper, "dump") as dump_mock: dumper.maybe_dump(0) @@ -732,7 +656,7 @@ def _new_rendezvous_dumper(self) -> DeltaEmbeddingDumper: dumper._device = torch.device("cpu") dumper._tracker = mock.MagicMock() dumper._pending_rendezvous = None - dumper._pending_upload_error = None + dumper._feature_store_enabled = False dumper._timed_dump_in_flight = False return dumper @@ -751,8 +675,7 @@ def test_time_dump_decision_is_reduced_from_rank_zero(self): def fake_all_reduce(state, op=None): self.assertIs(op, torch.distributed.ReduceOp.MAX) launched_states.append(state.tolist()) - if state.numel() == 2: - state[0] = 1 + state[0] = 1 with ( mock.patch.object( @@ -766,95 +689,11 @@ def fake_all_reduce(state, op=None): dump_mock.assert_not_called() dumper.maybe_dump(11) - dump_mock.assert_called_once_with(11, check_upload_error=False) + dump_mock.assert_called_once_with(11) self.assertEqual(dumper._last_dump_step, 11) - # Both steps launched [0, 0] locally; rank zero's vote arrived through - # the MAX reduce. The trailing [0] is the dump-failure rendezvous. - self.assertEqual(launched_states, [[0, 0], [0, 0], [0]]) + self.assertEqual(launched_states, [[0], [0]]) self.assertEqual(dumper._tracker.step.call_count, 2) - def test_timed_maybe_dump_reduces_local_upload_error_before_raising(self): - # Rank zero observes an async uploader failure immediately; the bit is - # launched on this step and consumed on the next so every rank raises - # together instead of rank zero raising alone. - dumper = self._new_rendezvous_dumper() - dumper._interval_steps = None - dumper._interval_secs = 60.0 - dumper._next_dump_time = 0.0 - dumper._rank = 0 - dumper._world_size = 2 - 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), - ): - dumper.maybe_dump(10) - with self.assertRaises(FeatureStoreUploadError) as context: - dumper.maybe_dump(11) - - self.assertTrue(collective_called) - self.assertIs(context.exception, local_error) - dump_mock.assert_not_called() - self.assertEqual(dumper._tracker.step.call_count, 1) - - 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 = self._new_rendezvous_dumper() - dumper._interval_steps = None - dumper._interval_secs = 60.0 - dumper._next_dump_time = 0.0 - dumper._rank = 1 - dumper._world_size = 2 - dumper._feature_store_enabled = True - dumper._feature_store_error_marker_path = marker_path - # The throttled poll hides the marker locally; only rank zero's - # rendezvoused failure bit surfaces the error, one step later. - 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[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), - ): - dumper.maybe_dump(10) - with self.assertRaisesRegex( - FeatureStoreUploadError, "another distributed worker" - ): - dumper.maybe_dump(11) - - dump_mock.assert_not_called() - self.assertEqual(dumper._tracker.step.call_count, 1) - def test_timed_maybe_dump_advances_after_all_workers_succeed(self): dumper = self._new_rendezvous_dumper() dumper._interval_steps = None @@ -885,14 +724,14 @@ def fake_all_reduce(state, op=None): dump_mock.assert_not_called() dumper.maybe_dump(11) - self.assertEqual(collective_states, [[1, 0], [0, 0], [0]]) - dump_mock.assert_called_once_with(11, check_upload_error=False) + self.assertEqual(collective_states, [[1], [0]]) + dump_mock.assert_called_once_with(11) self.assertEqual(dumper._last_dump_step, 11) self.assertEqual(dumper._next_dump_time, 62.0) self.assertFalse(dumper._timed_dump_in_flight) self.assertEqual(dumper._tracker.step.call_count, 2) - def test_timed_maybe_dump_reduces_local_dump_failure_before_raising(self): + def test_timed_maybe_dump_propagates_local_dump_failure(self): dumper = self._new_rendezvous_dumper() dumper._interval_steps = None dumper._interval_secs = 60.0 @@ -917,42 +756,9 @@ def fake_all_reduce(state, op=None): with self.assertRaises(RuntimeError) as context: dumper.maybe_dump(11) - self.assertEqual(collective_states, [[1, 0], [0, 0], [1]]) + self.assertEqual(collective_states, [[1], [0]]) self.assertIs(context.exception, dump_error) - dump_mock.assert_called_once_with(11, check_upload_error=False) - self.assertIsNone(dumper._last_dump_step) - self.assertEqual(dumper._tracker.step.call_count, 1) - - def test_timed_maybe_dump_surfaces_remote_dump_failure(self): - dumper = self._new_rendezvous_dumper() - dumper._interval_steps = None - dumper._interval_secs = 60.0 - dumper._next_dump_time = 0.0 - dumper._rank = 0 - dumper._world_size = 2 - collective_states = [] - - def fake_all_reduce(state, op=None): - self.assertIs(op, torch.distributed.ReduceOp.MAX) - collective_states.append(state.tolist()) - if state.numel() == 1: - state[0] = 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), - ): - dumper.maybe_dump(10) - with self.assertRaisesRegex(RuntimeError, "another distributed worker"): - dumper.maybe_dump(11) - - self.assertEqual(collective_states, [[1, 0], [0, 0], [0]]) - dump_mock.assert_called_once_with(11, check_upload_error=False) + dump_mock.assert_called_once_with(11) self.assertIsNone(dumper._last_dump_step) self.assertEqual(dumper._tracker.step.call_count, 1) @@ -979,7 +785,7 @@ def fake_all_reduce(state, op=None): ) as dump_mock, mock.patch( "tzrec.utils.delta_embedding_dump.time.monotonic", - side_effect=[100.0, 100.0, 100.0, 100.0], + side_effect=[100.0, 100.0, 100.0], ), mock.patch("torch.distributed.is_available", return_value=True), mock.patch("torch.distributed.is_initialized", return_value=True), @@ -989,192 +795,12 @@ def fake_all_reduce(state, op=None): dumper.maybe_dump(11) dumper.maybe_dump(12) - dump_mock.assert_called_once_with(11, check_upload_error=False) + dump_mock.assert_called_once_with(11) self.assertEqual(dumper._next_dump_time, 160.0) self.assertFalse(dumper._timed_dump_in_flight) - self.assertEqual(collective_states, [[1, 0], [0, 0], [0], [0, 0]]) + self.assertEqual(collective_states, [[1], [0], [0]]) self.assertEqual(dumper._tracker.step.call_count, 3) - def test_step_maybe_dump_reduces_local_upload_error_before_raising(self): - # Rank zero observes an async uploader failure immediately; the - # pipelined rendezvous launches the bit on this step and raises on the - # next so every rank stops together instead of rank zero raising alone. - dumper = self._new_rendezvous_dumper() - dumper._interval_steps = 50 - dumper._interval_secs = None - dumper._rank = 0 - dumper._world_size = 2 - dumper._feature_store_enabled = True - 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(), [0, 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), - ): - dumper.maybe_dump(10) - with self.assertRaises(FeatureStoreUploadError) as context: - dumper.maybe_dump(11) - - self.assertTrue(collective_called) - self.assertIs(context.exception, local_error) - dump_mock.assert_not_called() - self.assertEqual(dumper._tracker.step.call_count, 1) - - def test_step_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 = self._new_rendezvous_dumper() - dumper._interval_steps = 50 - dumper._interval_secs = None - dumper._rank = 1 - dumper._world_size = 2 - dumper._feature_store_enabled = True - dumper._feature_store_error_marker_path = marker_path - # The throttled poll hides the marker on this rank; only the - # rendezvoused failure bit from rank zero surfaces the error. - 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[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), - ): - dumper.maybe_dump(10) - with self.assertRaisesRegex( - FeatureStoreUploadError, "another distributed worker" - ): - dumper.maybe_dump(11) - - dump_mock.assert_not_called() - self.assertEqual(dumper._tracker.step.call_count, 1) - - def test_step_feature_store_boundary_dump_skips_rank_local_upload_recheck(self): - # Only the failure bit travels through the pipelined rendezvous; the - # boundary decision stays deterministic and on its exact boundary step. - dumper = self._new_rendezvous_dumper() - dumper._interval_steps = 50 - dumper._interval_secs = None - dumper._rank = 1 - dumper._world_size = 2 - dumper._feature_store_enabled = True - 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("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(49) - dump_mock.assert_not_called() - dumper.maybe_dump(50) - - self.assertEqual(collective_states, [[0, 0], [0, 0], [0]]) - dump_mock.assert_called_once_with(50, check_upload_error=False) - self.assertEqual(dumper._last_dump_step, 50) - self.assertEqual(dumper._tracker.step.call_count, 2) - - def test_step_boundary_dump_failure_is_reduced_before_raising(self): - # Even without FeatureStore, a rank-local boundary dump failure (e.g. a - # shard write error on one node) must be reduced so the failing rank - # does not raise alone and strand peers in the next collective. - dumper = object.__new__(DeltaEmbeddingDumper) - dumper._interval_steps = 50 - dumper._interval_secs = None - dumper._last_dump_step = None - dumper._rank = 1 - dumper._world_size = 2 - dumper._device = torch.device("cpu") - dumper._tracker = mock.MagicMock() - dump_error = RuntimeError("local dump 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]) - collective_called = True - - with ( - 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(50) - - self.assertTrue(collective_called) - self.assertIs(context.exception, dump_error) - dump_mock.assert_called_once_with(50) - self.assertIsNone(dumper._last_dump_step) - dumper._tracker.step.assert_not_called() - - def test_step_boundary_dump_surfaces_remote_failure(self): - dumper = object.__new__(DeltaEmbeddingDumper) - dumper._interval_steps = 50 - dumper._interval_secs = None - dumper._last_dump_step = None - dumper._rank = 1 - dumper._world_size = 2 - dumper._device = torch.device("cpu") - dumper._tracker = mock.MagicMock() - - def fake_all_reduce(state, op=None): - self.assertIs(op, torch.distributed.ReduceOp.MAX) - self.assertEqual(state.tolist(), [0]) - state[0] = 1 - - with ( - 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(50) - - dump_mock.assert_called_once_with(50) - 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 @@ -1247,33 +873,6 @@ 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_multi_rank_step_interval_requires_synced_dataloader_exhaustion(self): - # Regression: with ranks finishing at steps 49 and 50, final_dump synced - # both to the boundary 50, both took the boundary-step skip, and the - # rank that never ran maybe_dump(50) silently dropped its trailing - # tracked rows. Every multi-rank cadence must keep exhaustion aligned - # so all ranks participate in every boundary dump. - dumper = object.__new__(DeltaEmbeddingDumper) - dumper._interval_secs = None - dumper._interval_steps = 50 - dumper._world_size = 2 - - self.assertTrue(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_dump_does_not_ack_tracker_when_submission_is_rejected(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._feature_store_enabled = True @@ -1324,6 +923,11 @@ def test_multi_gpu_dump_writes_empty_shard_when_rank_has_no_delta(self): dumper._file_prefix = "delta_embedding" dumper._rank = 1 dumper._world_size = 2 + dumper._feature_store_enabled = False + dumper._uploader = None + dumper._run_generation = None + dumper._tracker = mock.MagicMock() + dumper._tracker.per_consumer_batch_idx = {_CONSUMER: 0} with ( mock.patch.object(dumper, "_collect_table_weights", return_value={}), mock.patch.object(dumper, "_collect_dynamic_modules", return_value={}), @@ -1350,6 +954,11 @@ def test_single_gpu_dump_skips_file_when_rank_has_no_delta(self): dumper._file_prefix = "delta_embedding" dumper._rank = 0 dumper._world_size = 1 + dumper._feature_store_enabled = False + dumper._uploader = None + dumper._run_generation = None + dumper._tracker = mock.MagicMock() + dumper._tracker.per_consumer_batch_idx = {_CONSUMER: 0} with ( mock.patch.object(dumper, "_collect_table_weights", return_value={}), mock.patch.object(dumper, "_collect_dynamic_modules", return_value={}), From fb5314376dba185cb0f50d088d0dabdfee7eaf22 Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 24 Jul 2026 10:24:26 +0800 Subject: [PATCH 25/50] [refactor] use credential provider chain instead of static config credentials Remove security_token from FeatureStoreConfig proto (now reserved) and remove all credential fields from FeatureStoreUploadSettings. Cloud credentials (AK/SK/STS) are now resolved at runtime through the alibabacloud_credentials default provider chain, enabling automatic STS token refresh for long-running training jobs. FeatureDB credentials remain env-var based (FEATUREDB_USERNAME/FEATUREDB_PASSWORD). Since the config no longer contains secrets, revert the save_pipeline_config_artifact sanitization and use save_message directly. --- tzrec/main.py | 2 +- tzrec/protos/train.proto | 10 +-- .../check_feature_store_delta.py | 20 ++++-- .../check_feature_store_delta_test.py | 23 ++++-- tzrec/utils/config_util.py | 57 +-------------- tzrec/utils/config_util_test.py | 41 ----------- tzrec/utils/feature_store_delta_uploader.py | 71 +++++++------------ .../feature_store_delta_uploader_test.py | 67 ++++++----------- 8 files changed, 86 insertions(+), 205 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index eeae060c1..1590f2eaf 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -853,7 +853,7 @@ def train_and_evaluate( dist.barrier() if is_rank_zero: - config_util.save_pipeline_config_artifact( + config_util.save_message( pipeline_config, os.path.join(pipeline_config.model_dir, "pipeline.config") ) with open(os.path.join(pipeline_config.model_dir, "version"), "w") as f: diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index 712174130..e1d5d46d5 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -30,10 +30,11 @@ 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"; + // Cloud credentials (AK/SK/STS) are resolved at runtime through the + // Alibaba Cloud default credential provider chain (alibabacloud_credentials). + // FeatureDB credentials are read from FEATUREDB_USERNAME/FEATUREDB_PASSWORD. + reserved 2 to 5, 11; + reserved "access_key_id", "access_key_secret", "featuredb_username", "featuredb_password", "security_token"; // FeatureStore control-plane region. An explicitly empty value falls back // to ALIBABA_CLOUD_REGION at runtime. required string region = 1; @@ -54,7 +55,6 @@ message FeatureStoreConfig { // required when a scheme is given) without userinfo, port, path, query or // fragment components, unless allow_custom_endpoint is set. 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; diff --git a/tzrec/tools/feature_store/check_feature_store_delta.py b/tzrec/tools/feature_store/check_feature_store_delta.py index 9bc4b960c..c105d2994 100644 --- a/tzrec/tools/feature_store/check_feature_store_delta.py +++ b/tzrec/tools/feature_store/check_feature_store_delta.py @@ -345,21 +345,29 @@ def create_feature_store_view(settings: FeatureStoreUploadSettings) -> Any: "feature_store_py is required; install requirements/feature_store.txt" ) from exc + try: + from alibabacloud_credentials.client import Client as CredClient + except ImportError as exc: + raise RuntimeError( + "alibabacloud_credentials is required; " + "install it via: pip install alibabacloud_credentials" + ) from exc + + credential = CredClient().get_credential() kwargs = { - "access_key_id": settings.access_key_id, - "access_key_secret": settings.access_key_secret, + "access_key_id": credential.access_key_id, + "access_key_secret": credential.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, + "security_token": credential.security_token or None, + "featuredb_username": os.environ.get("FEATUREDB_USERNAME") or None, + "featuredb_password": os.environ.get("FEATUREDB_PASSWORD") or None, } try: parameters = inspect.signature(FeatureStoreClient).parameters except (TypeError, ValueError): parameters = {} if "test_mode" in parameters: - # 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 7aa24960d..e62b8ac25 100644 --- a/tzrec/tools/feature_store/check_feature_store_delta_test.py +++ b/tzrec/tools/feature_store/check_feature_store_delta_test.py @@ -60,22 +60,33 @@ def __init__(self, test_mode=False, **kwargs): def get_project(self, name): return project + class FakeCredClient: + def get_credential(self): + return SimpleNamespace( + access_key_id="fake-ak", + access_key_secret="fake-sk", + security_token="fake-sts", + ) + 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 ) + cred_module = SimpleNamespace(Client=FakeCredClient) - with mock.patch.dict(sys.modules, {"feature_store_py": feature_store_module}): + with mock.patch.dict( + sys.modules, + { + "feature_store_py": feature_store_module, + "alibabacloud_credentials": SimpleNamespace(client=cred_module), + "alibabacloud_credentials.client": cred_module, + }, + ): actual = create_feature_store_view(settings) self.assertIs(actual, view) diff --git a/tzrec/utils/config_util.py b/tzrec/utils/config_util.py index cd757b002..cb7f2ff4c 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, train_pb2 +from tzrec.protos import data_pb2, pipeline_pb2 from tzrec.protos.data_pb2 import FgMode from tzrec.utils.logging_util import logger @@ -63,61 +63,6 @@ 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", -) - - -def sanitize_pipeline_config_for_artifact( - pipeline_config: pipeline_pb2.EasyRecConfig, -) -> pipeline_pb2.EasyRecConfig: - """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. - """ - 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)) - 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 c9721f01f..ab29f9d5d 100644 --- a/tzrec/utils/config_util_test.py +++ b/tzrec/utils/config_util_test.py @@ -9,8 +9,6 @@ # 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 @@ -42,45 +40,6 @@ def test_edit_config(self): ) self.assertEqual(pipeline_config.feature_configs[4].id_feature.num_buckets, 3) - def test_pipeline_artifact_redacts_feature_store_security_token(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.security_token = "SECRET_STS" - - 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()) - 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() - self.assertNotIn("SECRET_STS", artifact_text) - if __name__ == "__main__": unittest.main() diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index 2b2a36e3d..94a756ec3 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -21,7 +21,7 @@ import os import threading import time -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import ( Any, Callable, @@ -89,11 +89,6 @@ class FeatureStoreUploadSettings: 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 @@ -112,7 +107,7 @@ class FeatureStoreUploadSettings: @classmethod def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": - """Resolve environment-backed credentials without logging them.""" + """Validate configuration without resolving credentials.""" initialization_errors = config.FindInitializationErrors() if initialization_errors: raise ValueError( @@ -121,34 +116,6 @@ def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": ) region = config.region or os.environ.get("ALIBABA_CLOUD_REGION", "") endpoint = config.endpoint - 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( @@ -206,11 +173,6 @@ def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": 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, @@ -230,7 +192,7 @@ def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": def validate_feature_store_config(config: FeatureStoreConfig) -> None: - """Validate upload configuration and its environment-resolved credentials.""" + """Validate upload configuration without resolving credentials.""" FeatureStoreUploadSettings.from_proto(config) @@ -336,6 +298,7 @@ def __init__( world_size: int, embedding_dimensions: Mapping[str, int], client_factory: Optional[Callable[..., Any]] = None, + credentials_client: Optional[Any] = None, clock_ms: Optional[Callable[[], int]] = None, ) -> None: """Initialize the uploader with validated settings and in-memory state.""" @@ -363,6 +326,9 @@ def __init__( ) self._client_factory = client_factory + self._credentials_client = ( + credentials_client or self._create_credentials_client() + ) self._clock_ms = clock_ms or (lambda: time.time_ns() // 1_000_000) self._view = None self._condition = threading.Condition() @@ -795,6 +761,18 @@ def _cleanup_shard_files(self, shard_paths: List[str]) -> None: except OSError: pass + @staticmethod + def _create_credentials_client() -> Any: + """Create the Alibaba Cloud credential provider (default chain).""" + try: + from alibabacloud_credentials.client import Client as CredClient + except ImportError as exc: + raise RuntimeError( + "alibabacloud_credentials is required when feature_store_config " + "is set; install it via: pip install alibabacloud_credentials" + ) from exc + return CredClient() + def _get_view(self) -> Any: if self._view is not None: return self._view @@ -809,14 +787,15 @@ def _get_view(self) -> Any: else: client_factory = self._client_factory + credential = self._credentials_client.get_credential() kwargs = { - "access_key_id": self._settings.access_key_id, - "access_key_secret": self._settings.access_key_secret, + "access_key_id": credential.access_key_id, + "access_key_secret": credential.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, + "security_token": credential.security_token or None, + "featuredb_username": os.environ.get("FEATUREDB_USERNAME") or None, + "featuredb_password": os.environ.get("FEATUREDB_PASSWORD") or None, } 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 77fd053a5..fa3ae8fe8 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -42,7 +42,6 @@ def _schema_with_generation(generation: str = _TEST_DUMP_GENERATION) -> pa.Schem def _feature_store_config(**overrides) -> FeatureStoreConfig: config = FeatureStoreConfig( region="cn-test", - security_token="sts-secret", project_name="project_a", feature_entity_name="embedding_entity", feature_view_name="shared_embeddings", @@ -257,6 +256,17 @@ def create_dynamic_embedding_feature_view(self, **kwargs): return self._view +class _FakeCredential: + access_key_id = "fake-ak" + access_key_secret = "fake-sk" + security_token = "fake-sts" + + +class _FakeCredentialsClient: + def get_credential(self): + return _FakeCredential() + + class _FakeClient: def __init__(self, project, kwargs): self._project = project @@ -289,18 +299,13 @@ 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._cred_patch = mock.patch.object( + FeatureStoreDeltaUploader, + "_create_credentials_client", + return_value=_FakeCredentialsClient(), ) - self._credential_env.start() - self.addCleanup(self._credential_env.stop) + self._cred_patch.start() + self.addCleanup(self._cred_patch.stop) def test_proto_groups_required_fields_before_optional_fields(self): required_fields = [ @@ -312,7 +317,6 @@ def test_proto_groups_required_fields_before_optional_fields(self): ] optional_fields = [ "endpoint", - "security_token", "upload_batch_size", "max_retries", "retry_backoff_secs", @@ -324,7 +328,6 @@ def test_proto_groups_required_fields_before_optional_fields(self): "feature_view_shard_count", "feature_view_replication_count", "allow_custom_endpoint", - "upload_mode", ] fields = list(FeatureStoreConfig.DESCRIPTOR.fields) @@ -343,7 +346,10 @@ 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], [1, *list(range(6, 24))]) + self.assertEqual( + [field.number for field in fields], + [1, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22], + ) for field_name in required_fields: with self.subTest(field_name=field_name): config = _feature_store_config() @@ -376,38 +382,11 @@ def test_version_is_required_and_must_be_explicit(self): _feature_store_config(version="default") ) - def test_environment_resolution_and_required_credentials(self): + def test_region_fallback_and_config_validation(self): config = _feature_store_config(region="") - config.ClearField("security_token") - 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): + with mock.patch.dict(os.environ, {"ALIBABA_CLOUD_REGION": "cn-env"}): settings = FeatureStoreUploadSettings.from_proto(config) self.assertEqual(settings.region, "cn-env") - self.assertEqual(settings.access_key_id, "env-ak") - 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.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, "userinfo"): FeatureStoreUploadSettings.from_proto( From a432ca88087041da0e7ca47ffd76a38b4c0ecae5 Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 24 Jul 2026 11:07:32 +0800 Subject: [PATCH 26/50] [chore] move feature_store_py to runtime requirements with official wheel URL --- requirements/feature_store.txt | 1 - requirements/runtime.txt | 1 + setup.py | 1 - tzrec/tools/feature_store/check_feature_store_delta.py | 2 +- 4 files changed, 2 insertions(+), 3 deletions(-) delete mode 100644 requirements/feature_store.txt diff --git a/requirements/feature_store.txt b/requirements/feature_store.txt deleted file mode 100644 index d4390de4a..000000000 --- a/requirements/feature_store.txt +++ /dev/null @@ -1 +0,0 @@ -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/requirements/runtime.txt b/requirements/runtime.txt index c43dadbd2..b23a4a5e4 100644 --- a/requirements/runtime.txt +++ b/requirements/runtime.txt @@ -3,6 +3,7 @@ anytree common_io @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/common_io-0.4.1%2Btunnel-py2.py3-none-any.whl confluent-kafka fbgemm-gpu==1.7.0 +feature_store_py @ https://feature-store-py.oss-cn-beijing.aliyuncs.com/package/feature_store_py-2.2.7-py3-none-any.whl fsspec graphlearn @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/graphlearn/graphlearn-1.3.8-cp312-cp312-linux_x86_64.whl ; python_version=="3.12" graphlearn @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/graphlearn/graphlearn-1.3.8-cp311-cp311-linux_x86_64.whl ; python_version=="3.11" diff --git a/setup.py b/setup.py index 1134fb2e3..1fe47e75e 100644 --- a/setup.py +++ b/setup.py @@ -80,6 +80,5 @@ def parse_require_file(fpath): "tests": parse_requirements("requirements/test.txt"), "gpu": parse_requirements("requirements/cu130.txt"), "cpu": parse_requirements("requirements/cpu.txt"), - "feature_store": parse_requirements("requirements/feature_store.txt"), }, ) diff --git a/tzrec/tools/feature_store/check_feature_store_delta.py b/tzrec/tools/feature_store/check_feature_store_delta.py index c105d2994..5e4d46e85 100644 --- a/tzrec/tools/feature_store/check_feature_store_delta.py +++ b/tzrec/tools/feature_store/check_feature_store_delta.py @@ -342,7 +342,7 @@ def create_feature_store_view(settings: FeatureStoreUploadSettings) -> Any: from feature_store_py import FeatureStoreClient except ImportError as exc: raise RuntimeError( - "feature_store_py is required; install requirements/feature_store.txt" + "feature_store_py is required; install requirements/runtime.txt" ) from exc try: From 30621126ea03d89e49f679bb1c905718ac84b7ac Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 24 Jul 2026 11:11:00 +0800 Subject: [PATCH 27/50] [chore] drop reserved credential fields and renumber FeatureStoreConfig --- tzrec/protos/train.proto | 34 +++++++++---------- .../feature_store_delta_uploader_test.py | 2 +- 2 files changed, 17 insertions(+), 19 deletions(-) diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index e1d5d46d5..6bbeed601 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -33,20 +33,18 @@ message FeatureStoreConfig { // Cloud credentials (AK/SK/STS) are resolved at runtime through the // Alibaba Cloud default credential provider chain (alibabacloud_credentials). // FeatureDB credentials are read from FEATUREDB_USERNAME/FEATUREDB_PASSWORD. - reserved 2 to 5, 11; - reserved "access_key_id", "access_key_secret", "featuredb_username", "featuredb_password", "security_token"; // FeatureStore control-plane region. An explicitly empty value falls back // to ALIBABA_CLOUD_REGION at runtime. required string region = 1; // Existing FeatureStore project name and target DynamicEmbedding FeatureView // name. The view is validated at startup and created when it does not exist. - required string project_name = 6; - required string feature_view_name = 7; + required string project_name = 2; + required string feature_view_name = 3; // Existing FeatureStore entity associated with an automatically created // DynamicEmbedding FeatureView. - required string feature_entity_name = 8; + required string feature_entity_name = 4; // FeatureDB version for this incremental training run. - required string version = 9; + required string version = 5; // Optional FeatureStore control-plane endpoint override. Cloud and @@ -54,33 +52,33 @@ message FeatureStoreConfig { // *.aliyuncs.com FeatureStore endpoint (an explicit https:// scheme is // required when a scheme is given) without userinfo, port, path, query or // fragment components, unless allow_custom_endpoint is set. - optional string endpoint = 10; + optional string endpoint = 6; // Maximum records submitted before each SDK write_flush() completion gate. - optional uint32 upload_batch_size = 12 [default = 1000]; + optional uint32 upload_batch_size = 7 [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]; + optional uint32 max_retries = 8 [default = 3]; + optional uint32 retry_backoff_secs = 9 [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]; + optional uint32 shard_wait_timeout_secs = 10 [default = 600]; // Maximum time normal training shutdown waits for the uploader to drain. - optional uint32 shutdown_timeout_secs = 16 [default = 600]; + optional uint32 shutdown_timeout_secs = 11 [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]; + optional uint32 max_pending_steps = 12 [default = 32]; + optional uint32 poll_interval_secs = 13 [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]; + optional uint32 feature_view_ttl_secs = 14 [default = 1296000]; + optional uint32 feature_view_shard_count = 15 [default = 20]; + optional uint32 feature_view_replication_count = 16 [default = 1]; // Explicit opt-in allowing endpoint to name a non-*.aliyuncs.com host // (e.g. a vetted private FeatureStore deployment with a custom port). // Credentials are sent to that host; structural endpoint checks (HTTPS // scheme, no userinfo/path/query/fragment) still apply. - optional bool allow_custom_endpoint = 22 [default = false]; + optional bool allow_custom_endpoint = 17 [default = false]; } message DeltaEmbeddingDumpConfig { diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index fa3ae8fe8..604738dec 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -348,7 +348,7 @@ def test_proto_groups_required_fields_before_optional_fields(self): ) self.assertEqual( [field.number for field in fields], - [1, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22], + list(range(1, 18)), ) for field_name in required_fields: with self.subTest(field_name=field_name): From 2122a0185abd71dfaa9b2fea4468790d499c2220 Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 24 Jul 2026 11:16:09 +0800 Subject: [PATCH 28/50] [chore] restore save_message for exported pipeline config artifacts --- tzrec/utils/export_util.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tzrec/utils/export_util.py b/tzrec/utils/export_util.py index 0a76fbff7..377626271 100644 --- a/tzrec/utils/export_util.py +++ b/tzrec/utils/export_util.py @@ -347,7 +347,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_pipeline_config_artifact( + config_util.save_message( pipeline_config, os.path.join(save_dir, "pipeline.config") ) logger.info("saving fg json...") @@ -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_pipeline_config_artifact( + config_util.save_message( pipeline_config, os.path.join(save_dir, "pipeline.config") ) From 396af6e8a4f067946e1bfb08549b8e8eeeeec9d2 Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 24 Jul 2026 14:07:31 +0800 Subject: [PATCH 29/50] [refactor] upload FeatureStore deltas per rank from memory Rank shards no longer funnel through parquet files on a shared filesystem to a rank-zero uploader: each rank streams its own in-memory delta table to the SDK, coordinating only view creation and startup via a barrier. Per-rank monotonic timestamps suffice because sharding is fixed within a process, and data-parallel duplicate MERGE writes are idempotent. Local parquet files are now opt-in via feature_store_config.retain_local_dump, which the offline readback checker requires. --- tzrec/main.py | 13 +- tzrec/protos/train.proto | 34 +- .../check_feature_store_delta.py | 16 +- tzrec/utils/delta_embedding_dump.py | 137 +-- tzrec/utils/delta_embedding_dump_test.py | 47 +- tzrec/utils/feature_store_delta_uploader.py | 335 ++---- .../feature_store_delta_uploader_test.py | 1057 ++++++----------- 7 files changed, 576 insertions(+), 1063 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index 1590f2eaf..7f5e48aa9 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -892,9 +892,10 @@ def train_and_evaluate( ) 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. + # Keep the original training error primary. Pending in-memory + # deltas are abandoned; the restarted run re-dumps from the latest + # checkpoint. Do not wait up to the normal upload-drain timeout + # while unwinding a training failure. delta_embedding_dumper.close(raise_on_error=False, drain=False) raise else: @@ -904,9 +905,9 @@ def train_and_evaluate( 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. + # Each rank drains its own background queue at a different pace. + # Propagate a late drain/SDK error uniformly instead of letting a + # single rank fail alone at shutdown. _raise_if_any_worker_failed( delta_embedding_dumper_close_error, device, diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index 6bbeed601..625fe1ee1 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -55,30 +55,30 @@ message FeatureStoreConfig { optional string endpoint = 6; // Maximum records submitted before each SDK write_flush() completion gate. optional uint32 upload_batch_size = 7 [default = 1000]; - // Total attempts per 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. + // Total attempts per rank for one step. All retries reuse the version; + // each attempt reserves a newer monotonic ts range and fully replays the + // rank's delta so incremental readers cannot miss a partial earlier attempt. optional uint32 max_retries = 8 [default = 3]; optional uint32 retry_backoff_secs = 9 [default = 5]; - // Maximum time 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 = 10 [default = 600]; // Maximum time normal training shutdown waits for the uploader to drain. - optional uint32 shutdown_timeout_secs = 11 [default = 600]; - // Apply back-pressure when too many completed dumps await publication. - optional uint32 max_pending_steps = 12 [default = 32]; - optional uint32 poll_interval_secs = 13 [default = 1]; + optional uint32 shutdown_timeout_secs = 10 [default = 600]; + // Apply back-pressure when too many completed dumps await upload; bounds + // the per-rank memory retained by pending in-memory deltas. + optional uint32 max_pending_steps = 11 [default = 32]; + optional uint32 poll_interval_secs = 12 [default = 1]; // Creation settings used only when feature_view_name does not yet exist. - optional uint32 feature_view_ttl_secs = 14 [default = 1296000]; - optional uint32 feature_view_shard_count = 15 [default = 20]; - optional uint32 feature_view_replication_count = 16 [default = 1]; + optional uint32 feature_view_ttl_secs = 13 [default = 1296000]; + optional uint32 feature_view_shard_count = 14 [default = 20]; + optional uint32 feature_view_replication_count = 15 [default = 1]; // Explicit opt-in allowing endpoint to name a non-*.aliyuncs.com host // (e.g. a vetted private FeatureStore deployment with a custom port). // Credentials are sent to that host; structural endpoint checks (HTTPS // scheme, no userinfo/path/query/fragment) still apply. - optional bool allow_custom_endpoint = 17 [default = false]; + optional bool allow_custom_endpoint = 16 [default = false]; + // Also write the local per-rank delta parquet files while uploading to + // FeatureStore (e.g. for the offline readback checker); by default the + // delta is handed to the uploader in memory and never touches disk. + optional bool retain_local_dump = 17 [default = false]; } message DeltaEmbeddingDumpConfig { @@ -92,7 +92,7 @@ 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. + // Presence enables best-effort per-rank background upload to FeatureStore. optional FeatureStoreConfig feature_store_config = 4; // Dump after this many elapsed minutes. The timer starts when training starts. // Do not set this together with dump_interval_steps. diff --git a/tzrec/tools/feature_store/check_feature_store_delta.py b/tzrec/tools/feature_store/check_feature_store_delta.py index 5e4d46e85..541d22757 100644 --- a/tzrec/tools/feature_store/check_feature_store_delta.py +++ b/tzrec/tools/feature_store/check_feature_store_delta.py @@ -11,10 +11,13 @@ r"""Read back sampled delta embeddings from FeatureStore. -The tool scans the local delta parquet outbox for the latest (or specified) +The tool scans the local delta parquet output for the latest (or specified) step, samples keys from its shard set, and queries the configured explicit FeatureDB version through ``DynamicEmbeddingFeatureView.get_online_features``. +Local parquet files are only produced alongside FeatureStore uploads when +``feature_store_config.retain_local_dump`` is enabled. + Example:: export ALIBABA_CLOUD_ACCESS_KEY_ID=YOUR_ACCESS_KEY_ID @@ -47,7 +50,6 @@ FEATURE_STORE_SK_FIELD, FEATURE_STORE_VALUE_FIELD, FeatureStoreUploadSettings, - feature_store_delta_file_prefix, ) @@ -407,6 +409,11 @@ 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 + if not feature_store_config.retain_local_dump: + raise ValueError( + "feature_store_config.retain_local_dump must be enabled so the " + "training job keeps local delta parquet files for readback" + ) settings = FeatureStoreUploadSettings.from_proto(feature_store_config) output_dir = resolve_output_dir( args.pipeline_config, @@ -419,11 +426,10 @@ def run_check(args: argparse.Namespace) -> int: f"delta embedding output directory not found: {output_dir}" ) - base_prefix = dump_config.file_prefix or "delta_embedding" - scoped_prefix = feature_store_delta_file_prefix(feature_store_config, base_prefix) + file_prefix = dump_config.file_prefix or "delta_embedding" world_size = args.world_size global_step, parquet_paths = resolve_upload_step( - output_dir, scoped_prefix, world_size, args.global_step + output_dir, file_prefix, world_size, args.global_step ) samples = sample_local_records( parquet_paths, diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 0ea7feeef..d6be8839a 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -9,11 +9,9 @@ # 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 @@ -32,10 +30,7 @@ from tzrec.protos.train_pb2 import DeltaEmbeddingDumpConfig from tzrec.utils.feature_store_delta_uploader import ( - DELTA_DUMP_GENERATION_METADATA_KEY, - DELTA_DUMP_SCHEMA_VERSION, FeatureStoreDeltaUploader, - feature_store_delta_file_prefix, validate_feature_store_config, ) from tzrec.utils.logging_util import logger @@ -51,6 +46,7 @@ ) _CONSUMER = "delta_embedding_dump" +DELTA_DUMP_SCHEMA_VERSION = "2" _DELTA_DUMP_SCHEMA = pa.schema( [ ("global_step", pa.int64()), @@ -389,17 +385,14 @@ def __init__( self._output_dir = config.output_dir or os.path.join( model_dir, "delta_embedding_dump" ) - file_prefix = config.file_prefix or "delta_embedding" - self._file_prefix = file_prefix + self._file_prefix = config.file_prefix or "delta_embedding" self._rank, self._world_size = _distributed_rank_world_size() self._device = device self._tracking_pause_depth = 0 self._feature_store_enabled = config.HasField("feature_store_config") - self._run_generation: Optional[str] = None - if self._feature_store_enabled: - self._file_prefix = feature_store_delta_file_prefix( - config.feature_store_config, self._file_prefix - ) + self._retain_local_dump = self._feature_store_enabled and bool( + config.feature_store_config.retain_local_dump + ) os.makedirs(self._output_dir, exist_ok=True) self._table_shard_infos = self._collect_table_shard_infos() @@ -423,13 +416,12 @@ def __init__( self._build_sparse_embedding_contract() ) self._uploader: Optional[FeatureStoreDeltaUploader] = None - if self._feature_store_enabled and self._rank == 0: + if self._feature_store_enabled: 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, + rank=self._rank, + manage_remote_view=self._rank == 0, ) interval_name = "minutes" if self._interval_secs is not None else "steps" @@ -455,23 +447,47 @@ def clear(self) -> None: self._tracker.clear() def start(self) -> None: - """Start timed cadence and rank-zero publication after initialization.""" - if self._feature_store_enabled: - self._initialize_run_generation() + """Start timed cadence and per-rank FeatureStore publication. + + The rank-zero uploader creates and validates the remote view first; + the other ranks start their uploaders only after the barrier so they + can open the already-published view without control-plane races. A + rank-zero startup failure still joins the barrier before raising so + every rank keeps issuing the same collective sequence. + """ if self._uploader is not None: - self._uploader.start() + start_error: Optional[BaseException] = None + if self._rank == 0: + try: + self._uploader.start() + except BaseException as exc: + start_error = exc + 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" + ) + torch.distributed.barrier() + if start_error is not None: + raise start_error.with_traceback(start_error.__traceback__) + if self._rank != 0: + 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.""" + """Close this rank's 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 uploader error without changing rank control flow.""" + """Collect this rank's uploader error without changing control flow.""" if not self._feature_store_enabled: return None @@ -483,40 +499,11 @@ def _feature_store_upload_error( return None def _check_feature_store_upload_error(self, force: bool = False) -> None: - """Surface the rank-zero background failure through the shared outbox.""" + """Surface this rank's background upload failure to the trainer.""" 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 self._feature_store_enabled: - 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]]: @@ -799,7 +786,7 @@ def dump(self, global_step: int) -> Optional[str]: if global_step <= 0: raise ValueError("delta embedding dump global_step must be > 0") uploader = self._uploader - dump_generation = self._next_dump_generation(global_step) + write_local = not self._feature_store_enabled or self._retain_local_dump tracker_cursor = self._tracker_cursor_before_read() try: table_weights = self._collect_table_weights() @@ -811,27 +798,31 @@ def dump(self, global_step: int) -> Optional[str]: table_weights=table_weights, dynamic_modules=dynamic_modules, ) - if ( - num_rows == 0 - and self._world_size == 1 - and not self._feature_store_enabled - ): - 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, dump_generation=dump_generation - ) - if uploader is not None: - uploader.submit(global_step) + output_path: Optional[str] = None + if write_local and (num_rows > 0 or self._world_size > 1): + # Multi-rank shard sets stay complete even for an empty rank so + # per-step file consumers never observe a partial set. + output_path = self._output_path(global_step) + self._write_table_chunks(table_chunks, output_path) + if uploader is not None and num_rows > 0: + uploader.submit(global_step, pa.concat_tables(table_chunks)) except BaseException: self._rollback_tracker_read(tracker_cursor) raise if num_rows == 0: + if output_path is None: + logger.info("No delta embedding rows to dump at step %s.", global_step) + else: + logger.info( + "Dumped empty delta embedding shard to %s at step %s.", + output_path, + global_step, + ) + elif output_path is None: logger.info( - "Dumped empty delta embedding shard to %s at step %s.", - output_path, + "Submitted %s delta embedding rows for FeatureStore upload at step %s.", + num_rows, global_step, ) else: @@ -1190,18 +1181,10 @@ def _write_table_chunks( self, table_chunks: List[pa.Table], output_path: str, - dump_generation: Optional[str] = None, ) -> None: tmp_path = f"{output_path}.rank{self._rank}.tmp" try: - 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: + with pq.ParquetWriter(tmp_path, _DELTA_DUMP_SCHEMA) as writer: chunks = table_chunks or [_DELTA_DUMP_SCHEMA.empty_table()] for table_chunk in chunks: writer.write_table(table_chunk) diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index a2a80c798..17af163d0 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -56,9 +56,6 @@ ) from tzrec.utils.dist_util import create_train_pipeline from tzrec.utils.dynamicemb_util import has_dynamicemb -from tzrec.utils.feature_store_delta_uploader import ( - DELTA_DUMP_GENERATION_METADATA_KEY, -) from tzrec.utils.test_util import gpu_unavailable, make_test_dir, mark_ci_scope _SHARDED_TABLE_NAME = "table_1" @@ -471,19 +468,6 @@ 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 @@ -584,18 +568,6 @@ def fake_all_reduce(tensor, op=None): self.assertEqual(collective_count, 1) dump_mock.assert_not_called() - 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_steps = 50 @@ -876,26 +848,31 @@ def test_minutes_interval_is_converted_to_seconds(self): def test_dump_does_not_ack_tracker_when_submission_is_rejected(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._feature_store_enabled = True + dumper._retain_local_dump = False 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") + def _append_one_chunk(table_chunks, **kwargs): + table_chunks.append(_DELTA_DUMP_SCHEMA.empty_table()) + return 1 + 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, "_append_model_delta_rows", side_effect=_append_one_chunk + ), mock.patch.object(dumper, "_rollback_tracker_read") as rollback, ): with self.assertRaisesRegex(RuntimeError, "submission rejected"): dumper.dump(10) - dumper._uploader.submit.assert_called_once_with(10) + dumper._uploader.submit.assert_called_once() + self.assertEqual(dumper._uploader.submit.call_args.args[0], 10) rollback.assert_called_once_with(7) def test_multi_gpu_output_path_uses_step_underscore_dir(self): @@ -925,7 +902,7 @@ def test_multi_gpu_dump_writes_empty_shard_when_rank_has_no_delta(self): dumper._world_size = 2 dumper._feature_store_enabled = False dumper._uploader = None - dumper._run_generation = None + dumper._retain_local_dump = False dumper._tracker = mock.MagicMock() dumper._tracker.per_consumer_batch_idx = {_CONSUMER: 0} with ( @@ -956,7 +933,7 @@ def test_single_gpu_dump_skips_file_when_rank_has_no_delta(self): dumper._world_size = 1 dumper._feature_store_enabled = False dumper._uploader = None - dumper._run_generation = None + dumper._retain_local_dump = False dumper._tracker = mock.MagicMock() dumper._tracker.per_consumer_batch_idx = {_CONSUMER: 0} with ( diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index 94a756ec3..c651d4fe2 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -9,14 +9,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Ephemeral rank-zero uploader for delta-embedding parquet shards. +"""Ephemeral per-rank uploader for in-memory delta-embedding tables. Best-effort upload for the current live training process only. No cross-restart recovery, no durable state, no replay. A process crash means restart from the latest checkpoint and pending deltas are discarded. """ -import hashlib import json import os import threading @@ -36,7 +35,6 @@ 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 @@ -50,25 +48,8 @@ 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" _FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES = 100 -_VERSION_INITIALIZATION = "AUTO_CREATE_ON_FIRST_DELTA_MERGE" - -_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()), -} class FeatureStoreUploadError(RuntimeError): @@ -79,10 +60,6 @@ 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.""" - - @dataclass(frozen=True) class FeatureStoreUploadSettings: """Validated immutable settings copied from the runtime protobuf.""" @@ -99,7 +76,6 @@ class FeatureStoreUploadSettings: 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 @@ -145,7 +121,6 @@ def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": "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), @@ -183,7 +158,6 @@ def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": 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"], @@ -246,70 +220,35 @@ def _validate_feature_store_endpoint( ) -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 _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]}" - - class FeatureStoreDeltaUploader: - """Ephemeral in-process uploader for delta-embedding parquet shards. + """Ephemeral per-rank uploader for in-memory delta-embedding tables. Best-effort upload for the current live process only. No cross-restart recovery, no durable state, no replay. A process crash means restart from the latest checkpoint and pending deltas are discarded. - The training thread writes parquet shards and enqueues a step. This - object's single background worker waits for the complete shard set, - streams parquet batches directly to the FeatureStore SDK, and cleans up - local files on success. + Every rank owns one uploader and streams only its local shard rows, so no + cross-rank aggregation or deduplication is needed: table-wise and row-wise + sharding give each (embedding_name, key_id) a unique owner rank, while + data-parallel replicas issue identical idempotent MERGE writes. Publish + timestamps are allocated per rank and only need per-rank monotonicity, + because a key's owner rank is fixed for the lifetime of the process. """ def __init__( self, config: FeatureStoreConfig, - output_dir: str, - file_prefix: str, - world_size: int, embedding_dimensions: Mapping[str, int], + rank: int = 0, + manage_remote_view: bool = True, client_factory: Optional[Callable[..., Any]] = None, credentials_client: Optional[Any] = None, clock_ms: Optional[Callable[[], int]] = None, ) -> None: """Initialize the uploader with validated settings and in-memory state.""" 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._rank = int(rank) + self._manage_remote_view = bool(manage_remote_view) self._embedding_dimensions = { str(name): int(dimension) for name, dimension in embedding_dimensions.items() @@ -332,7 +271,7 @@ def __init__( 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._pending: Dict[int, pa.Table] = {} self._started = False self._closing = False self._aborting = False @@ -363,14 +302,15 @@ def start(self) -> None: self._worker.start() logger.info( "FeatureStore delta uploader started: project=%s feature_view=%s " - "version=%s", + "version=%s rank=%s", self._settings.project_name, self._settings.feature_view_name, self._settings.version, + self._rank, ) - def submit(self, global_step: int) -> None: - """Enqueue a completed step for background upload with back-pressure.""" + def submit(self, global_step: int, table: pa.Table) -> None: + """Enqueue one step's in-memory delta table with back-pressure.""" global_step = int(global_step) if global_step <= 0: raise ValueError("FeatureStore delta global_step must be > 0") @@ -388,7 +328,7 @@ def submit(self, global_step: int) -> None: ): self._condition.wait(self._settings.poll_interval_secs) self._raise_if_failed_locked() - self._pending.setdefault(global_step, time.monotonic()) + self._pending.setdefault(global_step, table) self._condition.notify_all() def check_error(self) -> None: @@ -441,28 +381,9 @@ def _run(self) -> None: self._condition.wait(self._settings.poll_interval_secs) continue current_step = min(self._pending) - pending_since = self._pending[current_step] - - shard_paths = self._expected_shard_paths(current_step) - if not all(os.path.isfile(path) for path in shard_paths): - elapsed = time.monotonic() - pending_since - if elapsed >= self._settings.shard_wait_timeout_secs: - raise TimeoutError( - "timed out waiting for a complete delta shard set " - f"at step {current_step}" - ) - with self._condition: - self._condition.wait(self._settings.poll_interval_secs) - continue + table = self._pending[current_step] - self._validate_shard_generation(shard_paths) - - with self._condition: - if self._aborting: - return - - self._upload_with_retries(current_step, shard_paths) - self._cleanup_shard_files(shard_paths) + self._upload_with_retries(current_step, table) with self._condition: self._pending.pop(current_step, None) @@ -499,43 +420,11 @@ def _raise_if_aborting(self) -> None: if self._aborting: raise _UploadAborted() - 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 _validate_shard_generation(self, shard_paths: List[str]) -> None: - """Verify all shards share the same dump generation.""" - generations = set() - for path in shard_paths: - parquet_file = pq.ParquetFile(path) - metadata = parquet_file.schema_arrow.metadata or {} - generation = metadata.get(DELTA_DUMP_GENERATION_METADATA_KEY) - generations.add(generation) - if len(generations) > 1: - raise _ShardSetNotReady( - "delta shard set has inconsistent dump generations; " - "a concurrent dump may be replacing files" - ) - - def _upload_with_retries(self, global_step: int, shard_paths: List[str]) -> None: + def _upload_with_retries(self, global_step: int, table: pa.Table) -> None: for attempt in range(1, self._settings.max_retries + 1): self._raise_if_aborting() try: - self._stream_upload(global_step, shard_paths) + self._stream_upload(global_step, table) return except _UploadAborted: raise @@ -556,19 +445,24 @@ def _upload_with_retries(self, global_step: int, shard_paths: List[str]) -> None raise AssertionError("unreachable FeatureStore retry state") def _allocate_timestamp_range(self, batch_count: int) -> Tuple[int, int]: - """Allocate a monotonically increasing timestamp range (in-memory only).""" + """Allocate a rank-locally monotonic timestamp range (in-memory only). + + Per-rank monotonicity is sufficient for Next-Ts incremental readers: + sharding is fixed for the lifetime of the process, so every key is + always republished by the same rank with a strictly newer timestamp. + """ reserved = max(batch_count, 1) range_start = max(int(self._clock_ms()), self._last_publish_ts + 1, 1) range_end = range_start + reserved - 1 self._last_publish_ts = range_end return range_start, range_end - def _stream_upload(self, global_step: int, shard_paths: List[str]) -> None: - """Stream parquet batches directly to the FeatureStore SDK.""" + def _stream_upload(self, global_step: int, table: pa.Table) -> None: + """Stream the in-memory delta table directly to the FeatureStore SDK.""" view = self._get_view() max_in_flight = int(getattr(view, "_max_workers", 1)) - total_batches = self._count_total_batches(shard_paths) + total_batches = self._count_total_batches(table) ts_range = self._allocate_timestamp_range(total_batches) range_start = ts_range[0] @@ -580,9 +474,10 @@ def _stream_upload(self, global_step: int, shard_paths: List[str]) -> None: logged_first_window = False logger.info( - "FeatureStore delta upload started: step=%s version=%s " + "FeatureStore delta upload started: step=%s rank=%s version=%s " "batches=%s ts_range=%s-%s", global_step, + self._rank, self._settings.version, total_batches, ts_range[0], @@ -590,59 +485,51 @@ def _stream_upload(self, global_step: int, shard_paths: List[str]) -> None: ) try: - for expected_rank, shard_path in enumerate(shard_paths): - parquet_file = pq.ParquetFile(shard_path) - self._validate_parquet_schema(parquet_file.schema_arrow, shard_path) - for batch in parquet_file.iter_batches( - batch_size=self._settings.upload_batch_size + for batch in table.to_batches( + max_chunksize=self._settings.upload_batch_size + ): + self._raise_if_aborting() + payload = self._validate_and_build_payload(batch) + if not payload: + continue + view.write_features( + data=payload, + version=self._settings.version, + write_mode=FEATURE_STORE_WRITE_MODE, + ts=range_start + completed_batches, + ) + completed_batches += 1 + window_batches += 1 + window_records += len(payload) + + if window_batches < max_in_flight and completed_batches < total_batches: + continue + summary = view.write_flush() + self._validate_flush_summary( + summary, + expected_records=window_records, + expected_batches=window_batches, + ) + if ( + not logged_first_window + or completed_batches >= next_progress_batch + or completed_batches == total_batches ): - self._raise_if_aborting() - payload = self._validate_and_build_payload( - batch, global_step, expected_rank - ) - if not payload: - continue - view.write_features( - data=payload, - version=self._settings.version, - write_mode=FEATURE_STORE_WRITE_MODE, - ts=range_start + completed_batches, + logger.info( + "FeatureStore delta upload progress: step=%s " + "batches=%s/%s elapsed_secs=%.1f", + global_step, + completed_batches, + total_batches, + time.monotonic() - started_at, ) - completed_batches += 1 - window_batches += 1 - window_records += len(payload) - - if ( - window_batches < max_in_flight - and completed_batches < total_batches - ): - continue - summary = view.write_flush() - self._validate_flush_summary( - summary, - expected_records=window_records, - expected_batches=window_batches, - ) - if ( - not logged_first_window - or completed_batches >= next_progress_batch - or completed_batches == total_batches - ): - logger.info( - "FeatureStore delta upload progress: step=%s " - "batches=%s/%s elapsed_secs=%.1f", - global_step, - completed_batches, - total_batches, - time.monotonic() - started_at, + logged_first_window = True + while next_progress_batch <= completed_batches: + next_progress_batch += ( + _FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES ) - logged_first_window = True - while next_progress_batch <= completed_batches: - next_progress_batch += ( - _FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES - ) - window_batches = 0 - window_records = 0 + window_batches = 0 + window_records = 0 if window_batches > 0: summary = view.write_flush() @@ -665,12 +552,9 @@ def _stream_upload(self, global_step: int, shard_paths: List[str]) -> None: time.monotonic() - started_at, ) - def _count_total_batches(self, shard_paths: List[str]) -> int: - """Count the total number of upload batches across all shards.""" - total_rows = 0 - for path in shard_paths: - parquet_file = pq.ParquetFile(path) - total_rows += int(parquet_file.metadata.num_rows) + def _count_total_batches(self, table: pa.Table) -> int: + """Count the number of upload batches in one step's delta table.""" + total_rows = int(table.num_rows) if total_rows == 0: return 1 return ( @@ -680,23 +564,11 @@ def _count_total_batches(self, shard_paths: List[str]) -> int: def _validate_and_build_payload( self, batch: pa.RecordBatch, - global_step: int, - expected_rank: int, ) -> List[Dict[str, Any]]: - """Validate one streamed shard batch and build the SDK payload.""" + """Validate one delta batch and build the SDK payload.""" num_rows = batch.num_rows if num_rows == 0: return [] - global_steps = batch.column("global_step").to_numpy(zero_copy_only=False) - if not bool((global_steps == global_step).all()): - raise ValueError("delta shard global_step mismatch") - ranks = batch.column("rank").to_numpy(zero_copy_only=False) - if not bool((ranks == expected_rank).all()): - raise ValueError("delta shard rank mismatch") - world_sizes = batch.column("world_size").to_numpy(zero_copy_only=False) - if not bool((world_sizes == self._world_size).all()): - raise ValueError("delta shard world_size mismatch") - batch_roles = set(batch.column("embedding_role").to_pylist()) if not batch_roles <= set(SPARSE_EMBEDDING_ROLES): raise ValueError("delta shard has an invalid embedding role") @@ -747,20 +619,6 @@ def _validate_and_build_payload( for i in range(num_rows) ] - def _cleanup_shard_files(self, shard_paths: List[str]) -> None: - """Remove uploaded shard files (best-effort).""" - for path in shard_paths: - try: - os.remove(path) - except OSError: - logger.warning("Failed to remove uploaded shard file: %s", path) - if self._world_size > 1 and shard_paths: - step_dir = os.path.dirname(shard_paths[0]) - try: - os.rmdir(step_dir) - except OSError: - pass - @staticmethod def _create_credentials_client() -> Any: """Create the Alibaba Cloud credential provider (default chain).""" @@ -850,7 +708,26 @@ def _reset_view(self, suppress_errors: bool = False) -> None: ) def _get_or_create_view(self, project: Any) -> Any: - """Return the configured DynamicEmbedding view, creating it if absent.""" + """Return the configured DynamicEmbedding view, creating it if absent. + + Only the primary (rank-zero) uploader creates the view and validates + control-plane metadata; other ranks open a handle to the view that the + primary published before they started. + """ + if not self._manage_remote_view: + view = project.get_dynamic_embedding_feature_view( + self._settings.feature_view_name + ) + if view is None: + view = self._wait_for_dynamic_embedding_view(project) + if view is None: + raise RuntimeError( + "configured DynamicEmbedding FeatureView was not found; " + "the rank-zero uploader must create it before other ranks " + "start" + ) + self._view = view + return view provisioned = False view = project.get_dynamic_embedding_feature_view( self._settings.feature_view_name @@ -1063,17 +940,3 @@ def _validate_flush_summary( or int(summary["total_records"]) != expected_records ): raise RuntimeError("FeatureStore write_flush reported incomplete writes") - - @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}" - ) diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index 604738dec..650ee7904 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -11,33 +11,21 @@ import json import os -import tempfile import threading import unittest 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" - - -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( @@ -49,7 +37,6 @@ def _feature_store_config(**overrides) -> FeatureStoreConfig: 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, @@ -80,42 +67,10 @@ def _row( } -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 +def _delta_table(rows) -> pa.Table: + if rows: + return pa.Table.from_pylist(rows, schema=_DELTA_DUMP_SCHEMA) + return _DELTA_DUMP_SCHEMA.empty_table() class _FakeView: @@ -307,6 +262,10 @@ def setUp(self): self._cred_patch.start() self.addCleanup(self._cred_patch.stop) + def _uploader(self, config=None, **kwargs): + kwargs.setdefault("embedding_dimensions", {"user_emb": 2}) + return FeatureStoreDeltaUploader(config or _feature_store_config(), **kwargs) + def test_proto_groups_required_fields_before_optional_fields(self): required_fields = [ "region", @@ -320,7 +279,6 @@ def test_proto_groups_required_fields_before_optional_fields(self): "upload_batch_size", "max_retries", "retry_backoff_secs", - "shard_wait_timeout_secs", "shutdown_timeout_secs", "max_pending_steps", "poll_interval_secs", @@ -328,6 +286,7 @@ def test_proto_groups_required_fields_before_optional_fields(self): "feature_view_shard_count", "feature_view_replication_count", "allow_custom_endpoint", + "retain_local_dump", ] fields = list(FeatureStoreConfig.DESCRIPTOR.fields) @@ -460,389 +419,265 @@ def test_allow_custom_endpoint_opts_into_vetted_hosts(self): ) 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, - ) + view = _FakeView() + factory = _FakeClientFactory(view) + uploader = self._uploader(client_factory=factory) - uploader.start() - uploader.close() + 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.assertNotIn("test_mode", factory.calls[0]) - self.assertEqual(view.closed, [True]) + 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): - 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, - ) + created_view = _FakeView() + factory = _FakeClientFactory(None, created_view=created_view) + uploader = self._uploader(client_factory=factory) - uploader.start() - uploader.close() + 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.assertEqual(created_view.closed, [True]) + 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.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, - ) + generic_view = mock.Mock(type="Batch") + factory = _FakeClientFactory(None, generic_view=generic_view) + uploader = self._uploader(client_factory=factory) - with self.assertRaisesRegex(RuntimeError, "incompatible type"): - uploader.start() + with self.assertRaisesRegex(RuntimeError, "incompatible type"): + uploader.start() - self.assertEqual(factory.project.create_calls, []) + 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, - ) + view = _FakeView() + generic_view = _FakeGenericFeatureView(entity="another_entity") + factory = _FakeClientFactory(view, generic_view=generic_view) + uploader = self._uploader(client_factory=factory) - with self.assertRaisesRegex(RuntimeError, "entity mismatch"): - uploader.start() + with self.assertRaisesRegex(RuntimeError, "entity mismatch"): + uploader.start() - self.assertEqual(factory.project.create_calls, []) - self.assertEqual(view.closed, [True]) + 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, - ) + view = _FakeView() + generic_view = _FakeGenericFeatureView() + generic_view.fields_dict["embedding"]["Type"] = "ARRAY" + factory = _FakeClientFactory(view, generic_view=generic_view) + uploader = self._uploader(client_factory=factory) - with self.assertRaisesRegex(RuntimeError, "field contract mismatch"): - uploader.start() + with self.assertRaisesRegex(RuntimeError, "field contract mismatch"): + uploader.start() - self.assertEqual(factory.project.create_calls, []) - self.assertEqual(view.closed, [True]) + 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, - ) + view = _FakeView() + generic_view = _FakeGenericFeatureView( + provisioning={ + "ttl": 60, + "shard_count": 20, + "replication_count": 1, + } + ) + factory = _FakeClientFactory(view, generic_view=generic_view) + uploader = self._uploader(client_factory=factory) - with self.assertRaisesRegex(RuntimeError, "provisioning mismatch"): - uploader.start() + with self.assertRaisesRegex(RuntimeError, "provisioning mismatch"): + uploader.start() - self.assertEqual(factory.project.create_calls, []) - self.assertEqual(view.closed, [True]) + 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, - ) + concurrent_view = _FakeView() + factory = _FakeClientFactory( + None, + create_error=RuntimeError("already exists"), + view_after_create_error=concurrent_view, + ) + uploader = self._uploader(client_factory=factory) - uploader.start() - uploader.close() + 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]) + 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, - ) + created_view = _FakeView() + created_view.pk_field = "wrong_pk" + factory = _FakeClientFactory(None, created_view=created_view) + uploader = self._uploader(client_factory=factory) - with self.assertRaisesRegex(RuntimeError, "schema mismatch"): - uploader.start() + with self.assertRaisesRegex(RuntimeError, "schema mismatch"): + uploader.start() - self.assertEqual(len(factory.project.create_calls), 1) - self.assertEqual(created_view.closed, [True]) + self.assertEqual(len(factory.project.create_calls), 1) + self.assertEqual(created_view.closed, [True]) def test_start_creates_missing_view_without_version_precheck(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, - ) + created_view = _FakeView() + factory = _FakeClientFactory(None, created_view=created_view) + uploader = self._uploader(client_factory=factory) - uploader.start() - uploader.close() + uploader.start() + uploader.close() - self.assertEqual(len(factory.project.create_calls), 1) - self.assertEqual(created_view.closed, [True]) + 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, - ) + factory = _FakeClientFactory(None, create_error=ValueError("Entity not found")) + uploader = self._uploader(client_factory=factory) + + with self.assertRaisesRegex(RuntimeError, "feature_entity_name"): + uploader.start() + + self.assertEqual(len(factory.project.create_calls), 1) + + def test_non_primary_uploader_opens_view_without_create_or_metadata_checks(self): + view = _FakeView() + factory = _FakeClientFactory(view) + uploader = self._uploader( + rank=1, + manage_remote_view=False, + client_factory=factory, + clock_ms=lambda: 100, + ) - with self.assertRaisesRegex(RuntimeError, "feature_entity_name"): - uploader.start() + uploader.start() + uploader.submit(10, _delta_table([_row(10, 1, 7, [1.0, 2.0])])) + uploader.close() + + self.assertEqual(factory.project.dynamic_get_calls, ["shared_embeddings"]) + self.assertEqual(factory.project.generic_get_calls, []) + self.assertEqual(factory.project.create_calls, []) + self.assertEqual(len(view.calls), 1) + self.assertEqual(view.calls[0]["data"][0]["key_id"], 7) + + def test_non_primary_uploader_fails_when_view_is_missing(self): + factory = _FakeClientFactory(None) + uploader = self._uploader( + rank=1, + manage_remote_view=False, + client_factory=factory, + ) + + with self.assertRaisesRegex(RuntimeError, "rank-zero uploader must create"): + uploader.start() - self.assertEqual(len(factory.project.create_calls), 1) + self.assertEqual(factory.project.create_calls, []) + self.assertEqual(factory.project.generic_get_calls, []) 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() + uploader = self._uploader(client_factory=_FakeClientFactory(_FakeView())) + with self.assertRaisesRegex(RuntimeError, "start.*before submit"): + uploader.submit(10, _delta_table([_row(10, 0, 1, [1.0, 2.0])])) + 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, + view = _FakeView() + factory = _FakeClientFactory(view) + uploader = self._uploader( + client_factory=factory, + clock_ms=lambda: 123456, + ) + uploader.start() + uploader.submit( + 10, + _delta_table( [ _row(10, 0, 1, [1.0, 2.0]), _row(10, 0, 2, [3.0, 4.0]), _row(10, 0, 3, [0.0, 0.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() + ] + ), + ) + uploader.close() - 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"} - ) - 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]) + 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"}) + 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]) 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, - ) + view = _FakeView(max_workers=2) + uploader = self._uploader( + _feature_store_config(upload_batch_size=1), + client_factory=_FakeClientFactory(view), + clock_ms=lambda: 100, + ) - uploader.start() - uploader.submit(10) - uploader.close() + uploader.start() + uploader.submit( + 10, _delta_table([_row(10, 0, key, [1.0, 2.0]) for key in range(1, 6)]) + ) + 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]]) + 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(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, - ) + view = _FakeView() + uploader = self._uploader( + client_factory=_FakeClientFactory(view), + clock_ms=lambda: 100, + ) - uploader.start() - uploader.submit(1) - uploader.close() + uploader.start() + uploader.submit(1, _delta_table([_row(1, 0, 1, [1.0, 2.0])])) + 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.assertEqual(len(view.calls), 1) + self.assertEqual(view.calls[0]["ts"], 100) + self.assertEqual(view.calls[0]["version"], "model_a@export_1") 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() - 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 = self._uploader(client_factory=_FakeClientFactory(_FakeView())) - uploader.start() + uploader.start() + try: + with self.assertRaisesRegex(ValueError, "global_step must be > 0"): + uploader.submit(0, _delta_table([])) + finally: uploader.close() - self.assertEqual(len(factory.calls), 1) - self.assertEqual(view.calls, []) - def test_flush_failure_raises_error(self): failed_summary = { "total_batches": 2, @@ -852,30 +687,25 @@ def test_flush_failure_raises_error(self): "failed_records": 1, "errors": ["failed future"], } - with tempfile.TemporaryDirectory() as output_dir: - _write_single_shard( - output_dir, - 10, + view = _FakeView([failed_summary]) + uploader = self._uploader( + _feature_store_config(max_retries=1), + client_factory=_FakeClientFactory(view), + ) + uploader.start() + uploader.submit( + 10, + _delta_table( [ _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, - "delta", - 1, - {"user_emb": 2}, - client_factory=_FakeClientFactory(view), - ) - uploader.start() - uploader.submit(10) - with self.assertRaises(FeatureStoreUploadError): - uploader.close() - self.assertEqual(view.flush_calls, [[2, 1], []]) + ] + ), + ) + with self.assertRaises(FeatureStoreUploadError): + uploader.close() + self.assertEqual(view.flush_calls, [[2, 1], []]) def test_retry_uses_fresh_view_and_newer_timestamp_range(self): failed_summary = { @@ -886,234 +716,118 @@ def test_retry_uses_fresh_view_and_newer_timestamp_range(self): "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() + first_view = _FakeView([failed_summary]) + second_view = _FakeView() + factory = _SequencedClientFactory([first_view, second_view]) + uploader = self._uploader( + _feature_store_config(max_retries=2), + client_factory=factory, + clock_ms=lambda: 777, + ) + uploader.start() + uploader.submit(10, _delta_table([_row(10, 0, 1, [1.0, 2.0])])) + 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]) + 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_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() - uploader = FeatureStoreDeltaUploader( - _feature_store_config(), - output_dir, - "delta", - 1, - {"user_emb": 2}, - client_factory=_FakeClientFactory(view), - ) - uploader.start() - uploader.submit(10) - uploader.close() + view = _FakeView() + uploader = self._uploader(client_factory=_FakeClientFactory(view)) + uploader.start() + uploader.submit(10, _delta_table([_row(10, 0, 1, [1.0, 2.0])])) + uploader.close() - self.assertEqual(len(view.calls), 1) - self.assertEqual(view.calls[0]["write_mode"], "MERGE") + self.assertEqual(len(view.calls), 1) + self.assertEqual(view.calls[0]["write_mode"], "MERGE") def test_non_draining_close_stops_without_commit(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, 10)], - ) - view = _BlockingView() - uploader = FeatureStoreDeltaUploader( - _feature_store_config(upload_batch_size=1), - output_dir, - "delta", - 1, - {"user_emb": 2}, - client_factory=_FakeClientFactory(view), - ) - uploader.start() - uploader.submit(10) - self.assertTrue(view.flush_started.wait(timeout=5)) - uploader.close(raise_on_error=False, drain=False) - view.release_flush.set() - self.assertTrue(view.close_finished.wait(timeout=5)) - self.assertTrue(len(view.calls) < 9) + view = _BlockingView() + uploader = self._uploader( + _feature_store_config(upload_batch_size=1), + client_factory=_FakeClientFactory(view), + ) + uploader.start() + uploader.submit( + 10, _delta_table([_row(10, 0, key, [1.0, 2.0]) for key in range(1, 10)]) + ) + self.assertTrue(view.flush_started.wait(timeout=5)) + uploader.close(raise_on_error=False, drain=False) + view.release_flush.set() + self.assertTrue(view.close_finished.wait(timeout=5)) + self.assertTrue(len(view.calls) < 9) def test_signed_int64_key_is_preserved(self): - with tempfile.TemporaryDirectory() as output_dir: - large_key = (1 << 63) - 1 - _write_single_shard(output_dir, 10, [_row(10, 0, large_key, [1.0, 2.0])]) - view = _FakeView() - uploader = FeatureStoreDeltaUploader( - _feature_store_config(), - output_dir, - "delta", - 1, - {"user_emb": 2}, - client_factory=_FakeClientFactory(view), - ) - uploader.start() - uploader.submit(10) - uploader.close() + large_key = (1 << 63) - 1 + view = _FakeView() + uploader = self._uploader(client_factory=_FakeClientFactory(view)) + uploader.start() + uploader.submit(10, _delta_table([_row(10, 0, large_key, [1.0, 2.0])])) + uploader.close() - self.assertEqual(view.calls[0]["data"][0]["key_id"], large_key) + self.assertEqual(view.calls[0]["data"][0]["key_id"], large_key) def test_reserved_invalid_key_is_rejected(self): - with tempfile.TemporaryDirectory() as output_dir: - _write_single_shard(output_dir, 10, [_row(10, 0, -1, [1.0, 2.0])]) - view = _FakeView() - 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.assertEqual(view.calls, []) - - def test_empty_step_validates_remote_contract(self): - with tempfile.TemporaryDirectory() as output_dir: - _write_single_shard(output_dir, 10, []) - view = _FakeView() - uploader = FeatureStoreDeltaUploader( - _feature_store_config(), - output_dir, - "delta", - 1, - {"user_emb": 2}, - client_factory=_FakeClientFactory(view), - ) - uploader.start() - uploader.submit(10) + view = _FakeView() + uploader = self._uploader( + _feature_store_config(max_retries=1), + client_factory=_FakeClientFactory(view), + ) + uploader.start() + uploader.submit(10, _delta_table([_row(10, 0, -1, [1.0, 2.0])])) + with self.assertRaises(FeatureStoreUploadError): uploader.close() + self.assertEqual(view.calls, []) - self.assertEqual(view.calls, []) - - def test_multi_rank_waits_for_all_shards(self): - with tempfile.TemporaryDirectory() as output_dir: - _write_rank_shard( - output_dir, 10, 0, 2, [_row(10, 0, 1, [1.0, 2.0], world_size=2)] - ) - _write_rank_shard( - output_dir, 10, 1, 2, [_row(10, 1, 2, [3.0, 4.0], world_size=2)] - ) - view = _FakeView() - uploader = FeatureStoreDeltaUploader( - _feature_store_config(), - output_dir, - "delta", - 2, - {"user_emb": 2}, - client_factory=_FakeClientFactory(view), - clock_ms=lambda: 100, - ) - uploader.start() - uploader.submit(10) - uploader.close() + def test_empty_table_upload_writes_nothing(self): + view = _FakeView() + uploader = self._uploader(client_factory=_FakeClientFactory(view)) + uploader.start() + uploader.submit(10, _delta_table([])) + uploader.close() - self.assertEqual(len(view.calls), 2) - self.assertEqual(view.calls[0]["data"][0]["key_id"], 1) - self.assertEqual(view.calls[1]["data"][0]["key_id"], 2) + self.assertEqual(view.calls, []) def test_dimension_and_finite_value_validation(self): - with tempfile.TemporaryDirectory() as output_dir: - _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0, 3.0])]) - view = _FakeView() - 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.assertEqual(view.calls, []) - - with tempfile.TemporaryDirectory() as output_dir: - _write_single_shard(output_dir, 10, [_row(10, 0, 1, [float("nan"), 2.0])]) - view = _FakeView() - 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.assertEqual(view.calls, []) - - def test_successful_upload_deletes_shard_files(self): - with tempfile.TemporaryDirectory() as output_dir: - shard_path = _write_single_shard( - output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])] - ) - self.assertTrue(os.path.isfile(shard_path)) - view = _FakeView() - uploader = FeatureStoreDeltaUploader( - _feature_store_config(), - output_dir, - "delta", - 1, - {"user_emb": 2}, - client_factory=_FakeClientFactory(view), - ) - uploader.start() - uploader.submit(10) + view = _FakeView() + uploader = self._uploader( + _feature_store_config(max_retries=1), + client_factory=_FakeClientFactory(view), + ) + uploader.start() + uploader.submit(10, _delta_table([_row(10, 0, 1, [1.0, 2.0, 3.0])])) + with self.assertRaises(FeatureStoreUploadError): uploader.close() + self.assertEqual(view.calls, []) - self.assertFalse(os.path.isfile(shard_path)) + view = _FakeView() + uploader = self._uploader( + _feature_store_config(max_retries=1), + client_factory=_FakeClientFactory(view), + ) + uploader.start() + uploader.submit(10, _delta_table([_row(10, 0, 1, [float("nan"), 2.0])])) + with self.assertRaises(FeatureStoreUploadError): + uploader.close() + self.assertEqual(view.calls, []) def test_in_memory_timestamp_monotonicity_across_steps(self): - with tempfile.TemporaryDirectory() as output_dir: - _write_single_shard(output_dir, 10, [_row(10, 0, 1, [1.0, 2.0])]) - _write_single_shard(output_dir, 20, [_row(20, 0, 2, [3.0, 4.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.submit(10) - uploader.submit(20) - uploader.close() + view = _FakeView() + uploader = self._uploader( + client_factory=_FakeClientFactory(view), + clock_ms=lambda: 100, + ) + uploader.start() + uploader.submit(10, _delta_table([_row(10, 0, 1, [1.0, 2.0])])) + uploader.submit(20, _delta_table([_row(20, 0, 2, [3.0, 4.0])])) + uploader.close() - ts_values = [call["ts"] for call in view.calls] - self.assertEqual(ts_values, [100, 101]) + ts_values = [call["ts"] for call in view.calls] + self.assertEqual(ts_values, [100, 101]) def test_in_memory_timestamp_monotonicity_across_retries(self): failed_summary = { @@ -1123,95 +837,64 @@ def test_in_memory_timestamp_monotonicity_across_retries(self): "success_records": 0, "failed_records": 1, } - 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: 500, - ) - uploader.start() - uploader.submit(10) - uploader.close() - - self.assertEqual([call["ts"] for call in first_view.calls], [500]) - self.assertEqual([call["ts"] for call in second_view.calls], [501]) - - def test_shard_wait_timeout_raises(self): - with tempfile.TemporaryDirectory() as output_dir: - view = _FakeView() - uploader = FeatureStoreDeltaUploader( - _feature_store_config(shard_wait_timeout_secs=1, poll_interval_secs=1), - output_dir, - "delta", - 1, - {"user_emb": 2}, - client_factory=_FakeClientFactory(view), - ) - uploader.start() - uploader.submit(10) - with self.assertRaises(FeatureStoreUploadError): - uploader.close() - self.assertEqual(view.calls, []) - - def test_data_parallel_duplicate_keys_uploaded_without_dedup(self): - with tempfile.TemporaryDirectory() as output_dir: - _write_rank_shard( - output_dir, - 10, - 0, - 2, - [_row(10, 0, 1, [1.0, 2.0], world_size=2)], - ) - _write_rank_shard( - output_dir, - 10, - 1, - 2, - [_row(10, 1, 1, [1.0, 2.0], world_size=2)], - ) - view = _FakeView() - uploader = FeatureStoreDeltaUploader( - _feature_store_config(), - output_dir, - "delta", - 2, - {"user_emb": 2}, - client_factory=_FakeClientFactory(view), - clock_ms=lambda: 100, - ) - uploader.start() - uploader.submit(10) - uploader.close() - - all_keys = [item["key_id"] for call in view.calls for item in call["data"]] - self.assertEqual(all_keys, [1, 1]) + first_view = _FakeView([failed_summary]) + second_view = _FakeView() + factory = _SequencedClientFactory([first_view, second_view]) + uploader = self._uploader( + _feature_store_config(max_retries=2), + client_factory=factory, + clock_ms=lambda: 500, + ) + uploader.start() + uploader.submit(10, _delta_table([_row(10, 0, 1, [1.0, 2.0])])) + uploader.close() + + self.assertEqual([call["ts"] for call in first_view.calls], [500]) + self.assertEqual([call["ts"] for call in second_view.calls], [501]) + + def test_data_parallel_ranks_upload_duplicate_keys_independently(self): + rank0_view = _FakeView() + rank1_view = _FakeView() + rank0 = self._uploader( + client_factory=_FakeClientFactory(rank0_view), + clock_ms=lambda: 100, + ) + rank1 = self._uploader( + rank=1, + manage_remote_view=False, + client_factory=_FakeClientFactory(rank1_view), + clock_ms=lambda: 100, + ) + rank0.start() + rank1.start() + rank0.submit(10, _delta_table([_row(10, 0, 1, [1.0, 2.0], world_size=2)])) + rank1.submit(10, _delta_table([_row(10, 1, 1, [1.0, 2.0], world_size=2)])) + rank0.close() + rank1.close() + + rank0_keys = [ + item["key_id"] for call in rank0_view.calls for item in call["data"] + ] + rank1_keys = [ + item["key_id"] for call in rank1_view.calls for item in call["data"] + ] + self.assertEqual(rank0_keys, [1]) + self.assertEqual(rank1_keys, [1]) + self.assertEqual({call["write_mode"] for call in rank0_view.calls}, {"MERGE"}) + self.assertEqual({call["write_mode"] for call in rank1_view.calls}, {"MERGE"}) def test_close_error_surfaces_via_check_error(self): - with tempfile.TemporaryDirectory() as output_dir: - _write_single_shard(output_dir, 10, [_row(10, 0, -1, [1.0, 2.0])]) - view = _FakeView() - 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() - with self.assertRaises(FeatureStoreUploadError): - uploader.check_error() + view = _FakeView() + uploader = self._uploader( + _feature_store_config(max_retries=1), + client_factory=_FakeClientFactory(view), + ) + uploader.start() + uploader.submit(10, _delta_table([_row(10, 0, -1, [1.0, 2.0])])) + with self.assertRaises(FeatureStoreUploadError): + uploader.close() + with self.assertRaises(FeatureStoreUploadError): + uploader.check_error() if __name__ == "__main__": From 877ece7af918bd5892551107c2430331ff882223 Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 24 Jul 2026 15:08:59 +0800 Subject: [PATCH 30/50] [bugfix] fan tracked delta ids out to every owning sparse module The TorchRec delta tracker keys its bookkeeping by raw table name, so a table name legitimately reused by several modules (an EC sequence path and an EBC deep path, or two data-group impls) collapsed into one entry and the dump either raised or published one owner's weights under the other's identity. Identities, weights and dynamicemb lookups are now keyed by (module_fqn, table_name), and the merged id set is fanned out to every owner so each publishes its own true rows under its own role-suffixed serving name. Same-role duplicates still raise for FeatureStore upload because the export naming contract cannot address them distinctly. --- tzrec/utils/delta_embedding_dump.py | 233 +++++++++--------- tzrec/utils/delta_embedding_dump_test.py | 97 +++++++- tzrec/utils/sparse_embedding_contract.py | 9 - tzrec/utils/sparse_embedding_contract_test.py | 16 -- 4 files changed, 211 insertions(+), 144 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index d6be8839a..2fabfb91e 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -42,7 +42,6 @@ SparseEmbeddingIdentity, build_sparse_embedding_name_map, resolve_sparse_embedding_name, - sparse_embedding_role_from_state_key, ) _CONSUMER = "delta_embedding_dump" @@ -179,6 +178,12 @@ def _feature_name(feature_names: Iterable[str]) -> str: return ",".join(names) +def _owner_table_fqn(module_fqn: str, identity: SparseEmbeddingIdentity) -> str: + """Build the state-dict style per-table FQN of one owning module.""" + segment = "embedding_bags" if identity.role == SPARSE_EBC_ROLE else "embeddings" + return f"{module_fqn}.{segment}.{identity.table_name}" + + def _int_attr(value: Any, name: str) -> int: attr = getattr(value, name, 0) return int(attr) if attr is not None else 0 @@ -319,9 +324,12 @@ def _local_table_weight( raise ValueError( "delta embedding dump only supports one local shard per table." ) + # The shard's own metadata is authoritative for offsets: the passed + # shard_info is merged by table name and may describe another module's + # same-named table. info = _merge_shard_info( - shard_info or _TableShardInfo(), _metadata_shard_info(getattr(shards[0], "metadata", None)), + shard_info or _TableShardInfo(), ) info = _table_shard_info_from_tensor(shards[0].tensor, info) return _TableWeight(tensor=shards[0].tensor, shard_info=info) @@ -334,8 +342,8 @@ def _local_table_weight( "delta embedding dump only supports one local shard per table." ) info = _merge_shard_info( - shard_info or _TableShardInfo(), _metadata_shard_info(getattr(shards[0], "metadata", None)), + shard_info or _TableShardInfo(), ) info = _table_shard_info_from_tensor(shards[0].tensor, info) return _TableWeight(tensor=shards[0].tensor, shard_info=info) @@ -410,11 +418,22 @@ def __init__( self._fqn_to_table: Dict[str, str] = { fqn: table for table, fqn in self._table_to_fqn.items() } - 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._identity_by_owner, embedding_dimensions = ( self._build_sparse_embedding_contract() ) + self._owners_by_table: Dict[str, List[str]] = {} + for module_fqn, table_name in self._identity_by_owner: + self._owners_by_table.setdefault(table_name, []).append(module_fqn) + unowned_tables = sorted( + table_name + for table_name in self._table_to_fqn + if table_name not in self._owners_by_table + ) + if unowned_tables: + raise ValueError( + "cannot resolve owning sparse embedding collections for tracked " + f"tables: {unowned_tables}" + ) self._uploader: Optional[FeatureStoreDeltaUploader] = None if self._feature_store_enabled: self._uploader = FeatureStoreDeltaUploader( @@ -506,11 +525,21 @@ def _check_feature_store_upload_error(self, force: bool = False) -> None: 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]] = {} + ) -> Tuple[Dict[Tuple[str, str], SparseEmbeddingIdentity], Dict[str, int]]: + """Build per-owner physical-table identities aligned with sparse export. + + The TorchRec tracker collapses same-named tables from different + modules into one entry, so identities are keyed by + ``(module_fqn, table_name)`` and ``dump`` fans the tracked ids out to + every owner, looking values up in each owner's own weights. Canonical + embedding names follow the sparse-export contract: a table name reused + across EC and EBC gets a role suffix, so both physical tables publish + under distinct serving names. + """ + identity_by_owner: Dict[Tuple[str, str], SparseEmbeddingIdentity] = {} + owner_roles: Dict[Tuple[str, str], str] = {} + owner_metadata: Dict[Tuple[str, str], Tuple[int, Tuple[str, ...]]] = {} + fqns_by_role_table: Dict[Tuple[str, str], Set[str]] = {} for module_fqn, module in self._tracker.get_tracked_modules().items(): if isinstance(module, ShardedEmbeddingCollection): role = SPARSE_EC_ROLE @@ -525,81 +554,49 @@ def _build_sparse_embedding_contract( 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: + if dimension <= 0: raise ValueError( - "inconsistent sparse embedding metadata for " - f"role={role} table={table_name}: {previous} vs {current}" + f"invalid embedding dimension for table {table_name!r}: " + f"{dimension}" ) - metadata_by_identity[identity_key] = current - roles_by_table.setdefault(table_name, set()).add(role) + feature_names = tuple(getattr(table_config, "feature_names", ())) + owner_key = (module_fqn, table_name) + owner_roles[owner_key] = role + owner_metadata[owner_key] = (dimension, feature_names) + fqns_by_role_table.setdefault((role, table_name), set()).add(module_fqn) - ambiguous_tables = sorted( - table_name for table_name, roles in roles_by_table.items() if len(roles) > 1 - ) - if ambiguous_tables: - # The TorchRec delta tracker keys bookkeeping by raw table name, so - # EC/EBC sharing a name collapse to one physical table and a wrong - # primary key could be published. Refuse for the shared FeatureStore - # (corrupting); the local parquet dump is re-derivable, so it warns - # and publishes the surviving table until the tracker is role-aware. - if self._feature_store_enabled: + if self._feature_store_enabled: + # Cross-role reuse gets distinct role-suffixed serving names, but + # two same-role physical tables sharing one table name (e.g. one + # embedding_name reused across data groups) would collide on one + # FeatureStore primary key with diverged weights. + duplicated = sorted( + f"{role}:{table_name}" + for (role, table_name), fqns in fqns_by_role_table.items() + if len(fqns) > 1 + ) + if duplicated: raise ValueError( - "delta embedding dump cannot safely upload to FeatureStore " - "while table names are reused by both EmbeddingCollection " - f"and EmbeddingBagCollection: {ambiguous_tables}. The " - "TorchRec tracker would publish a wrong primary key; " - "role-aware tracker identity is required." + "FeatureStore delta upload cannot address multiple physical " + "tables that share one (role, table_name) serving identity: " + f"{duplicated}. Give the tables distinct embedding_names." ) - logger.warning( - "Delta embedding dump cannot distinguish tables reused by both " - "EmbeddingCollection and EmbeddingBagCollection (%s); the TorchRec " - "tracker collapses them to one physical table, so only that " - "table's deltas are published.", - ambiguous_tables, - ) - name_by_identity = build_sparse_embedding_name_map(metadata_by_identity) - identity_by_fqn: Dict[str, SparseEmbeddingIdentity] = {} + name_by_identity = build_sparse_embedding_name_map(fqns_by_role_table.keys()) 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}" - ) + for owner_key, role in owner_roles.items(): + module_fqn, table_name = owner_key + dimension, feature_names = owner_metadata[owner_key] embedding_name = resolve_sparse_embedding_name( name_by_identity, table_name, role ) - identity = SparseEmbeddingIdentity( + identity_by_owner[owner_key] = 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( @@ -607,7 +604,7 @@ def _build_sparse_embedding_contract( f"dimensions: {previous_dimension} vs {dimension}" ) embedding_dimensions[embedding_name] = dimension - return identity_by_fqn, embedding_dimensions + return identity_by_owner, embedding_dimensions def _tracker_cursor_before_read(self) -> Optional[int]: if not self._feature_store_enabled: @@ -875,8 +872,8 @@ def _append_model_delta_rows( self, table_chunks: List[pa.Table], global_step: int, - table_weights: Dict[str, _TableWeight], - dynamic_modules: Dict[str, nn.Module], + table_weights: Dict[Tuple[str, str], _TableWeight], + dynamic_modules: Dict[Tuple[str, str], nn.Module], ) -> int: num_rows = 0 # A dynamic module hosting multiple tables is shared across their @@ -891,49 +888,60 @@ 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: + owners = self._owners_by_table.get(table_name) + if not owners: raise ValueError( - f"Missing sparse embedding contract for table fqn {fqn!r}" + f"Missing sparse embedding contract for table {table_name!r}" ) ids = ids.unique(sorted=True) - embeddings, key_ids = self._lookup_embeddings( - table_name, - ids, - table_weights=table_weights, - dynamic_modules=dynamic_modules, - flushed_module_ids=flushed_module_ids, - ) - feature_name = _feature_name(self._fqn_to_feature_names.get(fqn, [])) - 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, - embeddings=embeddings, - ) + # The tracker merges same-named tables into one id set; fan it out + # to every owning module and read each owner's own weights, so an + # id that only one owner touched still publishes that owner's true + # current row for the others. + for module_fqn in owners: + identity = self._identity_by_owner[(module_fqn, table_name)] + embeddings, key_ids = self._lookup_embeddings( + module_fqn, + table_name, + ids, + table_weights=table_weights, + dynamic_modules=dynamic_modules, + flushed_module_ids=flushed_module_ids, + ) + 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(identity.feature_names), + table_fqn=_owner_table_fqn(module_fqn, identity), + key_ids=key_ids, + embeddings=embeddings, + ) return num_rows def _lookup_embeddings( self, + module_fqn: str, table_name: str, ids: torch.Tensor, - table_weights: Dict[str, _TableWeight], - dynamic_modules: Dict[str, nn.Module], + table_weights: Dict[Tuple[str, str], _TableWeight], + dynamic_modules: Dict[Tuple[str, str], nn.Module], flushed_module_ids: Optional[Set[int]] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - dynamic_module = dynamic_modules.get(table_name) + owner_key = (module_fqn, table_name) + dynamic_module = dynamic_modules.get(owner_key) if dynamic_module is not None: return self._lookup_dynamic_embeddings( dynamic_module, table_name, ids, flushed_module_ids ) - if table_name not in table_weights: - raise KeyError(f"Embedding table {table_name} not found in sharded model.") - table_weight = table_weights[table_name] + if owner_key not in table_weights: + raise KeyError( + f"Embedding table {table_name} not found in sharded module " + f"{module_fqn}." + ) + table_weight = table_weights[owner_key] _validate_table_shard_info(table_name, table_weight.shard_info) self._validate_row_shard_metadata(table_name, table_weight.shard_info) weight = table_weight.tensor @@ -1063,10 +1071,10 @@ def _validate_row_shard_metadata( "TorchRec shard metadata is missing." ) - def _collect_table_weights(self) -> Dict[str, _TableWeight]: - table_weights: Dict[str, _TableWeight] = {} + def _collect_table_weights(self) -> Dict[Tuple[str, str], _TableWeight]: + table_weights: Dict[Tuple[str, str], _TableWeight] = {} table_shard_infos = self._table_shard_infos - for module in self._model.modules(): + for module_fqn, module in self._tracker.get_tracked_modules().items(): lookups = getattr(module, "_lookups", None) if lookups is None: continue @@ -1078,25 +1086,22 @@ def _collect_table_weights(self) -> Dict[str, _TableWeight]: if named_parameters_by_table is None: continue for table_name, table_value in named_parameters_by_table(): - table_weights[table_name] = _local_table_weight( + table_weights[(module_fqn, table_name)] = _local_table_weight( table_value, table_shard_infos.get(table_name), ) return table_weights - def _collect_dynamic_modules(self) -> Dict[str, nn.Module]: + def _collect_dynamic_modules(self) -> Dict[Tuple[str, str], nn.Module]: try: from dynamicemb.dump_load import get_dynamic_emb_module except ImportError: return {} - modules: Dict[str, nn.Module] = {} - seen = set() - for dynamic_module in get_dynamic_emb_module(self._model): - if id(dynamic_module) in seen: - continue - seen.add(id(dynamic_module)) - for table_name in dynamic_module.table_names: - modules[table_name] = dynamic_module + modules: Dict[Tuple[str, str], nn.Module] = {} + for module_fqn, module in self._tracker.get_tracked_modules().items(): + for dynamic_module in get_dynamic_emb_module(module): + for table_name in dynamic_module.table_names: + modules[(module_fqn, table_name)] = dynamic_module return modules def _append_table_chunk( diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 17af163d0..9b3a1ea45 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -56,6 +56,7 @@ ) from tzrec.utils.dist_util import create_train_pipeline from tzrec.utils.dynamicemb_util import has_dynamicemb +from tzrec.utils.sparse_embedding_contract import SparseEmbeddingIdentity from tzrec.utils.test_util import gpu_unavailable, make_test_dir, mark_ci_scope _SHARDED_TABLE_NAME = "table_1" @@ -164,7 +165,12 @@ def _assert_sharded_dump_file(rank: int, output_path: str, dumper) -> None: ) testcase.assertEqual(set(table["embedding_role"].to_pylist()), {"ebc"}) - table_weight = dumper._collect_table_weights()[_SHARDED_TABLE_NAME] + table_weights = dumper._collect_table_weights() + table_weight = next( + weight + for (_fqn, table_name), weight in table_weights.items() + if table_name == _SHARDED_TABLE_NAME + ) expected_key_ids = [ key_id for key_id in _SHARDED_INPUT_IDS @@ -414,6 +420,83 @@ def test_dump_rows_include_rank_metadata(self): ], ) + def test_shared_table_name_fans_out_to_all_owners(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._rank = 0 + dumper._world_size = 1 + dumper._fqn_to_table = {"sparse.shared": "shared"} + dumper._owners_by_table = {"shared": ["model.ec_dict.2", "model.ebc"]} + dumper._identity_by_owner = { + ("model.ec_dict.2", "shared"): SparseEmbeddingIdentity( + role="ec", + table_name="shared", + embedding_name="shared__ec", + dimension=2, + feature_names=("query_feat",), + ), + ("model.ebc", "shared"): SparseEmbeddingIdentity( + role="ebc", + table_name="shared", + embedding_name="shared__ebc", + dimension=2, + feature_names=("deep_feat",), + ), + } + dumper._tracker = mock.MagicMock() + dumper._tracker.get_unique.return_value = { + "sparse.shared": SimpleNamespace(ids=torch.tensor([0, 2])) + } + ec_weight = torch.tensor([[0.0, 0.1], [1.0, 1.1], [2.0, 2.1], [3.0, 3.1]]) + ebc_weight = ec_weight + 10.0 + shard_info = _TableShardInfo( + local_rows=4, local_cols=2, global_rows=4, global_cols=2 + ) + table_weights = { + ("model.ec_dict.2", "shared"): _TableWeight( + tensor=ec_weight, shard_info=shard_info + ), + ("model.ebc", "shared"): _TableWeight( + tensor=ebc_weight, shard_info=shard_info + ), + } + + table_chunks = [] + num_rows = dumper._append_model_delta_rows( + table_chunks, + global_step=10, + table_weights=table_weights, + dynamic_modules={}, + ) + + self.assertEqual(num_rows, 4) + table = pa.concat_tables(table_chunks) + rows_by_name = {} + for row in table.to_pylist(): + rows_by_name.setdefault(row["embedding_name"], []).append(row) + self.assertEqual(set(rows_by_name), {"shared__ec", "shared__ebc"}) + ec_rows = rows_by_name["shared__ec"] + self.assertEqual([row["key_id"] for row in ec_rows], [0, 2]) + self.assertEqual( + [row["embedding"] for row in ec_rows], + ec_weight[[0, 2]].tolist(), + ) + self.assertEqual( + {row["table_fqn"] for row in ec_rows}, + {"model.ec_dict.2.embeddings.shared"}, + ) + self.assertEqual({row["feature_name"] for row in ec_rows}, {"query_feat"}) + ebc_rows = rows_by_name["shared__ebc"] + self.assertEqual([row["key_id"] for row in ebc_rows], [0, 2]) + self.assertEqual( + [row["embedding"] for row in ebc_rows], + ebc_weight[[0, 2]].tolist(), + ) + self.assertEqual( + {row["table_fqn"] for row in ebc_rows}, + {"model.ebc.embedding_bags.shared"}, + ) + self.assertEqual({row["feature_name"] for row in ebc_rows}, {"deep_feat"}) + def test_dump_rows_reject_processor_invalid_key_sentinel(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._rank = 0 @@ -1061,10 +1144,11 @@ def test_row_wise_lookup_outputs_global_key_ids(self): dumper._world_size = 2 weight = torch.tensor([[0.0, 0.1], [1.0, 1.1], [2.0, 2.1], [3.0, 3.1]]) embeddings, key_ids = dumper._lookup_embeddings( + "model.ebc", "user_emb", torch.tensor([0, 2]), table_weights={ - "user_emb": _TableWeight( + ("model.ebc", "user_emb"): _TableWeight( tensor=weight, shard_info=_TableShardInfo( row_offset=32, @@ -1086,10 +1170,11 @@ def test_lookup_filters_out_of_range_ids(self): dumper._world_size = 2 weight = torch.tensor([[0.0, 0.1], [1.0, 1.1], [2.0, 2.1], [3.0, 3.1]]) embeddings, key_ids = dumper._lookup_embeddings( + "model.ebc", "user_emb", torch.tensor([0, 2, 99, -1]), table_weights={ - "user_emb": _TableWeight( + ("model.ebc", "user_emb"): _TableWeight( tensor=weight, shard_info=_TableShardInfo( row_offset=32, @@ -1111,10 +1196,11 @@ def test_lookup_handles_empty_ids(self): dumper._world_size = 2 weight = torch.tensor([[0.0, 0.1], [1.0, 1.1], [2.0, 2.1], [3.0, 3.1]]) embeddings, key_ids = dumper._lookup_embeddings( + "model.ebc", "user_emb", torch.tensor([], dtype=torch.long), table_weights={ - "user_emb": _TableWeight( + ("model.ebc", "user_emb"): _TableWeight( tensor=weight, shard_info=_TableShardInfo( row_offset=32, @@ -1136,10 +1222,11 @@ def test_row_wise_lookup_requires_shard_metadata(self): dumper._world_size = 2 with self.assertRaisesRegex(ValueError, "shard metadata"): dumper._lookup_embeddings( + "model.ebc", "user_emb", torch.tensor([0]), table_weights={ - "user_emb": _TableWeight( + ("model.ebc", "user_emb"): _TableWeight( tensor=torch.zeros(4, 2), shard_info=_TableShardInfo( local_rows=4, diff --git a/tzrec/utils/sparse_embedding_contract.py b/tzrec/utils/sparse_embedding_contract.py index f64f93733..3aae3b58d 100644 --- a/tzrec/utils/sparse_embedding_contract.py +++ b/tzrec/utils/sparse_embedding_contract.py @@ -94,12 +94,3 @@ def resolve_sparse_embedding_name( 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 index 2a9f54648..abe1f258a 100644 --- a/tzrec/utils/sparse_embedding_contract_test.py +++ b/tzrec/utils/sparse_embedding_contract_test.py @@ -14,7 +14,6 @@ from tzrec.utils.sparse_embedding_contract import ( build_sparse_embedding_name_map, resolve_sparse_embedding_name, - sparse_embedding_role_from_state_key, ) @@ -46,21 +45,6 @@ def test_resolver_requires_role_for_ambiguous_table(self): 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 e7d28547ebc94b5a64adb0e5d3c6a60d6f7779c0 Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 24 Jul 2026 15:22:32 +0800 Subject: [PATCH 31/50] [perf] decide timed delta dumps per rank without a per-step collective Minute-cadence dumps all-reduced a rank-zero clock vote on every training step so all ranks dumped together. With per-rank FeatureStore uploads and rank-local tracker windows, ranks no longer need to dump on the same step, so each rank now checks its own clock against a shared fixed-rate deadline sequence armed at start(), removing the hot-path collective entirely. Multi-rank timed local parquet sets are best-effort step-aligned. --- tzrec/utils/delta_embedding_dump.py | 106 ++++-------------- tzrec/utils/delta_embedding_dump_test.py | 136 ++++++----------------- 2 files changed, 58 insertions(+), 184 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 2fabfb91e..f7d8d349f 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -388,14 +388,11 @@ def __init__( self._interval_steps = int(config.dump_interval_steps) self._next_dump_time: Optional[float] = None self._last_dump_step: Optional[int] = None - self._pending_rendezvous: Optional[torch.Tensor] = None - self._timed_dump_in_flight = False self._output_dir = config.output_dir or os.path.join( model_dir, "delta_embedding_dump" ) self._file_prefix = config.file_prefix or "delta_embedding" self._rank, self._world_size = _distributed_rank_world_size() - self._device = device self._tracking_pause_depth = 0 self._feature_store_enabled = config.HasField("feature_store_config") self._retain_local_dump = self._feature_store_enabled and bool( @@ -503,9 +500,7 @@ def close(self, raise_on_error: bool = True, drain: bool = True) -> None: 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]: + def _feature_store_upload_error(self) -> Optional[BaseException]: """Collect this rank's uploader error without changing control flow.""" if not self._feature_store_enabled: return None @@ -517,9 +512,9 @@ def _feature_store_upload_error( return exc return None - def _check_feature_store_upload_error(self, force: bool = False) -> None: + def _check_feature_store_upload_error(self) -> None: """Surface this rank's background upload failure to the trainer.""" - error = self._feature_store_upload_error(force=force) + error = self._feature_store_upload_error() if error is not None: raise error.with_traceback(error.__traceback__) @@ -630,34 +625,30 @@ def maybe_dump(self, global_step: int) -> None: Args: global_step: Current training step. """ - self._check_feature_store_upload_error(force=True) - if self._requires_dump_state_rendezvous(): - should_dump = self._consume_dump_state_rendezvous() - self._launch_dump_state_rendezvous(global_step) - else: - should_dump = self._local_dump_decision(global_step) - if should_dump: + self._check_feature_store_upload_error() + if self._local_dump_decision(global_step): self.dump(global_step) self._last_dump_step = global_step - if self._interval_secs is not None and self._rank == 0: - self._next_dump_time = time.monotonic() + self._interval_secs - self._timed_dump_in_flight = False + if self._interval_secs is not None and self._next_dump_time is not None: + # Fixed-rate rescheduling keeps every rank's deadline sequence + # identical, so timed dumps stay step-aligned across ranks up + # to clock skew exactly astride a step boundary; missed + # deadlines are skipped instead of fired as a burst. + now = time.monotonic() + while self._next_dump_time <= now: + self._next_dump_time += self._interval_secs self._tracker.step() - def _requires_dump_state_rendezvous(self) -> bool: - """Return whether maybe_dump must all-reduce its dump vote. - - Timed dumps OR-reduce the rank-zero clock decision so every rank dumps - together. Step-based dumps are deterministic across ranks and need no - collective. - """ - return self._world_size > 1 and self._interval_secs is not None - def _local_dump_decision(self, global_step: int) -> bool: - """Return the dump decision when no cross-rank rendezvous is needed.""" - upload_error = self._feature_store_upload_error() - if upload_error is not None: - raise upload_error.with_traceback(upload_error.__traceback__) + """Return this rank's local dump decision for the current step. + + Timed dumps are decided per rank against its own clock: FeatureStore + uploads are per rank and tracker windows are rank-local, so ranks do + not need to dump on the same step and no cross-rank rendezvous is + required. All ranks arm the same deadline sequence in ``start`` and + advance it at a fixed rate, so their decisions can differ by at most + the steps their clocks disagree on. + """ if global_step <= 0: return False if self._interval_steps is not None: @@ -666,58 +657,7 @@ def _local_dump_decision(self, global_step: int) -> bool: raise RuntimeError( "time-based delta embedding dumper must be started before training" ) - return self._rank == 0 and time.monotonic() >= self._next_dump_time - - def _launch_dump_state_rendezvous(self, global_step: int) -> None: - """Asynchronously all-reduce this step's timed dump vote. - - The reduced tensor is consumed by the next ``maybe_dump`` call. NCCL - orders this collective ahead of the next training step's communication - on the same process group, so it has already completed by the time - the next call reads it back. - """ - if not ( - torch.distributed.is_available() and torch.distributed.is_initialized() - ): - raise RuntimeError( - "distributed delta embedding dump requires an initialized process group" - ) - state = torch.tensor( - [int(self._timed_dump_vote(global_step))], - dtype=torch.int32, - device=self._device, - ) - torch.distributed.all_reduce(state, op=torch.distributed.ReduceOp.MAX) - self._pending_rendezvous = state - - def _consume_dump_state_rendezvous(self) -> bool: - """Read back the previous step's reduced timed dump vote.""" - state = self._pending_rendezvous - self._pending_rendezvous = None - if state is None: - return False - return bool(state.tolist()[0]) - - def _timed_dump_vote(self, global_step: int) -> bool: - """Return rank zero's local vote for whether a timed dump is due. - - A fired vote arms ``_timed_dump_in_flight`` so the following step's - vote stays low until the dump executes and reschedules the timer from - its completion time; without the guard the still-elapsed deadline - would fire on the next launch and dump twice. - """ - if self._interval_secs is None or global_step <= 0: - return False - if self._next_dump_time is None: - raise RuntimeError( - "time-based delta embedding dumper must be started before training" - ) - if self._timed_dump_in_flight: - return False - if self._rank == 0 and time.monotonic() >= self._next_dump_time: - self._timed_dump_in_flight = True - return True - return False + return time.monotonic() >= self._next_dump_time 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 9b3a1ea45..d779d2700 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -674,7 +674,7 @@ 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): + def test_maybe_dump_uses_elapsed_time_with_fixed_rate_schedule(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._interval_steps = None dumper._interval_secs = 60.0 @@ -682,10 +682,9 @@ def test_maybe_dump_uses_elapsed_time_and_resets_from_completion(self): dumper._last_dump_step = None dumper._rank = 0 dumper._world_size = 1 - dumper._device = torch.device("cpu") + dumper._feature_store_enabled = False dumper._tracker = mock.MagicMock() with ( - 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", @@ -699,161 +698,96 @@ def test_maybe_dump_uses_elapsed_time_and_resets_from_completion(self): self.assertEqual( [call.args[0] for call in dump_mock.call_args_list], - [11, 13], + [11, 12], ) - self.assertEqual(dumper._next_dump_time, 283.0) - self.assertEqual(dumper._last_dump_step, 13) + # Deadlines advance at a fixed rate from the armed schedule + # (160 -> 220 -> 280), not from each dump's completion time. + self.assertEqual(dumper._next_dump_time, 280.0) + self.assertEqual(dumper._last_dump_step, 12) self.assertEqual(dumper._tracker.step.call_count, 4) - def _new_rendezvous_dumper(self) -> DeltaEmbeddingDumper: + def test_timed_dump_decides_locally_without_collectives(self): dumper = object.__new__(DeltaEmbeddingDumper) - dumper._last_dump_step = None - dumper._device = torch.device("cpu") - dumper._tracker = mock.MagicMock() - dumper._pending_rendezvous = None - dumper._feature_store_enabled = False - dumper._timed_dump_in_flight = False - return dumper - - def test_time_dump_decision_is_reduced_from_rank_zero(self): - # The rank-zero timer vote is launched on one step and consumed on the - # next: the pipelined rendezvous lands a timed dump at most one step - # after the timer fires without blocking any step on the collective. - dumper = self._new_rendezvous_dumper() dumper._interval_steps = None dumper._interval_secs = 60.0 dumper._next_dump_time = 160.0 + dumper._last_dump_step = None dumper._rank = 1 dumper._world_size = 2 - launched_states = [] - - def fake_all_reduce(state, op=None): - self.assertIs(op, torch.distributed.ReduceOp.MAX) - launched_states.append(state.tolist()) - state[0] = 1 - - with ( - 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), - ): - dumper.maybe_dump(10) - dump_mock.assert_not_called() - dumper.maybe_dump(11) - - dump_mock.assert_called_once_with(11) - self.assertEqual(dumper._last_dump_step, 11) - self.assertEqual(launched_states, [[0], [0]]) - self.assertEqual(dumper._tracker.step.call_count, 2) - - def test_timed_maybe_dump_advances_after_all_workers_succeed(self): - dumper = self._new_rendezvous_dumper() - dumper._interval_steps = None - dumper._interval_secs = 60.0 - dumper._next_dump_time = 0.0 - dumper._rank = 0 - dumper._world_size = 2 - collective_states = [] - - def fake_all_reduce(state, op=None): - self.assertIs(op, torch.distributed.ReduceOp.MAX) - collective_states.append(state.tolist()) - + dumper._feature_store_enabled = False + dumper._tracker = mock.MagicMock() 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.all_reduce") as all_reduce_mock, mock.patch( "tzrec.utils.delta_embedding_dump.time.monotonic", - side_effect=[1.0, 2.0, 3.0], + side_effect=[159.0, 160.5, 161.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) dump_mock.assert_not_called() dumper.maybe_dump(11) - self.assertEqual(collective_states, [[1], [0]]) + all_reduce_mock.assert_not_called() dump_mock.assert_called_once_with(11) self.assertEqual(dumper._last_dump_step, 11) - self.assertEqual(dumper._next_dump_time, 62.0) - self.assertFalse(dumper._timed_dump_in_flight) + self.assertEqual(dumper._next_dump_time, 220.0) self.assertEqual(dumper._tracker.step.call_count, 2) def test_timed_maybe_dump_propagates_local_dump_failure(self): - dumper = self._new_rendezvous_dumper() + 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._feature_store_enabled = False + dumper._tracker = mock.MagicMock() dump_error = RuntimeError("local dump failed") - 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", 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), + mock.patch( + "tzrec.utils.delta_embedding_dump.time.monotonic", return_value=1.0 + ), ): - dumper.maybe_dump(10) with self.assertRaises(RuntimeError) as context: - dumper.maybe_dump(11) + dumper.maybe_dump(10) - self.assertEqual(collective_states, [[1], [0]]) self.assertIs(context.exception, dump_error) - dump_mock.assert_called_once_with(11) + dump_mock.assert_called_once_with(10) self.assertIsNone(dumper._last_dump_step) - self.assertEqual(dumper._tracker.step.call_count, 1) + self.assertEqual(dumper._tracker.step.call_count, 0) - def test_timed_dump_vote_does_not_fire_twice_before_completion(self): - # Once rank zero's timer vote fires, the next step's vote stays low - # until the consumed dump reschedules from its completion time; the - # still-elapsed deadline must not launch a second dump. - dumper = self._new_rendezvous_dumper() + def test_timed_dump_skips_missed_deadlines_without_burst(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 - collective_states = [] - - def fake_all_reduce(state, op=None): - self.assertIs(op, torch.distributed.ReduceOp.MAX) - collective_states.append(state.tolist()) - + dumper._feature_store_enabled = False + dumper._tracker = mock.MagicMock() 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=[100.0, 100.0, 100.0], + return_value=100.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) dumper.maybe_dump(11) dumper.maybe_dump(12) - dump_mock.assert_called_once_with(11) - self.assertEqual(dumper._next_dump_time, 160.0) - self.assertFalse(dumper._timed_dump_in_flight) - self.assertEqual(collective_states, [[1], [0], [0]]) + dump_mock.assert_called_once_with(10) + # Deadlines 0 and 60 already elapsed at the dump; skip past them + # instead of firing a burst of catch-up dumps. + self.assertEqual(dumper._next_dump_time, 120.0) self.assertEqual(dumper._tracker.step.call_count, 3) def test_final_dump_skips_step_already_dumped_by_time_interval(self): From 031e7af5b543f2511f2bdbf51cb20412c1c04aa9 Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 24 Jul 2026 15:43:47 +0800 Subject: [PATCH 32/50] [chore] drop redundant value checks from the dump hot path --- tzrec/utils/delta_embedding_dump.py | 7 ------- tzrec/utils/delta_embedding_dump_test.py | 17 ----------------- 2 files changed, 24 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index f7d8d349f..0c570d46e 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -1082,13 +1082,6 @@ def _append_table_chunk( 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( [ diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index d779d2700..276569767 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -497,23 +497,6 @@ def test_shared_table_name_fans_out_to_all_owners(self): ) self.assertEqual({row["feature_name"] for row in ebc_rows}, {"deep_feat"}) - 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]]), - ) - def test_write_table_chunks_preserves_parquet_schema(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._rank = 0 From d0b5e629ef05414fc10ed79687ecb1f5ef05b25a Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 24 Jul 2026 21:24:47 +0800 Subject: [PATCH 33/50] [bugfix] drop stale security_token assertions from distributed export test fb53143 removed the credential fields from FeatureStoreConfig (3062112 renumbered the proto) and reverted the export-time sanitization to plain save_message, but this test -- added earlier by 70ee3b8 when sanitization existed -- still set feature_store_config.security_token and asserted it was scrubbed from the exported pipeline.config, so it fails on the CPU CI lane with AttributeError: Protocol message FeatureStoreConfig has no "security_token" field. Drop the dead credential assertions, keep the override and FeatureStoreConfig round-trip coverage the test still provides, and rename it to describe what it now verifies. Co-Authored-By: Claude --- tzrec/utils/export_util_test.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/tzrec/utils/export_util_test.py b/tzrec/utils/export_util_test.py index c253efb83..4e6c92012 100644 --- a/tzrec/utils/export_util_test.py +++ b/tzrec/utils/export_util_test.py @@ -175,7 +175,7 @@ def test_distributed_embedding_export_skips_nonzero_rank_before_pg_init( else: os.environ[key] = value - def test_distributed_embedding_export_uses_overrides_and_sanitizes_config( + def test_distributed_embedding_export_uses_overrides_and_preserves_config( self, ) -> None: class FakeBatch: @@ -232,7 +232,6 @@ def forward(self, data, device=None): # type: ignore[no-untyped-def] 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() @@ -300,16 +299,12 @@ 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 95a0baaaa40fb6531f9397263e6d1e7e7285291c Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 24 Jul 2026 21:25:44 +0800 Subject: [PATCH 34/50] [bugfix] restore synced dataloader exhaustion for multi-rank delta dumps eeaace1 removed requires_synced_dataloader_exhaustion and the main.py fail_on_uneven_data wiring as a side item of the error-bit rendezvous removal, but final_dump still MAX-all-reduces the final step across ranks and skips steps that fall on a dump boundary. Without the enforcement, a rank whose dataloader stops before the synced step (e.g. 49 batches vs 50) never ran maybe_dump on that boundary, adopts the boundary skip, and silently drops every trailing tracked row -- the exact data loss the regression test left behind by that commit guards against, which now fails with "RuntimeError not raised" on the GPU CI lane. Restore the property (True whenever world_size > 1, either cadence) and wire check_all_workers_data_status/fail_on_uneven_data from it in _train_and_evaluate, so uneven exhaustion fails loudly on every rank instead of dropping a lagging rank's delta; also restore the property's contract tests and the pipeline kwargs in the regression worker, and reword the final_dump call-site comment the removal left false. The error-bit removal itself stays. Co-Authored-By: Claude --- tzrec/main.py | 22 +++++++++++----- tzrec/utils/delta_embedding_dump.py | 12 +++++++++ tzrec/utils/delta_embedding_dump_test.py | 33 ++++++++++++++++++++++-- 3 files changed, 59 insertions(+), 8 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index 7f5e48aa9..c35b01232 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -469,11 +469,22 @@ def run_eval(step: int, epoch: int) -> None: # this rank's last consumed event-time, reused by the epoch / final saves data_timestamp = -1.0 + require_equal_train_batches = ( + delta_embedding_dumper is not None + and delta_embedding_dumper.requires_synced_dataloader_exhaustion + ) + sync_train_data_exhaustion = ( + check_all_workers_data_status or require_equal_train_batches + ) for i_epoch in epoch_iter: + # Multi-rank dumps synchronize their per-step state, so uneven workers + # must fail together instead of silently dropping a lagging rank's + # trailing delta rows or leaving collectives. pipeline = create_train_pipeline( model, optimizer, - check_all_workers_data_status=check_all_workers_data_status, + check_all_workers_data_status=sync_train_data_exhaustion, + fail_on_uneven_data=require_equal_train_batches, ) if plogger is not None: plogger.set_description(f"Training Epoch {i_epoch}") @@ -575,11 +586,10 @@ def run_eval(step: int, epoch: int) -> None: _model.on_train_end() if delta_embedding_dumper is not None: # Flush the trailing partial interval before the final checkpoint. - # final_dump skips dump-boundary steps already written by maybe_dump, - # so it never overwrites their shards with an empty file. Ranks can - # reach here at different i_step (independent dataloader exhaustion with - # check_all_workers_data_status=False), so final_dump all-reduces the - # step across ranks to keep one complete shard set per step dir. + # final_dump skips dump-boundary steps already written by maybe_dump + # (multi-rank dumps force synced exhaustion, so every rank participated + # in those dumps); the MAX all-reduce in final_dump keeps one complete + # shard set per step dir as a defensive guard. delta_embedding_dumper.final_dump(i_step) _log_train( diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 0c570d46e..031b7e0b3 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -462,6 +462,18 @@ 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. + + Multi-rank dumps of either cadence need every rank to reach the same + final step. ``final_dump`` skips steps already written on their dump + boundary, and a rank that stops before the synced MAX step never ran + that boundary's ``maybe_dump``, so an unsynced stop would make it + adopt the boundary skip and silently drop its trailing tracked rows. + """ + return self._world_size > 1 + def start(self) -> None: """Start timed cadence and per-rank FeatureStore publication. diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 276569767..b59ac843b 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -227,7 +227,7 @@ def _run_uneven_exhaustion_rejection(rank: int, world_size: int, output_dir: str # uneven stop must fail loudly on every rank instead. with MultiProcessContext(rank=rank, world_size=world_size, backend="nccl") as ctx: model = _build_sharded_delta_dump_model(rank, world_size, ctx) - DeltaEmbeddingDumper( + dumper = DeltaEmbeddingDumper( model, DeltaEmbeddingDumpConfig(dump_interval_steps=50, output_dir=output_dir), output_dir, @@ -235,9 +235,11 @@ def _run_uneven_exhaustion_rejection(rank: int, world_size: int, output_dir: str [], ) testcase = unittest.TestCase() + testcase.assertTrue(dumper.requires_synced_dataloader_exhaustion) pipeline = create_train_pipeline( model, - check_all_workers_data_status=True, + check_all_workers_data_status=dumper.requires_synced_dataloader_exhaustion, + fail_on_uneven_data=dumper.requires_synced_dataloader_exhaustion, ) # Rank 0's 50th fetch returns None while rank 1 still has batch 50; the # pipeline's collective data-status check must raise on both ranks. @@ -845,6 +847,33 @@ 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_multi_rank_step_interval_requires_synced_dataloader_exhaustion(self): + # Regression: with ranks finishing at steps 49 and 50, final_dump synced + # both to the boundary 50, both took the boundary-step skip, and the + # rank that never ran maybe_dump(50) silently dropped its trailing + # tracked rows. Every multi-rank cadence must keep exhaustion aligned + # so all ranks participate in every boundary dump. + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._interval_secs = None + dumper._interval_steps = 50 + dumper._world_size = 2 + + self.assertTrue(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_dump_does_not_ack_tracker_when_submission_is_rejected(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._feature_store_enabled = True From 39898d5ade676e1ea551672bb7ba9ef43dc59146 Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 24 Jul 2026 21:26:05 +0800 Subject: [PATCH 35/50] [bugfix] skip unhosted owners in delta dump fan-out 877ece7's per-owner fan-out calls _lookup_embeddings for every owner of a tracker-merged table name, but _collect_table_weights and _collect_dynamic_modules only gather this rank's local shards. An owner whose shard lives on another rank -- e.g. the table_wise EBC copy of item_id_emb, placed on rank 0 only, queried on rank 1 in the dynamicemb multi-GPU integration test -- is absent from both maps and raised KeyError, crashing the whole dump. Pre-filter owners this rank cannot resolve and let the hosting rank dump their rows, restoring the pre-877ece7 robustness (the model-wide collection then only ever found local shards) while keeping per-owner fan-out and identities for locally resolvable owners. The KeyError in _lookup_embeddings stays as a defensive net for genuine internal inconsistency. Co-Authored-By: Claude --- tzrec/utils/delta_embedding_dump.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 031b7e0b3..622ae0a9b 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -538,10 +538,10 @@ def _build_sparse_embedding_contract( The TorchRec tracker collapses same-named tables from different modules into one entry, so identities are keyed by ``(module_fqn, table_name)`` and ``dump`` fans the tracked ids out to - every owner, looking values up in each owner's own weights. Canonical - embedding names follow the sparse-export contract: a table name reused - across EC and EBC gets a role suffix, so both physical tables publish - under distinct serving names. + every owner this rank hosts, looking values up in each hosted owner's + own weights. Canonical embedding names follow the sparse-export + contract: a table name reused across EC and EBC gets a role suffix, so + both physical tables publish under distinct serving names. """ identity_by_owner: Dict[Tuple[str, str], SparseEmbeddingIdentity] = {} owner_roles: Dict[Tuple[str, str], str] = {} @@ -847,11 +847,17 @@ def _append_model_delta_rows( ) ids = ids.unique(sorted=True) # The tracker merges same-named tables into one id set; fan it out - # to every owning module and read each owner's own weights, so an - # id that only one owner touched still publishes that owner's true - # current row for the others. + # to every owning module this rank hosts and read each owner's own + # weights, so an id that only one owner touched still publishes + # that owner's true current row for the others. for module_fqn in owners: - identity = self._identity_by_owner[(module_fqn, table_name)] + owner_key = (module_fqn, table_name) + if owner_key not in dynamic_modules and owner_key not in table_weights: + # This rank hosts no shard for this owner (e.g. a + # table_wise table placed on another rank); the hosting + # rank dumps its rows. + continue + identity = self._identity_by_owner[owner_key] embeddings, key_ids = self._lookup_embeddings( module_fqn, table_name, From 0707194a0c1b518812be3c8e742848f724b1f6d1 Mon Sep 17 00:00:00 2001 From: gecheng Date: Sat, 25 Jul 2026 09:14:03 +0800 Subject: [PATCH 36/50] [refactor] drop getattr/_int_attr from delta embedding dump The dumper probed torchrec internals and proto fields with getattr-defaults and an _int_attr helper, hiding which attributes are type-guaranteed versus optional. Use direct access for fields the torchrec dataclasses always expose (ParameterSharding.ranks/sharding_spec, Shard.placement, shards[0].metadata, EmbeddingBagConfig.num_embeddings/embedding_dim/feature_names, the isinstance-narrowed _table_name_to_config) and hasattr guards for genuinely optional attributes on arbitrary modules (_lookups, config/_config, embedding_tables, module_sharding_plan, post_*_tracker_fn). _int_attr and _feature_config_name are inlined, _has_proto_field simplified. Semantics are unchanged: hasattr plus direct access is equivalent to the prior getattr defaults, and tests are updated to use faithful, real-type-shaped configs. The only remaining getattr accesses a proto oneof field by dynamic name (feature_config.WhichOneof -> getattr), which has no getattr-free equivalent. Co-Authored-By: Claude --- tzrec/utils/delta_embedding_dump.py | 117 +++++++++++------------ tzrec/utils/delta_embedding_dump_test.py | 23 ++++- 2 files changed, 77 insertions(+), 63 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 622ae0a9b..35e36c5d9 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -131,14 +131,9 @@ def validate_delta_embedding_dump_config( def _has_proto_field(config: Any, field_name: str) -> bool: - descriptor = getattr(config, "DESCRIPTOR", None) - if descriptor is None or field_name not in descriptor.fields_by_name: - return False - return config.HasField(field_name) - - -def _feature_config_name(config: Any) -> str: - return getattr(config, "feature_name", "") + return field_name in config.DESCRIPTOR.fields_by_name and config.HasField( + field_name + ) def _zch_feature_names(feature_configs: Iterable[Any]) -> Set[str]: @@ -147,10 +142,10 @@ def _zch_feature_names(feature_configs: Iterable[Any]) -> Set[str]: feature_type = feature_config.WhichOneof("feature") if feature_type is None: continue + # getattr is the idiomatic way to access a proto oneof field by name. config = getattr(feature_config, feature_type) if _has_proto_field(config, "zch"): - feature_name = _feature_config_name(config) or feature_type - zch_feature_names.add(feature_name) + zch_feature_names.add(config.feature_name or feature_type) return zch_feature_names @@ -184,16 +179,11 @@ def _owner_table_fqn(module_fqn: str, identity: SparseEmbeddingIdentity) -> str: return f"{module_fqn}.{segment}.{identity.table_name}" -def _int_attr(value: Any, name: str) -> int: - attr = getattr(value, name, 0) - return int(attr) if attr is not None else 0 - - def _metadata_shard_info(metadata: Any) -> _TableShardInfo: if metadata is None or not hasattr(metadata, "shard_offsets"): return _TableShardInfo() - offsets = getattr(metadata, "shard_offsets", []) - sizes = getattr(metadata, "shard_sizes", []) + offsets = metadata.shard_offsets + sizes = metadata.shard_sizes return _TableShardInfo( row_offset=int(offsets[0]) if len(offsets) > 0 else 0, column_offset=int(offsets[1]) if len(offsets) > 1 else 0, @@ -206,7 +196,7 @@ def _metadata_shard_info(metadata: Any) -> _TableShardInfo: def _placement_rank(placement: Any) -> Optional[int]: if placement is None: return None - rank_fn = getattr(placement, "rank", None) + rank_fn = placement.rank if hasattr(placement, "rank") else None if callable(rank_fn): rank = rank_fn() if rank is not None: @@ -220,14 +210,16 @@ def _placement_rank(placement: Any) -> Optional[int]: def _table_shard_info_from_parameter_sharding( parameter_sharding: Any, rank: int ) -> _TableShardInfo: - sharding_spec = getattr(parameter_sharding, "sharding_spec", None) - shards = getattr(sharding_spec, "shards", None) + sharding_spec = parameter_sharding.sharding_spec + if not hasattr(sharding_spec, "shards"): + return _TableShardInfo() + shards = sharding_spec.shards if not shards: return _TableShardInfo() - ranks = getattr(parameter_sharding, "ranks", None) + ranks = parameter_sharding.ranks for idx, shard in enumerate(shards): - placement_rank = _placement_rank(getattr(shard, "placement", None)) + placement_rank = _placement_rank(shard.placement) if placement_rank == rank: return _metadata_shard_info(shard) if ranks is not None and idx < len(ranks) and ranks[idx] == rank: @@ -256,13 +248,18 @@ def _merge_shard_info( def _table_shard_info_from_config(table_config: Any) -> _TableShardInfo: - metadata_info = _metadata_shard_info(getattr(table_config, "local_metadata", None)) + if not hasattr(table_config, "local_rows"): + return _TableShardInfo( + global_rows=int(table_config.num_embeddings), + global_cols=int(table_config.embedding_dim), + ) config_info = _TableShardInfo( - local_rows=_int_attr(table_config, "local_rows"), - local_cols=_int_attr(table_config, "local_cols"), - global_rows=_int_attr(table_config, "num_embeddings"), - global_cols=_int_attr(table_config, "embedding_dim"), + local_rows=int(table_config.local_rows), + local_cols=int(table_config.local_cols), + global_rows=int(table_config.num_embeddings), + global_cols=int(table_config.embedding_dim), ) + metadata_info = _metadata_shard_info(table_config.local_metadata) return _merge_shard_info(config_info, metadata_info) @@ -328,7 +325,7 @@ def _local_table_weight( # shard_info is merged by table name and may describe another module's # same-named table. info = _merge_shard_info( - _metadata_shard_info(getattr(shards[0], "metadata", None)), + _metadata_shard_info(shards[0].metadata), shard_info or _TableShardInfo(), ) info = _table_shard_info_from_tensor(shards[0].tensor, info) @@ -342,7 +339,7 @@ def _local_table_weight( "delta embedding dump only supports one local shard per table." ) info = _merge_shard_info( - _metadata_shard_info(getattr(shards[0], "metadata", None)), + _metadata_shard_info(shards[0].metadata), shard_info or _TableShardInfo(), ) info = _table_shard_info_from_tensor(shards[0].tensor, info) @@ -394,6 +391,7 @@ def __init__( self._file_prefix = config.file_prefix or "delta_embedding" self._rank, self._world_size = _distributed_rank_world_size() self._tracking_pause_depth = 0 + self._guarded_tracking_modules: Set[int] = set() self._feature_store_enabled = config.HasField("feature_store_config") self._retain_local_dump = self._feature_store_enabled and bool( config.feature_store_config.retain_local_dump @@ -554,9 +552,9 @@ def _build_sparse_embedding_contract( role = SPARSE_EBC_ROLE else: continue - table_name_to_config = getattr(module, "_table_name_to_config", {}) + table_name_to_config = module._table_name_to_config for table_name, table_config in table_name_to_config.items(): - dimension = _int_attr(table_config, "embedding_dim") + dimension = int(table_config.embedding_dim) if dimension <= 0: dimension = self._table_shard_infos.get( table_name, _TableShardInfo() @@ -566,7 +564,7 @@ def _build_sparse_embedding_contract( f"invalid embedding dimension for table {table_name!r}: " f"{dimension}" ) - feature_names = tuple(getattr(table_config, "feature_names", ())) + feature_names = tuple(table_config.feature_names) owner_key = (module_fqn, table_name) owner_roles[owner_key] = role owner_metadata[owner_key] = (dimension, feature_names) @@ -794,23 +792,30 @@ def _output_path(self, global_step: int) -> str: ) def _install_tracking_pause_guard(self) -> None: - guarded_modules = getattr(self, "_guarded_tracking_modules", set()) + guarded_modules = self._guarded_tracking_modules for module in self._tracker.get_tracked_modules().values(): if id(module) in guarded_modules: continue has_tracker_fn = False - post_lookup_fn = getattr(module, "post_lookup_tracker_fn", None) + post_lookup_fn = ( + module.post_lookup_tracker_fn + if hasattr(module, "post_lookup_tracker_fn") + else None + ) if post_lookup_fn is not None: module.post_lookup_tracker_fn = self._wrap_tracker_fn(post_lookup_fn) has_tracker_fn = True - post_odist_fn = getattr(module, "post_odist_tracker_fn", None) + post_odist_fn = ( + module.post_odist_tracker_fn + if hasattr(module, "post_odist_tracker_fn") + else None + ) if post_odist_fn is not None: module.post_odist_tracker_fn = self._wrap_tracker_fn(post_odist_fn) has_tracker_fn = True if not has_tracker_fn: continue guarded_modules.add(id(module)) - self._guarded_tracking_modules = guarded_modules def _wrap_tracker_fn(self, tracker_fn: Callable[..., Any]) -> Callable[..., Any]: def guarded_tracker_fn(*args: Any, **kwargs: Any) -> Any: @@ -965,25 +970,23 @@ def _lookup_dynamic_embeddings( def _collect_table_shard_infos(self) -> Dict[str, _TableShardInfo]: table_shard_infos: Dict[str, _TableShardInfo] = {} for module in self._model.modules(): - table_name_to_config = getattr(module, "_table_name_to_config", None) - if table_name_to_config is not None: - for table_name, table_config in table_name_to_config.items(): + if hasattr(module, "_table_name_to_config"): + for table_name, table_config in module._table_name_to_config.items(): table_shard_infos[table_name] = _merge_table_shard_info( table_shard_infos.get(table_name), _table_shard_info_from_config(table_config), ) for table_config in self._grouped_embedding_table_configs(module): - table_name = getattr(table_config, "name", "") + table_name = table_config.name if not table_name: continue table_shard_infos[table_name] = _merge_table_shard_info( table_shard_infos.get(table_name), _table_shard_info_from_config(table_config), ) - module_sharding_plan = getattr(module, "module_sharding_plan", None) - if module_sharding_plan is None: + if not hasattr(module, "module_sharding_plan"): continue - for table_name, parameter_sharding in module_sharding_plan.items(): + for table_name, parameter_sharding in module.module_sharding_plan.items(): table_shard_infos[table_name] = _merge_table_shard_info( table_shard_infos.get(table_name), _table_shard_info_from_parameter_sharding( @@ -993,19 +996,19 @@ def _collect_table_shard_infos(self) -> Dict[str, _TableShardInfo]: return table_shard_infos def _grouped_embedding_table_configs(self, module: nn.Module) -> Iterable[Any]: + module_config = module.config if hasattr(module, "config") else None grouped_configs = [] - module_config = getattr(module, "config", None) if module_config is not None: grouped_configs.append(module_config) - private_config = getattr(module, "_config", None) - if private_config is not None and private_config is not module_config: - grouped_configs.append(private_config) + if hasattr(module, "_config"): + private_config = module._config + if private_config is not module_config: + grouped_configs.append(private_config) for grouped_config in grouped_configs: - embedding_tables = getattr(grouped_config, "embedding_tables", None) - if embedding_tables is None: + if not hasattr(grouped_config, "embedding_tables"): continue - yield from embedding_tables + yield from grouped_config.embedding_tables def _validate_supported_table_sharding( self, table_shard_infos: Dict[str, _TableShardInfo] @@ -1033,17 +1036,13 @@ def _collect_table_weights(self) -> Dict[Tuple[str, str], _TableWeight]: table_weights: Dict[Tuple[str, str], _TableWeight] = {} table_shard_infos = self._table_shard_infos for module_fqn, module in self._tracker.get_tracked_modules().items(): - lookups = getattr(module, "_lookups", None) - if lookups is None: + if not hasattr(module, "_lookups"): continue - for lookup in lookups: - lookup = getattr(lookup, "module", lookup) - named_parameters_by_table = getattr( - lookup, "named_parameters_by_table", None - ) - if named_parameters_by_table is None: + for lookup in module._lookups: + lookup = lookup.module if hasattr(lookup, "module") else lookup + if not hasattr(lookup, "named_parameters_by_table"): continue - for table_name, table_value in named_parameters_by_table(): + for table_name, table_value in lookup.named_parameters_by_table(): table_weights[(module_fqn, table_name)] = _local_table_weight( table_value, table_shard_infos.get(table_name), diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index b59ac843b..fb612984e 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -381,6 +381,24 @@ def test_column_wise_shard_info_fails_fast(self): with self.assertRaisesRegex(ValueError, "column-wise"): _validate_table_shard_info("user_emb", shard_info) + def test_table_shard_info_from_config_without_local_dims(self): + # A non-sharded table config (e.g. EmbeddingBagConfig from + # _table_name_to_config) carries no local_rows/local_cols/ + # local_metadata; the shard info must fall back to global dims only + # and leave local dims at zero for the tensor/sharding-plan paths. + table_config = EmbeddingBagConfig( + num_embeddings=64, + embedding_dim=8, + name="user_emb", + feature_names=[_SHARDED_FEATURE_NAME], + ) + shard_info = _table_shard_info_from_config(table_config) + self.assertEqual(shard_info.local_rows, 0) + self.assertEqual(shard_info.local_cols, 0) + self.assertEqual(shard_info.global_rows, 64) + self.assertEqual(shard_info.global_cols, 8) + self.assertFalse(shard_info.has_shard_metadata) + def test_dump_rows_include_rank_metadata(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._rank = 1 @@ -983,6 +1001,7 @@ def test_pause_tracking_suppresses_post_lookup_recording(self): ) dumper = object.__new__(DeltaEmbeddingDumper) dumper._tracking_pause_depth = 0 + dumper._guarded_tracking_modules = set() dumper._tracker = SimpleNamespace( get_tracked_modules=lambda: {"user_emb": sharded_module} ) @@ -1014,8 +1033,6 @@ def test_collect_table_shard_infos_prefers_grouped_embedding_metadata(self): original_config_module = torch.nn.Module() original_config_module._table_name_to_config = { "user_emb": SimpleNamespace( - local_rows=16, - local_cols=8, num_embeddings=64, embedding_dim=8, ) @@ -1048,8 +1065,6 @@ def test_collect_table_shard_infos_falls_back_to_sharding_plan(self): sharded_module = torch.nn.Module() sharded_module._table_name_to_config = { "adgroup_id_emb": SimpleNamespace( - local_rows=16, - local_cols=8, num_embeddings=64, embedding_dim=8, ) From 272fdb86923127e4673dcac87535e6dfc31fcba2 Mon Sep 17 00:00:00 2001 From: gecheng Date: Sat, 25 Jul 2026 11:17:19 +0800 Subject: [PATCH 37/50] [refactor] rename identity tuple-key vars in sparse contract `identity` was overloaded: it denoted both a `(role, table_name)` dict key (the `table_identities` param / `name_by_identity` local in the contract and delta dump) and a `SparseEmbeddingIdentity` instance, confusing readers. Rename the tuple-keyed identifiers to `role_table_pairs` / `name_by_role_table` so the key shape is explicit, and note that `identity_by_owner`'s key is `(module_fqn, table_name)` -- distinct from the contract's `(role, table_name)`. The `SparseEmbeddingIdentity` dataclass and its instance locals stay (the legitimate singular). Pure rename; all callers are positional, so behavior and the exported embedding_name are unchanged. Co-Authored-By: Claude --- tzrec/utils/delta_embedding_dump.py | 5 +++-- tzrec/utils/sparse_embedding_contract.py | 20 ++++++++++---------- 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 35e36c5d9..d4a6ed513 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -541,6 +541,7 @@ def _build_sparse_embedding_contract( contract: a table name reused across EC and EBC gets a role suffix, so both physical tables publish under distinct serving names. """ + # Keyed by (module_fqn, table_name); cf. the role-keyed name map below. identity_by_owner: Dict[Tuple[str, str], SparseEmbeddingIdentity] = {} owner_roles: Dict[Tuple[str, str], str] = {} owner_metadata: Dict[Tuple[str, str], Tuple[int, Tuple[str, ...]]] = {} @@ -587,13 +588,13 @@ def _build_sparse_embedding_contract( f"{duplicated}. Give the tables distinct embedding_names." ) - name_by_identity = build_sparse_embedding_name_map(fqns_by_role_table.keys()) + name_by_role_table = build_sparse_embedding_name_map(fqns_by_role_table.keys()) embedding_dimensions: Dict[str, int] = {} for owner_key, role in owner_roles.items(): module_fqn, table_name = owner_key dimension, feature_names = owner_metadata[owner_key] embedding_name = resolve_sparse_embedding_name( - name_by_identity, table_name, role + name_by_role_table, table_name, role ) identity_by_owner[owner_key] = SparseEmbeddingIdentity( role=role, diff --git a/tzrec/utils/sparse_embedding_contract.py b/tzrec/utils/sparse_embedding_contract.py index 3aae3b58d..23edbb01b 100644 --- a/tzrec/utils/sparse_embedding_contract.py +++ b/tzrec/utils/sparse_embedding_contract.py @@ -35,7 +35,7 @@ class SparseEmbeddingIdentity: def build_sparse_embedding_name_map( - table_identities: Iterable[Tuple[str, str]], + role_table_pairs: Iterable[Tuple[str, str]], ) -> Dict[Tuple[str, str], str]: """Allocate canonical names for ``(collection role, table name)`` pairs. @@ -45,7 +45,7 @@ def build_sparse_embedding_name_map( raw table names are resolved deterministically with a numeric suffix. """ roles_by_name = defaultdict(set) - for role, table_name in table_identities: + for role, table_name in role_table_pairs: if role not in SPARSE_EMBEDDING_ROLES: raise ValueError(f"unsupported sparse embedding role: {role!r}") if not table_name: @@ -53,11 +53,11 @@ def build_sparse_embedding_name_map( roles_by_name[table_name].add(role) used_names = set(roles_by_name.keys()) - name_by_identity: Dict[Tuple[str, str], str] = {} + name_by_role_table: 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 + name_by_role_table[(role, table_name)] = table_name continue for role in sorted(roles): @@ -68,22 +68,22 @@ def build_sparse_embedding_name_map( candidate = f"{base_candidate}_{suffix}" suffix += 1 used_names.add(candidate) - name_by_identity[(role, table_name)] = candidate - return name_by_identity + name_by_role_table[(role, table_name)] = candidate + return name_by_role_table def resolve_sparse_embedding_name( - name_by_identity: Dict[Tuple[str, str], str], + name_by_role_table: 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)] + if role is not None and (role, table_name) in name_by_role_table: + return name_by_role_table[(role, table_name)] candidates = [ embedding_name - for (_role, name), embedding_name in name_by_identity.items() + for (_role, name), embedding_name in name_by_role_table.items() if name == table_name ] if len(candidates) == 1: From f62d86163bd97899adcfa9eeac513eceb2b84252 Mon Sep 17 00:00:00 2001 From: gecheng Date: Sat, 25 Jul 2026 11:17:22 +0800 Subject: [PATCH 38/50] [refactor] source sparse export names from shared contract export_util duplicated the EC/EBC role-suffix naming algorithm that sparse_embedding_contract already owns, risking drift between the exported feat_meta["embedding_name"] (read by serving) and the delta-dump name. Drop the private role constants and the _build_sparse_export_name_map / _resolve_sparse_export_name copies; import the shared build_sparse_embedding_name_map / resolve_sparse_embedding_name instead. _build_sparse_export_name_map stays as a thin adapter (two config lists -> (role, name) pairs) and the four resolve call sites call the shared function directly. Output is byte-identical for every realistic config (differential fuzz over 610 cases incl. shared-table and raw-candidate collisions); the only behavior change is raising on an empty table name, which real configs never produce (feature.py `embedding_name or f"{name}_emb"` fallback). This addresses the instinct behind review #2 without the unsafe FQN switch: embedding_name is a cross-system serving key (export + delta dump + serving share it) and the export path has no module FQN, so a delta-side FQN would diverge and orphan every delta row. Co-Authored-By: Claude --- tzrec/utils/export_util.py | 83 ++++++++++---------------------------- 1 file changed, 22 insertions(+), 61 deletions(-) diff --git a/tzrec/utils/export_util.py b/tzrec/utils/export_util.py index 377626271..395c93919 100644 --- a/tzrec/utils/export_util.py +++ b/tzrec/utils/export_util.py @@ -77,6 +77,12 @@ ) from tzrec.utils.logging_util import logger from tzrec.utils.plan_util import create_planner, get_default_sharders +from tzrec.utils.sparse_embedding_contract import ( + SPARSE_EBC_ROLE, + SPARSE_EC_ROLE, + build_sparse_embedding_name_map, + resolve_sparse_embedding_name, +) from tzrec.utils.state_dict_util import fix_mch_state, init_parameters @@ -1744,10 +1750,6 @@ def _get_sparse_feature_to_embedding_info( return feature_to_embedding_bag_info, feature_to_embedding_info -_SPARSE_EC_ROLE = "ec" -_SPARSE_EBC_ROLE = "ebc" - - def _build_sparse_export_name_map( embedding_infos: List[BaseEmbeddingConfig], embedding_bag_info: List[BaseEmbeddingConfig], @@ -1759,38 +1761,20 @@ def _build_sparse_export_name_map( but those are still distinct checkpoint tensors. Keep the historical name unless it appears in multiple sparse collection kinds; then suffix the exported table name so serving can address the correct physical table. - """ - roles_by_name = defaultdict(set) - for emb_info in embedding_infos: - roles_by_name[emb_info.name].add(_SPARSE_EC_ROLE) - for emb_info in embedding_bag_info: - roles_by_name[emb_info.name].add(_SPARSE_EBC_ROLE) - - used_names = set(roles_by_name.keys()) - export_name_by_role: Dict[Tuple[str, str], str] = {} - for emb_name, roles in roles_by_name.items(): - if len(roles) == 1: - role = next(iter(roles)) - export_name_by_role[(role, emb_name)] = emb_name - continue - for role in sorted(roles): - base_candidate = f"{emb_name}__{role}" - candidate = base_candidate - suffix = 1 - while candidate in used_names: - candidate = f"{base_candidate}_{suffix}" - suffix += 1 - used_names.add(candidate) - export_name_by_role[(role, emb_name)] = candidate - return export_name_by_role + The naming algorithm is shared with delta dump via + ``sparse_embedding_contract`` so export and serving stay aligned. + """ + role_table_pairs = [(SPARSE_EC_ROLE, e.name) for e in embedding_infos] + role_table_pairs += [(SPARSE_EBC_ROLE, e.name) for e in embedding_bag_info] + return build_sparse_embedding_name_map(role_table_pairs) def _sparse_export_role_from_state_key(state_key: str) -> Optional[str]: if ".embedding_bags." in state_key: - return _SPARSE_EBC_ROLE + return SPARSE_EBC_ROLE if ".embeddings." in state_key: - return _SPARSE_EC_ROLE + return SPARSE_EC_ROLE return None @@ -1802,35 +1786,12 @@ def _sparse_export_role_from_dynamic_path(path: str) -> Optional[str]: or ".mc_ec" in parent or ".ec." in parent ): - return _SPARSE_EC_ROLE + return SPARSE_EC_ROLE if ".ebc" in parent or ".mc_ebc" in parent: - return _SPARSE_EBC_ROLE + return SPARSE_EBC_ROLE return None -def _resolve_sparse_export_name( - export_name_by_role: Dict[Tuple[str, str], str], - emb_name: str, - role: Optional[str], -) -> str: - if role is not None and (role, emb_name) in export_name_by_role: - return export_name_by_role[(role, emb_name)] - - candidates = [ - export_name - for (_role, name), export_name in export_name_by_role.items() - if name == emb_name - ] - if len(candidates) == 1: - return candidates[0] - if not candidates: - raise KeyError(f"sparse embedding {emb_name} is not in export metadata") - raise ValueError( - f"sparse embedding {emb_name} appears in multiple sparse collection " - f"kinds; cannot resolve export name without role, got role={role}" - ) - - def _remove_torch_prefix(src: str) -> str: prefix = "torch." if src.startswith(prefix): @@ -1917,8 +1878,8 @@ def _get_sparse_embedding_tensor( ) for emb_info in embedding_infos: - export_emb_name = _resolve_sparse_export_name( - export_name_by_role, emb_info.name, _SPARSE_EC_ROLE + export_emb_name = resolve_sparse_embedding_name( + export_name_by_role, emb_info.name, SPARSE_EC_ROLE ) emb_name_to_emb_dim[export_emb_name] = emb_info.embedding_dim emb_name_to_feat_name_impl[export_emb_name] = [] @@ -1933,8 +1894,8 @@ def _get_sparse_embedding_tensor( else: feat_name_to_pooling[feat_name_impl] = "NONE" for emb_info in embedding_bag_info: - export_emb_name = _resolve_sparse_export_name( - export_name_by_role, emb_info.name, _SPARSE_EBC_ROLE + export_emb_name = resolve_sparse_embedding_name( + export_name_by_role, emb_info.name, SPARSE_EBC_ROLE ) emb_name_to_emb_dim[export_emb_name] = emb_info.embedding_dim emb_name_to_feat_name_impl[export_emb_name] = [] @@ -1966,7 +1927,7 @@ def _get_sparse_embedding_tensor( world_size = int(os.environ.get("WORLD_SIZE", 1)) for name, values in model.state_dict().items(): emb_name = name.split(".")[-2] - export_emb_name = _resolve_sparse_export_name( + export_emb_name = resolve_sparse_embedding_name( export_name_by_role, emb_name, _sparse_export_role_from_state_key(name), @@ -2020,7 +1981,7 @@ def _get_sparse_embedding_tensor( if not match: continue emb_name = match.group("emb_name") - export_emb_name = _resolve_sparse_export_name( + export_emb_name = resolve_sparse_embedding_name( export_name_by_role, emb_name, _sparse_export_role_from_dynamic_path(key_file), From 6505e8186f1bfeb8e2ed5ae1f3e4f4e03732941b Mon Sep 17 00:00:00 2001 From: gecheng Date: Sun, 26 Jul 2026 12:15:22 +0800 Subject: [PATCH 39/50] [refactor] drop over-defensive FeatureStore endpoint and view validation Remove the endpoint allowlist and control-plane FeatureView metadata validation that review flagged as over-engineered: endpoint selection is an explicit operator action owned by the FeatureStore SDK, and view schema compatibility is already enforced on the data-plane writer. Also drop the redundant invalid_dimensions check and validate_feature_store_config helper, which duplicate validation performed when the uploader settings are built from the proto. Co-Authored-By: Claude --- tzrec/protos/train.proto | 14 +- tzrec/utils/delta_embedding_dump.py | 7 +- tzrec/utils/feature_store_delta_uploader.py | 207 +----------------- .../feature_store_delta_uploader_test.py | 174 +-------------- 4 files changed, 13 insertions(+), 389 deletions(-) diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index 625fe1ee1..16904c093 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -47,11 +47,8 @@ message FeatureStoreConfig { required string version = 5; - // Optional FeatureStore control-plane endpoint override. Cloud and - // FeatureDB credentials are sent to this host, so it must be a trusted - // *.aliyuncs.com FeatureStore endpoint (an explicit https:// scheme is - // required when a scheme is given) without userinfo, port, path, query or - // fragment components, unless allow_custom_endpoint is set. + // Optional FeatureStore control-plane endpoint override, passed directly to + // the FeatureStore SDK, which owns endpoint handling and connection errors. optional string endpoint = 6; // Maximum records submitted before each SDK write_flush() completion gate. optional uint32 upload_batch_size = 7 [default = 1000]; @@ -70,11 +67,8 @@ message FeatureStoreConfig { optional uint32 feature_view_ttl_secs = 13 [default = 1296000]; optional uint32 feature_view_shard_count = 14 [default = 20]; optional uint32 feature_view_replication_count = 15 [default = 1]; - // Explicit opt-in allowing endpoint to name a non-*.aliyuncs.com host - // (e.g. a vetted private FeatureStore deployment with a custom port). - // Credentials are sent to that host; structural endpoint checks (HTTPS - // scheme, no userinfo/path/query/fragment) still apply. - optional bool allow_custom_endpoint = 16 [default = false]; + reserved 16; + reserved "allow_custom_endpoint"; // Also write the local per-rank delta parquet files while uploading to // FeatureStore (e.g. for the offline readback checker); by default the // delta is handed to the uploader in memory and never touches disk. diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index d4a6ed513..1ef2c0e0d 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -29,10 +29,7 @@ from torchrec.distributed.model_tracker.types import TrackingMode from tzrec.protos.train_pb2 import DeltaEmbeddingDumpConfig -from tzrec.utils.feature_store_delta_uploader import ( - FeatureStoreDeltaUploader, - validate_feature_store_config, -) +from tzrec.utils.feature_store_delta_uploader import FeatureStoreDeltaUploader from tzrec.utils.logging_util import logger from tzrec.utils.sparse_embedding_contract import ( SPARSE_EBC_ROLE, @@ -126,8 +123,6 @@ def validate_delta_embedding_dump_config( ) 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) def _has_proto_field(config: Any, field_name: str) -> bool: diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index c651d4fe2..ecb66f232 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -16,7 +16,6 @@ latest checkpoint and pending deltas are discarded. """ -import json import os import threading import time @@ -31,7 +30,6 @@ Tuple, cast, ) -from urllib.parse import urlsplit import numpy as np import pyarrow as pa @@ -79,7 +77,6 @@ class FeatureStoreUploadSettings: shutdown_timeout_secs: int max_pending_steps: int poll_interval_secs: int - allow_custom_endpoint: bool @classmethod def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": @@ -98,8 +95,6 @@ def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": "feature_store_config.region must not be empty " "(it may come from ALIBABA_CLOUD_REGION)" ) - allow_custom_endpoint = bool(config.allow_custom_endpoint) - _validate_feature_store_endpoint(endpoint, allow_custom_endpoint) project_name = config.project_name.strip() feature_entity_name = config.feature_entity_name.strip() feature_view_name = config.feature_view_name.strip() @@ -161,65 +156,9 @@ def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": shutdown_timeout_secs=positive_values["shutdown_timeout_secs"], max_pending_steps=positive_values["max_pending_steps"], poll_interval_secs=positive_values["poll_interval_secs"], - allow_custom_endpoint=allow_custom_endpoint, ) -def validate_feature_store_config(config: FeatureStoreConfig) -> None: - """Validate upload configuration without resolving credentials.""" - FeatureStoreUploadSettings.from_proto(config) - - -def _validate_feature_store_endpoint( - endpoint: str, allow_custom_endpoint: bool -) -> None: - """Reject endpoints that could divert cloud or FeatureDB credentials. - - The endpoint receives the access keys, STS token and FeatureDB - credentials, so only trusted Alibaba Cloud hosts are accepted unless the - operator explicitly opts in to a vetted custom deployment. - """ - if not endpoint: - return - parsed_endpoint = urlsplit(endpoint if "://" in endpoint else f"//{endpoint}") - if parsed_endpoint.scheme and parsed_endpoint.scheme != "https": - raise ValueError( - "feature_store_config.endpoint must use HTTPS when set; " - f"got scheme={parsed_endpoint.scheme!r}" - ) - if parsed_endpoint.username or parsed_endpoint.password: - raise ValueError( - "feature_store_config.endpoint must not contain userinfo credentials" - ) - if parsed_endpoint.path not in ("", "/"): - raise ValueError( - "feature_store_config.endpoint must not contain a path, query, or fragment" - ) - if parsed_endpoint.query or parsed_endpoint.fragment: - raise ValueError( - "feature_store_config.endpoint must not contain a path, query, or fragment" - ) - if parsed_endpoint.port is not None and not allow_custom_endpoint: - raise ValueError( - "feature_store_config.endpoint must not contain an explicit port" - ) - hostname = (parsed_endpoint.hostname or "").lower() - if not hostname: - raise ValueError("feature_store_config.endpoint must contain a hostname") - trusted_suffixes = ( - ".aliyuncs.com", - ".alibabacloud.com", - ".aliyun.com", - ) - if not hostname.endswith(trusted_suffixes): - if not allow_custom_endpoint: - raise ValueError( - "feature_store_config.endpoint must be a trusted Alibaba Cloud " - f"host (got {hostname!r}); set allow_custom_endpoint only for a " - "vetted private deployment" - ) - - class FeatureStoreDeltaUploader: """Ephemeral per-rank uploader for in-memory delta-embedding tables. @@ -253,16 +192,6 @@ def __init__( 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._credentials_client = ( @@ -710,9 +639,11 @@ def _reset_view(self, suppress_errors: bool = False) -> None: def _get_or_create_view(self, project: Any) -> Any: """Return the configured DynamicEmbedding view, creating it if absent. - Only the primary (rank-zero) uploader creates the view and validates - control-plane metadata; other ranks open a handle to the view that the - primary published before they started. + Only the primary (rank-zero) uploader creates the view; other ranks open + a handle to the view that the primary published before they started. + Schema compatibility is checked once on the data-plane writer in + ``_get_view``; creation-time provisioning (TTL/shard/replication) does + not affect upload compatibility and is not re-validated here. """ if not self._manage_remote_view: view = project.get_dynamic_embedding_feature_view( @@ -732,19 +663,7 @@ def _get_or_create_view(self, project: Any) -> Any: 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: + if view is None: create_error: Optional[Exception] = None try: view = project.create_dynamic_embedding_feature_view( @@ -772,15 +691,6 @@ def _get_or_create_view(self, project: Any) -> Any: raise error from create_error raise error self._view = view - metadata = self._wait_for_feature_view_metadata(project) - - if metadata is None: - 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", @@ -814,111 +724,6 @@ def _wait_for_dynamic_embedding_view(self, project: Any) -> Any: ) 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}" - ) - @staticmethod def _validate_flush_summary( summary: Any, expected_records: int, expected_batches: int diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index 650ee7904..ddedf7781 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -9,7 +9,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json import os import threading import unittest @@ -131,83 +130,32 @@ def close(self, wait=True): 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 @@ -285,7 +233,6 @@ def test_proto_groups_required_fields_before_optional_fields(self): "feature_view_ttl_secs", "feature_view_shard_count", "feature_view_replication_count", - "allow_custom_endpoint", "retain_local_dump", ] fields = list(FeatureStoreConfig.DESCRIPTOR.fields) @@ -307,7 +254,7 @@ def test_proto_groups_required_fields_before_optional_fields(self): ) self.assertEqual( [field.number for field in fields], - list(range(1, 18)), + list(range(1, 16)) + [17], ) for field_name in required_fields: with self.subTest(field_name=field_name): @@ -347,10 +294,6 @@ def test_region_fallback_and_config_validation(self): settings = FeatureStoreUploadSettings.from_proto(config) self.assertEqual(settings.region, "cn-env") - with self.assertRaisesRegex(ValueError, "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) @@ -364,60 +307,6 @@ def test_region_fallback_and_config_validation(self): _feature_store_config(feature_view_replication_count=4) ) - def test_endpoint_must_be_a_trusted_https_host(self): - for endpoint in ( - "featurestore.cn-hangzhou.aliyuncs.com", - "https://featurestore.cn-hangzhou.aliyuncs.com", - "https://featurestore-vpc.cn-beijing.aliyuncs.com/", - ): - with self.subTest(endpoint=endpoint): - settings = FeatureStoreUploadSettings.from_proto( - _feature_store_config(endpoint=endpoint) - ) - self.assertEqual(settings.endpoint, endpoint) - self.assertFalse(settings.allow_custom_endpoint) - - rejected_endpoints = { - "http://featurestore.cn-hangzhou.aliyuncs.com": "HTTPS", - "ftp://featurestore.cn-hangzhou.aliyuncs.com": "HTTPS", - "evil.example.com": "trusted", - "https://aliyuncs.com": "trusted", - "https://featurestore.cn-hangzhou.aliyuncs.com.evil.com": "trusted", - "featurestore.cn-hangzhou.aliyuncs.com:8443": "port", - "https://featurestore.cn-hangzhou.aliyuncs.com/v1": "path, query", - "https://featurestore.cn-hangzhou.aliyuncs.com?x=1": "path, query", - "https://featurestore.cn-hangzhou.aliyuncs.com#frag": "fragment", - "https://": "host", - } - for endpoint, message in rejected_endpoints.items(): - with self.subTest(endpoint=endpoint): - with self.assertRaisesRegex(ValueError, message): - FeatureStoreUploadSettings.from_proto( - _feature_store_config(endpoint=endpoint) - ) - - def test_allow_custom_endpoint_opts_into_vetted_hosts(self): - settings = FeatureStoreUploadSettings.from_proto( - _feature_store_config( - endpoint="https://featurestore.internal.example.com:8443", - allow_custom_endpoint=True, - ) - ) - self.assertTrue(settings.allow_custom_endpoint) - - for endpoint in ( - "http://featurestore.internal.example.com", - "https://user:secret@featurestore.internal.example.com", - "https://featurestore.internal.example.com/v1", - ): - with self.subTest(endpoint=endpoint): - with self.assertRaisesRegex(ValueError, "endpoint"): - FeatureStoreUploadSettings.from_proto( - _feature_store_config( - endpoint=endpoint, allow_custom_endpoint=True - ) - ) - def test_start_reuses_existing_dynamic_embedding_feature_view(self): view = _FakeView() factory = _FakeClientFactory(view) @@ -427,7 +316,6 @@ def test_start_reuses_existing_dynamic_embedding_feature_view(self): 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.assertNotIn("test_mode", factory.calls[0]) self.assertEqual(view.closed, [True]) @@ -440,10 +328,7 @@ def test_start_creates_missing_dynamic_embedding_feature_view(self): uploader.start() uploader.close() - self.assertEqual( - factory.project.generic_get_calls, - ["shared_embeddings", "shared_embeddings"], - ) + self.assertEqual(factory.project.dynamic_get_calls, ["shared_embeddings"]) self.assertEqual( factory.project.create_calls, [ @@ -463,59 +348,6 @@ def test_start_creates_missing_dynamic_embedding_feature_view(self): ) self.assertEqual(created_view.closed, [True]) - def test_start_rejects_same_name_non_dynamic_feature_view(self): - generic_view = mock.Mock(type="Batch") - factory = _FakeClientFactory(None, generic_view=generic_view) - uploader = self._uploader(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): - view = _FakeView() - generic_view = _FakeGenericFeatureView(entity="another_entity") - factory = _FakeClientFactory(view, generic_view=generic_view) - uploader = self._uploader(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): - view = _FakeView() - generic_view = _FakeGenericFeatureView() - generic_view.fields_dict["embedding"]["Type"] = "ARRAY" - factory = _FakeClientFactory(view, generic_view=generic_view) - uploader = self._uploader(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): - view = _FakeView() - generic_view = _FakeGenericFeatureView( - provisioning={ - "ttl": 60, - "shard_count": 20, - "replication_count": 1, - } - ) - factory = _FakeClientFactory(view, generic_view=generic_view) - uploader = self._uploader(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): concurrent_view = _FakeView() factory = _FakeClientFactory( @@ -582,7 +414,6 @@ def test_non_primary_uploader_opens_view_without_create_or_metadata_checks(self) uploader.close() self.assertEqual(factory.project.dynamic_get_calls, ["shared_embeddings"]) - self.assertEqual(factory.project.generic_get_calls, []) self.assertEqual(factory.project.create_calls, []) self.assertEqual(len(view.calls), 1) self.assertEqual(view.calls[0]["data"][0]["key_id"], 7) @@ -599,7 +430,6 @@ def test_non_primary_uploader_fails_when_view_is_missing(self): uploader.start() self.assertEqual(factory.project.create_calls, []) - self.assertEqual(factory.project.generic_get_calls, []) def test_submit_requires_started_uploader(self): uploader = self._uploader(client_factory=_FakeClientFactory(_FakeView())) From a9dffb6d1383dffda9f5d78f32fb40cc19f1e98e Mon Sep 17 00:00:00 2001 From: gecheng Date: Sun, 26 Jul 2026 12:22:07 +0800 Subject: [PATCH 40/50] [refactor] use a deque for pending uploads and quiet per-dump logs Submissions arrive in training order, so the pending queue is a plain in-process FIFO: replace the step-keyed dict (with min() and a setdefault dedup guard) with a collections.deque. Also demote the per-rank, per-dump log lines to DEBUG, since their volume scaled with world size and dump frequency; periodic upload progress and the final completion stay at INFO. Co-Authored-By: Claude --- tzrec/utils/delta_embedding_dump.py | 8 ++++---- tzrec/utils/feature_store_delta_uploader.py | 21 ++++++++++----------- 2 files changed, 14 insertions(+), 15 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 1ef2c0e0d..52e66c055 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -755,21 +755,21 @@ def dump(self, global_step: int) -> Optional[str]: if num_rows == 0: if output_path is None: - logger.info("No delta embedding rows to dump at step %s.", global_step) + logger.debug("No delta embedding rows to dump at step %s.", global_step) else: - logger.info( + logger.debug( "Dumped empty delta embedding shard to %s at step %s.", output_path, global_step, ) elif output_path is None: - logger.info( + logger.debug( "Submitted %s delta embedding rows for FeatureStore upload at step %s.", num_rows, global_step, ) else: - logger.info("Dumped %s delta embedding rows to %s.", num_rows, output_path) + logger.debug("Dumped %s delta embedding rows to %s.", num_rows, output_path) return output_path def _output_path(self, global_step: int) -> str: diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index ecb66f232..342185f2a 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -19,10 +19,12 @@ import os import threading import time +from collections import deque from dataclasses import dataclass from typing import ( Any, Callable, + Deque, Dict, List, Mapping, @@ -200,7 +202,7 @@ def __init__( self._clock_ms = clock_ms or (lambda: time.time_ns() // 1_000_000) self._view = None self._condition = threading.Condition() - self._pending: Dict[int, pa.Table] = {} + self._pending: Deque[Tuple[int, pa.Table]] = deque() self._started = False self._closing = False self._aborting = False @@ -251,13 +253,10 @@ def submit(self, global_step: int, table: pa.Table) -> None: ) if self._closing or self._closed: raise RuntimeError("cannot submit to a closing FeatureStore uploader") - while ( - global_step not in self._pending - and len(self._pending) >= self._settings.max_pending_steps - ): + while 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, table) + self._pending.append((global_step, table)) self._condition.notify_all() def check_error(self) -> None: @@ -309,13 +308,12 @@ def _run(self) -> None: return self._condition.wait(self._settings.poll_interval_secs) continue - current_step = min(self._pending) - table = self._pending[current_step] + current_step, table = self._pending[0] self._upload_with_retries(current_step, table) with self._condition: - self._pending.pop(current_step, None) + self._pending.popleft() self._condition.notify_all() current_step = None except _UploadAborted: @@ -402,7 +400,7 @@ def _stream_upload(self, global_step: int, table: pa.Table) -> None: next_progress_batch = _FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES logged_first_window = False - logger.info( + logger.debug( "FeatureStore delta upload started: step=%s rank=%s version=%s " "batches=%s ts_range=%s-%s", global_step, @@ -444,7 +442,8 @@ def _stream_upload(self, global_step: int, table: pa.Table) -> None: or completed_batches >= next_progress_batch or completed_batches == total_batches ): - logger.info( + log_progress = logger.info if logged_first_window else logger.debug + log_progress( "FeatureStore delta upload progress: step=%s " "batches=%s/%s elapsed_secs=%.1f", global_step, From a066e080eddbc60780501813c4aff19c12ac3b8e Mon Sep 17 00:00:00 2001 From: gecheng Date: Sun, 26 Jul 2026 12:35:12 +0800 Subject: [PATCH 41/50] [perf] cache sparse table weights and dynamic modules across dumps Each dump traversed every tracked module to rebuild the owner-to-weight and owner-to-dynamic-module maps, even though those maps never change during a run: the owning modules and their local shard tensors keep a stable identity and training updates the storage in place. Discover both once on the first dump and reuse the cached references, so later dumps read the current rows through the same handles without re-traversing. Co-Authored-By: Claude --- tzrec/utils/delta_embedding_dump.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 52e66c055..423c70eda 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -424,6 +424,8 @@ def __init__( "cannot resolve owning sparse embedding collections for tracked " f"tables: {unowned_tables}" ) + self._cached_table_weights: Optional[Dict[Tuple[str, str], _TableWeight]] = None + self._cached_dynamic_modules: Optional[Dict[Tuple[str, str], nn.Module]] = None self._uploader: Optional[FeatureStoreDeltaUploader] = None if self._feature_store_enabled: self._uploader = FeatureStoreDeltaUploader( @@ -1029,6 +1031,14 @@ def _validate_row_shard_metadata( ) def _collect_table_weights(self) -> Dict[Tuple[str, str], _TableWeight]: + """Return per-owner local table weights, discovered once and cached. + + The owning modules and their local shard tensors keep a stable identity + for the life of the process (training updates the storage in place), so + the references discovered on the first dump stay valid for later dumps. + """ + if self._cached_table_weights is not None: + return self._cached_table_weights table_weights: Dict[Tuple[str, str], _TableWeight] = {} table_shard_infos = self._table_shard_infos for module_fqn, module in self._tracker.get_tracked_modules().items(): @@ -1043,18 +1053,28 @@ def _collect_table_weights(self) -> Dict[Tuple[str, str], _TableWeight]: table_value, table_shard_infos.get(table_name), ) + self._cached_table_weights = table_weights return table_weights def _collect_dynamic_modules(self) -> Dict[Tuple[str, str], nn.Module]: + """Return per-owner dynamicemb modules, discovered once and cached. + + Only the module handles are cached; each dump still flushes and reads + the current values through them. + """ + if self._cached_dynamic_modules is not None: + return self._cached_dynamic_modules try: from dynamicemb.dump_load import get_dynamic_emb_module except ImportError: - return {} + self._cached_dynamic_modules = {} + return self._cached_dynamic_modules modules: Dict[Tuple[str, str], nn.Module] = {} for module_fqn, module in self._tracker.get_tracked_modules().items(): for dynamic_module in get_dynamic_emb_module(module): for table_name in dynamic_module.table_names: modules[(module_fqn, table_name)] = dynamic_module + self._cached_dynamic_modules = modules return modules def _append_table_chunk( From 483ed66ddb1b10f2aefe3444537e871543a4c319 Mon Sep 17 00:00:00 2001 From: gecheng Date: Sun, 26 Jul 2026 15:03:44 +0800 Subject: [PATCH 42/50] [refactor] drop redundant final-step sync and tracker rollback in delta dump Multi-rank dumps already force synced dataloader exhaustion, so every rank reaches the same final step before final_dump(); the MAX all-reduce that re-aligned it was redundant and is removed in favor of the local step. The tracker read rollback is also removed: get_unique() deletes the consumed rows, so restoring only the consumer cursor could never recover them, and under the best-effort contract a dump failure simply propagates and torchrun tears down the job. Co-Authored-By: Claude --- tzrec/utils/delta_embedding_dump.py | 65 +++++------------ tzrec/utils/delta_embedding_dump_test.py | 89 ++---------------------- 2 files changed, 24 insertions(+), 130 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 423c70eda..4d3d3c898 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -609,15 +609,6 @@ def _build_sparse_embedding_contract( embedding_dimensions[embedding_name] = dimension return identity_by_owner, embedding_dimensions - def _tracker_cursor_before_read(self) -> Optional[int]: - if not self._feature_store_enabled: - 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 - @contextmanager def pause_tracking(self) -> Iterator[None]: """Temporarily skip delta tracking for non-training forward passes.""" @@ -680,7 +671,6 @@ def final_dump(self, global_step: int) -> Optional[str]: Returns: 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) @@ -702,22 +692,6 @@ def final_dump(self, global_step: int) -> Optional[str]: return None return self.dump(global_step) - def _sync_final_step(self, global_step: int) -> int: - """Align the final step across ranks before the trailing flush. - - The MAX all-reduce ensures every rank takes the same skip/dump - decision into the same ``step_/`` directory. - """ - synced_step = global_step - if self._world_size > 1 and ( - torch.distributed.is_available() and torch.distributed.is_initialized() - ): - device = torch.device(f"cuda:{torch.cuda.current_device()}") - final_state = torch.tensor([global_step], dtype=torch.long, device=device) - torch.distributed.all_reduce(final_state, op=torch.distributed.ReduceOp.MAX) - synced_step = int(final_state[0].item()) - return synced_step - def dump(self, global_step: int) -> Optional[str]: """Dump currently tracked sparse ids and embeddings to a parquet file. @@ -732,28 +706,23 @@ def dump(self, global_step: int) -> Optional[str]: raise ValueError("delta embedding dump global_step must be > 0") uploader = self._uploader write_local = not self._feature_store_enabled or self._retain_local_dump - 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, - ) - output_path: Optional[str] = None - if write_local and (num_rows > 0 or self._world_size > 1): - # Multi-rank shard sets stay complete even for an empty rank so - # per-step file consumers never observe a partial set. - output_path = self._output_path(global_step) - self._write_table_chunks(table_chunks, output_path) - if uploader is not None and num_rows > 0: - uploader.submit(global_step, pa.concat_tables(table_chunks)) - except BaseException: - self._rollback_tracker_read(tracker_cursor) - raise + 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, + ) + output_path: Optional[str] = None + if write_local and (num_rows > 0 or self._world_size > 1): + # Multi-rank shard sets stay complete even for an empty rank so + # per-step file consumers never observe a partial set. + output_path = self._output_path(global_step) + self._write_table_chunks(table_chunks, output_path) + if uploader is not None and num_rows > 0: + uploader.submit(global_step, pa.concat_tables(table_chunks)) if num_rows == 0: if output_path is None: diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index fb612984e..977d1d3dc 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -593,66 +593,21 @@ def test_final_dump_skips_boundary_step_to_avoid_overwrite(self): [73], ) - def test_final_dump_syncs_step_across_ranks_before_flush(self): - # A lagging rank reaches final_dump at a boundary step (50) while the - # furthest rank stopped at 73. Without syncing, the lagging rank would - # 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. + def test_final_dump_uses_local_step_without_cross_rank_collective(self): + # Multi-rank dumps force synced dataloader exhaustion, so every rank + # reaches the same final step; final_dump decides on the local step and + # issues no cross-rank collective. dumper = object.__new__(DeltaEmbeddingDumper) dumper._interval_steps = 50 dumper._interval_secs = None dumper._world_size = 2 - - def fake_all_reduce(tensor, op=None): - self.assertIs(op, torch.distributed.ReduceOp.MAX) - tensor[0] = 73 - 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), + mock.patch("torch.distributed.all_reduce") as all_reduce_mock, ): - dumper.final_dump(50) + dumper.final_dump(73) dump_mock.assert_called_once_with(73) - - def test_final_dump_syncs_ranks_before_skipping_step_zero(self): - dumper = object.__new__(DeltaEmbeddingDumper) - dumper._interval_steps = 50 - dumper._interval_secs = None - 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 - 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() + all_reduce_mock.assert_not_called() def test_maybe_dump_uses_checkpoint_aligned_global_step(self): dumper = object.__new__(DeltaEmbeddingDumper) @@ -892,36 +847,6 @@ def test_single_rank_minutes_does_not_require_synced_dataloader_exhaustion(self) self.assertFalse(dumper.requires_synced_dataloader_exhaustion) - def test_dump_does_not_ack_tracker_when_submission_is_rejected(self): - dumper = object.__new__(DeltaEmbeddingDumper) - dumper._feature_store_enabled = True - dumper._retain_local_dump = False - 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") - - def _append_one_chunk(table_chunks, **kwargs): - table_chunks.append(_DELTA_DUMP_SCHEMA.empty_table()) - return 1 - - with ( - mock.patch.object(dumper, "_check_feature_store_upload_error"), - 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", side_effect=_append_one_chunk - ), - mock.patch.object(dumper, "_rollback_tracker_read") as rollback, - ): - with self.assertRaisesRegex(RuntimeError, "submission rejected"): - dumper.dump(10) - - dumper._uploader.submit.assert_called_once() - self.assertEqual(dumper._uploader.submit.call_args.args[0], 10) - rollback.assert_called_once_with(7) - def test_multi_gpu_output_path_uses_step_underscore_dir(self): with tempfile.TemporaryDirectory() as tmp_dir: dumper = object.__new__(DeltaEmbeddingDumper) From 8840664b659f533f4cac0e1b23519297a842600a Mon Sep 17 00:00:00 2001 From: gecheng Date: Sun, 26 Jul 2026 15:04:18 +0800 Subject: [PATCH 43/50] [refactor] rely on torchrun teardown for delta dump failure paths Stop catching training failures to close the dumper and stop rendezvousing rank-local startup/shutdown errors across workers via all-reduce. torchrun terminates the whole job when any rank exits with an uncaught exception, and pending in-memory deltas are intentionally discarded on restart from the latest checkpoint, so the failure propagation added nothing. The dumper is now started directly and closed only after training succeeds, to drain background uploads; _raise_if_any_worker_failed is removed entirely. Co-Authored-By: Claude --- tzrec/main.py | 118 ++++++++++++++------------------------------------ 1 file changed, 33 insertions(+), 85 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index c35b01232..eb7130b2a 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -129,21 +129,6 @@ 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], @@ -588,8 +573,7 @@ def run_eval(step: int, epoch: int) -> None: # Flush the trailing partial interval before the final checkpoint. # final_dump skips dump-boundary steps already written by maybe_dump # (multi-rank dumps force synced exhaustion, so every rank participated - # in those dumps); the MAX all-reduce in final_dump keeps one complete - # shard set per step dir as a defensive guard. + # in those dumps and reaches the same final step). delta_embedding_dumper.final_dump(i_step) _log_train( @@ -775,22 +759,13 @@ def train_and_evaluate( plan=plan, ) delta_embedding_dumper = None - delta_embedding_dumper_error: Optional[BaseException] = None if enable_delta_embedding_dump: - 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, + delta_embedding_dumper = DeltaEmbeddingDumper( + model, + train_config.delta_embedding_dump_config, + pipeline_config.model_dir, device, - "delta embedding dumper initialization", + pipeline_config.feature_configs, ) dense_optim_cls, dense_optim_kwargs = optimizer_builder.create_dense_optimizer( @@ -869,60 +844,33 @@ def train_and_evaluate( with open(os.path.join(pipeline_config.model_dir, "version"), "w") as f: f.write(tzrec_version + "\n") - try: - if delta_embedding_dumper is not None: - delta_embedding_dumper_start_error: Optional[BaseException] = None - try: - delta_embedding_dumper.start() - except BaseException as exc: - delta_embedding_dumper_start_error = exc - _raise_if_any_worker_failed( - delta_embedding_dumper_start_error, - device, - "FeatureStore delta uploader startup", - ) - # when slice batch by sample cost, data on all workers may not be balanced - check_all_workers_data_status = data_config.HasField("batch_cost_size") - _train_and_evaluate( - model, - optimizer, - train_dataloader, - eval_dataloader, - [sparse_lr, dense_lr, *part_lrs], - pipeline_config.model_dir, - train_config=train_config, - eval_config=pipeline_config.eval_config, - ckpt_manager=ckpt_manager, - skip_steps=skip_steps, - ckpt_path=ckpt_path, - check_all_workers_data_status=check_all_workers_data_status, - ignore_restore_optimizer=ignore_restore_optimizer, - dataloader_state=dataloader_state, - delta_embedding_dumper=delta_embedding_dumper, - ) - except BaseException: - if delta_embedding_dumper is not None: - # Keep the original training error primary. Pending in-memory - # deltas are abandoned; the restarted run re-dumps from the latest - # checkpoint. Do not wait up to the normal upload-drain timeout - # while unwinding a training failure. - delta_embedding_dumper.close(raise_on_error=False, drain=False) - raise - else: - if delta_embedding_dumper is not None: - delta_embedding_dumper_close_error: Optional[BaseException] = None - try: - delta_embedding_dumper.close() - except BaseException as exc: - delta_embedding_dumper_close_error = exc - # Each rank drains its own background queue at a different pace. - # Propagate a late drain/SDK error uniformly instead of letting a - # single rank fail alone at shutdown. - _raise_if_any_worker_failed( - delta_embedding_dumper_close_error, - device, - "FeatureStore delta uploader shutdown", - ) + if delta_embedding_dumper is not None: + delta_embedding_dumper.start() + # when slice batch by sample cost, data on all workers may not be balanced + check_all_workers_data_status = data_config.HasField("batch_cost_size") + _train_and_evaluate( + 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, + ) + # Drain background uploads only after training succeeds. A training failure + # terminates the whole job (torchrun tears down every rank) and pending + # in-memory deltas are intentionally abandoned: the restarted run re-dumps + # from the latest checkpoint, so there is nothing to roll back or undo. + if delta_embedding_dumper is not None: + delta_embedding_dumper.close() if is_local_rank_zero: logger.info("Train and Evaluate Finished.") From f0ec6b8d6f09121a82369a8bf073d9c8809c4ade Mon Sep 17 00:00:00 2001 From: gecheng Date: Mon, 27 Jul 2026 09:55:07 +0800 Subject: [PATCH 44/50] [refactor] move rank-zero view rendezvous into uploader start The barrier sequencing rank-zero view creation ahead of non-primary opens lived in DeltaEmbeddingDumper.start, leaking the uploader's manage_remote_view startup protocol to the caller. Move the create-barrier-open sequence into FeatureStoreDeltaUploader.start so the dumper only delegates and arms the timed cadence; the failure rendezvous (rank-zero create error still joins the barrier before re-raising) is preserved verbatim, with world_size plumbed through the uploader ctor. Co-Authored-By: Claude --- tzrec/utils/delta_embedding_dump.py | 31 ++++----------------- tzrec/utils/feature_store_delta_uploader.py | 28 ++++++++++++++++++- 2 files changed, 33 insertions(+), 26 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 4d3d3c898..20f9414d6 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -432,6 +432,7 @@ def __init__( config.feature_store_config, embedding_dimensions=embedding_dimensions, rank=self._rank, + world_size=self._world_size, manage_remote_view=self._rank == 0, ) @@ -472,33 +473,13 @@ def requires_synced_dataloader_exhaustion(self) -> bool: def start(self) -> None: """Start timed cadence and per-rank FeatureStore publication. - The rank-zero uploader creates and validates the remote view first; - the other ranks start their uploaders only after the barrier so they - can open the already-published view without control-plane races. A - rank-zero startup failure still joins the barrier before raising so - every rank keeps issuing the same collective sequence. + The rank-zero view rendezvous (create the DynamicEmbedding view, + barrier, then non-primary ranks open it) lives in + :meth:`FeatureStoreDeltaUploader.start`; this method only delegates and + arms the timed cadence. """ if self._uploader is not None: - start_error: Optional[BaseException] = None - if self._rank == 0: - try: - self._uploader.start() - except BaseException as exc: - start_error = exc - 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" - ) - torch.distributed.barrier() - if start_error is not None: - raise start_error.with_traceback(start_error.__traceback__) - if self._rank != 0: - self._uploader.start() + self._uploader.start() if self._interval_secs is not None: self._next_dump_time = time.monotonic() + self._interval_secs diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index 342185f2a..76d86bb20 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -181,6 +181,7 @@ def __init__( config: FeatureStoreConfig, embedding_dimensions: Mapping[str, int], rank: int = 0, + world_size: int = 1, manage_remote_view: bool = True, client_factory: Optional[Callable[..., Any]] = None, credentials_client: Optional[Any] = None, @@ -189,6 +190,7 @@ def __init__( """Initialize the uploader with validated settings and in-memory state.""" self._settings = FeatureStoreUploadSettings.from_proto(config) self._rank = int(rank) + self._world_size = int(world_size) self._manage_remote_view = bool(manage_remote_view) self._embedding_dimensions = { str(name): int(dimension) @@ -212,13 +214,37 @@ def __init__( self._last_publish_ts: int = 0 def start(self) -> None: - """Start the background upload worker after validating the remote view.""" + """Start the worker after the rank-zero view rendezvous. + + Rank zero creates and validates the DynamicEmbedding view first; every + rank then barriers so non-primary ranks open the already-published view + without control-plane races. A rank-zero startup failure still joins the + barrier before re-raising so every rank issues the same collective order. + """ with self._condition: if self._started: return if self._closed: raise RuntimeError("FeatureStoreDeltaUploader is already closed") self._raise_if_failed_locked() + start_error: Optional[BaseException] = None + if self._manage_remote_view: + try: + self._get_view() + except BaseException as exc: + self._reset_view(suppress_errors=True) + start_error = exc + if self._world_size > 1: + import torch.distributed as dist + + if not (dist.is_available() and dist.is_initialized()): + raise RuntimeError( + "distributed FeatureStore delta dump requires an " + "initialized process group" + ) + dist.barrier() + if start_error is not None: + raise start_error.with_traceback(start_error.__traceback__) try: self._get_view() except BaseException: From defc8045db7044ec7368ea339379c9a04ff1a831 Mon Sep 17 00:00:00 2001 From: gecheng Date: Mon, 27 Jul 2026 10:02:50 +0800 Subject: [PATCH 45/50] [refactor] encapsulate feature store client construction in _create_client FeatureStoreClient construction and credential resolution were inlined in _get_view and reached from __init__ via client_factory/credentials_client ctor params, leaking an SDK-specific concern into the constructor. Move both into a private _create_client() seam and drop the two ctor params; tests now inject a fake by patching that one method. The test_mode kwarg hygiene check moves to a focused unit test asserting _create_client forwards exactly the fixed credential allowlist. Co-Authored-By: Claude --- tzrec/utils/feature_store_delta_uploader.py | 53 +++++++++---------- .../feature_store_delta_uploader_test.py | 42 ++++++++++++++- 2 files changed, 65 insertions(+), 30 deletions(-) diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index 76d86bb20..364b668ed 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -183,8 +183,6 @@ def __init__( rank: int = 0, world_size: int = 1, manage_remote_view: bool = True, - client_factory: Optional[Callable[..., Any]] = None, - credentials_client: Optional[Any] = None, clock_ms: Optional[Callable[[], int]] = None, ) -> None: """Initialize the uploader with validated settings and in-memory state.""" @@ -197,10 +195,7 @@ def __init__( for name, dimension in embedding_dimensions.items() } - self._client_factory = client_factory - self._credentials_client = ( - credentials_client or self._create_credentials_client() - ) + self._credentials_client = self._create_credentials_client() self._clock_ms = clock_ms or (lambda: time.time_ns() // 1_000_000) self._view = None self._condition = threading.Condition() @@ -585,31 +580,33 @@ def _create_credentials_client() -> Any: ) from exc return CredClient() + def _create_client(self) -> Any: + """Construct a FeatureStoreClient with refreshed credentials. + + Single seam for credential resolution and client construction; tests + patch this method to inject a fake client. + """ + 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 + credential = self._credentials_client.get_credential() + return FeatureStoreClient( + access_key_id=credential.access_key_id, + access_key_secret=credential.access_key_secret, + region=self._settings.region or None, + endpoint=self._settings.endpoint or None, + security_token=credential.security_token or None, + featuredb_username=os.environ.get("FEATUREDB_USERNAME") or None, + featuredb_password=os.environ.get("FEATUREDB_PASSWORD") or None, + ) + 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 - - credential = self._credentials_client.get_credential() - kwargs = { - "access_key_id": credential.access_key_id, - "access_key_secret": credential.access_key_secret, - "region": self._settings.region or None, - "endpoint": self._settings.endpoint or None, - "security_token": credential.security_token or None, - "featuredb_username": os.environ.get("FEATUREDB_USERNAME") or None, - "featuredb_password": os.environ.get("FEATUREDB_PASSWORD") or None, - } - client = client_factory(**kwargs) + client = self._create_client() project = client.get_project(self._settings.project_name) if project is None: raise RuntimeError("configured FeatureStore project was not found") diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index ddedf7781..ff91cedff 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -10,7 +10,9 @@ # limitations under the License. import os +import sys import threading +import types import unittest from unittest import mock @@ -211,8 +213,17 @@ def setUp(self): self.addCleanup(self._cred_patch.stop) def _uploader(self, config=None, **kwargs): + client_factory = kwargs.pop("client_factory", None) kwargs.setdefault("embedding_dimensions", {"user_emb": 2}) - return FeatureStoreDeltaUploader(config or _feature_store_config(), **kwargs) + uploader = FeatureStoreDeltaUploader( + config or _feature_store_config(), **kwargs + ) + if client_factory is not None: + # The production ctor no longer accepts a client factory; install the + # test's fake on the _create_client instance seam so each uploader + # keeps its own fake project/view wiring and construction counts. + uploader._create_client = lambda *a, **k: client_factory() + return uploader def test_proto_groups_required_fields_before_optional_fields(self): required_fields = [ @@ -317,9 +328,36 @@ def test_start_reuses_existing_dynamic_embedding_feature_view(self): self.assertEqual(factory.project.dynamic_get_calls, ["shared_embeddings"]) self.assertEqual(factory.project.create_calls, []) - self.assertNotIn("test_mode", factory.calls[0]) self.assertEqual(view.closed, [True]) + def test_create_client_forwards_only_credential_kwargs(self): + """_create_client forwards the fixed credential allowlist, no extras.""" + recorded = {} + + class _RecordingClient: + def __init__(self, **kwargs): + recorded.update(kwargs) + + fake_module = types.ModuleType("feature_store_py") + fake_module.FeatureStoreClient = _RecordingClient + with mock.patch.dict(sys.modules, {"feature_store_py": fake_module}): + uploader = self._uploader() + client = uploader._create_client() + self.assertIsInstance(client, _RecordingClient) + self.assertEqual( + set(recorded), + { + "access_key_id", + "access_key_secret", + "region", + "endpoint", + "security_token", + "featuredb_username", + "featuredb_password", + }, + ) + self.assertNotIn("test_mode", recorded) + def test_start_creates_missing_dynamic_embedding_feature_view(self): created_view = _FakeView() factory = _FakeClientFactory(None, created_view=created_view) From db266ac2af8181322a38eca3cf03d743d1f31e2b Mon Sep 17 00:00:00 2001 From: gecheng Date: Tue, 28 Jul 2026 15:39:35 +0800 Subject: [PATCH 46/50] [feat] auto-create default entity for dynamic embedding feature view Drop feature_entity_name from FeatureStoreConfig; the rank-zero uploader now provisions a default entity (default_dynemb_entity) on demand when it creates the DynamicEmbedding feature view. Co-Authored-By: Claude --- tzrec/protos/train.proto | 10 +- tzrec/utils/export_util_test.py | 1 - tzrec/utils/feature_store_delta_uploader.py | 48 ++++++-- .../feature_store_delta_uploader_test.py | 109 +++++++++++++++--- 4 files changed, 134 insertions(+), 34 deletions(-) diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index 16904c093..559d9a7ac 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -37,12 +37,10 @@ message FeatureStoreConfig { // to ALIBABA_CLOUD_REGION at runtime. required string region = 1; // Existing FeatureStore project name and target DynamicEmbedding FeatureView - // name. The view is validated at startup and created when it does not exist. + // name. The view is validated at startup and created when it does not exist, + // together with a default FeatureStore entity that is auto-created on demand. required string project_name = 2; required string feature_view_name = 3; - // Existing FeatureStore entity associated with an automatically created - // DynamicEmbedding FeatureView. - required string feature_entity_name = 4; // FeatureDB version for this incremental training run. required string version = 5; @@ -67,8 +65,8 @@ message FeatureStoreConfig { optional uint32 feature_view_ttl_secs = 13 [default = 1296000]; optional uint32 feature_view_shard_count = 14 [default = 20]; optional uint32 feature_view_replication_count = 15 [default = 1]; - reserved 16; - reserved "allow_custom_endpoint"; + reserved 4, 16; + reserved "feature_entity_name", "allow_custom_endpoint"; // Also write the local per-rank delta parquet files while uploading to // FeatureStore (e.g. for the offline readback checker); by default the // delta is handed to the uploader in memory and never touches disk. diff --git a/tzrec/utils/export_util_test.py b/tzrec/utils/export_util_test.py index ed7ef199f..29f643366 100644 --- a/tzrec/utils/export_util_test.py +++ b/tzrec/utils/export_util_test.py @@ -366,7 +366,6 @@ def forward(self, data, device=None): # type: ignore[no-untyped-def] 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" model_acc = {"SPARSE_INT64": "1", "cand_seq_pk": "cand_seq"} diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index c785af1db..924fc6c2b 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -47,6 +47,12 @@ FEATURE_STORE_WRITE_MODE = "MERGE" FEATURE_STORE_SDK_BATCH_SIZE = 1000 +# Default FeatureStore entity used when a DynamicEmbedding FeatureView has to be +# created. The entity is provisioned on demand by the rank-zero uploader, so +# users never configure it; the join_id only needs to be a stable non-empty key. +FEATURE_STORE_DEFAULT_ENTITY_NAME = "default_dynemb_entity" +FEATURE_STORE_DEFAULT_ENTITY_JOIN_ID = "default_dynemb_join_id" + _FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES = 100 @@ -65,7 +71,6 @@ class FeatureStoreUploadSettings: region: str endpoint: str project_name: str - feature_entity_name: str feature_view_name: str feature_view_ttl_secs: int feature_view_shard_count: int @@ -96,15 +101,10 @@ def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": "(it may come from ALIBABA_CLOUD_REGION)" ) 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": @@ -144,7 +144,6 @@ def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": region=region, endpoint=endpoint, 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, @@ -683,9 +682,10 @@ def _get_or_create_view(self, project: Any) -> Any: if view is None: create_error: Optional[Exception] = None try: + entity_name = self._get_or_create_entity(project) view = project.create_dynamic_embedding_feature_view( name=self._settings.feature_view_name, - entity=self._settings.feature_entity_name, + entity=entity_name, pk_field_name=FEATURE_STORE_PK_FIELD, sk_field_name=FEATURE_STORE_SK_FIELD, embedding_field_name=FEATURE_STORE_VALUE_FIELD, @@ -701,8 +701,7 @@ def _get_or_create_view(self, project: Any) -> Any: 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" + "failed to create configured DynamicEmbedding FeatureView" ) if create_error is not None: raise error from create_error @@ -712,11 +711,38 @@ def _get_or_create_view(self, project: Any) -> Any: logger.info( "Created DynamicEmbedding FeatureView: project=%s entity=%s view=%s", self._settings.project_name, - self._settings.feature_entity_name, + FEATURE_STORE_DEFAULT_ENTITY_NAME, self._settings.feature_view_name, ) return view + def _get_or_create_entity(self, project: Any) -> str: + """Return the default DynamicEmbedding entity name, creating it if absent. + + Only the rank-zero uploader runs this, immediately before it creates the + view, so there is no cross-rank race; a concurrent external creator is + handled by re-getting the entity if the create call fails. + """ + if project.get_entity(FEATURE_STORE_DEFAULT_ENTITY_NAME) is not None: + return FEATURE_STORE_DEFAULT_ENTITY_NAME + try: + project.create_entity( + FEATURE_STORE_DEFAULT_ENTITY_NAME, + FEATURE_STORE_DEFAULT_ENTITY_JOIN_ID, + ) + except Exception as exc: + if project.get_entity(FEATURE_STORE_DEFAULT_ENTITY_NAME) is None: + raise RuntimeError( + "failed to create default DynamicEmbedding entity " + f"{FEATURE_STORE_DEFAULT_ENTITY_NAME!r}" + ) from exc + logger.info( + "Created FeatureStore entity: project=%s entity=%s", + self._settings.project_name, + FEATURE_STORE_DEFAULT_ENTITY_NAME, + ) + return FEATURE_STORE_DEFAULT_ENTITY_NAME + 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 diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index d6de5d923..721ea8d7f 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -22,6 +22,8 @@ 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 ( + FEATURE_STORE_DEFAULT_ENTITY_JOIN_ID, + FEATURE_STORE_DEFAULT_ENTITY_NAME, FeatureStoreDeltaUploader, FeatureStoreUploadError, FeatureStoreUploadSettings, @@ -32,7 +34,6 @@ def _feature_store_config(**overrides) -> FeatureStoreConfig: config = FeatureStoreConfig( region="cn-test", project_name="project_a", - feature_entity_name="embedding_entity", feature_view_name="shared_embeddings", version="model_a@export_1", upload_batch_size=2, @@ -131,6 +132,11 @@ def close(self, wait=True): self.close_finished.set() +class _FakeEntity: + def __init__(self, name): + self.feature_entity_name = name + + class _FakeProject: def __init__( self, @@ -139,13 +145,19 @@ def __init__( created_view=None, create_error=None, view_after_create_error=None, + entity="existing_entity", + entity_create_error=None, ): self._view = view self._created_view = created_view self._create_error = create_error self._view_after_create_error = view_after_create_error + self._entity = entity + self._entity_create_error = entity_create_error self.dynamic_get_calls = [] self.create_calls = [] + self.entity_get_calls = [] + self.entity_create_calls = [] def get_dynamic_embedding_feature_view(self, name): self.dynamic_get_calls.append(name) @@ -159,6 +171,17 @@ def create_dynamic_embedding_feature_view(self, **kwargs): self._view = self._created_view or _FakeView() return self._view + def get_entity(self, name): + self.entity_get_calls.append(name) + return self._entity + + def create_entity(self, name, join_id, parent_feature_entity_name=None): + self.entity_create_calls.append((name, join_id)) + if self._entity_create_error is not None: + raise self._entity_create_error + self._entity = _FakeEntity(name) + return self._entity + class _FakeCredential: access_key_id = "fake-ak" @@ -231,7 +254,6 @@ def test_proto_groups_required_fields_before_optional_fields(self): "region", "project_name", "feature_view_name", - "feature_entity_name", "version", ] optional_fields = [ @@ -266,7 +288,7 @@ def test_proto_groups_required_fields_before_optional_fields(self): ) self.assertEqual( [field.number for field in fields], - list(range(1, 16)) + [17], + [1, 2, 3] + list(range(5, 16)) + [17], ) for field_name in required_fields: with self.subTest(field_name=field_name): @@ -277,15 +299,6 @@ def test_proto_groups_required_fields_before_optional_fields(self): 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") @@ -373,7 +386,7 @@ def test_start_creates_missing_dynamic_embedding_feature_view(self): [ { "name": "shared_embeddings", - "entity": "embedding_entity", + "entity": FEATURE_STORE_DEFAULT_ENTITY_NAME, "pk_field_name": "embedding_name", "sk_field_name": "key_id", "embedding_field_name": "embedding", @@ -385,6 +398,7 @@ def test_start_creates_missing_dynamic_embedding_feature_view(self): } ], ) + self.assertEqual(factory.project.entity_create_calls, []) self.assertEqual(created_view.closed, [True]) def test_start_recovers_from_concurrent_feature_view_creation(self): @@ -429,14 +443,77 @@ def test_start_creates_missing_view_without_version_precheck(self): 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): - factory = _FakeClientFactory(None, create_error=ValueError("Entity not found")) + def test_start_raises_when_view_creation_fails_and_view_never_appears(self): + factory = _FakeClientFactory(None, create_error=ValueError("boom")) uploader = self._uploader(client_factory=factory) - with self.assertRaisesRegex(RuntimeError, "feature_entity_name"): + with self.assertRaisesRegex( + RuntimeError, "failed to create configured DynamicEmbedding FeatureView" + ): uploader.start() self.assertEqual(len(factory.project.create_calls), 1) + self.assertEqual(factory.project.entity_create_calls, []) + + def test_start_creates_default_entity_when_it_does_not_exist(self): + created_view = _FakeView() + factory = _FakeClientFactory(None, created_view=created_view, entity=None) + uploader = self._uploader(client_factory=factory) + + uploader.start() + uploader.close() + + self.assertEqual( + factory.project.entity_get_calls, [FEATURE_STORE_DEFAULT_ENTITY_NAME] + ) + self.assertEqual( + factory.project.entity_create_calls, + [ + ( + FEATURE_STORE_DEFAULT_ENTITY_NAME, + FEATURE_STORE_DEFAULT_ENTITY_JOIN_ID, + ) + ], + ) + self.assertEqual( + factory.project.create_calls[0]["entity"], + FEATURE_STORE_DEFAULT_ENTITY_NAME, + ) + self.assertEqual(created_view.closed, [True]) + + def test_start_recovers_from_concurrent_default_entity_creation(self): + created_view = _FakeView() + factory = _FakeClientFactory( + None, + created_view=created_view, + entity=None, + entity_create_error=RuntimeError("entity already exists"), + ) + # A concurrent creator wins the race: the first get_entity sees no entity, + # create_entity then fails, and the entity is visible on the retry. + entity_results = [None, _FakeEntity(FEATURE_STORE_DEFAULT_ENTITY_NAME)] + + def get_entity(name): + factory.project.entity_get_calls.append(name) + return entity_results.pop(0) + + factory.project.get_entity = get_entity + uploader = self._uploader(client_factory=factory) + + uploader.start() + uploader.close() + + self.assertEqual( + factory.project.entity_create_calls, + [ + ( + FEATURE_STORE_DEFAULT_ENTITY_NAME, + FEATURE_STORE_DEFAULT_ENTITY_JOIN_ID, + ) + ], + ) + self.assertEqual(len(factory.project.create_calls), 1) + self.assertEqual(created_view.closed, [True]) def test_non_primary_uploader_opens_view_without_create_or_metadata_checks(self): view = _FakeView() From b49157e0a7e61ec4619138b926ace91193b4318d Mon Sep 17 00:00:00 2001 From: gecheng Date: Tue, 28 Jul 2026 20:33:36 +0800 Subject: [PATCH 47/50] [feat] default delta dump upload to pyarrow Arrow IPC Add FeatureStoreConfig.upload_format (default ARROW) to stream delta batches through the SDK's write_features_arrow() columnar IPC path, avoiding the JSON path's per-row dict construction and embedding deep-copy. The legacy JSON write_features() path stays behind upload_format="JSON"; both reuse a shared _validate_delta_batch and the existing MERGE/ts/window/flush/retry semantics. Two pre-existing defects surfaced and are fixed in passing: streaming now materializes the actual to_batches() list (ceil(total/batch_size) undercounted multi-FQN multi-chunk tables, letting a stuck clock reuse a prior step's ts so Next-Ts readers miss updates), and the NaN/Inf check is bounded to each batch's value range instead of re-scanning the whole chunk buffer per batch. Co-Authored-By: Claude --- tzrec/protos/train.proto | 5 + tzrec/utils/feature_store_delta_uploader.py | 188 +++++++++++--- .../feature_store_delta_uploader_test.py | 232 +++++++++++++++++- 3 files changed, 390 insertions(+), 35 deletions(-) diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index 559d9a7ac..d254199eb 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -71,6 +71,11 @@ message FeatureStoreConfig { // FeatureStore (e.g. for the offline readback checker); by default the // delta is handed to the uploader in memory and never touches disk. optional bool retain_local_dump = 17 [default = false]; + // Wire format for delta upload. "ARROW" (default) streams a columnar Arrow + // IPC batch through write_features_arrow(), avoiding the JSON path's per-row + // dict construction and embedding deep-copy; "JSON" keeps the legacy + // write_features() per-row payload. Both paths use MERGE write_mode. + optional string upload_format = 18 [default = "ARROW"]; } message DeltaEmbeddingDumpConfig { diff --git a/tzrec/utils/feature_store_delta_uploader.py b/tzrec/utils/feature_store_delta_uploader.py index 924fc6c2b..7ab0019f8 100644 --- a/tzrec/utils/feature_store_delta_uploader.py +++ b/tzrec/utils/feature_store_delta_uploader.py @@ -46,6 +46,8 @@ FEATURE_STORE_VALUE_FIELD = "embedding" FEATURE_STORE_WRITE_MODE = "MERGE" FEATURE_STORE_SDK_BATCH_SIZE = 1000 +FEATURE_STORE_UPLOAD_FORMAT_DEFAULT = "ARROW" +FEATURE_STORE_UPLOAD_FORMATS = ("ARROW", "JSON") # Default FeatureStore entity used when a DynamicEmbedding FeatureView has to be # created. The entity is provisioned on demand by the rank-zero uploader, so @@ -56,6 +58,22 @@ _FEATURE_STORE_PROGRESS_LOG_INTERVAL_BATCHES = 100 +@dataclass(frozen=True) +class _DeltaBatch: + """Decoded view of one delta RecordBatch shared by both upload paths. + + PK values are pre-remapped (INPUT_TILE=3 user-side keys -> non-user twin + keys); the raw table_fqn is retained only for contract/dimension checks. + """ + + num_rows: int + remapped_fqns: List[str] + key_ids: np.ndarray + embedding_column: pa.Array + flat_embeddings: np.ndarray + offsets: np.ndarray + + class FeatureStoreUploadError(RuntimeError): """Safe, credential-free error propagated from the uploader thread.""" @@ -82,6 +100,7 @@ class FeatureStoreUploadSettings: shutdown_timeout_secs: int max_pending_steps: int poll_interval_secs: int + upload_format: str @classmethod def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": @@ -140,6 +159,16 @@ def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": "exactly one FeatureStore SDK HTTP batch" ) + upload_format = ( + (config.upload_format or FEATURE_STORE_UPLOAD_FORMAT_DEFAULT) + .strip() + .upper() + ) + if upload_format not in FEATURE_STORE_UPLOAD_FORMATS: + raise ValueError( + "feature_store_config.upload_format must be one of " + f"{FEATURE_STORE_UPLOAD_FORMATS}, got {upload_format!r}" + ) return cls( region=region, endpoint=endpoint, @@ -155,6 +184,7 @@ def from_proto(cls, config: FeatureStoreConfig) -> "FeatureStoreUploadSettings": shutdown_timeout_secs=positive_values["shutdown_timeout_secs"], max_pending_steps=positive_values["max_pending_steps"], poll_interval_secs=positive_values["poll_interval_secs"], + upload_format=upload_format, ) @@ -407,7 +437,12 @@ def _stream_upload(self, global_step: int, table: pa.Table) -> None: view = self._get_view() max_in_flight = int(getattr(view, "_max_workers", 1)) - total_batches = self._count_total_batches(table) + # Materialize the actual batch list so the ts range covers every batch + # the SDK will see. to_batches() splits each physical chunk independently, + # so ceil(total_rows / batch_size) undercounts multi-chunk (multi-FQN) + # tables and lets consecutive steps' timestamps collide. + batches = list(table.to_batches(max_chunksize=self._settings.upload_batch_size)) + total_batches = len(batches) or 1 ts_range = self._allocate_timestamp_range(total_batches) range_start = ts_range[0] @@ -430,22 +465,16 @@ def _stream_upload(self, global_step: int, table: pa.Table) -> None: ) try: - for batch in table.to_batches( - max_chunksize=self._settings.upload_batch_size - ): + for batch in batches: self._raise_if_aborting() - payload = self._validate_and_build_payload(batch) - if not payload: - continue - view.write_features( - data=payload, - version=self._settings.version, - write_mode=FEATURE_STORE_WRITE_MODE, - ts=range_start + completed_batches, + num_rows = self._submit_one_batch( + view, batch, range_start + completed_batches ) + if num_rows == 0: + continue completed_batches += 1 window_batches += 1 - window_records += len(payload) + window_records += num_rows if window_batches < max_in_flight and completed_batches < total_batches: continue @@ -498,23 +527,64 @@ def _stream_upload(self, global_step: int, table: pa.Table) -> None: time.monotonic() - started_at, ) - def _count_total_batches(self, table: pa.Table) -> int: - """Count the number of upload batches in one step's delta table.""" - total_rows = int(table.num_rows) - if total_rows == 0: - return 1 - return ( - total_rows + self._settings.upload_batch_size - 1 - ) // self._settings.upload_batch_size + def _submit_one_batch(self, view: Any, batch: pa.RecordBatch, ts: int) -> int: + """Validate, build, and submit one batch; return submitted rows (0 skip). - def _validate_and_build_payload( - self, - batch: pa.RecordBatch, - ) -> List[Dict[str, Any]]: - """Validate one delta batch and build the SDK payload.""" + Dispatches on upload_format: ARROW streams a columnar RecordBatch through + write_features_arrow(); JSON falls back to the per-row write_features() + payload. Both reuse _validate_delta_batch so invariants are identical. + """ + if self._settings.upload_format == FEATURE_STORE_UPLOAD_FORMAT_DEFAULT: + wire_batch, num_rows = self._validate_and_build_arrow_batch(batch) + if num_rows == 0: + return 0 + view.write_features_arrow( + batch=wire_batch, + version=self._settings.version, + write_mode=FEATURE_STORE_WRITE_MODE, + ts=ts, + ) + else: + payload = self._validate_and_build_payload(batch) + num_rows = len(payload) + if num_rows == 0: + return 0 + view.write_features( + data=payload, + version=self._settings.version, + write_mode=FEATURE_STORE_WRITE_MODE, + ts=ts, + ) + return num_rows + + def _validate_delta_batch(self, batch: pa.RecordBatch) -> _DeltaBatch: + """Validate one delta batch and decode it for payload construction. + + Enforces the table_fqn / key_id / embedding / dimension invariants shared + by the JSON and Arrow upload paths. PK values are pre-remapped here + (INPUT_TILE=3 user-side keys -> non-user twin keys); dimension checks key + against the raw table_fqn, which is how the model contract is indexed. + + Args: + batch: One delta RecordBatch from the in-memory delta table. + + Returns: + Decoded _DeltaBatch; num_rows==0 when the batch carries no rows. + + Raises: + ValueError: On empty table_fqn, contract mismatch, reserved + key_id=-1, NaN/Inf embeddings, or dimension mismatch. + """ num_rows = batch.num_rows if num_rows == 0: - return [] + return _DeltaBatch( + num_rows=0, + remapped_fqns=[], + key_ids=np.empty(0, dtype=np.int64), + embedding_column=batch.column("embedding"), + flat_embeddings=np.empty(0, dtype=np.float32), + offsets=np.array([0], dtype=np.int32), + ) table_fqns = batch.column("table_fqn").to_pylist() for table_fqn in set(table_fqns): if not table_fqn: @@ -533,10 +603,15 @@ def _validate_and_build_payload( ) embedding_column = cast(pa.ListArray, batch.column("embedding")) + offsets = embedding_column.offsets.to_numpy() flat_embeddings = embedding_column.values.to_numpy(zero_copy_only=False) - if not bool(np.isfinite(flat_embeddings).all()): + # A sliced ListArray shares the whole chunk's child buffer, so scanning + # flat_embeddings in full would re-scan the chunk on every batch. Bound + # the NaN/Inf check to this batch's value range [offsets[0], offsets[-1]). + value_start = int(offsets[0]) + value_end = int(offsets[-1]) + if not bool(np.isfinite(flat_embeddings[value_start:value_end]).all()): raise ValueError("delta embedding contains NaN or Inf") - offsets = embedding_column.offsets.to_numpy() lengths = np.diff(offsets) expected_dims = np.array( [self._embedding_dimensions[fqn] for fqn in table_fqns], @@ -551,17 +626,62 @@ def _validate_and_build_payload( f"actual={int(lengths[row])}" ) + remap_cache = {fqn: remap_input_tile_user_key(fqn) for fqn in set(table_fqns)} + remapped_fqns = [remap_cache[fqn] for fqn in table_fqns] + return _DeltaBatch( + num_rows=num_rows, + remapped_fqns=remapped_fqns, + key_ids=key_ids, + embedding_column=embedding_column, + flat_embeddings=flat_embeddings, + offsets=offsets, + ) + + def _validate_and_build_payload( + self, + batch: pa.RecordBatch, + ) -> List[Dict[str, Any]]: + """Validate one delta batch and build the JSON SDK payload.""" + delta = self._validate_delta_batch(batch) + if delta.num_rows == 0: + return [] return [ { - FEATURE_STORE_PK_FIELD: remap_input_tile_user_key(table_fqns[i]), - FEATURE_STORE_SK_FIELD: int(key_ids[i]), - FEATURE_STORE_VALUE_FIELD: flat_embeddings[ - int(offsets[i]) : int(offsets[i + 1]) + FEATURE_STORE_PK_FIELD: delta.remapped_fqns[i], + FEATURE_STORE_SK_FIELD: int(delta.key_ids[i]), + FEATURE_STORE_VALUE_FIELD: delta.flat_embeddings[ + int(delta.offsets[i]) : int(delta.offsets[i + 1]) ].copy(), } - for i in range(num_rows) + for i in range(delta.num_rows) ] + def _validate_and_build_arrow_batch( + self, + batch: pa.RecordBatch, + ) -> Tuple[Optional[pa.RecordBatch], int]: + """Validate one delta batch and build the Arrow IPC wire batch. + + Returns a RecordBatch with the configured PK/SK/embedding field names so + the SDK remaps them to its wire (pk/sk/embedding) columns. The embedding + column is reused zero-copy; only the string PK column and the int64 SK + column are rebuilt, avoiding the JSON path's per-row embedding deep-copy. + """ + delta = self._validate_delta_batch(batch) + if delta.num_rows == 0: + return None, 0 + pk_column = pa.array(delta.remapped_fqns, type=pa.string()) + sk_column = pa.array(delta.key_ids, type=pa.int64()) + wire_batch = pa.RecordBatch.from_arrays( + [pk_column, sk_column, delta.embedding_column], + names=[ + FEATURE_STORE_PK_FIELD, + FEATURE_STORE_SK_FIELD, + FEATURE_STORE_VALUE_FIELD, + ], + ) + return wire_batch, delta.num_rows + @staticmethod def _create_credentials_client() -> Any: """Create the Alibaba Cloud credential provider (default chain).""" diff --git a/tzrec/utils/feature_store_delta_uploader_test.py b/tzrec/utils/feature_store_delta_uploader_test.py index 721ea8d7f..d19492eba 100644 --- a/tzrec/utils/feature_store_delta_uploader_test.py +++ b/tzrec/utils/feature_store_delta_uploader_test.py @@ -16,6 +16,7 @@ import unittest from unittest import mock +import numpy as np import pyarrow as pa from google.protobuf.descriptor import FieldDescriptor @@ -81,6 +82,7 @@ class _FakeView: def __init__(self, summaries=None, close_error=None, max_workers=4): self.calls = [] + self.arrow_calls = [] self.closed = [] self.flush_calls = [] self._summaries = list(summaries or []) @@ -93,6 +95,31 @@ def write_features(self, **kwargs): self.calls.append(kwargs) self._pending_sizes.append(len(kwargs["data"])) + def write_features_arrow(self, *, batch, version, write_mode, ts): + # Decode the Arrow wire batch into the same {data, version, write_mode, + # ts} call shape as write_features, so the existing JSON-path assertions + # also exercise the default Arrow path unchanged. The raw batch is kept + # on arrow_calls for column-type / column-name assertions. + self.arrow_calls.append( + {"batch": batch, "version": version, "write_mode": write_mode, "ts": ts} + ) + data = [ + { + self.pk_field: pk, + self.sk_field: int(sk), + self.embedding_field: np.asarray(emb, dtype=np.float32), + } + for pk, sk, emb in zip( + batch.column(self.pk_field).to_pylist(), + batch.column(self.sk_field).to_pylist(), + batch.column(self.embedding_field).to_pylist(), + ) + ] + self.calls.append( + {"data": data, "version": version, "write_mode": write_mode, "ts": ts} + ) + self._pending_sizes.append(len(data)) + def write_flush(self): pending_sizes = self._pending_sizes self._pending_sizes = [] @@ -268,6 +295,7 @@ def test_proto_groups_required_fields_before_optional_fields(self): "feature_view_shard_count", "feature_view_replication_count", "retain_local_dump", + "upload_format", ] fields = list(FeatureStoreConfig.DESCRIPTOR.fields) @@ -288,7 +316,7 @@ def test_proto_groups_required_fields_before_optional_fields(self): ) self.assertEqual( [field.number for field in fields], - [1, 2, 3] + list(range(5, 16)) + [17], + [1, 2, 3] + list(range(5, 16)) + [17, 18], ) for field_name in required_fields: with self.subTest(field_name=field_name): @@ -761,6 +789,19 @@ def test_dimension_and_finite_value_validation(self): uploader.close() self.assertEqual(view.calls, []) + # Inf is also rejected (np.isfinite covers both; a regression to + # np.isnan would let Inf embeddings through undetected). + view = _FakeView() + uploader = self._uploader( + _feature_store_config(max_retries=1), + client_factory=_FakeClientFactory(view), + ) + uploader.start() + uploader.submit(10, _delta_table([_row(10, 0, 1, [float("inf"), 2.0])])) + with self.assertRaises(FeatureStoreUploadError): + uploader.close() + self.assertEqual(view.calls, []) + def test_in_memory_timestamp_monotonicity_across_steps(self): view = _FakeView() uploader = self._uploader( @@ -842,6 +883,195 @@ def test_close_error_surfaces_via_check_error(self): with self.assertRaises(FeatureStoreUploadError): uploader.check_error() + def test_default_upload_format_is_arrow(self): + # An unset upload_format inherits the proto default ARROW. + settings = FeatureStoreUploadSettings.from_proto(_feature_store_config()) + self.assertEqual(settings.upload_format, "ARROW") + + def test_rejects_unknown_upload_format(self): + config = _feature_store_config() + config.upload_format = "protobuf" + with self.assertRaisesRegex(ValueError, "upload_format must be one of"): + FeatureStoreUploadSettings.from_proto(config) + + def test_json_path_routes_through_write_features(self): + # Explicit JSON keeps the legacy per-row write_features payload path. + view = _FakeView() + factory = _FakeClientFactory(view) + uploader = self._uploader( + _feature_store_config(upload_format="JSON"), + client_factory=factory, + clock_ms=lambda: 123456, + ) + uploader.start() + uploader.submit( + 10, + _delta_table( + [ + _row(10, 0, 1, [1.0, 2.0]), + _row(10, 0, 2, [3.0, 4.0]), + _row(10, 0, 3, [0.0, 0.0]), + ] + ), + ) + uploader.close() + + self.assertEqual(view.arrow_calls, []) + 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"}) + 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[0]["data"][0]["key_id"], 1) + self.assertEqual( + view.calls[0]["data"][0]["embedding_name"], + "model.ebc.embedding_bags.user_emb", + ) + self.assertTrue( + np.array_equal( + view.calls[0]["data"][0]["embedding"], + np.array([1.0, 2.0], dtype=np.float32), + ) + ) + # The sliced second batch exercises the offset-indexed embedding slice. + self.assertTrue( + np.array_equal( + view.calls[1]["data"][0]["embedding"], + np.array([0.0, 0.0], dtype=np.float32), + ) + ) + self.assertEqual(view.closed, [True]) + + def test_arrow_path_builds_wire_batch_columns(self): + # Default ARROW path builds a wire RecordBatch whose configured field + # names the SDK remaps to its pk/sk/embedding wire columns. + view = _FakeView() + factory = _FakeClientFactory(view) + uploader = self._uploader( + _feature_store_config(upload_batch_size=2), + client_factory=factory, + clock_ms=lambda: 100, + ) + uploader.start() + uploader.submit( + 10, + _delta_table( + [ + _row(10, 0, 1, [1.0, 2.0]), + _row(10, 0, 2, [3.0, 4.0]), + _row(10, 0, 3, [5.0, 6.0]), + ] + ), + ) + uploader.close() + + self.assertEqual(len(view.arrow_calls), 2) + self.assertEqual([c["ts"] for c in view.arrow_calls], [100, 101]) + self.assertEqual({c["version"] for c in view.arrow_calls}, {"model_a@export_1"}) + self.assertEqual({c["write_mode"] for c in view.arrow_calls}, {"MERGE"}) + + batch0 = view.arrow_calls[0]["batch"] + self.assertEqual(batch0.schema.names, ["embedding_name", "key_id", "embedding"]) + self.assertEqual(batch0.num_rows, 2) + # PK is the remapped table_fqn (string); SK stays int64 (the SDK casts to + # the string wire type); embedding is list reused zero-copy. + self.assertEqual(batch0.column("embedding_name").type, pa.string()) + self.assertEqual(batch0.column("key_id").type, pa.int64()) + self.assertEqual(batch0.column("embedding").type, pa.list_(pa.float32())) + self.assertEqual( + batch0.column("embedding_name").to_pylist(), + ["model.ebc.embedding_bags.user_emb"] * 2, + ) + self.assertEqual(batch0.column("key_id").to_pylist(), [1, 2]) + self.assertEqual( + batch0.column("embedding").to_pylist(), + [[1.0, 2.0], [3.0, 4.0]], + ) + self.assertEqual(view.closed, [True]) + + def test_multi_chunk_table_keeps_timestamps_monotonic_across_steps(self): + # A multi-FQN delta table concatenates one chunk per FQN; to_batches() + # splits each chunk independently, so the ts range must cover every + # actual batch or a stuck clock reuses a prior step's timestamps and + # Next-Ts incremental readers miss updates. + view = _FakeView() + factory = _FakeClientFactory(view) + uploader = self._uploader( + _feature_store_config(upload_batch_size=1000), + client_factory=factory, + clock_ms=lambda: 100, + embedding_dimensions={ + "model.ebc.embedding_bags.user_emb": 2, + "model.ebc.embedding_bags.item_emb": 2, + }, + ) + uploader.start() + rows_a = [_row(10, 0, k, [1.0, 2.0]) for k in range(5)] + rows_b = [_row(10, 0, k, [3.0, 4.0], name="item_emb") for k in range(5)] + table = pa.concat_tables([_delta_table(rows_a), _delta_table(rows_b)]) + uploader.submit(10, table) + uploader.submit(20, table) + uploader.close() + + ts_values = [call["ts"] for call in view.calls] + # Two chunks -> two batches per step; the stuck clock forces step 2 to + # start strictly after step 1's last ts (101), i.e. 102. + self.assertEqual(ts_values, [100, 101, 102, 103]) + self.assertEqual([len(call["data"]) for call in view.calls], [5, 5, 5, 5]) + self.assertEqual(view.closed, [True]) + + def test_mixed_fqn_dimensions_validated_per_row(self): + # _validate_delta_batch keys the expected dimension per row, so a batch + # carrying multiple FQNs with different dimensions must not be + # mis-flagged as a dimension mismatch. + view = _FakeView() + factory = _FakeClientFactory(view) + uploader = self._uploader( + _feature_store_config(upload_batch_size=1000), + client_factory=factory, + clock_ms=lambda: 100, + embedding_dimensions={ + "model.ebc.embedding_bags.user_emb": 2, + "model.ebc.embedding_bags.item_emb": 3, + }, + ) + uploader.start() + uploader.submit( + 10, + _delta_table( + [ + _row(10, 0, 1, [1.0, 2.0]), + _row(10, 0, 2, [3.0, 4.0, 5.0], name="item_emb"), + ] + ), + ) + uploader.close() + + self.assertEqual(len(view.calls), 1) + self.assertEqual(len(view.calls[0]["data"]), 2) + self.assertEqual( + view.calls[0]["data"][0]["embedding_name"], + "model.ebc.embedding_bags.user_emb", + ) + self.assertEqual( + view.calls[0]["data"][1]["embedding_name"], + "model.ebc.embedding_bags.item_emb", + ) + self.assertTrue( + np.array_equal( + view.calls[0]["data"][0]["embedding"], + np.array([1.0, 2.0], dtype=np.float32), + ) + ) + self.assertTrue( + np.array_equal( + view.calls[0]["data"][1]["embedding"], + np.array([3.0, 4.0, 5.0], dtype=np.float32), + ) + ) + self.assertEqual(view.closed, [True]) + if __name__ == "__main__": unittest.main() From fd8673a346642bd09d968a3e6a693083d8cff7a7 Mon Sep 17 00:00:00 2001 From: gecheng Date: Wed, 29 Jul 2026 19:46:42 +0800 Subject: [PATCH 48/50] [refactor] remove fail_on_uneven_data from create_train_pipeline The dataloader already guarantees even batches across workers, so the fail_on_uneven_data guard is unreachable and can be removed. --- tzrec/main.py | 4 -- tzrec/utils/delta_embedding_dump_test.py | 53 ------------------------ tzrec/utils/dist_util.py | 13 ------ 3 files changed, 70 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index be8ffa23b..6eb3632ea 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -473,14 +473,10 @@ def run_eval(step: int, epoch: int) -> None: ) try: for i_epoch in epoch_iter: - # Multi-rank dumps synchronize their per-step state, so uneven workers - # must fail together instead of silently dropping a lagging rank's - # trailing delta rows or leaving collectives. pipeline = create_train_pipeline( model, optimizer, check_all_workers_data_status=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/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 5e7fa0f7e..9d5f99307 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -72,7 +72,6 @@ validate_delta_embedding_dump_config, validate_delta_embedding_dump_no_zch_features, ) -from tzrec.utils.dist_util import create_train_pipeline from tzrec.utils.dynamicemb_util import has_dynamicemb from tzrec.utils.test_util import gpu_unavailable, make_test_dir, mark_ci_scope @@ -372,38 +371,6 @@ def _run_shared_table_fqn_delta_embedding_dump( ) -def _run_uneven_exhaustion_rejection(rank: int, world_size: int, output_dir: str): - # Regression: with rank 0 stopping at step 49 and rank 1 at step 50, - # final_dump used to sync both ranks to the boundary 50, both took the - # boundary-step skip, and rank 0's trailing tracked rows were silently - # lost. Multi-rank dumps now force aligned dataloader exhaustion, so the - # uneven stop must fail loudly on every rank instead. - with MultiProcessContext(rank=rank, world_size=world_size, backend="nccl") as ctx: - model = _build_sharded_delta_dump_model(rank, world_size, ctx) - dumper = DeltaEmbeddingDumper( - model, - DeltaEmbeddingDumpConfig(dump_interval_steps=50, output_dir=output_dir), - output_dir, - torch.device(f"cuda:{rank}"), - [], - ) - testcase = unittest.TestCase() - testcase.assertTrue(dumper.requires_synced_dataloader_exhaustion) - pipeline = create_train_pipeline( - model, - check_all_workers_data_status=dumper.requires_synced_dataloader_exhaustion, - fail_on_uneven_data=dumper.requires_synced_dataloader_exhaustion, - ) - # Rank 0's 50th fetch returns None while rank 1 still has batch 50; the - # pipeline's collective data-status check must raise on both ranks. - num_batches = 49 if rank == 0 else 50 - batches = iter([_sharded_features(rank) for _ in range(num_batches)]) - with testcase.assertRaisesRegex(RuntimeError, "exhausted unevenly"): - for _ in range(num_batches + 1): - pipeline._next_batch(batches) - torch.distributed.barrier() - - class DeltaEmbeddingDumpValidationTest(unittest.TestCase): def test_missing_config_skips_runtime_validation(self): with mock.patch.dict(os.environ, {"WORLD_SIZE": "2"}): @@ -1546,26 +1513,6 @@ def test_shared_table_name_tracks_ec_and_ebc_fqns_independently(self): output_dir=tmp_dir, ) - @unittest.skipIf(torch.cuda.device_count() < 2, "test requires 2+ GPUs") - @mark_ci_scope("gpu") - def test_uneven_exhaustion_is_rejected_for_step_interval_dump(self): - with ( - tempfile.TemporaryDirectory() as tmp_dir, - mock.patch.dict( - os.environ, - { - "NCCL_DEBUG": "WARN", - "FORCED_NCCL_DEBUG": "WARN", - "NCCL_DEBUG_SUBSYS": "", - }, - ), - ): - self._run_multi_process_test( - callable=_run_uneven_exhaustion_rejection, - world_size=self.world_size, - output_dir=tmp_dir, - ) - class DeltaEmbeddingDumpDynamicembIntegrationTest(unittest.TestCase): """End-to-end multi-process delta dump over a sharded dynamicemb model. diff --git a/tzrec/utils/dist_util.py b/tzrec/utils/dist_util.py index 1ab6859ec..88bb0f182 100644 --- a/tzrec/utils/dist_util.py +++ b/tzrec/utils/dist_util.py @@ -237,7 +237,6 @@ 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, @@ -252,7 +251,6 @@ 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() @@ -291,13 +289,6 @@ def _next_batch(self, dataloader_iter: Iterator[In]) -> Optional[In]: ) dist.all_reduce(has_batch, dist.ReduceOp.AVG) 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 @@ -346,7 +337,6 @@ 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. @@ -355,8 +345,6 @@ 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. @@ -386,5 +374,4 @@ 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, ) From 488b60111639a27da616293fa86a5d9af9b1fd56 Mon Sep 17 00:00:00 2001 From: gecheng Date: Thu, 30 Jul 2026 10:55:23 +0800 Subject: [PATCH 49/50] bumpup to 1.3.9 --- tzrec/version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tzrec/version.py b/tzrec/version.py index af5c54e6f..ca1bee616 100644 --- a/tzrec/version.py +++ b/tzrec/version.py @@ -9,4 +9,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "1.3.7" +__version__ = "1.3.9" From f7ec4dfb2a4f0eb53951a4f18ab66133458c7db9 Mon Sep 17 00:00:00 2001 From: gecheng Date: Thu, 30 Jul 2026 14:22:18 +0800 Subject: [PATCH 50/50] [refactor] drop delta-dumper wiring into check_all_workers_data_status check_all_workers_data_status is a fallback that only belongs on when the dataloader itself cannot guarantee even batch counts across ranks; the whole distributed training chain already requires even exhaustion. The delta embedding dumper is not a source of unevenness, so deriving this flag from requires_synced_dataloader_exhaustion conflated a downstream consumer's sync need with a dataloader-capability knob. Restore the plain passthrough of check_all_workers_data_status and remove the now-unused property, its contract tests, and two comments that still claimed dumps force synced exhaustion. Co-Authored-By: Claude --- tzrec/main.py | 13 +++--------- tzrec/utils/delta_embedding_dump.py | 14 +----------- tzrec/utils/delta_embedding_dump_test.py | 27 ------------------------ 3 files changed, 4 insertions(+), 50 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index 6eb3632ea..1fe66cdf9 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -464,19 +464,12 @@ def run_eval(step: int, epoch: int) -> None: # this rank's last consumed event-time, reused by the epoch / final saves data_timestamp = -1.0 - require_equal_train_batches = ( - delta_embedding_dumper is not None - and delta_embedding_dumper.requires_synced_dataloader_exhaustion - ) - sync_train_data_exhaustion = ( - check_all_workers_data_status or require_equal_train_batches - ) try: for i_epoch in epoch_iter: pipeline = create_train_pipeline( model, optimizer, - check_all_workers_data_status=sync_train_data_exhaustion, + check_all_workers_data_status=check_all_workers_data_status, ) if plogger is not None: plogger.set_description(f"Training Epoch {i_epoch}") @@ -588,8 +581,8 @@ def run_eval(step: int, epoch: int) -> None: if delta_embedding_dumper is not None: # Flush the trailing partial interval before the final checkpoint. # final_dump skips dump-boundary steps already written by maybe_dump - # (multi-rank dumps force synced exhaustion, so every rank participated - # in those dumps and reaches the same final step). + # (all ranks run the same step count, so every rank participated in + # those dumps and reaches the same final step). delta_embedding_dumper.final_dump(i_step) _log_train( diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index f03bb850f..fbea3ee3d 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -650,18 +650,6 @@ def pause_tracking(self) -> Iterator[None]: finally: self._tracking_pause_depth -= 1 - @property - def requires_synced_dataloader_exhaustion(self) -> bool: - """Return whether input exhaustion must stay aligned across ranks. - - Multi-rank dumps of either cadence need every rank to reach the same - final step. ``final_dump`` skips boundary steps already written on their - dump boundary, and a rank that stops before the synced step never ran - that boundary's ``maybe_dump``, so an unsynced stop would make it adopt - the boundary skip and silently drop its trailing tracked rows. - """ - return self._world_size > 1 - def start(self) -> None: """Start timed cadence and per-rank FeatureStore publication. @@ -746,7 +734,7 @@ def final_dump(self, global_step: int) -> Optional[str]: # 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 - # delta (multi-rank dumps force synced exhaustion, so every rank + # delta (all ranks run the same step count, so every rank # participated in the boundary dump) -- and torchrec's # ``get_unique`` raises ``torch.cat(): expected a non-empty list of # Tensors`` on the empty consumer window. Re-dumping would also diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 9d5f99307..4272b8539 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -1077,33 +1077,6 @@ def test_collect_dynamic_modules_uses_owner_fqn_keys(self): self.assertIs(dynamic_modules["model.ebc.shared"], ebc_dynamic_module) self.assertIs(dynamic_modules["model.ec.shared"], ec_dynamic_module) - 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_multi_rank_step_interval_requires_synced_dataloader_exhaustion(self): - # Regression: with ranks finishing at steps 49 and 50, final_dump synced - # both to the boundary 50, both took the boundary-step skip, and the - # rank that never ran maybe_dump(50) silently dropped its trailing - # tracked rows. Every multi-rank cadence must keep exhaustion aligned - # so all ranks participate in every boundary dump. - dumper = object.__new__(DeltaEmbeddingDumper) - dumper._interval_secs = None - dumper._interval_steps = 50 - dumper._world_size = 2 - - self.assertTrue(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_multi_gpu_output_path_uses_step_underscore_dir(self): with tempfile.TemporaryDirectory() as tmp_dir: dumper = object.__new__(DeltaEmbeddingDumper)