Merge pull request #1609 from langchain-ai/nc/4sep/benchmark

Add benchmark script

.github/workflows/bench.yml

@@ -0,0 +1,63 @@
+name: bench
+
+on:
+  pull_request:
+    paths:
+      - "libs/**"
+
+env:
+  POETRY_VERSION: "1.7.1"
+
+jobs:
+  baseline:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: libs/langgraph
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: main
+      - name: Set up Python 3.11 + Poetry ${{ env.POETRY_VERSION }}
+        uses: "./.github/actions/poetry_setup"
+        with:
+          python-version: "3.11"
+          poetry-version: ${{ env.POETRY_VERSION }}
+          cache-key: bench
+      - name: Install dependencies
+        run: poetry install --with dev
+      - name: Run benchmarks
+        run: make benchmark
+      - name: Upload benchmark baseline
+        uses: actions/upload-artifact@v4
+        with:
+          name: bench-baseline.json
+          path: out/bench.json
+  compare:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: libs/langgraph
+    needs: [baseline]
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python 3.11 + Poetry ${{ env.POETRY_VERSION }}
+        uses: "./.github/actions/poetry_setup"
+        with:
+          python-version: "3.11"
+          poetry-version: ${{ env.POETRY_VERSION }}
+          working-directory: libs/langgraph
+          cache-key: bench
+      - name: Install dependencies
+        run: poetry install --with dev
+      - name: Run benchmarks
+        run: make benchmark
+      - name: Download benchmark baseline
+        uses: actions/download-artifact@v4
+        with:
+          name: bench-baseline.json
+          path: out
+      - name: Compare benchmarks
+        run: poetry run pyperf compare_to out/bench-baseline.json out/bench.json --table --group-by-speed >> $GITHUB_OUTPUT
+      - name: Annotation
+        run: echo "::notice file=libs/langgraph/bench/__main__.py::$GITHUB_OUTPUT"

.gitignore

@@ -177,3 +177,5 @@ docs/docs_skeleton/yarn.lock
 Untitled*.ipynb
 
 Chinook.db
+
+libs/langgraph/out

libs/langgraph/Makefile

@@ -1,4 +1,4 @@
-.PHONY: all format lint test test_watch integration_tests spell_check spell_fix
+.PHONY: all format lint test test_watch integration_tests spell_check spell_fix benchmark profile
 
 # Default target executed when no arguments are given to make.
 all: help
@@ -7,6 +7,20 @@ all: help
 # TESTING AND COVERAGE
 ######################
 
+# Benchmarks
+
+OUTPUT ?= out/bench.json
+
+benchmark:
+	mkdir -p out
+	poetry run python -m bench -o $(OUTPUT) --rigorous
+
+GRAPH ?= bench/fanout_to_subgraph.py
+
+profile:
+	mkdir -p out
+	sudo poetry run py-spy record -g -o out/profile.svg -- python $(GRAPH)
+
 # Run unit tests and generate a coverage report.
 coverage:
 	poetry run pytest --cov \

libs/langgraph/bench/__main__.py

@@ -0,0 +1,63 @@
+import random
+from typing import Optional
+
+from pyperf._runner import Runner
+from uvloop import new_event_loop
+
+from bench.fanout_to_subgraph import fanout_to_subgraph
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.pregel import Pregel
+
+
+async def run(graph: Pregel, input: dict, config: Optional[dict]):
+    len([c async for c in graph.astream(input, config=config)])
+
+
+benchmarks = (
+    (
+        "fanout_to_subgraph_10x",
+        fanout_to_subgraph().compile(checkpointer=None),
+        {
+            "subjects": [
+                random.choices("abcdefghijklmnopqrstuvwxyz", k=1000) for _ in range(10)
+            ]
+        },
+        None,
+    ),
+    (
+        "fanout_to_subgraph_10x_checkpoint",
+        fanout_to_subgraph().compile(checkpointer=MemorySaver()),
+        {
+            "subjects": [
+                random.choices("abcdefghijklmnopqrstuvwxyz", k=1000) for _ in range(10)
+            ]
+        },
+        {"configurable": {"thread_id": "1"}},
+    ),
+    (
+        "fanout_to_subgraph_100x",
+        fanout_to_subgraph().compile(checkpointer=None),
+        {
+            "subjects": [
+                random.choices("abcdefghijklmnopqrstuvwxyz", k=1000) for _ in range(100)
+            ]
+        },
+        None,
+    ),
+    (
+        "fanout_to_subgraph_100x_checkpoint",
+        fanout_to_subgraph().compile(checkpointer=MemorySaver()),
+        {
+            "subjects": [
+                random.choices("abcdefghijklmnopqrstuvwxyz", k=1000) for _ in range(100)
+            ]
+        },
+        {"configurable": {"thread_id": "1"}},
+    ),
+)
+
+
+r = Runner()
+
+for name, graph, input, config in benchmarks:
+    r.bench_async_func(name, run, graph, input, config, loop_factory=new_event_loop)

libs/langgraph/bench/fanout_to_subgraph.py

@@ -0,0 +1,76 @@
+import asyncio
+import operator
+from typing import Annotated, TypedDict
+
+from langgraph.constants import END, START, Send
+from langgraph.graph.state import StateGraph
+
+
+def fanout_to_subgraph() -> StateGraph:
+    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)
+
+    return builder
+
+
+if __name__ == "__main__":
+    import random
+
+    import uvloop
+
+    from langgraph.checkpoint.memory import MemorySaver
+
+    graph = fanout_to_subgraph().compile(checkpointer=MemorySaver())
+    input = {
+        "subjects": [
+            random.choices("abcdefghijklmnopqrstuvwxyz", k=1000) for _ in range(1000)
+        ]
+    }
+    config = {"configurable": {"thread_id": "1"}}
+
+    async def run():
+        len([c async for c in graph.astream(input, config=config)])
+
+    uvloop.install()
+    asyncio.run(run())

libs/langgraph/poetry.lock

@@ -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"
@@ -1949,6 +1949,22 @@ files = [
 [package.extras]
 tests = ["pytest"]
 
+[[package]]
+name = "py-spy"
+version = "0.3.14"
+description = "Sampling profiler for Python programs"
+optional = false
+python-versions = "*"
+files = [
+    {file = "py_spy-0.3.14-py2.py3-none-macosx_10_7_x86_64.whl", hash = "sha256:5b342cc5feb8d160d57a7ff308de153f6be68dcf506ad02b4d67065f2bae7f45"},
+    {file = "py_spy-0.3.14-py2.py3-none-macosx_10_9_x86_64.macosx_11_0_arm64.macosx_10_9_universal2.whl", hash = "sha256:fe7efe6c91f723442259d428bf1f9ddb9c1679828866b353d539345ca40d9dd2"},
+    {file = "py_spy-0.3.14-py2.py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:590905447241d789d9de36cff9f52067b6f18d8b5e9fb399242041568d414461"},
+    {file = "py_spy-0.3.14-py2.py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:fd6211fe7f587b3532ba9d300784326d9a6f2b890af7bf6fff21a029ebbc812b"},
+    {file = "py_spy-0.3.14-py2.py3-none-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:3e8e48032e71c94c3dd51694c39e762e4bbfec250df5bf514adcdd64e79371e0"},
+    {file = "py_spy-0.3.14-py2.py3-none-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:f59b0b52e56ba9566305236375e6fc68888261d0d36b5addbe3cf85affbefc0e"},
+    {file = "py_spy-0.3.14-py2.py3-none-win_amd64.whl", hash = "sha256:8f5b311d09f3a8e33dbd0d44fc6e37b715e8e0c7efefafcda8bfd63b31ab5a31"},
+]
+
 [[package]]
 name = "pycparser"
 version = "2.22"
@@ -2098,6 +2114,23 @@ files = [
 [package.extras]
 diagrams = ["jinja2", "railroad-diagrams"]
 
+[[package]]
+name = "pyperf"
+version = "2.7.0"
+description = "Python module to run and analyze benchmarks"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "pyperf-2.7.0-py3-none-any.whl", hash = "sha256:dce63053b916b73d8736a77404309328f938851b5c2c5e8493cde910ce37e362"},
+    {file = "pyperf-2.7.0.tar.gz", hash = "sha256:4201c6601032f374e9c900c6d2544a2f5891abedc1a96eec0e7b2338a6247589"},
+]
+
+[package.dependencies]
+psutil = ">=5.9.0"
+
+[package.extras]
+dev = ["importlib-metadata", "tox"]
+
 [[package]]
 name = "pytest"
 version = "8.3.2"
@@ -2956,6 +2989,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 +3159,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 = "ea51ea9d9a33ea282b72e48ab84e1c07533ab5fee04ff7b1672599e38902e74d"

libs/langgraph/pyproject.toml

@@ -31,6 +31,9 @@ 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"
+pyperf = "^2.7.0"
+py-spy = "^0.3.14"
 
 [tool.ruff]
 lint.select = [ "E", "F", "I" ]