Merge branch 'fix_model_dump_without_infer_graph' into 'master'

fix model dump

See merge request data/monolith!2128

GitOrigin-RevId: 2448868d25c75c3063556ea933b09ee8016dd9f1

monolith/native_training/cpu_training.py

@@ -1190,6 +1190,8 @@ def create_saver():
       # should_do_first_save = self.config.partial_recovery and ckpt_state is None
       # Here we just make it false because there are issues with uninitialized iterator.
       should_do_first_save = False
+      if self.config.enable_model_dump:
+        save_utils.NoFirstSaveCheckpointSaverHook._in_model_dump_mode = True
       saver_hook = save_utils.NoFirstSaveCheckpointSaverHook(
           model_dir,
           save_secs=save_checkpoints_secs,

monolith/native_training/save_utils.py

@@ -249,6 +249,7 @@ class NoFirstSaveCheckpointSaverHook(tf.estimator.CheckpointSaverHook):
   """A saver hook which won't perform the first save (which happened on after_create_session)."""
 
   _has_dense_only: bool = False
+  _in_model_dump_mode: bool = False
   _last_triggered_step: int = 0
 
   def __init__(self,
@@ -416,11 +417,11 @@ def _create_or_update_monolith_ckpt_state(self, do_update=False):
     logging.info("monolith ckpt state saved")
 
   def end(self, session):
+    last_step = session.run(self._global_step_tensor)
     if self._is_dense_only:
       pass
-    elif self._has_dense_only:
+    elif self._has_dense_only or self._in_model_dump_mode:
       # force save
-      last_step = session.run(self._global_step_tensor)
       self._timer.update_last_triggered_step(last_step)
       super()._save(session, last_step)
       for l in self._listeners: