Thread‑safe inference configuration for tf.keras.Model under concurrent load

Implementing a thread_safe argument in tf.keras.Model would likely function as a wrapper around the predict method, utilizing a reentrant lock to serialize access to the model's internal state. While this would provide convenience, it would not fundamentally change the underlying execution model of TensorFlow 2.x, which is not natively designed for concurrent predict calls on a single model instance.

Expected Latency and Performance Impact

An internal thread_safe flag would introduce latency overhead nearly identical to that of an external threading.Lock. Because the lock would serialize all inference requests, the primary overhead is not the lock acquisition itself, but the contention created under high load. As the number of concurrent threads increases, latency for individual requests will grow linearly as they queue for the lock.

Comparison of Locking Strategies

Approach	Implementation	Overhead	Control
External Lock	with lock: model.predict(...)	Low (Lock acquisition)	High (Developer controls scope)
Internal Flag	model = Model(..., thread_safe=True)	Low (Lock acquisition)	Low (Automatic serialization)
Model Duplication	One model per thread/process	High (VRAM/RAM usage)	Maximum (Parallel execution)

Behavior with tf.function and Custom Layers

The interaction between a thread-safety flag and TensorFlow's optimization layers depends on where the lock is placed:

• tf.function: If the model is wrapped in @tf.function, the lock must reside outside the traced graph. Placing a lock inside the compiled graph is generally not possible as tf.function converts Python logic into a static graph. The flag would protect the call to the compiled function, ensuring only one thread triggers the graph execution at a time.
• Custom Layers: The flag would protect the predict call sequence, but it would not automatically make custom layers thread-safe if those layers manipulate external Python state or non-TensorFlow variables. If a custom layer updates a shared Python list or dictionary during inference, additional synchronization is still required.

Verification Method

To verify the impact of such a mechanism in a current environment, you can simulate the flag using a wrapper class:

import threading
import tensorflow as tf

class ThreadSafeModel(tf.keras.Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._lock = threading.Lock()

    def predict(self, *args, **kwargs):
        with self._lock:
            return super().predict(*args, **kwargs)

To provide a more tailored recommendation, are you currently seeing specific ResourceExhaustedError crashes or nondeterministic output values during your concurrent runs?