flwr-nightly 1.10.0.dev20240624__tar.gz → 1.10.0.dev20240707__tar.gz

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/PKG-INFO

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

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "flwr-nightly"
-version = "1.10.0.dev20240624"
+version = "1.10.0.dev20240707"
 description = "Flower: A Friendly Federated Learning Framework"
 license = "Apache-2.0"
 authors = ["The Flower Authors <hello@flower.ai>"]
@@ -122,7 +122,7 @@ mdformat-gfm = "==0.3.5"
 mdformat-frontmatter = "==2.0.1"
 mdformat-beautysh = "==0.1.1"
 mdformat-myst = "==0.1.5"
-twine = "==4.0.2"
+twine = "==5.1.1"
 pyroma = "==4.2"
 check-wheel-contents = "==0.4.0"
 GitPython = "==3.1.32"

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/cli/build.py

@@ -31,7 +31,7 @@ from .utils import get_sha256_hash, is_valid_project_name
 def build(
     directory: Annotated[
         Optional[Path],
-        typer.Option(help="The Flower project directory to bundle into a FAB"),
+        typer.Option(help="Path of the Flower project to bundle into a FAB"),
     ] = None,
 ) -> str:
     """Build a Flower project into a Flower App Bundle (FAB).
@@ -118,7 +118,7 @@ def build(
         fab_file.writestr(".info/CONTENT", list_file_content)
 
     typer.secho(
-        f"🎊 Successfully built {fab_filename}.", fg=typer.colors.GREEN, bold=True
+        f"🎊 Successfully built {fab_filename}", fg=typer.colors.GREEN, bold=True
     )
 
     return fab_filename

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/cli/run/run.py

@@ -17,12 +17,14 @@
 import sys
 from enum import Enum
 from logging import DEBUG
+from pathlib import Path
 from typing import Optional
 
 import typer
 from typing_extensions import Annotated
 
 from flwr.cli import config_utils
+from flwr.cli.build import build
 from flwr.common.constant import SUPEREXEC_DEFAULT_ADDRESS
 from flwr.common.grpc import GRPC_MAX_MESSAGE_LENGTH, create_channel
 from flwr.common.logger import log
@@ -52,10 +54,14 @@ def run(
             case_sensitive=False, help="Use this flag to use the new SuperExec API"
         ),
     ] = False,
+    directory: Annotated[
+        Optional[Path],
+        typer.Option(help="Path of the Flower project to run"),
+    ] = None,
 ) -> None:
     """Run Flower project."""
     if use_superexec:
-        _start_superexec_run()
+        _start_superexec_run(directory)
         return
 
     typer.secho("Loading project configuration... ", fg=typer.colors.BLUE)
@@ -109,7 +115,7 @@ def run(
         )
 
 
-def _start_superexec_run() -> None:
+def _start_superexec_run(directory: Optional[Path]) -> None:
     def on_channel_state_change(channel_connectivity: str) -> None:
         """Log channel connectivity."""
         log(DEBUG, channel_connectivity)
@@ -124,5 +130,8 @@ def _start_superexec_run() -> None:
     channel.subscribe(on_channel_state_change)
     stub = ExecStub(channel)
 
-    req = StartRunRequest()
-    stub.StartRun(req)
+    fab_path = build(directory)
+
+    req = StartRunRequest(fab_file=Path(fab_path).read_bytes())
+    res = stub.StartRun(req)
+    typer.secho(f"🎊 Successfully started run {res.run_id}", fg=typer.colors.GREEN)

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/client/__init__.py

@@ -23,11 +23,13 @@ from .numpy_client import NumPyClient as NumPyClient
 from .supernode import run_client_app as run_client_app
 from .supernode import run_supernode as run_supernode
 from .typing import ClientFn as ClientFn
+from .typing import ClientFnExt as ClientFnExt
 
 __all__ = [
     "Client",
     "ClientApp",
     "ClientFn",
+    "ClientFnExt",
     "NumPyClient",
     "mod",
     "run_client_app",

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/client/app.py

@@ -26,7 +26,7 @@ from grpc import RpcError
 
 from flwr.client.client import Client
 from flwr.client.client_app import ClientApp, LoadClientAppError
-from flwr.client.typing import ClientFn
+from flwr.client.typing import ClientFnExt
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH, EventType, Message, event
 from flwr.common.address import parse_address
 from flwr.common.constant import (
@@ -51,7 +51,7 @@ from .numpy_client import NumPyClient
 
 
 def _check_actionable_client(
-    client: Optional[Client], client_fn: Optional[ClientFn]
+    client: Optional[Client], client_fn: Optional[ClientFnExt]
 ) -> None:
     if client_fn is None and client is None:
         raise ValueError(
@@ -72,7 +72,7 @@ def _check_actionable_client(
 def start_client(
     *,
     server_address: str,
-    client_fn: Optional[ClientFn] = None,
+    client_fn: Optional[ClientFnExt] = None,
     client: Optional[Client] = None,
     grpc_max_message_length: int = GRPC_MAX_MESSAGE_LENGTH,
     root_certificates: Optional[Union[bytes, str]] = None,
@@ -92,7 +92,7 @@ def start_client(
         The IPv4 or IPv6 address of the server. If the Flower
         server runs on the same machine on port 8080, then `server_address`
         would be `"[::]:8080"`.
-    client_fn : Optional[ClientFn]
+    client_fn : Optional[ClientFnExt]
         A callable that instantiates a Client. (default: None)
     client : Optional[flwr.client.Client]
         An implementation of the abstract base
@@ -136,7 +136,7 @@ def start_client(
 
     Starting an SSL-enabled gRPC client using system certificates:
 
-    >>> def client_fn(cid: str):
+    >>> def client_fn(node_id: int, partition_id: Optional[int]):
     >>>     return FlowerClient()
     >>>
     >>> start_client(
@@ -180,7 +180,7 @@ def _start_client_internal(
     *,
     server_address: str,
     load_client_app_fn: Optional[Callable[[str, str], ClientApp]] = None,
-    client_fn: Optional[ClientFn] = None,
+    client_fn: Optional[ClientFnExt] = None,
     client: Optional[Client] = None,
     grpc_max_message_length: int = GRPC_MAX_MESSAGE_LENGTH,
     root_certificates: Optional[Union[bytes, str]] = None,
@@ -191,6 +191,7 @@ def _start_client_internal(
     ] = None,
     max_retries: Optional[int] = None,
     max_wait_time: Optional[float] = None,
+    partition_id: Optional[int] = None,
 ) -> None:
     """Start a Flower client node which connects to a Flower server.
 
@@ -202,7 +203,7 @@ def _start_client_internal(
         would be `"[::]:8080"`.
     load_client_app_fn : Optional[Callable[[], ClientApp]] (default: None)
         A function that can be used to load a `ClientApp` instance.
-    client_fn : Optional[ClientFn]
+    client_fn : Optional[ClientFnExt]
         A callable that instantiates a Client. (default: None)
     client : Optional[flwr.client.Client]
         An implementation of the abstract base
@@ -234,6 +235,9 @@ def _start_client_internal(
         The maximum duration before the client stops trying to
         connect to the server in case of connection error.
         If set to None, there is no limit to the total time.
+    partitioni_id: Optional[int] (default: None)
+        The data partition index associated with this node. Better suited for
+        prototyping purposes.
     """
     if insecure is None:
         insecure = root_certificates is None
@@ -244,7 +248,8 @@ def _start_client_internal(
         if client_fn is None:
             # Wrap `Client` instance in `client_fn`
             def single_client_factory(
-                cid: str,  # pylint: disable=unused-argument
+                node_id: int,  # pylint: disable=unused-argument
+                partition_id: Optional[int],  # pylint: disable=unused-argument
             ) -> Client:
                 if client is None:  # Added this to keep mypy happy
                     raise ValueError(
@@ -293,7 +298,7 @@ def _start_client_internal(
     retry_invoker = RetryInvoker(
         wait_gen_factory=exponential,
         recoverable_exceptions=connection_error_type,
-        max_tries=max_retries,
+        max_tries=max_retries + 1 if max_retries is not None else None,
         max_time=max_wait_time,
         on_giveup=lambda retry_state: (
             log(
@@ -309,7 +314,7 @@ def _start_client_internal(
         on_backoff=_on_backoff,
     )
 
-    node_state = NodeState()
+    node_state = NodeState(partition_id=partition_id)
     # run_id -> (fab_id, fab_version)
     run_info: Dict[int, Tuple[str, str]] = {}
 

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/client/client_app.py

@@ -15,19 +15,42 @@
 """Flower ClientApp."""
 
 
+import inspect
 from typing import Callable, List, Optional
 
+from flwr.client.client import Client
 from flwr.client.message_handler.message_handler import (
     handle_legacy_message_from_msgtype,
 )
 from flwr.client.mod.utils import make_ffn
-from flwr.client.typing import ClientFn, Mod
+from flwr.client.typing import ClientFnExt, Mod
 from flwr.common import Context, Message, MessageType
-from flwr.common.logger import warn_preview_feature
+from flwr.common.logger import warn_deprecated_feature, warn_preview_feature
 
 from .typing import ClientAppCallable
 
 
+def _inspect_maybe_adapt_client_fn_signature(client_fn: ClientFnExt) -> ClientFnExt:
+    client_fn_args = inspect.signature(client_fn).parameters
+
+    if not all(key in client_fn_args for key in ["node_id", "partition_id"]):
+        warn_deprecated_feature(
+            "`client_fn` now expects a signature `def client_fn(node_id: int, "
+            "partition_id: Optional[int])`.\nYou provided `client_fn` with signature: "
+            f"{dict(client_fn_args.items())}"
+        )
+
+        # Wrap depcreated client_fn inside a function with the expected signature
+        def adaptor_fn(
+            node_id: int, partition_id: Optional[int]  # pylint: disable=unused-argument
+        ) -> Client:
+            return client_fn(str(partition_id))  # type: ignore
+
+        return adaptor_fn
+
+    return client_fn
+
+
 class ClientAppException(Exception):
     """Exception raised when an exception is raised while executing a ClientApp."""
 
@@ -48,7 +71,7 @@ class ClientApp:
     >>> class FlowerClient(NumPyClient):
     >>>     # ...
     >>>
-    >>> def client_fn(cid):
+    >>> def client_fn(node_id: int, partition_id: Optional[int]):
     >>>    return FlowerClient().to_client()
     >>>
     >>> app = ClientApp(client_fn)
@@ -65,7 +88,7 @@ class ClientApp:
 
     def __init__(
         self,
-        client_fn: Optional[ClientFn] = None,  # Only for backward compatibility
+        client_fn: Optional[ClientFnExt] = None,  # Only for backward compatibility
         mods: Optional[List[Mod]] = None,
     ) -> None:
         self._mods: List[Mod] = mods if mods is not None else []
@@ -74,6 +97,8 @@ class ClientApp:
         self._call: Optional[ClientAppCallable] = None
         if client_fn is not None:
 
+            client_fn = _inspect_maybe_adapt_client_fn_signature(client_fn)
+
             def ffn(
                 message: Message,
                 context: Context,

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/client/message_handler/message_handler.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 """Client-side message handler."""
 
-
 from logging import WARN
 from typing import Optional, Tuple, cast
 
@@ -25,7 +24,7 @@ from flwr.client.client import (
     maybe_call_get_properties,
 )
 from flwr.client.numpy_client import NumPyClient
-from flwr.client.typing import ClientFn
+from flwr.client.typing import ClientFnExt
 from flwr.common import ConfigsRecord, Context, Message, Metadata, RecordSet, log
 from flwr.common.constant import MessageType, MessageTypeLegacy
 from flwr.common.recordset_compat import (
@@ -90,10 +89,10 @@ def handle_control_message(message: Message) -> Tuple[Optional[Message], int]:
 
 
 def handle_legacy_message_from_msgtype(
-    client_fn: ClientFn, message: Message, context: Context
+    client_fn: ClientFnExt, message: Message, context: Context
 ) -> Message:
     """Handle legacy message in the inner most mod."""
-    client = client_fn(str(message.metadata.partition_id))
+    client = client_fn(message.metadata.dst_node_id, context.partition_id)
 
     # Check if NumPyClient is returend
     if isinstance(client, NumPyClient):

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/client/node_state.py

@@ -15,7 +15,7 @@
 """Node state."""
 
 
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 from flwr.common import Context, RecordSet
 
@@ -23,14 +23,17 @@ from flwr.common import Context, RecordSet
 class NodeState:
     """State of a node where client nodes execute runs."""
 
-    def __init__(self) -> None:
+    def __init__(self, partition_id: Optional[int]) -> None:
         self._meta: Dict[str, Any] = {}  # holds metadata about the node
         self.run_contexts: Dict[int, Context] = {}
+        self._partition_id = partition_id
 
     def register_context(self, run_id: int) -> None:
         """Register new run context for this node."""
         if run_id not in self.run_contexts:
-            self.run_contexts[run_id] = Context(state=RecordSet())
+            self.run_contexts[run_id] = Context(
+                state=RecordSet(), partition_id=self._partition_id
+            )
 
     def retrieve_context(self, run_id: int) -> Context:
         """Get run context given a run_id."""

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/client/node_state_tests.py

@@ -41,7 +41,7 @@ def test_multirun_in_node_state() -> None:
     expected_values = {0: "1", 1: "1" * 3, 2: "1" * 2, 3: "1", 5: "1"}
 
     # NodeState
-    node_state = NodeState()
+    node_state = NodeState(partition_id=None)
 
     for task in tasks:
         run_id = task.run_id

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/client/supernode/app.py

@@ -67,6 +67,7 @@ def run_supernode() -> None:
         authentication_keys=authentication_keys,
         max_retries=args.max_retries,
         max_wait_time=args.max_wait_time,
+        partition_id=args.partition_id,
     )
 
     # Graceful shutdown
@@ -344,8 +345,8 @@ def _parse_args_common(parser: argparse.ArgumentParser) -> None:
         "--max-retries",
         type=int,
         default=None,
-        help="The maximum number of times the client will try to connect to the"
-        "server before giving up in case of a connection error. By default,"
+        help="The maximum number of times the client will try to reconnect to the"
+        "SuperLink before giving up in case of a connection error. By default,"
         "it is set to None, meaning there is no limit to the number of tries.",
     )
     parser.add_argument(
@@ -353,7 +354,7 @@ def _parse_args_common(parser: argparse.ArgumentParser) -> None:
         type=float,
         default=None,
         help="The maximum duration before the client stops trying to"
-        "connect to the server in case of connection error. By default, it"
+        "connect to the SuperLink in case of connection error. By default, it"
         "is set to None, meaning there is no limit to the total time.",
     )
     parser.add_argument(
@@ -373,6 +374,13 @@ def _parse_args_common(parser: argparse.ArgumentParser) -> None:
         type=str,
         help="The SuperNode's public key (as a path str) to enable authentication.",
     )
+    parser.add_argument(
+        "--partition-id",
+        type=int,
+        help="The data partition index associated with this SuperNode. Better suited "
+        "for prototyping purposes where a SuperNode might only load a fraction of an "
+        "artificially partitioned dataset (e.g. using `flwr-datasets`)",
+    )
 
 
 def _try_setup_client_authentication(

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/client/typing.py

@@ -15,7 +15,7 @@
 """Custom types for Flower clients."""
 
 
-from typing import Callable
+from typing import Callable, Optional
 
 from flwr.common import Context, Message
 
@@ -23,6 +23,7 @@ from .client import Client as Client
 
 # Compatibility
 ClientFn = Callable[[str], Client]
+ClientFnExt = Callable[[int, Optional[int]], Client]
 
 ClientAppCallable = Callable[[Message, Context], Message]
 Mod = Callable[[Message, Context, ClientAppCallable], Message]

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/common/constant.py

@@ -46,6 +46,9 @@ PING_BASE_MULTIPLIER = 0.8
 PING_RANDOM_RANGE = (-0.1, 0.1)
 PING_MAX_INTERVAL = 1e300
 
+# IDs
+RUN_ID_NUM_BYTES = 8
+NODE_ID_NUM_BYTES = 8
 GRPC_ADAPTER_METADATA_FLOWER_VERSION_KEY = "flower-version"
 GRPC_ADAPTER_METADATA_SHOULD_EXIT_KEY = "should-exit"
 

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/common/context.py

@@ -16,13 +16,14 @@
 
 
 from dataclasses import dataclass
+from typing import Optional
 
 from .record import RecordSet
 
 
 @dataclass
 class Context:
-    """State of your run.
+    """Context of your run.
 
     Parameters
     ----------
@@ -33,6 +34,15 @@ class Context:
         executing mods. It can also be used as a memory to access
         at different points during the lifecycle of this entity (e.g. across
         multiple rounds)
+    partition_id : Optional[int] (default: None)
+        An index that specifies the data partition that the ClientApp using this Context
+        object should make use of. Setting this attribute is better suited for
+        simulation or proto typing setups.
     """
 
     state: RecordSet
+    partition_id: Optional[int]
+
+    def __init__(self, state: RecordSet, partition_id: Optional[int] = None) -> None:
+        self.state = state
+        self.partition_id = partition_id

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/common/logger.py

@@ -197,6 +197,19 @@ def warn_deprecated_feature(name: str) -> None:
     )
 
 
+def warn_unsupported_feature(name: str) -> None:
+    """Warn the user when they use an unsupported feature."""
+    log(
+        WARN,
+        """UNSUPPORTED FEATURE: %s
+
+            This is an unsupported feature. It will be removed
+            entirely in future versions of Flower.
+        """,
+        name,
+    )
+
+
 def set_logger_propagation(
     child_logger: logging.Logger, value: bool = True
 ) -> logging.Logger:

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/common/message.py

@@ -48,10 +48,6 @@ class Metadata:  # pylint: disable=too-many-instance-attributes
     message_type : str
         A string that encodes the action to be executed on
         the receiving end.
-    partition_id : Optional[int]
-        An identifier that can be used when loading a particular
-        data partition for a ClientApp. Making use of this identifier
-        is more relevant when conducting simulations.
     """
 
     def __init__(  # pylint: disable=too-many-arguments
@@ -64,7 +60,6 @@ class Metadata:  # pylint: disable=too-many-instance-attributes
         group_id: str,
         ttl: float,
         message_type: str,
-        partition_id: int | None = None,
     ) -> None:
         var_dict = {
             "_run_id": run_id,
@@ -75,7 +70,6 @@ class Metadata:  # pylint: disable=too-many-instance-attributes
             "_group_id": group_id,
             "_ttl": ttl,
             "_message_type": message_type,
-            "_partition_id": partition_id,
         }
         self.__dict__.update(var_dict)
 
@@ -149,16 +143,6 @@ class Metadata:  # pylint: disable=too-many-instance-attributes
         """Set message_type."""
         self.__dict__["_message_type"] = value
 
-    @property
-    def partition_id(self) -> int | None:
-        """An identifier telling which data partition a ClientApp should use."""
-        return cast(int, self.__dict__["_partition_id"])
-
-    @partition_id.setter
-    def partition_id(self, value: int) -> None:
-        """Set partition_id."""
-        self.__dict__["_partition_id"] = value
-
     def __repr__(self) -> str:
         """Return a string representation of this instance."""
         view = ", ".join([f"{k.lstrip('_')}={v!r}" for k, v in self.__dict__.items()])
@@ -398,5 +382,4 @@ def _create_reply_metadata(msg: Message, ttl: float) -> Metadata:
         group_id=msg.metadata.group_id,
         ttl=ttl,
         message_type=msg.metadata.message_type,
-        partition_id=msg.metadata.partition_id,
     )

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/server/superlink/fleet/vce/backend/raybackend.py

@@ -168,7 +168,7 @@ class RayBackend(Backend):
 
         Return output message and updated context.
         """
-        partition_id = message.metadata.partition_id
+        partition_id = context.partition_id
 
         try:
             # Submit a task to the pool

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/server/superlink/fleet/vce/vce_api.py

@@ -57,7 +57,6 @@ async def worker(
     queue: "asyncio.Queue[TaskIns]",
     node_states: Dict[int, NodeState],
     state_factory: StateFactory,
-    nodes_mapping: NodeToPartitionMapping,
     backend: Backend,
 ) -> None:
     """Get TaskIns from queue and pass it to an actor in the pool to execute it."""
@@ -74,8 +73,6 @@ async def worker(
 
             # Convert TaskIns to Message
             message = message_from_taskins(task_ins)
-            # Set partition_id
-            message.metadata.partition_id = nodes_mapping[node_id]
 
             # Let backend process message
             out_mssg, updated_context = await backend.process_message(
@@ -187,9 +184,7 @@ async def run(
         # Add workers (they submit Messages to Backend)
         worker_tasks = [
             asyncio.create_task(
-                worker(
-                    app_fn, queue, node_states, state_factory, nodes_mapping, backend
-                )
+                worker(app_fn, queue, node_states, state_factory, backend)
             )
             for _ in range(backend.num_workers)
         ]
@@ -291,8 +286,8 @@ def start_vce(
 
     # Construct mapping of NodeStates
     node_states: Dict[int, NodeState] = {}
-    for node_id in nodes_mapping:
-        node_states[node_id] = NodeState()
+    for node_id, partition_id in nodes_mapping.items():
+        node_states[node_id] = NodeState(partition_id=partition_id)
 
     # Load backend config
     log(DEBUG, "Supported backends: %s", list(supported_backends.keys()))

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/server/superlink/state/in_memory_state.py

@@ -15,7 +15,6 @@
 """In-memory State implementation."""
 
 
-import os
 import threading
 import time
 from logging import ERROR
@@ -23,12 +22,13 @@ from typing import Dict, List, Optional, Set, Tuple
 from uuid import UUID, uuid4
 
 from flwr.common import log, now
+from flwr.common.constant import NODE_ID_NUM_BYTES, RUN_ID_NUM_BYTES
 from flwr.common.typing import Run
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.state.state import State
 from flwr.server.utils import validate_task_ins_or_res
 
-from .utils import make_node_unavailable_taskres
+from .utils import generate_rand_int_from_bytes, make_node_unavailable_taskres
 
 
 class InMemoryState(State):  # pylint: disable=R0902,R0904
@@ -216,7 +216,7 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
     ) -> int:
         """Create, store in state, and return `node_id`."""
         # Sample a random int64 as node_id
-        node_id: int = int.from_bytes(os.urandom(8), "little", signed=True)
+        node_id = generate_rand_int_from_bytes(NODE_ID_NUM_BYTES)
 
         with self.lock:
             if node_id in self.node_ids:
@@ -279,7 +279,7 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
         """Create a new run for the specified `fab_id` and `fab_version`."""
         # Sample a random int64 as run_id
         with self.lock:
-            run_id: int = int.from_bytes(os.urandom(8), "little", signed=True)
+            run_id = generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
 
             if run_id not in self.run_ids:
                 self.run_ids[run_id] = Run(

{flwr_nightly-1.10.0.dev20240624 → flwr_nightly-1.10.0.dev20240707}/src/py/flwr/server/superlink/state/sqlite_state.py

@@ -15,7 +15,6 @@
 """SQLite based implemenation of server state."""
 
 
-import os
 import re
 import sqlite3
 import time
@@ -24,6 +23,7 @@ from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union, cast
 from uuid import UUID, uuid4
 
 from flwr.common import log, now
+from flwr.common.constant import NODE_ID_NUM_BYTES, RUN_ID_NUM_BYTES
 from flwr.common.typing import Run
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.recordset_pb2 import RecordSet  # pylint: disable=E0611
@@ -31,7 +31,7 @@ from flwr.proto.task_pb2 import Task, TaskIns, TaskRes  # pylint: disable=E0611
 from flwr.server.utils.validator import validate_task_ins_or_res
 
 from .state import State
-from .utils import make_node_unavailable_taskres
+from .utils import generate_rand_int_from_bytes, make_node_unavailable_taskres
 
 SQL_CREATE_TABLE_NODE = """
 CREATE TABLE IF NOT EXISTS node(
@@ -541,7 +541,7 @@ class SqliteState(State):  # pylint: disable=R0904
     ) -> int:
         """Create, store in state, and return `node_id`."""
         # Sample a random int64 as node_id
-        node_id: int = int.from_bytes(os.urandom(8), "little", signed=True)
+        node_id = generate_rand_int_from_bytes(NODE_ID_NUM_BYTES)
 
         query = "SELECT node_id FROM node WHERE public_key = :public_key;"
         row = self.query(query, {"public_key": public_key})
@@ -616,7 +616,7 @@ class SqliteState(State):  # pylint: disable=R0904
     def create_run(self, fab_id: str, fab_version: str) -> int:
         """Create a new run for the specified `fab_id` and `fab_version`."""
         # Sample a random int64 as run_id
-        run_id: int = int.from_bytes(os.urandom(8), "little", signed=True)
+        run_id = generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
 
         # Check conflicts
         query = "SELECT COUNT(*) FROM run WHERE run_id = ?;"