From 950585da349e7f55b4d1fc2b9ac8f7212a0c8951 Mon Sep 17 00:00:00 2001 From: Anton Dubovik Date: Mon, 20 Jan 2025 18:07:51 +0000 Subject: [PATCH] fix: fixed issues with lost Context --- aidial_analytics_realtime/utils/concurrency.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/aidial_analytics_realtime/utils/concurrency.py b/aidial_analytics_realtime/utils/concurrency.py index 0026675..d60be11 100644 --- a/aidial_analytics_realtime/utils/concurrency.py +++ b/aidial_analytics_realtime/utils/concurrency.py @@ -1,4 +1,6 @@ import asyncio +import contextvars +import functools from concurrent.futures import ThreadPoolExecutor from typing import Callable, ParamSpec, TypeVar @@ -12,4 +14,6 @@ async def run_in_cpu_tasks_executor( func: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs ) -> _T: loop = asyncio.get_event_loop() - return await loop.run_in_executor(cpu_task_executor, func, *args) # type: ignore + ctx = contextvars.copy_context() + func_call = functools.partial(ctx.run, func, *args, **kwargs) + return await loop.run_in_executor(cpu_task_executor, func_call) # type: ignore