tasks.py

import datetime
import functools
import logging
import time
from collections.abc import Callable

from asgiref.sync import async_to_sync
from celery.signals import task_failure, task_postrun, task_prerun
from django.db import transaction

from ami.ml.orchestration.nats_queue import TaskQueueManager
from ami.ml.orchestration.task_state import TaskStateManager
from ami.ml.schemas import PipelineResultsResponse
from ami.tasks import default_soft_time_limit, default_time_limit
from config import celery_app

logger = logging.getLogger(__name__)


@celery_app.task(bind=True, soft_time_limit=default_soft_time_limit, time_limit=default_time_limit)
def run_job(self, job_id: int) -> None:
    from ami.jobs.models import Job

    try:
        job = Job.objects.get(pk=job_id)
    except Job.DoesNotExist as e:
        raise e
        # self.retry(exc=e, countdown=1, max_retries=1)
    else:
        job.logger.info(f"Running job {job}")
        try:
            job.run()
        except Exception as e:
            job.logger.error(f'Job #{job.pk} "{job.name}" failed: {e}')
            raise
        else:
            job.refresh_from_db()
            job.logger.info(f"Finished job {job}")


@celery_app.task(
    bind=True,
    max_retries=0,  # don't retry since we already have retry logic in the NATS queue
    soft_time_limit=300,  # 5 minutes
    time_limit=360,  # 6 minutes
)
def process_pipeline_result(self, job_id: int, result_data: dict, reply_subject: str) -> None:
    """
    Process a single pipeline result asynchronously.

    This task:
    1. Deserializes the pipeline result
    2. Saves it to the database
    3. Updates progress by removing processed image IDs from Redis
    4. Acknowledges the task via NATS

    Args:
        job_id: The job ID
        result_data: Dictionary containing the pipeline result
        reply_subject: NATS reply subject for acknowledgment
    """
    from ami.jobs.models import Job  # avoid circular import

    _, t = log_time()
    error = result_data.get("error")
    pipeline_result = None
    if not error:
        pipeline_result = PipelineResultsResponse(**result_data)
        processed_image_ids = {str(img.id) for img in pipeline_result.source_images}
    else:
        image_id = result_data.get("image_id")
        processed_image_ids = {str(image_id)} if image_id else set()
        logger.error(f"Pipeline returned error for job {job_id}, image {image_id}: {error}")

    state_manager = TaskStateManager(job_id)

    progress_info = state_manager.update_state(processed_image_ids, stage="process", request_id=self.request.id)
    if not progress_info:
        logger.warning(
            f"Another task is already processing results for job {job_id}. "
            f"Retrying task {self.request.id} in 5 seconds..."
        )
        raise self.retry(countdown=5, max_retries=10)

    try:
        _update_job_progress(job_id, "process", progress_info.percentage)

        _, t = t(f"TIME: Updated job {job_id} progress in PROCESS stage progress to {progress_info.percentage*100}%")
        job = Job.objects.get(pk=job_id)
        job.logger.info(f"Processing pipeline result for job {job_id}, reply_subject: {reply_subject}")
        job.logger.info(
            f" Job {job_id} progress: {progress_info.processed}/{progress_info.total} images processed "
            f"({progress_info.percentage*100}%), {progress_info.remaining} remaining, {len(processed_image_ids)} just "
            "processed"
        )
    except Job.DoesNotExist:
        # don't raise and ack so that we don't retry since the job doesn't exists
        logger.error(f"Job {job_id} not found")
        _ack_task_via_nats(reply_subject, logger)
        return

    try:
        # Save to database (this is the slow operation)
        if pipeline_result:
            # should never happen since otherwise we could not be processing results here
            assert job.pipeline is not None, "Job pipeline is None"
            job.pipeline.save_results(results=pipeline_result, job_id=job.pk)
            job.logger.info(f"Successfully saved results for job {job_id}")

            _, t = t(
                f"Saved pipeline results to database with {len(pipeline_result.detections)} detections"
                f", percentage: {progress_info.percentage*100}%"
            )

        _ack_task_via_nats(reply_subject, job.logger)
        # Update job stage with calculated progress
        progress_info = state_manager.update_state(processed_image_ids, stage="results", request_id=self.request.id)

        if not progress_info:
            logger.warning(
                f"Another task is already processing results for job {job_id}. "
                f"Retrying task {self.request.id} in 5 seconds..."
            )
            raise self.retry(countdown=5, max_retries=10)
        _update_job_progress(job_id, "results", progress_info.percentage)

    except Exception as e:
        job.logger.error(
            f"Failed to process pipeline result for job {job_id}: {e}. NATS will redeliver the task message."
        )


def _ack_task_via_nats(reply_subject: str, job_logger: logging.Logger) -> None:
    try:

        async def ack_task():
            async with TaskQueueManager() as manager:
                return await manager.acknowledge_task(reply_subject)

        ack_success = async_to_sync(ack_task)()

        if ack_success:
            job_logger.info(f"Successfully acknowledged task via NATS: {reply_subject}")
        else:
            job_logger.warning(f"Failed to acknowledge task via NATS: {reply_subject}")
    except Exception as ack_error:
        job_logger.error(f"Error acknowledging task via NATS: {ack_error}")
        # Don't fail the task if ACK fails - data is already saved


def _update_job_progress(job_id: int, stage: str, progress_percentage: float) -> None:
    from ami.jobs.models import Job, JobState  # avoid circular import

    with transaction.atomic():
        job = Job.objects.select_for_update().get(pk=job_id)
        job.progress.update_stage(
            stage,
            status=JobState.SUCCESS if progress_percentage >= 1.0 else JobState.STARTED,
            progress=progress_percentage,
        )
        if stage == "results" and progress_percentage >= 1.0:
            job.status = JobState.SUCCESS
            job.progress.summary.status = JobState.SUCCESS
            job.finished_at = datetime.datetime.now()  # Use naive datetime in local time
        job.logger.info(f"Updated job {job_id} progress in stage '{stage}' to {progress_percentage*100}%")
        job.save()


@task_prerun.connect(sender=run_job)
def pre_update_job_status(sender, task_id, task, **kwargs):
    # in the prerun signal, set the job status to PENDING
    update_job_status(sender, task_id, task, "PENDING", **kwargs)


@task_postrun.connect(sender=run_job)
def update_job_status(sender, task_id, task, state: str, retval=None, **kwargs):
    from ami.jobs.models import Job

    job_id = task.request.kwargs["job_id"]
    if job_id is None:
        logger.error(f"Job id is None for task {task_id}")
        return
    try:
        job = Job.objects.get(pk=job_id)
    except Job.DoesNotExist:
        try:
            job = Job.objects.get(task_id=task_id)
        except Job.DoesNotExist:
            logger.error(f"No job found for task {task_id} or job_id {job_id}")
            return

    job.update_status(state)


@task_failure.connect(sender=run_job, retry=False)
def update_job_failure(sender, task_id, exception, *args, **kwargs):
    from ami.jobs.models import Job, JobState

    job = Job.objects.get(task_id=task_id)
    job.update_status(JobState.FAILURE, save=False)

    job.logger.error(f'Job #{job.pk} "{job.name}" failed: {exception}')

    job.save()


def log_time(start: float = 0, msg: str | None = None) -> tuple[float, Callable]:
    """
    Small helper to measure time between calls.

    Returns: elapsed time since the last call, and a partial function to measure from the current call
    Usage:

    _, tlog = log_time()
    # do something
    _, tlog = tlog("Did something") # will log the time taken by 'something'
    # do something else
    t, tlog = tlog("Did something else") # will log the time taken by 'something else', returned as 't'
    """

    end = time.perf_counter()
    if start == 0:
        dur = 0.0
    else:
        dur = end - start
    if msg and start > 0:
        logger.info(f"{msg}: {dur:.3f}s")
    new_start = time.perf_counter()
    return dur, functools.partial(log_time, new_start)