More intelligent stream handling.

Author: ellisonbg

jupyter_rtc_core/outputs/manager.py

@@ -3,6 +3,7 @@
 from pathlib import Path, PurePath
 import shutil
 
+from pycrdt import Map
 
 from traitlets.config import LoggingConfigurable
 from traitlets import (
@@ -14,6 +15,7 @@
     TraitError,
     Type,
     Unicode,
+    Int,
     default,
     validate,
 )
@@ -22,8 +24,11 @@
 
 class OutputsManager(LoggingConfigurable):
 
-    outputs_path = Instance(PurePath, help="The local runtime dir")
     _last_output_index = Dict(default_value={})
+    _stream_count = Dict(default_value={})
+
+    outputs_path = Instance(PurePath, help="The local runtime dir")
+    stream_limit = Int(default_value=10, config=True, allow_none=True)
 
     @default("outputs_path")
     def _default_outputs_path(self):
@@ -64,11 +69,10 @@ def get_stream(self, file_id, cell_id):
     def write(self, file_id, cell_id, output):
         """Write a new output for file_id and cell_id."""
         self.log.info(f"OutputsManager.write: {file_id} {cell_id} {output}")
-        if output["output_type"] == "stream":
-            url = self.write_stream(file_id, cell_id, output)
-        else:
-            url = self.write_output(file_id, cell_id, output)
-        return url
+        result = self.write_output(file_id, cell_id, output)
+        if output["output_type"] == "stream" and self.stream_limit is not None:
+            result = self.write_stream(file_id, cell_id, output)
+        return result
 
     def write_output(self, file_id, cell_id, output):
         self._ensure_path(file_id, cell_id)
@@ -80,17 +84,46 @@ def write_output(self, file_id, cell_id, output):
         with open(path, "w", encoding="utf-8") as f:
             f.write(data)
         url = f"/api/outputs/{file_id}/{cell_id}/{index}.output"
-        return url
+        return Map({})
 
-    def write_stream(self, file_id, cell_id, output):
+    def write_stream(self, file_id, cell_id, output) -> Map:
+        # How many stream outputs have been written for this cell previously
+        count = self._stream_count.get(cell_id, 0)
+
+        # Go ahead and write the incoming stream
         self._ensure_path(file_id, cell_id)
         path = self._build_path(file_id, cell_id) / "stream"
         text = output["text"]
         mode = 'a' if os.path.isfile(path) else 'w'
         with open(path, "a", encoding="utf-8") as f:
             f.write(text)
         url = f"/api/outputs/{file_id}/{cell_id}/stream"
-        return url
+
+        # Increment the count
+        count = count + 1
+        self._stream_count[cell_id] = count
+
+        # Now create the placeholder output
+        if count < self.stream_limit:
+            # Return the original if we haven't reached the limit
+            placeholder = Map(output)
+        elif count == self.stream_limit:
+            # Return a link to the full stream output
+            placeholder = Map({
+                "output_type": "display_data",
+                "data": {
+                    'text/html': f'<a href="{url}">Click this link to see the full stream output</a>'
+                }
+            })
+        elif count > self.stream_limit:
+            placeholder = None
+        return placeholder
+
+
+
+
+
+   
 
     def clear(self, file_id, cell_id=None):
         path = self._build_path(file_id, cell_id)

jupyter_rtc_core/outputs/output_processor.py

@@ -3,7 +3,7 @@
 
 from pycrdt import Map
 
-from traitlets import Dict, Unicode
+from traitlets import Dict, Unicode, Bool
 from traitlets.config import LoggingConfigurable
 
 
@@ -13,6 +13,12 @@ class OutputProcessor(LoggingConfigurable):
     _cell_indices = Dict(default_value={}) # a map from cell_id -> cell index in notebook
     _file_id = Unicode(default_value=None, allow_none=True)
 
+    use_outputs_service = Bool(
+        default_value=True,
+        config=True,
+        help="Should outputs be routed to the outputs service to minimize the in memory ydoc size."
+    )
+
     @property
     def settings(self):
         """A shortcut for the Tornado web app settings."""
@@ -99,7 +105,8 @@ def process_incoming_message(self, channel: str, msg: list[bytes]):
         if existing_msg_id != msg_id:  # cell is being re-run, clear output state
             self.clear(cell_id)
             if self._file_id is not None:
-                self.outputs_manager.clear(file_id=self._file_id, cell_id=cell_id)
+                if self.use_outputs_service:
+                    self.outputs_manager.clear(file_id=self._file_id, cell_id=cell_id)
         self.log.info(f"Saving (msg_id, cell_id): ({msg_id} {cell_id})")
         self.set_cell_id(msg_id, cell_id)
 
@@ -146,18 +153,13 @@ async def output_task(self, msg_type, cell_id, content):
             return
 
         # Convert from the message spec to the nbformat output structure
-        output = self.transform_output(msg_type, content, ydoc=False)
-        output_url = self.outputs_manager.write(file_id, cell_id, output)
-        nb_output = Map({
-                "output_type": "display_data",
-                "data": {
-                    'text/html': f'<a href="{output_url}">Output</a>'
-                },
-                "metadata": {
-                    "outputs_service": True
-                }
-            })
-        target_cell["outputs"].append(nb_output)
+        if self.use_outputs_service:
+            output = self.transform_output(msg_type, content, ydoc=False)
+            output = self.outputs_manager.write(file_id, cell_id, output)
+        else:
+            output = self.transform_output(msg_type, content, ydoc=True)
+        if output is not None:
+            target_cell["outputs"].append(output)
 
     def find_cell(self, cell_id, cells):
         """Find a cell with a given cell_id in the list of cells.