From 7993ed3b705113ad1bfa031da770ab22d0943b8d Mon Sep 17 00:00:00 2001 From: Jacob Hatchett <152253273+newmattock@users.noreply.github.com> Date: Mon, 18 May 2026 12:40:42 -0700 Subject: [PATCH] fix: avoid pickling vertex component instances (#13172) * Avoid pickling vertex component runtime state * Add docstrings for vertex serialization regression helper --------- Co-authored-by: Gabriel Luiz Freitas Almeida --- src/lfx/src/lfx/graph/vertex/base.py | 1 + .../unit/graph/vertex/test_vertex_base.py | 24 +++++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/src/lfx/src/lfx/graph/vertex/base.py b/src/lfx/src/lfx/graph/vertex/base.py index f18b4f1d69..d5b8acfe3d 100644 --- a/src/lfx/src/lfx/graph/vertex/base.py +++ b/src/lfx/src/lfx/graph/vertex/base.py @@ -227,6 +227,7 @@ class Vertex: def __getstate__(self): state = self.__dict__.copy() state["_lock"] = None # Locks are not serializable + state["custom_component"] = None # Component instances are rebuilt and may hold non-serializable runtime state state["built_object"] = None if isinstance(self.built_object, UnbuiltObject) else self.built_object state["built_result"] = None if isinstance(self.built_result, UnbuiltResult) else self.built_result return state diff --git a/src/lfx/tests/unit/graph/vertex/test_vertex_base.py b/src/lfx/tests/unit/graph/vertex/test_vertex_base.py index 11b47857bb..93359fa148 100644 --- a/src/lfx/tests/unit/graph/vertex/test_vertex_base.py +++ b/src/lfx/tests/unit/graph/vertex/test_vertex_base.py @@ -4,6 +4,7 @@ This module contains tests for verifying the functionality of the ParameterHandl which is responsible for processing and managing parameters in vertices. """ +import pickle from unittest.mock import Mock from uuid import uuid4 @@ -18,6 +19,29 @@ from lfx.services.storage.service import StorageService from lfx.utils.util import unescape_string +def test_vertex_getstate_drops_custom_component_runtime_state(): + """Graph cache serialization should rebuild component instances instead of pickling live runtime state.""" + + class UnpickleableComponent: + """Component stand-in that fails if live runtime state is serialized.""" + + def __getstate__(self): + """Raise when pickle tries to serialize the component instance.""" + msg = "cannot pickle component runtime state" + raise TypeError(msg) + + vertex = object.__new__(Vertex) + vertex._lock = Mock() + vertex.custom_component = UnpickleableComponent() + vertex.built_object = object() + vertex.built_result = object() + + state = vertex.__getstate__() + + assert state["custom_component"] is None + pickle.dumps(state) + + @pytest.fixture def mock_storage_service() -> Mock: """Create a mock storage service for testing."""