From 82e6b38ca4941eafb6430e3e03d08f882d183465 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 17 Aug 2023 10:16:00 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(worker.py):=20add=20exceptio?= =?UTF-8?q?n=20handling=20and=20retry=20mechanism=20to=20build=5Fvertex=20?= =?UTF-8?q?task=20to=20handle=20SoftTimeLimitExceeded=20error=20and=20retr?= =?UTF-8?q?y=20the=20task=20up=20to=203=20times=20with=20a=20countdown=20o?= =?UTF-8?q?f=202=20seconds=20=E2=9C=A8=20feat(worker.py):=20add=20import?= =?UTF-8?q?=20statement=20for=20SoftTimeLimitExceeded=20exception=20from?= =?UTF-8?q?=20celery.exceptions=20module=20to=20handle=20SoftTimeLimitExce?= =?UTF-8?q?eded=20error=20in=20build=5Fvertex=20task?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/worker.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/worker.py b/src/backend/langflow/worker.py index 67eaf6b3b..f8831df76 100644 --- a/src/backend/langflow/worker.py +++ b/src/backend/langflow/worker.py @@ -1,6 +1,7 @@ from langflow.core.celery_app import celery_app from typing import Any, Dict, Optional from typing import TYPE_CHECKING +from celery.exceptions import SoftTimeLimitExceeded if TYPE_CHECKING: from langflow.graph.vertex.base import Vertex @@ -11,13 +12,19 @@ def test_celery(word: str) -> str: return f"test task return {word}" -@celery_app.task -def build_vertex(vertex: "Vertex") -> "Vertex": +@celery_app.task(bind=True, soft_time_limit=30, max_retries=3) +def build_vertex(self, vertex: "Vertex") -> "Vertex": """ Build a vertex """ - vertex.build() - return vertex + try: + vertex.task_id = self.request.id + vertex.build() + return vertex + except SoftTimeLimitExceeded as e: + raise self.retry( + exc=SoftTimeLimitExceeded("Task took too long"), countdown=2 + ) from e @celery_app.task(acks_late=True)