A minimal celery worker client example

minimal_celery_worker_example.py

import fire
import pika
import redis
import os
from typing import Any
from typing import Dict
from typing import List
from typing import Tuple

import celery
from celery import Celery
from celery import Signature
from funcy import chunks
from halo import Halo
from kombu import Exchange
from kombu import Queue
from loguru import logger
from rich import print as rprint

POINT = float | int
BBOX = Tuple[POINT, POINT, POINT, POINT]
TOKEN_BOX = Tuple[BBOX, str]

"""
This requires: python 3.10
Dependencies can be installed with: 
pip install celery logguru rich fire 
"""


def connect_rabbit(host:str, port:int):
    parameters = pika.ConnectionParameters(host=host, port=port)
    try:
        connection = pika.BlockingConnection(parameters)
        if connection.is_open:
            rprint('[green]OK[/green]')
            connection.close()
    except Exception as error:
        rprint(f'[red]Could not connect to RabbitMQ @ {host}:{port} -> {error}[/red]')
        
        
def connect_redis(host:str, port:int):
    try:
        r = redis.Redis(host=host, port=port)
        r.ping()
        rprint('[green]OK[/green]')
    except Exception as error:
        rprint(f'[red]Could not connect to Redis @ {host}:{port} -> {error}[/red]')
        
        

# ============================== CELERY ==========================

# >>>>>>>>>>>> Set this two values <<<<<<<<<<<<<<<<<<<
CELERY_BROKER_URI = os.getenv("AMQP_URI")
CELERY_BACKEND_URI = os.getenv("REDIS_URI")
# >>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<<<<

CELERY_DEFAULT_TIMEOUT = 600
CELERY_DOC_2_PDF = "workers.preproc.conv.doc2pdf"
CELERY_OCR_MY_PDF = "workers.preproc.ocr.ocr_pdf"
CELERY_TESSERACT_PDF = "workers.preproc.ocr.tesseract_pdf"
CELERY_SPACY_NER = "workers.processors.ner.extract"
CELERY_CUAD_CLASSIFY = "workers.processors.cuad.classify"
CELERY_CUAD_EXTRACT = "workers.processors.cuad.extract"
CELERY_TEXT_DIFF = "workers.processors.diff.sdiff"
CELERY_PARAPHRASE_PREDICT = "workers.processors.paraphrase.predict"
CELERY_ENTAILMENT_PREDICT = "workers.processors.entailment.predict"


def route_task(name, args, kwargs, options, task=None, **kw):
    """Defines a router policy, assuming all task functions live in
    a structure like:
        workers.<worker-name>.<task-group>.<task-name>
    Here we take the task-group name as the Q-name
    """
    try:
        _, _, q_name, tname = name.split(".")
        logger.debug(f"Routing task: '{tname}' -> Q={q_name}")
    except ValueError:
        logger.warning(
            f"Received task name ({name}) does not conform to "
            "'workers.<worker-name>.<task-group>.<task-name>'. "
            "Routing by default to 'celery' queue"
        )
        q_name = "celery"

    return q_name



def remote_call(
    signature: str,
    args: Tuple[Any, ...] | None = (),
    kwargs: Dict[str, Any] | None = {},
    link_signature: Signature | str | None = None,
    timeout: int | None = CELERY_DEFAULT_TIMEOUT,
) -> Any:
    """Makes a remote call given the functions signature (task registered name).
    if a 'link_signature' is provided an AsyncResult is returned otherwise
    blocks until a result is ready and returns it.
    Args:
        signature (str): task registered name. e.g.: 'workers.processors.ner.run_ner'
        link_signature (str | None, optional): task callback signature's. Defaults to None.
        link_args (Tuple[Any] | None, optional): callback args. Defaults to None.
        link_kwargs (Dict[str, Any] | None, optional): callback kwargs. Defaults to None.
    """
    logger.info(f"Calling: {signature} (App: {app.main})")
    with Halo(f"Waiting for: {signature}"):
        try:
            # NOTE: When calling from a thread, celery won't see the app,
            # so we need to explicitly pass it along the signature creation.
            # Note that here we are using the **global app**
            fsig = celery.signature(signature, app=app)

            ares = fsig.apply_async(args=args, kwargs=kwargs, link=link_signature)
            if link_signature:
                return ares

            return ares.get(timeout=timeout)
        except Exception as e:
            logger.error(f"💥 Error in remote call: {e}")
            raise e



app = Celery(
    "lint",
    broker_url=os.getenv(CELERY_BROKER_URI),
    result_backend=os.getenv(CELERY_BACKEND_URI),
)

# Configure the exchange and queue for callbacks
app.conf.task_queues = [
    Queue("callbacks", Exchange("workers", type="topic"), routing_key="callbacks")
]

# Default route task
app.conf.task_routes = (route_task,)


# ============================== WORKERS ==========================


def ocr(pdf_file: str) -> Tuple[List[str], List[List[TOKEN_BOX]]]:
    """This functions performs OCR on a given PDF file. The file path
    must be common to both the client (this file) and the worker.
    i.e.: The EFS
    """
    rprint(f"🔍️ Running OCR for '{pdf_file}'")
    pages_text, per_page_boxes = remote_call(CELERY_TESSERACT_PDF, args=(pdf_file,))

    rprint(f"TEXT: {pages_text}")
    rprint(f"BBOXES: {per_page_boxes}") 


def clf(text:str):
    rprint(f"🔍️ Running CUAD classification in '{text}'")
    res = remote_call(CELERY_CUAD_CLASSIFY, args=(text))

    rprint(f"Result: {res}")


if __name__ == "__main__":
    fire.Fire({
        "ocr": ocr,
        "clf": clf,
        "rabbit": connect_rabbit,
        "redis": connect_redis
    })