Repository: mozilla/telemetry-analysis-service

Issue:
Use list_cluster EMR API for updating job run status
We currently facing lots of AWS API throttling when updating individual job run statuses before deciding to run them or not. We basically run one `describe_cluster` call per jobs per call against the `run_jobs` job.


Discussion:

Resolving patch:
diff --git a/atmo/jobs/jobs.py b/atmo/jobs/jobs.py
--- a/atmo/jobs/jobs.py
+++ b/atmo/jobs/jobs.py
@@ -1,9 +1,15 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, you can obtain one at http://mozilla.org/MPL/2.0/.
+from django.db import transaction
+import logging
 import newrelic.agent
 
 from .models import SparkJob
+from atmo.clusters.models import Cluster
+from atmo.clusters.provisioners import ClusterProvisioner
+
+logger = logging.getLogger(__name__)
 
 
 @newrelic.agent.background_task(group='RQ')
@@ -11,20 +17,52 @@ def run_jobs():
     """
     Run all the scheduled tasks that are supposed to run.
     """
+    # first let's update the job statuses if there are prior runs
     run_jobs = []
-    for job in SparkJob.objects.all():
-        # first let's update the status if there is a prior run
-        if job.latest_run:
-            job.latest_run.update_status()
-
-        # then let's check if the job should be run at all
-        if job.should_run():
-            job.run()
-            run_jobs.append(job.identifier)
-
-        # and then check if the job is expired and terminate it if needed
-        if job.is_expired:
-            # This shouldn't be required as we set a timeout in the bootstrap script,
-            # but let's keep it as a guard.
-            job.terminate()
+    jobs = SparkJob.objects.all()
+
+    # get the jobs with prior runs
+    jobs_with_active_runs = jobs.filter(
+        runs__status__in=Cluster.ACTIVE_STATUS_LIST,
+    ).prefetch_related('runs')
+    logger.debug('Updating Spark jobs: %s', jobs_with_active_runs)
+
+    # create a map between the jobflow ids of the latest runs and the jobs
+    jobflow_job_map = {
+        job.latest_run.jobflow_id: job
+        for job in jobs_with_active_runs
+    }
+    # get the created dates of the job runs to limit the ListCluster API call
+    provisioner = ClusterProvisioner()
+    runs_created_at = jobs_with_active_runs.datetimes('runs__created_at', 'day')
+    logger.debug('Fetching clusters older than %s', runs_created_at[0])
+
+    cluster_list = provisioner.list(created_after=runs_created_at[0])
+    logger.debug('Clusters found: %s', cluster_list)
+
+    for cluster_info in cluster_list:
+        # filter out the clusters that don't relate to the job run ids
+        job = jobflow_job_map.get(cluster_info['jobflow_id'], None)
+        if job is None:
+            continue
+        logger.debug('Updating job status for %s, latest run %s', job, job.latest_run)
+        # update the latest run status
+        with transaction.atomic():
+            job.latest_run.update_status(cluster_info)
+
+    for job in jobs:
+        with transaction.atomic():
+            # then let's check if the job should be run at all
+            should_run = job.should_run()
+            logger.debug('Checking if job %s should run: %s', job, should_run)
+            if should_run:
+                job.run()
+                run_jobs.append(job.identifier)
+
+            # and then check if the job is expired and terminate it if needed
+            if job.is_expired:
+                logger.debug('Job %s is expired and is terminated', job)
+                # This shouldn't be required as we set a timeout in the bootstrap script,
+                # but let's keep it as a guard.
+                job.terminate()
     return run_jobs
diff --git a/atmo/jobs/models.py b/atmo/jobs/models.py
--- a/atmo/jobs/models.py
+++ b/atmo/jobs/models.py
@@ -154,6 +154,21 @@ def notebook_s3_object(self):
     def get_absolute_url(self):
         return reverse('jobs-detail', kwargs={'id': self.id})
 
+    def is_due(self, now=None):
+        """
+        Whether the scheduled Spark job is due to be run based on the
+        latest run and the configured interval in hours.
+        """
+        if now is None:
+            now = timezone.now()
+        if not self.latest_run or self.latest_run.scheduled_date is None:
+            # job has never run before
+            hours_since_last_run = float('inf')
+        else:
+            hours_since_last_run = (now - self.latest_run.scheduled_date).total_seconds() // 3600
+
+        return hours_since_last_run >= self.interval_in_hours
+
     def should_run(self):
         """Whether the scheduled Spark job should run."""
         if not self.is_runnable:
@@ -162,14 +177,11 @@ def should_run(self):
         active = self.start_date <= now
         if self.end_date is not None:
             active = active and self.end_date >= now
-        if not self.latest_run or self.latest_run.scheduled_date is None:
-            # job has never run before
-            hours_since_last_run = float('inf')
-        else:
-            hours_since_last_run = (
-                (now - self.latest_run.scheduled_date).total_seconds() // 3600)
-        can_run_now = hours_since_last_run >= self.interval_in_hours
-        return self.is_enabled and active and can_run_now
+        return (
+            self.is_enabled and
+            active and
+            self.is_due(now)
+        )
 
     def run(self):
         """Actually run the scheduled Spark job."""
@@ -189,7 +201,7 @@ def run(self):
         run = self.runs.create(
             spark_job=self,
             jobflow_id=jobflow_id,
-            scheduled_date = timezone.now(),
+            scheduled_date=timezone.now(),
         )
         # Remove the cached latest run to this objects will requery it.
         try:
@@ -265,18 +277,18 @@ def __repr__(self):
     def get_info(self):
         return self.spark_job.cluster_provisioner.info(self.jobflow_id)
 
-    def update_status(self):
+    def update_status(self, info=None):
         """
         Updates latest status and life cycle datetimes.
         """
-        info = self.get_info()
-        if info is not None:
-            if self.status != info['state']:
-                self.status = info['state']
-                if self.status == Cluster.STATUS_RUNNING:
-                    self.run_date = timezone.now()
-                elif self.status in (Cluster.STATUS_TERMINATED,
-                                     Cluster.STATUS_TERMINATED_WITH_ERRORS):
-                    self.terminated_date = timezone.now()
-                self.save()
+        if info is None:
+            info = self.get_info()
+        if self.status != info['state']:
+            self.status = info['state']
+            if self.status == Cluster.STATUS_RUNNING:
+                self.run_date = timezone.now()
+            elif self.status in (Cluster.STATUS_TERMINATED,
+                                 Cluster.STATUS_TERMINATED_WITH_ERRORS):
+                self.terminated_date = timezone.now()
+            self.save()
         return self.status


Test patch: