pythonram

How to measure RAM usage of each part of code in python?


I want to measure the RAM usage of each for loop in my code. I searched internet and find process = psutil.Process(os.getpid()) and print(process.memory_info().rss) for measuring RAM. But this code gets the pid of the whole process and not a specific part. Is there any way to measure RAM usage of each part of the code? For example in code below we have 3 for loops which fill 3 different dictionaries. I want to print RAM usage of each for loop and between the processing of each loop, if the RAM exceed a threshold, i want to break that for loop.

dict1 = {}
dict2 = {}
dict3 = {}

for i in range (200):
  do something with dict1
  if RAM usage of this block exceeds 1GB then break
  this loop used: x Mb

for i in range (500):
  do something with dict2
  if RAM usage of this block exceeds 1GB then break
  this loop used: x2 Mb

for i in range (800):
  do something with dict3
  if RAM usage of this block exceeds 1GB then break
  this loop used: x3 Mb

I appreciate answers which can help me a lot


Solution

  • You can read memory usage just before loop and then read it again inside loop. Then you can calculate loop memory usage as a difference between these two values. If it exceeds some threshold, break the loop.

    Here is sample code:

    import numpy as np
    import psutil
    import os
    
    process = psutil.Process(os.getpid())
    a = []
    threshhold = 64*1024*1024
    
    base_memory_usage = process.memory_info().rss
    
    for i in range(10):
        memory_usage = process.memory_info().rss
        loop_memory_usage = memory_usage - base_memory_usage
        
        print(loop_memory_usage)
        
        if loop_memory_usage > threshhold:
            print('exceeded threshold')
            break
    
        a.append(np.random.random((1000, 1000)))
    

    Result:

    0
    8028160
    16031744
    24035328
    32038912
    40042496
    48046080
    56049664
    64053248
    72056832
    exceeded threshold
    

    As you can see, before any action the loop uses 0 bytes of memory.