def create_dir_and_save_flags(flags_module, log_dir, json_filename):
  """Creates log directory and saves flags to a JSON file.

  Args:
    flags_module: An absl.flags module.
    log_dir: A string for log directory.
    json_filename: A string for output JSON file name.
  """
  # Create log directory if necessary.
  if not tf.io.gfile.exists(log_dir):
    tf.io.gfile.makedirs(log_dir)

  # Save all key flags.
  key_flag_dict = {
      flag.name: flag.value
      for flag in flags_module.FLAGS.get_key_flags_for_module(sys.argv[0])
  }
  json_path = os.path.join(log_dir, json_filename)
  with tf.io.gfile.GFile(json_path, 'w') as f:
    json.dump(key_flag_dict, f, indent=2, sort_keys=True)