Add profling/tracing via VizTracer (#2850)

* Add profiling via VizTracer

* add viztracer to requirements.txt

* pyright

* don't open after automatically + use asyncio mode

.gitignore

@@ -241,3 +241,6 @@ dmypy.json
 
 # vite
 .vite/
+
+# traces
+backend/traces

backend/src/server.py

@@ -186,7 +186,18 @@ async def run(request: Request):
         logger.warning(message)
         return json(already_running_response(message), status=500)
 
+    executor_id = ExecutionId("main-executor " + uuid.uuid4().hex)
+
+    tracer = None
     try:
+        if ctx.config.trace:
+            logger.info("Starting VizTracer...")
+            from viztracer import VizTracer
+
+            tracer = VizTracer()
+            tracer.log_async = True
+            tracer.start()
+
         # wait until all previews are done
         await run_individual_counter.wait_zero()
 
@@ -197,7 +208,7 @@ async def run(request: Request):
 
         logger.info("Running new executor...")
         executor = Executor(
-            id=ExecutionId("main-executor " + uuid.uuid4().hex),
+            id=executor_id,
             chain=chain,
             send_broadcast_data=full_data["sendBroadcastData"],
             options=ExecutionOptions.parse(full_data["options"]),
@@ -249,6 +260,11 @@ async def run(request: Request):
 
         await ctx.queue.put({"event": "execution-error", "data": error})
         return json(error_response("Error running nodes!", exception), status=500)
+    finally:
+        if ctx.config.trace and tracer is not None:
+            logger.info("Stopping VizTracer...")
+            tracer.stop()
+            tracer.save(f"../traces/trace_{executor_id}.json")
 
 
 class RunIndividualRequest(TypedDict):

backend/src/server_config.py

@@ -39,6 +39,13 @@ class ServerConfig:
     Usage: `--storage-dir /foo/bar`
     """
 
+    trace: bool
+    """
+    Whether to enable tracing using VizTracer.
+
+    Usage: `--trace`
+    """
+
     @staticmethod
     def parse_argv() -> ServerConfig:
         parser = argparse.ArgumentParser(description="ChaiNNer's server.")
@@ -69,6 +76,11 @@ def parse_argv() -> ServerConfig:
             type=str,
             help="Directory to store for nodes to store files in.",
         )
+        parser.add_argument(
+            "--trace",
+            action="store_true",
+            help="Enable tracing using VizTracer.",
+        )
 
         parsed = parser.parse_args()
 
@@ -78,4 +90,5 @@ def parse_argv() -> ServerConfig:
             install_builtin_packages=parsed.install_builtin_packages,
             error_on_failed_node=parsed.error_on_failed_node,
             storage_dir=parsed.storage_dir or None,
+            trace=parsed.trace,
         )

backend/src/server_host.py

@@ -40,6 +40,8 @@ def __init__(self):
         worker_flags: list[str] = []
         if self.config.storage_dir is not None:
             worker_flags.extend(["--storage-dir", self.config.storage_dir])
+        if self.config.trace:
+            worker_flags.append("--trace")
 
         self._worker: Final[WorkerServer] = WorkerServer(worker_flags)
         self.pool: Final[ThreadPoolExecutor] = ThreadPoolExecutor(max_workers=4)

package.json

@@ -10,7 +10,9 @@
     "start": "electron-forge start -- --devtools",
     "frontend": "electron-forge start -- --remote-backend=http://127.0.0.1:8000 --devtools",
     "dev": "concurrently \"npm run dev:py\" \"electron-forge start -- --remote-backend=http://127.0.0.1:8000 --refresh --devtools\"",
+    "dev:profile": "concurrently \"npm run dev:py:profile\" \"electron-forge start -- --remote-backend=http://127.0.0.1:8000 --refresh --devtools\"",
     "dev:py": "cd backend/src && cross-env CHECK_LEVEL=fix nodemon --exec \"python -m debugpy --listen 5678\" ./run.py 8000",
+    "dev:py:profile": "cd backend/src && cross-env CHECK_LEVEL=fix nodemon --exec \"python ./run.py 8000 --trace",
     "package": "cross-env NODE_ENV=production electron-forge package",
     "make": "cross-env NODE_ENV=production electron-forge make",
     "make-linux-zip": "cross-env NODE_ENV=production electron-forge make --targets @electron-forge/maker-zip --platform linux",

requirements.txt

@@ -3,3 +3,4 @@ debugpy
 pyright==1.1.338
 pytest
 appdirs==1.4.4
+viztracer