controller.py

# Copyright 2024 The Orbit Authors. All Rights Reserved.
#
# 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.

"""Provides a `Controller` class for managing the outer training loop."""

import pprint
import time

from typing import Callable, Iterable, Optional, Union

from absl import logging

from orbit import runner
from orbit import utils

import tensorflow as tf, tf_keras

# pylint: disable=g-direct-tensorflow-import
from tensorflow.python.eager import monitoring
# pylint: enable=g-direct-tensorflow-import

_orbit_api_gauge = monitoring.BoolGauge(
    "/tensorflow/api/orbit", "orbit api usage"
)


def _log(message: str):
  """Logs `message` to the `info` log, and also prints to stdout."""
  logging.info(message)
  print(message)


logging.ABSLLogger.register_frame_to_skip(__file__, _log.__name__)


def _format_output(output, indent=4):
  """Formats `output`, either on one line, or indented across multiple lines."""
  formatted = pprint.pformat(output)
  lines = formatted.splitlines()
  if len(lines) == 1:
    return formatted
  lines = [" " * indent + line for line in lines]
  return "\n" + "\n".join(lines)


Action = Callable[[runner.Output], None]


class Controller:
  """Class that controls the outer loop of model training and evaluation.

  Orbit divides training and evaluation into "inner" and "outer" loops. Inner
  loops are implemented by users in the form of `AbstractTrainer` and
  `AbstractEvaluator` subclasses, and define how to run a given number of
  training or evaluation steps. The outer loop is provided by this `Controller`,
  and interleaves calls to the user-provided inner loops with additional actions
  such as saving checkpoints, running evaluations, writing summaries, as well as
  (optionally) user provided `Action`s (see below).

  There are four top-level "outer loops" provided:

    - `train`, which trains until a specified number of global steps is reached;
    - `evaluate`, for one-off model evaluation;
    - `train_and_evaluate`, for interleaved training and evaluation;
    - `evaluate_continuously`, for monitoring a given directory and running
      evaluations on new model checkpoints.

  While this class attempts to provide out-of-the-box solutions for common
  training and evaluation use cases, the internal details and method
  implementations are also intended to be simple enough to make subclassing or
  other custom outer loop implementations easy to achieve.

  Some additional customization can be achieved by supplying `train_actions` or
  `eval_actions` when constructing the `Controller`. Actions arbitrary callables
  that are applied by the `Controller` to the output of train steps (after each
  inner loop of `steps_per_loop` steps) or an evaluation. This provides a hook
  mechanism, enabling things like reporting metrics to Vizier, model exporting,
  additional logging, etc. See the `orbit.actions` package for a small handful
  of predefined actions and some utility classes that may be useful in defining
  your own.
  """

  def __init__(
      self,
      *,  # Makes all args keyword only.
      global_step: tf.Variable,
      trainer: Optional[runner.AbstractTrainer] = None,
      evaluator: Optional[runner.AbstractEvaluator] = None,
      strategy: Optional[tf.distribute.Strategy] = None,
      # Actions
      train_actions: Optional[Iterable[Action]] = None,
      eval_actions: Optional[Iterable[Action]] = None,
      # Train related
      steps_per_loop: Optional[Union[int, Callable[[int], int]]] = None,
      checkpoint_manager: Optional[tf.train.CheckpointManager] = None,
      enable_async_checkpointing: bool = False,
      # Summary related
      summary_interval: Optional[int] = None,
      summary_dir: Optional[str] = None,
      # Evaluation related
      eval_summary_dir: Optional[str] = None,
      summary_manager: Optional[utils.SummaryManagerInterface] = None,
      eval_summary_manager: Optional[utils.SummaryManagerInterface] = None):
    """Initializes a `Controller` instance.

    Note that if `checkpoint_manager` is provided and there are checkpoints in
    the associated model directory, the model will be restored from the most
    recent checkpoint during this `__init__` method.

    Args:
      global_step: An integer `tf.Variable` storing the global training step
        number. Usually this can be obtained from the `iterations` property of
        the model's optimizer (e.g. `trainer.optimizer.iterations`). In cases
        where multiple optimizers are used, or if one model "step" corresponds
        to more than one update to model parameters, users can create and
        increment their own global step variable as well. In this case it is
        recommended to create the `tf.Variable` inside the distribution strategy
        scope, with `aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA` (see
        also `orbit.utils.create_global_step()`).
      trainer: An instance of `orbit.AbstractTrainer`, which implements the
        inner training loop.
      evaluator: An instance of `orbit.AbstractEvaluator`, which implements
        evaluation.
      strategy: An instance of `tf.distribute.Strategy`. If not provided, the
        strategy will be initialized from the current in-scope strategy using
        `tf.distribute.get_strategy()`.
      train_actions: Optional `orbit.Action`s to call after each block of
        `steps_per_loop` training steps are run. These will be called with the
        output of `trainer.train`.
      eval_actions: Optional `orbit.Action`s to call after each evaluation.
        These will be called with the output of `evaluator.evaluate`.
      steps_per_loop: Optional integer to indicate the number of steps to run in
        each inner loop of training (passed as the `num_steps` parameter of
        `trainer.train`). It can be also a callable which takes the current
        global step value as input and returns the number of steps to run as
        output.
      checkpoint_manager: An instance of `tf.train.CheckpointManager`. If
        provided and there are checkpoints in the associated model directory,
        the model will be restored from the most recent checkpoint inside this
        `__init__` method. If not provided, the `Controller` will not
        automatically save to or restore from checkpoints.
      enable_async_checkpointing: Optional bool indicating whether to enable
        async checkpoint saving.
      summary_interval: Step interval for training summaries. Note that this
        argument only applies to `tf.summary` calls inside the `trainer.train`
        function. Summaries written by the `Controller` (specifically
        "steps_per_second" and output from the `trainer.train` method) will
        always be enabled unless the `summary_dir` parameter is `None`. If set,
        the value must be divisible by `steps_per_loop`.
      summary_dir: The directory to write summaries to. To use the same
        directory as for checkpointing, pass `checkpoint_manager.directory`. If
        `None`, no training summaries will be written.
      eval_summary_dir: The directory to write eval summaries to. If `None`, it
        will be set to `summary_dir`. If both `summary_dir` and
        `eval_summary_dir` are `None`, no eval summaries will be written.
      summary_manager: Instance of the summary manager. If set, the
        `summary_dir` will be ignored. Otherwise the summary manager will be
        created internally for TensorBoard summaries by default from the
        `summary_dir`.
      eval_summary_manager: Instance of the eval summary manager. If set, the
        `eval_summary_dir` will be ignored. Otherwise the eval summary manager
        will be created internally for TensorBoard summaries by default from the
        `eval_summary_dir`.

    Raises:
      ValueError: If both `trainer` and `evaluator` are `None`.
      ValueError: If `steps_per_loop` is not a positive integer or a callable.
      ValueError: If `summary_interval` is not a positive integer or is not
        divisible by `steps_per_loop`.
    """
    if trainer is None and evaluator is None:
      raise ValueError("`trainer` and `evaluator` should not both be `None`.")

    if trainer is not None:
      if steps_per_loop is None:
        raise ValueError(
            "`steps_per_loop` is required when `trainer` is provided.")
      elif not callable(steps_per_loop) and (
          not isinstance(steps_per_loop, int) or steps_per_loop < 1):
        raise ValueError(
            f"`steps_per_loop` ({steps_per_loop}) must be a positive integer "
            "or a callable.")

      if summary_interval is not None:
        if summary_interval <= 0:
          raise ValueError(
              f"`summary_interval` ({summary_interval}) must be larger than 0.")
        elif not callable(steps_per_loop) and (summary_interval % steps_per_loop
                                               != 0):
          raise ValueError(
              f"`summary interval` ({summary_interval}) must be a multiple "
              f"of `steps_per_loop` ({steps_per_loop}).")

    if not isinstance(global_step, tf.Variable):
      raise ValueError("`global_step` must be a `tf.Variable`.")

    self.trainer = trainer
    self.evaluator = evaluator

    self.strategy = strategy or tf.distribute.get_strategy()

    self.train_actions = () if train_actions is None else tuple(train_actions)
    self.eval_actions = () if eval_actions is None else tuple(eval_actions)

    self.global_step = global_step
    self.checkpoint_manager = checkpoint_manager
    self._enable_async_checkpoint_saving = enable_async_checkpointing
    self._checkpoint_options = tf.train.CheckpointOptions(
        enable_async=enable_async_checkpointing
    )

    if self.trainer is not None:
      self.step_timer = None
      self.summary_interval = summary_interval
      if summary_manager:
        self.summary_manager = summary_manager
      else:
        self.summary_manager = utils.SummaryManager(
            summary_dir, tf.summary.scalar, global_step=self.global_step)
      self._steps_per_loop = steps_per_loop

    if self.evaluator is not None:
      eval_summary_dir = eval_summary_dir or summary_dir
      if eval_summary_dir == summary_dir and self.trainer is not None:
        # Reuse the summary writer if train and evaluation summary directory
        # are the same.
        self.eval_summary_manager = self.summary_manager
      else:
        if eval_summary_manager:
          self.eval_summary_manager = eval_summary_manager
        else:
          self.eval_summary_manager = utils.SummaryManager(
              eval_summary_dir, tf.summary.scalar, global_step=self.global_step)

    tf.summary.experimental.set_step(self.global_step)

    # Restores the model if needed.
    if self.checkpoint_manager is not None:
      restored_path = self.restore_checkpoint()
      if restored_path:
        _log(f"restored from checkpoint: {restored_path}")

    # Set Orbit framework gauge to True value
    _orbit_api_gauge.get_cell().set(True)

  def train(self, steps: int, checkpoint_at_completion: bool = True):
    """Runs training until the specified global step count has been reached.

    This method makes calls to `self.trainer.train()` until the global step
    count is equal to `steps`. It will additionally save checkpoints (if a
    `CheckpointManager` was passed to `Controller.__init__`) and summarize
    training output (if `summary_dir` is set).

    When async checkpointing is enabled, a sync is triggered at the end of this
    method to make sure any ongoing async checkpoint saving is finished before
    returning.

    Args:
      steps: The global step count to train up to.
      checkpoint_at_completion: Whether to save a checkpoint when this method
        returns (regardless of the checkpointing interval). Defaults to `True`.
    """
    self._require("trainer", for_method="train")

    # TODO(momernick): Support steps=None or -1 (training to exhaustion).
    current_step = self.global_step.numpy()  # Cache, since this is expensive.
    _log(f"train | step: {current_step: 6d} | training until step {steps}...")
    while current_step < steps:
      # Calculates steps to run for the next train loop.
      num_steps = min(steps - current_step, self.steps_per_loop)
      self._train_n_steps(num_steps)
      self._maybe_save_checkpoint()
      current_step = self.global_step.numpy()

    if checkpoint_at_completion:
      self._maybe_save_checkpoint(check_interval=False)

    self._sync_on_async_checkpointing()

  def evaluate(self, steps: int = -1) -> Optional[runner.Output]:
    """Runs evaluation for the given number of steps.

    This method calls `self.evaluator.evaluate(steps)`, then writes the returned
    summaries (if any).

    Args:
      steps: The number of evaluation steps to run. The value `-1` is reserved
        as a special sentinel to indicate a "complete" evaluation that runs
        until the underlying dataset is exhausted. Support for this is dependent
        on the specific `evaluator` being used.

    Returns:
      The evaluation results as a dictionary mapping names to NumPy values.

    Raises:
      ValueError: If `evaluator` was not provided to `Controller.__init__`.
      ValueError: If no checkpoint is present in `checkpoint_manager.directory`.
      ValueError: If `steps` is not a positive value or -1.
    """
    self._require("evaluator", for_method="evaluate")

    if steps > 0:
      steps_msg = f"running {steps} steps of evaluation..."
    elif steps == -1:
      steps_msg = "running complete evaluation..."
    else:
      raise ValueError(f"`steps` ({steps}) should be > 0, or == -1.")

    current_step = self.global_step.numpy()
    _log(f" eval | step: {current_step: 6d} | {steps_msg}")

    start = time.time()
    assert isinstance(self.evaluator, runner.AbstractEvaluator)
    with self.eval_summary_manager.summary_writer().as_default():
      steps_tensor = tf.convert_to_tensor(steps, dtype=tf.int32)
      eval_output = self.evaluator.evaluate(steps_tensor)
    elapsed = time.time() - start

    eval_output = eval_output or {}
    for action in self.eval_actions:
      action(eval_output)
    eval_output = tf.nest.map_structure(utils.get_value, eval_output)

    if steps > 0:
      # Only log if steps has been specified.
      steps_per_second = steps / elapsed
      eval_output["steps_per_second"] = steps_per_second
      steps_per_second_log = f"steps/sec: {steps_per_second: 6.1f} | "
    else:
      steps_per_second_log = ""

    _log(f" eval | step: {current_step: 6d} | "
         f"{steps_per_second_log}"
         f"eval time: {elapsed: 6.1f} sec | "
         f"output: {_format_output(eval_output)}")

    self.eval_summary_manager.write_summaries(eval_output)
    self.eval_summary_manager.flush()

    return eval_output

  def train_and_evaluate(
      self,
      train_steps: int,
      eval_steps: int = -1,
      eval_interval: Optional[int] = None,
  ) -> Optional[runner.Output]:
    """Runs interleaved training and evaluation.

    This method interleaves calls to `self.train()` and `self.evaluate()`,
    training the model until the global step count equals `train_steps`, and
    running an evaluation for `eval_steps` every `eval_interval` training steps.
    In addition, this method will run a final evaluation at the end of the
    training sequence.

    When async checkpointing is enabled, a sync is triggered at the end of this
    method to make sure any ongoing async checkpoint saving is finished before
    returning.

    Args:
      train_steps: The global step count to train up to.
      eval_steps: The number of steps to run during an evaluation. If -1, this
        method will evaluate over the entire evaluation dataset.
      eval_interval: The number of training steps to run between evaluations. If
        set, training will always stop every `eval_interval` steps, even if this
        results in a shorter inner loop than specified by `steps_per_loop`
        setting. If None, evaluation will only be performed after training is
        complete.

    Returns:
      The evaluation results as a dictionary mapping names to NumPy values.
    """
    self._require("trainer", for_method="train_and_evaluate")
    self._require("evaluator", for_method="train_and_evaluate")

    output = None
    current_step = self.global_step.numpy()  # Cache, since this is expensive.
    eval_interval = eval_interval or (train_steps - current_step)
    while current_step < train_steps:
      interval = min(train_steps - current_step, eval_interval)
      num_steps = current_step + interval
      self.train(steps=num_steps, checkpoint_at_completion=False)
      output = self.evaluate(steps=eval_steps)
      current_step = self.global_step.numpy()
    self._maybe_save_checkpoint(check_interval=False)
    self._sync_on_async_checkpointing()
    return output

  def evaluate_continuously(
      self,
      steps: int = -1,
      timeout: Optional[Union[int, float]] = None,
      timeout_fn: Optional[Callable[[], bool]] = None,
  ) -> Optional[runner.Output]:
    """Continuously monitors a directory and evaluates new checkpoints in it.

    This method continuously monitors a directory as specified by this
    Controller's CheckpointManager init arg and runs evaluation on the
    checkpoints found there.

    Args:
      steps: The number of steps to run when evaluating. If -1, this method will
        evaluate over the entire evaluation dataset.
      timeout: The maximum number of seconds to wait between checkpoints. See
        tf.train.checkpoints_iterator documentation.
      timeout_fn: Optional callable to call after a timeout. If the function
        returns True, then it means that no new checkpoints will be generated
        and the iterator will exit.

    Returns:
      The evaluation results as a dictionary mapping names to NumPy values.

    Raises:
      ValueError: If no checkpoint found in `self.checkpoint_manager.directory`.
      ValueError: If `evaluator` was not provided as a controller init arg.
    """
    self._require("evaluator", for_method="evaluate_continuously")
    self._require("checkpoint_manager", for_method="evaluate_continuously")

    output = None
    assert isinstance(self.checkpoint_manager, tf.train.CheckpointManager)
    for checkpoint_path in tf.train.checkpoints_iterator(
        self.checkpoint_manager.directory,
        timeout=timeout,
        timeout_fn=timeout_fn):
      self.restore_checkpoint(checkpoint_path)
      output = self.evaluate(steps)
    return output

  def restore_checkpoint(self, checkpoint_path: Optional[str] = None):
    """Restores the model from a checkpoint.

    Args:
      checkpoint_path: An optional string specifying the checkpoint path to
        restore from. If `None`, will restore from the most recent checkpoint
        (or initialize the model using a custom `init_fn` if no checkpoints can
        be found) using `self.checkpoint_manager.restore_or_initialize()`.

    Returns:
      The path to the restored checkpoint if a restore happened, or `None` if no
      restore occurred.
    """
    self._require("checkpoint_manager", for_method="restore_checkpoint")

    assert isinstance(self.checkpoint_manager, tf.train.CheckpointManager)
    with self.strategy.scope():
      # Checkpoint restoring should be inside scope (b/139450638).
      if checkpoint_path is not None:
        _log(f"restoring model from {checkpoint_path}...")
        self.checkpoint_manager.checkpoint.restore(checkpoint_path)
      else:
        _log("restoring or initializing model...")
        checkpoint_path = self.checkpoint_manager.restore_or_initialize()

    if checkpoint_path is not None:
      _log(f"restored model from {checkpoint_path}.")

    return checkpoint_path

  def save_checkpoint(self):
    """Saves the model to a checkpoint.

    This method will save a checkpoint containing the current state of the
    model.

    Raises:
      ValueError: If no `checkpoint_manager` was provided to
        `Controller.__init__`.
    """
    self._require("checkpoint_manager", for_method="save_checkpoint")
    self._maybe_save_checkpoint(check_interval=False)

  @property
  def steps_per_loop(self):
    """Returns current steps_per_loop value in a training loop."""
    if callable(self._steps_per_loop):
      return self._steps_per_loop(self.global_step.numpy())
    return self._steps_per_loop

  def _train_n_steps(self, num_steps: int):
    """Runs training for `num_steps` steps.

    Also prints/logs updates about training progress, and summarizes training
    output (if output is returned from `self.trainer.train()`, and if
    `self.summary_dir` is set).

    Args:
      num_steps: An integer specifying how many steps of training to run.

    Raises:
      RuntimeError: If `global_step` is not properly incremented by `num_steps`
        after calling `self.trainer.train(num_steps)`.
    """
    if not self.step_timer:
      self.step_timer = StepTimer(self.global_step)
    current_step = self.global_step.numpy()

    with self.summary_manager.summary_writer().as_default():
      should_record = False  # Allows static optimization in no-summary cases.
      if self.summary_interval:
        # Create a predicate to determine when summaries should be written.
        should_record = lambda: (self.global_step % self.summary_interval == 0)
      assert isinstance(self.trainer, runner.AbstractTrainer)
      with tf.summary.record_if(should_record):
        num_steps_tensor = tf.convert_to_tensor(num_steps, dtype=tf.int32)
        train_output = self.trainer.train(num_steps_tensor)

    # Verify that global_step was updated properly, then update current_step.
    expected_step = current_step + num_steps
    if self.global_step.numpy() != expected_step:
      message = (
          f"`trainer.train({num_steps})` did not update `global_step` by "
          f"{num_steps}. Old value was {current_step}, expected updated value "
          f"to be {expected_step}, but it was {self.global_step.numpy()}.")
      logging.warning(message)

    train_output = train_output or {}
    for action in self.train_actions:
      action(train_output)
    train_output = tf.nest.map_structure(utils.get_value, train_output)

    current_step = self.global_step.numpy()
    steps_per_second = self.step_timer.steps_per_second()
    _log(f"train | step: {current_step: 6d} | "
         f"steps/sec: {steps_per_second: 6.1f} | "
         f"output: {_format_output(train_output)}")

    train_output["steps_per_second"] = steps_per_second
    self.summary_manager.write_summaries(train_output)
    self.summary_manager.flush()

  def _maybe_save_checkpoint(self, check_interval: bool = True):
    """Conditionally saves a checkpoint.

    A checkpoint is saved if a `CheckpointManager` is available, and if the
    required number of steps has elapsed since the last checkpoint was saved
    (although this condition can be disabled by setting `check_interval=False`).

    Args:
      check_interval: Whether to check if the checkpoint interval has fully
        elapsed. If `False`, a checkpoint is saved regardless of the elapsed
        steps since the most recent checkpoint, unless no `checkpoint_manager`
        was provided to `Controller.__init__`.

    Returns:
      A boolean indicating whether a checkpoint was saved.
    """
    if self.checkpoint_manager and self.checkpoint_manager.checkpoint_interval:
      ckpt_path = self.checkpoint_manager.save(
          checkpoint_number=self.global_step.numpy(),
          check_interval=check_interval,
          options=self._checkpoint_options)
      if ckpt_path is not None:
        _log(f"saved checkpoint to {ckpt_path}.")
        return True
    return False

  def _require(self, attribute, for_method):
    """Utility method to raise an error if the given `attribute` is not set."""
    if getattr(self, attribute, None) is None:
      raise ValueError(
          f"`{attribute}` is not set. Pass `{attribute}` to "
          f"`Controller.__init__` before calling `{for_method}()`.")

  def _sync_on_async_checkpointing(self):
    """Force to wait for the async checkpoint saving (if any) to finish."""
    # pylint: disable=protected-access
    if self.checkpoint_manager:
      logging.info("Sync on async checkpoint saving.")
      self.checkpoint_manager.sync()


class StepTimer:
  """Utility class for measuring steps/second."""

  def __init__(self, step):
    self.step = step
    self.start()

  def start(self):
    self.last_iteration = self.step.numpy()
    self.last_time = time.time()

  def steps_per_second(self, restart=True):
    value = ((self.step.numpy() - self.last_iteration) /
             (time.time() - self.last_time))
    if restart:
      self.start()
    return value