cmhub/apps/api/management/commands/run_image_tasks.py

import time
import uuid

from django.conf import settings
from django.core.management.base import BaseCommand

from apps.api.image_tasks import (
    image_task_max_attempts,
    reap_stale_image_tasks,
    run_one_image_task,
)
from apps.api.models import ImageGenerationTask


class Command(BaseCommand):
    help = "Run asynchronous image generation tasks."

    def add_arguments(self, parser):
        parser.add_argument(
            "--once",
            action="store_true",
            help="Process at most one queued task and exit.",
        )
        parser.add_argument(
            "--worker-id",
            default="",
            help="Stable worker identifier recorded on claimed tasks.",
        )
        parser.add_argument(
            "--sleep-seconds",
            type=float,
            default=1.0,
            help="Sleep interval when no queued task is available.",
        )
        parser.add_argument(
            "--reap-only",
            action="store_true",
            help="Only reap stale running tasks and exit.",
        )

    def handle(self, *args, **options):
        worker_id = options["worker_id"] or f"image-worker-{uuid.uuid4().hex[:8]}"
        once = bool(options["once"])
        reap_only = bool(options["reap_only"])
        sleep_seconds = max(0.1, float(options["sleep_seconds"]))
        reaper_interval = max(
            1,
            int(getattr(settings, "IMAGE_TASK_REAPER_INTERVAL_SECONDS", 60)),
        )
        last_reap_at = 0.0

        while True:
            now = time.monotonic()
            if reap_only or now - last_reap_at >= reaper_interval:
                reaped = reap_stale_image_tasks()
                if reaped:
                    self.stdout.write(f"reaped={reaped}")
                last_reap_at = now
                if reap_only:
                    return

            task = run_one_image_task(worker_id=worker_id)
            if task is not None:
                self.stdout.write(format_task_log_line(task, started_at=now))
                if once:
                    return
                continue

            if once:
                return
            time.sleep(sleep_seconds)


def format_task_log_line(task: ImageGenerationTask, *, started_at: float) -> str:
    duration_ms = max(0, int((time.monotonic() - started_at) * 1000))
    alias = str((task.request_payload or {}).get("model") or "")
    retrying = bool(
        task.status == ImageGenerationTask.Status.QUEUED
        and task.error_code
        and task.next_attempt_at
    )
    fields = {
        "event": "image_task_processed",
        "task_id": str(task.task_id),
        "status": task.status,
        "alias": alias,
        "attempt": str(task.attempt_count),
        "max_attempts": str(image_task_max_attempts()),
        "retrying": str(retrying).lower(),
        "next_attempt_at": task.next_attempt_at.isoformat() if task.next_attempt_at else "",
        "duration_ms": str(duration_ms),
    }
    if retrying or task.status in {
        ImageGenerationTask.Status.FAILED,
        ImageGenerationTask.Status.EXPIRED,
    }:
        fields["error_code"] = task.error_code or "upstream_error"
    return " ".join(f"{key}={value}" for key, value in fields.items())