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

flwr/cli/new/templates/app/code/flwr_tune/server.py.tpl

@@ -0,0 +1,48 @@
+"""$project_name: A Flower / FlowerTune app."""
+
+from $import_name.client import set_parameters
+from $import_name.models import get_model
+
+
+# Get function that will be executed by the strategy's evaluate() method
+# Here we use it to save global model checkpoints
+def get_evaluate_fn(model_cfg, save_every_round, total_round, save_path):
+    """Return an evaluation function for saving global model."""
+
+    def evaluate(server_round: int, parameters, config):
+        # Save model
+        if server_round != 0 and (
+            server_round == total_round or server_round % save_every_round == 0
+        ):
+            # Init model
+            model = get_model(model_cfg)
+            set_parameters(model, parameters)
+
+            model.save_pretrained(f"{save_path}/peft_{server_round}")
+
+        return 0.0, {}
+
+    return evaluate
+
+
+def get_on_fit_config():
+    """
+    Return a function that will be used to construct the config
+    that the client's fit() method will receive.
+    """
+
+    def fit_config_fn(server_round: int):
+        fit_config = {"current_round": server_round}
+        return fit_config
+
+    return fit_config_fn
+
+
+def fit_weighted_average(metrics):
+    """Aggregate (federated) evaluation metrics."""
+    # Multiply accuracy of each client by number of examples used
+    losses = [num_examples * m["train_loss"] for num_examples, m in metrics]
+    examples = [num_examples for num_examples, _ in metrics]
+
+    # Aggregate and return custom metric (weighted average)
+    return {"train_loss": sum(losses) / sum(examples)}

flwr/cli/new/templates/app/code/flwr_tune/static_config.yaml.tpl

@@ -0,0 +1,11 @@
+# Federated Instruction Tuning (static)
+---
+dataset:
+  name: $dataset_name
+
+# FL experimental settings
+num_clients: $num_clients # total number of clients
+num_rounds: 200
+partitioner:
+  _target_: flwr_datasets.partitioner.IidPartitioner
+  num_partitions: $num_clients

flwr/cli/new/templates/app/pyproject.flowertune.toml.tpl

@@ -0,0 +1,42 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "$package_name"
+version = "1.0.0"
+description = ""
+authors = [
+    { name = "The Flower Authors", email = "hello@flower.ai" },
+]
+license = { text = "Apache License (2.0)" }
+dependencies = [
+    "flwr[simulation]>=1.9.0,<2.0",
+    "flwr-datasets>=0.1.0,<1.0.0",
+    "hydra-core==1.3.2",
+    "trl==0.8.1",
+    "bitsandbytes==0.43.0",
+    "scipy==1.13.0",
+    "peft==0.6.2",
+    "transformers==4.39.3",
+    "sentencepiece==0.2.0",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["."]
+
+[flower]
+publisher = "$username"
+
+[flower.components]
+serverapp = "$import_name.app:server"
+clientapp = "$import_name.app:client"
+
+[flower.engine]
+name = "simulation"
+
+[flower.engine.simulation.supernode]
+num = $num_clients
+
+[flower.engine.simulation]
+backend_config = { client_resources = { num_cpus = 8, num_gpus = 1.0 } }

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
@@ -41,7 +43,10 @@ class Engine(str, Enum):
 def run(
     engine: Annotated[
         Optional[Engine],
-        typer.Option(case_sensitive=False, help="The execution engine to run the app"),
+        typer.Option(
+            case_sensitive=False,
+            help="The engine to run FL with (currently only simulation is supported).",
+        ),
     ] = None,
     use_superexec: Annotated[
         bool,
@@ -49,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)
@@ -87,12 +96,16 @@ def run(
 
     if engine == Engine.SIMULATION:
         num_supernodes = config["flower"]["engine"]["simulation"]["supernode"]["num"]
+        backend_config = config["flower"]["engine"]["simulation"].get(
+            "backend_config", None
+        )
 
         typer.secho("Starting run... ", fg=typer.colors.BLUE)
         _run_simulation(
             server_app_attr=server_app_ref,
             client_app_attr=client_app_ref,
             num_supernodes=num_supernodes,
+            backend_config=backend_config,
         )
     else:
         typer.secho(
@@ -102,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)
@@ -117,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/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/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/client/client_app.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,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/client/dpfedavg_numpy_client.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/client/grpc_rere_client/__init__.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/client/grpc_rere_client/connection.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/client/message_handler/__init__.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/client/message_handler/message_handler.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.
@@ -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/client/mod/__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/client/mod/secure_aggregation/__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/client/mod/utils.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/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/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/client/rest_client/__init__.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/client/rest_client/connection.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/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
@@ -267,7 +268,7 @@ def _parse_args_run_supernode() -> argparse.ArgumentParser:
         "--flwr-dir",
         default=None,
         help="""The path containing installed Flower Apps.
-    By default, this value isequal to:
+    By default, this value is equal to:
 
         - `$FLWR_HOME/` if `$FLWR_HOME` is defined
         - `$XDG_DATA_HOME/.flwr/` if `$XDG_DATA_HOME` is defined
@@ -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/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/common/address.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/common/config.py

@@ -24,14 +24,16 @@ from flwr.cli.config_utils import validate_fields
 from flwr.common.constant import APP_DIR, FAB_CONFIG_FILE, FLWR_HOME
 
 
-def get_flwr_dir() -> Path:
+def get_flwr_dir(provided_path: Optional[str] = None) -> Path:
     """Return the Flower home directory based on env variables."""
-    return Path(
-        os.getenv(
-            FLWR_HOME,
-            f"{os.getenv('XDG_DATA_HOME', os.getenv('HOME'))}/.flwr",
+    if provided_path is None or not Path(provided_path).is_dir():
+        return Path(
+            os.getenv(
+                FLWR_HOME,
+                f"{os.getenv('XDG_DATA_HOME', os.getenv('HOME'))}/.flwr",
+            )
         )
-    )
+    return Path(provided_path).absolute()
 
 
 def get_project_dir(

flwr/common/constant.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.
@@ -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/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