67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
import time
|
|
import uuid
|
|
|
|
from django.conf import settings
|
|
from django.core.management.base import BaseCommand
|
|
|
|
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(f"task={task.task_id} status={task.status}")
|
|
if once:
|
|
return
|
|
continue
|
|
|
|
if once:
|
|
return
|
|
time.sleep(sleep_seconds)
|