from pynvml import *
import psutil

"""
To track memory allocation, let's take advantage of the nvidia-ml-py3 package and GPU memory allocation from python.

ref: https://huggingface.co/docs/transformers/v4.20.1/en/perf_train_gpu_one
"""


def print_gpu_initialization():
    nvmlInit()
    handle = nvmlDeviceGetHandleByIndex(0)
    info = nvmlDeviceGetMemoryInfo(handle)
    print(f"GPU memory occupied: {info.used//1024**2} MB.")
    return info.used // 1024**2


def print_summary(result):
    print(f"Time: {result.metrics['train_runtime']:.2f}")
    print(f"Samples/second: {result.metrics['train_samples_per_second']:.2f}")
    print_gpu_utilization()


def print_memory_usage():
    process = psutil.Process(os.getpid())
    memory_in_bytes = process.memory_info().rss
    memory_in_megabytes = memory_in_bytes / (1024 ** 2)
    #print(f"Memory used by this script: {memory_in_megabytes:.2f} MB")

    return memory_in_megabytes