diff --git a/lms/djangoapps/instructor_task/models.py b/lms/djangoapps/instructor_task/models.py index c38282ca54c6..2bc44e2d8325 100644 --- a/lms/djangoapps/instructor_task/models.py +++ b/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) diff --git a/lms/djangoapps/instructor_task/tasks_helper/grades.py b/lms/djangoapps/instructor_task/tasks_helper/grades.py index 9b6a764137db..f8a5f1c23191 100644 --- a/lms/djangoapps/instructor_task/tasks_helper/grades.py +++ b/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): diff --git a/lms/djangoapps/instructor_task/tests/test_models.py b/lms/djangoapps/instructor_task/tests/test_models.py index 204f4d1fac90..b9ea380f299b 100644 --- a/lms/djangoapps/instructor_task/tests/test_models.py +++ b/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