Readd profiler.

Author: ayjayt

src/py/kaleido/_kaleido_tab/_tab.py

@@ -113,8 +113,9 @@ async def _calc_fig(
         self,
         spec: _fig_tools.Spec,
         *,
-        topojson: str | None = None,
-        stepper=False,
+        topojson: str | None,
+        render_prof,
+        stepper,
     ) -> bytes:
         # js script
         kaleido_js_fn = (
@@ -123,7 +124,7 @@ async def _calc_fig(
             r"return kaleido_scopes.plotly(spec, ...args).then(JSON.stringify);"
             r"}"
         )
-
+        render_prof.profile_log.tick("sending javascript")
         result = await _dtools.exec_js_fn(
             self.tab,
             self._current_js_id,
@@ -133,16 +134,23 @@ async def _calc_fig(
             stepper,
         )
         _raise_error(result)
+        render_prof.profile_log.tick("javascript sent")
 
         _logger.debug2(f"Result of function call: {result}")
         js_response = _dtools.check_kaleido_js_response(result)
 
         if (response_format := js_response.get("format")) == "pdf":
+            render_prof.profile_log.tick("printing pdf")
             img_raw = await _dtools.print_pdf(self.tab)
+            render_prof.profile_log.tick("pdf printed")
         else:
             img_raw = js_response["result"]
 
         if response_format not in _TEXT_FORMATS:
-            return base64.b64decode(img_raw)
+            res = base64.b64decode(img_raw)
         else:
-            return str.encode(img_raw)
+            res = str.encode(img_raw)
+
+        render_prof.data_out_size = len(res)
+        render_prof.js_log = self.js_logger.log
+        return res

src/py/kaleido/_profiler.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pathlib import Path
+    from typing import Any
+
+    from . import _fig_tools
+
+    Event = str
+
+
+class WriteCall:
+    name: str
+    renders: list[RenderTaskProfile]
+
+    __slots__ = tuple(__annotations__)
+
+    def __init__(self, name: str):
+        self.name = name
+        self.renders = []
+
+
+class RenderTaskProfile:
+    info: dict[str, Any]  # literal?
+    error: None | BaseException
+    js_log: list[str]
+    profile_log: ProfileLog
+    data_in_size: int | None
+    data_out_size: int | None
+
+    __slots__ = tuple(__annotations__)
+
+    def __init__(
+        self,
+        spec: _fig_tools.Spec,
+        full_path: Path | None,
+        tab_id: str,
+    ) -> None:
+        self.info = {}
+        self.error = None
+        self.js_log = []
+        self.profile_log = ProfileLog()
+        self.data_in_size = None  # need to get this from choreographer
+        self.data_out_size = None
+
+        self.info.update(
+            {k: v for k, v in spec.items() if k != "data"},
+        )
+        self.info["path"] = full_path
+        self.info["tab"] = tab_id
+
+
+class ProfileLog:
+    _logs: dict[Event, float]
+
+    __slots__ = tuple(__annotations__)
+
+    def __init__(self) -> None:
+        self._logs = {}
+
+    def tick(self, name: str) -> None:
+        self._logs[name] = time.perf_counter()
+
+    def get_logs(self) -> dict[Event, float]:
+        return self._logs

src/py/kaleido/kaleido.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+from collections import deque
 from collections.abc import AsyncIterable, Iterable
 from pathlib import Path
 from typing import TYPE_CHECKING, TypedDict, cast, overload
@@ -12,7 +13,7 @@
 from choreographer.errors import ChromeNotFoundError
 from choreographer.utils import TmpDirectory
 
-from . import _fig_tools, _path_tools, _utils
+from . import _fig_tools, _path_tools, _profiler, _utils
 from ._kaleido_tab import _KaleidoTab
 from ._page_generator import PageGenerator
 
@@ -155,6 +156,7 @@ def __init__(
         self.tabs_ready = asyncio.Queue(maxsize=0)
         self._total_tabs = 0  # tabs properly registered
         self._html_tmp_dir = None
+        self._task_profiler: deque[_profiler.WriteCall] = deque(maxlen=5)
 
         # Kaleido Config
         page = page_generator
@@ -305,8 +307,10 @@ async def _render_task(
         self,
         fig_arg: FigureDict,
         *,
+        topojson: str | None,
         _write: bool,
-        **kwargs: Any,
+        profiler: _profiler.WriteCall,
+        stepper: bool,
     ) -> None | bytes:
         spec = _fig_tools.coerce_for_js(
             fig_arg.get("fig"),
@@ -324,25 +328,40 @@ async def _render_task(
 
         tab = await self._get_kaleido_tab()
 
+        render_prof = _profiler.RenderTaskProfile(
+            spec,
+            full_path if _write else None,
+            tab.tab.target_id,
+        )
+        profiler.renders.append(render_prof)
+
         try:
             img_bytes = await asyncio.wait_for(
                 tab._calc_fig(  # noqa: SLF001
                     spec,
-                    **kwargs,
+                    topojson=topojson,
+                    render_prof=render_prof,
+                    stepper=stepper,
                 ),
                 self._timeout,
             )
             if _write:
+                render_prof.profile_log.tick("starting file write")
                 await _utils.to_thread(full_path.write_bytes, img_bytes)
+                render_prof.profile_log.tick("file write done")
                 return None
             else:
                 return img_bytes
-        except:
+        except BaseException as e:
+            render_prof.profile_log.tick("errored out")
             if _write:
                 full_path.unlink()  # failure, no write
+            render_prof.error = e
             raise
         finally:
+            render_prof.profile_log.tick("returning tab")
             await self._return_kaleido_tab(tab)
+            render_prof.profile_log.tick("tab returned")
 
     ### API ###
     @overload
@@ -400,8 +419,13 @@ async def write_fig_from_object(
         if _is_figuredict(fig_dicts):
             fig_dicts = [fig_dicts]
 
+        name = "No Name"
         if main_task := asyncio.current_task():
             self._main_render_coroutines.add(main_task)
+            name = main_task.get_name()
+
+        profiler = _profiler.WriteCall(name)
+        self._task_profiler.append(profiler)
 
         tasks: set[asyncio.Task] = set()
 
@@ -412,6 +436,7 @@ async def write_fig_from_object(
                         fig_arg=fig_arg,
                         topojson=fig_arg.get("topojson"),
                         _write=_write,  # backwards compatibility
+                        profiler=profiler,
                         stepper=stepper,
                     ),
                 )