Add a job cancelled status, shutdown behavior, some atomic usage

Author: davegaeddert

bolt-jobs/bolt/jobs/admin.py

@@ -30,7 +30,7 @@ def _td_format(td_object):
         if seconds > period_seconds:
             period_value, seconds = divmod(seconds, period_seconds)
             has_s = "s" if period_value > 1 else ""
-            strings.append("%s %s%s" % (period_value, period_name, has_s))
+            strings.append(f"{period_value} {period_name}{has_s}")
 
     return ", ".join(strings)
 
@@ -147,6 +147,7 @@ class ListView(AdminModelListView):
         filters = [
             "Successful",
             "Errored",
+            "Cancelled",
             "Lost",
             "Retried",
         ]
@@ -166,6 +167,8 @@ def get_initial_queryset(self):
                 return queryset.successful()
             if self.filter == "Errored":
                 return queryset.errored()
+            if self.filter == "Cancelled":
+                return queryset.cancelled()
             if self.filter == "Lost":
                 return queryset.lost()
             if self.filter == "Retried":

bolt-jobs/bolt/jobs/cli.py

@@ -1,5 +1,6 @@
 import datetime
 import logging
+import signal
 
 import click
 
@@ -40,11 +41,22 @@ def cli():
     envvar="BOLT_JOBS_STATS_EVERY",
 )
 def worker(max_processes, max_jobs_per_process, stats_every):
-    Worker(
+    worker = Worker(
         max_processes=max_processes,
         max_jobs_per_process=max_jobs_per_process,
         stats_every=stats_every,
-    ).run()
+    )
+
+    def _shutdown(signalnum, _):
+        logger.info("Job worker shutdown signal received signalnum=%s", signalnum)
+        worker.shutdown()
+
+    # Allow the worker to be stopped gracefully on SIGTERM
+    signal.signal(signal.SIGTERM, _shutdown)
+    signal.signal(signal.SIGINT, _shutdown)
+
+    # Start processing jobs
+    worker.run()
 
 
 @cli.command()

bolt-jobs/bolt/jobs/migrations/0015_job_worker_uuid_jobresult_worker_uuid_and_more.py

@@ -0,0 +1,36 @@
+# Generated by Bolt 5.0.dev20240118212224 on 2024-01-18 21:53
+
+from bolt.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("boltqueue", "0014_job_unique_key_jobrequest_unique_key_and_more"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="job",
+            name="worker_uuid",
+            field=models.UUIDField(blank=True, db_index=True, null=True),
+        ),
+        migrations.AddField(
+            model_name="jobresult",
+            name="worker_uuid",
+            field=models.UUIDField(blank=True, db_index=True, null=True),
+        ),
+        migrations.AlterField(
+            model_name="jobresult",
+            name="status",
+            field=models.CharField(
+                choices=[
+                    ("SUCCESSFUL", "Successful"),
+                    ("ERRORED", "Errored"),
+                    ("CANCELLED", "Cancelled"),
+                    ("LOST", "Lost"),
+                ],
+                db_index=True,
+                max_length=20,
+            ),
+        ),
+    ]

bolt-jobs/bolt/jobs/models.py

@@ -61,24 +61,26 @@ class Meta:
     def __str__(self):
         return f"{self.job_class} [{self.uuid}]"
 
-    def convert_to_job(self):
+    def convert_to_job(self, *, worker_uuid=None):
         """
         JobRequests are the pending jobs that are waiting to be executed.
         We immediately convert them to JobResults when they are picked up.
         """
-        result = Job.objects.create(
-            job_request_uuid=self.uuid,
-            job_class=self.job_class,
-            parameters=self.parameters,
-            priority=self.priority,
-            source=self.source,
-            retries=self.retries,
-            retry_attempt=self.retry_attempt,
-            unique_key=self.unique_key,
-        )
+        with transaction.atomic():
+            result = Job.objects.create(
+                job_request_uuid=self.uuid,
+                job_class=self.job_class,
+                parameters=self.parameters,
+                priority=self.priority,
+                source=self.source,
+                retries=self.retries,
+                retry_attempt=self.retry_attempt,
+                unique_key=self.unique_key,
+                worker_uuid=worker_uuid,
+            )
 
-        # Delete the pending JobRequest now
-        self.delete()
+            # Delete the pending JobRequest now
+            self.delete()
 
         return result
 
@@ -110,6 +112,9 @@ class Job(models.Model):
     created_at = models.DateTimeField(auto_now_add=True, db_index=True)
     started_at = models.DateTimeField(blank=True, null=True, db_index=True)
 
+    # To associate with a worker
+    worker_uuid = models.UUIDField(blank=True, null=True, db_index=True)
+
     # From the JobRequest
     job_request_uuid = models.UUIDField(db_index=True)
     job_class = models.CharField(max_length=255, db_index=True)
@@ -152,26 +157,29 @@ def convert_to_result(self, *, status, error=""):
         """
         Convert this Job to a JobResult.
         """
-        result = JobResult.objects.create(
-            ended_at=timezone.now(),
-            error=error,
-            status=status,
-            # From the Job
-            job_uuid=self.uuid,
-            started_at=self.started_at,
-            # From the JobRequest
-            job_request_uuid=self.job_request_uuid,
-            job_class=self.job_class,
-            parameters=self.parameters,
-            priority=self.priority,
-            source=self.source,
-            retries=self.retries,
-            retry_attempt=self.retry_attempt,
-            unique_key=self.unique_key,
-        )
+        with transaction.atomic():
+            result = JobResult.objects.create(
+                ended_at=timezone.now(),
+                error=error,
+                status=status,
+                # From the worker
+                worker_uuid=self.worker_uuid,
+                # From the Job
+                job_uuid=self.uuid,
+                started_at=self.started_at,
+                # From the JobRequest
+                job_request_uuid=self.job_request_uuid,
+                job_class=self.job_class,
+                parameters=self.parameters,
+                priority=self.priority,
+                source=self.source,
+                retries=self.retries,
+                retry_attempt=self.retry_attempt,
+                unique_key=self.unique_key,
+            )
 
-        # Delete the Job now
-        self.delete()
+            # Delete the Job now
+            self.delete()
 
         return result
 
@@ -180,6 +188,9 @@ class JobResultQuerySet(models.QuerySet):
     def successful(self):
         return self.filter(status=JobResultStatuses.SUCCESSFUL)
 
+    def cancelled(self):
+        return self.filter(status=JobResultStatuses.CANCELLED)
+
     def lost(self):
         return self.filter(status=JobResultStatuses.LOST)
 
@@ -192,19 +203,31 @@ def retried(self):
             | models.Q(retry_attempt__gt=0)
         )
 
-    def retry_failed_jobs(self):
-        for result in self.filter(
-            status__in=[JobResultStatuses.ERRORED, JobResultStatuses.LOST],
+    def failed(self):
+        return self.filter(
+            status__in=[
+                JobResultStatuses.ERRORED,
+                JobResultStatuses.LOST,
+                JobResultStatuses.CANCELLED,
+            ]
+        )
+
+    def retryable(self):
+        return self.failed().filter(
             retry_job_request_uuid__isnull=True,
             retries__gt=0,
             retry_attempt__lt=models.F("retries"),
-        ):
+        )
+
+    def retry_failed_jobs(self):
+        for result in self.retryable():
             result.retry_job()
 
 
 class JobResultStatuses(models.TextChoices):
     SUCCESSFUL = "SUCCESSFUL", "Successful"
     ERRORED = "ERRORED", "Errored"  # Threw an error
+    CANCELLED = "CANCELLED", "Cancelled"  # Cancelled (probably by deploy)
     LOST = (
         "LOST",
         "Lost",
@@ -219,6 +242,9 @@ class JobResult(models.Model):
     uuid = models.UUIDField(default=uuid.uuid4, editable=False, unique=True)
     created_at = models.DateTimeField(auto_now_add=True, db_index=True)
 
+    # To associate with a worker
+    worker_uuid = models.UUIDField(blank=True, null=True, db_index=True)
+
     # From the Job
     job_uuid = models.UUIDField(db_index=True)
     started_at = models.DateTimeField(blank=True, null=True, db_index=True)

bolt-jobs/bolt/jobs/workers.py

@@ -2,13 +2,14 @@
 import multiprocessing
 import os
 import time
+import uuid
 from concurrent.futures import ProcessPoolExecutor
 
 from bolt.db import transaction
 from bolt.logs import app_logger
 from bolt.signals import request_finished, request_started
 
-from .models import Job, JobRequest, JobResult
+from .models import Job, JobRequest, JobResult, JobResultStatuses
 
 logger = logging.getLogger("bolt.jobs")
 
@@ -27,49 +28,71 @@ def __init__(self, max_processes=None, max_jobs_per_process=None, stats_every=No
         self.max_processes = self.executor._max_workers
         self.max_jobs_per_process = max_jobs_per_process
 
+        self.uuid = uuid.uuid4()
+
     def run(self):
         logger.info(
             "Starting job worker with %s max processes",
             self.max_processes,
         )
 
-        try:
-            while True:
-                try:
-                    self.maybe_log_stats()
-                    self.maybe_check_job_results()
-                except Exception as e:
-                    # Log the issue, but don't stop the worker
-                    # (these tasks are kind of ancilarry to the main job processing)
-                    logger.exception(e)
-
-                with transaction.atomic():
-                    job_request = JobRequest.objects.next_up()
-                    if not job_request:
-                        # Potentially no jobs to process (who knows for how long)
-                        # but sleep for a second to give the CPU and DB a break
-                        time.sleep(1)
-                        continue
-
-                    logger.info(
-                        'Preparing to execute job job_class=%s job_request_uuid=%s job_priority=%s job_source="%s"',
-                        job_request.job_class,
-                        job_request.uuid,
-                        job_request.priority,
-                        job_request.source,
-                    )
-
-                    job = job_request.convert_to_job()
-
-                job_uuid = str(job.uuid)  # Make a str copy
-
-                # Release these now
-                del job_request
-                del job
-
-                self.executor.submit(process_job, job_uuid)
-        except (KeyboardInterrupt, SystemExit):
-            self.executor.shutdown(wait=True, cancel_futures=True)
+        while True:
+            try:
+                self.maybe_log_stats()
+                self.maybe_check_job_results()
+            except Exception as e:
+                # Log the issue, but don't stop the worker
+                # (these tasks are kind of ancilarry to the main job processing)
+                logger.exception(e)
+
+            with transaction.atomic():
+                job_request = JobRequest.objects.next_up()
+                if not job_request:
+                    # Potentially no jobs to process (who knows for how long)
+                    # but sleep for a second to give the CPU and DB a break
+                    time.sleep(1)
+                    continue
+
+                logger.info(
+                    'Preparing to execute job job_class=%s job_request_uuid=%s job_priority=%s job_source="%s"',
+                    job_request.job_class,
+                    job_request.uuid,
+                    job_request.priority,
+                    job_request.source,
+                )
+
+                job = job_request.convert_to_job(worker_uuid=self.uuid)
+
+            job_uuid = str(job.uuid)  # Make a str copy
+
+            # Release these now
+            del job_request
+            del job
+
+            self.executor.submit(process_job, job_uuid)
+
+    def shutdown(self):
+        # Prevent duplicate keyboard and sigterm calls
+        # (the way this works also only lets us shutdown once per instance)
+        if getattr(self, "_is_shutting_down", False):
+            return
+        self._is_shutting_down = True
+
+        logger.info("Job worker shutdown started")
+
+        # Make an attept to immediatelly move any unstarted jobs from this worker
+        # to the JobResults as cancelled. If they have a retry, they'll be picked
+        # up by the next worker process.
+        for job in Job.objects.filter(worker_uuid=self.uuid, started_at__isnull=True):
+            job.convert_to_result(status=JobResultStatuses.CANCELLED)
+
+        # Now shutdown the process pool.
+        # There's still some chance of a race condition here where we
+        # just deleted a Job and it was still picked up, but we'll log
+        # missing jobs as warnings instead of exceptions.
+        self.executor.shutdown(wait=True, cancel_futures=True)
+
+        logger.info("Job worker shutdown complete")
 
     def maybe_log_stats(self):
         if not self.stats_every:
@@ -124,7 +147,16 @@ def process_job(job_uuid):
         worker_pid = os.getpid()
 
         request_started.send(sender=None)
-        job = Job.objects.get(uuid=job_uuid)
+
+        try:
+            job = Job.objects.get(uuid=job_uuid)
+        except Job.DoesNotExist:
+            logger.warning(
+                "Job not found worker_pid=%s job_uuid=%s",
+                worker_pid,
+                job_uuid,
+            )
+            return
 
         logger.info(
             'Executing job worker_pid=%s job_class=%s job_request_uuid=%s job_priority=%s job_source="%s"',