GPU memory management

Hi guys,

How is GPU memory managed by Ray? Here is how I see the picture. An example:

import pandas
import ray


df = pandas.DataFrame()
o_ref = ray.put(df)

def foo(ref):
    # processing with ref


When ray.put(df) is called, the data is put into plasma. When calling foo the data is transferred to GPU memory to do some processing.

The question is: is the data left on GPU memory alive or cleaned up from there right away?

It would be great if there is some docs describing GPU memory management in Ray.

Thanks in advance!

@sangcho , any thoughts?

Have you read this part of the doc? GPU Support — Ray v2.0.0.dev0

Is it helpful? I believe it is related

@sangcho , thank you for pointing out to the docs! It is very helpful. As I understand, either a task itself should manage/clean up GPU memory or we can specify max_calls=1 to directly release the GPU resources, right?

@YarShev Yes that seems like how this works now! IIUC, Ray itself doesn’t have special handling in GPU memory usage!

1 Like