From 3454ede5a56c1ce7a4578560aca18bd426d26945 Mon Sep 17 00:00:00 2001 From: Christophe Bornet Date: Sat, 28 Dec 2024 05:54:55 +0100 Subject: [PATCH] ref: Use better task cancellation logic (#5397) Use better task cancellation logic Co-authored-by: Gabriel Luiz Freitas Almeida --- .../base/langflow/services/telemetry/service.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/src/backend/base/langflow/services/telemetry/service.py b/src/backend/base/langflow/services/telemetry/service.py index bb8d27e45..908b50094 100644 --- a/src/backend/base/langflow/services/telemetry/service.py +++ b/src/backend/base/langflow/services/telemetry/service.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio import os import platform -import sys from datetime import datetime, timezone from typing import TYPE_CHECKING @@ -135,15 +134,11 @@ class TelemetryService(Service): async def _cancel_task(self, task: asyncio.Task, cancel_msg: str) -> None: task.cancel(cancel_msg) - try: - await task - except asyncio.CancelledError: - current_task = asyncio.current_task() - if sys.version_info >= (3, 11): - if current_task and current_task.cancelling() > 0: - raise - elif current_task and hasattr(current_task, "_must_cancel") and current_task._must_cancel: - raise + await asyncio.wait([task]) + if not task.cancelled(): + exc = task.exception() + if exc is not None: + raise exc async def stop(self) -> None: if self.do_not_track or self._stopping: