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 <gabriel@logspace.ai>
This commit is contained in:
Jacob Hatchett
2026-05-18 12:40:42 -07:00
committed by GitHub
parent c78e10e85a
commit 7993ed3b70
2 changed files with 25 additions and 0 deletions

View File

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

View File

@ -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."""