c5a6003ef5
✨ feat(manager.py): add debug logs to check Celery availability and backend being used 🔧 fix(manager.py): move check_celery_availability function definition above its usage to improve code readability ✨ feat(manager.py): add debug logs to show task launch details and backend being used
69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
from typing import Any, Callable, Coroutine, Union
|
|
from loguru import logger
|
|
from langflow.services.base import Service
|
|
from langflow.services.task.backends.anyio import AnyIOBackend
|
|
from langflow.services.task.backends.base import TaskBackend
|
|
from langflow.services.task.utils import get_celery_worker_status
|
|
|
|
|
|
def check_celery_availability():
|
|
from langflow.worker import celery_app
|
|
|
|
try:
|
|
status = get_celery_worker_status(celery_app)
|
|
logger.debug(f"Celery status: {status}")
|
|
except Exception as e:
|
|
logger.error(f"An error occurred: {e}")
|
|
status = {"availability": None}
|
|
return status
|
|
|
|
|
|
try:
|
|
status = check_celery_availability()
|
|
|
|
USE_CELERY = status.get("availability") is not None
|
|
except ImportError:
|
|
USE_CELERY = False
|
|
|
|
|
|
class TaskService(Service):
|
|
name = "task_service"
|
|
|
|
def __init__(self):
|
|
self.backend = self.get_backend()
|
|
self.use_celery = USE_CELERY
|
|
|
|
def get_backend(self) -> TaskBackend:
|
|
if USE_CELERY:
|
|
from langflow.services.task.backends.celery import CeleryBackend
|
|
|
|
logger.debug("Using Celery backend")
|
|
return CeleryBackend()
|
|
logger.debug("Using AnyIO backend")
|
|
return AnyIOBackend()
|
|
|
|
# In your TaskService class
|
|
async def launch_and_await_task(
|
|
self,
|
|
task_func: Callable[..., Any],
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
if not self.use_celery:
|
|
return None, await task_func(*args, **kwargs)
|
|
if not hasattr(task_func, "apply"):
|
|
raise ValueError(f"Task function {task_func} does not have an apply method")
|
|
task = task_func.apply(args=args, kwargs=kwargs)
|
|
result = task.get()
|
|
return task.id, result
|
|
|
|
async def launch_task(
|
|
self, task_func: Callable[..., Any], *args: Any, **kwargs: Any
|
|
) -> Any:
|
|
logger.debug(f"Launching task {task_func} with args {args} and kwargs {kwargs}")
|
|
logger.debug(f"Using backend {self.backend}")
|
|
task = self.backend.launch_task(task_func, *args, **kwargs)
|
|
return await task if isinstance(task, Coroutine) else task
|
|
|
|
def get_task(self, task_id: Union[int, str]) -> Any:
|
|
return self.backend.get_task(task_id)
|