feat(framework): Add context lock

framework/py/flwr/app/message/context.py

@@ -15,7 +15,10 @@
 """Context."""
 
 
+from collections.abc import Iterator
+from contextlib import contextmanager
 from dataclasses import dataclass
+from threading import RLock
 
 from flwr.app.user_config import UserConfig
 
@@ -72,3 +75,10 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-positional-argumen
         self.state = state
         self.run_config = run_config
         self.series_id = series_id
+        self._lock = RLock()
+
+    @contextmanager
+    def locked(self) -> Iterator[None]:
+        """Lock this context for an atomic in-process operation."""
+        with self._lock:
+            yield

framework/py/flwr/common/serde.py

@@ -606,15 +606,15 @@ def message_from_proto(message_proto: ProtoMessage) -> Message:
 
 def context_to_proto(context: Context) -> ProtoContext:
     """Serialize `Context` to ProtoBuf."""
-    proto = ProtoContext(
-        run_id=context.run_id,
-        node_id=context.node_id,
-        node_config=user_config_to_proto(context.node_config),
-        state=recorddict_to_proto(context.state),
-        run_config=user_config_to_proto(context.run_config),
-        series_id=context.series_id,
-    )
-    return proto
+    with context.locked():
+        return ProtoContext(
+            run_id=context.run_id,
+            node_id=context.node_id,
+            node_config=user_config_to_proto(context.node_config),
+            state=recorddict_to_proto(context.state),
+            run_config=user_config_to_proto(context.run_config),
+            series_id=context.series_id,
+        )
 
 
 def context_from_proto(context_proto: ProtoContext) -> Context:

framework/py/flwr/server/superlink/linkstate/linkstate_test.py

@@ -218,6 +218,65 @@ def test_create_run_reuses_series_id_in_same_federation(self) -> None:
         runs = state.get_run_info(run_ids=[run_id_1, run_id_2])
         self.assertEqual({run.series_id for run in runs}, {first_run.series_id})
 
+    def test_claim_task_serializes_agentapp_runs_in_series(self) -> None:
+        """Only one AgentApp task in a run series should be claimable at a time."""
+        state = self.state_factory()
+        run_id_1 = create_dummy_run(state, primary_task_type=TaskType.AGENT_APP)
+        run_1 = state.get_run_info(run_ids=[run_id_1])[0]
+        run_id_2 = create_dummy_run(
+            state,
+            primary_task_type=TaskType.AGENT_APP,
+            series_id=run_1.series_id,
+        )
+        task_id_1 = get_primary_task_id(state, run_id_1)
+        task_id_2 = get_primary_task_id(state, run_id_2)
+
+        assert state.claim_task(task_id_1) is not None
+        assert state.claim_task(task_id_2) is None
+        assert task_id_2 in {
+            task.task_id for task in state.get_tasks(statuses=[Status.PENDING])
+        }
+        assert task_id_2 not in {
+            task.task_id
+            for task in state.get_tasks(statuses=[Status.PENDING], claimable=True)
+        }
+
+        assert state.finish_task(task_id_1, SubStatus.FAILED, "done")
+        assert state.claim_task(task_id_2) is not None
+
+    def test_claim_task_allows_parallel_non_agent_runs_in_series(self) -> None:
+        """Series serialization should not affect non-AgentApp tasks."""
+        state = self.state_factory()
+        run_id_1 = create_dummy_run(state, primary_task_type=TaskType.SERVER_APP)
+        run_1 = state.get_run_info(run_ids=[run_id_1])[0]
+        run_id_2 = create_dummy_run(
+            state,
+            primary_task_type=TaskType.SERVER_APP,
+            series_id=run_1.series_id,
+        )
+
+        assert state.claim_task(get_primary_task_id(state, run_id_1)) is not None
+        assert state.claim_task(get_primary_task_id(state, run_id_2)) is not None
+
+    def test_expired_agentapp_claim_releases_run_series(self) -> None:
+        """An expired AgentApp claim should let the next series run start."""
+        state = self.state_factory()
+        run_id_1 = create_dummy_run(state, primary_task_type=TaskType.AGENT_APP)
+        run_1 = state.get_run_info(run_ids=[run_id_1])[0]
+        run_id_2 = create_dummy_run(
+            state,
+            primary_task_type=TaskType.AGENT_APP,
+            series_id=run_1.series_id,
+        )
+        task_id_1 = get_primary_task_id(state, run_id_1)
+        task_id_2 = get_primary_task_id(state, run_id_2)
+        assert state.claim_task(task_id_1) is not None
+
+        patched_dt = now() + timedelta(seconds=HEARTBEAT_DEFAULT_INTERVAL + 1)
+        with patch("datetime.datetime") as mock_dt:
+            mock_dt.now.return_value = patched_dt
+            assert state.claim_task(task_id_2) is not None
+
     @parameterized.expand(  # type: ignore[untyped-decorator]
         [
             (TaskType.AGENT_APP, True),

framework/py/flwr/supercore/corestate/corestate.py

@@ -402,6 +402,7 @@ def set_run_series_context(self, series_id: int, context: Context) -> None:
             The ID of the RunSeries for which to persist shared context.
         context : Context
             The shared context to store.
+
         """
 
     @abstractmethod
@@ -699,6 +700,7 @@ def get_tasks(  # pylint: disable=too-many-arguments
         order_by: Literal["pending_at"] | None = None,
         ascending: bool = True,
         limit: int | None = None,
+        claimable: bool = False,
     ) -> Sequence[Task]:
         """Retrieve information about tasks based on the specified filters.
 
@@ -720,6 +722,8 @@ def get_tasks(  # pylint: disable=too-many-arguments
             Whether sorting should be in ascending order.
         limit : Optional[int] (default: None)
             Maximum number of tasks to return. If `None`, no limit is applied.
+        claimable : bool (default: False)
+            If `True`, exclude AgentApp tasks whose run series is already active.
 
         Returns
         -------

framework/py/flwr/supercore/corestate/in_memory_corestate.py

@@ -50,7 +50,11 @@
     TaskUsage,
 )
 from flwr.supercore import log
-from flwr.supercore.constant import OBJECT_PUSH_SESSION_TTL_SECONDS, AutomationStatus
+from flwr.supercore.constant import (
+    OBJECT_PUSH_SESSION_TTL_SECONDS,
+    AutomationStatus,
+    TaskType,
+)
 from flwr.supercore.date import now
 from flwr.supercore.fab import Fab
 from flwr.supercore.typing import ConnectorOAuthSessionRecord, ConnectorRecord
@@ -135,6 +139,8 @@ def __init__(self, object_store: ObjectStore) -> None:
         self.lock_nonce_store = Lock()
         self.run_series_store: dict[int, RunSeries] = {}
         self.agent_run_series_ids: set[int] = set()
+        self.run_id_to_series_id: dict[int, int] = {}
+        self.active_agent_tasks: dict[int, int] = {}
         self.lock_run_series_store = Lock()
         self.run_series_context_store: dict[int, Context] = {}
         self.lock_run_series_context_store = Lock()
@@ -657,6 +663,7 @@ def store_run_in_series(  # pylint: disable=too-many-arguments,too-many-position
             if run_id in run_series.run_ids:
                 return None
             run_series.run_ids.append(run_id)
+            self.run_id_to_series_id[run_id] = resolved_series_id
             if series_id is not None:
                 run_series.updated_at = now().isoformat()
             return resolved_series_id
@@ -938,6 +945,7 @@ def get_tasks(  # pylint: disable=too-many-arguments
         order_by: Literal["pending_at"] | None = None,
         ascending: bool = True,
         limit: int | None = None,
+        claimable: bool = False,
     ) -> Sequence[Task]:
         """Retrieve information about tasks based on the specified filters."""
         if order_by not in (None, "pending_at"):
@@ -980,6 +988,13 @@ def get_tasks(  # pylint: disable=too-many-arguments
                     if self.task_store[task_id].status.status in status_set
                 }
 
+            if claimable:
+                matched_task_ids &= {
+                    task_id
+                    for task_id in matched_task_ids
+                    if self._is_task_claimable_locked(self.task_store[task_id])
+                }
+
             tasks = [self.task_store[task_id] for task_id in matched_task_ids]
 
             if order_by is not None:
@@ -999,6 +1014,13 @@ def get_tasks(  # pylint: disable=too-many-arguments
                 result.append(task_copy)
             return result
 
+    def _is_task_claimable_locked(self, task: Task) -> bool:
+        """Return whether a task can acquire its run series under the task lock."""
+        if task.type != TaskType.AGENT_APP:
+            return True
+        series_id = self.run_id_to_series_id.get(task.run_id)
+        return series_id is None or series_id not in self.active_agent_tasks
+
     def add_task_usage(self, task_id: int, usage: TaskUsage) -> None:
         """Record usage for the specified task."""
         with self.lock_task_store:
@@ -1045,11 +1067,14 @@ def claim_task(self, task_id: int) -> str | None:
         """Atomically claim a pending task."""
         token = secrets.token_hex(FLWR_TASK_TOKEN_LENGTH)
         with self.lock_task_store:
+            self._cleanup_expired_task_tokens_locked()
             task = self.task_store.get(task_id)
             if task is None or task_id in self.task_token_store:
                 return None
             if task.status.status != Status.PENDING:
                 return None
+            if not self._reserve_agent_run_series_locked(task):
+                return None
 
             # Claiming moves the task into STARTING and records the heartbeat state.
             claimed_at = now()
@@ -1118,6 +1143,7 @@ def finish_task(self, task_id: int, sub_status: str, details: str) -> bool:
             # Revoke any existing task token now that the task is finished.
             if (record := self.task_token_store.pop(task_id, None)) is not None:
                 self.task_token_to_task_id.pop(record.token, None)
+            self._release_agent_run_series_locked(task)
             return True
 
     def acknowledge_task_heartbeat(self, task_id: int) -> bool:
@@ -1311,6 +1337,7 @@ def _cleanup_expired_task_tokens_locked(self) -> None:
                     task.status.CopyFrom(
                         TaskStatus(status=Status.PENDING, sub_status="", details="")
                     )
+                    self._release_agent_run_series_locked(task)
                 elif task and task.status.status == Status.RUNNING:
                     task.finished_at = record.active_until.isoformat()
                     task.status.CopyFrom(
@@ -1323,12 +1350,36 @@ def _cleanup_expired_task_tokens_locked(self) -> None:
                     expired_task = Task()
                     expired_task.CopyFrom(task)
                     expired_tasks.append(expired_task)
+                    self._release_agent_run_series_locked(task)
                 del self.task_token_store[task_id]
                 self.task_token_to_task_id.pop(record.token, None)
 
         if expired_tasks:
             self._on_task_tokens_expired(expired_tasks)
 
+    def _reserve_agent_run_series_locked(self, task: Task) -> bool:
+        """Reserve an AgentApp run series for a task while holding the task lock."""
+        if task.type != TaskType.AGENT_APP:
+            return True
+        series_id = self.run_id_to_series_id.get(task.run_id)
+        if series_id is None:
+            return True
+        if series_id in self.active_agent_tasks:
+            return False
+        self.active_agent_tasks[series_id] = task.task_id
+        return True
+
+    def _release_agent_run_series_locked(self, task: Task) -> None:
+        """Release an AgentApp run series while holding the task lock."""
+        if task.type != TaskType.AGENT_APP:
+            return
+        series_id = self.run_id_to_series_id.get(task.run_id)
+        if (
+            series_id is not None
+            and self.active_agent_tasks.get(series_id) == task.task_id
+        ):
+            del self.active_agent_tasks[series_id]
+
     def _cleanup_invalid_task_messages_locked(self, current: float) -> None:
         """Remove expired Messages and Messages for invalid destination tasks."""
         for message_id, message in list(self.task_message_store.items()):