mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-13 13:17:52 +02:00
Add benchmark/profile commands
This commit is contained in:
@@ -177,3 +177,5 @@ docs/docs_skeleton/yarn.lock
|
||||
Untitled*.ipynb
|
||||
|
||||
Chinook.db
|
||||
|
||||
libs/langgraph/out
|
||||
|
||||
+15
-1
@@ -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 bench 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/results.json
|
||||
|
||||
bench:
|
||||
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 \
|
||||
|
||||
@@ -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)
|
||||
@@ -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())
|
||||
Generated
+34
-1
@@ -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"
|
||||
@@ -3126,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 = "f4ed3ad03f1ed1db10fcd18edd6665fc691411473026e2ddb5c8634043b6ef96"
|
||||
content-hash = "ea51ea9d9a33ea282b72e48ab84e1c07533ab5fee04ff7b1672599e38902e74d"
|
||||
|
||||
@@ -32,6 +32,8 @@ 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" ]
|
||||
|
||||
Binary file not shown.
@@ -1,81 +0,0 @@
|
||||
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())
|
||||
Reference in New Issue
Block a user