Internal change

recml/core/training/keras_trainer.py

@@ -126,6 +126,8 @@ def __init__(
       legacy_checkpoint_format: bool | None = None,
       train_callbacks: Sequence[keras.callbacks.Callback] = (),
       eval_callbacks: Sequence[keras.callbacks.Callback] = (),
+      verbose: int | str = "auto",
+      steps_per_execution: int | None = None,
   ):
     """Initializes the instance.
 
@@ -152,6 +154,17 @@ def __init__(
         trainer's own callbacks. Use `keras_utils.MetricsCallback` to report
         metrics, because Keras does not return changes made to `logs` in
         `on_test_end`.
+      verbose: The verbosity passed to `model.fit` and `model.evaluate`. "auto"
+        (the Keras default) and 1 show a progress bar that updates after every
+        step, 2 prints one line per epoch in `fit` and a single line in
+        `evaluate`, and 0 prints nothing.
+      steps_per_execution: If set, the number of training steps that each `jit`
+        call runs, which reduces the per-step host work of dispatching steps and
+        running batch-level callbacks. Batch-level callbacks run once per call
+        and see the logs of its last step. `steps_per_loop` must be divisible by
+        it. It applies to `train` and `train_and_evaluate`, including the
+        validation pass that `train_and_evaluate` runs inside `fit`. If None,
+        the model keeps the value that its `compile` call set.
       resume_training_launcher: Launcher to resume training.
       enable_xmanager_measurements: Whether to enable XManager measurements.
       enable_autoxprof: Whether to enable AutoXprof.
@@ -164,6 +177,17 @@ def __init__(
     if model_dir is None:
       model_dir = "/tmp"
 
+    if steps_per_execution is not None:
+      if steps_per_execution < 1:
+        raise ValueError(
+            f"`steps_per_execution` ({steps_per_execution}) must be at least 1."
+        )
+      if steps_per_loop % steps_per_execution != 0:
+        raise ValueError(
+            f"`steps_per_loop` ({steps_per_loop}) must be divisible by "
+            f"`steps_per_execution` ({steps_per_execution})."
+        )
+
     # This should be set before any layers are constructed and this is a
     # fallback in case the trainer binary doesn't already do this.
     if (
@@ -180,6 +204,7 @@ def __init__(
     self._steps_per_eval = steps_per_eval
     self._continuous_eval_timeout = continuous_eval_timeout
     self._steps_per_loop = steps_per_loop
+    self._steps_per_execution = steps_per_execution
     self._marker_path = os.path.join(
         model_dir, core.TRAINING_COMPLETE_MARKER_FILE
     )
@@ -188,6 +213,7 @@ def __init__(
     self._checkpoint_save_interval_epochs = checkpoint_save_interval_epochs
     self._custom_train_callbacks = list(train_callbacks)
     self._custom_eval_callbacks = list(eval_callbacks)
+    self._verbose = verbose
     if checkpoint_version is not None:
       self._checkpoint_version = keras_utils.CheckpointVersion(
           checkpoint_version
@@ -322,6 +348,17 @@ def _maybe_get_model_kws(
 
     return kws
 
+  def _maybe_set_steps_per_execution(self, model: keras.Model):
+    """Sets steps_per_execution and resets cached execution functions."""
+    if self._steps_per_execution is None:
+      return
+    model.steps_per_execution = self._steps_per_execution
+    # Clear functions so they are rebuilt by fit with the new value, mirroring
+    # Keras Trainer.compile (third_party/py/keras/src/trainers/trainer.py:223).
+    model.train_function = None
+    model.test_function = None
+    model.predict_function = None
+
   def _val_logs(
       self,
       history: Mapping[str, Any],
@@ -352,12 +389,15 @@ def train(self, task: KerasTask) -> core.Logs:
         **self._maybe_get_model_kws(task, dataset, val=False)
     )
 
+    self._maybe_set_steps_per_execution(model)
+
     history = model.fit(
         dataset,
         epochs=self._train_epochs + 1,
         steps_per_epoch=self._steps_per_loop,
         callbacks=self.train_callbacks,
         initial_epoch=1,
+        verbose=self._verbose,
     )
     model.summary(print_fn=logging.info)
 
@@ -387,6 +427,7 @@ def evaluate(self, task: KerasTask) -> core.Logs:
           steps=self._steps_per_eval,
           callbacks=self.eval_callbacks,
           return_dict=True,
+          verbose=self._verbose,
       )
       epoch_dt = time.time() - epoch_start_time
       steps_per_second = self._steps_per_eval / epoch_dt  # pyrefly: ignore[unsupported-operation]
@@ -399,6 +440,7 @@ def evaluate(self, task: KerasTask) -> core.Logs:
         dataset,
         steps=self._steps_per_eval,
         callbacks=self.eval_callbacks,
+        verbose=self._verbose,
     )
 
   def train_and_evaluate(self, task: KerasTask) -> core.Logs:
@@ -417,6 +459,8 @@ def train_and_evaluate(self, task: KerasTask) -> core.Logs:
         **self._maybe_get_model_kws(task, train_dataset, val=False)
     )
 
+    self._maybe_set_steps_per_execution(model)
+
     history = model.fit(
         train_dataset,
         validation_data=eval_dataset,
@@ -426,6 +470,7 @@ def train_and_evaluate(self, task: KerasTask) -> core.Logs:
         validation_steps=None,
         callbacks=self.train_callbacks,
         initial_epoch=1,
+        verbose=self._verbose,
     )
     model.summary(print_fn=logging.info)
 
@@ -514,6 +559,7 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None):
             steps=self._steps_per_eval,
             callbacks=callbacks,
             return_dict=True,
+            verbose=self._verbose,
         )
 
         logging.info(

recml/core/training/keras_trainer_test.py

@@ -158,6 +158,198 @@ def test_keras_task_and_trainer(
           [7.0] * len(history.history["loss"]),
       )
 
+  def test_verbose_passed_to_keras(self):
+    original_fit = keras.Model.fit
+    original_evaluate = keras.Model.evaluate
+
+    with mock.patch.object(
+        keras.Model, "fit", autospec=True, side_effect=original_fit
+    ) as mock_fit, mock.patch.object(
+        keras.Model, "evaluate", autospec=True, side_effect=original_evaluate
+    ) as mock_evaluate:
+      trainer = keras_trainer.KerasTrainer(
+          model_dir=self.create_tempdir().full_path,
+          train_steps=2,
+          steps_per_loop=2,
+          steps_per_eval=1,
+          verbose=2,
+          continuous_eval_timeout=5,
+          distribution=keras.distribution.DataParallel()
+          if keras.backend.backend() == "jax"
+          else None,
+      )
+      experiment = core.Experiment(_KerasTask(), trainer)
+
+      mock_fit.reset_mock()
+      mock_evaluate.reset_mock()
+      core.run_experiment(experiment, core.Trainer.Mode.TRAIN)
+      mock_fit.assert_called()
+      self.assertEqual(mock_fit.call_args.kwargs.get("verbose"), 2)
+
+      mock_fit.reset_mock()
+      mock_evaluate.reset_mock()
+      core.run_experiment(experiment, core.Trainer.Mode.EVAL)
+      mock_evaluate.assert_called()
+      self.assertEqual(mock_evaluate.call_args.kwargs.get("verbose"), 2)
+
+      mock_fit.reset_mock()
+      mock_evaluate.reset_mock()
+      core.run_experiment(experiment, core.Trainer.Mode.TRAIN_AND_EVAL)
+      mock_fit.assert_called()
+      self.assertEqual(mock_fit.call_args.kwargs.get("verbose"), 2)
+
+      if keras.backend.backend() == "jax":
+        keras.backend.clear_session()
+        mock_fit.reset_mock()
+        mock_evaluate.reset_mock()
+        core.run_experiment(experiment, core.Trainer.Mode.CONTINUOUS_EVAL)
+        mock_evaluate.assert_called()
+        self.assertEqual(mock_evaluate.call_args.kwargs.get("verbose"), 2)
+
+  def test_default_verbose_passed_to_keras(self):
+    original_fit = keras.Model.fit
+    original_evaluate = keras.Model.evaluate
+
+    with mock.patch.object(
+        keras.Model, "fit", autospec=True, side_effect=original_fit
+    ) as mock_fit, mock.patch.object(
+        keras.Model, "evaluate", autospec=True, side_effect=original_evaluate
+    ) as mock_evaluate:
+      trainer = keras_trainer.KerasTrainer(
+          model_dir=self.create_tempdir().full_path,
+          train_steps=2,
+          steps_per_loop=2,
+          steps_per_eval=1,
+          continuous_eval_timeout=5,
+          distribution=keras.distribution.DataParallel()
+          if keras.backend.backend() == "jax"
+          else None,
+      )
+      experiment = core.Experiment(_KerasTask(), trainer)
+
+      mock_fit.reset_mock()
+      mock_evaluate.reset_mock()
+      core.run_experiment(experiment, core.Trainer.Mode.TRAIN)
+      mock_fit.assert_called()
+      self.assertEqual(mock_fit.call_args.kwargs.get("verbose"), "auto")
+
+      mock_fit.reset_mock()
+      mock_evaluate.reset_mock()
+      core.run_experiment(experiment, core.Trainer.Mode.EVAL)
+      mock_evaluate.assert_called()
+      self.assertEqual(mock_evaluate.call_args.kwargs.get("verbose"), "auto")
+
+      mock_fit.reset_mock()
+      mock_evaluate.reset_mock()
+      core.run_experiment(experiment, core.Trainer.Mode.TRAIN_AND_EVAL)
+      mock_fit.assert_called()
+      self.assertEqual(mock_fit.call_args.kwargs.get("verbose"), "auto")
+
+      if keras.backend.backend() == "jax":
+        keras.backend.clear_session()
+        mock_fit.reset_mock()
+        mock_evaluate.reset_mock()
+        core.run_experiment(experiment, core.Trainer.Mode.CONTINUOUS_EVAL)
+        mock_evaluate.assert_called()
+        self.assertEqual(mock_evaluate.call_args.kwargs.get("verbose"), "auto")
+
+  @parameterized.named_parameters(
+      ("set", 2, 2),
+      ("default", None, 1),
+  )
+  def test_steps_per_execution_passed_to_keras(
+      self, configured_value, expected_value
+  ):
+    model_steps_per_execution = []
+    original_fit = keras.Model.fit
+
+    def original_fit_with_record(model, *args, **kwargs):
+      model_steps_per_execution.append(
+          getattr(model, "steps_per_execution", None)
+      )
+      return original_fit(model, *args, **kwargs)
+
+    with mock.patch.object(
+        keras.Model, "fit", autospec=True, side_effect=original_fit_with_record
+    ) as mock_fit:
+      trainer = keras_trainer.KerasTrainer(
+          model_dir=self.create_tempdir().full_path,
+          train_steps=4,
+          steps_per_loop=4,
+          steps_per_eval=1,
+          steps_per_execution=configured_value,
+          distribution=keras.distribution.DataParallel()
+          if keras.backend.backend() == "jax"
+          else None,
+      )
+      experiment = core.Experiment(_KerasTask(), trainer)
+
+      model_steps_per_execution.clear()
+      core.run_experiment(experiment, core.Trainer.Mode.TRAIN)
+      mock_fit.assert_called()
+      self.assertEqual(model_steps_per_execution, [expected_value])
+
+      model_steps_per_execution.clear()
+      core.run_experiment(experiment, core.Trainer.Mode.TRAIN_AND_EVAL)
+      mock_fit.assert_called()
+      self.assertEqual(model_steps_per_execution, [expected_value])
+
+  def test_steps_per_execution_resets_functions(self):
+    model_train_functions = []
+    original_fit = keras.Model.fit
+    task = _KerasTask()
+    original_create_model = task.create_model
+
+    def mock_create_model(*args, **kwargs):
+      model = original_create_model(*args, **kwargs)
+      model.train_function = "stale_train"
+      model.test_function = "stale_test"
+      model.predict_function = "stale_predict"
+      return model
+
+    def original_fit_with_record(model, *args, **kwargs):
+      model_train_functions.append(getattr(model, "train_function", None))
+      return original_fit(model, *args, **kwargs)
+
+    with mock.patch.object(
+        task, "create_model", autospec=True, side_effect=mock_create_model
+    ), mock.patch.object(
+        keras.Model, "fit", autospec=True, side_effect=original_fit_with_record
+    ) as mock_fit:
+      trainer = keras_trainer.KerasTrainer(
+          model_dir=self.create_tempdir().full_path,
+          train_steps=4,
+          steps_per_loop=4,
+          steps_per_eval=1,
+          steps_per_execution=2,
+          distribution=keras.distribution.DataParallel()
+          if keras.backend.backend() == "jax"
+          else None,
+      )
+      experiment = core.Experiment(task, trainer)
+      core.run_experiment(experiment, core.Trainer.Mode.TRAIN)
+      mock_fit.assert_called()
+      self.assertEqual(model_train_functions, [None])
+
+  def test_steps_per_execution_validation(self):
+    with self.assertRaisesRegex(
+        ValueError, r"`steps_per_execution` \(0\) must be at least 1\."
+    ):
+      keras_trainer.KerasTrainer(
+          steps_per_loop=4,
+          steps_per_execution=0,
+      )
+
+    with self.assertRaisesRegex(
+        ValueError,
+        r"`steps_per_loop` \(5\) must be divisible by `steps_per_execution`"
+        r" \(2\)\.",
+    ):
+      keras_trainer.KerasTrainer(
+          steps_per_loop=5,
+          steps_per_execution=2,
+      )
+
   def _run_eval_with_callbacks(self, model_dir: str, mode: core.Trainer.Mode):
     if keras.backend.backend() != "jax":
       self.skipTest("`EpochSummaryCallback` is only used on the Jax backend.")

recml/core/utils/keras_utils.py

@@ -1744,6 +1744,20 @@ def __init__(
     # asynchronously.
     self.async_safe = True
 
+  def on_train_batch_begin(
+      self, batch: int, logs: dict[str, Any] | None = None
+  ):
+    # Skips the per-batch bookkeeping of the inherited hook;
+    # see `on_train_batch_end`.
+    del batch, logs
+
+  def on_train_batch_end(self, batch: int, logs: dict[str, Any] | None = None):
+    # The inherited `TensorBoard.on_train_batch_end` calls `summary.scalar` for
+    # every log on every step, even with `update_freq="epoch"`, when no summary
+    # writer records them. This callback only reports epoch metrics, and epoch
+    # summaries and `steps_per_second` don't use the batch hooks.
+    del batch, logs
+
   def _get_or_create_writer(
       self, key: str, log_dir: str
   ) -> tf.summary.SummaryWriter:

recml/core/utils/keras_utils_test.py

@@ -923,6 +923,23 @@ def test_epoch_orbax_checkpoint_and_restore_callback_saves_at_epoch_0(
     checkpoint_manager.wait_until_finished()
     self.assertTrue(os.path.exists(os.path.join(checkpoint_dir, "0")))
 
+  def test_epoch_summary_callback_batch_hooks_are_noops(self):
+    log_dir = self.create_tempdir().full_path
+    with mock.patch.object(keras_utils.jax, "process_index", return_value=0):
+      callback = keras_utils.EpochSummaryCallback(
+          log_dir=log_dir,
+          steps_per_epoch=10,
+          write_steps_per_second=False,
+      )
+
+    with mock.patch.object(
+        callback.summary, "scalar", autospec=True
+    ) as mock_scalar:
+      logs = {"loss": 0.5}
+      callback.on_train_batch_begin(0, logs=logs)
+      callback.on_train_batch_end(0, logs=logs)
+      mock_scalar.assert_not_called()
+
   def test_epoch_summary_callback_writes_on_process_0(self):
     log_dir = self.create_tempdir().full_path
     dummy_inputs = _create_dummy_inputs()