fix: avoid pickling vertex component runtime state

[gKev1n]

Summary

• Fixes Flows not working when using celery rabitmq and redis #8476 on main by removing live custom_component runtime instances from Vertex.__getstate__() output.
• Keeps serializable graph state while allowing component instances to be rebuilt after deserialization.
• Adds a focused regression test with an intentionally unpickleable component to ensure vertex state remains pickleable.

Root cause

Distributed/Celery execution can pickle vertex state. A live component instance may contain runtime-only objects that cannot be pickled, so keeping custom_component in the serialized vertex state can fail with errors like cannot pickle ... object.

Verification

• uv run --package lfx pytest src/lfx/tests/unit/graph/vertex/test_vertex_base.py -k getstate -q
• uv run --package lfx pytest src/lfx/tests/unit/graph/vertex/test_vertex_base.py -q
• uv run ruff check src/lfx/src/lfx/graph/vertex/base.py src/lfx/tests/unit/graph/vertex/test_vertex_base.py
• git diff --check

Result: targeted test passed, full vertex test file passed (23 passed), ruff passed, and git diff --check passed.

Summary by CodeRabbit

• Bug Fixes

  • Vertex objects now properly serialize without including runtime component state, preventing serialization errors.

• Tests

  • Added regression test to verify vertex serialization behavior.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: f81bcefd-6a0f-44a1-8e98-061f4543ef90

📥 Commits

Reviewing files that changed from the base of the PR and between 500b5e9 and ce2f796.

📒 Files selected for processing (2)

• src/lfx/src/lfx/graph/vertex/base.py
• src/lfx/tests/unit/graph/vertex/test_vertex_base.py

---

Walkthrough

Vertex.__getstate__ now excludes the runtime custom_component field during serialization by setting it to None. A regression test validates that the resulting state is safely pickleable, addressing serialization failures when component instances cannot be pickled.

Changes

Vertex Serialization Fix

Layer / File(s)	Summary
Exclude runtime component from serialization src/lfx/src/lfx/graph/vertex/base.py	Vertex.__getstate__ clears custom_component to None during serialization, preventing runtime component instances from being pickled.
Pickling behavior test src/lfx/tests/unit/graph/vertex/test_vertex_base.py	New test test_vertex_getstate_drops_custom_component_runtime_state verifies that Vertex.__getstate__() excludes the custom_component field and that the result is safely pickleable.

---

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 7 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Test Quality And Coverage	⚠️ Warning	Test validates main functionality (preventing component pickling) but lacks comprehensive coverage: missing setstate and round-trip serialization tests.	Add tests for setstate, round-trip pickle/unpickle, assertions on cleared attributes, and tests with fully initialized Vertex objects for realistic behavior validation.

✅ Passed checks (7 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main change: preventing pickling of vertex component runtime state by clearing the custom_component field during serialization.
Linked Issues check	✅ Passed	The PR directly addresses issue #8476 by implementing the required fix to make vertex serialization pickleable for Celery/distributed execution.
Out of Scope Changes check	✅ Passed	All changes are directly scoped to the objective: modifying Vertex.getstate() to clear custom_component and adding a regression test.
Test Coverage For New Implementations	✅ Passed	Regression test added in test_vertex_base.py validates getstate() clears custom_component and state is pickleable. Test uses unpickleable component to verify fix without placeholder assertions.
Test File Naming And Structure	✅ Passed	Test file follows required patterns: test_*.py naming, pytest structure with fixtures, descriptive test names, organized logically, includes edge cases and both positive/negative scenarios.
Excessive Mock Usage Warning	✅ Passed	New test uses only 1 Mock (for _lock which isn't accessed), real objects elsewhere, and a clever test double (UnpickleableComponent) to validate the fix—demonstrates minimal, appropriate mock usage.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gKev1n]

Thanks CodeRabbit, I see the latest review has no actionable comments.

The PR is currently blocked by an unrelated CI failure in Run Frontend Tests / Playwright Tests - Shard 24/70 (Final attempt failed. Child_process exited with error code 1). This PR only touches:

• src/lfx/src/lfx/graph/vertex/base.py
• src/lfx/tests/unit/graph/vertex/test_vertex_base.py

Local verification for the affected backend path still passes:

uv run --package lfx pytest src/lfx/tests/unit/graph/vertex/test_vertex_base.py -k getstate -q
uv run --package lfx pytest src/lfx/tests/unit/graph/vertex/test_vertex_base.py -q
uv run ruff check src/lfx/src/lfx/graph/vertex/base.py src/lfx/tests/unit/graph/vertex/test_vertex_base.py
git diff --check

I do not have permission to rerun failed jobs on this repo from the fork. Could a maintainer rerun the failed frontend shard if CI is still required?