diff --git a/.gitignore b/.gitignore index c69829e3a..770e2f194 100644 --- a/.gitignore +++ b/.gitignore @@ -177,3 +177,5 @@ docs/docs_skeleton/yarn.lock Untitled*.ipynb Chinook.db + +libs/langgraph/out diff --git a/libs/langgraph/Makefile b/libs/langgraph/Makefile index d5c6d7ee3..ee219870b 100644 --- a/libs/langgraph/Makefile +++ b/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 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 \ diff --git a/libs/langgraph/bench/__init__.py b/libs/langgraph/bench/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/libs/langgraph/bench/__main__.py b/libs/langgraph/bench/__main__.py new file mode 100644 index 000000000..45c9b3dfd --- /dev/null +++ b/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) diff --git a/libs/langgraph/bench/fanout_to_subgraph.py b/libs/langgraph/bench/fanout_to_subgraph.py new file mode 100644 index 000000000..53852609d --- /dev/null +++ b/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()) diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index 93d04e327..085c27366 100644 --- a/libs/langgraph/poetry.lock +++ b/libs/langgraph/poetry.lock @@ -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" diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index e5e3b7a16..fe29bbbf5 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -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" ] diff --git a/perf_test.prof b/perf_test.prof deleted file mode 100644 index 7ddbc252c..000000000 Binary files a/perf_test.prof and /dev/null differ diff --git a/perf_test.py b/perf_test.py deleted file mode 100644 index 7c4506525..000000000 --- a/perf_test.py +++ /dev/null @@ -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())