83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
import time
|
|
import uuid
|
|
|
|
from django.conf import settings
|
|
from django.core.management.base import BaseCommand
|
|
|
|
from apps.api.models import ImageGenerationTask
|
|
from apps.api.image_tasks import reap_stale_image_tasks, run_one_image_task
|
|
|
|
|
|
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 "")
|
|
fields = {
|
|
"event": "image_task_processed",
|
|
"task_id": str(task.task_id),
|
|
"status": task.status,
|
|
"alias": alias,
|
|
"duration_ms": str(duration_ms),
|
|
}
|
|
if 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())
|