FORMAT_CONVERSION_MAXCOMPUTE/format_convert/get_memory_info.py

import os
import time
from functools import wraps
import logging
import psutil
from format_convert.judge_platform import get_platform
if get_platform() == "Linux":
    import resource

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')


def memory_decorator(func):
    @wraps(func)
    def get_memory_info(*args, **kwargs):
        if get_platform() == "Windows":
            return func(*args, **kwargs)

        # 只有linux有resource包
        # usage = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
        usage = psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024
        start_time = time.time()
        logging.info("----- memory info start - " + func.__name__
                     + " - " + str(round(usage, 2)) + " GB"
                     + " - " + str(round(time.time()-start_time, 2)) + " sec")

        result = func(*args, **kwargs)

        # usage = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
        usage = psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024
        logging.info("----- memory info end - " + func.__name__
                     + " - " + str(round(usage, 2)) + " GB"
                     + " - " + str(round(time.time()-start_time, 2)) + " sec")
        return result

    return get_memory_info