Module memory leaks on every execution of widget

jleibs commented 2 months ago

It appears that every time we execute our widget, a new copy of the module is instantiated, and subsequently leaked.

You can replicate this with a module that includes a large const such as:

import anywidget
import traitlets

class LeakyWidget(anywidget.AnyWidget):
    _esm = """
    const buffer = new ArrayBuffer(100000000);

    function render({ model, el }) {
      let div = document.createElement("div");

      div.innerHTML = `Buffer is is ${buffer.byteLength} bytes.`;

      el.classList.add("leaky-widget");
      el.appendChild(div);
    }
    export default { render };
    """
    _css = """
    .leaky-widget button { background-color: #ea580c; }
    """
    value = traitlets.Int(0).tag(sync=True)

Repeatedly executing LeakyWidget() in a cell will leak ~100MB per run.

You can see this in a browser tool inspector by taking snapshots after each execution of the cell and seeing newly created ModuleEnvironmentObjects:

Ideally, we would have some mechanism to simply cache and re-use the same module across all widget instances, but it would be acceptable for them to at least get freed correctly.

manzt commented 2 months ago

Thanks for opening the issue and for the clear, reproducible example. This is certainly a bug and something I want to fix. I've been preparing for an anywidget tutuorial and talk at SciPy, which has been taking some time.

Ideally, we would have some mechanism to simply cache and re-use the same module across all widget instances, but it would be acceptable for them to at least get freed correctly.

Completely agree. I've been intending to migrate to this type of execution model since introducing the lifecycle hooks in v0.9.0. At that point we need to officially remove support for the prior named export however. So planning for v0.10.0, but we could fix the memory leak in v0.9.

jleibs commented 2 months ago

Thanks for the context, and good luck at SciPy.

FYI, for now we ended up working around this on our end by pushing our widget.js to our CDN in our build process and then by default referring to it via url instead of using the normal inlined approach. This hits the load-from-url codepath and correct caching behavior ensues.

Here's the PR in case anyone finds the approach informative: https://github.com/rerun-io/rerun/pull/6680

Given how large our widget.js bundle is, this also seems to have other benefits, in particular, in Colab notebooks where we were seeing a pretty lengthy load-process on every execution. If every instance of anywidget is isolated it may be tricky to implement caching using any approach other than a hosted asset.

manzt commented 2 months ago

Ok I got some headway with this experimenting in a standalone HTML file. Below is an example of basically what anywidget does behind the scenes:

<!doctype html>
<html lang="en">
    <head>
        <meta charset="utf-8">
        <title>Memory Leak</title>
    </head>
    <script type="module">
        add.addEventListener("click", async () => {
            let esm = `
            let buffer = new ArrayBuffer(100_000_000);
            export default {
                render() {
                    buffer; // just referencing here leaks memory
                },
            }
            `;
            let blob = new Blob([esm], { type: "application/javascript" });
            let url = URL.createObjectURL(blob);
            let mod = await import(url);
            mod.default.render();
            URL.revokeObjectURL(url);
            items.appendChild(
                Object.assign(document.createElement('p'), {
                    textContent: "Item",
                }),
            )
        });
    </script>
    <body>
        <h1>Memory Leak</h1>
        <button id="add">Add</button>
        <div id="items"></div>
    </body>
</html>

Clicking the button will leak memory for every click.

I was able to support freeing memory by having the module export a "free" method that sets buffer = null:

<!doctype html>
<html lang="en">
    <head>
        <meta charset="utf-8">
        <title>Memory Leak</title>
    </head>
    <script type="module">
        window.imports = [];
        add.addEventListener("click", async () => {
            let esm = `
            let buffer = new ArrayBuffer(100_000_000);
            export default {
                render() {
                    buffer; // just referencing here leaks memory
                },
                free() {
                    buffer = null;
                }
            }
            `;
            let blob = new Blob([esm], { type: "application/javascript" });
            let url = URL.createObjectURL(blob);
            let mod = await import(url);
            mod.default.render();
            URL.revokeObjectURL(url);
            items.appendChild(
                Object.assign(document.createElement('p'), {
                    textContent: "Item",
                }),
            )
            imports.push(mod);
        });
        remove.addEventListener("click", () => {
            imports.forEach(mod => mod.default.free());
            imports = [];
            items.innerHTML = "";
        });
    </script>
    <body>
        <h1>Memory Leak</h1>
        <button id="add">Add</button>
        <button id="remove">Clear</button>
        <div id="items"></div>
    </body>
</html>

Right now we support cleaning up resources during model initializaiton and rending with by returning functions from the callback hooks:

import anywidget
import traitlets

class LeakyWidget(anywidget.AnyWidget):
    _esm = """
    let buffer = new ArrayBuffer(100000000);
    function render({ model, el }) {
      let div = document.createElement("div");

      div.innerHTML = `Buffer is is ${buffer.byteLength} bytes.`;

      el.classList.add("leaky-widget");
      el.appendChild(div);
      return () => {
         // clean up resources created in render
      }
    }
    export default { 
        render,
        initialize() {
            return () => {
                buffer = null;
            }
        }
    };
    """
    _css = """
    .leaky-widget button { background-color: #ea580c; }
    """
    value = traitlets.Int(0).tag(sync=True)

widget = LeakyWidget()
widget

 widget.close() # closes the front-end comm, will trigger the dispose func from initialize and set `buffer = null`.

Since we reload the ESM per model instance, this might be a workaround at the moment. Perhaps we could consider a module-level API for "freeing" resource because that seems neccesary if we are going to reload the ESM, but I still think the goal should be to load the ESM once and cache.

manzt commented 2 months ago

Also thank you for sharing the idea CDN work around. I think that's a good idea for others experiencing this kind of issue or who have "large" widgets on the frontned

jprochazk commented 2 months ago

I was very confused by this behavior when investigating it, because anywidget doesn't actually do anything that would result in a memory leak. It turns out that this is a behavior of all major JS engines.

Modules are not collected until their realm is, even if the module and its exports are not reachable anymore. It's hard to dig up any concrete information about this, but here's one thread I found: https://esdiscuss.org/topic/are-es-modules-garbage-collected-if-so-do-they-re-execute-on-next-import It makes sense, because modules may be re-imported using the same specifier at any time. If they were evicted from cache, the browser would have to re-evaluate the module, including its side-effects.

For URLs created using URL.createObjectURL which are subsequently revoked, the cache could evict the module, because it isn't possible for it to be re-imported anymore. No browser does that, though. You can actually continue to import an object URL even after it's been revoked:

<html>
  <body>
    <h1>Memory Leak</h1>
    <script>
      async function leak() {
        let esm = `
          let buffer = new Uint8Array(100_000_000);
          export function render(msg) {
            console.log("leaked bytes: ", msg, buffer.byteLength);
          }
        `;
        let blob = new Blob([esm], { type: "application/javascript" });
        let url = URL.createObjectURL(blob);
        const m0 = await import(url);
        m0.render("first");
        URL.revokeObjectURL(url);
        const m1 = await import(url);
        m1.render("second");
      }
    </script>
    <button onclick="leak()">Add</button>
  </body>
</html>

To fix this in anywidget, the ESM should not be reloaded every time a widget is initialized. Because anywidget supports hot-reloading, it would still have to reload it in case the source has changed, which will continue to leak memory.

manzt commented 2 months ago

Thanks for the response.

To fix this in anywidget, the ESM should not be reloaded every time a widget is initialized.

Completely agree. We should load ESM once, upfront for the widget class. The changes I have for this may break packages that depend on Widget._esm (i.e., not many widgets but libraries like ipyreact that build on anywidget's core). But I'd like to postpone that release until after our SciPy tutorial.

Because anywidget supports hot-reloading, it would still have to reload it in case the source has changed, which will continue to leak memory.

Unfortunately, I think we'll have to accept that anywidget's built-in hot-reloading for development has a memory leak. I can add a section to the documentation explaining this caveat and suggest that you use a separate development server for very large bundles (e.g., Vite with the anywidget vite plugin). The ergonomics of having built-in HMR for prototyping is just too nice, and perhaps the memory leak is an acceptable trade off for the time being since it happens only in development.

jleibs commented 2 months ago

Agreed. A documented memory leak to support hot-reloading during development cycles seems well worth it. Also for probably 95% of users developing small JS-only packages the memory leak is not nearly as problematic.

Granted, as wasm continues to gain popularity this kind of issue will probably start to show up more often.

keller-mark commented 2 months ago

It might be helpful to think of this less like a memory leak and more of an issue of documenting / clarifying the widget lifecycle. A basic counter example demonstrates the current behavior as well.

class Widget(anywidget.AnyWidget):
    _esm = """
    if(!window.numModules) window.numModules = 0;
    if(!window.numGlobalInit) window.numGlobalInit = 0;
    if(!window.numGlobalRender) window.numGlobalRender = 0;

    window.numModules += 1;

    let numLocalInit = 0;
    let numLocalRender = 0;

    function initialize() {
        numLocalInit += 1;
        window.numGlobalInit += 1;
    }

    function render(view) {{
        numLocalRender += 1;
        window.numGlobalRender += 1;

        const h1 = document.createElement("h1");
        h1.innerText = JSON.stringify({ numModules, numLocalInit, numLocalRender, numGlobalInit, numGlobalRender });
        view.el.appendChild(h1);
    }}
    export default { initialize, render }
    """

keller-mark commented 2 months ago

We should load ESM once, upfront for the widget class

It might need to involve some kind of string hashing to ensure that widget development can still occur directly in a notebook

jprochazk commented 1 month ago

For URLs created using URL.createObjectURL which are subsequently revoked, the cache could evict the module, because it isn't possible for it to be re-imported anymore. No browser does that, though. You can actually continue to import an object URL even after it's been revoked

I created an issue on the chromium bug tracker ~2 weeks ago (https://issues.chromium.org/issues/350426234), and finally got a response from someone who understands the spec. This behavior is actually spec compliant, and evicting the module from the module map would be a regression.

revokeObjectURL will cause any subsequent dereferences of the Blob URL to fail with a network error, according to the spec - but the blob URL is just a plain string, and any subsequent imports of the same blob URL do not dereference the blob URL, they only retrieve the same module from the document's module map using the blob URL as the module specifier. So that does not result in a network error, and there is no way to be sure that the module is never going to be imported again.

keller-mark commented 1 month ago

Is there a reason this cannot be resolved by reusing the same widget instance on the Python side?

manzt commented 1 month ago

I created an issue on the chromium bug tracker ~2 weeks ago

Thanks for doing that and sharing the response.

Is there a reason this cannot be resolved by reusing the same widget instance on the Python side?

Widget ESM is defined on the class but we grab it from the front end from each view of each class instance. I have branch to change the behavior to reuse ESM (i.e., load once) and share it across all instances in the front end. For HMR we will just have documentation that there is a memory leak and point to this issue for explanation.

manzt / anywidget

Module memory leaks on every execution of widget #613