celery_worker.py

from celery import Celery
import subprocess
import os


# Configure Celery to connect to the Redis broker
celery_app = Celery(
    "orion_worker",
    broker=os.getenv("CELERY_BROKER_URL", "redis://redis:6379/0"),
    backend=os.getenv("CELERY_RESULT_BACKEND", "redis://redis:6379/0"),
)

celery_app.conf.update(
    task_serializer="json",
    accept_content=["json"],
    result_serializer="json",
    task_track_started=True,
    broker_connection_retry_on_startup=True
)

# Define the task handler for orion.data_ingestion
# Make sure the name aligns with the task name used by celery in the upstream app
# so that the data ingestion task can be picked up from the task queue
@celery_app.task(name="orion.data_ingestion", queue="orion")
def run_build_manager(task_data):
    print(f'task_data: {task_data}', flush=True)
    # Resolve the additional graph spec path from the shared data directory and the supplied filename.
    shared_data_dir = os.getenv('SHARED_SOURCE_DATA_PATH')
    if not shared_data_dir:
        raise RuntimeError('SHARED_SOURCE_DATA_PATH is not set; cannot locate the additional graph spec.')
    additional_graph_spec = os.path.join(shared_data_dir, task_data["graph_spec_filename"])
    # no need to catch CalledProcessError exception, but rather let it propogate to Celery task handling
    result = subprocess.run(
        ["orion-build", task_data["graph_id"], "--graph_spec", additional_graph_spec],
        capture_output=True,
        text=True,
        check=True
    )
    return {"status": "success", "output": f'graph {task_data["graph_id"]} is created successfully'}