From 1b339144553018eb16e80d386e81cfea511efba1 Mon Sep 17 00:00:00 2001 From: Dev Aggarwal Date: Tue, 30 Jul 2024 03:45:24 +0530 Subject: [PATCH] Log warnings for UserError exceptions + handle UserError from GPU --- celeryapp/tasks.py | 2 ++ daras_ai_v2/gpu_server.py | 7 +++++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/celeryapp/tasks.py b/celeryapp/tasks.py index 439f9114e..794b5a061 100644 --- a/celeryapp/tasks.py +++ b/celeryapp/tasks.py @@ -11,6 +11,7 @@ from django.db.models import Sum from django.utils import timezone from fastapi import HTTPException +from loguru import logger from app_users.models import AppUser, AppUserTransaction from bots.admin_links import change_obj_url @@ -95,6 +96,7 @@ def save_on_step(yield_val: str | tuple[str, dict] = None, *, done: bool = False except Exception as e: if isinstance(e, UserError): sentry_level = e.sentry_level + logger.warning(e) else: sentry_level = "error" traceback.print_exc() diff --git a/daras_ai_v2/gpu_server.py b/daras_ai_v2/gpu_server.py index 1a59b12a5..03ce72855 100644 --- a/daras_ai_v2/gpu_server.py +++ b/daras_ai_v2/gpu_server.py @@ -6,7 +6,7 @@ from daras_ai.image_input import storage_blob_for from daras_ai_v2 import settings -from daras_ai_v2.exceptions import GPUError +from daras_ai_v2.exceptions import GPUError, UserError from gooeysite.bg_db_conn import get_celery_result_db_safe @@ -94,7 +94,10 @@ def call_celery_task( try: result.maybe_throw() except Exception as e: - raise GPUError(f"Error in GPU Task {queue}:{task_name} - {e}") from e + if type(e).__name__ == "UserError" and e.args and isinstance(e.args[0], dict): + raise UserError(**e.args[0]) + else: + raise GPUError(f"Error in GPU Task {queue}:{task_name} - {e}") from e record_cost_auto( model=queue, sku=ModelSku.gpu_ms, quantity=int((time() - s) * 1000) )