Reduce preview size for large Jupyter notebooks (#1822)

lambdas/preview/index.py

@@ -4,6 +4,7 @@
 
 Lambda functions can have up to 3GB of RAM and only 512MB of disk.
 """
+import os
 import io
 from contextlib import redirect_stderr
 from urllib.parse import urlparse
@@ -27,6 +28,8 @@
 # Number of bytes for read routines like decompress() and
 # response.content.iter_content()
 CHUNK = 1024*8
+# We can pump a max of 6MB out of Lambda
+LAMBDA_MAX_OUT = 6_000_000
 MIN_VCF_COLS = 8  # per 4.2 spec on header and data lines
 
 S3_DOMAIN_SUFFIX = '.amazonaws.com'
@@ -212,7 +215,7 @@ def extract_excel(file_):
     return html, {}
 
 
-def extract_ipynb(file_, exclude_output):
+def extract_ipynb(file_, exclude_output: bool):
     """
     parse and extract ipynb files
 
@@ -226,6 +229,13 @@ def extract_ipynb(file_, exclude_output):
     # local import reduces amortized latency, saves memory
     from nbconvert import HTMLExporter
     import nbformat
+    # get the file size
+    file_.seek(0, os.SEEK_END)
+    size = file_.tell()
+    if size > LAMBDA_MAX_OUT:
+        exclude_output = True
+    # rewind
+    file_.seek(0, os.SEEK_SET)
 
     html_exporter = HTMLExporter()
     html_exporter.template_file = 'basic'

lambdas/preview/test/test_index.py

@@ -2,6 +2,7 @@
 Test functions for preview endpoint
 """
 import json
+import math
 import os
 from pathlib import Path
 import re
@@ -241,6 +242,25 @@ def test_ipynb(self):
             '<span class="p">'
         ) in body_html, 'Last cell output missing'
 
+    @patch(__name__ + '.index.LAMBDA_MAX_OUT', 89_322)
+    @responses.activate
+    def test_ipynb_chop(self):
+        """test that we eliminate output cells when we're in danger of breaking
+        Lambda's invocation limit"""
+        notebook = BASE_DIR / 'nb_1200727.ipynb'
+        responses.add(
+            responses.GET,
+            self.FILE_URL,
+            body=notebook.read_bytes(),
+            status=200)
+        event = self._make_event({'url': self.FILE_URL, 'input': 'ipynb'})
+        resp = index.lambda_handler(event, None)
+        body = json.loads(read_body(resp))
+        assert resp['statusCode'] == 200, 'preview failed on nb_1200727.ipynb'
+        body_html = body['html']
+        # isclose bc string sizes differ, e.g. on Linux
+        assert math.isclose(len(body_html), 18084, abs_tol=200), "Hmm, didn't chop nb_1200727.ipynb"
+
     @responses.activate
     def test_ipynb_exclude(self):
         """test sending ipynb bytes"""