-
Notifications
You must be signed in to change notification settings - Fork 17
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
add notebook functions for starting dashboard and cloud functions
- Loading branch information
Showing
7 changed files
with
208 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
from ._cloud_functions import run_cloud_function # noqa: F401 | ||
from ._dashboards import run_dashboard # noqa: F401 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,63 @@ | ||
import asyncio | ||
import typing | ||
import urllib.parse | ||
|
||
from ._config import AppConfig | ||
from ._process import kill_process_by_port | ||
|
||
if typing.TYPE_CHECKING: | ||
from fastapi import FastAPI | ||
from ._uvicorn_util import AwaitableUvicornServer | ||
|
||
# Keep track of currently running server | ||
_running_server: 'typing.Optional[AwaitableUvicornServer]' = None | ||
|
||
|
||
async def run_cloud_function( | ||
app: 'FastAPI', | ||
kill_existing_server: bool = True, | ||
host: str = '0.0.0.0', | ||
log_level: str = 'error', | ||
) -> None: | ||
global _running_server | ||
from ._uvicorn_util import AwaitableUvicornServer | ||
|
||
try: | ||
import uvicorn | ||
except ImportError: | ||
raise ImportError('package uvicorn is required to run cloud functions') | ||
|
||
app_config = AppConfig.from_env() | ||
|
||
if kill_existing_server: | ||
# Shutdown the server gracefully if it was started by us. | ||
# Since the uvicorn server doesn't start a new subprocess | ||
# killing the process would result in kernel dying. | ||
if _running_server is not None: | ||
await _running_server.shutdown() | ||
_running_server = None | ||
|
||
# Kill if any other process is occupying the port | ||
kill_process_by_port(app_config.listen_port) | ||
|
||
# Add `GET /` route, used for liveness check | ||
@app.get('/') | ||
def ping() -> str: | ||
return 'Success!' | ||
|
||
base_path = urllib.parse.urlparse(app_config.url).path | ||
app.root_path = base_path | ||
|
||
config = uvicorn.Config( | ||
app, | ||
host=host, | ||
port=app_config.listen_port, | ||
log_level=log_level, | ||
) | ||
_running_server = AwaitableUvicornServer(config) | ||
|
||
asyncio.create_task(_running_server.serve()) | ||
await _running_server.wait_for_startup() | ||
|
||
if app_config.running_interactively: | ||
print(f'Cloud function available at {app_config.url}') |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,33 @@ | ||
import os | ||
from dataclasses import dataclass | ||
|
||
|
||
@dataclass | ||
class AppConfig: | ||
listen_port: int | ||
url: str | ||
running_interactively: bool | ||
|
||
@classmethod | ||
def from_env(cls) -> 'AppConfig': | ||
port = os.environ.get('SINGLESTOREDB_APP_LISTEN_PORT') | ||
if port is None: | ||
raise RuntimeError( | ||
'Missing SINGLESTOREDB_APP_LISTEN_PORT environment variable. ' | ||
'Is the code running outside SingleStoreDB notebook environment?', | ||
) | ||
url = os.environ.get('SINGLESTOREDB_APP_URL') | ||
if url is None: | ||
raise RuntimeError( | ||
'Missing SINGLESTOREDB_APP_URL environment variable. ' | ||
'Is the code running outside SingleStoreDB notebook environment?', | ||
) | ||
running_interactively = ( | ||
os.environ.get('SINGLESTOREDB_WORKLOAD_TYPE') == 'InteractiveNotebook' | ||
) | ||
|
||
return cls( | ||
listen_port=int(port), | ||
url=url, | ||
running_interactively=running_interactively, | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,44 @@ | ||
import typing | ||
import urllib.parse | ||
|
||
from ._config import AppConfig | ||
from ._process import kill_process_by_port | ||
|
||
if typing.TYPE_CHECKING: | ||
from plotly.graph_objs import Figure | ||
|
||
|
||
def run_dashboard( | ||
figure: 'Figure', | ||
kill_existing_server: bool = True, | ||
host: str = '0.0.0.0', | ||
debug: bool = False, | ||
) -> None: | ||
try: | ||
import dash | ||
except ImportError: | ||
raise ImportError('package dash is required to run dashboards') | ||
|
||
app_config = AppConfig.from_env() | ||
|
||
if kill_existing_server: | ||
kill_process_by_port(app_config.listen_port) | ||
|
||
base_path = urllib.parse.urlparse(app_config.url).path | ||
|
||
app = dash.Dash(requests_pathname_prefix=base_path) | ||
app.layout = dash.html.Div( | ||
[ | ||
dash.dcc.Graph(figure=figure), | ||
], | ||
) | ||
|
||
app.run( | ||
host=host, | ||
debug=debug, | ||
port=str(app_config.listen_port), | ||
jupyter_mode='external', | ||
) | ||
|
||
if app_config.running_interactively: | ||
print(f'Dash app available at {app_config.url}') |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,32 @@ | ||
import os | ||
import signal | ||
import typing | ||
if typing.TYPE_CHECKING: | ||
from psutil import Process | ||
|
||
|
||
def kill_process_by_port(port: int) -> None: | ||
existing_process = _find_process_by_port(port) | ||
kernel_pid = os.getpid() | ||
# Make sure we are not killing current kernel | ||
if existing_process is not None and kernel_pid != existing_process.pid: | ||
print(f'Killing process {existing_process.pid} which is using port {port}') | ||
os.kill(existing_process.pid, signal.SIGKILL) | ||
|
||
|
||
def _find_process_by_port(port: int) -> 'Process | None': | ||
try: | ||
import psutil | ||
except ImportError: | ||
raise ImportError('package psutil is required') | ||
|
||
for proc in psutil.process_iter(['pid']): | ||
try: | ||
connections = proc.connections() | ||
for conn in connections: | ||
if conn.laddr.port == port: | ||
return proc | ||
except psutil.AccessDenied: | ||
pass | ||
|
||
return None |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,32 @@ | ||
import asyncio | ||
import socket | ||
from typing import List | ||
from typing import Optional | ||
try: | ||
import uvicorn | ||
except ImportError: | ||
raise ImportError('package uvicorn is required') | ||
|
||
|
||
class AwaitableUvicornServer(uvicorn.Server): | ||
""" | ||
Adds `wait_for_startup` method. | ||
The function (asynchronously) blocks until the server | ||
starts listening or throws an error. | ||
""" | ||
|
||
def __init__(self, config: 'uvicorn.Config') -> None: | ||
super().__init__(config) | ||
self._startup_future = asyncio.get_event_loop().create_future() | ||
|
||
async def startup(self, sockets: Optional[List[socket.socket]] = None) -> None: | ||
try: | ||
result = await super().startup(sockets) | ||
self._startup_future.set_result(True) | ||
return result | ||
except Exception as error: | ||
self._startup_future.set_exception(error) | ||
raise error | ||
|
||
async def wait_for_startup(self) -> None: | ||
await self._startup_future |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,4 +1,6 @@ | ||
coverage | ||
dash | ||
fastapi | ||
pandas | ||
parameterized | ||
polars | ||
|