perf: stream report uploads instead of buffering the whole file in memory

lms/djangoapps/instructor_task/models.py

@@ -17,13 +17,14 @@
 import json
 import logging
 import os.path
+from io import BytesIO, TextIOWrapper
 from uuid import uuid4
 
 from botocore.exceptions import ClientError
 from django.apps import apps
 from django.conf import settings
 from django.contrib.auth.models import User  # pylint: disable=imported-auth-user
-from django.core.files.base import ContentFile
+from django.core.files.base import ContentFile, File
 from django.db import models, transaction
 from django.utils.translation import gettext as _
 from model_utils.models import TimeStampedModel
@@ -293,27 +294,57 @@ def store(self, course_id, filename, buff, parent_dir=''):
         Store the contents of `buff` in a directory determined by hashing
         `course_id`, and name the file `filename`. `buff` can be any file-like
         object, ready to be read from the beginning.
+
+        A binary buffer is handed to the storage backend as-is, so the backend
+        streams it -- S3Boto3Storage uploads in parts and FileSystemStorage
+        writes in chunks -- and peak memory stays at one chunk rather than the
+        whole report. This matters for grade reports, which can run to hundreds
+        of megabytes on large courses; previously the entire file was read into
+        memory, re-encoded, and copied into a ContentFile before upload, so an
+        on-disk report still cost roughly 3x its size in RAM at the final step.
+
+        A text buffer still works, but has to be read and encoded in full
+        because the storage backends require bytes. Callers handling
+        potentially large reports should pass a binary file.
         """
         path = self.path_to(course_id, filename, parent_dir)
+
+        if self._yields_bytes(buff):
+            self.storage.save(path, File(buff, name=filename))
+            return
+
         # See https://github.com/boto/boto/issues/2868
         # Boto doesn't play nice with unicode in python3
-        buff_contents = buff.read()
-
-        if not isinstance(buff_contents, bytes):
-            buff_contents = buff_contents.encode('utf-8')
+        self.storage.save(path, ContentFile(buff.read().encode('utf-8')))
 
-        buff = ContentFile(buff_contents)
+    @staticmethod
+    def _yields_bytes(buff):
+        """
+        Return True if reading from `buff` produces bytes rather than str.
 
-        self.storage.save(path, buff)
+        Probes with a zero-length read so the buffer is left positioned exactly
+        where it was, rather than inspecting the type -- callers pass a mix of
+        raw file objects, ContentFile and BytesIO, and mode attributes are not
+        consistently present across them.
+        """
+        try:
+            return isinstance(buff.read(0), bytes)
+        except (AttributeError, TypeError, ValueError):
+            return False
 
     def store_rows(self, course_id, filename, rows, parent_dir=''):
         """
         Given a course_id, filename, and rows (each row is an iterable of
         strings), write the rows to the storage backend in csv format.
         """
-        output_buffer = ContentFile('')
-        csvwriter = csv.writer(output_buffer)
+        output_buffer = BytesIO()
+        # newline='' per the csv module's contract; the writer emits its own
+        # line terminators and must not have them translated again.
+        text_wrapper = TextIOWrapper(output_buffer, encoding='utf-8', newline='', write_through=True)
+        csvwriter = csv.writer(text_wrapper)
         csvwriter.writerows(self._get_utf8_encoded_rows(rows))
+        # Detach so closing the wrapper does not close the buffer underneath it.
+        text_wrapper.detach()
         output_buffer.seek(0)
         self.store(course_id, filename, output_buffer, parent_dir)
 

lms/djangoapps/instructor_task/tasks_helper/grades.py

@@ -7,6 +7,7 @@
 import re
 from collections import OrderedDict, defaultdict
 from datetime import datetime
+from io import TextIOWrapper
 from itertools import chain
 from tempfile import TemporaryFile
 from time import time
@@ -351,7 +352,11 @@ def _generate(self):
         self.context.update_status('TemporaryFileReportMixin - 1: Starting grade report')
         batched_rows = self._batched_rows()
 
-        with TemporaryFile('r+') as success_file, TemporaryFile('r+') as error_file:
+        # Binary temp files, written through a text wrapper. The report store
+        # streams a binary buffer straight to the storage backend, where a text
+        # one has to be read into memory and encoded in full -- which would undo
+        # most of the benefit of spilling to disk in the first place.
+        with TemporaryFile('w+b') as success_file, TemporaryFile('w+b') as error_file:
             self.context.update_status('TemporaryFileReportMixin - 2: Compiling grades into temp files')
             has_errors = self.iter_and_write_batched_rows(batched_rows, success_file, error_file)
 
@@ -360,13 +365,24 @@ def _generate(self):
 
         return self.context.update_status('TemporaryFileReportMixin - 4: Completed grades')
 
+    @staticmethod
+    def _csv_writer_for(binary_file):
+        """
+        Return a csv.writer over a binary file, plus the wrapper to flush.
+
+        newline='' is the csv module's documented requirement: the writer emits
+        its own line terminators and they must not be translated again.
+        """
+        wrapper = TextIOWrapper(binary_file, encoding='utf-8', newline='', write_through=True)
+        return csv.writer(wrapper), wrapper
+
     def iter_and_write_batched_rows(self, batched_rows, success_file, error_file):
         """
         Iterate through batched rows, writing returned chunks to disk as we go.
         This should hopefully help us avoid out of memory errors.
         """
-        success_writer = csv.writer(success_file)
-        error_writer = csv.writer(error_file)
+        success_writer, success_wrapper = self._csv_writer_for(success_file)
+        error_writer, error_wrapper = self._csv_writer_for(error_file)
 
         # Write headers
         success_writer.writerow(self._success_headers())
@@ -386,6 +402,12 @@ def iter_and_write_batched_rows(self, batched_rows, success_file, error_file):
         self.context.task_progress.attempted = succeeded + failed
         self.context.task_progress.total = self.context.task_progress.attempted
 
+        # Detach rather than close: the wrappers must flush their buffered text
+        # into the temp files before those are read back for upload, but the
+        # temp files themselves stay open and are closed by the caller.
+        success_wrapper.detach()
+        error_wrapper.detach()
+
         return self.context.task_progress.failed > 0
 
     def upload_temp_files(self, success_file, error_file, has_errors):

lms/djangoapps/instructor_task/tests/test_models.py

@@ -5,7 +5,7 @@
 
 import copy
 import time
-from io import StringIO
+from io import BytesIO, StringIO
 
 import pytest
 from django.conf import settings
@@ -51,6 +51,69 @@ def create_report_store(self):
         """
         pass  # pylint: disable=unnecessary-pass
 
+    def test_store_streams_binary_buffer(self):
+        """
+        A binary buffer should reach the storage backend without first being
+        read into memory in its entirety.
+
+        Asserted behaviourally rather than by inspecting call args: a streaming
+        backend reads in bounded chunks (File.chunks / upload_fileobj both pass
+        an explicit size), so an unbounded read() is the signature of the whole
+        report being slurped into RAM before upload.
+        """
+        unbounded_reads = []
+
+        class _RecordingBytesIO(BytesIO):
+            """BytesIO that notes any read() not bounded by an explicit size."""
+            def read(self, size=-1, /):
+                if size is None or size < 0:
+                    unbounded_reads.append(size)
+                return super().read(size)
+
+        report_store = self.create_report_store()  # pylint: disable=assignment-from-no-return
+        payload = b'student_id,grade\n' + b'1,0.5\n' * 5000
+
+        report_store.store(self.course_id, 'streamed.csv', _RecordingBytesIO(payload))
+
+        assert not unbounded_reads, f'buffer was read without a size bound: {unbounded_reads}'
+        with report_store.storage.open(report_store.path_to(self.course_id, 'streamed.csv')) as stored:
+            assert stored.read() == payload
+
+    def test_store_text_buffer_round_trips(self):
+        """
+        Text buffers are still accepted, and are utf-8 encoded on the way out.
+
+        Not every caller has a binary file to hand, so this path has to keep
+        working even though it cannot stream.
+        """
+        report_store = self.create_report_store()  # pylint: disable=assignment-from-no-return
+        contents = 'student_id,grade\n1,0.5\nüser,1.0\n'
+
+        report_store.store(self.course_id, 'text.csv', StringIO(contents))
+
+        with report_store.storage.open(report_store.path_to(self.course_id, 'text.csv')) as stored:
+            assert stored.read() == contents.encode('utf-8')
+
+    def test_store_rows_round_trips(self):
+        """
+        store_rows() builds its CSV in a binary buffer, so it takes the
+        streaming path too. Verify the bytes on disk are unchanged by that.
+        """
+        report_store = self.create_report_store()  # pylint: disable=assignment-from-no-return
+
+        report_store.store_rows(
+            self.course_id,
+            'rows.csv',
+            [['student_id', 'grade'], [1, 0.5], ['üser', 'Not Attempted']],
+        )
+
+        with report_store.storage.open(report_store.path_to(self.course_id, 'rows.csv')) as stored:
+            assert stored.read().decode('utf-8').splitlines() == [
+                'student_id,grade',
+                '1,0.5',
+                'üser,Not Attempted',
+            ]
+
     def test_links_for_order(self):
         """
         Test that ReportStore.links_for() returns file download links