diff --git a/recml/core/utils/keras_utils.py b/recml/core/utils/keras_utils.py index dbf8a38..e271edc 100644 --- a/recml/core/utils/keras_utils.py +++ b/recml/core/utils/keras_utils.py @@ -1694,6 +1694,12 @@ def restore_keras_model( model.optimizer.iterations.assign(0) +# One process writes TensorBoard event files and XManager measurements so +# hosts don't issue concurrent CNS flushes. +_SUMMARY_PROCESS_INDEX = 0 +_SUMMARY_MAX_QUEUE = 1000 # Fits a full epoch of scalars (~300 per writer). + + # TODO(b/343544467): Support logging metrics more frequently. class EpochSummaryCallback(keras.callbacks.TensorBoard): """A custom summary callback that only reports epoch metrics.""" @@ -1711,6 +1717,7 @@ def __init__( update_freq="epoch", write_graph=False, ) + self._is_summary_writer = jax.process_index() == _SUMMARY_PROCESS_INDEX self._steps_per_epoch = steps_per_epoch self._num_params = None self._eval_subdir = eval_subdir @@ -1718,6 +1725,30 @@ def __init__( # asynchronously. self.async_safe = True + def _get_or_create_writer( + self, key: str, log_dir: str + ) -> tf.summary.SummaryWriter: + # The base `TensorBoard.on_train_begin` / `on_test_begin` read + # `_train_writer` and `_val_writer` even with `update_freq="epoch"`, so + # returning no-op writers here keeps other processes from creating event + # files. + if key not in self._writers: + if self._is_summary_writer: + self._writers[key] = self.summary.create_file_writer( + log_dir, max_queue=_SUMMARY_MAX_QUEUE + ) + else: + self._writers[key] = self.summary.create_noop_writer() + return self._writers[key] + + @property + def _train_writer(self) -> tf.summary.SummaryWriter: + return self._get_or_create_writer("train", self._train_dir) + + @property + def _val_writer(self) -> tf.summary.SummaryWriter: + return self._get_or_create_writer("val", self._val_dir) + def set_model(self, model: keras.Model): """Sets Keras model and writes graph if specified.""" super().set_model(model) diff --git a/recml/core/utils/keras_utils_test.py b/recml/core/utils/keras_utils_test.py index 9900cc9..88ae6bb 100644 --- a/recml/core/utils/keras_utils_test.py +++ b/recml/core/utils/keras_utils_test.py @@ -31,6 +31,7 @@ import numpy as np import orbax.checkpoint as ocp from recml.core.utils import keras_utils +import tensorflow as tf def _create_dummy_inputs() -> dict[str, jax.Array]: @@ -890,6 +891,115 @@ 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_writes_on_process_0(self): + log_dir = self.create_tempdir().full_path + dummy_inputs = _create_dummy_inputs() + model = _create_model(jax.tree.map(jnp.shape, dummy_inputs)) + 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, + ) + callback.set_model(model) + with mock.patch.object( + callback.summary, + "create_file_writer", + wraps=callback.summary.create_file_writer, + ) as mock_create_file_writer: + callback.on_train_begin() + logs = {f"metric_{i}": float(i) for i in range(15)} + logs["val_loss"] = 0.4 + callback.on_epoch_end(epoch=1, logs=logs) + + self.assertIn("num_params/total", logs) + self.assertIn("val_num_params/total", logs) + train_dir = os.path.join(log_dir, "train") + val_dir = os.path.join(log_dir, "validation") + mock_create_file_writer.assert_has_calls([ + mock.call(train_dir, max_queue=1000), + mock.call(val_dir, max_queue=1000), + ]) + + # Leaving `with writer.as_default():` flushes events before + # `on_train_end` (needed by `evaluate_continuously`). + [train_event_file] = os.listdir(train_dir) + train_events = list( + tf.compat.v1.train.summary_iterator( + os.path.join(train_dir, train_event_file) + ) + ) + expected_train_tags = {f"metric_{i}" for i in range(15)} | { + "learning_rate", + "num_params/trainable", + "num_params/non_trainable", + "num_params/optimizer", + "num_params/total", + } + self.assertEqual( + { + value.tag + for event in train_events + for value in event.summary.value + }, + expected_train_tags, + ) + self.assertEqual( + {event.step for event in train_events if event.summary.value}, + {10}, + ) + + [val_event_file] = os.listdir(val_dir) + val_events = list( + tf.compat.v1.train.summary_iterator( + os.path.join(val_dir, val_event_file) + ) + ) + self.assertEqual( + {value.tag for event in val_events for value in event.summary.value}, + { + "loss", + "num_params/trainable", + "num_params/non_trainable", + "num_params/optimizer", + "num_params/total", + }, + ) + self.assertEqual( + {event.step for event in val_events if event.summary.value}, + {10}, + ) + callback.on_train_end() + + def test_epoch_summary_callback_skips_writes_on_other_processes(self): + log_dir = self.create_tempdir().full_path + dummy_inputs = _create_dummy_inputs() + model = _create_model(jax.tree.map(jnp.shape, dummy_inputs)) + with mock.patch.object(keras_utils.jax, "process_index", return_value=1): + callback = keras_utils.EpochSummaryCallback( + log_dir=log_dir, + steps_per_epoch=10, + write_steps_per_second=False, + ) + callback.set_model(model) + with mock.patch.object( + callback.summary, + "create_file_writer", + wraps=callback.summary.create_file_writer, + ) as mock_create_file_writer: + callback.on_train_begin() + logs = {f"metric_{i}": float(i) for i in range(15)} + logs["val_loss"] = 0.4 + callback.on_epoch_end(epoch=1, logs=logs) + + # `num_params` must still be populated in `logs` on all processes. + self.assertIn("num_params/total", logs) + self.assertIn("val_num_params/total", logs) + mock_create_file_writer.assert_not_called() + self.assertFalse(os.path.exists(os.path.join(log_dir, "train"))) + self.assertFalse(os.path.exists(os.path.join(log_dir, "validation"))) + callback.on_train_end() + class KerasOrbaxCheckpointUtilsTest(absltest.TestCase):