chore(breadbox): Celery checker

breadbox/breadbox/api/compute.py

@@ -7,11 +7,11 @@
 from breadbox.config import Settings, get_settings
 from depmap_compute import models
 
-from breadbox.schemas.custom_http_exception import UserError, HTTPError
+from breadbox.schemas.custom_http_exception import UserError, CeleryConnectionError
 from ..schemas.compute import ComputeParams, ComputeResponse
 from ..compute import analysis_tasks
 from .dependencies import get_user
-from ..celery_task.utils import format_task_status, cast_celery_task
+from ..celery_task import utils
 
 
 router = APIRouter(prefix="/compute", tags=["compute"])
@@ -69,6 +69,8 @@ def compute_univariate_associations(
     user: str = Depends(get_user),
     settings: Settings = Depends(get_settings),
 ):
+    utils.check_celery()
+
     resultsDirPrefix = settings.compute_results_location
     dataset_id = computeParams.datasetId
     vector_variable_type = computeParams.vectorVariableType
@@ -97,7 +99,7 @@ def compute_univariate_associations(
     )
 
     try:
-        result = cast_celery_task(analysis_tasks.run_custom_analysis).delay(
+        result = utils.cast_celery_task(analysis_tasks.run_custom_analysis).delay(
             user=user,
             analysis_type=analysis_type,
             query_node_id=computeParams.queryId,
@@ -113,9 +115,11 @@ def compute_univariate_associations(
     except PermissionError as e:
         raise HTTPException(403, detail=str(e))
 
-    return format_task_status(result)
+    return utils.format_task_status(result)
 
 
 @router.get("/test_task", operation_id="test_task")
 def test_task(message):
-    cast_celery_task(analysis_tasks.test_task).delay(message)
+    utils.check_celery()
+
+    utils.cast_celery_task(analysis_tasks.test_task).delay(message)

breadbox/breadbox/api/dataset_uploads.py

@@ -6,6 +6,7 @@
 from .dependencies import get_user
 
 from ..celery_task import utils
+from breadbox.schemas.custom_http_exception import CeleryConnectionError
 
 router = APIRouter(prefix="/dataset-v2", tags=["datasets"])
 
@@ -64,6 +65,8 @@ def add_dataset_uploads(
         - `col_type`: Annotation type for the column. Annotation types may include: `continuous`, `categorical`, `binary`, `text`, or `list_strings`
 
     """
+    utils.check_celery()
+
     # Converts a data type (like a Pydantic model) to something compatible with JSON, in this case a dict. Although Celery uses a JSON serializer to serialize arguments to tasks by default, pydantic models are too complex for their default serializer. Pydantic models have a built-in .dict() method but it turns out it doesn't convert enums to strings which celery can't JSON serialize, so I opted to use fastapi's jsonable_encoder() which appears to successfully json serialize enums
     dataset_json = jsonable_encoder(dataset)
     result = run_dataset_upload.delay(dataset_json, user)  # pyright: ignore

breadbox/breadbox/api/downloads.py

@@ -20,6 +20,8 @@
 
 from ..config import Settings, get_settings
 from .dependencies import get_user, get_db_with_user
+from breadbox.schemas.custom_http_exception import CeleryConnectionError
+
 
 router = APIRouter(prefix="/downloads", tags=["downloads"])
 log = getLogger(__name__)
@@ -67,6 +69,7 @@ def export_dataset(
     user: str = Depends(get_user),
     settings: Settings = Depends(get_settings),
 ):
+    utils.check_celery()
 
     dataset_id = exportParams.datasetId
     feature_labels = exportParams.featureLabels
@@ -103,6 +106,8 @@ def export_merged_dataset(
     user: str = Depends(get_user),
     settings: Settings = Depends(get_settings),
 ):
+    utils.check_celery()
+
     dataset_ids = exportParams.datasetIds
     feature_labels = exportParams.featureLabels
     sample_ids = exportParams.cellLineIds

breadbox/breadbox/celery_task/utils.py

@@ -11,7 +11,7 @@
 from fastapi import HTTPException
 from typing import Any, Optional
 from ..compute.celery import app
-from breadbox.schemas.custom_http_exception import UserError
+from breadbox.schemas.custom_http_exception import UserError, CeleryConnectionError
 from typing import Protocol, cast, Callable
 from celery.result import AsyncResult, EagerResult
 
@@ -204,3 +204,25 @@ def update_state(
         meta["message"] = message
 
     task.update_state(state=state, meta=meta)
+
+
+def check_celery():
+    """
+    Checks to see if celery redis broker is connected.
+    Check worker stats to see if any workers are running
+    """
+    inspect = app.control.inspect()
+    try:
+        # Tries to connect to celery broker
+        conn = app.broker_connection().ensure_connection(max_retries=3)
+    except Exception as exc:
+        raise CeleryConnectionError(
+            "Failed to connect to celery redis broker!"
+        ) from exc
+    # Pings workers to see if any of them respond. Returns None if no response
+    stats = inspect.stats()
+    # NOTE: app.control.broadcast("ping", reply=True, limit=1) or inspect.ping() pings all workers but will not return if all workers are busy
+    if stats is None:
+        raise CeleryConnectionError(
+            "Celery workers are not responding. Check if workers are running!"
+        )

breadbox/breadbox/health_check/health_check.py

@@ -1,6 +1,8 @@
 from fastapi import APIRouter
 from breadbox.health_check import site_check_task
-from breadbox.celery_task.utils import format_task_status
+from breadbox.celery_task.utils import format_task_status, check_celery
+from breadbox.schemas.custom_http_exception import CeleryConnectionError
+
 import logging
 
 router = APIRouter(prefix="/health_check", tags=["health_check"])
@@ -26,6 +28,8 @@ def log_test():
     "/ok", operation_id="ok",
 )
 def ok():
+    check_celery()
+
     task = site_check_task.is_ok.delay()
     task.wait(timeout=60, interval=0.5)
 
@@ -35,3 +39,10 @@ def ok():
 @router.get("/simulate-error", operation_id="simulate_error")
 def simulate_error():
     raise Exception("Simulated error")
+
+
+@router.get("/celery", operation_id="celery_check")
+def celery_check():
+    check_celery()
+
+    return {"message": "ok"}

breadbox/breadbox/schemas/custom_http_exception.py

@@ -16,6 +16,7 @@ class HTTPError(BaseModel):
     status.HTTP_403_FORBIDDEN,
     status.HTTP_404_NOT_FOUND,
     status.HTTP_409_CONFLICT,
+    status.HTTP_503_SERVICE_UNAVAILABLE,
 ]
 ERROR_RESPONSES = dict.fromkeys(ERROR_CODES, {"model": HTTPError})
 
@@ -64,3 +65,8 @@ def __init__(self, msg):
 class ComputeLinearFitError(UserError):
     def __init__(self, msg):
         super().__init__(msg)
+
+
+class CeleryConnectionError(HTTPException):
+    def __init__(self, msg, error_code=503):
+        super().__init__(error_code, msg)

breadbox/tests/api/test_celery.py

@@ -12,6 +12,7 @@
 from breadbox.models.dataset import ValueType
 from breadbox.compute import analysis_tasks
 from breadbox.compute.analysis_tasks import _subset_feature_df
+from breadbox.celery_task import utils
 from tests import factories
 from breadbox.compute.analysis_tasks import (
     run_custom_analysis,
@@ -88,6 +89,12 @@ def mock_db_context(user, commit=True):
     def get_test_settings():
         return settings
 
+    def mock_check_celery():
+        return True
+
+    # Monkeypatch check_celery and pretend celery is running for test
+    monkeypatch.setattr(utils, "check_celery", mock_check_celery)
+
     monkeypatch.setattr(
         analysis_tasks, "get_settings", get_test_settings,
     )
@@ -226,6 +233,12 @@ def mock_db_context(user, commit=True):
     def get_test_settings():
         return settings
 
+    def mock_check_celery():
+        return True
+
+    # Monkeypatch check_celery and pretend celery is running for test
+    monkeypatch.setattr(utils, "check_celery", mock_check_celery)
+
     monkeypatch.setattr(
         analysis_tasks, "get_settings", get_test_settings,
     )

breadbox/tests/api/test_compute.py

@@ -21,6 +21,7 @@
 )
 from breadbox.compute import analysis_tasks
 from breadbox.compute.analysis_tasks import get_feature_data_slice_values
+from breadbox.celery_task import utils
 from breadbox.crud.dataset import get_dataset
 from tests import factories
 
@@ -691,6 +692,12 @@ def mock_db_context(user, **kwargs):
     def get_test_settings():
         return settings
 
+    def mock_check_celery():
+        return True
+
+    # Monkeypatch check_celery and pretend celery is running for test
+    monkeypatch.setattr(utils, "check_celery", mock_check_celery)
+
     monkeypatch.setattr(
         analysis_tasks, "get_settings", get_test_settings,
     )

breadbox/tests/conftest.py

@@ -226,6 +226,12 @@ def mock_db_context(user, commit=False):
     def get_test_settings():
         return settings
 
+    def mock_check_celery():
+        return True
+
+    # Monkeypatch check_celery and pretend celery is running for test
+    monkeypatch.setattr(utils, "check_celery", mock_check_celery)
+
     # The endpoint uses celery, and needs monkeypatching to replace db_context and get_settings,
     # which are not passed in as params due to the limits of redis serialization.
     monkeypatch.setattr(dataset_tasks, "db_context", mock_db_context)