mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-24 08:32:24 +02:00
fix(langgraph): detect subgraphs from bytecode instead of source (#8569)
Fixes langchain-ai/langgraph#8559 `find_subgraph_pregel` runs once per node at build time and recovers each node function's reachable values with `inspect.getsource` + `ast.parse`; langchain-ai/langgraph#8559 measures that source parsing at 80% of `StateGraph.compile()`. This replaces it with a `dis` walk over `func.__code__`, which reads the same information already in memory. Closure cells and the globals named in `co_names` supply the values directly, and a walk over the instruction stream recovers the attribute paths the function actually takes. The closure alone cannot express those: a captured `holder` whose graph lives at `holder.graph` is reachable only if something records that `graph` is loaded off `holder`. Both implementations over-declare — a node can reference a graph it never invokes — and this one over-declares a different set. It no longer reports a graph named only along an attribute path in code the compiler removed; such a path cannot execute, so that entry was always a phantom. Nothing that can actually run stopped being detected. `find_subgraph_pregel` returns the first `PregelProtocol` it finds, so the remaining extra candidates only widen detection. ## Release note Subgraph auto-detection now reads node functions' bytecode instead of parsing their source. Graph builds with many function-backed nodes are substantially faster, and subgraphs are now detected inside functions with no retrievable source — defined in a REPL or notebook cell, or via `exec` — where detection previously failed silently and returned nothing. ## Performance Against [langgraph-build-bench](<https://github.com/soarez/langgraph-build-bench>) — 713 nodes, 500 state fields, 264 tools; CPython 3.12.8, Apple silicon: ``` before build 1412 ms subgraph detection 1132 ms (80%) after build 286 ms subgraph detection 4 ms (2%) ``` **4.9x on total build, \~266x on detection.** The bench prints whole ms; at full precision detection is 1139 ms -> 4.23 ms. Co-authored-by: Elior Nataf Lackritz <elior.nataflackritz@langchain.dev>
This commit is contained in:
Igor Soarez
co-authored by
Elior Nataf Lackritz
Elior Nataf Lackritz
parent
644815f9e5
commit
1e44bda48f
@@ -1,11 +1,10 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import inspect
|
||||
import dis
|
||||
import re
|
||||
import textwrap
|
||||
from collections.abc import Callable, Sequence
|
||||
from functools import partial
|
||||
from types import CodeType, FunctionType
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.runnables import (
|
||||
@@ -17,7 +16,6 @@ from langchain_core.runnables import (
|
||||
from langchain_core.runnables.base import RunnableBindingBase
|
||||
from langchain_core.runnables.config import run_in_executor
|
||||
from langgraph.checkpoint.base import ChannelVersions
|
||||
from typing_extensions import override
|
||||
|
||||
from langgraph._internal._runnable import RunnableCallable, RunnableSeq
|
||||
from langgraph._internal._timeout import sync_timeout_unsupported
|
||||
@@ -137,155 +135,87 @@ def validate_timeout_supported(runnable: Runnable, *, name: str) -> None:
|
||||
raise sync_timeout_unsupported(name)
|
||||
|
||||
|
||||
# Values treated as dead ends when deciding whether to walk a function's
|
||||
# bytecode. A container can hold a graph, but `find_subgraph_pregel` does not
|
||||
# look inside one, so skipping it costs nothing while that holds. Matched by
|
||||
# exact type, since a subclass of a builtin can carry attributes.
|
||||
_LEAF_TYPES = frozenset(
|
||||
{
|
||||
int,
|
||||
float,
|
||||
complex,
|
||||
bool,
|
||||
str,
|
||||
bytes,
|
||||
bytearray,
|
||||
list,
|
||||
tuple,
|
||||
dict,
|
||||
set,
|
||||
frozenset,
|
||||
type(None),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def get_function_nonlocals(func: Callable) -> list[Any]:
|
||||
"""Get the nonlocal variables accessed by a function.
|
||||
"""Get the values a function reaches from outside its own scope.
|
||||
|
||||
Args:
|
||||
func: The function to check.
|
||||
|
||||
Returns:
|
||||
List[Any]: The nonlocal variables accessed by the function.
|
||||
Every captured cell value, the globals the function names, and each
|
||||
value along an attribute path it loads. Over-approximates: a value can
|
||||
come back without the function reaching it at runtime.
|
||||
"""
|
||||
try:
|
||||
code = inspect.getsource(func)
|
||||
tree = ast.parse(textwrap.dedent(code))
|
||||
visitor = FunctionNonLocals()
|
||||
visitor.visit(tree)
|
||||
values: list[Any] = []
|
||||
closure = (
|
||||
inspect.getclosurevars(func.__wrapped__)
|
||||
if hasattr(func, "__wrapped__") and callable(func.__wrapped__)
|
||||
else inspect.getclosurevars(func)
|
||||
)
|
||||
candidates = {**closure.globals, **closure.nonlocals}
|
||||
for k, v in candidates.items():
|
||||
if k in visitor.nonlocals:
|
||||
values.append(v)
|
||||
for kk in visitor.nonlocals:
|
||||
if "." in kk and kk.startswith(k):
|
||||
vv = v
|
||||
for part in kk.split(".")[1:]:
|
||||
if vv is None:
|
||||
break
|
||||
else:
|
||||
try:
|
||||
vv = getattr(vv, part)
|
||||
except AttributeError:
|
||||
break
|
||||
else:
|
||||
values.append(vv)
|
||||
except (SyntaxError, TypeError, OSError, SystemError):
|
||||
func = getattr(func, "__func__", func) # bound method -> function
|
||||
wrapped = getattr(func, "__wrapped__", None)
|
||||
if callable(wrapped):
|
||||
func = getattr(wrapped, "__func__", wrapped)
|
||||
if not isinstance(func, FunctionType):
|
||||
return []
|
||||
code = func.__code__
|
||||
|
||||
cells: dict[str, Any] = {}
|
||||
for name, cell in zip(code.co_freevars, func.__closure__ or ()):
|
||||
try:
|
||||
cells[name] = cell.cell_contents
|
||||
except ValueError:
|
||||
continue # empty cell: a recursive def not yet bound
|
||||
|
||||
# Every captured value counts, referenced or not: over-declaring costs an
|
||||
# introspection entry, under-declaring drops the subgraph's checkpoints and
|
||||
# stream events. Checking each cell against the bytecode would cost more and
|
||||
# only trade the cheap error for the expensive one.
|
||||
values: list[Any] = list(cells.values())
|
||||
global_ns = func.__globals__
|
||||
globals_ = {name: global_ns[name] for name in code.co_names if name in global_ns}
|
||||
if all(type(v) in _LEAF_TYPES for v in (*cells.values(), *globals_.values())):
|
||||
return values
|
||||
|
||||
# Nested code objects hold the references made by inner defs, lambdas and
|
||||
# comprehensions, which resolve against the namespaces gathered above.
|
||||
codes = [code]
|
||||
for c in codes:
|
||||
codes.extend(k for k in c.co_consts if isinstance(k, CodeType))
|
||||
value: Any = None
|
||||
for instruction in dis.get_instructions(c):
|
||||
opname = instruction.opname
|
||||
if opname == "LOAD_GLOBAL":
|
||||
value = globals_.get(instruction.argval)
|
||||
elif opname == "LOAD_DEREF":
|
||||
value = cells.get(instruction.argval)
|
||||
elif opname in ("LOAD_ATTR", "LOAD_METHOD"):
|
||||
value = getattr(value, instruction.argval, None)
|
||||
else:
|
||||
value = None # anything else ends the chain: `a, b.c` is not `a.c`
|
||||
continue
|
||||
if value is not None:
|
||||
values.append(value)
|
||||
return values
|
||||
|
||||
|
||||
class FunctionNonLocals(ast.NodeVisitor):
|
||||
"""Get the nonlocal variables accessed of a function."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.nonlocals: set[str] = set()
|
||||
|
||||
@override
|
||||
def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:
|
||||
"""Visit a function definition.
|
||||
|
||||
Args:
|
||||
node: The node to visit.
|
||||
|
||||
Returns:
|
||||
Any: The result of the visit.
|
||||
"""
|
||||
visitor = NonLocals()
|
||||
visitor.visit(node)
|
||||
self.nonlocals.update(visitor.loads - visitor.stores)
|
||||
|
||||
@override
|
||||
def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any:
|
||||
"""Visit an async function definition.
|
||||
|
||||
Args:
|
||||
node: The node to visit.
|
||||
|
||||
Returns:
|
||||
Any: The result of the visit.
|
||||
"""
|
||||
visitor = NonLocals()
|
||||
visitor.visit(node)
|
||||
self.nonlocals.update(visitor.loads - visitor.stores)
|
||||
|
||||
@override
|
||||
def visit_Lambda(self, node: ast.Lambda) -> Any:
|
||||
"""Visit a lambda function.
|
||||
|
||||
Args:
|
||||
node: The node to visit.
|
||||
|
||||
Returns:
|
||||
Any: The result of the visit.
|
||||
"""
|
||||
visitor = NonLocals()
|
||||
visitor.visit(node)
|
||||
self.nonlocals.update(visitor.loads - visitor.stores)
|
||||
|
||||
|
||||
class NonLocals(ast.NodeVisitor):
|
||||
"""Get nonlocal variables accessed."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.loads: set[str] = set()
|
||||
self.stores: set[str] = set()
|
||||
|
||||
@override
|
||||
def visit_Name(self, node: ast.Name) -> Any:
|
||||
"""Visit a name node.
|
||||
|
||||
Args:
|
||||
node: The node to visit.
|
||||
|
||||
Returns:
|
||||
Any: The result of the visit.
|
||||
"""
|
||||
if isinstance(node.ctx, ast.Load):
|
||||
self.loads.add(node.id)
|
||||
elif isinstance(node.ctx, ast.Store):
|
||||
self.stores.add(node.id)
|
||||
|
||||
@override
|
||||
def visit_Attribute(self, node: ast.Attribute) -> Any:
|
||||
"""Visit an attribute node.
|
||||
|
||||
Args:
|
||||
node: The node to visit.
|
||||
|
||||
Returns:
|
||||
Any: The result of the visit.
|
||||
"""
|
||||
if isinstance(node.ctx, ast.Load):
|
||||
parent = node.value
|
||||
attr_expr = node.attr
|
||||
while isinstance(parent, ast.Attribute):
|
||||
attr_expr = parent.attr + "." + attr_expr
|
||||
parent = parent.value
|
||||
if isinstance(parent, ast.Name):
|
||||
self.loads.add(parent.id + "." + attr_expr)
|
||||
self.loads.discard(parent.id)
|
||||
elif isinstance(parent, ast.Call):
|
||||
if isinstance(parent.func, ast.Name):
|
||||
self.loads.add(parent.func.id)
|
||||
else:
|
||||
parent = parent.func
|
||||
attr_expr = ""
|
||||
while isinstance(parent, ast.Attribute):
|
||||
if attr_expr:
|
||||
attr_expr = parent.attr + "." + attr_expr
|
||||
else:
|
||||
attr_expr = parent.attr
|
||||
parent = parent.value
|
||||
if isinstance(parent, ast.Name):
|
||||
self.loads.add(parent.id + "." + attr_expr)
|
||||
|
||||
|
||||
def is_xxh3_128_hexdigest(value: str) -> bool:
|
||||
"""Check if the given string matches the format of xxh3_128_hexdigest."""
|
||||
return bool(re.fullmatch(r"[0-9a-f]{32}", value))
|
||||
|
||||
@@ -0,0 +1,286 @@
|
||||
"""Tests for subgraph auto-detection (`pregel/_utils.py`).
|
||||
|
||||
Detection failing is silent — the graph still runs, only introspection goes
|
||||
quiet — so every shape a node can hold a graph in is pinned here. The expected
|
||||
values are what the source-parsing implementation this replaced produced for
|
||||
the same shapes, except for `sourceless`, whose source it could not read,
|
||||
`empty_closure_cell`, on which it raised, and `unreachable_attribute_chain`,
|
||||
where it reported a graph that dropped code could never invoke.
|
||||
"""
|
||||
|
||||
import functools
|
||||
import operator
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.pregel._utils import get_function_nonlocals
|
||||
|
||||
|
||||
class State(TypedDict):
|
||||
log: Annotated[list, operator.add]
|
||||
|
||||
|
||||
def _leaf(tag: str) -> Any:
|
||||
"""Return a compiled graph that reports itself as `tag`."""
|
||||
builder = StateGraph(State)
|
||||
builder.add_node(tag, lambda s: {"log": [tag]})
|
||||
builder.add_edge(START, tag)
|
||||
builder.add_edge(tag, END)
|
||||
compiled = builder.compile()
|
||||
compiled.name = tag
|
||||
return compiled
|
||||
|
||||
|
||||
def _detect(node: Any) -> str | None:
|
||||
"""Return the name of the subgraph detected for `node`, or None."""
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("n", node)
|
||||
builder.add_edge(START, "n")
|
||||
builder.add_edge("n", END)
|
||||
subgraphs = builder.compile().nodes["n"].subgraphs
|
||||
return getattr(subgraphs[0], "name", "?") if subgraphs else None
|
||||
|
||||
|
||||
class _Box:
|
||||
def __init__(self, payload: Any) -> None:
|
||||
self.payload = payload
|
||||
|
||||
|
||||
class _ListSubclass(list):
|
||||
pass
|
||||
|
||||
|
||||
class _MethodHolder:
|
||||
def __init__(self) -> None:
|
||||
self.graph = _leaf("via_self")
|
||||
|
||||
def as_node(self, state: State) -> Any:
|
||||
return self.graph.invoke(state)
|
||||
|
||||
|
||||
MODULE_GRAPH = _leaf("module_global")
|
||||
CHAIN = _Box(_Box(_leaf("attr_chain")))
|
||||
GRAPH_IN_PLAIN_LIST = [_leaf("in_list")]
|
||||
METHOD_HOLDER = _MethodHolder()
|
||||
|
||||
|
||||
def closure_capture() -> Any:
|
||||
sub = _leaf("closure")
|
||||
|
||||
def node(state: State) -> Any:
|
||||
return sub.invoke(state)
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def module_global() -> Any:
|
||||
def node(state: State) -> Any:
|
||||
return MODULE_GRAPH.invoke(state)
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def attribute_chain() -> Any:
|
||||
def node(state: State) -> Any:
|
||||
return CHAIN.payload.payload.invoke(state)
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def nested_def_captured_attribute() -> Any:
|
||||
"""A chain on a captured holder, named only inside a nested code object.
|
||||
|
||||
The captured value is the holder, not the graph, so the chain itself has to
|
||||
be recovered from the nested scope.
|
||||
"""
|
||||
holder = _Box(_leaf("nested_captured"))
|
||||
|
||||
def node(state: State) -> Any:
|
||||
def inner() -> Any:
|
||||
return holder.payload.invoke(state)
|
||||
|
||||
return inner()
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def unreachable_branch() -> Any:
|
||||
"""A captured graph referenced only from code the compiler removes."""
|
||||
sub = _leaf("unreachable")
|
||||
|
||||
def node(state: State) -> Any:
|
||||
if False:
|
||||
sub.invoke(state)
|
||||
return {"log": []}
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def unreachable_attribute_chain() -> Any:
|
||||
"""A graph named only along an attribute path the compiler dropped.
|
||||
|
||||
The closure keeps `holder`, but the `.payload` load is gone. The source
|
||||
parser reported this one; dropped code cannot invoke anything, so that was
|
||||
a phantom rather than a detection.
|
||||
"""
|
||||
holder = _Box(_leaf("unreachable_attr"))
|
||||
|
||||
def node(state: State) -> Any:
|
||||
if False:
|
||||
holder.payload.invoke(state)
|
||||
return {"log": []}
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def wrapper_referencing_nothing() -> Any:
|
||||
"""A wrapper whose own scope holds nothing, so only `__wrapped__` leads on.
|
||||
|
||||
`functools.wraps` would leave the wrapper closing over the inner function;
|
||||
setting the attribute by hand does not.
|
||||
"""
|
||||
sub = _leaf("via_wrapped")
|
||||
|
||||
def inner(state: State) -> Any:
|
||||
return sub.invoke(state)
|
||||
|
||||
def wrapper(state: State) -> Any:
|
||||
return {"log": []}
|
||||
|
||||
wrapper.__wrapped__ = inner
|
||||
return wrapper
|
||||
|
||||
|
||||
def captured_list_subclass() -> Any:
|
||||
"""A `list` subclass is not a leaf: it can carry a graph as an attribute."""
|
||||
holder = _ListSubclass()
|
||||
holder.payload = _leaf("list_subclass")
|
||||
|
||||
def node(state: State) -> Any:
|
||||
return holder.payload.invoke(state)
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def empty_closure_cell() -> Any:
|
||||
"""An unassigned closure variable leaves a cell that cannot be read."""
|
||||
sub = _leaf("beside_empty_cell")
|
||||
|
||||
def node(state: State) -> Any:
|
||||
return unassigned, sub.invoke(state)
|
||||
|
||||
return node
|
||||
unassigned = 1 # never runs, so the cell it creates is never filled
|
||||
|
||||
|
||||
def sourceless() -> Any:
|
||||
"""A node compiled without a source file, which `getsource` could not read."""
|
||||
namespace: dict[str, Any] = {"SOURCELESS": _leaf("sourceless")}
|
||||
exec(
|
||||
compile(
|
||||
"def node(state):\n return SOURCELESS.invoke(state)", "<test>", "exec"
|
||||
),
|
||||
namespace,
|
||||
)
|
||||
return namespace["node"]
|
||||
|
||||
|
||||
async def _async_node(state: State) -> Any:
|
||||
return await MODULE_GRAPH.ainvoke(state)
|
||||
|
||||
|
||||
def async_node() -> Any:
|
||||
return _async_node
|
||||
|
||||
|
||||
def no_subgraph() -> Any:
|
||||
"""Nothing but leaf values in reach, so the bytecode walk is skipped."""
|
||||
|
||||
def node(state: State) -> Any:
|
||||
return {"log": [len("abc") + 1]}
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def recombined_names() -> Any:
|
||||
"""Loads `CHAIN.payload` and `local.payload`, never `CHAIN.payload.payload`."""
|
||||
|
||||
def node(state: State) -> Any:
|
||||
local = _Box("not a graph")
|
||||
return {"log": [CHAIN.payload, local.payload]}
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def broken_attribute_chain() -> Any:
|
||||
holder = _Box("a string, so `.payload.missing` cannot resolve")
|
||||
|
||||
def node(state: State) -> Any:
|
||||
return holder.payload.missing.invoke(state)
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def nested_def_global() -> Any:
|
||||
"""A global named only in a nested code object: out of reach, as before."""
|
||||
|
||||
def node(state: State) -> Any:
|
||||
def inner() -> Any:
|
||||
return MODULE_GRAPH.invoke(state)
|
||||
|
||||
return inner()
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def graph_in_plain_list() -> Any:
|
||||
def node(state: State) -> Any:
|
||||
return GRAPH_IN_PLAIN_LIST[0].invoke(state)
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def bound_method_self() -> Any:
|
||||
return METHOD_HOLDER.as_node
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("factory", "expected"),
|
||||
[
|
||||
(closure_capture, "closure"),
|
||||
(module_global, "module_global"),
|
||||
(attribute_chain, "attr_chain"),
|
||||
(nested_def_captured_attribute, "nested_captured"),
|
||||
(unreachable_branch, "unreachable"),
|
||||
(wrapper_referencing_nothing, "via_wrapped"),
|
||||
(captured_list_subclass, "list_subclass"),
|
||||
(empty_closure_cell, "beside_empty_cell"),
|
||||
(sourceless, "sourceless"),
|
||||
(async_node, "module_global"),
|
||||
# Shapes no reference chain reaches: a subscript, an instance attribute
|
||||
# of `self`, a global named only in a nested scope, and an attribute
|
||||
# path the compiler dropped.
|
||||
(no_subgraph, None),
|
||||
(recombined_names, None),
|
||||
(broken_attribute_chain, None),
|
||||
(nested_def_global, None),
|
||||
(graph_in_plain_list, None),
|
||||
(bound_method_self, None),
|
||||
(unreachable_attribute_chain, None),
|
||||
],
|
||||
ids=lambda value: value.__name__ if callable(value) else str(value),
|
||||
)
|
||||
def test_subgraph_detection(factory: Any, expected: str | None) -> None:
|
||||
assert _detect(factory()) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"candidate",
|
||||
[functools.partial(lambda state, extra: {"log": [extra]}, extra="x"), len],
|
||||
ids=["partial", "builtin"],
|
||||
)
|
||||
def test_callables_without_a_code_object_are_handled(candidate: Any) -> None:
|
||||
assert get_function_nonlocals(candidate) == []
|
||||
Reference in New Issue
Block a user