Apply patch [skip ci]

This commit is contained in:
open-swe[bot]
2025-09-09 23:44:59 +00:00
parent a35ea66722
commit bd5d458ad4
2 changed files with 63 additions and 79 deletions

View File

@@ -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)

View File

@@ -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():