[WIP] [feat]online export of dense models

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.")

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()

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()