Repository: truera/trulens
Issue: [BUG] Context not copied in custom TP multi-threading class
Root cause: Architectural Constraint

**Bug Description**
Context is not copied over in threads in Threadpool of feedback functions. This is needed for persistent contextvars over the whole operation that might be used for example in logging or storing conversation info.

Trulens has an implementation on ThreadPoolExecutor that copies over the context, but the one from `concurrent.futures` is used in the `TP` class here, not copying over the context:
https://github.com/truera/trulens/blob/e22c75cef2da7222da893ca5bbf526e0ae96ab2e/src/core/trulens/core/utils/threading.py#L140

The issue is resolved by changing `fThreadPoolExecutor` into `ThreadPoolExecutor` here the context is coppied over and everything works as intended

**To Reproduce**
1. set contextvar on main thread
2. add multiple feedback functions with use_async
3. print the context var to std_out
4. Ask 2 queries after each other in the same process.

First time context var will be copied over, second query will use the threads with the old contextvar (context didn't get coppied).

**Expected behavior**
Context should be copied over when using `ThreadPoolExecutor`

**Environment:**
- OS: MacOS
- 3.10
- latest


<!-- Greeting -->
Hey @glennfeys! I'm here to help you with any bugs, questions, or contributions you have. Let's work together to solve this issue.

<!-- Answer -->
To resolve the issue of context not being copied over in threads when using `concurrent.futures.ThreadPoolExecutor` in the `TP` class, you should replace `fThreadPoolExecutor` with the custom `ThreadPoolExecutor` that correctly copies the context. Here is the relevant implementation:

1. **Custom `ThreadPoolExecutor` Implementation**:
    ```python
    class ThreadPoolExecutor(fThreadPoolExecutor):
        """A ThreadPoolExecutor that keeps track of the stack prior to each thread's
        invocation.

        Apps that do not use this thread pool might not be properly tracked.
        """

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)

        def submit(self, fn, /, *args, **kwargs):
            present_stack = stack()
            present_context = contextvars.copy_context()
            return super().submit(
                _future_target_wrapper,
                present_stack,
                present_context,
                fn,
                *args,
                **kwargs,
            )
    ```

2. **Update the `TP` Class**:
    ```python
    class TP(SingletonPerName):  # "thread processing"
        """Manager of thread pools.

        Singleton.
        """

        MAX_THREADS: int = 128
        """Maximum number of threads to run concurrently."""

        DEBUG_TIMEOUT: Optional[float] = 600.0  # [seconds], None to disable
        """How long to wait (seconds) for any task before restarting it."""

        def __init__(self):
            if safe_hasattr(self, "thread_pool"):
                # Already initialized as per SingletonPerName mechanism.
                return

            # Run tasks started with this class using this pool.
            self.thread_pool = ThreadPoolExecutor(
                max_workers=TP.MAX_THREADS, thread_name_prefix="TP.submit"
            )

            # Keep a separate pool for threads whose function is only to wait for
            # the tasks executed in the above pool. Keeping this separate to prevent
            # the deadlock whereas the wait thread waits for a task which will
            # never be run because the thread pool is filled with wait threads.
            self.thread_pool_debug_tasks = ThreadPoolExecutor(
                max_workers=TP.MAX_THREADS,
                thread_name_prefix="TP.submit with debug timeout",
            )

            self.completed_tasks = 0
            self.timedout_tasks = 0
            self.failed_tasks = 0

        def _run_with_timeout(
            self,
            func: Callable[[A], T],
            *args,
            timeout: Optional[float] = None,
            **kwargs
        ) -> T:
            if timeout is None:
                timeout = TP.DEBUG_TIMEOUT

            fut: Future[T] = self.thread_pool.submit(func, *args, **kwargs)

            try:
                res: T = fut.result(timeout=timeout)
                return res

            except TimeoutError as e:
                logger.error(
                    f"Run of {func.__name__} in {threading.current_thread()} timed out after {TP.DEBUG_TIMEOUT} second(s).\n"
                    f"{code_line(func)}"
                )

                raise e

            except Exception as e:
                logger.warning(
                    f"Run of {func.__name__} in {threading.current_thread()} failed with: {e}"
                )
                raise e

        def submit(
            self,
            func: Callable[[A], T],
            *args,
            timeout: Optional[float] = None,
            **kwargs
        ) -> Future[T]:
            if timeout is None:
                timeout = TP.DEBUG_TIMEOUT

            return self._submit(func, *args, timeout=timeout, **kwargs)

        def _submit(
            self,
            func: Callable[[A], T],
            *args,
            timeout: Optional[float] = None,
            **kwargs
        ) -> Future[T]:
            if timeout is None:
                timeout = TP.DEBUG_TIMEOUT

            return self.thread_pool_debug_tasks.submit(
                self._run_with_timeout, func, *args, timeout=timeout, **kwargs
            )
    ```

By using this custom `ThreadPoolExecutor`, the context will be correctly copied over in threads, ensuring that context variables are consistent across multiple queries <sup>[[1]](https://github.com/truera/trulens/blob/main/src/core/trulens/core/utils/threading.py)</sup><sup>[[2]](https://github.com/truera/trulens/blob/main/trulens_eval/trulens_eval/utils/threading.py)</sup>.


<!-- Dosu Comment Footer -->
*To continue talking to [Dosu](https://dosu.dev), mention @dosu.*


Thanks for finding this.