Function maybe_record_goodput

src/MaxText/utils/goodput_utils.py:74–86

Record goodput if `enable_goodput_recording=True`.

(recorder, event_name, *args)

@contextlib.contextmanager
def maybe_record_goodput(recorder, event_name, *args):
  """Record goodput if `enable_goodput_recording=True`."""
  try:
    start_event_name = f"record_{event_name.value}_start_time"
    record_goodput(recorder, start_event_name, *args)
    yield
  except BaseException:  # pylint: disable=W0706
    raise
  else:
    end_event_name = f"record_{event_name.value}_end_time"
    record_goodput(recorder, end_event_name, *args)
  finally:
    pass


def record_goodput(recorder, event_name, *args):