diff --git a/src/lfx/src/lfx/custom/custom_component/component.py b/src/lfx/src/lfx/custom/custom_component/component.py index e56f340002..df13b5ec86 100644 --- a/src/lfx/src/lfx/custom/custom_component/component.py +++ b/src/lfx/src/lfx/custom/custom_component/component.py @@ -382,12 +382,34 @@ class Component(CustomComponent): return memo[id(self)] # Shallow-copy config/inputs: they may contain non-picklable services # (e.g. _tracing_service holds ServiceManager with threading.RLock). - kwargs = dict(self.__config) - kwargs["inputs"] = dict(self.__inputs) + # use the mangled names to access the private attributes + config = getattr(self, "_Component__config", {}) + inputs_raw = getattr(self, "_Component__inputs", {}) + + kwargs = dict(config) + kwargs["inputs"] = dict(inputs_raw) new_component = type(self)(**kwargs) new_component._code = self._code new_component._outputs_map = self._outputs_map - new_component._inputs = deepcopy(self._inputs, memo) + + # Safe deepcopy of inputs + new_inputs = {} + for k, v in self._inputs.items(): + try: + # Attempt to deepcopy the entire input object + new_inputs[k] = deepcopy(v, memo) + except Exception: # noqa: BLE001 + # If deepcopy fails (e.g. due to RLock), handle the value carefully + # Pydantic's model_copy(deep=False) creates a shallow copy + input_copy = v.model_copy() + try: + input_copy.value = deepcopy(v.value, memo) + except Exception: # noqa: BLE001 + # Keep the original value (shallow copy) if it can't be deepcopied + input_copy.value = v.value + new_inputs[k] = input_copy + + new_component._inputs = new_inputs new_component._edges = self._edges new_component._components = self._components new_component._parameters = dict(self._parameters) diff --git a/src/lfx/tests/unit/custom/component/test_concurrent_tool_invocation.py b/src/lfx/tests/unit/custom/component/test_concurrent_tool_invocation.py index 5f10a52047..ad67196645 100644 --- a/src/lfx/tests/unit/custom/component/test_concurrent_tool_invocation.py +++ b/src/lfx/tests/unit/custom/component/test_concurrent_tool_invocation.py @@ -11,7 +11,7 @@ from copy import deepcopy from lfx.base.tools.component_tool import ComponentToolkit from lfx.custom.custom_component.component import Component -from lfx.inputs.inputs import MessageTextInput +from lfx.inputs.inputs import DataInput, MessageTextInput from lfx.io import Output from lfx.schema.data import Data @@ -149,6 +149,77 @@ def test_should_isolate_inputs_when_component_has_non_picklable_state(): assert product_ids == {"P1", "P2"} +def test_deepcopy_with_non_picklable_input_value(): + """Deepcopy must not fail when an input value is non-picklable. + + Some components receive complex objects (e.g. LangChain models/clients) + in their inputs that hold threading.RLock instances. + """ + + class MockComponentWithLockValue(Component): + inputs = [DataInput(name="input_with_lock")] + + def build(self): + return "ok" + + component = MockComponentWithLockValue() + lock_val = _FakeServiceWithLock() + component.set(input_with_lock=lock_val) + + # Must not raise "cannot pickle '_thread.RLock' object" + clone = deepcopy(component) + + assert clone is not component + # The non-picklable value should be shared (shallow-copied fallback) + assert clone.input_with_lock is lock_val + + +def test_should_isolate_inputs_when_input_has_non_picklable_value(): + """End-to-end: concurrent tool invocation must work even with non-picklable input values. + + Verified that the tool isolation still works when one of the inputs + carries a non-picklable object (forcing the shallow-copy fallback in deepcopy). + """ + + class SlowToolWithLock(SlowLabelComponent): + inputs = [ + *SlowLabelComponent.inputs, + DataInput(name="lock_input", tool_mode=True), + ] + + # Arrange + lock_val = _FakeServiceWithLock() + component = SlowToolWithLock() + component.set(lock_input=lock_val) # Set the non-picklable value on the component + toolkit = ComponentToolkit(component=component) + tools = toolkit.get_tools() + tool = tools[0] + + results = [] + + def invoke_tool(product_id: str, label: str) -> None: + # We don't pass lock_input here to avoid Pydantic validation of tool arguments + # The deepcopy(component) inside tool.invoke() will still encounter lock_val + result = tool.invoke({"product_id": product_id, "label": label}) + results.append(result) + + # Act + with ThreadPoolExecutor(max_workers=2) as executor: + future1 = executor.submit(invoke_tool, "P1", "L1") + future2 = executor.submit(invoke_tool, "P2", "L2") + future1.result() + future2.result() + + # Assert + assert len(results) == 2 + for result in results: + assert result["product_id_before"] == result["product_id_after"] + assert result["label_before"] == result["label_after"] + + product_ids = {r["product_id_before"] for r in results} + assert product_ids == {"P1", "P2"} + + # --------------------------------------------------------------------------- # Helpers # ---------------------------------------------------------------------------