Add PyInstrument Profiling

notebooks/experimental/Profiling.ipynb

@@ -0,0 +1,93 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a810d2e7-86f4-4764-a153-5a65ab704889",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# syft absolute\n",
+    "import syft as sy"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4d3b2e87-0488-45bb-9395-d2027815556a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Works only for in-memory python workers with uvicorn"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9a38eedb-849b-4cc3-8077-4a347ab8db34",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "canada_datasite = sy.orchestra.launch(\n",
+    "    \"canada-datasite\", port=\"auto\", dev_mode=True, profile=True, profile_interval=0.001\n",
+    ")\n",
+    "\n",
+    "italy_datasite = sy.orchestra.launch(\n",
+    "    \"italy-datasite\", port=\"auto\", dev_mode=True, profile=True, profile_interval=0.001\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "300b2180-d611-465c-8d13-c02fedf8a959",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "do_canada_client = canada_datasite.login(\n",
+    "    email=\"[email protected]\", password=\"changethis\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "abb6481b-22c0-4cce-a8c9-4f0cea63b564",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%pyinstrument\n",
+    "do_italy_client = italy_datasite.login(email=\"[email protected]\", password=\"changethis\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6c030558-0a16-4194-9397-0261ecb40b9b",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

packages/syft/setup.cfg

@@ -94,6 +94,7 @@ dev =
     %(test_plugins)s
     %(telemetry)s
     bandit==1.7.8
+    pyinstrument==4.6.2
     debugpy==1.8.2
     importlib-metadata==7.1.0
     isort==5.13.2

packages/syft/src/syft/orchestra.py

@@ -186,6 +186,9 @@ def deploy_to_python(
     background_tasks: bool = False,
     debug: bool = False,
     migrate: bool = False,
+    profile: bool = False,
+    profile_interval: float = 0.001,
+    profile_dir: str | None = None,
 ) -> ServerHandle:
     worker_classes = {
         ServerType.DATASITE: Datasite,
@@ -215,6 +218,9 @@ def deploy_to_python(
         "background_tasks": background_tasks,
         "debug": debug,
         "migrate": migrate,
+        "profile": profile,
+        "profile_interval": profile_interval,
+        "profile_dir": profile_dir,
     }
 
     if port:
@@ -325,6 +331,10 @@ def launch(
         background_tasks: bool = False,
         debug: bool = False,
         migrate: bool = False,
+        # Profiling Related Input for in-memory fastapi server
+        profile: bool = False,
+        profile_interval: float = 0.001,
+        profile_dir: str | None = None,
     ) -> ServerHandle:
         if dev_mode is True:
             thread_workers = True
@@ -363,6 +373,9 @@ def launch(
                 background_tasks=background_tasks,
                 debug=debug,
                 migrate=migrate,
+                profile=profile,
+                profile_interval=profile_interval,
+                profile_dir=profile_dir,
             )
             display(
                 SyftInfo(

packages/syft/src/syft/server/routes.py

@@ -2,7 +2,10 @@
 import base64
 import binascii
 from collections.abc import AsyncGenerator
+from collections.abc import Callable
+from datetime import datetime
 import logging
+from pathlib import Path
 from typing import Annotated
 
 # third party
@@ -34,12 +37,13 @@
 from ..util.telemetry import TRACE_MODE
 from .credentials import SyftVerifyKey
 from .credentials import UserLoginCredentials
+from .uvicorn_settings import UvicornSettings
 from .worker import Worker
 
 logger = logging.getLogger(__name__)
 
 
-def make_routes(worker: Worker) -> APIRouter:
+def make_routes(worker: Worker, settings: UvicornSettings | None = None) -> APIRouter:
     if TRACE_MODE:
         # third party
         try:
@@ -49,6 +53,34 @@ def make_routes(worker: Worker) -> APIRouter:
         except Exception as e:
             logger.error("Failed to import opentelemetry", exc_info=e)
 
+    def _handle_profile(
+        request: Request, handler_func: Callable, *args: list, **kwargs: dict
+    ) -> Response:
+        if not settings:
+            raise Exception("Server Settings are required to enable profiling")
+        # third party
+        from pyinstrument import Profiler  # Lazy Load
+
+        profiles_dir = Path(settings.profile_dir or Path.cwd()) / "profiles"
+        profiles_dir.mkdir(parents=True, exist_ok=True)
+
+        with Profiler(
+            interval=settings.profile_interval, async_mode="enabled"
+        ) as profiler:
+            response = handler_func(*args, **kwargs)
+
+        timestamp = datetime.now().strftime("%d-%m-%Y-%H:%M:%S")
+        url_path = request.url.path.replace("/api/v2", "").replace("/", "-")
+        profile_output_path = (
+            profiles_dir / f"{settings.name}-{timestamp}{url_path}.html"
+        )
+        profiler.write_html(profile_output_path)
+
+        logger.info(
+            f"Request to {request.url.path} took {profiler.last_session.duration:.2f} seconds"
+        )
+        return response
+
     router = APIRouter()
 
     async def get_body(request: Request) -> bytes:
@@ -165,6 +197,13 @@ def syft_new_api(
                 kind=trace.SpanKind.SERVER,
             ):
                 return handle_syft_new_api(user_verify_key, communication_protocol)
+        elif settings and settings.profile:
+            return _handle_profile(
+                request,
+                handle_syft_new_api,
+                user_verify_key,
+                communication_protocol,
+            )
         else:
             return handle_syft_new_api(user_verify_key, communication_protocol)
 
@@ -188,6 +227,8 @@ def syft_new_api_call(
                 kind=trace.SpanKind.SERVER,
             ):
                 return handle_new_api_call(data)
+        elif settings and settings.profile:
+            return _handle_profile(request, handle_new_api_call, data)
         else:
             return handle_new_api_call(data)
 
@@ -255,6 +296,8 @@ def login(
                 kind=trace.SpanKind.SERVER,
             ):
                 return handle_login(email, password, worker)
+        elif settings and settings.profile:
+            return _handle_profile(request, handle_login, email, password, worker)
         else:
             return handle_login(email, password, worker)
 
@@ -269,6 +312,8 @@ def register(
                 kind=trace.SpanKind.SERVER,
             ):
                 return handle_register(data, worker)
+        elif settings and settings.profile:
+            return _handle_profile(request, handle_register, data, worker)
         else:
             return handle_register(data, worker)