mirror of
https://github.com/hwchase17/langchain.git
synced 2026-07-15 07:00:38 +00:00
Apply patch [skip ci]
This commit is contained in:
@@ -549,19 +549,19 @@ class ContextThreadPoolExecutor(ThreadPoolExecutor):
|
||||
"""
|
||||
# Convert iterables to lists to get their length and allow indexing
|
||||
iterables_as_lists = [list(iterable) for iterable in iterables]
|
||||
|
||||
|
||||
# Create contexts for each item in the first iterable
|
||||
contexts = [copy_context() for _ in range(len(iterables_as_lists[0]))]
|
||||
|
||||
|
||||
# Create a wrapper that includes the index with the arguments
|
||||
def _wrapped_fn_with_index(index: int, *args: Any) -> T:
|
||||
# Use the index to get the corresponding context
|
||||
return contexts[index].run(fn, *args)
|
||||
|
||||
|
||||
# Create indexed versions: first element is the index, rest are the actual values
|
||||
indices = list(range(len(iterables_as_lists[0])))
|
||||
indexed_iterables = [indices] + iterables_as_lists
|
||||
|
||||
|
||||
# Map the wrapped function over the indexed iterables
|
||||
return super().map(
|
||||
_wrapped_fn_with_index,
|
||||
@@ -625,6 +625,3 @@ async def run_in_executor(
|
||||
)
|
||||
|
||||
return await asyncio.get_running_loop().run_in_executor(executor_or_config, wrapper)
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1,22 +1,18 @@
|
||||
"""Test batch processing order preservation and thread-safety."""
|
||||
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any, Optional
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from langchain_core.runnables import Runnable, RunnableConfig, RunnableSequence
|
||||
from langchain_core.runnables import Runnable, RunnableConfig
|
||||
from langchain_core.runnables.config import ContextThreadPoolExecutor
|
||||
|
||||
|
||||
class OrderTrackingRunnable(Runnable[str, str]):
|
||||
"""A runnable that tracks processing order and adds unique identifiers."""
|
||||
|
||||
|
||||
def __init__(self, name: str, delay: float = 0.0, fail_on: Optional[str] = None):
|
||||
"""Initialize the runnable.
|
||||
|
||||
|
||||
Args:
|
||||
name: Name of this runnable for identification.
|
||||
delay: Artificial delay to simulate processing time.
|
||||
@@ -27,25 +23,22 @@ class OrderTrackingRunnable(Runnable[str, str]):
|
||||
self.fail_on = fail_on
|
||||
self.processed_items = []
|
||||
self.processing_order = []
|
||||
|
||||
|
||||
def invoke(
|
||||
self,
|
||||
input: str,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
**kwargs: Any
|
||||
self, input: str, config: Optional[RunnableConfig] = None, **kwargs: Any
|
||||
) -> str:
|
||||
"""Process input and track order."""
|
||||
# Track when this item started processing
|
||||
self.processing_order.append(input)
|
||||
|
||||
|
||||
# Simulate processing delay
|
||||
if self.delay > 0:
|
||||
time.sleep(self.delay)
|
||||
|
||||
|
||||
# Check if we should fail
|
||||
if self.fail_on and input.startswith(self.fail_on):
|
||||
raise ValueError(f"{self.name} failed on input: {input}")
|
||||
|
||||
|
||||
# Process and track
|
||||
result = f"{input}_{self.name}"
|
||||
self.processed_items.append((input, result))
|
||||
@@ -54,26 +47,23 @@ class OrderTrackingRunnable(Runnable[str, str]):
|
||||
|
||||
class ContextCapturingRunnable(Runnable[str, str]):
|
||||
"""A runnable that captures context information to verify uniqueness."""
|
||||
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the runnable."""
|
||||
self.contexts_seen = []
|
||||
self.input_context_map = {}
|
||||
|
||||
|
||||
def invoke(
|
||||
self,
|
||||
input: str,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
**kwargs: Any
|
||||
self, input: str, config: Optional[RunnableConfig] = None, **kwargs: Any
|
||||
) -> str:
|
||||
"""Process input and capture context."""
|
||||
import threading
|
||||
from contextvars import copy_context
|
||||
|
||||
|
||||
# Capture current context and thread info
|
||||
context_id = id(copy_context())
|
||||
thread_id = threading.current_thread().ident
|
||||
|
||||
|
||||
# Store context info
|
||||
context_info = {
|
||||
"input": input,
|
||||
@@ -82,7 +72,7 @@ class ContextCapturingRunnable(Runnable[str, str]):
|
||||
}
|
||||
self.contexts_seen.append(context_info)
|
||||
self.input_context_map[input] = context_info
|
||||
|
||||
|
||||
return f"{input}_ctx{context_id}"
|
||||
|
||||
|
||||
@@ -90,14 +80,14 @@ def test_batch_preserves_order():
|
||||
"""Test that batch processing preserves input/output order."""
|
||||
runnable = OrderTrackingRunnable("test", delay=0.01)
|
||||
inputs = [f"input_{i}" for i in range(10)]
|
||||
|
||||
|
||||
# Process batch
|
||||
outputs = runnable.batch(inputs)
|
||||
|
||||
|
||||
# Verify order is preserved
|
||||
expected_outputs = [f"input_{i}_test" for i in range(10)]
|
||||
assert outputs == expected_outputs, "Output order should match input order"
|
||||
|
||||
|
||||
# Verify all items were processed
|
||||
assert len(runnable.processed_items) == 10
|
||||
processed_inputs = [item[0] for item in runnable.processed_items]
|
||||
@@ -108,13 +98,13 @@ def test_batch_with_return_exceptions():
|
||||
"""Test that batch with return_exceptions=True preserves order."""
|
||||
runnable = OrderTrackingRunnable("test", delay=0.01, fail_on="fail")
|
||||
inputs = ["input_0", "fail_1", "input_2", "fail_3", "input_4"]
|
||||
|
||||
|
||||
# Process batch with return_exceptions=True
|
||||
outputs = runnable.batch(inputs, return_exceptions=True)
|
||||
|
||||
|
||||
# Verify we got 5 outputs
|
||||
assert len(outputs) == 5, "Should have one output per input"
|
||||
|
||||
|
||||
# Check that outputs are in correct positions
|
||||
assert outputs[0] == "input_0_test", "First output should be processed normally"
|
||||
assert isinstance(outputs[1], ValueError), "Second output should be an exception"
|
||||
@@ -126,18 +116,20 @@ def test_batch_with_return_exceptions():
|
||||
def test_batch_high_concurrency():
|
||||
"""Test batch processing with high concurrency to expose race conditions."""
|
||||
runnable = OrderTrackingRunnable("test", delay=0.001)
|
||||
|
||||
|
||||
# Create many inputs to increase chance of race conditions
|
||||
inputs = [f"input_{i:03d}" for i in range(100)]
|
||||
|
||||
|
||||
# Process with high concurrency
|
||||
config = {"max_concurrency": 20}
|
||||
outputs = runnable.batch(inputs, config=config)
|
||||
|
||||
|
||||
# Verify order is preserved despite high concurrency
|
||||
expected_outputs = [f"input_{i:03d}_test" for i in range(100)]
|
||||
assert outputs == expected_outputs, "Order should be preserved even with high concurrency"
|
||||
|
||||
assert outputs == expected_outputs, (
|
||||
"Order should be preserved even with high concurrency"
|
||||
)
|
||||
|
||||
# Verify no duplicates
|
||||
assert len(set(outputs)) == len(outputs), "No duplicate outputs should exist"
|
||||
assert len(outputs) == 100, "All inputs should produce outputs"
|
||||
@@ -147,13 +139,13 @@ def test_batch_unique_contexts():
|
||||
"""Test that each input gets a unique context without duplicates."""
|
||||
runnable = ContextCapturingRunnable()
|
||||
inputs = [f"input_{i}" for i in range(20)]
|
||||
|
||||
|
||||
# Process batch
|
||||
outputs = runnable.batch(inputs)
|
||||
|
||||
|
||||
# Verify each input was processed
|
||||
assert len(runnable.contexts_seen) == 20, "Each input should be processed once"
|
||||
|
||||
|
||||
# Verify each input got a unique processing slot
|
||||
inputs_seen = [ctx["input"] for ctx in runnable.contexts_seen]
|
||||
assert sorted(inputs_seen) == sorted(inputs), "All inputs should be processed"
|
||||
@@ -166,39 +158,38 @@ def test_sequence_batch_with_exceptions():
|
||||
step1 = OrderTrackingRunnable("step1", delay=0.001)
|
||||
step2 = OrderTrackingRunnable("step2", delay=0.001, fail_on="input_1")
|
||||
step3 = OrderTrackingRunnable("step3", delay=0.001)
|
||||
|
||||
|
||||
sequence = step1 | step2 | step3
|
||||
|
||||
|
||||
inputs = ["input_0", "input_1", "input_2"]
|
||||
outputs = sequence.batch(inputs, return_exceptions=True)
|
||||
|
||||
|
||||
# Verify we got 3 outputs
|
||||
assert len(outputs) == 3, "Should have one output per input"
|
||||
|
||||
|
||||
# First input should process through all steps
|
||||
assert outputs[0] == "input_0_step1_step2_step3"
|
||||
|
||||
|
||||
# Second input should fail at step2
|
||||
assert isinstance(outputs[1], ValueError)
|
||||
assert "step2 failed" in str(outputs[1])
|
||||
|
||||
|
||||
# Third input should process through all steps
|
||||
assert outputs[2] == "input_2_step1_step2_step3"
|
||||
|
||||
|
||||
def test_context_thread_pool_executor_map_order():
|
||||
"""Test ContextThreadPoolExecutor.map preserves order."""
|
||||
from contextvars import copy_context
|
||||
|
||||
|
||||
def process_item(item: str) -> str:
|
||||
# Simulate some work
|
||||
time.sleep(0.001)
|
||||
return f"processed_{item}"
|
||||
|
||||
|
||||
with ContextThreadPoolExecutor(max_workers=4) as executor:
|
||||
inputs = [f"item_{i}" for i in range(20)]
|
||||
results = list(executor.map(process_item, inputs))
|
||||
|
||||
|
||||
# Verify order is preserved
|
||||
expected = [f"processed_item_{i}" for i in range(20)]
|
||||
assert results == expected, "Map should preserve order"
|
||||
@@ -206,18 +197,17 @@ def test_context_thread_pool_executor_map_order():
|
||||
|
||||
def test_context_thread_pool_executor_map_with_multiple_iterables():
|
||||
"""Test ContextThreadPoolExecutor.map with multiple iterables."""
|
||||
from contextvars import copy_context
|
||||
|
||||
|
||||
def combine_items(a: str, b: int, c: float) -> str:
|
||||
return f"{a}_{b}_{c}"
|
||||
|
||||
|
||||
with ContextThreadPoolExecutor(max_workers=4) as executor:
|
||||
list_a = [f"a{i}" for i in range(10)]
|
||||
list_b = list(range(10))
|
||||
list_c = [float(i) * 0.5 for i in range(10)]
|
||||
|
||||
|
||||
results = list(executor.map(combine_items, list_a, list_b, list_c))
|
||||
|
||||
|
||||
# Verify order and correctness
|
||||
expected = [f"a{i}_{i}_{i * 0.5}" for i in range(10)]
|
||||
assert results == expected, "Map with multiple iterables should preserve order"
|
||||
@@ -225,18 +215,16 @@ def test_context_thread_pool_executor_map_with_multiple_iterables():
|
||||
|
||||
def test_batch_no_race_conditions():
|
||||
"""Test that batch processing has no race conditions with shared state."""
|
||||
|
||||
class SharedStateRunnable(Runnable[int, int]):
|
||||
"""A runnable that modifies shared state."""
|
||||
|
||||
|
||||
def __init__(self):
|
||||
self.counter = 0
|
||||
self.results = []
|
||||
|
||||
|
||||
def invoke(
|
||||
self,
|
||||
input: int,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
**kwargs: Any
|
||||
self, input: int, config: Optional[RunnableConfig] = None, **kwargs: Any
|
||||
) -> int:
|
||||
# Simulate race condition scenario
|
||||
current = self.counter
|
||||
@@ -245,18 +233,18 @@ def test_batch_no_race_conditions():
|
||||
result = input * 2
|
||||
self.results.append((input, result))
|
||||
return result
|
||||
|
||||
|
||||
runnable = SharedStateRunnable()
|
||||
inputs = list(range(50))
|
||||
|
||||
|
||||
# Process with high concurrency
|
||||
config = {"max_concurrency": 10}
|
||||
outputs = runnable.batch(inputs, config=config)
|
||||
|
||||
|
||||
# Verify outputs are correct and in order
|
||||
expected = [i * 2 for i in range(50)]
|
||||
assert outputs == expected, "Outputs should be correct and in order"
|
||||
|
||||
|
||||
# Note: The counter might not be 50 due to race conditions in the test runnable itself,
|
||||
# but the batch processing order should still be preserved
|
||||
|
||||
@@ -264,29 +252,28 @@ def test_batch_no_race_conditions():
|
||||
def test_batch_with_varying_processing_times():
|
||||
"""Test that batch preserves order even with varying processing times."""
|
||||
import random
|
||||
|
||||
|
||||
class VariableDelayRunnable(Runnable[str, str]):
|
||||
"""A runnable with variable processing delays."""
|
||||
|
||||
|
||||
def invoke(
|
||||
self,
|
||||
input: str,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
**kwargs: Any
|
||||
self, input: str, config: Optional[RunnableConfig] = None, **kwargs: Any
|
||||
) -> str:
|
||||
# Random delay to simulate varying processing times
|
||||
delay = random.uniform(0.001, 0.01)
|
||||
time.sleep(delay)
|
||||
return f"processed_{input}"
|
||||
|
||||
|
||||
runnable = VariableDelayRunnable()
|
||||
inputs = [f"item_{i:02d}" for i in range(20)]
|
||||
|
||||
|
||||
# Run multiple times to ensure consistency
|
||||
for _ in range(3):
|
||||
outputs = runnable.batch(inputs)
|
||||
expected = [f"processed_item_{i:02d}" for i in range(20)]
|
||||
assert outputs == expected, "Order should be preserved regardless of processing time"
|
||||
assert outputs == expected, (
|
||||
"Order should be preserved regardless of processing time"
|
||||
)
|
||||
|
||||
|
||||
def test_batch_empty_input():
|
||||
|
||||
Reference in New Issue
Block a user