메모리 사용량 추적 객체

메모리 확인 - psutil

import os
import psutil

pid = os.getpid()
process = psutil.Process(pid)
memory = process.memory_info().rss

print(f"사용 중인 메모리: {memory / 1024**2}MiB")

현재 할당된 pid 값을 넘겨받은 psutil.Process는 memory_info를 통해 물리적 메모리의 사용량을 알려준다. 

처음에는 데코레이터를 활용해 함수의 시작과 끝에서 메모리 확인을 통해 메모리 사용량을 측정하면 되지 않을까 생각했다. 하지만 Python은 내부적으로 Garbage Collector를 활용해 자동으로 메모리가 관리되기 때문에 함수가 실행되는 과정에서 메모리의 변화가 발생할 수 있다. 그래서 함수의 실행 과정을 추적하는 동작이 추가로 필요하다고 판단했다. 

---

추적

import sys
  
def tracer(frame, event, arg):
    # 필요한 작업 수행
    return tracer
   
sys.settrace(my_tracer)
# 추적할 코드
sys.settrace(None)

추적을 실행하기 위해서는 tracer 함수를 만들어주어야 한다. 함수는 frame, event, arg 3개의 값을 파라미터를 통해 전달받아야 한다. 그리고 tracer 함수를 settrace로 넘겨주며 추적을 시작하고, None을 넘겨주며 추적을 종료한다. 자세한  내용은 아래와 같다.

• frame: 현재의 스택 프레임(stack frame)
• event: 'call', 'line', 'return', 'exception', 'opcode' 중 하나이다.
• arg는 event에 따라 달라진다. 아래 문서 참고

sys — System-specific parameters and functions — Python 3.10.5 documentation

예시:

import sys

def my_tracer(frame, event, arg):
    """tracer 함수"""
    print(f"{event}\t{frame.f_code.co_name}")
    return my_tracer

def test():
    """테스트를 위한 함수"""
    list_ = [i for i in range(2)]
    print(" --출력:", list_)
    return list_

print("event\tfunction")
print("-----------------")

# 추적 시작
sys.settrace(my_tracer)
list_ = test()

# 추적 종료
sys.settrace(None)

event   function
-----------------
call    test
line    test
call    <listcomp>
line    <listcomp>
line    <listcomp>
line    <listcomp>
return  <listcomp>
line    test
 --출력: [0, 1]
return  test

*위 결과를 확인해보면 함수 내의 코드가 실행되기 전 tracer 함수가 먼저 실행된다는 것을 알 수 있다.  

---

메모리 추적 객체

import os
import sys
import psutil

class Tracer(object):
    def __init__(self, *object_names):
        self.__object_names = object_names
        self.__process = psutil.Process(os.getpid())
        self.__record = [0 for _ in range(10)]  # 임시로 설정
        self.__count = 0

    def __enter__(self):
        """with문 추적 시작"""
        sys.settrace(self.trace)
        return self

    def __exit__(self, *args):
        """with문 추적 종료"""
        sys.settrace(None)

    def trace(self, frame, event, arg):
        """추적 내용"""
        if (frame.f_code.co_name in self.__object_names) and (
            event in ("call", "line", "return")
        ):
            self.__record[self.__count] = self.__process.memory_info().rss
            self.__count += 1
        return self.trace

    @property
    def record(self):
        """추적된 메모리 반환"""
        return self.__record[: self.__count]

우선 Tracer라는 class를 선언해 보았다. trace 함수를 살펴보면 특정 함수의 코드 라인이 실행되었고, 이벤트가 call · line · return일 때 __record에 메모리를 기록한다. 이렇게 기록된 메모리 정보는 record를 통해 확인할 수 있다. 사용법은 아래와 같다.

* class의 __record의 길이는 실행되는 함수의 라인 수보다 많아야 한다. 

with Tracer(추적하고자 하는 객체명1, 객체명2 ... ) as tracer:
    # 함수 실행
    
memory = tracer.record  # 메모리 기록

---

실행 예시

Tracer 선언 (위 코드와 동일)

class Tracer(object):
    """메모리 정보를 가져올 Tracer"""
    ...

 

Tracer를 이용해 메모리 정보 확인하기

import matplotlib.pyplot as plt

def temp():
    """테스트를 위해 만든 함수"""
    a = [i for i in range(10 ** 5)]
    b = [i for i in range(10 ** 2)]
    del a
    c = [i for i in range(10 ** 4)]
    d = [i for i in range(3)]
    final = b + d
    return final

with Tracer("temp") as tracer:
    # tracer 실행
    temp()

# 메모리 기록 가져오기
record = [x / 1024 ** 2 for x in tracer.record]
# 시각화
plt.figure()
plt.plot(record)
plt.ylabel("Memory (MiB)")
plt.show()

결과:

     event   내용
----------------------------------------------
0~1: call    `temp` 호출
1~2: line    `a = [i for i in range(10 ** 5)]`
2~3: line    `b = [i for i in range(10 ** 2)]`
3~4: line    `del a`
4~5: line    `c = [i for i in range(10 ** 4)]`
5~6: line    `d = [i for i in range(3)]`
6~7: line    `final = b + d`
7~8: line    `return final`
8~ : return  `final` 반환

시각화 외에도 최대, 최소를 구하거나 평균을 구하는 등 다양하게 응용할 수 있다.