Add benchmark/profile commands

This commit is contained in:
Nuno Campos
2024-09-11 11:44:18 -07:00
parent 26f23ea26c
commit f577529356
9 changed files with 192 additions and 83 deletions
+2
View File
@@ -177,3 +177,5 @@ docs/docs_skeleton/yarn.lock
Untitled*.ipynb
Chinook.db
libs/langgraph/out
+15 -1
View File
@@ -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 \
View File
+63
View File
@@ -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())
+34 -1
View File
@@ -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"
+2
View File
@@ -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" ]
BIN
View File
Binary file not shown.
-81
View File
@@ -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())