flwr-nightly 1.10.0.dev20240612__py3-none-any.whl → 1.10.0.dev20240624__py3-none-any.whl

flwr/proto/fab_pb2.pyi

@@ -0,0 +1,56 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""
+import builtins
+import google.protobuf.descriptor
+import google.protobuf.message
+import typing
+import typing_extensions
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+class Fab(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    HASH_FIELD_NUMBER: builtins.int
+    CONTENT_FIELD_NUMBER: builtins.int
+    hash: typing.Text
+    """This field is the hash of the data field. It is used to identify the data.
+    The hash is calculated using the SHA-256 algorithm and is represented as a
+    hex string (sha256hex).
+    """
+
+    content: builtins.bytes
+    """This field contains the fab file contents a one bytes blob."""
+
+    def __init__(self,
+        *,
+        hash: typing.Text = ...,
+        content: builtins.bytes = ...,
+        ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["content",b"content","hash",b"hash"]) -> None: ...
+global___Fab = Fab
+
+class GetFabRequest(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    HASH_FIELD_NUMBER: builtins.int
+    hash: typing.Text
+    def __init__(self,
+        *,
+        hash: typing.Text = ...,
+        ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["hash",b"hash"]) -> None: ...
+global___GetFabRequest = GetFabRequest
+
+class GetFabResponse(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    FAB_FIELD_NUMBER: builtins.int
+    @property
+    def fab(self) -> global___Fab: ...
+    def __init__(self,
+        *,
+        fab: typing.Optional[global___Fab] = ...,
+        ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["fab",b"fab"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["fab",b"fab"]) -> None: ...
+global___GetFabResponse = GetFabResponse

flwr/proto/fab_pb2_grpc.py

@@ -0,0 +1,4 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+

flwr/proto/fab_pb2_grpc.pyi

@@ -0,0 +1,4 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""

flwr/server/__init__.py

@@ -34,12 +34,12 @@ __all__ = [
     "Driver",
     "History",
     "LegacyContext",
-    "run_server_app",
-    "run_superlink",
     "Server",
     "ServerApp",
     "ServerConfig",
     "SimpleClientManager",
+    "run_server_app",
+    "run_superlink",
     "start_server",
     "strategy",
     "workflow",

flwr/server/app.py

@@ -36,6 +36,7 @@ from flwr.common import GRPC_MAX_MESSAGE_LENGTH, EventType, event
 from flwr.common.address import parse_address
 from flwr.common.constant import (
     MISSING_EXTRA_REST,
+    TRANSPORT_TYPE_GRPC_ADAPTER,
     TRANSPORT_TYPE_GRPC_RERE,
     TRANSPORT_TYPE_REST,
 )
@@ -48,6 +49,7 @@ from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
 from flwr.proto.fleet_pb2_grpc import (  # pylint: disable=E0611
     add_FleetServicer_to_server,
 )
+from flwr.proto.grpcadapter_pb2_grpc import add_GrpcAdapterServicer_to_server
 
 from .client_manager import ClientManager
 from .history import History
@@ -55,6 +57,7 @@ from .server import Server, init_defaults, run_fl
 from .server_config import ServerConfig
 from .strategy import Strategy
 from .superlink.driver.driver_grpc import run_driver_api_grpc
+from .superlink.fleet.grpc_adapter.grpc_adapter_servicer import GrpcAdapterServicer
 from .superlink.fleet.grpc_bidi.grpc_server import (
     generic_create_grpc_server,
     start_grpc_server,
@@ -200,15 +203,7 @@ def run_superlink() -> None:
     args = _parse_args_run_superlink().parse_args()
 
     # Parse IP address
-    parsed_driver_address = parse_address(args.driver_api_address)
-    if not parsed_driver_address:
-        sys.exit(f"Driver IP address ({args.driver_api_address}) cannot be parsed.")
-    driver_host, driver_port, driver_is_v6 = parsed_driver_address
-    driver_address = (
-        f"[{driver_host}]:{driver_port}"
-        if driver_is_v6
-        else f"{driver_host}:{driver_port}"
-    )
+    driver_address, _, _ = _format_address(args.driver_api_address)
 
     # Obtain certificates
     certificates = _try_obtain_certificates(args)
@@ -226,18 +221,15 @@ def run_superlink() -> None:
     grpc_servers = [driver_server]
     bckg_threads = []
     if not args.fleet_api_address:
-        args.fleet_api_address = (
-            ADDRESS_FLEET_API_GRPC_RERE
-            if args.fleet_api_type == TRANSPORT_TYPE_GRPC_RERE
-            else ADDRESS_FLEET_API_REST
-        )
-    parsed_fleet_address = parse_address(args.fleet_api_address)
-    if not parsed_fleet_address:
-        sys.exit(f"Fleet IP address ({args.fleet_api_address}) cannot be parsed.")
-    fleet_host, fleet_port, fleet_is_v6 = parsed_fleet_address
-    fleet_address = (
-        f"[{fleet_host}]:{fleet_port}" if fleet_is_v6 else f"{fleet_host}:{fleet_port}"
-    )
+        if args.fleet_api_type in [
+            TRANSPORT_TYPE_GRPC_RERE,
+            TRANSPORT_TYPE_GRPC_ADAPTER,
+        ]:
+            args.fleet_api_address = ADDRESS_FLEET_API_GRPC_RERE
+        elif args.fleet_api_type == TRANSPORT_TYPE_REST:
+            args.fleet_api_address = ADDRESS_FLEET_API_REST
+
+    fleet_address, host, port = _format_address(args.fleet_api_address)
 
     num_workers = args.fleet_api_num_workers
     if num_workers != 1:
@@ -267,8 +259,8 @@ def run_superlink() -> None:
         fleet_thread = threading.Thread(
             target=_run_fleet_api_rest,
             args=(
-                fleet_host,
-                fleet_port,
+                host,
+                port,
                 ssl_keyfile,
                 ssl_certfile,
                 state_factory,
@@ -306,6 +298,13 @@ def run_superlink() -> None:
             interceptors=interceptors,
         )
         grpc_servers.append(fleet_server)
+    elif args.fleet_api_type == TRANSPORT_TYPE_GRPC_ADAPTER:
+        fleet_server = _run_fleet_api_grpc_adapter(
+            address=fleet_address,
+            state_factory=state_factory,
+            certificates=certificates,
+        )
+        grpc_servers.append(fleet_server)
     else:
         raise ValueError(f"Unknown fleet_api_type: {args.fleet_api_type}")
 
@@ -325,6 +324,16 @@ def run_superlink() -> None:
         driver_server.wait_for_termination(timeout=1)
 
 
+def _format_address(address: str) -> Tuple[str, str, int]:
+    parsed_address = parse_address(address)
+    if not parsed_address:
+        sys.exit(
+            f"Address ({address}) cannot be parsed (expected: URL or IPv4 or IPv6)."
+        )
+    host, port, is_v6 = parsed_address
+    return (f"[{host}]:{port}" if is_v6 else f"{host}:{port}", host, port)
+
+
 def _try_setup_client_authentication(
     args: argparse.Namespace,
     certificates: Optional[Tuple[bytes, bytes, bytes]],
@@ -422,7 +431,7 @@ def _try_obtain_certificates(
         log(WARN, "Option `--insecure` was set. Starting insecure HTTP server.")
         return None
     # Check if certificates are provided
-    if args.fleet_api_type == TRANSPORT_TYPE_GRPC_RERE:
+    if args.fleet_api_type in [TRANSPORT_TYPE_GRPC_RERE, TRANSPORT_TYPE_GRPC_ADAPTER]:
         if args.ssl_certfile and args.ssl_keyfile and args.ssl_ca_certfile:
             if not isfile(args.ssl_ca_certfile):
                 sys.exit("Path argument `--ssl-ca-certfile` does not point to a file.")
@@ -494,6 +503,30 @@ def _run_fleet_api_grpc_rere(
     return fleet_grpc_server
 
 
+def _run_fleet_api_grpc_adapter(
+    address: str,
+    state_factory: StateFactory,
+    certificates: Optional[Tuple[bytes, bytes, bytes]],
+) -> grpc.Server:
+    """Run Fleet API (GrpcAdapter)."""
+    # Create Fleet API gRPC server
+    fleet_servicer = GrpcAdapterServicer(
+        state_factory=state_factory,
+    )
+    fleet_add_servicer_to_server_fn = add_GrpcAdapterServicer_to_server
+    fleet_grpc_server = generic_create_grpc_server(
+        servicer_and_add_fn=(fleet_servicer, fleet_add_servicer_to_server_fn),
+        server_address=address,
+        max_message_length=GRPC_MAX_MESSAGE_LENGTH,
+        certificates=certificates,
+    )
+
+    log(INFO, "Flower ECE: Starting Fleet API (GrpcAdapter) on %s", address)
+    fleet_grpc_server.start()
+
+    return fleet_grpc_server
+
+
 # pylint: disable=import-outside-toplevel,too-many-arguments
 def _run_fleet_api_rest(
     host: str,
@@ -609,7 +642,11 @@ def _add_args_fleet_api(parser: argparse.ArgumentParser) -> None:
         "--fleet-api-type",
         default=TRANSPORT_TYPE_GRPC_RERE,
         type=str,
-        choices=[TRANSPORT_TYPE_GRPC_RERE, TRANSPORT_TYPE_REST],
+        choices=[
+            TRANSPORT_TYPE_GRPC_RERE,
+            TRANSPORT_TYPE_GRPC_ADAPTER,
+            TRANSPORT_TYPE_REST,
+        ],
         help="Start a gRPC-rere or REST (experimental) Fleet API server.",
     )
     parser.add_argument(

flwr/server/compat/app.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/compat/app_utils.py

@@ -91,7 +91,7 @@ def _update_client_manager(
                 node_id=node_id,
                 driver=driver,
                 anonymous=False,
-                run_id=driver.run_id,  # type: ignore
+                run_id=driver.run.run_id,
             )
             if client_manager.register(client_proxy):
                 registered_nodes[node_id] = client_proxy

flwr/server/compat/driver_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/driver/driver.py

@@ -19,11 +19,17 @@ from abc import ABC, abstractmethod
 from typing import Iterable, List, Optional
 
 from flwr.common import Message, RecordSet
+from flwr.common.typing import Run
 
 
 class Driver(ABC):
     """Abstract base Driver class for the Driver API."""
 
+    @property
+    @abstractmethod
+    def run(self) -> Run:
+        """Run information."""
+
     @abstractmethod
     def create_message(  # pylint: disable=too-many-arguments
         self,

flwr/server/driver/grpc_driver.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.
@@ -17,7 +17,7 @@
 import time
 import warnings
 from logging import DEBUG, ERROR, WARNING
-from typing import Iterable, List, Optional, Tuple
+from typing import Iterable, List, Optional, Tuple, cast
 
 import grpc
 
@@ -25,6 +25,7 @@ from flwr.common import DEFAULT_TTL, EventType, Message, Metadata, RecordSet, ev
 from flwr.common.grpc import create_channel
 from flwr.common.logger import log
 from flwr.common.serde import message_from_taskres, message_to_taskins
+from flwr.common.typing import Run
 from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
     CreateRunRequest,
     CreateRunResponse,
@@ -37,6 +38,7 @@ from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
 )
 from flwr.proto.driver_pb2_grpc import DriverStub  # pylint: disable=E0611
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
 
 from .driver import Driver
@@ -46,13 +48,24 @@ DEFAULT_SERVER_ADDRESS_DRIVER = "[::]:9091"
 ERROR_MESSAGE_DRIVER_NOT_CONNECTED = """
 [Driver] Error: Not connected.
 
-Call `connect()` on the `GrpcDriverHelper` instance before calling any of the other
-`GrpcDriverHelper` methods.
+Call `connect()` on the `GrpcDriverStub` instance before calling any of the other
+`GrpcDriverStub` methods.
 """
 
 
-class GrpcDriverHelper:
-    """`GrpcDriverHelper` provides access to the gRPC Driver API/service."""
+class GrpcDriverStub:
+    """`GrpcDriverStub` provides access to the gRPC Driver API/service.
+
+    Parameters
+    ----------
+    driver_service_address : Optional[str]
+        The IPv4 or IPv6 address of the Driver API server.
+        Defaults to `"[::]:9091"`.
+    root_certificates : Optional[bytes] (default: None)
+        The PEM-encoded root certificates as a byte string.
+        If provided, a secure connection using the certificates will be
+        established to an SSL-enabled Flower server.
+    """
 
     def __init__(
         self,
@@ -64,6 +77,10 @@ class GrpcDriverHelper:
         self.channel: Optional[grpc.Channel] = None
         self.stub: Optional[DriverStub] = None
 
+    def is_connected(self) -> bool:
+        """Return True if connected to the Driver API server, otherwise False."""
+        return self.channel is not None
+
     def connect(self) -> None:
         """Connect to the Driver API."""
         event(EventType.DRIVER_CONNECT)
@@ -95,18 +112,29 @@ class GrpcDriverHelper:
         # Check if channel is open
         if self.stub is None:
             log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverHelper` instance not connected")
+            raise ConnectionError("`GrpcDriverStub` instance not connected")
 
         # Call Driver API
         res: CreateRunResponse = self.stub.CreateRun(request=req)
         return res
 
+    def get_run(self, req: GetRunRequest) -> GetRunResponse:
+        """Get run information."""
+        # Check if channel is open
+        if self.stub is None:
+            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
+            raise ConnectionError("`GrpcDriverStub` instance not connected")
+
+        # Call gRPC Driver API
+        res: GetRunResponse = self.stub.GetRun(request=req)
+        return res
+
     def get_nodes(self, req: GetNodesRequest) -> GetNodesResponse:
         """Get client IDs."""
         # Check if channel is open
         if self.stub is None:
             log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverHelper` instance not connected")
+            raise ConnectionError("`GrpcDriverStub` instance not connected")
 
         # Call gRPC Driver API
         res: GetNodesResponse = self.stub.GetNodes(request=req)
@@ -117,7 +145,7 @@ class GrpcDriverHelper:
         # Check if channel is open
         if self.stub is None:
             log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverHelper` instance not connected")
+            raise ConnectionError("`GrpcDriverStub` instance not connected")
 
         # Call gRPC Driver API
         res: PushTaskInsResponse = self.stub.PushTaskIns(request=req)
@@ -128,7 +156,7 @@ class GrpcDriverHelper:
         # Check if channel is open
         if self.stub is None:
             log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverHelper` instance not connected")
+            raise ConnectionError("`GrpcDriverStub` instance not connected")
 
         # Call Driver API
         res: PullTaskResResponse = self.stub.PullTaskRes(request=req)
@@ -140,56 +168,52 @@ class GrpcDriver(Driver):
 
     Parameters
     ----------
-    driver_service_address : Optional[str]
-        The IPv4 or IPv6 address of the Driver API server.
-        Defaults to `"[::]:9091"`.
-    certificates : bytes (default: None)
-        Tuple containing root certificate, server certificate, and private key
-        to start a secure SSL-enabled server. The tuple is expected to have
-        three bytes elements in the following order:
-
-            * CA certificate.
-            * server certificate.
-            * server private key.
-    fab_id : str (default: None)
-        The identifier of the FAB used in the run.
-    fab_version : str (default: None)
-        The version of the FAB used in the run.
+    run_id : int
+        The identifier of the run.
+    stub : Optional[GrpcDriverStub] (default: None)
+        The ``GrpcDriverStub`` instance used to communicate with the SuperLink.
+        If None, an instance connected to "[::]:9091" will be created.
     """
 
-    def __init__(
+    def __init__(  # pylint: disable=too-many-arguments
         self,
-        driver_service_address: str = DEFAULT_SERVER_ADDRESS_DRIVER,
-        root_certificates: Optional[bytes] = None,
-        fab_id: Optional[str] = None,
-        fab_version: Optional[str] = None,
+        run_id: int,
+        stub: Optional[GrpcDriverStub] = None,
     ) -> None:
-        self.addr = driver_service_address
-        self.root_certificates = root_certificates
-        self.driver_helper: Optional[GrpcDriverHelper] = None
-        self.run_id: Optional[int] = None
-        self.fab_id = fab_id if fab_id is not None else ""
-        self.fab_version = fab_version if fab_version is not None else ""
+        self._run_id = run_id
+        self._run: Optional[Run] = None
+        self.stub = stub if stub is not None else GrpcDriverStub()
         self.node = Node(node_id=0, anonymous=True)
 
-    def _get_grpc_driver_helper_and_run_id(self) -> Tuple[GrpcDriverHelper, int]:
-        # Check if the GrpcDriverHelper is initialized
-        if self.driver_helper is None or self.run_id is None:
-            # Connect and create run
-            self.driver_helper = GrpcDriverHelper(
-                driver_service_address=self.addr,
-                root_certificates=self.root_certificates,
+    @property
+    def run(self) -> Run:
+        """Run information."""
+        self._get_stub_and_run_id()
+        return Run(**vars(cast(Run, self._run)))
+
+    def _get_stub_and_run_id(self) -> Tuple[GrpcDriverStub, int]:
+        # Check if is initialized
+        if self._run is None:
+            # Connect
+            if not self.stub.is_connected():
+                self.stub.connect()
+            # Get the run info
+            req = GetRunRequest(run_id=self._run_id)
+            res = self.stub.get_run(req)
+            if not res.HasField("run"):
+                raise RuntimeError(f"Cannot find the run with ID: {self._run_id}")
+            self._run = Run(
+                run_id=res.run.run_id,
+                fab_id=res.run.fab_id,
+                fab_version=res.run.fab_version,
             )
-            self.driver_helper.connect()
-            req = CreateRunRequest(fab_id=self.fab_id, fab_version=self.fab_version)
-            res = self.driver_helper.create_run(req)
-            self.run_id = res.run_id
-        return self.driver_helper, self.run_id
+
+        return self.stub, self._run.run_id
 
     def _check_message(self, message: Message) -> None:
         # Check if the message is valid
         if not (
-            message.metadata.run_id == self.run_id
+            message.metadata.run_id == cast(Run, self._run).run_id
             and message.metadata.src_node_id == self.node.node_id
             and message.metadata.message_id == ""
             and message.metadata.reply_to_message == ""
@@ -210,7 +234,7 @@ class GrpcDriver(Driver):
         This method constructs a new `Message` with given content and metadata.
         The `run_id` and `src_node_id` will be set automatically.
         """
-        _, run_id = self._get_grpc_driver_helper_and_run_id()
+        _, run_id = self._get_stub_and_run_id()
         if ttl:
             warnings.warn(
                 "A custom TTL was set, but note that the SuperLink does not enforce "
@@ -234,9 +258,9 @@ class GrpcDriver(Driver):
 
     def get_node_ids(self) -> List[int]:
         """Get node IDs."""
-        grpc_driver_helper, run_id = self._get_grpc_driver_helper_and_run_id()
-        # Call GrpcDriverHelper method
-        res = grpc_driver_helper.get_nodes(GetNodesRequest(run_id=run_id))
+        stub, run_id = self._get_stub_and_run_id()
+        # Call GrpcDriverStub method
+        res = stub.get_nodes(GetNodesRequest(run_id=run_id))
         return [node.node_id for node in res.nodes]
 
     def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
@@ -245,7 +269,7 @@ class GrpcDriver(Driver):
         This method takes an iterable of messages and sends each message
         to the node specified in `dst_node_id`.
         """
-        grpc_driver_helper, _ = self._get_grpc_driver_helper_and_run_id()
+        stub, _ = self._get_stub_and_run_id()
         # Construct TaskIns
         task_ins_list: List[TaskIns] = []
         for msg in messages:
@@ -255,10 +279,8 @@ class GrpcDriver(Driver):
             taskins = message_to_taskins(msg)
             # Add to list
             task_ins_list.append(taskins)
-        # Call GrpcDriverHelper method
-        res = grpc_driver_helper.push_task_ins(
-            PushTaskInsRequest(task_ins_list=task_ins_list)
-        )
+        # Call GrpcDriverStub method
+        res = stub.push_task_ins(PushTaskInsRequest(task_ins_list=task_ins_list))
         return list(res.task_ids)
 
     def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
@@ -267,9 +289,9 @@ class GrpcDriver(Driver):
         This method is used to collect messages from the SuperLink that correspond to a
         set of given message IDs.
         """
-        grpc_driver, _ = self._get_grpc_driver_helper_and_run_id()
+        stub, _ = self._get_stub_and_run_id()
         # Pull TaskRes
-        res = grpc_driver.pull_task_res(
+        res = stub.pull_task_res(
             PullTaskResRequest(node=self.node, task_ids=message_ids)
         )
         # Convert TaskRes to Message
@@ -308,8 +330,8 @@ class GrpcDriver(Driver):
 
     def close(self) -> None:
         """Disconnect from the SuperLink if connected."""
-        # Check if GrpcDriverHelper is initialized
-        if self.driver_helper is None:
+        # Check if `connect` was called before
+        if not self.stub.is_connected():
             return
         # Disconnect
-        self.driver_helper.disconnect()
+        self.stub.disconnect()

flwr/server/driver/inmemory_driver.py

@@ -17,11 +17,12 @@
 
 import time
 import warnings
-from typing import Iterable, List, Optional
+from typing import Iterable, List, Optional, cast
 from uuid import UUID
 
 from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
 from flwr.common.serde import message_from_taskres, message_to_taskins
+from flwr.common.typing import Run
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.server.superlink.state import StateFactory
 
@@ -33,30 +34,27 @@ class InMemoryDriver(Driver):
 
     Parameters
     ----------
+    run_id : int
+        The identifier of the run.
     state_factory : StateFactory
         A StateFactory embedding a state that this driver can interface with.
-    fab_id : str (default: None)
-        The identifier of the FAB used in the run.
-    fab_version : str (default: None)
-        The version of the FAB used in the run.
     """
 
     def __init__(
         self,
+        run_id: int,
         state_factory: StateFactory,
-        fab_id: Optional[str] = None,
-        fab_version: Optional[str] = None,
     ) -> None:
-        self.run_id: Optional[int] = None
-        self.fab_id = fab_id if fab_id is not None else ""
-        self.fab_version = fab_version if fab_version is not None else ""
-        self.node = Node(node_id=0, anonymous=True)
+        self._run_id = run_id
+        self._run: Optional[Run] = None
         self.state = state_factory.state()
+        self.node = Node(node_id=0, anonymous=True)
 
     def _check_message(self, message: Message) -> None:
+        self._init_run()
         # Check if the message is valid
         if not (
-            message.metadata.run_id == self.run_id
+            message.metadata.run_id == cast(Run, self._run).run_id
             and message.metadata.src_node_id == self.node.node_id
             and message.metadata.message_id == ""
             and message.metadata.reply_to_message == ""
@@ -64,16 +62,20 @@ class InMemoryDriver(Driver):
         ):
             raise ValueError(f"Invalid message: {message}")
 
-    def _get_run_id(self) -> int:
-        """Return run_id.
-
-        If unset, create a new run.
-        """
-        if self.run_id is None:
-            self.run_id = self.state.create_run(
-                fab_id=self.fab_id, fab_version=self.fab_version
-            )
-        return self.run_id
+    def _init_run(self) -> None:
+        """Initialize the run."""
+        if self._run is not None:
+            return
+        run = self.state.get_run(self._run_id)
+        if run is None:
+            raise RuntimeError(f"Cannot find the run with ID: {self._run_id}")
+        self._run = run
+
+    @property
+    def run(self) -> Run:
+        """Run ID."""
+        self._init_run()
+        return Run(**vars(cast(Run, self._run)))
 
     def create_message(  # pylint: disable=too-many-arguments
         self,
@@ -88,7 +90,7 @@ class InMemoryDriver(Driver):
         This method constructs a new `Message` with given content and metadata.
         The `run_id` and `src_node_id` will be set automatically.
         """
-        run_id = self._get_run_id()
+        self._init_run()
         if ttl:
             warnings.warn(
                 "A custom TTL was set, but note that the SuperLink does not enforce "
@@ -99,7 +101,7 @@ class InMemoryDriver(Driver):
         ttl_ = DEFAULT_TTL if ttl is None else ttl
 
         metadata = Metadata(
-            run_id=run_id,
+            run_id=cast(Run, self._run).run_id,
             message_id="",  # Will be set by the server
             src_node_id=self.node.node_id,
             dst_node_id=dst_node_id,
@@ -112,8 +114,8 @@ class InMemoryDriver(Driver):
 
     def get_node_ids(self) -> List[int]:
         """Get node IDs."""
-        run_id = self._get_run_id()
-        return list(self.state.get_nodes(run_id))
+        self._init_run()
+        return list(self.state.get_nodes(cast(Run, self._run).run_id))
 
     def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
         """Push messages to specified node IDs.