VRAM usage with adding and removing WorldObjects

[kushalkolar]

I've been looking into VRAM usage with adding and removing WorldObjects, here's some things that I found:

With a Qt canvas it seems like WorldObjects are garbage collected and GPU VRAM is freed only after mouse movements on the canvas even after the WorldObject has visibly been removed from the scene.

This example uses nvidia-smi:

1. Creates a 4096x4096 image when you double click on the canvas.
2. Click on the image to remove it from the scene and call del on it. GPU VRAM is not freed.
3. Make a mouse movement on the canvas after the image has been removed, GPU VRAM gets freed. But the GPU VRAM remains utilized until you make this mouse event.

import numpy as np
from wgpu.gui.auto import WgpuCanvas, run
import pygfx as gfx
import subprocess

canvas = WgpuCanvas()
renderer = gfx.WgpuRenderer(canvas)
scene = gfx.Scene()
camera = gfx.OrthographicCamera(5000, 5000)
camera.position.x = 2048
camera.position.y = 2048


def make_image():
    data = np.random.rand(4096, 4096).astype(np.float32)

    return gfx.Image(
        gfx.Geometry(grid=gfx.Texture(data, dim=2)),
        gfx.ImageBasicMaterial(clim=(0, 1)),
    )


def draw():
    renderer.render(scene, camera)
    canvas.request_draw()


def print_nvidia(msg=""):
    print(msg)
    print(
        subprocess.check_output(["nvidia-smi", "--format=csv", "--query-gpu=memory.used"]).decode().split("\n")[1]
    )
    print()


def add_img(*args):
    print_nvidia("Before creating image")
    img = make_image()
    print_nvidia("After creating image")
    scene.add(img)
    img.add_event_handler(remove_img, "click")
    draw()
    print_nvidia("After add image to scene")


def remove_img(*args):
    img = scene.children[0]
    scene.remove(img)
    draw()
    print_nvidia("After remove image from scene")
    del img
    draw()
    print_nvidia("After del image")
    renderer.add_event_handler(print_nvidia, "pointer_move")


renderer.add_event_handler(add_img, "double_click")

draw()
run()

This outputs:

Before creating image
73 MiB

After creating image
73 MiB

After add image to scene
137 MiB

After remove image from scene
137 MiB

After del image
137 MiB

<pygfx.objects._events.PointerEvent object at 0x7fa186d8f970>
73 MiB

<pygfx.objects._events.PointerEvent object at 0x7fa186d8f5e0>
73 MiB

To follow up from #382 , with jupyter the situation is more messy. As Korijn mentioned here: https://github.com/pygfx/pygfx/issues/382#issuecomment-1310186390

If anyone runs into issues with jupyter in the future, in fastplotlib I'm working on a workaround where all WorldObjects are stored in a dict that fastplotlib uses internally, and only weakref proxies are used to access the WorldObject so that they are never directly exposed to any references on the jupyter side: https://github.com/kushalkolar/fastplotlib/pull/160

[almarklein]

Thanks for looking into this!

I rewrote the example a bit to make the draw calls async as would normally happen, with the mem usage continuously being updated in the terminal:

import psutil
import time

import numpy as np
# import PySide6
from wgpu.gui.auto import WgpuCanvas, run
import pygfx as gfx
import subprocess

canvas = WgpuCanvas()
renderer = gfx.WgpuRenderer(canvas)
scene = gfx.Scene()
camera = gfx.OrthographicCamera(5000, 5000)
camera.position.x = 2048
camera.position.y = 2048

p = psutil.Process()

def make_image():
    data = np.random.rand(4096, 4096).astype(np.float32)

    return gfx.Image(
        gfx.Geometry(grid=gfx.Texture(data, dim=2)),
        gfx.ImageBasicMaterial(clim=(0, 1)),
    )


def draw():
    renderer.render(scene, camera)
    canvas.request_draw()
    print(f"\r{get_mem_usage()}", end='')


def get_mem_usage():
    # return p.memory_info().rss / 1024**2
    return subprocess.check_output(["nvidia-smi", "--format=csv", "--query-gpu=memory.used"]).decode().split("\n")[1].strip()


def add_img(*args):
    img = make_image()
    scene.add(img)
    print("\nCreated image")
    img.add_event_handler(remove_img, "click")


def remove_img(*args):
    img = scene.children[0]
    scene.remove(img)
    print("\nRemoved image")


renderer.add_event_handler(add_img, "double_click")

renderer.request_draw(draw)
run()

When I test this on Windows 11, with the glfw backend, it works fine: the memory goes up and then down. So maybe it just needs a bit of time settle, which is why you saw it reduce in the mouse move event. When the GPU decides to clean up its memory exactly, is not in our control.

BTW: If anyone knows how to get the GPU memory consumption on MacOS, that'd be great :)

[Korijn]

BTW: If anyone knows how to get the GPU memory consumption on MacOS, that'd be great :)

Or for any GPU that is not NVidia, for that matter!

[kushalkolar]

If I use this then the VRAM is cleared immediately after the object is removed from the scene:

renderer.request_draw(draw)
run()

However with this it requires an event to clear, otherwise VRAM will linger even for minutes:

draw()
run()

The same code also works with garbage collection in jupyter. However if you run a jupyter cell such that it ends with a WorldObject (calls the __repr__), then it does not get garbage collected. I can help create a note in the jupyter-rfb repo with some examples and suggest using weakreferences if you need to interact with WorldObjects on the jupyter side (unless you think there's a better way?). Anyways I'll open up an issue there.

[almarklein]

However with this it requires an event to clear, otherwise VRAM will linger even for minutes:

During that time, are there any new draws? I suspect it's more related to a draw being performed (which then triggers GPU''s GC in some way), rather than events.

However if you run a jupyter cell such that it ends with a WorldObject (calls the __repr__), then it does not get garbage collected.

I can help create a note in the jupyter-rfb repo with some examples and suggest using weakreferences if you need to interact with WorldObjects on the jupyter side (unless you think there's a better way?). Anyways I'll open up an issue there.

Yeah, a PR or issue would be nice. I think explaining the gotcha's and some ways to prevent that would be nice, but suggestion weakrefs in user code may be a bit much? :)

[kushalkolar]

BTW: If anyone knows how to get the GPU memory consumption on MacOS, that'd be great :)

Or for any GPU that is not NVidia, for that matter!

nvtop works for AMD GPUs on linux at least!

And I just learned that I actually have an RX 570 not a 470. :laughing: