Skip checking if the eval context is active if we're in the background thread of async execution.

Signed-off-by: Rostan Tabet <rtabet@nvidia.com>

Author: rostan-t

dali/python/nvidia/dali/experimental/dynamic/_async.py

@@ -78,6 +78,7 @@ def __init__(self):
         self._queue = queue.SimpleQueue[Optional[_Future]]()
         self._thread: Optional[threading.Thread] = None
         self._event = threading.Event()
+        self._main_thread = threading.current_thread()
 
     def _worker(self):
         while True:

dali/python/nvidia/dali/experimental/dynamic/_eval_context.py

@@ -171,6 +171,9 @@ def __init__(self, *, num_threads=None, device_id=None, cuda_stream=None):
 
         # Used to disallow the EvalContext to be active in two threads simultaneously
         self._lock = threading.RLock()
+        # Python's RLock doesn't expose the owner nor the number of locks held so we track it here
+        self._lock_count = 0
+        self._locking_thread: threading.Thread | None = None
 
     def _purge_operator_cache(self):
         """Empties the operator instance cache"""
@@ -210,14 +213,13 @@ def _is_current(self) -> bool:
         return self is _tls.default.get(current_device_id)
 
     def __enter__(self):
-        if not self._lock.acquire(blocking=False):
-            raise RuntimeError("An EvalContext cannot be active in two threads simultaneously.")
+        self._try_acquire_lock()
         try:
             _tls.stack.append(self)
             if self._device:
                 self._device.__enter__()
         except Exception:
-            self._lock.release()
+            self._try_release_lock()
             raise
         return self
 
@@ -235,7 +237,7 @@ def __exit__(self, exc_type, exc_value, traceback):
             if self._device:
                 self._device.__exit__(exc_type, exc_value, traceback)
         finally:
-            self._lock.release()
+            self._try_release_lock()
 
     def evaluate_all(self):
         """Evaluates all pending invocations."""
@@ -331,6 +333,52 @@ def _snapshot(self):
     def _is_in_background_thread(self):
         return threading.current_thread() is self._async_executor._thread
 
+    def _should_acquire_lock(self):
+        # Skip checking if the eval context is already active if either:
+        # - We're in the background thread for async execution and it's active in the main thread
+        # - We're in the main thread and it's active in the background thread
+        background_thread = self._async_executor._thread
+        main_thread = self._async_executor._main_thread
+
+        if self._locking_thread is None or background_thread is None:
+            return True
+
+        return (threading.current_thread, self._locking_thread) not in (
+            (main_thread, background_thread),
+            (background_thread, main_thread),
+        )
+
+    def _try_acquire_lock(self):
+        # Skip checking if the eval context is already active if either:
+        # - We're in the background thread for async execution and it's active in the main thread
+        # - We're in the main thread and it's active in the background thread
+        background_thread = self._async_executor._thread
+        main_thread = self._async_executor._main_thread
+        if (
+            self._locking_thread is not None
+            and background_thread is not None
+            and (threading.current_thread(), self._locking_thread)
+            in (
+                (main_thread, background_thread),
+                (background_thread, main_thread),
+            )
+        ):
+            return
+
+        if not self._lock.acquire(blocking=False):
+            raise RuntimeError("An EvalContext cannot be active in two threads simultaneously.")
+        self._locking_thread = threading.current_thread()
+        self._lock_count += 1
+
+    def _try_release_lock(self):
+        if self._locking_thread is not threading.current_thread():
+            return
+
+        self._lock_count -= 1
+        if self._lock_count == 0:
+            self._locking_thread = None
+        self._lock.release()
+
 
 __all__ = [
     "EvalContext",