code: def print_details_info(timestamps, model_config, args, whole_end2end):
      if dist.get_rank() == 0:
          prefill = []
          encoder = []
          end2end = []
          for timestamp in timestamps:
              prefill.append(timestamp[1] - timestamp[0])
              encoder.append(
                  sum(
                      timestamp[i + 1] - timestamp[i]
                      for i in range(1, len(timestamp) - 1)
                  )
                  / (len(timestamp) - 2)
              )
              end2end.append(timestamp[-1] - timestamp[0])
          print(whole_end2end)
          with open(
              f"{args.log_path}/llama-{args.model}{args.dtype}_pp{args.pp_size}_{args.seq_len}_{args.new_length}_bsz{args.batch_size}_mbsz{args.mb_size}.log",
              "w+",
          ) as f:
              mb_avg_end2end = sum(end2end) / len(end2end)
              mb_avg_latency = mb_avg_end2end / (
                  args.new_length * args.mb_size
              )
              whole_avg_latency = whole_end2end / (
                  args.new_length * args.batch_size
              )
              num_layers = getattr(
                  model_config,
                  "num_layers",
                  model_config.num_hidden_layers,
              )
              num_parameters = (
                  num_layers
                  * model_config.hidden_size
                  * model_config.hidden_size
                  * 12
                  / args.pp_size
              )
              if args.dtype in ["fp16", "bf16"]:
                  num_bytes = 2
              else:
                  num_bytes = 4
  
              f.write(
                  f"llama-{args.model}{args.dtype}_pp{args.pp_size}, input_len:{args.seq_len}, output_len:{args.new_length}, bsz:{args.batch_size}, mbsz:{args.mb_size}\n"
              )
              f.write(
                  "Average prefill time: {0:8.2f} ms\n".format(
                      sum(prefill) / len(prefill) * 1000
                  )
              )
              f.write(
                  "Average encode time: {0:8.2f} ms\n".format(
                      sum(encoder) / len(encoder) * 1000
                  )
              )
              f.write(
                  "Average micro batch end2end time: {0:8.2f} ms\n".format(
                      mb_avg_end2end * 1000
                  )
              )
              f.write(
                  "Average micro batch Per Token Latency: {0:8.2f} ms\n".format(
                      mb_avg_latency * 1000
                  )
              )
              f.write(
                  "Whole batch end2end time: {0:8.2f} ms\n".format(
                      whole_end2end * 1000
                  )
              )
              f.write(
                  "Whole batch Per Token Latency: {0:8.2f} ms\n".format(
                      whole_avg_latency * 1000
                  )
              )
              f.write(
                  "Throughput: {} tokens/s\n".format(
                      (1000 / (whole_avg_latency * 1000))
                  )
              )
              f.write(
                  "flops: {0:8.2f} TFlops/s\n".format(
                      1
                      / whole_avg_latency
                      * num_parameters
                      * num_bytes
                      / 1e12
                  )
              )
              f.write(
                  "----------------------------------------------------------\n"
              )
  
      if torch.cuda.is_available():
          current_device = torch.cuda.current_device()
  
          # free memory and the total available memory in bytes
          (
              global_free_memory,
              total_GPU_memory_occupied,
          ) = torch.cuda.mem_get_info()
          memory_allocated = torch.cuda.memory_allocated()
          max_memory_allocated = torch.cuda.max_memory_allocated()
          memory_reserved = torch.cuda.memory_reserved()
          max_memory_reserved = torch.cuda.max_memory_reserved()
          with open(
              f"{args.log_path}/llama-{args.model}{args.dtype}_pp{args.pp_size}_{args.seq_len}_{args.new_length}_bsz{args.batch_size}_mbsz{args.mb_size}.log",
              "a",
          ) as f:
              f.write(
                  f"\nCurrently using GPU: {current_device}\n"
                  f"free memory : {global_free_memory / GIGABYTE:.4f} GB,\n"
                  f"total memory: {total_GPU_memory_occupied / GIGABYTE:.4f} GB,\n"
                  f"memory allocated: {memory_allocated / GIGABYTE:.4f} GB,\n"
                  f"Max CUDA memory allocated: {max_memory_allocated / GIGABYTE:.4f} GB,\n"
                  f"memory reserved/cached: {memory_reserved / GIGABYTE:.4f} GB,\n"
                  f"Max CUDA memory reserved/cached: {max_memory_reserved / GIGABYTE:.4f} GB,\n"
              )
  
file_path: colossalai/colossalai/legacy/inference/pipeline/benchmark/benchmark.py
imports: import argparse
  import time
  import torch
  import torch.distributed
  import transformers
  import colossalai
  from colossalai.inference import PPInferEngine
  from colossalai.inference.pipeline.policies import LlamaModelInferPolicy
label_desc: Use image node_id 2 for calling a global function with example usage: print_details_info(timestamps, model_config, args, whole_end2end) without return types