code: def benchmark_context(flag_obj):
      """Context of benchmark, which will update status of the run accordingly."""
      benchmark_logger = config_benchmark_logger(flag_obj)
      try:
          yield
          benchmark_logger.on_finish(RUN_STATUS_SUCCESS)
      except Exception:  # pylint: disable=broad-except
          # Catch all the exception, update the run status to be failure, and re-raise
          benchmark_logger.on_finish(RUN_STATUS_FAILURE)
          raise
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import contextlib
  import datetime
  import json
  import multiprocessing
  import numbers
  import os
  import threading
  import uuid
  from six.moves import _thread
  from absl import flags
  import tensorflow
  from tensorflow.python.client import device_lib
  from utils.logs import cloud_lib
label_desc: Use image node_id 3 for calling a global function with example usage: benchmark_context(flag_obj) without return types