feat(framework): Add context lock - #7975
Conversation
github-actions
Bot
added
the
Maintainer
There was a problem hiding this comment.
Pull request overview
This PR adds concurrency protection for Context across both in-process mutations (via an RLock) and persisted storage (via a version field and optimistic compare-and-set updates), with corresponding API errors, schema/migration updates, protobuf regeneration, and tests.
Changes:
- Add
Context.versionand serialize/deserialize it through protobuf andserde. - Implement optimistic CAS semantics for persisted run-series context updates (in-memory + SQL) and surface conflicts as HTTP 409 / gRPC ABORTED.
- Add/adjust tests and state/schema artifacts (SQL schema docs + Alembic migration).
Critical issues
- SQL
set_run_series_contextcan currently INSERT a new row even whencontext.version != 0, which violates the optimistic CAS contract and differs from the in-memory behavior (see stored PR comment ID: 001). _refresh_run_series_contextraisesRuntimeErroron CAS failure, which can fail run creation on benign concurrent updates; a retry (re-read latest + retry write) is safer (see stored PR comment ID: 002).
Simplicity/readability suggestions
- None.
Consistency concerns
_pull_and_store_messagenow returnsNoneon CAS conflict after receiving a message; the caller treats this as “no message” and sleeps, which contradicts the function’s docstring/contract (see stored PR comment ID: 003).
Whether the PR should be split
- No.
Overall verdict
Changes are directionally correct and well-covered by tests, but the SQL CAS edge case and the run-creation crash-on-race behavior should be addressed before approval.
Reviewed changes
Copilot reviewed 22 out of 23 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| framework/py/flwr/supernode/start_client_internal.py | Preserve and CAS-check run-series context version during run initialization. |
| framework/py/flwr/supernode/servicer/runtime/runtime_handlers.py | Persist context before finishing task; fail on CAS conflict with explicit API error. |
| framework/py/flwr/supernode/servicer/runtime/runtime_handlers_test.py | Add test ensuring stale context updates are rejected and task is failed. |
| framework/py/flwr/superlink/servicer/runtime/runtime_handlers.py | Persist primary-task context pre-finish and fail stale concurrent writes. |
| framework/py/flwr/superlink/servicer/runtime/runtime_handlers_test.py | Update tests for versioned context behavior and stale-write rejection. |
| framework/py/flwr/superlink/servicer/control/control_servicer_test.py | Update series-context tests to account for versioning/copy semantics. |
| framework/py/flwr/supercore/task_process/agent/context_items.py | Guard context state mutations with context.locked() for in-process safety. |
| framework/py/flwr/supercore/state/schema/README.md | Document new series_context.version column. |
| framework/py/flwr/supercore/state/schema/corestate_models.py | Add version column to SeriesContext ORM model. |
| framework/py/flwr/supercore/state/alembic/versions/rev_2026_08_22_add_run_series_context_version.py | Alembic migration adding version to series_context. |
| framework/py/flwr/supercore/error/catalog.py | Map new conflict error code to gRPC ABORTED / HTTP 409. |
| framework/py/flwr/supercore/error/base.py | Introduce RUNTIME_RUN_SERIES_CONTEXT_CONFLICT. |
| framework/py/flwr/supercore/corestate/sql_corestate.py | Implement versioned CAS update for persisted series context (SQL). |
| framework/py/flwr/supercore/corestate/in_memory_corestate.py | Implement versioned CAS update and deep-copy semantics (in-memory). |
| framework/py/flwr/supercore/corestate/corestate.py | Update CoreState interface: set_run_series_context now returns bool. |
| framework/py/flwr/supercore/corestate/corestate_test.py | Add roundtrip + stale-write tests for versioned context CAS behavior. |
| framework/py/flwr/server/superlink/linkstate/linkstate.py | Preserve version during refresh; enforce CAS on refresh update. |
| framework/py/flwr/proto/message_pb2.pyi | Regenerated stubs reflecting Context.version. |
| framework/py/flwr/proto/message_pb2.py | Regenerated protobuf module reflecting Context.version. |
| framework/py/flwr/common/serde.py | Lock during context serialization and include version field. |
| framework/py/flwr/common/serde_test.py | Extend serde tests to include Context.version. |
| framework/py/flwr/app/message/context.py | Add version and locked() (RLock) to Context. |
| framework/proto/flwr/proto/message.proto | Add version field to Context protobuf definition. |
Files not reviewed (1)
- framework/py/flwr/proto/message_pb2.py: Generated file
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| stmt = self.dialect_insert(SeriesContextModel).values( | ||
| series_id=sint_series_id, | ||
| context=context_bytes, | ||
| version=next_version, | ||
| ) |
| if existing_context := self.get_run_series_context(series_id): | ||
| context.state = existing_context.state | ||
| self.set_run_series_context(series_id=series_id, context=context) | ||
| context.version = existing_context.version | ||
| if not self.set_run_series_context(series_id=series_id, context=context): | ||
| raise RuntimeError("Run series context changed while refreshing it.") |
| if not state.set_run_series_context(run_info.series_id, context): | ||
| log(ERROR, "Run series context changed while initializing the run.") | ||
| return None |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 02a5e311c2
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| self.run_config = run_config | ||
| self.series_id = series_id | ||
| self.version = version | ||
| self._lock = RLock() |
There was a problem hiding this comment.
Every Context now has _lock in vars(context), but LegacyContext.__init__ still calls super().__init__(**vars(context)); consequently, constructing the public LegacyContext used by the repository's ServerApp examples and e2e apps raises TypeError: Context.__init__() got an unexpected keyword argument '_lock'. Copy only declared context fields or otherwise exclude _lock.
AGENTS.md reference: framework/AGENTS.md:L113-L115
Useful? React with 👍 / 👎.
| if not state.set_run_series_context( | ||
| run.series_id, | ||
| context_from_proto(request.context), | ||
| ): |
There was a problem hiding this comment.
When two PushTaskOutput requests authenticate concurrently for the same task, one can persist the context here while the other observes the new version, marks the task failed, and raises a conflict before the first request calls finish_task; the first completion then fails because the task is already finished, leaving a valid context persisted but its task marked failed. The SuperNode handler has the same ordering, so the context CAS and task transition need one atomic operation or equivalent coordination.
Useful? React with 👍 / 👎.
| if not self.set_run_series_context(series_id=series_id, context=context): | ||
| raise RuntimeError("Run series context changed while refreshing it.") |
There was a problem hiding this comment.
When a previous task updates the same series context between the read above and this CAS, a normal same-series StartRun raises this uncaught RuntimeError instead of retrying or returning the established creation failure result. Both state implementations have already persisted the run-series membership before calling this method, so the failed request can also leave an association to a run ID that was never inserted; retry the refresh with the latest context or roll back the membership.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 00d3033de6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| self.state = state | ||
| self.run_config = run_config | ||
| self.series_id = series_id | ||
| self._lock = RLock() |
There was a problem hiding this comment.
When a simulation or VCE Ray backend submits a Context to ClientAppActor.run (for example, ray_actor.py passes it as a remote argument), Ray must cloudpickle the object, but _thread.RLock is not serializable; ray.cloudpickle.dumps(Context(...)) now raises TypeError: cannot pickle '_thread.RLock' object before the ClientApp runs. Exclude the lock from serialized state and recreate it when unpickling so existing Ray-based simulations continue to work.
AGENTS.md reference: framework/AGENTS.md:L113-L115
Useful? React with 👍 / 👎.
Summary
Implemented concurrency protections for Context.
compound mutations can use with context.locked():.
stale executions from overwriting newer context.