diff --git a/tzrec/main.py b/tzrec/main.py index e1905ae7..4d3864fd 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -92,6 +92,7 @@ ) from tzrec.utils.filesystem_util import url_to_fs from tzrec.utils.logging_util import ProgressLogger, logger +from tzrec.utils.online_dense_export_util import OnlineDenseExportManager from tzrec.utils.plan_util import create_planner, get_default_sharders from tzrec.version import __version__ as tzrec_version @@ -336,6 +337,7 @@ def _train_and_evaluate( ignore_restore_optimizer: bool = False, dataloader_state: Optional[Dict[str, Any]] = None, delta_embedding_dumper: Optional[DeltaEmbeddingDumper] = None, + pipeline_config_path: Optional[str] = None, ) -> None: """Train and evaluate the model.""" is_rank_zero = int(os.environ.get("RANK", 0)) == 0 @@ -450,6 +452,16 @@ def run_eval(step: int, epoch: int) -> None: ) model.train() + online_dense_exporter = OnlineDenseExportManager( + model_dir, + pipeline_config_path or os.path.join(model_dir, "pipeline.config"), + ckpt_manager, + ) + + def after_checkpoint_saved(step: int, data_ts: float) -> None: + checkpoint_path = os.path.join(model_dir, f"model.ckpt-{step}") + online_dense_exporter.submit(step, checkpoint_path, data_ts) + # this rank's last consumed event-time, reused by the epoch / final saves data_timestamp = -1.0 for i_epoch in epoch_iter: @@ -530,6 +542,7 @@ def run_eval(step: int, epoch: int) -> None: dataloader_state, data_timestamp=data_timestamp, ): + after_checkpoint_saved(i_step, data_timestamp) run_eval(i_step, i_epoch) if train_config.is_profiling: prof.step() @@ -542,6 +555,7 @@ def run_eval(step: int, epoch: int) -> None: epoch=i_epoch, data_timestamp=data_timestamp, ): + after_checkpoint_saved(i_step, data_timestamp) run_eval(i_step, i_epoch) if use_step and i_step >= train_config.num_steps - 1: @@ -586,7 +600,9 @@ def run_eval(step: int, epoch: int) -> None: data_timestamp=data_timestamp, final=True, ): + after_checkpoint_saved(i_step, data_timestamp) run_eval(i_step, i_epoch) + online_dense_exporter.close() ckpt_manager.close() @@ -851,6 +867,7 @@ def train_and_evaluate( ignore_restore_optimizer=ignore_restore_optimizer, dataloader_state=dataloader_state, delta_embedding_dumper=delta_embedding_dumper, + pipeline_config_path=os.path.join(pipeline_config.model_dir, "pipeline.config"), ) if is_local_rank_zero: logger.info("Train and Evaluate Finished.") diff --git a/tzrec/tools/online_dense_export.py b/tzrec/tools/online_dense_export.py new file mode 100644 index 00000000..31143a12 --- /dev/null +++ b/tzrec/tools/online_dense_export.py @@ -0,0 +1,154 @@ +# 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. + +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); + +import argparse +import datetime +import json +import os +import shutil +from typing import Any, Dict, Optional + +from tzrec.main import _create_features, _create_model +from tzrec.models.model import ScriptWrapper +from tzrec.utils import config_util, env_util +from tzrec.utils.export_util import ( + ensure_input_tile_for_distributed_embedding, + export_dense_model_cpu, +) +from tzrec.utils.logging_util import logger +from tzrec.utils.online_dense_export_util import make_version + +DENSE_HOT_EXPORT_DIR = "dense_hot_export" +VERSIONS_DIR = "versions" +CURRENT_JSON = "current.json" + + +def _utc_now() -> str: + return datetime.datetime.now(datetime.timezone.utc).isoformat() + + +def _atomic_write_json(path: str, data: Dict[str, Any]) -> None: + os.makedirs(os.path.dirname(path), exist_ok=True) + tmp_path = f"{path}.tmp.{os.getpid()}" + with open(tmp_path, "w") as f: + json.dump(data, f, indent=2, sort_keys=True) + f.write("\n") + os.replace(tmp_path, path) + + +def _publish_current(current_path: str, payload: Dict[str, Any]) -> None: + _atomic_write_json(current_path, payload) + + +def export_online_dense_model( + pipeline_config_path: str, + checkpoint_path: str, + model_dir: str, + version: Optional[str] = None, + checkpoint_step: Optional[int] = None, + data_timestamp: Optional[float] = None, +) -> Dict[str, Any]: + """Export and publish one online-learning dense model version.""" + if not env_util.use_distributed_embedding(): + raise RuntimeError("ONLINE_DENSE_EXPORT requires USE_DISTRIBUTED_EMBEDDING=1.") + + ensure_input_tile_for_distributed_embedding() + + version = version or make_version() + export_root = os.path.join(model_dir, DENSE_HOT_EXPORT_DIR) + versions_root = os.path.join(export_root, VERSIONS_DIR) + version_dir = os.path.join(versions_root, version) + tmp_dir = f"{version_dir}.tmp.{os.getpid()}" + + if os.path.exists(version_dir): + raise RuntimeError(f"dense version already exists: {version_dir}") + if os.path.exists(tmp_dir): + shutil.rmtree(tmp_dir) + os.makedirs(versions_root, exist_ok=True) + + pipeline_config = config_util.load_pipeline_config(pipeline_config_path) + features = _create_features( + list(pipeline_config.feature_configs), pipeline_config.data_config + ) + model = _create_model( + pipeline_config.model_config, + features, + list(pipeline_config.data_config.label_fields), + sampler_type=None, + ) + model.set_is_inference(True) + scripted_model = ScriptWrapper(model) + + try: + export_dense_model_cpu( + pipeline_config=pipeline_config, + model=scripted_model, + checkpoint_path=checkpoint_path, + save_dir=tmp_dir, + ) + + required_files = ["scripted_model.pt", "dense_meta.json"] + for file_name in required_files: + file_path = os.path.join(tmp_dir, file_name) + if not os.path.exists(file_path): + raise RuntimeError(f"missing dense export artifact: {file_path}") + + ready_path = os.path.join(tmp_dir, "READY") + with open(ready_path, "w") as f: + f.write(_utc_now()) + f.write("\n") + + os.rename(tmp_dir, version_dir) + except BaseException: + if os.path.exists(tmp_dir): + shutil.rmtree(tmp_dir) + raise + + current_payload: Dict[str, Any] = { + "version": version, + "checkpoint_path": os.path.abspath(checkpoint_path), + "created_at": _utc_now(), + } + + # Keep the service-facing pointer beside the immutable dense export versions. + _publish_current(os.path.join(export_root, CURRENT_JSON), current_payload) + logger.info("published online dense export version %s to %s", version, version_dir) + return current_payload + + +def main() -> None: + """Run the online dense export command-line entrypoint.""" + parser = argparse.ArgumentParser( + description="Export one online-learning dense model version." + ) + parser.add_argument("--pipeline_config_path", required=True) + parser.add_argument("--checkpoint_path", required=True) + parser.add_argument("--model_dir", required=True) + parser.add_argument("--version", default=None) + parser.add_argument("--checkpoint_step", type=int, default=None) + parser.add_argument("--data_timestamp", type=float, default=None) + args = parser.parse_args() + + export_online_dense_model( + pipeline_config_path=args.pipeline_config_path, + checkpoint_path=args.checkpoint_path, + model_dir=args.model_dir, + version=args.version, + checkpoint_step=args.checkpoint_step, + data_timestamp=args.data_timestamp, + ) + + +if __name__ == "__main__": + main() diff --git a/tzrec/tools/online_dense_export_test.py b/tzrec/tools/online_dense_export_test.py new file mode 100644 index 00000000..2cf9184e --- /dev/null +++ b/tzrec/tools/online_dense_export_test.py @@ -0,0 +1,116 @@ +# 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. + +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); + +import json +import os +import tempfile +import unittest +from types import SimpleNamespace +from unittest import mock + +from tzrec.tools.online_dense_export import export_online_dense_model + + +class OnlineDenseExportTest(unittest.TestCase): + def test_export_online_dense_model_publishes_ready_version(self) -> None: + with tempfile.TemporaryDirectory() as tmp_dir: + checkpoint_path = os.path.join(tmp_dir, "model.ckpt-10") + os.makedirs(checkpoint_path) + pipeline_config_path = os.path.join(tmp_dir, "pipeline.config") + open(pipeline_config_path, "w").close() + + def fake_export_dense_model_cpu(**kwargs): + self.assertNotIn("dense_only", kwargs) + save_dir = kwargs["save_dir"] + os.makedirs(save_dir, exist_ok=True) + with open(os.path.join(save_dir, "scripted_model.pt"), "w") as f: + f.write("pt") + with open(os.path.join(save_dir, "dense_meta.json"), "w") as f: + json.dump({"group": ["feature__ebc"]}, f) + + dummy_config = SimpleNamespace( + feature_configs=[], + data_config=SimpleNamespace(label_fields=[]), + model_config=SimpleNamespace(), + ) + dummy_model = mock.Mock() + + with ( + mock.patch.dict( + os.environ, + {"USE_DISTRIBUTED_EMBEDDING": "1"}, + clear=False, + ), + mock.patch( + "tzrec.tools.online_dense_export.config_util.load_pipeline_config", + return_value=dummy_config, + ), + mock.patch( + "tzrec.tools.online_dense_export._create_features", + return_value=[], + ), + mock.patch( + "tzrec.tools.online_dense_export._create_model", + return_value=dummy_model, + ), + mock.patch( + "tzrec.tools.online_dense_export.ScriptWrapper", + side_effect=lambda model: model, + ), + mock.patch( + "tzrec.tools.online_dense_export.export_dense_model_cpu", + side_effect=fake_export_dense_model_cpu, + ), + ): + payload = export_online_dense_model( + pipeline_config_path=pipeline_config_path, + checkpoint_path=checkpoint_path, + model_dir=tmp_dir, + version="20260623174703", + checkpoint_step=10, + data_timestamp=42.0, + ) + + version_dir = os.path.join( + tmp_dir, "dense_hot_export", "versions", "20260623174703" + ) + self.assertTrue(os.path.exists(os.path.join(version_dir, "READY"))) + self.assertTrue( + os.path.exists(os.path.join(version_dir, "scripted_model.pt")) + ) + self.assertTrue( + os.path.exists(os.path.join(version_dir, "dense_meta.json")) + ) + + current_path = os.path.join(tmp_dir, "dense_hot_export", "current.json") + with open(current_path) as f: + current = json.load(f) + self.assertEqual( + set(current.keys()), {"checkpoint_path", "created_at", "version"} + ) + self.assertEqual(current["version"], "20260623174703") + self.assertEqual( + current["checkpoint_path"], os.path.abspath(checkpoint_path) + ) + self.assertTrue(current["created_at"]) + self.assertFalse(os.path.exists(os.path.join(tmp_dir, "dense_hot_update"))) + + self.assertEqual( + set(payload.keys()), {"checkpoint_path", "created_at", "version"} + ) + self.assertEqual(payload["version"], "20260623174703") + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/utils/checkpoint_util.py b/tzrec/utils/checkpoint_util.py index 672ecd4c..d2c1415c 100644 --- a/tzrec/utils/checkpoint_util.py +++ b/tzrec/utils/checkpoint_util.py @@ -20,7 +20,7 @@ import threading import weakref from dataclasses import replace -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Set, Tuple import torch import torch.distributed as dist @@ -331,6 +331,7 @@ def __init__( self._eval_result_filename = eval_result_filename self._prune_queue: "queue.Queue[object]" = queue.Queue() self._prune_pending = False + self._protected_checkpoints: Set[str] = set() self._lock = threading.Lock() self._prune_worker: Optional[threading.Thread] = None self._finalizer: Optional[weakref.finalize] = None @@ -364,6 +365,18 @@ def save( self.prune() return ckpt_dir + def protect_checkpoint(self, ckpt_path: str) -> None: + """Prevent an in-progress consumer from being pruned.""" + with self._lock: + self._protected_checkpoints.add(self._canonical_checkpoint_path(ckpt_path)) + + def unprotect_checkpoint(self, ckpt_path: str) -> None: + """Allow a previously protected checkpoint to be pruned again.""" + with self._lock: + self._protected_checkpoints.discard( + self._canonical_checkpoint_path(ckpt_path) + ) + def set_save_policy( self, save_steps: int, @@ -621,7 +634,10 @@ def _run_prune(self) -> None: if len(ckpt_metas) <= self._keep_checkpoint_max: return ckpt_metas.sort(key=_get_checkpoint_step) - protected = set(ckpt_metas[-self._keep_checkpoint_max :]) + protected = { + self._canonical_checkpoint_path(path) + for path in ckpt_metas[-self._keep_checkpoint_max :] + } if ( self._export_config is not None and self._export_config.exporter_type == "best" @@ -630,15 +646,21 @@ def _run_prune(self) -> None: self._model_dir, self._export_config, self._eval_result_filename ) if best_ckpt_path is not None: - protected.add(best_ckpt_path.rstrip(os.path.sep)) + protected.add(self._canonical_checkpoint_path(best_ckpt_path)) + with self._lock: + protected.update(self._protected_checkpoints) for ckpt_path in ckpt_metas: - if ckpt_path not in protected: + if self._canonical_checkpoint_path(ckpt_path) not in protected: logger.info(f"Removing old checkpoint {ckpt_path}...") try: shutil.rmtree(ckpt_path) except Exception as e: # noqa: BLE001 logger.warning(f"Failed to remove checkpoint {ckpt_path}: {e}") + @staticmethod + def _canonical_checkpoint_path(ckpt_path: str) -> str: + return os.path.abspath(ckpt_path.rstrip(os.path.sep)) + _DISTCP_RANK_RE = re.compile(r"__(\d+)_\d+\.distcp$") diff --git a/tzrec/utils/checkpoint_util_test.py b/tzrec/utils/checkpoint_util_test.py index 2be844d1..c64e70b0 100644 --- a/tzrec/utils/checkpoint_util_test.py +++ b/tzrec/utils/checkpoint_util_test.py @@ -273,6 +273,25 @@ def test_checkpoint_manager_prune_non_rank_zero(self): manager.close() self.assertEqual(self._remaining_ckpt_steps(), [0, 10, 20, 30]) + def test_checkpoint_manager_prune_keeps_protected_checkpoint(self): + for step in [0, 10, 20, 30]: + os.makedirs(os.path.join(self.test_dir, f"model.ckpt-{step}")) + protected_ckpt = os.path.join(self.test_dir, "model.ckpt-10") + manager = checkpoint_util.CheckpointManager( + self.test_dir, keep_checkpoint_max=2 + ) + manager.protect_checkpoint(protected_ckpt) + with mock.patch.dict(os.environ, {"RANK": "0"}): + manager.prune() + manager.close() + self.assertEqual(self._remaining_ckpt_steps(), [10, 20, 30]) + + manager.unprotect_checkpoint(protected_ckpt) + with mock.patch.dict(os.environ, {"RANK": "0"}): + manager.prune() + manager.close() + self.assertEqual(self._remaining_ckpt_steps(), [20, 30]) + def test_checkpoint_manager_prune_idempotent(self): for step in [0, 10, 20, 30]: os.makedirs(os.path.join(self.test_dir, f"model.ckpt-{step}")) diff --git a/tzrec/utils/env_util.py b/tzrec/utils/env_util.py index 649b8368..4f931276 100644 --- a/tzrec/utils/env_util.py +++ b/tzrec/utils/env_util.py @@ -32,6 +32,11 @@ def use_rtp() -> bool: return flag +def use_distributed_embedding() -> bool: + """Export model for distributed embedding mode of EAS processor.""" + return os.environ.get("USE_DISTRIBUTED_EMBEDDING", "0") == "1" + + def enable_tma() -> bool: """Enable TMA (Tensor Memory Accelerator) for triton ops.""" flag = os.environ.get("ENABLE_TMA", "0") == "1" diff --git a/tzrec/utils/export_util.py b/tzrec/utils/export_util.py index 37762627..33474c20 100644 --- a/tzrec/utils/export_util.py +++ b/tzrec/utils/export_util.py @@ -445,6 +445,78 @@ def _get_dense_embedding_leaf_module_names(model: torch.nn.Module) -> List[str]: return names +def _get_sparse_embedding_leaf_module_names(model: torch.nn.Module) -> List[str]: + """Get sparse embedding modules to keep as FX leaf modules during export.""" + names = [] + for path, module in model.named_modules(): + if isinstance( + module, + (EmbeddingBagCollectionInterface, EmbeddingCollectionInterface), + ): + names.append(path) + return names + + +def _resolve_keyed_tensor_source_module( + model: torch.nn.Module, node: Any +) -> Optional[torch.nn.Module]: + """Resolve the module producing a KeyedTensor FX node.""" + if getattr(node, "op", None) == "call_module": + return model.get_submodule(node.target) + if getattr(node, "op", None) == "call_function": + if node.target == operator.getitem: + return _resolve_keyed_tensor_source_module(model, node.args[0]) + if node.target == KeyedTensor: + for attr_name in ("keys", "length_per_key", "values"): + attr_node = node.kwargs.get(attr_name) + source_module = _resolve_keyed_tensor_source_module(model, attr_node) + if source_module is not None: + return source_module + if getattr(node, "op", None) == "call_method" and node.target in ( + "keys", + "length_per_key", + "values", + "tile", + ): + return _resolve_keyed_tensor_source_module(model, node.args[0]) + return None + + +def _get_embedding_bag_configs(module: torch.nn.Module) -> Optional[List[Any]]: + """Get EmbeddingBag configs from EBC or MC_EBC modules by duck typing.""" + if hasattr(module, "embedding_bag_configs"): + return list(module.embedding_bag_configs()) + inner = getattr(module, "_embedding_module", None) + if inner is not None and hasattr(inner, "embedding_bag_configs"): + return list(inner.embedding_bag_configs()) + return None + + +def _infer_keyed_tensor_attrs_from_module( + module: torch.nn.Module, +) -> Optional[Tuple[List[str], List[int]]]: + """Infer KeyedTensor keys/length_per_key from embedding module configs.""" + configs = _get_embedding_bag_configs(module) + if configs is None: + return None + + keys = [] + length_per_key = [] + for config in configs: + embedding_names = getattr(config, "embedding_names", None) + if not embedding_names: + embedding_names = [config.name] + embedding_dim = config.embedding_dim + keys.extend(embedding_names) + length_per_key.extend([embedding_dim] * len(embedding_names)) + return keys, length_per_key + + +def _is_fx_node(value: Any) -> bool: + """Whether value is an FX node from the traced graph.""" + return getattr(value, "op", None) is not None and hasattr(value, "target") + + def _get_rtp_feature_to_embedding_info( model: nn.Module, ) -> Dict[str, BaseEmbeddingConfig]: @@ -1640,6 +1712,166 @@ def export_distributed_embedding( json.dump(merged_emb_json, f, indent=4) +def export_dense_model_cpu( + pipeline_config: EasyRecConfig, + model: BaseModule, + checkpoint_path: Optional[str], + save_dir: str, + assets: Optional[List[str]] = None, + use_local_cache_dir: bool = False, + data_input_path: Optional[str] = None, + **kwargs: Any, +) -> None: + """Export only the dense model on CPU without DMP or GPU usage.""" + del pipeline_config, assets, use_local_cache_dir, data_input_path, kwargs + if not checkpoint_path: + raise ValueError("checkpoint path should be specified.") + + device = torch.device("cpu") + graph_dir = os.path.join(save_dir, "graph") + os.makedirs(save_dir, exist_ok=True) + os.makedirs(graph_dir, exist_ok=True) + + model.set_is_inference(True) + model.eval() + + leaf_modules = _get_sparse_embedding_leaf_module_names( + model + ) + _get_dense_embedding_leaf_module_names(model) + tracer = Tracer(leaf_modules=leaf_modules) + full_graph = tracer.trace(model) + with open(os.path.join(graph_dir, "gm_full.graph"), "w") as f: + f.write(str(full_graph)) + + logger.info("collecting sparse attrs statically for CPU dense export...") + sparse_attrs = {} + for node in list(full_graph.nodes): + if node.op != "call_function" or node.target != fx_mark_keyed_tensor: + continue + name = node.args[0] + if node.kwargs.get("is_dense", False): + continue + node_kt = node.args[1] + node_kt_kwargs = getattr(node_kt, "kwargs", {}) + keys = node_kt_kwargs.get("keys") + length_per_key = node_kt_kwargs.get("length_per_key") + need_infer_attrs = ( + keys is None + or length_per_key is None + or _is_fx_node(keys) + or _is_fx_node(length_per_key) + ) + if need_infer_attrs: + source_module = _resolve_keyed_tensor_source_module(model, node_kt) + inferred_attrs = ( + _infer_keyed_tensor_attrs_from_module(source_module) + if source_module is not None + else None + ) + if inferred_attrs is not None: + keys, length_per_key = inferred_attrs + if ( + keys is None + or length_per_key is None + or _is_fx_node(keys) + or _is_fx_node(length_per_key) + ): + raise RuntimeError( + "CPU dense export cannot statically infer KeyedTensor attrs " + f"for feature group [{name}]." + ) + sparse_attrs[name + "__keys"] = keys + sparse_attrs[name + "__length_per_key"] = length_per_key + + logger.info("exporting dense model on CPU...") + graph = copy.deepcopy(full_graph) + output_keys = [] + output_values = [] + dense_graph_config = defaultdict() + for node in graph.nodes: + if node.op == "output": + for k, v in sorted(node.args[0].items()): + if k == TARGET_REPEAT_INTERLEAVE_KEY: + continue + output_keys.append(k) + output_values.append(v) + graph.erase_node(node) + input_node = next(node for node in graph.nodes if node.op == "placeholder") + + dense_graph_config["sequence__ec"] = [] + for node in list(graph.nodes): + if node.op == "call_function" and node.target == fx_mark_keyed_tensor: + name = node.args[0] + if node.kwargs.get("is_dense", False): + continue + node_kt = node.args[1] + with graph.inserting_before(node_kt): + getitem_node = graph.call_function( + operator.getitem, args=(input_node, name) + ) + values_node = getitem_node + if _is_input_tile_user_keyed_tensor(name): + batch_size_node = graph.call_function( + operator.getitem, args=(input_node, "batch_size") + ) + tile_size_node = graph.call_function( + _tile_size, args=(batch_size_node,) + ) + values_node = graph.call_method( + "tile", args=(getitem_node, tile_size_node, 1) + ) + new_node = graph.call_function( + KeyedTensor, + kwargs={ + "keys": sparse_attrs[name + "__keys"], + "length_per_key": sparse_attrs[name + "__length_per_key"], + "values": values_node, + }, + ) + dense_graph_config[name] = [ + k + "__ebc" for k in new_node.kwargs["keys"] + ] + node_kt.replace_all_uses_with(new_node) + elif node.op == "call_function" and node.target == fx_mark_seq_ec_jt: + name = node.args[0] + node_jt = node.args[1] + with graph.inserting_before(node_jt): + getitem_node = graph.call_function( + operator.getitem, args=(input_node, name) + ) + getitem_lengths = graph.call_function( + operator.getitem, args=(input_node, name + "__lengths") + ) + new_node = graph.call_function( + JaggedTensor, + kwargs={"values": getitem_node, "lengths": getitem_lengths}, + ) + node_jt.replace_all_uses_with(new_node) + emb_name = name + "__ec" + dense_graph_config["sequence__ec"].append(emb_name) + dense_graph_config["sequence__ec"].append(name + "__lengths") + + graph.output(dict(zip(output_keys, output_values))) + gm = torch.fx.GraphModule(model, graph) + gm.graph.eliminate_dead_code() + gm = _prune_unused_param_and_buffer(gm) + + init_parameters(gm, device) + gm.to(device) + checkpoint_util.restore_model(checkpoint_path, gm) + + with open(os.path.join(save_dir, "dense_meta.json"), "w") as f: + json.dump(dense_graph_config, f, indent=4) + with open(os.path.join(graph_dir, "gm_dense.graph"), "w") as f: + f.write(str(gm.graph)) + + dense_model_traced = symbolic_trace(gm) + with open(os.path.join(save_dir, "gm_dense.code"), "w") as f: + f.write(dense_model_traced.code) + dense_model_scripted = torch.jit.script(dense_model_traced) + dense_model_scripted.save(os.path.join(save_dir, "scripted_model.pt")) + + def _merge_sharded_embedding_json( emb_json_files: List[Dict[str, Any]], ) -> Dict[str, Any]: diff --git a/tzrec/utils/online_dense_export_util.py b/tzrec/utils/online_dense_export_util.py new file mode 100644 index 00000000..0cc41c4a --- /dev/null +++ b/tzrec/utils/online_dense_export_util.py @@ -0,0 +1,236 @@ +# 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. + +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); + +import datetime +import os +import socket +import subprocess +import sys +from queue import Queue +from threading import Thread +from typing import Any, Dict, Optional + +from tzrec.utils import checkpoint_util, env_util +from tzrec.utils.logging_util import logger + +_DISTRIBUTED_ENV_KEYS = { + "GROUP_RANK", + "GROUP_WORLD_SIZE", + "LOCAL_RANK", + "LOCAL_WORLD_SIZE", + "MASTER_ADDR", + "MASTER_PORT", + "RANK", + "ROLE_NAME", + "ROLE_RANK", + "ROLE_WORLD_SIZE", + "WORLD_SIZE", +} +_DISTRIBUTED_ENV_PREFIXES = ("TORCHELASTIC_",) +_VERSION_TIME_FORMAT = "%Y%m%d%H%M%S" + + +def _format_version(now: datetime.datetime) -> str: + return now.strftime(_VERSION_TIME_FORMAT) + + +def make_version(now: Optional[datetime.datetime] = None) -> str: + """Build a yyyyMMddHHmmss dense export version name.""" + now = now or datetime.datetime.now() + return _format_version(now) + + +def _make_monotonic_version( + last_version: str, now: Optional[datetime.datetime] = None +) -> str: + version = make_version(now) + if not last_version or version > last_version: + return version + last_version_dt = datetime.datetime.strptime(last_version, _VERSION_TIME_FORMAT) + return _format_version(last_version_dt + datetime.timedelta(seconds=1)) + + +def _online_dense_export_enabled() -> bool: + return os.environ.get("ONLINE_DENSE_EXPORT", "0") == "1" + + +def _get_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _build_export_subprocess_env(repo_root: str) -> Dict[str, str]: + env = os.environ.copy() + for key in list(env): + if key in _DISTRIBUTED_ENV_KEYS or key.startswith(_DISTRIBUTED_ENV_PREFIXES): + del env[key] + env.update( + { + "USE_DISTRIBUTED_EMBEDDING": "1", + "INPUT_TILE": "3", + "CUDA_VISIBLE_DEVICES": "", + "RANK": "0", + "LOCAL_RANK": "0", + "WORLD_SIZE": "1", + "LOCAL_WORLD_SIZE": "1", + "MASTER_ADDR": "127.0.0.1", + "MASTER_PORT": str(_get_free_port()), + } + ) + env["PYTHONPATH"] = ( + repo_root + if not env.get("PYTHONPATH") + else repo_root + os.pathsep + env["PYTHONPATH"] + ) + return env + + +class OnlineDenseExportManager: + """Background launcher for online-learning dense model export.""" + + def __init__( + self, + model_dir: str, + pipeline_config_path: str, + ckpt_manager: checkpoint_util.CheckpointManager, + ) -> None: + self._enabled = _online_dense_export_enabled() + self._rank = int(os.environ.get("RANK", 0)) + self._model_dir = os.path.abspath(model_dir) + self._pipeline_config_path = os.path.abspath(pipeline_config_path) + self._ckpt_manager = ckpt_manager + self._queue: "Queue[Optional[Dict[str, Any]]]" = Queue() + self._worker: Optional[Thread] = None + self._last_version = "" + + if not self._enabled: + return + if not env_util.use_distributed_embedding(): + raise RuntimeError( + "ONLINE_DENSE_EXPORT=1 requires USE_DISTRIBUTED_EMBEDDING=1." + ) + if self._rank == 0: + self._worker = Thread( + target=self._worker_loop, + name="online-dense-export", + daemon=True, + ) + self._worker.start() + logger.info( + "ONLINE_DENSE_EXPORT enabled; dense versions will be exported under %s", + os.path.join(model_dir, "dense_hot_export"), + ) + + def submit( + self, + step: int, + checkpoint_path: str, + data_timestamp: float, + ) -> None: + """Queue a dense export task for one saved checkpoint.""" + if not self._enabled or self._rank != 0: + return + checkpoint_path = os.path.abspath(checkpoint_path) + version = _make_monotonic_version(self._last_version) + self._last_version = version + self._ckpt_manager.protect_checkpoint(checkpoint_path) + self._queue.put( + { + "step": step, + "checkpoint_path": checkpoint_path, + "data_timestamp": data_timestamp, + "version": version, + } + ) + + def close(self) -> None: + """Wait for queued dense export tasks to finish.""" + if self._worker is None: + return + self._queue.put(None) + self._worker.join() + + def _worker_loop(self) -> None: + while True: + task = self._queue.get() + try: + if task is None: + return + self._run_task(task) + finally: + self._queue.task_done() + + def _run_task(self, task: Dict[str, Any]) -> None: + checkpoint_path = task["checkpoint_path"] + try: + if not os.path.exists(checkpoint_path): + logger.error( + "skip online dense export version %s: checkpoint missing: %s", + task["version"], + checkpoint_path, + ) + return + + repo_root = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ) + env = _build_export_subprocess_env(repo_root) + cmd = [ + sys.executable, + "-m", + "tzrec.tools.online_dense_export", + "--pipeline_config_path", + self._pipeline_config_path, + "--checkpoint_path", + checkpoint_path, + "--model_dir", + self._model_dir, + "--version", + task["version"], + "--checkpoint_step", + str(task["step"]), + "--data_timestamp", + str(task["data_timestamp"]), + ] + logger.info( + "start online dense export version %s from %s", + task["version"], + checkpoint_path, + ) + log_dir = os.path.join(self._model_dir, "dense_hot_export", "logs") + os.makedirs(log_dir, exist_ok=True) + log_path = os.path.join(log_dir, f"{task['version']}.log") + try: + with open(log_path, "w") as log_file: + subprocess.run( + cmd, + check=True, + env=env, + cwd=repo_root, + stdout=log_file, + stderr=subprocess.STDOUT, + ) + except subprocess.CalledProcessError as e: + logger.error( + "online dense export version %s failed with return code %s, see %s", + task["version"], + e.returncode, + log_path, + ) + return + logger.info("online dense export version %s finished", task["version"]) + finally: + self._ckpt_manager.unprotect_checkpoint(checkpoint_path) + self._ckpt_manager.prune() diff --git a/tzrec/utils/online_dense_export_util_test.py b/tzrec/utils/online_dense_export_util_test.py new file mode 100644 index 00000000..e707f570 --- /dev/null +++ b/tzrec/utils/online_dense_export_util_test.py @@ -0,0 +1,83 @@ +# 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. + +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); + +import datetime +import os +import unittest +from unittest import mock + +from tzrec.utils.online_dense_export_util import ( + _build_export_subprocess_env, + _make_monotonic_version, + make_version, +) + + +class OnlineDenseExportUtilTest(unittest.TestCase): + """Tests for online dense export utilities.""" + + def test_make_version_uses_yyyymmddhhmmss(self) -> None: + version = make_version(datetime.datetime(2026, 6, 23, 17, 47, 3)) + + self.assertEqual(version, "20260623174703") + + def test_make_monotonic_version_keeps_timestamp_format(self) -> None: + version = _make_monotonic_version( + "20260623174703", datetime.datetime(2026, 6, 23, 17, 47, 3) + ) + + self.assertEqual(version, "20260623174704") + + def test_build_export_subprocess_env_removes_torchelastic_env(self) -> None: + with ( + mock.patch.dict( + os.environ, + { + "GROUP_RANK": "3", + "LOCAL_RANK": "2", + "MASTER_ADDR": "elastic-master", + "MASTER_PORT": "123", + "PATH": "/usr/bin", + "PYTHONPATH": "/old/path", + "RANK": "2", + "TORCHELASTIC_RUN_ID": "job", + "TORCHELASTIC_USE_AGENT_STORE": "True", + "WORLD_SIZE": "4", + }, + clear=True, + ), + mock.patch( + "tzrec.utils.online_dense_export_util._get_free_port", + return_value=45678, + ), + ): + env = _build_export_subprocess_env("/repo") + + self.assertNotIn("GROUP_RANK", env) + self.assertNotIn("TORCHELASTIC_RUN_ID", env) + self.assertNotIn("TORCHELASTIC_USE_AGENT_STORE", env) + self.assertEqual(env["RANK"], "0") + self.assertEqual(env["LOCAL_RANK"], "0") + self.assertEqual(env["WORLD_SIZE"], "1") + self.assertEqual(env["LOCAL_WORLD_SIZE"], "1") + self.assertEqual(env["MASTER_ADDR"], "127.0.0.1") + self.assertEqual(env["MASTER_PORT"], "45678") + self.assertEqual(env["USE_DISTRIBUTED_EMBEDDING"], "1") + self.assertEqual(env["INPUT_TILE"], "3") + self.assertEqual(env["CUDA_VISIBLE_DEVICES"], "") + self.assertEqual(env["PYTHONPATH"], "/repo:/old/path") + + +if __name__ == "__main__": + unittest.main()