diff --git a/setup.cfg b/setup.cfg index 81cf1ce0f..7085bb36c 100644 --- a/setup.cfg +++ b/setup.cfg @@ -114,6 +114,8 @@ doc = pydantic==1.10 quartodoc==0.7.2 griffe==0.33.0 +profile = + pyinstrument [options.packages.find] include = shiny, shiny.* diff --git a/shiny/_main.py b/shiny/_main.py index 17cc1c8dc..36f44ee7d 100644 --- a/shiny/_main.py +++ b/shiny/_main.py @@ -1,5 +1,6 @@ from __future__ import annotations +import contextlib import copy import importlib import importlib.util @@ -20,6 +21,7 @@ from . import _autoreload, _hostenv, _static, _utils from ._docstring import no_example +from ._profiler import check_profiler_dependencies, profiler from ._typing_extensions import NotRequired, TypedDict from .express import is_express_app from .express._utils import escape_to_var_name @@ -150,6 +152,15 @@ def main() -> None: help="Dev mode", show_default=True, ) +@click.option( + "--profile", + is_flag=True, + default=False, + help="Run the app in profiling mode. This will launch the app under a " + "sampling profiler, and after Shiny is stopped, open a Speedscope.app flamegraph " + "in a web browser.", + show_default=True, +) @no_example() def run( app: str | shiny.App, @@ -167,8 +178,17 @@ def run( factory: bool, launch_browser: bool, dev_mode: bool, + profile: bool, **kwargs: object, ) -> None: + if profile: + if reload: + print( + "Error: --profile and --reload cannot be used together", file=sys.stderr + ) + sys.exit(1) + check_profiler_dependencies() + reload_includes_list = reload_includes.split(",") reload_excludes_list = reload_excludes.split(",") return run_app( @@ -186,6 +206,7 @@ def run( factory=factory, launch_browser=launch_browser, dev_mode=dev_mode, + profile=profile, **kwargs, ) @@ -206,6 +227,7 @@ def run_app( factory: bool = False, launch_browser: bool = False, dev_mode: bool = True, + profile: bool = False, **kwargs: object, ) -> None: """ @@ -250,6 +272,10 @@ def run_app( Treat ``app`` as an application factory, i.e. a () -> callable. launch_browser Launch app browser after app starts, using the Python webbrowser module. + profile + Run the app in profiling mode. This will launch the app under a sampling + profiler, and after Shiny is stopped, open a Speedscope.app flamegraph in a web + browser. **kwargs Additional keyword arguments which are passed to ``uvicorn.run``. For more information see [Uvicorn documentation](https://www.uvicorn.org/). @@ -353,19 +379,20 @@ def run_app( maybe_setup_rsw_proxying(log_config) - uvicorn.run( # pyright: ignore[reportUnknownMemberType] - app, - host=host, - port=port, - ws_max_size=ws_max_size, - log_level=log_level, - log_config=log_config, - app_dir=app_dir, - factory=factory, - lifespan="on", - **reload_args, # pyright: ignore[reportArgumentType] - **kwargs, - ) + with profiler() if profile else contextlib.nullcontext(): + uvicorn.run( # pyright: ignore[reportUnknownMemberType] + app, + host=host, + port=port, + ws_max_size=ws_max_size, + log_level=log_level, + log_config=log_config, + app_dir=app_dir, + factory=factory, + lifespan="on", + **reload_args, # pyright: ignore[reportArgumentType] + **kwargs, + ) def setup_hot_reload( diff --git a/shiny/_profiler.py b/shiny/_profiler.py new file mode 100644 index 000000000..ee26ae051 --- /dev/null +++ b/shiny/_profiler.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import contextlib +import sys +from typing import Any, Generator + + +def check_profiler_dependencies() -> None: + try: + import pyinstrument # pyright: ignore[reportUnusedImport] + except ImportError: + print( + "Error: Profiler is not installed. You can install it with " + "'pip install shiny[profile]'.", + file=sys.stderr, + ) + sys.exit(1) + + +@contextlib.contextmanager +def profiler() -> Generator[None, Any, None]: + import base64 + import os + import time + import webbrowser + from urllib.parse import quote_plus + + import pyinstrument + + prof = pyinstrument.Profiler() + prof.start() + + epoch_time = int(time.time()) + output_filename = f"profile-{epoch_time}.json" + output_filename_abs = os.path.join(os.getcwd(), output_filename) + print(f"Profiling to {output_filename_abs}", file=sys.stderr) + + try: + yield + finally: + prof_session = prof.stop() + + import pyinstrument.renderers.speedscope + + renderer = pyinstrument.renderers.speedscope.SpeedscopeRenderer() + output_str = renderer.render(prof_session) + + with open(output_filename_abs, "w", encoding="utf-8") as f: + f.write(output_str) + + print(f"Profile saved to {output_filename_abs}", file=sys.stderr) + + b64_str = base64.b64encode(output_str.encode("utf-8")).decode("utf-8") + data_uri = f"data:application/json;base64,{b64_str}" + full_url = ( + "https://speedscope.app/#profileURL=" + + quote_plus(data_uri) + # + "&title=" + # + quote_plus(output_filename) + ) + + webbrowser.open(full_url)