Add benchmark script

2026-09-05 09:17:47 +02:00 · 2024-09-10 17:32:47 -07:00
parent e3ca7bb3e9
commit 26f23ea26c
6 changed files with 623 additions and 2 deletions
@@ -1255,7 +1255,7 @@ typing-extensions = ">=4.7"

 [[package]]
 name = "langgraph-checkpoint"
-version = "1.0.8"
+version = "1.0.9"
 description = "Library with base interfaces for LangGraph checkpoint savers."
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -2956,6 +2956,50 @@ h2 = ["h2 (>=4,<5)"]
 socks = ["pysocks (>=1.5.6,!=1.5.7,<2.0)"]
 zstd = ["zstandard (>=0.18.0)"]

+[[package]]
+name = "uvloop"
+version = "0.20.0"
+description = "Fast implementation of asyncio event loop on top of libuv"
+optional = false
+python-versions = ">=3.8.0"
+files = [
+    {file = "uvloop-0.20.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:9ebafa0b96c62881d5cafa02d9da2e44c23f9f0cd829f3a32a6aff771449c996"},
+    {file = "uvloop-0.20.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:35968fc697b0527a06e134999eef859b4034b37aebca537daeb598b9d45a137b"},
+    {file = "uvloop-0.20.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b16696f10e59d7580979b420eedf6650010a4a9c3bd8113f24a103dfdb770b10"},
+    {file = "uvloop-0.20.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9b04d96188d365151d1af41fa2d23257b674e7ead68cfd61c725a422764062ae"},
+    {file = "uvloop-0.20.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:94707205efbe809dfa3a0d09c08bef1352f5d3d6612a506f10a319933757c006"},
+    {file = "uvloop-0.20.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:89e8d33bb88d7263f74dc57d69f0063e06b5a5ce50bb9a6b32f5fcbe655f9e73"},
+    {file = "uvloop-0.20.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:e50289c101495e0d1bb0bfcb4a60adde56e32f4449a67216a1ab2750aa84f037"},
+    {file = "uvloop-0.20.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:e237f9c1e8a00e7d9ddaa288e535dc337a39bcbf679f290aee9d26df9e72bce9"},
+    {file = "uvloop-0.20.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:746242cd703dc2b37f9d8b9f173749c15e9a918ddb021575a0205ec29a38d31e"},
+    {file = "uvloop-0.20.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:82edbfd3df39fb3d108fc079ebc461330f7c2e33dbd002d146bf7c445ba6e756"},
+    {file = "uvloop-0.20.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:80dc1b139516be2077b3e57ce1cb65bfed09149e1d175e0478e7a987863b68f0"},
+    {file = "uvloop-0.20.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:4f44af67bf39af25db4c1ac27e82e9665717f9c26af2369c404be865c8818dcf"},
+    {file = "uvloop-0.20.0-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:4b75f2950ddb6feed85336412b9a0c310a2edbcf4cf931aa5cfe29034829676d"},
+    {file = "uvloop-0.20.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:77fbc69c287596880ecec2d4c7a62346bef08b6209749bf6ce8c22bbaca0239e"},
+    {file = "uvloop-0.20.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6462c95f48e2d8d4c993a2950cd3d31ab061864d1c226bbf0ee2f1a8f36674b9"},
+    {file = "uvloop-0.20.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:649c33034979273fa71aa25d0fe120ad1777c551d8c4cd2c0c9851d88fcb13ab"},
+    {file = "uvloop-0.20.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:3a609780e942d43a275a617c0839d85f95c334bad29c4c0918252085113285b5"},
+    {file = "uvloop-0.20.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:aea15c78e0d9ad6555ed201344ae36db5c63d428818b4b2a42842b3870127c00"},
+    {file = "uvloop-0.20.0-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:f0e94b221295b5e69de57a1bd4aeb0b3a29f61be6e1b478bb8a69a73377db7ba"},
+    {file = "uvloop-0.20.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:fee6044b64c965c425b65a4e17719953b96e065c5b7e09b599ff332bb2744bdf"},
+    {file = "uvloop-0.20.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:265a99a2ff41a0fd56c19c3838b29bf54d1d177964c300dad388b27e84fd7847"},
+    {file = "uvloop-0.20.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b10c2956efcecb981bf9cfb8184d27d5d64b9033f917115a960b83f11bfa0d6b"},
+    {file = "uvloop-0.20.0-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:e7d61fe8e8d9335fac1bf8d5d82820b4808dd7a43020c149b63a1ada953d48a6"},
+    {file = "uvloop-0.20.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:2beee18efd33fa6fdb0976e18475a4042cd31c7433c866e8a09ab604c7c22ff2"},
+    {file = "uvloop-0.20.0-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:d8c36fdf3e02cec92aed2d44f63565ad1522a499c654f07935c8f9d04db69e95"},
+    {file = "uvloop-0.20.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:a0fac7be202596c7126146660725157d4813aa29a4cc990fe51346f75ff8fde7"},
+    {file = "uvloop-0.20.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9d0fba61846f294bce41eb44d60d58136090ea2b5b99efd21cbdf4e21927c56a"},
+    {file = "uvloop-0.20.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:95720bae002ac357202e0d866128eb1ac82545bcf0b549b9abe91b5178d9b541"},
+    {file = "uvloop-0.20.0-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:36c530d8fa03bfa7085af54a48f2ca16ab74df3ec7108a46ba82fd8b411a2315"},
+    {file = "uvloop-0.20.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:e97152983442b499d7a71e44f29baa75b3b02e65d9c44ba53b10338e98dedb66"},
+    {file = "uvloop-0.20.0.tar.gz", hash = "sha256:4603ca714a754fc8d9b197e325db25b2ea045385e8a3ad05d3463de725fdf469"},
+]
+
+[package.extras]
+docs = ["Sphinx (>=4.1.2,<4.2.0)", "sphinx-rtd-theme (>=0.5.2,<0.6.0)", "sphinxcontrib-asyncio (>=0.3.0,<0.4.0)"]
+test = ["Cython (>=0.29.36,<0.30.0)", "aiohttp (==3.9.0b0)", "aiohttp (>=3.8.1)", "flake8 (>=5.0,<6.0)", "mypy (>=0.800)", "psutil", "pyOpenSSL (>=23.0.0,<23.1.0)", "pycodestyle (>=2.9.0,<2.10.0)"]
+
 [[package]]
 name = "watchdog"
 version = "4.0.1"
@@ -3082,4 +3126,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9.0,<4.0"
-content-hash = "3c3d4b9ce6b0609d0399ca9aece50495d0a29b978042908dd981ea267f541934"
+content-hash = "f4ed3ad03f1ed1db10fcd18edd6665fc691411473026e2ddb5c8634043b6ef96"
@@ -31,6 +31,7 @@ langgraph-checkpoint = {path = "../checkpoint", develop = true}
 langgraph-checkpoint-sqlite = {path = "../checkpoint-sqlite", develop = true}
 langgraph-checkpoint-postgres = {path = "../checkpoint-postgres", develop = true}
 psycopg = {extras = ["binary"], version = ">=3.0.0"}
+uvloop = "^0.20.0"

 [tool.ruff]
 lint.select = [ "E", "F", "I" ]
@@ -0,0 +1,81 @@
+import asyncio
+import operator
+import random
+from time import perf_counter_ns
+from typing import Annotated, TypedDict
+
+import uvloop
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.constants import END, START, Send
+from langgraph.graph.state import StateGraph
+
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+
+
+class OverallState(TypedDict):
+    subjects: list[str]
+    jokes: Annotated[list[str], operator.add]
+
+
+async def continue_to_jokes(state: OverallState):
+    return [Send("generate_joke", {"subject": s}) for s in state["subjects"]]
+
+
+class JokeInput(TypedDict):
+    subject: str
+
+
+class JokeOutput(TypedDict):
+    jokes: list[str]
+
+
+async def bump(state: JokeOutput):
+    return {"jokes": [state["jokes"][0] + " a"]}
+
+
+async def generate(state: JokeInput):
+    return {"jokes": [f"Joke about {state['subject']}"]}
+
+
+async def edit(state: JokeInput):
+    subject = state["subject"]
+    return {"subject": f"{subject} - hohoho"}
+
+
+async def bump_loop(state: JokeOutput):
+    return END if state["jokes"][0].endswith(" a" * 10) else "bump"
+
+
+# subgraph
+subgraph = StateGraph(input=JokeInput, output=JokeOutput)
+subgraph.add_node("edit", edit)
+subgraph.add_node("generate", generate)
+subgraph.add_node("bump", bump)
+subgraph.set_entry_point("edit")
+subgraph.add_edge("edit", "generate")
+subgraph.add_edge("generate", "bump")
+subgraph.add_conditional_edges("bump", bump_loop)
+subgraph.set_finish_point("generate")
+subgraphc = subgraph.compile()
+
+# parent graph
+builder = StateGraph(OverallState)
+builder.add_node("generate_joke", subgraphc)
+builder.add_conditional_edges(START, continue_to_jokes)
+builder.add_edge("generate_joke", END)
+
+
+async def main():
+    graph = builder.compile(checkpointer=None)
+    config = {"configurable": {"thread_id": "1"}}
+    input = {
+        "subjects": [random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(1000)]
+    }
+
+    # invoke and pause at nested interrupt
+    s = perf_counter_ns()
+    len([c async for c in graph.astream(input, config=config)])
+    print("Time taken:", (perf_counter_ns() - s) / 1e9)
+
+
+asyncio.run(main())
@@ -0,0 +1,80 @@
+import operator
+import random
+from time import perf_counter_ns
+from typing import Annotated, TypedDict
+
+from langgraph.checkpoint.sqlite import SqliteSaver
+from langgraph.constants import END, START, Send
+from langgraph.graph.state import StateGraph
+
+
+class OverallState(TypedDict):
+    subjects: list[str]
+    jokes: Annotated[list[str], operator.add]
+
+
+def continue_to_jokes(state: OverallState):
+    return [Send("generate_joke", {"subject": s}) for s in state["subjects"]]
+
+
+class JokeInput(TypedDict):
+    subject: str
+
+
+class JokeOutput(TypedDict):
+    jokes: list[str]
+
+
+def bump(state: JokeOutput):
+    return {"jokes": [state["jokes"][0] + " a"]}
+
+
+def generate(state: JokeInput):
+    return {"jokes": [f"Joke about {state['subject']}"]}
+
+
+def edit(state: JokeInput):
+    subject = state["subject"]
+    return {"subject": f"{subject} - hohoho"}
+
+
+def bump_loop(state: JokeOutput):
+    return END if state["jokes"][0].endswith(" a" * 10) else "bump"
+
+
+# subgraph
+subgraph = StateGraph(input=JokeInput, output=JokeOutput)
+subgraph.add_node("edit", edit)
+subgraph.add_node("generate", generate)
+subgraph.add_node("bump", bump)
+subgraph.set_entry_point("edit")
+subgraph.add_edge("edit", "generate")
+subgraph.add_edge("generate", "bump")
+subgraph.add_conditional_edges("bump", bump_loop)
+subgraph.set_finish_point("generate")
+subgraphc = subgraph.compile()
+
+# parent graph
+builder = StateGraph(OverallState)
+builder.add_node("generate_joke", subgraphc)
+builder.add_conditional_edges(START, continue_to_jokes)
+builder.add_edge("generate_joke", END)
+
+
+def main():
+    with SqliteSaver.from_conn_string(":memory:") as checkpointer:
+        graph = builder.compile(checkpointer=checkpointer)
+        config = {"configurable": {"thread_id": "1"}}
+        input = {
+            "subjects": [
+                random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(100)
+            ]
+        }
+
+        # invoke and pause at nested interrupt
+        s = perf_counter_ns()
+        assert len([c for c in graph.stream(input, config=config)]) == 100
+        print("Time taken:", (perf_counter_ns() - s) / 1e9)
+
+
+main()