Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement async generator based Voila get handler #1025

Merged
merged 1 commit into from
Dec 23, 2021
Merged

Implement async generator based Voila get handler #1025

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 34 additions & 18 deletions voila/handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#############################################################################


import asyncio
import os
from typing import Dict

Expand All @@ -33,8 +34,7 @@ def initialize(self, **kwargs):
# we want to avoid starting multiple kernels due to template mistakes
self.kernel_started = False

@tornado.web.authenticated
async def get(self, path=None):
async def get_generator(self, path=None):
# if the handler got a notebook_path argument, always serve that
notebook_path = self.notebook_path or path

Expand Down Expand Up @@ -101,22 +101,18 @@ async def get(self, path=None):
QueryStringSocketHandler.send_updates({'kernel_id': kernel_id, 'payload': self.request.query})
# Send rendered cell to frontend
if len(rendered_cache) > 0:
self.write(''.join(rendered_cache))
self.flush()
yield ''.join(rendered_cache)

# Wait for current running cell finish and send this cell to
# frontend.
rendered, rendering = await render_task
if len(rendered) > len(rendered_cache):
html_snippet = ''.join(rendered[len(rendered_cache):])
self.write(html_snippet)
self.flush()
yield html_snippet

# Continue render cell from generator.
async for html_snippet, _ in rendering:
self.write(html_snippet)
self.flush()
self.flush()
yield html_snippet

else:
# All kernels are used or pre-heated kernel is disabled, start a normal kernel.
Expand All @@ -139,8 +135,7 @@ def time_out():
can be used in a template to give feedback to a user
"""

self.write('<script>voila_heartbeat()</script>\n')
self.flush()
return '<script>voila_heartbeat()</script>\n'

kernel_env[ENV_VARIABLE.VOILA_PREHEAT] = 'False'
kernel_env[ENV_VARIABLE.VOILA_BASE_URL] = self.base_url
Expand All @@ -154,13 +149,34 @@ def time_out():
)
)
kernel_future = self.kernel_manager.get_kernel(kernel_id)
async for html_snippet, _ in gen.generate_content_generator(
kernel_id, kernel_future, time_out
):
self.write(html_snippet)
self.flush()
# we may not want to consider not flusing after each snippet, but add an explicit flush function to the jinja context
# yield # give control back to tornado's IO loop, so it can handle static files or other requests
queue = asyncio.Queue()

async def put_html():
async for html_snippet, _ in gen.generate_content_generator(kernel_id, kernel_future):
await queue.put(html_snippet)

await queue.put(None)

asyncio.ensure_future(put_html())

# If not done within the timeout, we send a heartbeat
# this is fundamentally to avoid browser/proxy read-timeouts, but
# can be used in a template to give feedback to a user
while True:
try:
html_snippet = await asyncio.wait_for(queue.get(), self.voila_configuration.http_keep_alive_timeout)
except asyncio.TimeoutError:
yield time_out()
else:
if html_snippet is None:
break
yield html_snippet

@tornado.web.authenticated
async def get(self, path=None):
gen = self.get_generator(path=path)
async for html in gen:
self.write(html)
self.flush()

def redirect_to_file(self, path):
Expand Down
28 changes: 5 additions & 23 deletions voila/notebook_renderer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,10 @@
#############################################################################


import asyncio
import os
import sys
import traceback
from typing import Callable, Generator, Tuple, Union, List
from typing import Generator, Tuple, Union, List

import nbformat
import tornado.web
Expand Down Expand Up @@ -150,13 +149,12 @@ def generate_content_generator(
self,
kernel_id: Union[str, None] = None,
kernel_future=None,
timeout_callback: Union[Callable, None] = None,
) -> Generator:
async def inner_kernel_start(nb):
return await self._jinja_kernel_start(nb, kernel_id, kernel_future)

def inner_cell_generator(nb, kernel_id):
return self._jinja_cell_generator(nb, kernel_id, timeout_callback)
return self._jinja_cell_generator(nb, kernel_id)

# These functions allow the start of a kernel and execution of the
# notebook after (parts of) the template has been rendered and send
Expand Down Expand Up @@ -248,32 +246,16 @@ async def _jinja_notebook_execute(self, nb, kernel_id):

await self._cleanup_resources()

async def _jinja_cell_generator(self, nb, kernel_id, timeout_callback):
async def _jinja_cell_generator(self, nb, kernel_id):
"""Generator that will execute a single notebook cell at a time"""
nb, _ = ClearOutputPreprocessor().preprocess(
nb, {'metadata': {'path': self.cwd}}
)
for cell_idx, input_cell in enumerate(nb.cells):
try:
task = asyncio.ensure_future(
self.executor.execute_cell(
input_cell, None, cell_idx, store_history=False
)
output_cell = await self.executor.execute_cell(
input_cell, None, cell_idx, store_history=False
)
while True:
_, pending = await asyncio.wait(
{task}, timeout=self.voila_configuration.http_keep_alive_timeout
)
if pending:
# If not done within the timeout, we send a heartbeat
# this is fundamentally to avoid browser/proxy read-timeouts, but
# can be used in a template to give feedback to a user
if timeout_callback is not None:
timeout_callback()

continue
output_cell = await task
break
except TimeoutError:
output_cell = input_cell
break
Expand Down