perf(grades): optimize database queries for large-scale grade recalcu…

lms/djangoapps/grades/management/commands/compute_grades.py

@@ -111,7 +111,7 @@ def _shuffled_task_kwargs(self, options):
         for args in all_args:
             yield {
                 'course_key': args[0],
-                'offset': args[1],
+                'start_id': args[1],
                 'batch_size': args[2],
                 'estimate_first_attempted': estimate_first_attempted,
             }

lms/djangoapps/grades/management/commands/tests/test_compute_grades.py

@@ -79,11 +79,19 @@ def test_tasks_fired(self, estimate_first_attempted, mock_task):
             command.append('--no_estimate_first_attempted')
         call_command(*(command + courses))
 
-        def _kwargs(course_key, offset):
+        def get_enrollment_ids(course_key):
+            return list(CourseEnrollment.objects.filter(
+                course_id=course_key
+            ).order_by('id').values_list('id', flat=True))
+
+        ids_course_0 = get_enrollment_ids(self.course_keys[0])
+        ids_course_3 = get_enrollment_ids(self.course_keys[3])
+
+        def _kwargs(course_key, start_id):
             return {
                 'course_key': course_key,
                 'batch_size': 2,
-                'offset': offset,
+                'start_id': start_id,
                 'estimate_first_attempted': estimate_first_attempted,
                 'seq_id': ANY
             }
@@ -93,19 +101,19 @@ def _kwargs(course_key, offset):
         expected = [
             ({
                 'queue': 'key',
-                'kwargs': _kwargs(self.course_keys[0], 0)
+                'kwargs': _kwargs(self.course_keys[0], ids_course_0[0])
             },),
             ({
                 'queue': 'key',
-                'kwargs': _kwargs(self.course_keys[0], 2)
+                'kwargs': _kwargs(self.course_keys[0], ids_course_0[2])
             },),
             ({
                 'queue': 'key',
-                'kwargs': _kwargs(self.course_keys[3], 0)
+                'kwargs': _kwargs(self.course_keys[3], ids_course_3[0])
             },),
             ({
                 'queue': 'key',
-                'kwargs': _kwargs(self.course_keys[3], 2)
+                'kwargs': _kwargs(self.course_keys[3], ids_course_3[2])
             },),
         ]
         assert len(expected) == len(actual)
@@ -116,14 +124,18 @@ def _kwargs(course_key, offset):
     def test_tasks_fired_from_settings(self, mock_task):
         ComputeGradesSetting.objects.create(course_ids=self.course_keys[1], batch_size=2)
         call_command('compute_grades', '--from_settings')
+
+        ids_course_1 = list(CourseEnrollment.objects.filter(
+            course_id=self.course_keys[1]
+        ).order_by('id').values_list('id', flat=True))
         actual = mock_task.apply_async.call_args_list
         # Order doesn't matter, but can't use a set because dicts aren't hashable
         expected = [
             ({
                 'kwargs': {
                     'course_key': self.course_keys[1],
                     'batch_size': 2,
-                    'offset': 0,
+                    'start_id': ids_course_1[0],
                     'estimate_first_attempted': True,
                     'seq_id': ANY,
                 },
@@ -132,7 +144,7 @@ def test_tasks_fired_from_settings(self, mock_task):
                 'kwargs': {
                     'course_key': self.course_keys[1],
                     'batch_size': 2,
-                    'offset': 2,
+                    'start_id': ids_course_1[2],
                     'estimate_first_attempted': True,
                     'seq_id': ANY,
                 },

lms/djangoapps/grades/tasks.py

@@ -69,10 +69,10 @@ def compute_all_grades_for_course(**kwargs):
         if are_grades_frozen(course_key):
             log.info("Attempted compute_all_grades_for_course for course '%s', but grades are frozen.", course_key)
             return
-        for course_key_string, offset, batch_size in _course_task_args(course_key=course_key, **kwargs):
+        for course_key_string, start_id, batch_size in _course_task_args(course_key=course_key, **kwargs):
             kwargs.update({
                 'course_key': course_key_string,
-                'offset': offset,
+                'start_id': start_id,
                 'batch_size': batch_size,
             })
             compute_grades_for_course_v2.apply_async(
@@ -107,14 +107,14 @@ def compute_grades_for_course_v2(self, **kwargs):
         set_event_transaction_type(kwargs['event_transaction_type'])
 
     try:
-        return compute_grades_for_course(kwargs['course_key'], kwargs['offset'], kwargs['batch_size'])
+        return compute_grades_for_course(kwargs['course_key'], kwargs['start_id'], kwargs['batch_size'])
     except Exception as exc:
         raise self.retry(kwargs=kwargs, exc=exc)  # noqa: B904
 
 
 @shared_task(base=LoggedPersistOnFailureTask)
 @set_code_owner_attribute
-def compute_grades_for_course(course_key, offset, batch_size, **kwargs):  # pylint: disable=unused-argument
+def compute_grades_for_course(course_key, start_id, batch_size, **kwargs):  # pylint: disable=unused-argument
     """
     Compute and save grades for a set of students in the specified course.
 
@@ -127,8 +127,12 @@ def compute_grades_for_course(course_key, offset, batch_size, **kwargs):  # pyli
         log.info("Attempted compute_grades_for_course for course '%s', but grades are frozen.", course_key)
         return
 
-    enrollments = CourseEnrollment.objects.filter(course_id=course_key).order_by('created')
-    student_iter = (enrollment.user for enrollment in enrollments[offset:offset + batch_size])
+    enrollments = CourseEnrollment.objects.filter(
+        course_id=course_key,
+        id__gte=start_id,
+    ).select_related('user').order_by('id')[:batch_size]
+    student_iter = (enrollment.user for enrollment in enrollments)
+
     for result in CourseGradeFactory().iter(users=student_iter, course_key=course_key, force_update=True):
         if result.error is not None:
             raise result.error
@@ -364,14 +368,20 @@ def _course_task_args(course_key, **kwargs):
     Helper function to generate course-grade task args.
     """
     from_settings = kwargs.pop('from_settings', True)
-    enrollment_count = CourseEnrollment.objects.filter(course_id=course_key).count()
-    if enrollment_count == 0:
+    enrollment_ids = list(
+        CourseEnrollment.objects.filter(course_id=course_key)
+        .order_by('id')
+        .values_list('id', flat=True)
+    )
+    total_count = len(enrollment_ids)
+
+    if total_count == 0:
         log.warning(f"No enrollments found for {course_key}")
 
     if from_settings is False:
         batch_size = kwargs.pop('batch_size', 100)
     else:
         batch_size = ComputeGradesSetting.current().batch_size
 
-    for offset in range(0, enrollment_count, batch_size):
-        yield (str(course_key), offset, batch_size)
+    for batch_start_id in enrollment_ids[::batch_size]:
+        yield (str(course_key), batch_start_id, batch_size)

lms/djangoapps/grades/tests/test_tasks.py

@@ -429,26 +429,34 @@ def setUp(self):
 
     @ddt.data(*range(0, 12, 3))
     def test_behavior(self, batch_size):
+        enrollment_ids = list(CourseEnrollment.objects.filter(
+            course_id=self.course.id
+        ).order_by('id').values_list('id', flat=True))
+        start_id = enrollment_ids[4]
+
         with mock_get_score(1, 2):
             result = compute_grades_for_course_v2.delay(
                 course_key=str(self.course.id),
                 batch_size=batch_size,
-                offset=4,
+                start_id=start_id,
             )
         assert result.successful
         assert PersistentCourseGrade.objects.filter(course_id=self.course.id).count() == min(batch_size, 8)
         assert PersistentSubsectionGrade.objects.filter(course_id=self.course.id).count() == min(batch_size, 8)
 
     @ddt.data(*range(1, 12, 3))
     def test_course_task_args(self, test_batch_size):
-        offset_expected = 0
-        for course_key, offset, batch_size in _course_task_args(
+        enrollment_ids = list(CourseEnrollment.objects.filter(
+            course_id=self.course.id
+        ).order_by('id').values_list('id', flat=True))
+        current_index = 0
+        for course_key, start_id, batch_size in _course_task_args(
             batch_size=test_batch_size, course_key=self.course.id, from_settings=False
         ):
             assert course_key == str(self.course.id)
             assert batch_size == test_batch_size
-            assert offset == offset_expected
-            offset_expected += test_batch_size
+            assert start_id == enrollment_ids[current_index]
+            current_index += test_batch_size
 
 
 class RecalculateGradesForUserTest(HasCourseWithProblemsMixin, ModuleStoreTestCase):
@@ -575,6 +583,11 @@ def test_compute_grades_for_course(self, freeze_flag_value, end_date_adjustment,
         for user in self.users:
             CourseEnrollment.enroll(user, self.course.id)
 
+        enrollment_ids = list(CourseEnrollment.objects.filter(
+            course_id=self.course.id
+        ).order_by('id').values_list('id', flat=True))
+        start_id = enrollment_ids[4]
+
         with override_waffle_flag(self.freeze_grade_flag, active=freeze_flag_value):
             with patch('lms.djangoapps.grades.tasks.CourseGradeFactory') as mock_factory:
                 factory = mock_factory.return_value
@@ -583,7 +596,7 @@ def test_compute_grades_for_course(self, freeze_flag_value, end_date_adjustment,
                         kwargs={
                             'course_key': str(self.course.id),
                             'batch_size': 2,
-                            'offset': 4,
+                            'start_id': start_id,
                         }
                     )
                     self._assert_for_freeze_grade_flag(