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

Question

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

Answer 1

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.