diff --git a/libs/core/langchain_core/runnables/config.py b/libs/core/langchain_core/runnables/config.py index 28b8eb80b94..fc492cfc325 100644 --- a/libs/core/langchain_core/runnables/config.py +++ b/libs/core/langchain_core/runnables/config.py @@ -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) - - - diff --git a/libs/core/tests/unit_tests/runnables/test_batch_ordering.py b/libs/core/tests/unit_tests/runnables/test_batch_ordering.py index 7e2ad6f1d39..719968fb5d5 100644 --- a/libs/core/tests/unit_tests/runnables/test_batch_ordering.py +++ b/libs/core/tests/unit_tests/runnables/test_batch_ordering.py @@ -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():