Add support for multi-threaded use of reraise_with_context.

PiperOrigin-RevId: 673306041

tensorflow_datasets/core/read_only_builder.py

@@ -356,9 +356,20 @@ def _find_builder_dir(name: str, **builder_kwargs: Any) -> str | None:
         all_builder_dirs.add(builder_dir)
   else:
     with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
-      for builder_dir in executor.map(find_builder_fn, all_data_dirs):
+      # Keep track of each new thread's error context, and add it to the main
+      # error context when each thread finishes.
+      def wrapped_find_builder_fn(data_dir):
+        with error_utils.record_error_context() as thread_context:
+          builder_dir = find_builder_fn(data_dir)
+        return thread_context, builder_dir
+
+      for context, builder_dir in executor.map(
+          wrapped_find_builder_fn, all_data_dirs
+      ):
         if builder_dir:
           all_builder_dirs.add(builder_dir)
+        for msg in context.messages:
+          error_utils.add_context(msg)
 
   if not all_builder_dirs:
     all_dirs_str = '\n\t- '.join([''] + [str(dir) for dir in all_data_dirs])

tensorflow_datasets/core/utils/error_utils.py

@@ -41,12 +41,29 @@ class ErrorContext:
 @edc.dataclass
 @dataclasses.dataclass
 class ContextHolder:
-  current_context_msg: ErrorContext | None = None
+  # Each thread will use its own instance of current_context_msg.
+  current_context_msg: edc.ContextVar[ErrorContext | None] = None
 
 
 context_holder = ContextHolder()
 
 
+@contextlib.contextmanager
+def record_error_context() -> Iterator[ErrorContext]:
+  """Contextmanager which captures the error context for a thread."""
+
+  if context_holder.current_context_msg is not None:
+    raise ValueError(
+        'Cannot record error context within the scope of another error context.'
+    )
+
+  context_holder.current_context_msg = ErrorContext()
+  try:
+    yield context_holder.current_context_msg
+  finally:
+    context_holder.current_context_msg = None
+
+
 @contextlib.contextmanager
 def reraise_with_context(error_cls: Type[Exception]) -> Iterator[None]:
   """Contextmanager which reraises an exception with an additional message.