flwr-nightly 1.8.0.dev20240310__tar.gz → 1.8.0.dev20240312__tar.gz

{flwr_nightly-1.8.0.dev20240310 → flwr_nightly-1.8.0.dev20240312}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flwr-nightly
-Version: 1.8.0.dev20240310
+Version: 1.8.0.dev20240312
 Summary: Flower: A Friendly Federated Learning Framework
 Home-page: https://flower.ai
 License: Apache-2.0

{flwr_nightly-1.8.0.dev20240310 → flwr_nightly-1.8.0.dev20240312}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "flwr-nightly"
-version = "1.8.0-dev20240310"
+version = "1.8.0-dev20240312"
 description = "Flower: A Friendly Federated Learning Framework"
 license = "Apache-2.0"
 authors = ["The Flower Authors <hello@flower.ai>"]

{flwr_nightly-1.8.0.dev20240310 → flwr_nightly-1.8.0.dev20240312}/src/py/flwr/cli/new/new.py

@@ -22,7 +22,7 @@ from typing import Dict, Optional
 import typer
 from typing_extensions import Annotated
 
-from ..utils import prompt_options
+from ..utils import prompt_options, prompt_text
 
 
 class MlFramework(str, Enum):
@@ -72,9 +72,9 @@ def render_and_create(file_path: str, template: str, context: Dict[str, str]) ->
 
 def new(
     project_name: Annotated[
-        str,
+        Optional[str],
         typer.Argument(metavar="project_name", help="The name of the project"),
-    ],
+    ] = None,
     framework: Annotated[
         Optional[MlFramework],
         typer.Option(case_sensitive=False, help="The ML framework to use"),
@@ -83,6 +83,9 @@ def new(
     """Create new Flower project."""
     print(f"Creating Flower project {project_name}...")
 
+    if project_name is None:
+        project_name = prompt_text("Please provide project name")
+
     if framework is not None:
         framework_str = str(framework.value)
     else:

{flwr_nightly-1.8.0.dev20240310 → flwr_nightly-1.8.0.dev20240312}/src/py/flwr/cli/utils.py

@@ -14,11 +14,24 @@
 # ==============================================================================
 """Flower command line interface utils."""
 
-from typing import List
+from typing import List, cast
 
 import typer
 
 
+def prompt_text(text: str) -> str:
+    """Ask user to enter text input."""
+    while True:
+        result = typer.prompt(
+            typer.style(f"\n💬 {text}", fg=typer.colors.MAGENTA, bold=True)
+        )
+        if len(result) > 0:
+            break
+        print(typer.style("❌ Invalid entry", fg=typer.colors.RED, bold=True))
+
+    return cast(str, result)
+
+
 def prompt_options(text: str, options: List[str]) -> str:
     """Ask user to select one of the given options and return the selected item."""
     # Turn options into a list with index as in " [ 0] quickstart-pytorch"

{flwr_nightly-1.8.0.dev20240310 → flwr_nightly-1.8.0.dev20240312}/src/py/flwr/client/app.py

@@ -456,7 +456,19 @@ def _start_client_internal(
                     time.sleep(3)  # Wait for 3s before asking again
                     continue
 
-                log(INFO, "Received message")
+                log(INFO, "")
+                log(
+                    INFO,
+                    "[RUN %s, ROUND %s]",
+                    message.metadata.run_id,
+                    message.metadata.group_id,
+                )
+                log(
+                    INFO,
+                    "Received: %s message %s",
+                    message.metadata.message_type,
+                    message.metadata.message_id,
+                )
 
                 # Handle control message
                 out_message, sleep_duration = handle_control_message(message)
@@ -484,7 +496,18 @@ def _start_client_internal(
 
                 # Send
                 send(out_message)
-                log(INFO, "Sent reply")
+                log(
+                    INFO,
+                    "[RUN %s, ROUND %s]",
+                    out_message.metadata.run_id,
+                    out_message.metadata.group_id,
+                )
+                log(
+                    INFO,
+                    "Sent: %s reply to message %s",
+                    out_message.metadata.message_type,
+                    message.metadata.message_id,
+                )
 
             # Unregister node
             if delete_node is not None:

{flwr_nightly-1.8.0.dev20240310 → flwr_nightly-1.8.0.dev20240312}/src/py/flwr/client/mod/__init__.py

@@ -17,7 +17,7 @@
 
 from .centraldp_mods import adaptiveclipping_mod, fixedclipping_mod
 from .localdp_mod import LocalDpMod
-from .secure_aggregation.secaggplus_mod import secaggplus_mod
+from .secure_aggregation import secagg_mod, secaggplus_mod
 from .utils import make_ffn
 
 __all__ = [
@@ -25,5 +25,6 @@ __all__ = [
     "fixedclipping_mod",
     "LocalDpMod",
     "make_ffn",
+    "secagg_mod",
     "secaggplus_mod",
 ]

{flwr_nightly-1.8.0.dev20240310 → flwr_nightly-1.8.0.dev20240312}/src/py/flwr/client/mod/secure_aggregation/__init__.py

@@ -15,8 +15,10 @@
 """Secure Aggregation mods."""
 
 
+from .secagg_mod import secagg_mod
 from .secaggplus_mod import secaggplus_mod
 
 __all__ = [
+    "secagg_mod",
     "secaggplus_mod",
 ]

flwr_nightly-1.8.0.dev20240312/src/py/flwr/client/mod/secure_aggregation/secagg_mod.py

@@ -0,0 +1,30 @@
+# Copyright 2024 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.
+# ==============================================================================
+"""Modifier for the SecAgg protocol."""
+
+
+from flwr.client.typing import ClientAppCallable
+from flwr.common import Context, Message
+
+from .secaggplus_mod import secaggplus_mod
+
+
+def secagg_mod(
+    msg: Message,
+    ctxt: Context,
+    call_next: ClientAppCallable,
+) -> Message:
+    """Handle incoming message and return results, following the SecAgg protocol."""
+    return secaggplus_mod(msg, ctxt, call_next)

{flwr_nightly-1.8.0.dev20240310 → flwr_nightly-1.8.0.dev20240312}/src/py/flwr/client/mod/secure_aggregation/secaggplus_mod.py

@@ -12,19 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Message handler for the SecAgg+ protocol."""
+"""Modifier for the SecAgg+ protocol."""
 
 
 import os
 from dataclasses import dataclass, field
-from logging import INFO, WARNING
-from typing import Any, Callable, Dict, List, Tuple, cast
+from logging import DEBUG, WARNING
+from typing import Any, Dict, List, Tuple, cast
 
 from flwr.client.typing import ClientAppCallable
 from flwr.common import (
     ConfigsRecord,
     Context,
     Message,
+    Parameters,
     RecordSet,
     ndarray_to_bytes,
     parameters_to_ndarrays,
@@ -62,7 +63,7 @@ from flwr.common.secure_aggregation.secaggplus_utils import (
     share_keys_plaintext_concat,
     share_keys_plaintext_separate,
 )
-from flwr.common.typing import ConfigsRecordValues, FitRes
+from flwr.common.typing import ConfigsRecordValues
 
 
 @dataclass
@@ -132,18 +133,6 @@ class SecAggPlusState:
         return ret
 
 
-def _get_fit_fn(
-    msg: Message, ctxt: Context, call_next: ClientAppCallable
-) -> Callable[[], FitRes]:
-    """Get the fit function."""
-
-    def fit() -> FitRes:
-        out_msg = call_next(msg, ctxt)
-        return compat.recordset_to_fitres(out_msg.content, keep_input=False)
-
-    return fit
-
-
 def secaggplus_mod(
     msg: Message,
     ctxt: Context,
@@ -173,25 +162,32 @@ def secaggplus_mod(
     check_configs(state.current_stage, configs)
 
     # Execute
+    out_content = RecordSet()
     if state.current_stage == Stage.SETUP:
         state.nid = msg.metadata.dst_node_id
         res = _setup(state, configs)
     elif state.current_stage == Stage.SHARE_KEYS:
         res = _share_keys(state, configs)
-    elif state.current_stage == Stage.COLLECT_MASKED_INPUT:
-        fit = _get_fit_fn(msg, ctxt, call_next)
-        res = _collect_masked_input(state, configs, fit)
+    elif state.current_stage == Stage.COLLECT_MASKED_VECTORS:
+        out_msg = call_next(msg, ctxt)
+        out_content = out_msg.content
+        fitres = compat.recordset_to_fitres(out_content, keep_input=True)
+        res = _collect_masked_vectors(
+            state, configs, fitres.num_examples, fitres.parameters
+        )
+        for p_record in out_content.parameters_records.values():
+            p_record.clear()
     elif state.current_stage == Stage.UNMASK:
         res = _unmask(state, configs)
     else:
-        raise ValueError(f"Unknown secagg stage: {state.current_stage}")
+        raise ValueError(f"Unknown SecAgg/SecAgg+ stage: {state.current_stage}")
 
     # Save state
     ctxt.state.configs_records[RECORD_KEY_STATE] = ConfigsRecord(state.to_dict())
 
     # Return message
-    content = RecordSet(configs_records={RECORD_KEY_CONFIGS: ConfigsRecord(res, False)})
-    return msg.create_reply(content, ttl="")
+    out_content.configs_records[RECORD_KEY_CONFIGS] = ConfigsRecord(res, False)
+    return msg.create_reply(out_content, ttl="")
 
 
 def check_stage(current_stage: str, configs: ConfigsRecord) -> None:
@@ -199,7 +195,7 @@ def check_stage(current_stage: str, configs: ConfigsRecord) -> None:
     # Check the existence of Config.STAGE
     if Key.STAGE not in configs:
         raise KeyError(
-            f"The required key '{Key.STAGE}' is missing from the input `named_values`."
+            f"The required key '{Key.STAGE}' is missing from the ConfigsRecord."
         )
 
     # Check the value type of the Config.STAGE
@@ -215,7 +211,7 @@ def check_stage(current_stage: str, configs: ConfigsRecord) -> None:
         if current_stage != Stage.UNMASK:
             log(WARNING, "Restart from the setup stage")
     # If stage is not "setup",
-    # the stage from `named_values` should be the expected next stage
+    # the stage from configs should be the expected next stage
     else:
         stages = Stage.all()
         expected_next_stage = stages[(stages.index(current_stage) + 1) % len(stages)]
@@ -229,7 +225,7 @@ def check_stage(current_stage: str, configs: ConfigsRecord) -> None:
 # pylint: disable-next=too-many-branches
 def check_configs(stage: str, configs: ConfigsRecord) -> None:
     """Check the validity of the configs."""
-    # Check `named_values` for the setup stage
+    # Check configs for the setup stage
     if stage == Stage.SETUP:
         key_type_pairs = [
             (Key.SAMPLE_NUMBER, int),
@@ -243,7 +239,7 @@ def check_configs(stage: str, configs: ConfigsRecord) -> None:
             if key not in configs:
                 raise KeyError(
                     f"Stage {Stage.SETUP}: the required key '{key}' is "
-                    "missing from the input `named_values`."
+                    "missing from the ConfigsRecord."
                 )
             # Bool is a subclass of int in Python,
             # so `isinstance(v, int)` will return True even if v is a boolean.
@@ -266,7 +262,7 @@ def check_configs(stage: str, configs: ConfigsRecord) -> None:
                     f"Stage {Stage.SHARE_KEYS}: "
                     f"the value for the key '{key}' must be a list of two bytes."
                 )
-    elif stage == Stage.COLLECT_MASKED_INPUT:
+    elif stage == Stage.COLLECT_MASKED_VECTORS:
         key_type_pairs = [
             (Key.CIPHERTEXT_LIST, bytes),
             (Key.SOURCE_LIST, int),
@@ -274,9 +270,9 @@ def check_configs(stage: str, configs: ConfigsRecord) -> None:
         for key, expected_type in key_type_pairs:
             if key not in configs:
                 raise KeyError(
-                    f"Stage {Stage.COLLECT_MASKED_INPUT}: "
+                    f"Stage {Stage.COLLECT_MASKED_VECTORS}: "
                     f"the required key '{key}' is "
-                    "missing from the input `named_values`."
+                    "missing from the ConfigsRecord."
                 )
             if not isinstance(configs[key], list) or any(
                 elm
@@ -285,7 +281,7 @@ def check_configs(stage: str, configs: ConfigsRecord) -> None:
                 if type(elm) is not expected_type
             ):
                 raise TypeError(
-                    f"Stage {Stage.COLLECT_MASKED_INPUT}: "
+                    f"Stage {Stage.COLLECT_MASKED_VECTORS}: "
                     f"the value for the key '{key}' "
                     f"must be of type List[{expected_type.__name__}]"
                 )
@@ -299,7 +295,7 @@ def check_configs(stage: str, configs: ConfigsRecord) -> None:
                 raise KeyError(
                     f"Stage {Stage.UNMASK}: "
                     f"the required key '{key}' is "
-                    "missing from the input `named_values`."
+                    "missing from the ConfigsRecord."
                 )
             if not isinstance(configs[key], list) or any(
                 elm
@@ -322,7 +318,7 @@ def _setup(
     # Assigning parameter values to object fields
     sec_agg_param_dict = configs
     state.sample_num = cast(int, sec_agg_param_dict[Key.SAMPLE_NUMBER])
-    log(INFO, "Node %d: starting stage 0...", state.nid)
+    log(DEBUG, "Node %d: starting stage 0...", state.nid)
 
     state.share_num = cast(int, sec_agg_param_dict[Key.SHARE_NUMBER])
     state.threshold = cast(int, sec_agg_param_dict[Key.THRESHOLD])
@@ -347,7 +343,7 @@ def _setup(
 
     state.sk1, state.pk1 = private_key_to_bytes(sk1), public_key_to_bytes(pk1)
     state.sk2, state.pk2 = private_key_to_bytes(sk2), public_key_to_bytes(pk2)
-    log(INFO, "Node %d: stage 0 completes. uploading public keys...", state.nid)
+    log(DEBUG, "Node %d: stage 0 completes. uploading public keys...", state.nid)
     return {Key.PUBLIC_KEY_1: state.pk1, Key.PUBLIC_KEY_2: state.pk2}
 
 
@@ -357,7 +353,7 @@ def _share_keys(
 ) -> Dict[str, ConfigsRecordValues]:
     named_bytes_tuples = cast(Dict[str, Tuple[bytes, bytes]], configs)
     key_dict = {int(sid): (pk1, pk2) for sid, (pk1, pk2) in named_bytes_tuples.items()}
-    log(INFO, "Node %d: starting stage 1...", state.nid)
+    log(DEBUG, "Node %d: starting stage 1...", state.nid)
     state.public_keys_dict = key_dict
 
     # Check if the size is larger than threshold
@@ -409,17 +405,18 @@ def _share_keys(
             dsts.append(nid)
             ciphertexts.append(ciphertext)
 
-    log(INFO, "Node %d: stage 1 completes. uploading key shares...", state.nid)
+    log(DEBUG, "Node %d: stage 1 completes. uploading key shares...", state.nid)
     return {Key.DESTINATION_LIST: dsts, Key.CIPHERTEXT_LIST: ciphertexts}
 
 
 # pylint: disable-next=too-many-locals
-def _collect_masked_input(
+def _collect_masked_vectors(
     state: SecAggPlusState,
     configs: ConfigsRecord,
-    fit: Callable[[], FitRes],
+    num_examples: int,
+    updated_parameters: Parameters,
 ) -> Dict[str, ConfigsRecordValues]:
-    log(INFO, "Node %d: starting stage 2...", state.nid)
+    log(DEBUG, "Node %d: starting stage 2...", state.nid)
     available_clients: List[int] = []
     ciphertexts = cast(List[bytes], configs[Key.CIPHERTEXT_LIST])
     srcs = cast(List[int], configs[Key.SOURCE_LIST])
@@ -447,26 +444,20 @@ def _collect_masked_input(
         state.rd_seed_share_dict[src] = rd_seed_share
         state.sk1_share_dict[src] = sk1_share
 
-    # Fit client
-    fit_res = fit()
-    if len(fit_res.metrics) > 0:
-        log(
-            WARNING,
-            "The metrics in FitRes will not be preserved or sent to the server.",
-        )
-    ratio = fit_res.num_examples / state.max_weight
+    # Fit
+    ratio = num_examples / state.max_weight
     if ratio > 1:
         log(
             WARNING,
             "Potential overflow warning: the provided weight (%s) exceeds the specified"
             " max_weight (%s). This may lead to overflow issues.",
-            fit_res.num_examples,
+            num_examples,
             state.max_weight,
         )
     q_ratio = round(ratio * state.target_range)
     dq_ratio = q_ratio / state.target_range
 
-    parameters = parameters_to_ndarrays(fit_res.parameters)
+    parameters = parameters_to_ndarrays(updated_parameters)
     parameters = parameters_multiply(parameters, dq_ratio)
 
     # Quantize parameter update (vector)
@@ -500,7 +491,7 @@ def _collect_masked_input(
 
     # Take mod of final weight update vector and return to server
     quantized_parameters = parameters_mod(quantized_parameters, state.mod_range)
-    log(INFO, "Node %d: stage 2 completed, uploading masked parameters...", state.nid)
+    log(DEBUG, "Node %d: stage 2 completed, uploading masked parameters...", state.nid)
     return {
         Key.MASKED_PARAMETERS: [ndarray_to_bytes(arr) for arr in quantized_parameters]
     }
@@ -509,7 +500,7 @@ def _collect_masked_input(
 def _unmask(
     state: SecAggPlusState, configs: ConfigsRecord
 ) -> Dict[str, ConfigsRecordValues]:
-    log(INFO, "Node %d: starting stage 3...", state.nid)
+    log(DEBUG, "Node %d: starting stage 3...", state.nid)
 
     active_nids = cast(List[int], configs[Key.ACTIVE_NODE_ID_LIST])
     dead_nids = cast(List[int], configs[Key.DEAD_NODE_ID_LIST])
@@ -523,5 +514,5 @@ def _unmask(
     shares += [state.rd_seed_share_dict[nid] for nid in active_nids]
     shares += [state.sk1_share_dict[nid] for nid in dead_nids]
 
-    log(INFO, "Node %d: stage 3 completes. uploading key shares...", state.nid)
+    log(DEBUG, "Node %d: stage 3 completes. uploading key shares...", state.nid)
     return {Key.NODE_ID_LIST: all_nids, Key.SHARE_LIST: shares}

{flwr_nightly-1.8.0.dev20240310 → flwr_nightly-1.8.0.dev20240312}/src/py/flwr/common/logger.py

@@ -168,11 +168,10 @@ def warn_experimental_feature(name: str) -> None:
     """Warn the user when they use an experimental feature."""
     log(
         WARN,
-        """
-        EXPERIMENTAL FEATURE: %s
+        """EXPERIMENTAL FEATURE: %s
 
-        This is an experimental feature. It could change significantly or be removed
-        entirely in future versions of Flower.
+            This is an experimental feature. It could change significantly or be removed
+            entirely in future versions of Flower.
         """,
         name,
     )
@@ -182,11 +181,10 @@ def warn_deprecated_feature(name: str) -> None:
     """Warn the user when they use a deprecated feature."""
     log(
         WARN,
-        """
-        DEPRECATED FEATURE: %s
+        """DEPRECATED FEATURE: %s
 
-        This is a deprecated feature. It will be removed
-        entirely in future versions of Flower.
+            This is a deprecated feature. It will be removed
+            entirely in future versions of Flower.
         """,
         name,
     )

flwr_nightly-1.8.0.dev20240312/src/py/flwr/common/pyproject.py

@@ -0,0 +1,41 @@
+# Copyright 2024 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.
+# ==============================================================================
+"""Validates the project's name property."""
+
+import re
+
+
+def validate_project_name(name: str) -> bool:
+    """Validate the project name against PEP 621 and PEP 503 specifications.
+
+    Conventions at a glance:
+    - Must be lowercase
+    - Must not contain special characters
+    - Must use hyphens(recommended) or underscores. No spaces.
+    - Recommended to be no more than 40 characters long (But it can be)
+
+    Parameters
+    ----------
+    name : str
+        The project name to validate.
+
+    Returns
+    -------
+    bool
+        True if the name is valid, False otherwise.
+    """
+    if not name or len(name) > 40 or not re.match(r"^[a-z0-9-_]+$", name):
+        return False
+    return True

{flwr_nightly-1.8.0.dev20240310 → flwr_nightly-1.8.0.dev20240312}/src/py/flwr/common/secure_aggregation/secaggplus_constants.py

@@ -27,9 +27,9 @@ class Stage:
 
     SETUP = "setup"
     SHARE_KEYS = "share_keys"
-    COLLECT_MASKED_INPUT = "collect_masked_input"
+    COLLECT_MASKED_VECTORS = "collect_masked_vectors"
     UNMASK = "unmask"
-    _stages = (SETUP, SHARE_KEYS, COLLECT_MASKED_INPUT, UNMASK)
+    _stages = (SETUP, SHARE_KEYS, COLLECT_MASKED_VECTORS, UNMASK)
 
     @classmethod
     def all(cls) -> tuple[str, str, str, str]:

{flwr_nightly-1.8.0.dev20240310 → flwr_nightly-1.8.0.dev20240312}/src/py/flwr/server/superlink/state/in_memory_state.py

@@ -122,7 +122,8 @@ class InMemoryState(State):
         task_res.task_id = str(task_id)
         task_res.task.created_at = created_at.isoformat()
         task_res.task.ttl = ttl.isoformat()
-        self.task_res_store[task_id] = task_res
+        with self.lock:
+            self.task_res_store[task_id] = task_res
 
         # Return the new task_id
         return task_id
@@ -132,46 +133,47 @@ class InMemoryState(State):
         if limit is not None and limit < 1:
             raise AssertionError("`limit` must be >= 1")
 
-        # Find TaskRes that were not delivered yet
-        task_res_list: List[TaskRes] = []
-        for _, task_res in self.task_res_store.items():
-            if (
-                UUID(task_res.task.ancestry[0]) in task_ids
-                and task_res.task.delivered_at == ""
-            ):
-                task_res_list.append(task_res)
-            if limit and len(task_res_list) == limit:
-                break
+        with self.lock:
+            # Find TaskRes that were not delivered yet
+            task_res_list: List[TaskRes] = []
+            for _, task_res in self.task_res_store.items():
+                if (
+                    UUID(task_res.task.ancestry[0]) in task_ids
+                    and task_res.task.delivered_at == ""
+                ):
+                    task_res_list.append(task_res)
+                if limit and len(task_res_list) == limit:
+                    break
 
-        # Mark all of them as delivered
-        delivered_at = now().isoformat()
-        for task_res in task_res_list:
-            task_res.task.delivered_at = delivered_at
+            # Mark all of them as delivered
+            delivered_at = now().isoformat()
+            for task_res in task_res_list:
+                task_res.task.delivered_at = delivered_at
 
-        # Return TaskRes
-        return task_res_list
+            # Return TaskRes
+            return task_res_list
 
     def delete_tasks(self, task_ids: Set[UUID]) -> None:
         """Delete all delivered TaskIns/TaskRes pairs."""
         task_ins_to_be_deleted: Set[UUID] = set()
         task_res_to_be_deleted: Set[UUID] = set()
 
-        for task_ins_id in task_ids:
-            # Find the task_id of the matching task_res
-            for task_res_id, task_res in self.task_res_store.items():
-                if UUID(task_res.task.ancestry[0]) != task_ins_id:
-                    continue
-                if task_res.task.delivered_at == "":
-                    continue
-
-                task_ins_to_be_deleted.add(task_ins_id)
-                task_res_to_be_deleted.add(task_res_id)
-
-        for task_id in task_ins_to_be_deleted:
-            with self.lock:
+        with self.lock:
+            for task_ins_id in task_ids:
+                # Find the task_id of the matching task_res
+                for task_res_id, task_res in self.task_res_store.items():
+                    if UUID(task_res.task.ancestry[0]) != task_ins_id:
+                        continue
+                    if task_res.task.delivered_at == "":
+                        continue
+
+                    task_ins_to_be_deleted.add(task_ins_id)
+                    task_res_to_be_deleted.add(task_res_id)
+
+            for task_id in task_ins_to_be_deleted:
                 del self.task_ins_store[task_id]
-        for task_id in task_res_to_be_deleted:
-            del self.task_res_store[task_id]
+            for task_id in task_res_to_be_deleted:
+                del self.task_res_store[task_id]
 
     def num_task_ins(self) -> int:
         """Calculate the number of task_ins in store.

{flwr_nightly-1.8.0.dev20240310 → flwr_nightly-1.8.0.dev20240312}/src/py/flwr/server/workflow/__init__.py

@@ -16,9 +16,10 @@
 
 
 from .default_workflows import DefaultWorkflow
-from .secure_aggregation import SecAggPlusWorkflow
+from .secure_aggregation import SecAggPlusWorkflow, SecAggWorkflow
 
 __all__ = [
     "DefaultWorkflow",
     "SecAggPlusWorkflow",
+    "SecAggWorkflow",
 ]