diff --git a/framework/py/flwr/app/message/context.py b/framework/py/flwr/app/message/context.py index 4102e79aba3b..d622a0a02647 100644 --- a/framework/py/flwr/app/message/context.py +++ b/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 diff --git a/framework/py/flwr/common/serde.py b/framework/py/flwr/common/serde.py index baae88fccf73..893286c6f0e7 100644 --- a/framework/py/flwr/common/serde.py +++ b/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: diff --git a/framework/py/flwr/server/superlink/linkstate/linkstate_test.py b/framework/py/flwr/server/superlink/linkstate/linkstate_test.py index 1e6f07971c21..f7c2e97e105b 100644 --- a/framework/py/flwr/server/superlink/linkstate/linkstate_test.py +++ b/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), diff --git a/framework/py/flwr/supercore/corestate/corestate.py b/framework/py/flwr/supercore/corestate/corestate.py index e4b2710baec3..d974a8eed82e 100644 --- a/framework/py/flwr/supercore/corestate/corestate.py +++ b/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 ------- diff --git a/framework/py/flwr/supercore/corestate/in_memory_corestate.py b/framework/py/flwr/supercore/corestate/in_memory_corestate.py index 157cec0e0cba..d5cf475b3e8b 100644 --- a/framework/py/flwr/supercore/corestate/in_memory_corestate.py +++ b/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()): diff --git a/framework/py/flwr/supercore/corestate/sql_corestate.py b/framework/py/flwr/supercore/corestate/sql_corestate.py index 3cf0356e8ece..1adde75f50b1 100644 --- a/framework/py/flwr/supercore/corestate/sql_corestate.py +++ b/framework/py/flwr/supercore/corestate/sql_corestate.py @@ -40,6 +40,7 @@ from sqlalchemy.dialects.sqlite import Insert as SQLiteInsert from sqlalchemy.dialects.sqlite import insert as sqlite_insert from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm import Session from flwr.app import Context, Message from flwr.app.message import make_message @@ -74,7 +75,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.sql_mixin import SqlMixin @@ -135,6 +140,10 @@ } +class _AgentRunSeriesBusyError(Exception): + """Signal that an AgentApp run series is already reserved.""" + + class SqlCoreState(CoreState, SqlMixin): # pylint: disable=R0904 """SQLAlchemy-based CoreState implementation.""" @@ -1231,6 +1240,7 @@ def get_tasks( # pylint: disable=too-many-arguments,too-many-locals,too-many-br 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"): @@ -1267,6 +1277,22 @@ def get_tasks( # pylint: disable=too-many-arguments,too-many-locals,too-many-br return [] query = query.where(or_(*status_conditions)) + if claimable: + active_series_exists = exists( + select(SeriesRunsModel.id) + .join( + RunSeriesModel, + RunSeriesModel.series_id == SeriesRunsModel.series_id, + ) + .where( + SeriesRunsModel.run_id == TaskModel.run_id, + RunSeriesModel.active_agent_task_id.is_not(None), + ) + ) + query = query.where( + or_(TaskModel.type != TaskType.AGENT_APP, ~active_series_exists) + ) + if order_by is not None: order_column = ( TaskModel.pending_at.asc() if ascending else TaskModel.pending_at.desc() @@ -1347,11 +1373,12 @@ def claim_task(self, task_id: int) -> str | None: claimed_at = now() active_until = claimed_at + timedelta(seconds=HEARTBEAT_DEFAULT_INTERVAL) sint64_task_id = uint64_to_int64(task_id) + self._cleanup_expired_task_tokens() try: # The conditional UPDATE is the atomic claim: exactly one caller can # move a pending, unclaimed task to STARTING and attach a token. with self.session() as session: - claimed_task_id = session.scalar( + claimed_task = session.execute( update(TaskModel) .where( TaskModel.task_id == sint64_task_id, @@ -1363,14 +1390,33 @@ def claim_task(self, task_id: int) -> str | None: active_until=active_until, starting_at=claimed_at, ) - .returning(TaskModel.task_id) - ) - if claimed_task_id is None: - return None + .returning(TaskModel.task_id, TaskModel.type, TaskModel.run_id) + ).one_or_none() + if claimed_task is None: + return None + + if claimed_task.type == TaskType.AGENT_APP: + series_id = session.scalar( + select(SeriesRunsModel.series_id).where( + SeriesRunsModel.run_id == claimed_task.run_id + ) + ) + if series_id is not None: + reserved_series_id = session.scalar( + update(RunSeriesModel) + .where( + RunSeriesModel.series_id == series_id, + RunSeriesModel.active_agent_task_id.is_(None), + ) + .values(active_agent_task_id=sint64_task_id) + .returning(RunSeriesModel.series_id) + ) + if reserved_series_id is None: + raise _AgentRunSeriesBusyError return token - except IntegrityError: - # Rare failure: generated token already exists (duplicate) + except (IntegrityError, _AgentRunSeriesBusyError): + # The token collided or another AgentApp owns the run series. return None def activate_task(self, task_id: int) -> bool: @@ -1414,16 +1460,24 @@ def finish_task(self, task_id: int, sub_status: str, details: str) -> bool: if sub_status == SubStatus.COMPLETED: query = query.where(TaskModel.running_at.is_not(None)) - finished_task_id = session.scalar( + finished_task = session.execute( query.values( finished_at=now(), sub_status=sub_status, details=details, active_until=None, token=None, - ).returning(TaskModel.task_id) - ) - return finished_task_id is not None + ).returning(TaskModel.task_id, TaskModel.type, TaskModel.run_id) + ).one_or_none() + if finished_task is None: + return False + if finished_task.type == TaskType.AGENT_APP: + self._release_agent_run_series( + session, + task_id=finished_task.task_id, + run_id=finished_task.run_id, + ) + return True def acknowledge_task_heartbeat(self, task_id: int) -> bool: """Extend heartbeat state for the claimed task.""" @@ -1682,21 +1736,25 @@ def _cleanup_expired_task_tokens(self) -> None: expired_at = now() with self.session() as session: # Claims that never reached RUNNING are retryable launch failures. - session.execute( - update(TaskModel) - .where( - TaskModel.token.is_not(None), - TaskModel.active_until < expired_at, - _task_status_filter(Status.STARTING), - ) - .values( - token=None, - active_until=None, - starting_at=None, - sub_status="", - details="", - ) - ) + expired_starting_tasks = [ + task_from_model(row) + for row in session.scalars( + update(TaskModel) + .where( + TaskModel.token.is_not(None), + TaskModel.active_until < expired_at, + _task_status_filter(Status.STARTING), + ) + .values( + token=None, + active_until=None, + starting_at=None, + sub_status="", + details="", + ) + .returning(TaskModel) + ).all() + ] # Expired running task claims are terminal failures and lose their token. expired_tasks = [ @@ -1718,9 +1776,36 @@ def _cleanup_expired_task_tokens(self) -> None: .returning(TaskModel) ).all() ] + for task in [*expired_starting_tasks, *expired_tasks]: + if task.type == TaskType.AGENT_APP: + self._release_agent_run_series( + session, + task_id=uint64_to_int64(task.task_id), + run_id=uint64_to_int64(task.run_id), + ) if expired_tasks: self._on_task_tokens_expired(expired_tasks) + @staticmethod + def _release_agent_run_series( + session: Session, + *, + task_id: int, + run_id: int, + ) -> None: + """Release an AgentApp run-series reservation in this transaction.""" + series_id = select(SeriesRunsModel.series_id).where( + SeriesRunsModel.run_id == run_id + ) + session.execute( + update(RunSeriesModel) + .where( + RunSeriesModel.series_id == series_id.scalar_subquery(), + RunSeriesModel.active_agent_task_id == task_id, + ) + .values(active_agent_task_id=None) + ) + def _cleanup_invalid_task_messages(self) -> None: """Remove expired task Messages.""" with self.session() as session: diff --git a/framework/py/flwr/supercore/servicer/runtime/runtime_handlers.py b/framework/py/flwr/supercore/servicer/runtime/runtime_handlers.py index 03f557a04dce..8bcec0c50047 100644 --- a/framework/py/flwr/supercore/servicer/runtime/runtime_handlers.py +++ b/framework/py/flwr/supercore/servicer/runtime/runtime_handlers.py @@ -63,7 +63,10 @@ def pull_pending_tasks( log(DEBUG, "Runtime.PullPendingTasks") tasks = state.get_tasks( - statuses=[Status.PENDING], order_by="pending_at", ascending=True + statuses=[Status.PENDING], + order_by="pending_at", + ascending=True, + claimable=True, ) return PullPendingTasksResponse(tasks=tasks) diff --git a/framework/py/flwr/supercore/servicer/runtime/runtime_handlers_test.py b/framework/py/flwr/supercore/servicer/runtime/runtime_handlers_test.py index f6d681f375b8..f7b674bd13de 100644 --- a/framework/py/flwr/supercore/servicer/runtime/runtime_handlers_test.py +++ b/framework/py/flwr/supercore/servicer/runtime/runtime_handlers_test.py @@ -87,7 +87,10 @@ def test_pull_pending_tasks_returns_pending_tasks(self) -> None: # Assert self.state.get_tasks.assert_called_once_with( - statuses=[Status.PENDING], order_by="pending_at", ascending=True + statuses=[Status.PENDING], + order_by="pending_at", + ascending=True, + claimable=True, ) self.assertEqual(len(response.tasks), 1) self.assertEqual(response.tasks[0].task_id, 123) diff --git a/framework/py/flwr/supercore/state/alembic/versions/rev_2026_08_22_serialize_agentapp_runs_per_series.py b/framework/py/flwr/supercore/state/alembic/versions/rev_2026_08_22_serialize_agentapp_runs_per_series.py new file mode 100644 index 000000000000..968240d80438 --- /dev/null +++ b/framework/py/flwr/supercore/state/alembic/versions/rev_2026_08_22_serialize_agentapp_runs_per_series.py @@ -0,0 +1,52 @@ +# Copyright 2026 Flower Labs GmbH. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Serialize AgentApp runs per series. + +Revision ID: 33157d2e33ec +Revises: 03f4cfe3ff15 +Create Date: 2026-08-22 17:10:32.343816 +""" +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +# pylint: disable=no-member + +# revision identifiers, used by Alembic. +revision: str = "33157d2e33ec" +down_revision: str | Sequence[str] | None = "03f4cfe3ff15" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + """Upgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("run_series", schema=None) as batch_op: + batch_op.add_column( + sa.Column("active_agent_task_id", sa.BigInteger(), nullable=True) + ) + + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("run_series", schema=None) as batch_op: + batch_op.drop_column("active_agent_task_id") + + # ### end Alembic commands ### diff --git a/framework/py/flwr/supercore/state/schema/README.md b/framework/py/flwr/supercore/state/schema/README.md index 64447ec2ba75..9eca6e0c00db 100644 --- a/framework/py/flwr/supercore/state/schema/README.md +++ b/framework/py/flwr/supercore/state/schema/README.md @@ -172,6 +172,7 @@ erDiagram run_series { BIGINT series_id PK + BIGINT active_agent_task_id "nullable" TIMESTAMP created_at VARCHAR description "nullable" VARCHAR federation_id diff --git a/framework/py/flwr/supercore/state/schema/corestate_models.py b/framework/py/flwr/supercore/state/schema/corestate_models.py index ca0e054dc30f..94ee6b72bed9 100644 --- a/framework/py/flwr/supercore/state/schema/corestate_models.py +++ b/framework/py/flwr/supercore/state/schema/corestate_models.py @@ -74,6 +74,7 @@ class RunSeries(FlwrBase): description: Mapped[str | None] = mapped_column(String, nullable=True) created_at: Mapped[datetime] = mapped_column(UTCDateTime(), nullable=False) updated_at: Mapped[datetime] = mapped_column(UTCDateTime(), nullable=False) + active_agent_task_id: Mapped[int | None] = mapped_column(BigInteger, nullable=True) class SeriesContext(FlwrBase): diff --git a/framework/py/flwr/supercore/task_process/agent/context_items.py b/framework/py/flwr/supercore/task_process/agent/context_items.py index 2ecdb3284492..b5d3312dee02 100644 --- a/framework/py/flwr/supercore/task_process/agent/context_items.py +++ b/framework/py/flwr/supercore/task_process/agent/context_items.py @@ -29,9 +29,10 @@ def append_items(context: Context, new_items: list[JSONObject]) -> None: """Append OpenResponses items to ``context.state``.""" - # Initialize the items storage if it doesn't exist yet - record = context.state.setdefault(ITEMS_KEY, ConfigRecord({JSON_KEY: []})) - items = cast(list[str], record[JSON_KEY]) + with context.locked(): + # Initialize the items storage if it doesn't exist yet + record = context.state.setdefault(ITEMS_KEY, ConfigRecord({JSON_KEY: []})) + items = cast(list[str], record[JSON_KEY]) - # Add the new items to the list - items.extend(strict_json_dumps(item, compact=True) for item in new_items) + # Add the new items to the list + items.extend(strict_json_dumps(item, compact=True) for item in new_items) diff --git a/framework/py/flwr/superlink/servicer/runtime/runtime_handlers.py b/framework/py/flwr/superlink/servicer/runtime/runtime_handlers.py index 9a0839120f5c..4c57e07ae28c 100644 --- a/framework/py/flwr/superlink/servicer/runtime/runtime_handlers.py +++ b/framework/py/flwr/superlink/servicer/runtime/runtime_handlers.py @@ -89,7 +89,10 @@ def pull_pending_tasks( log(DEBUG, "Runtime.PullPendingTasks") process_due_automations(state, limit=AUTOMATION_BATCH_LIMIT) tasks = state.get_tasks( - statuses=[Status.PENDING], order_by="pending_at", ascending=True + statuses=[Status.PENDING], + order_by="pending_at", + ascending=True, + claimable=True, ) return PullPendingTasksResponse(tasks=tasks) @@ -256,6 +259,8 @@ def pull_task_input( if run and run.series_id: series_context = state.get_run_series_context(run.series_id) if run and fab and series_context and state.activate_task(task.task_id): + series_context.run_id = run_id + series_context.series_id = run.series_id log(INFO, "Started task %d of run %d", task.task_id, run_id) return PullTaskInputResponse( context=context_to_proto(series_context), @@ -283,18 +288,31 @@ def push_task_output( if request.HasField("clientapp_runtime"): state.add_clientapp_runtime(run_id, request.clientapp_runtime) + series_id = None + output_context = None + if request.HasField("context"): + runs = state.get_run_info(run_ids=[run_id]) + run = runs[0] if runs else None + if run and run.series_id and run.primary_task_id == task.task_id: + series_id = run.series_id + output_context = context_from_proto(request.context) + + # Finishing releases an AgentApp's run-series reservation, so persist its + # context first. Other task types retain the existing finish-then-store order. + if series_id is not None and output_context is not None: + if task.type == TaskType.AGENT_APP: + state.set_run_series_context(series_id, output_context) + if state.finish_task( task.task_id, sub_status=request.sub_status, details=request.details ): log(INFO, "Finished task %d of run %d", task.task_id, run_id) - if request.HasField("context"): - runs = state.get_run_info(run_ids=[run_id]) - run = runs[0] if runs else None - if run and run.series_id and run.primary_task_id == task.task_id: - state.set_run_series_context( - run.series_id, - context_from_proto(request.context), - ) + if ( + series_id is not None + and output_context is not None + and task.type != TaskType.AGENT_APP + ): + state.set_run_series_context(series_id, output_context) else: log(ERROR, "Failed to finish task %d of run %s", task.task_id, run_id) return PushTaskOutputResponse() diff --git a/framework/py/flwr/superlink/servicer/runtime/runtime_handlers_test.py b/framework/py/flwr/superlink/servicer/runtime/runtime_handlers_test.py index 1069dca38347..35c68a9bd6c5 100644 --- a/framework/py/flwr/superlink/servicer/runtime/runtime_handlers_test.py +++ b/framework/py/flwr/superlink/servicer/runtime/runtime_handlers_test.py @@ -921,7 +921,7 @@ def test_run_status_transitions(self) -> None: # Assert: Response is successful and run status is now RUNNING assert isinstance(response, PullTaskInputResponse) - assert response.context.run_id == 123 + assert response.context.run_id == run_id assert response.context.series_id == run.series_id run_status = self.state.get_run_status({run_id})[run_id] assert run_status.status == Status.RUNNING