Internal change

recml/core/training/keras_trainer.py

@@ -16,7 +16,7 @@
 from __future__ import annotations
 
 import abc
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 import dataclasses
 import functools
 import gc
@@ -124,6 +124,8 @@ def __init__(
       rng_seed: int = core.DEFAULT_RNG_SEED,
       checkpoint_version: keras_utils.CheckpointVersion | str | None = None,
       legacy_checkpoint_format: bool | None = None,
+      train_callbacks: Sequence[keras.callbacks.Callback] = (),
+      eval_callbacks: Sequence[keras.callbacks.Callback] = (),
   ):
     """Initializes the instance.
 
@@ -143,6 +145,13 @@ def __init__(
       legacy_checkpoint_format: Deprecated. Use checkpoint_version instead. If
         set, True maps to V1, and False maps to V2. TODO(b/542602169): Remove
         this in v2.
+      train_callbacks: Additional callbacks for training. They run before the
+        trainer's own callbacks, so metrics that they add to `logs` are written
+        to TensorBoard.
+      eval_callbacks: Additional callbacks for evaluation. They run before the
+        trainer's own callbacks. Use `keras_utils.MetricsCallback` to report
+        metrics, because Keras does not return changes made to `logs` in
+        `on_test_end`.
       resume_training_launcher: Launcher to resume training.
       enable_xmanager_measurements: Whether to enable XManager measurements.
       enable_autoxprof: Whether to enable AutoXprof.
@@ -177,6 +186,8 @@ def __init__(
     self._checkpoint_dir = os.path.join(model_dir, core.CHECKPOINT_DIR)
     self._max_checkpoints_to_keep = max_checkpoints_to_keep
     self._checkpoint_save_interval_epochs = checkpoint_save_interval_epochs
+    self._custom_train_callbacks = list(train_callbacks)
+    self._custom_eval_callbacks = list(eval_callbacks)
     if checkpoint_version is not None:
       self._checkpoint_version = keras_utils.CheckpointVersion(
           checkpoint_version
@@ -216,6 +227,7 @@ def train_callbacks(self) -> list[keras.callbacks.Callback]:
       )
 
       callbacks = [
+          *self._custom_train_callbacks,
           keras_utils.EpochSummaryCallback(
               log_dir=os.path.join(self._model_dir, core.LOG_DIR),
               steps_per_epoch=self._steps_per_loop,
@@ -230,6 +242,7 @@ def train_callbacks(self) -> list[keras.callbacks.Callback]:
       return callbacks  # pyrefly: ignore[bad-return]
 
     return [
+        *self._custom_train_callbacks,
         keras.callbacks.TensorBoard(
             log_dir=os.path.join(self._model_dir, core.LOG_DIR),
             write_steps_per_second=True,
@@ -253,6 +266,7 @@ def eval_callbacks(self) -> list[keras.callbacks.Callback]:
     """Returns the evaluation callbacks."""
     if keras.backend.backend() == "jax":
       return [
+          *self._custom_eval_callbacks,
           keras_utils.EpochSummaryCallback(
               log_dir=os.path.join(self._model_dir, core.LOG_DIR),
               steps_per_epoch=self._steps_per_loop,
@@ -266,6 +280,7 @@ def eval_callbacks(self) -> list[keras.callbacks.Callback]:
           ),
       ]
     return [
+        *self._custom_eval_callbacks,
         keras.callbacks.TensorBoard(
             log_dir=os.path.join(self._model_dir, core.LOG_DIR),
             write_steps_per_second=True,
@@ -307,6 +322,27 @@ def _maybe_get_model_kws(
 
     return kws
 
+  def _val_logs(
+      self,
+      history: Mapping[str, Any],
+      callbacks: Sequence[keras.callbacks.Callback],
+      steps_per_second: float,
+  ) -> dict[str, Any]:
+    """Returns the logs for `EpochSummaryCallback` after an evaluation."""
+    val_logs = {f"val_{k}": v for k, v in history.items()}
+    for cbk in callbacks:
+      if isinstance(cbk, keras_utils.MetricsCallback):
+        for k, v in cbk.latest_metrics.items():
+          val_key = f"val_{k}"
+          if val_key in val_logs:
+            logging.warning(
+                f"{type(cbk).__name__} overwrites existing evaluation log key"
+                f" {val_key}."
+            )
+          val_logs[val_key] = v
+    val_logs["val_steps_per_second"] = steps_per_second
+    return val_logs
+
   def train(self, task: KerasTask) -> core.Logs:
     """Trains a Keras model."""
     dataset = task.create_dataset(
@@ -354,9 +390,9 @@ def evaluate(self, task: KerasTask) -> core.Logs:
       )
       epoch_dt = time.time() - epoch_start_time
       steps_per_second = self._steps_per_eval / epoch_dt  # pyrefly: ignore[unsupported-operation]
-      val_logs = {"val_" + k: v for k, v in history.items()}
-      val_logs["val_steps_per_second"] = steps_per_second
-      tb_cbk.on_epoch_end(0, val_logs)
+      tb_cbk.on_epoch_end(
+          0, self._val_logs(history, self.eval_callbacks, steps_per_second)
+      )
       return history
 
     return model.evaluate(
@@ -464,6 +500,7 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None):
           for cbk in self.eval_callbacks
           if isinstance(cbk, keras_utils.EpochSummaryCallback)
       ]
+      callbacks = [restore_callback, *self.eval_callbacks]
       try:
         logging.info(f"eval | epoch: {epoch: 6d} | {steps_msg}")
         epoch_start_time = time.time()
@@ -475,7 +512,7 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None):
         history = model.evaluate(
             eval_dataset,
             steps=self._steps_per_eval,
-            callbacks=[restore_callback] + self.eval_callbacks,
+            callbacks=callbacks,
             return_dict=True,
         )
 
@@ -491,9 +528,10 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None):
         epoch_dt = time.time() - epoch_start_time
         steps_per_second = self._steps_per_eval / epoch_dt  # pyrefly: ignore[unsupported-operation]
 
-        val_logs = {"val_" + k: v for k, v in history.items()}
-        val_logs["val_steps_per_second"] = steps_per_second
-        tb_cbk.on_epoch_end(epoch, val_logs)
+        tb_cbk.on_epoch_end(
+            epoch,
+            self._val_logs(history, self.eval_callbacks, steps_per_second),
+        )
       except FileNotFoundError:
         logging.info("Checkpoint epoch: %s did not finish writing...", epoch)
 

recml/core/training/keras_trainer_test.py

@@ -52,6 +52,16 @@ def create_model(self) -> keras.Model:
     return model
 
 
+class _MarkerCallback(keras_utils.MetricsCallback):
+  """Reports a fixed metric at the end of each training epoch and eval pass."""
+
+  def on_epoch_end(self, epoch, logs=None):
+    self._record(logs, "custom/marker", 7.0)
+
+  def on_test_end(self, logs=None):
+    self._record(logs, "custom/marker", 7.0)
+
+
 class KerasTrainerTest(parameterized.TestCase):
 
   def setUp(self):
@@ -63,12 +73,22 @@ def setUp(self):
 
   @parameterized.named_parameters(
       {"testcase_name": "train", "mode": core.Trainer.Mode.TRAIN},
+      {
+          "testcase_name": "train_with_callbacks",
+          "mode": core.Trainer.Mode.TRAIN,
+          "with_train_callbacks": True,
+      },
       {"testcase_name": "eval", "mode": core.Trainer.Mode.EVAL},
       {
           "testcase_name": "train_and_eval_v3",
           "mode": core.Trainer.Mode.TRAIN_AND_EVAL,
           "checkpoint_version": "v3",
       },
+      {
+          "testcase_name": "train_and_eval_with_callbacks",
+          "mode": core.Trainer.Mode.TRAIN_AND_EVAL,
+          "with_train_callbacks": True,
+      },
       {
           "testcase_name": "train_and_eval_v2",
           "mode": core.Trainer.Mode.TRAIN_AND_EVAL,
@@ -96,7 +116,10 @@ def setUp(self):
       },
   )
   def test_keras_task_and_trainer(
-      self, mode: str, checkpoint_version: str = "v3"
+      self,
+      mode: str,
+      checkpoint_version: str = "v3",
+      with_train_callbacks: bool = False,
   ):
     if keras.backend.backend() == "jax":
       distribution = keras.distribution.DataParallel()
@@ -113,6 +136,7 @@ def test_keras_task_and_trainer(
         model_dir=self.create_tempdir().full_path,
         continuous_eval_timeout=5,
         checkpoint_version=checkpoint_version,
+        train_callbacks=[_MarkerCallback()] if with_train_callbacks else (),
     )
     experiment = core.Experiment(_KerasTask(), trainer)
 
@@ -128,6 +152,72 @@ def test_keras_task_and_trainer(
         and keras.backend.backend() == "jax"
     ):
       self.assertEqual(history.history["num_params/trainable"][0], 2)
+    if with_train_callbacks:
+      self.assertEqual(
+          history.history["custom/marker"],
+          [7.0] * len(history.history["loss"]),
+      )
+
+  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.")
+    trainer = keras_trainer.KerasTrainer(
+        distribution=keras.distribution.DataParallel(),
+        train_steps=2,
+        steps_per_eval=1,
+        steps_per_loop=2,
+        model_dir=model_dir,
+        continuous_eval_timeout=1,
+        eval_callbacks=[_MarkerCallback()],
+    )
+    with mock.patch.object(
+        keras_utils.EpochSummaryCallback, "on_epoch_end", autospec=True
+    ) as on_epoch_end:
+      core.run_experiment(core.Experiment(_KerasTask(), trainer), mode)
+
+    val_logs = on_epoch_end.call_args.args[2]
+    self.assertEqual(val_logs["val_custom/marker"], 7.0)
+    self.assertIn("val_loss", val_logs)
+    self.assertIn("val_steps_per_second", val_logs)
+
+  def test_eval_metrics_callbacks_are_written_in_eval(self):
+    self._run_eval_with_callbacks(
+        self.create_tempdir().full_path, core.Trainer.Mode.EVAL
+    )
+
+  def test_eval_metrics_callbacks_are_written_in_continuous_eval(self):
+    if keras.backend.backend() != "jax":
+      self.skipTest("`EpochSummaryCallback` is only used on the Jax backend.")
+    model_dir = self.create_tempdir().full_path
+    # Produce one checkpoint so there is something to evaluate.
+    train_trainer = keras_trainer.KerasTrainer(
+        distribution=keras.distribution.DataParallel(),
+        train_steps=2,
+        steps_per_eval=1,
+        steps_per_loop=2,
+        model_dir=model_dir,
+    )
+    core.run_experiment(
+        core.Experiment(_KerasTask(), train_trainer), core.Trainer.Mode.TRAIN
+    )
+    keras.backend.clear_session()
+
+    self._run_eval_with_callbacks(model_dir, core.Trainer.Mode.CONTINUOUS_EVAL)
+
+  def test_val_logs_warns_when_metrics_callback_overwrites_key(self):
+    trainer = keras_trainer.KerasTrainer(
+        model_dir=self.create_tempdir().full_path
+    )
+    callback = _MarkerCallback()
+    callback.on_test_end()
+
+    with self.assertLogs(level="WARNING") as log:
+      val_logs = trainer._val_logs(
+          {"custom/marker": 1.0}, [callback], steps_per_second=2.0
+      )
+
+    self.assertIn("overwrites existing evaluation log key", log.output[0])
+    self.assertEqual(val_logs["val_custom/marker"], 7.0)
 
   def test_eval_name(self):
     if keras.backend.backend() != "jax":