flwr-nightly 1.10.0.dev20240619__py3-none-any.whl → 1.10.0.dev20240707__py3-none-any.whl

flwr/server/run_serverapp.py

@@ -22,10 +22,13 @@ from pathlib import Path
 from typing import Optional
 
 from flwr.common import Context, EventType, RecordSet, event
+from flwr.common.config import get_flwr_dir, get_project_config, get_project_dir
 from flwr.common.logger import log, update_console_handler, warn_deprecated_feature
 from flwr.common.object_ref import load_app
+from flwr.proto.driver_pb2 import CreateRunRequest  # pylint: disable=E0611
 
-from .driver import Driver, GrpcDriver
+from .driver import Driver
+from .driver.grpc_driver import GrpcDriver, GrpcDriverStub
 from .server_app import LoadServerAppError, ServerApp
 
 ADDRESS_DRIVER_API = "0.0.0.0:9091"
@@ -41,7 +44,7 @@ def run(
     if not (server_app_attr is None) ^ (loaded_server_app is None):
         raise ValueError(
             "Either `server_app_attr` or `loaded_server_app` should be set "
-            "but not both. "
+            "but not both."
         )
 
     if server_app_dir is not None:
@@ -74,7 +77,7 @@ def run(
     log(DEBUG, "ServerApp finished running.")
 
 
-def run_server_app() -> None:
+def run_server_app() -> None:  # pylint: disable=too-many-branches
     """Run Flower server app."""
     event(EventType.RUN_SERVER_APP_ENTER)
 
@@ -134,11 +137,43 @@ def run_server_app() -> None:
             cert_path,
         )
 
-    log(
-        DEBUG,
-        "Flower will load ServerApp `%s`",
-        getattr(args, "server-app"),
+    server_app_attr: Optional[str] = getattr(args, "server-app")
+    if not (server_app_attr is None) ^ (args.run_id is None):
+        raise sys.exit(
+            "Please provide either a ServerApp reference or a Run ID, but not both. "
+            "For more details, use: ``flower-server-app -h``"
+        )
+
+    stub = GrpcDriverStub(
+        driver_service_address=args.superlink, root_certificates=root_certificates
     )
+    if args.run_id is not None:
+        # User provided `--run-id`, but not `server-app`
+        run_id = args.run_id
+    else:
+        # User provided `server-app`, but not `--run-id`
+        # Create run if run_id is not provided
+        stub.connect()
+        req = CreateRunRequest(fab_id=args.fab_id, fab_version=args.fab_version)
+        res = stub.create_run(req)
+        run_id = res.run_id
+
+    # Initialize GrpcDriver
+    driver = GrpcDriver(run_id=run_id, stub=stub)
+
+    # Dynamically obtain ServerApp path based on run_id
+    if args.run_id is not None:
+        # User provided `--run-id`, but not `server-app`
+        flwr_dir = get_flwr_dir(args.flwr_dir)
+        run_ = driver.run
+        server_app_dir = str(get_project_dir(run_.fab_id, run_.fab_version, flwr_dir))
+        config = get_project_config(server_app_dir)
+        server_app_attr = config["flower"]["components"]["serverapp"]
+    else:
+        # User provided `server-app`, but not `--run-id`
+        server_app_dir = str(Path(args.dir).absolute())
+
+    log(DEBUG, "Flower will load ServerApp `%s` in %s", server_app_attr, server_app_dir)
 
     log(
         DEBUG,
@@ -146,17 +181,6 @@ def run_server_app() -> None:
         root_certificates,
     )
 
-    server_app_dir = args.dir
-    server_app_attr = getattr(args, "server-app")
-
-    # Initialize GrpcDriver
-    driver = GrpcDriver(
-        driver_service_address=args.superlink,
-        root_certificates=root_certificates,
-        fab_id=args.fab_id,
-        fab_version=args.fab_version,
-    )
-
     # Run the ServerApp with the Driver
     run(driver=driver, server_app_dir=server_app_dir, server_app_attr=server_app_attr)
 
@@ -174,6 +198,8 @@ def _parse_args_run_server_app() -> argparse.ArgumentParser:
 
     parser.add_argument(
         "server-app",
+        nargs="?",
+        default=None,
         help="For example: `server:app` or `project.package.module:wrapper.app`",
     )
     parser.add_argument(
@@ -223,5 +249,22 @@ def _parse_args_run_server_app() -> argparse.ArgumentParser:
         type=str,
         help="The version of the FAB used in the run.",
     )
+    parser.add_argument(
+        "--run-id",
+        default=None,
+        type=int,
+        help="The identifier of the run.",
+    )
+    parser.add_argument(
+        "--flwr-dir",
+        default=None,
+        help="""The path containing installed Flower Apps.
+    By default, this value is equal to:
+
+        - `$FLWR_HOME/` if `$FLWR_HOME` is defined
+        - `$XDG_DATA_HOME/.flwr/` if `$XDG_DATA_HOME` is defined
+        - `$HOME/.flwr/` in all other cases
+    """,
+    )
 
     return parser

flwr/server/strategy/bulyan.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2023 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.

flwr/server/strategy/dpfedavg_adaptive.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2022 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.

flwr/server/strategy/dpfedavg_fixed.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2022 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.

flwr/server/strategy/fedadagrad.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2021 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.

flwr/server/strategy/fedadam.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2021 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.

flwr/server/strategy/fedavg_android.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2021 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.

flwr/server/strategy/fedavgm.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2022 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.

flwr/server/strategy/fedmedian.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2022 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.

flwr/server/strategy/fedopt.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2021 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.

flwr/server/strategy/fedprox.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2023 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.

flwr/server/strategy/fedxgb_bagging.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2023 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.

flwr/server/strategy/fedxgb_cyclic.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2023 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.

flwr/server/strategy/fedxgb_nn_avg.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2023 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.

flwr/server/strategy/fedyogi.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2021 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.

flwr/server/strategy/krum.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2022 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.

flwr/server/strategy/qfedavg.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2021 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.

flwr/server/superlink/driver/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/driver/driver_grpc.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/driver/driver_servicer.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# 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.
@@ -35,7 +35,11 @@ from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
     PushTaskInsResponse,
 )
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
-from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
+from flwr.proto.run_pb2 import (  # pylint: disable=E0611
+    GetRunRequest,
+    GetRunResponse,
+    Run,
+)
 from flwr.proto.task_pb2 import TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.state import State, StateFactory
 from flwr.server.utils.validator import validate_task_ins_or_res
@@ -134,7 +138,15 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         self, request: GetRunRequest, context: grpc.ServicerContext
     ) -> GetRunResponse:
         """Get run information."""
-        raise NotImplementedError
+        log(DEBUG, "DriverServicer.GetRun")
+
+        # Init state
+        state: State = self.state_factory.state()
+
+        # Retrieve run information
+        run = state.get_run(request.run_id)
+        run_proto = None if run is None else Run(**vars(run))
+        return GetRunResponse(run=run_proto)
 
 
 def _raise_if(validation_error: bool, detail: str) -> None:

flwr/server/superlink/fleet/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/fleet/grpc_bidi/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/fleet/grpc_bidi/flower_service_servicer.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/fleet/grpc_bidi/grpc_bridge.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/fleet/grpc_bidi/grpc_client_proxy.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/fleet/grpc_bidi/grpc_server.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/fleet/grpc_rere/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/fleet/grpc_rere/fleet_servicer.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/fleet/message_handler/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/fleet/message_handler/message_handler.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/fleet/rest_rere/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/fleet/rest_rere/rest_api.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/fleet/vce/backend/raybackend.py

@@ -15,7 +15,7 @@
 """Ray backend for the Fleet API using the Simulation Engine."""
 
 import pathlib
-from logging import DEBUG, ERROR, WARNING
+from logging import DEBUG, ERROR
 from typing import Callable, Dict, List, Tuple, Union
 
 import ray
@@ -24,16 +24,15 @@ from flwr.client.client_app import ClientApp
 from flwr.common.context import Context
 from flwr.common.logger import log
 from flwr.common.message import Message
-from flwr.simulation.ray_transport.ray_actor import (
-    BasicActorPool,
-    ClientAppActor,
-    init_ray,
-)
+from flwr.common.typing import ConfigsRecordValues
+from flwr.simulation.ray_transport.ray_actor import BasicActorPool, ClientAppActor
 from flwr.simulation.ray_transport.utils import enable_tf_gpu_growth
 
 from .backend import Backend, BackendConfig
 
 ClientResourcesDict = Dict[str, Union[int, float]]
+ActorArgsDict = Dict[str, Union[int, float, Callable[[], None]]]
+RunTimeEnvDict = Dict[str, Union[str, List[str]]]
 
 
 class RayBackend(Backend):
@@ -51,40 +50,29 @@ class RayBackend(Backend):
         if not pathlib.Path(work_dir).exists():
             raise ValueError(f"Specified work_dir {work_dir} does not exist.")
 
-        # Init ray and append working dir if needed
-        runtime_env = (
-            self._configure_runtime_env(work_dir=work_dir) if work_dir else None
-        )
-
-        if backend_config.get("mute_logging", False):
-            init_ray(
-                logging_level=WARNING, log_to_driver=False, runtime_env=runtime_env
-            )
-        elif backend_config.get("silent", False):
-            init_ray(logging_level=WARNING, log_to_driver=True, runtime_env=runtime_env)
-        else:
-            init_ray(runtime_env=runtime_env)
+        # Initialise ray
+        self.init_args_key = "init_args"
+        self.init_ray(backend_config, work_dir)
 
         # Validate client resources
         self.client_resources_key = "client_resources"
+        client_resources = self._validate_client_resources(config=backend_config)
 
         # Create actor pool
-        use_tf = backend_config.get("tensorflow", False)
-        actor_kwargs = {"on_actor_init_fn": enable_tf_gpu_growth} if use_tf else {}
+        actor_kwargs = self._validate_actor_arguments(config=backend_config)
 
-        client_resources = self._validate_client_resources(config=backend_config)
         self.pool = BasicActorPool(
             actor_type=ClientAppActor,
             client_resources=client_resources,
             actor_kwargs=actor_kwargs,
         )
 
-    def _configure_runtime_env(self, work_dir: str) -> Dict[str, Union[str, List[str]]]:
+    def _configure_runtime_env(self, work_dir: str) -> RunTimeEnvDict:
         """Return list of files/subdirectories to exclude relative to work_dir.
 
         Without this, Ray will push everything to the Ray Cluster.
         """
-        runtime_env: Dict[str, Union[str, List[str]]] = {"working_dir": work_dir}
+        runtime_env: RunTimeEnvDict = {"working_dir": work_dir}
 
         excludes = []
         path = pathlib.Path(work_dir)
@@ -125,6 +113,37 @@ class RayBackend(Backend):
 
         return client_resources
 
+    def _validate_actor_arguments(self, config: BackendConfig) -> ActorArgsDict:
+        actor_args_config = config.get("actor", False)
+        actor_args: ActorArgsDict = {}
+        if actor_args_config:
+            use_tf = actor_args.get("tensorflow", False)
+            if use_tf:
+                actor_args["on_actor_init_fn"] = enable_tf_gpu_growth
+        return actor_args
+
+    def init_ray(self, backend_config: BackendConfig, work_dir: str) -> None:
+        """Intialises Ray if not already initialised."""
+        if not ray.is_initialized():
+            # Init ray and append working dir if needed
+            runtime_env = (
+                self._configure_runtime_env(work_dir=work_dir) if work_dir else None
+            )
+
+            ray_init_args: Dict[
+                str,
+                Union[ConfigsRecordValues, RunTimeEnvDict],
+            ] = {}
+
+            if backend_config.get(self.init_args_key):
+                for k, v in backend_config[self.init_args_key].items():
+                    ray_init_args[k] = v
+
+            if runtime_env is not None:
+                ray_init_args["runtime_env"] = runtime_env
+
+            ray.init(**ray_init_args)
+
     @property
     def num_workers(self) -> int:
         """Return number of actors in pool."""
@@ -149,10 +168,10 @@ class RayBackend(Backend):
 
         Return output message and updated context.
         """
-        partition_id = message.metadata.partition_id
+        partition_id = context.partition_id
 
         try:
-            # Submite a task to the pool
+            # Submit a task to the pool
             future = await self.pool.submit(
                 lambda a, a_fn, mssg, cid, state: a.run.remote(a_fn, mssg, cid, state),
                 (app, message, str(partition_id), context),

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/server/superlink/state/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/server/superlink/state/in_memory_state.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# 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.
@@ -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/server/superlink/state/sqlite_state.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# 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.
@@ -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 = ?;"