[40171] Limit size of kernel cache

[frreiss]

The EagerContext class caches instances of kernels that have been used by the current session. Because kernel state is opaque to EagerContext, the caching mechanism is very conservative about reusing kernels. Also, there is no limit on the number of kernels cached. As a result of these two factors, the size of the kernel cache grows in an unbounded fashion if the application calls tf.saved_model.load() repeatedly. This growth is the root cause of one of the memory leaks observed in issue #40171 .

This PR puts an upper limit on the size of the kernel cache in EagerContext. If the number of kernels exceeds this limit, the cache discards the least recently used kernel. I have hard-coded the capacity to a conservative number of 10000 kernels, which should be enough to prevent thrashing in normal usage while still providing protection against memory leaks.

I also added a RemoveKernelFromCache() method to go with the AddKernelToCache() method.

Here is a graph showing memory usage for a Python script that repeatedly loads and unloads a toy Keras model:

The blue points show the memory footprint of that script before applying these changes; and the orange points show the memory footprint after applying these changes.

After the changes in this PR, memory usage of my test script increases until the kernel cache reaches its maximum size. After that, the memory usage increases more slowly, because there are additional memory leaks in saved_model.load() that are not addressed in this PR.

[jaingaurav]

Thanks for this, but I do not believe this is the correct fix here. The kernel cache constantly growing is a sign of a different bug, i.e. kernels that should not be cached are being cached. By limiting the kernel cache we risk hiding actual bugs with poor performance.

Looking at your thorough investigation, it seems the problem is VarHandleOp is being cached due to the unique shared_name. If so, we should simply remove it from being a kernel cache candidate like we do for the tf.data ops. The change can be made here: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/common_runtime/eager/execute.cc#L92

[frreiss]

Thanks for the advice. I'll try blacklisting VarHandleOp and see if that fixes the problem. It will take a day or two to recompile.

[frreiss]

Opened a new PR #42377 that excludes VarHandleOp from kernel caching. I'll close the current PR in favor of the new one.