Memory leak when repeatedly saving a variable during a host_callback call.

I have a use case where I want to use host_callback to save a python object (specifically a scipy SuperLU object) and update it over and over again within a jitted function. These should persist between jit calls. Unfortunately I am running into a memory leak. I have a minimal example below, where the memory consumptions blows up linearly as the code runs. Even if I insert gc.collect() in the host_callback function, it seems that the SuperLU objects are somehow being kept around.

import numpy as np
import scipy.sparse
import scipy.sparse.linalg

import jax.experimental.host_callback as hcb
import jax
import jax.numpy as jnp

class SolveFnContainer:
    lu = None

def host_do_something(_):
    print('solving')

    n = 100000
    A = scipy.sparse.diags((-1*np.ones(n-1), -1*np.ones(n-1), 2*np.ones(n)), offsets=(-1, 1, 0), format='csc')
    SolveFnContainer.lu = scipy.sparse.linalg.spilu(A)
    return 0.0

def device_do_something():
    inputs = (0.0,)
    return hcb.call(host_do_something, inputs,
                    result_shape=jax.ShapeDtypeStruct((), np.float64))

def main():
    def body_fn(x):
        return device_do_something()

    def cond_fn(x):
        return x == 0.0

    return jax.lax.while_loop(cond_fn, body_fn, 0.0)

if __name__ == '__main__':
    jit_main = jax.jit(main)
    print(jit_main())

[hawkinsp]

I think this is a scipy bug. Here’s a reproduction without any JAX pieces:

from time import sleep

import numpy as np
import scipy.sparse
import scipy.sparse.linalg

import os
import psutil
import gc
import threading


class Container:
    A = None


def host_do_something(_):
    n = 100000
    A = scipy.sparse.diags((-1*np.ones(n-1), -1*np.ones(n-1), 2*np.ones(n)), offsets=(-1, 1, 0), format='csc')
    Container.A = scipy.sparse.linalg.spilu(A)
    #Container.A = A
    return 0.0


def thread_do_something(_):
    t = threading.Thread(target=host_do_something, args=(0.0,))
    t.start()
    t.join()

if __name__ == '__main__':
    while True:
        thread_do_something(0.0)

        print(psutil.Process(os.getpid()).memory_info().rss / 1024 ** 2)
        sleep(0.2)
        gc.collect()

It appears the leak occurs because JAX is running your callback on another thread. For some reason this triggers a scipy leak. I don’t know what’s happening without debugging scipy, but at this point it’s not really a JAX problem.

As to why this happens under JAX on GPU but not CPU: JAX uses different callback implementations on CPU and GPU, and only on GPU is it calling back on another thread. (As it happens, @sharadmv is working on a GPU callback implementation that uses only one thread, but that is ultimately papering over the problem I think.)

Hope that helps!

[denizokt]

Thanks so much! I will look into this and redirect to scipy if needed.