mirror of
https://github.com/langflow-ai/langflow.git
synced 2026-07-25 20:20:08 +08:00
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:
@ -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
|
||||
|
||||
@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user