Monitoring memory usage with /proc

memory.py

import numpy as np
import ray
import os


@ray.remote
def f():
    return np.random.rand(1000_000_000 // 8)


ray.init()
pid = os.getpid()
procfile = f"/proc/{pid}/statm"


def get_memory():
    with open(procfile, 'r') as f:
        out = f.read()
        out = [int(m) for m in out.split(' ')]
        output = list(zip(['size', 'resident', 'shared'], out))
    return dict(output)

def run(fn):
    mem_init = get_memory()
    x = fn()
    mem = get_memory()
    for key, m in mem.items():
        print(f"{(m - mem_init[key]) * 4 / 1000}MB additional {key}")
    print("")

def numpy():
    print("x = np.random")
    return np.random.rand(1000_000_000 // 8)

def ray_put():
    print("x = ray.put(np.random)")
    return ray.put(np.random.rand(1000_000_000 // 8))

def ray_put_get():
    print("x = ray.get(ray.put(np.random))")
    return ray.get(ray.put(np.random.rand(1000_000_000 // 8)))

def ray_remote():
    print("x = ray.get(f.remote())")
    return ray.get(f.remote())


run(numpy)
run(ray_put)
run(ray_put_get)
run(ray_remote)