diff --git a/inference/signals.py b/inference/signals.py index d662ff16..89e83f70 100644 --- a/inference/signals.py +++ b/inference/signals.py @@ -7,34 +7,52 @@ def create_periodic_tasks(sender, **kwargs): if sender.name == "inference": from django_celery_beat.models import CrontabSchedule, PeriodicTask - # Create schedule for every 5 minutes between 6pm-7am - crontab, _ = CrontabSchedule.objects.get_or_create( + # Weekday evening schedule (6pm-7am) + weekday_crontab, _ = CrontabSchedule.objects.get_or_create( minute="*/5", hour="18-23,0-6", - day_of_week="*", + day_of_week="1-5", # Monday-Friday day_of_month="*", month_of_year="*", ) - # Create the periodic task if it doesn't exist - PeriodicTask.objects.get_or_create( - crontab=crontab, - name="Process inference queue (6pm-7am)", - task="inference.tasks.process_inference_job_queue", - ) - - # Create schedule the weekend + # Weekend schedule (all day) weekend_crontab, _ = CrontabSchedule.objects.get_or_create( minute="*/5", hour="*", - day_of_week="0,6", + day_of_week="0,6", # Saturday-Sunday day_of_month="*", month_of_year="*", ) - # Create the periodic task if it doesn't exist - PeriodicTask.objects.get_or_create( - crontab=weekend_crontab, - name="Process inference queue (Weekends)", - task="inference.tasks.process_inference_job_queue", - ) + # Check if weekday task exists + weekday_task_name = "Process inference queue (6pm-7am)" + weekday_task_exists = PeriodicTask.objects.filter(name=weekday_task_name).exists() + + if not weekday_task_exists: + PeriodicTask.objects.create( + crontab=weekday_crontab, + name=weekday_task_name, + task="inference.tasks.process_inference_job_queue", + ) + else: + weekday_task = PeriodicTask.objects.get(name=weekday_task_name) + weekday_task.crontab = weekday_crontab + weekday_task.task = "inference.tasks.process_inference_job_queue" + weekday_task.save() + + # Check if weekend task exists + weekend_task_name = "Process inference queue (Weekends)" + weekend_task_exists = PeriodicTask.objects.filter(name=weekend_task_name).exists() + + if not weekend_task_exists: + PeriodicTask.objects.create( + crontab=weekend_crontab, + name=weekend_task_name, + task="inference.tasks.process_inference_job_queue", + ) + else: + weekend_task = PeriodicTask.objects.get(name=weekend_task_name) + weekend_task.crontab = weekend_crontab + weekend_task.task = "inference.tasks.process_inference_job_queue" + weekend_task.save()