tf.function deadlock with multiple multiprocess/threading

[jonas-eschle]

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

binary

TensorFlow version

2.16

Custom code

Yes

Current behavior?

TensorFlow gets stuck when using multiprocessing/threading more than once.

I've observed it in more complicated situations with only once multithreading, but the following is a reproducible, standalone example that illustrates the point.

The code works correctly if either the tf.function decorator is removed or if xla compilation jit_compile is enabled (!).

Standalone code to reproduce the issue

import multiprocessing as mp  # can also be another module for multiprocess/threading
import random

import tensorflow as tf


# if we use jit_compile=True, it will work, magically
# it also works if there is no decorator at all, i.e. only eager mode
@tf.function(jit_compile=False)
def testjit(x):
    return tf.math.reduce_sum(x)


def make_zdata(_=None):
    print('Making data')
    # just to make sure that we recompile the function (different shapes)
    rnd = tf.random.uniform([random.randint(100, 10000)], -1, 1)
    zdata = testjit(rnd)
    print('Made data')
    return zdata

with mp.Pool(1) as executor:
    executor.map(make_zdata, [1])
executor.terminate()

# if we run this, it will fail (if jit_compile=False)
with mp.Pool(1) as executor:
    executor.map(make_zdata, [1])  # here, the code will be stuck
executor.terminate()

Relevant log output

will be, approximately:

Making data
Made data
Making data

and then it's stuck. Otherwise, if using jit_compile=True or not using the tf.function decorator at all, another Made data will be printed

[Venkat6871]

Hi @jonas-eschle ,
Sorry for the delay, I tried to run your code on colab using TF v2.16.1, nightly and faced the same issue. Please find the gist here for reference.

Thank you!