flwr-nightly 1.8.0.dev20240314__py3-none-any.whl → 1.15.0.dev20250114__py3-none-any.whl

flwr/common/serde.py

@@ -15,12 +15,19 @@
 """ProtoBuf serialization and deserialization."""
 
 
-from typing import Any, Dict, List, MutableMapping, OrderedDict, Type, TypeVar, cast
+from collections import OrderedDict
+from collections.abc import MutableMapping
+from typing import Any, TypeVar, cast
 
 from google.protobuf.message import Message as GrpcMessage
 
 # pylint: disable=E0611
+from flwr.proto.clientappio_pb2 import ClientAppOutputCode, ClientAppOutputStatus
 from flwr.proto.error_pb2 import Error as ProtoError
+from flwr.proto.fab_pb2 import Fab as ProtoFab
+from flwr.proto.message_pb2 import Context as ProtoContext
+from flwr.proto.message_pb2 import Message as ProtoMessage
+from flwr.proto.message_pb2 import Metadata as ProtoMetadata
 from flwr.proto.node_pb2 import Node
 from flwr.proto.recordset_pb2 import Array as ProtoArray
 from flwr.proto.recordset_pb2 import BoolList, BytesList
@@ -31,7 +38,9 @@ from flwr.proto.recordset_pb2 import MetricsRecord as ProtoMetricsRecord
 from flwr.proto.recordset_pb2 import MetricsRecordValue as ProtoMetricsRecordValue
 from flwr.proto.recordset_pb2 import ParametersRecord as ProtoParametersRecord
 from flwr.proto.recordset_pb2 import RecordSet as ProtoRecordSet
-from flwr.proto.recordset_pb2 import Sint64List, StringList
+from flwr.proto.recordset_pb2 import SintList, StringList, UintList
+from flwr.proto.run_pb2 import Run as ProtoRun
+from flwr.proto.run_pb2 import RunStatus as ProtoRunStatus
 from flwr.proto.task_pb2 import Task, TaskIns, TaskRes
 from flwr.proto.transport_pb2 import (
     ClientMessage,
@@ -44,7 +53,15 @@ from flwr.proto.transport_pb2 import (
 )
 
 # pylint: enable=E0611
-from . import Array, ConfigsRecord, MetricsRecord, ParametersRecord, RecordSet, typing
+from . import (
+    Array,
+    ConfigsRecord,
+    Context,
+    MetricsRecord,
+    ParametersRecord,
+    RecordSet,
+    typing,
+)
 from .message import Error, Message, Metadata
 from .record.typeddict import TypedDict
 
@@ -58,7 +75,7 @@ def parameters_to_proto(parameters: typing.Parameters) -> Parameters:
 
 def parameters_from_proto(msg: Parameters) -> typing.Parameters:
     """Deserialize `Parameters` from ProtoBuf."""
-    tensors: List[bytes] = list(msg.tensors)
+    tensors: list[bytes] = list(msg.tensors)
     return typing.Parameters(tensors=tensors, tensor_type=msg.tensor_type)
 
 
@@ -324,6 +341,7 @@ def metrics_from_proto(proto: Any) -> typing.Metrics:
 
 
 # === Scalar messages ===
+INT64_MAX_VALUE = 9223372036854775807  # (1 << 63) - 1
 
 
 def scalar_to_proto(scalar: typing.Scalar) -> Scalar:
@@ -338,6 +356,9 @@ def scalar_to_proto(scalar: typing.Scalar) -> Scalar:
         return Scalar(double=scalar)
 
     if isinstance(scalar, int):
+        # Use uint64 for integers larger than the maximum value of sint64
+        if scalar > INT64_MAX_VALUE:
+            return Scalar(uint64=scalar)
         return Scalar(sint64=scalar)
 
     if isinstance(scalar, str):
@@ -358,16 +379,16 @@ def scalar_from_proto(scalar_msg: Scalar) -> typing.Scalar:
 # === Record messages ===
 
 
-_type_to_field = {
+_type_to_field: dict[type, str] = {
     float: "double",
     int: "sint64",
     bool: "bool",
     str: "string",
     bytes: "bytes",
 }
-_list_type_to_class_and_field = {
+_list_type_to_class_and_field: dict[type, tuple[type[GrpcMessage], str]] = {
     float: (DoubleList, "double_list"),
-    int: (Sint64List, "sint64_list"),
+    int: (SintList, "sint_list"),
     bool: (BoolList, "bool_list"),
     str: (StringList, "string_list"),
     bytes: (BytesList, "bytes_list"),
@@ -375,8 +396,13 @@ _list_type_to_class_and_field = {
 T = TypeVar("T")
 
 
+def _is_uint64(value: Any) -> bool:
+    """Check if a value is uint64."""
+    return isinstance(value, int) and value > INT64_MAX_VALUE
+
+
 def _record_value_to_proto(
-    value: Any, allowed_types: List[type], proto_class: Type[T]
+    value: Any, allowed_types: list[type], proto_class: type[T]
 ) -> T:
     """Serialize `*RecordValue` to ProtoBuf.
 
@@ -387,12 +413,18 @@ def _record_value_to_proto(
         # Single element
         # Note: `isinstance(False, int) == True`.
         if isinstance(value, t):
-            arg[_type_to_field[t]] = value
+            fld = _type_to_field[t]
+            if t is int and _is_uint64(value):
+                fld = "uint64"
+            arg[fld] = value
             return proto_class(**arg)
         # List
         if isinstance(value, list) and all(isinstance(item, t) for item in value):
-            list_class, field_name = _list_type_to_class_and_field[t]
-            arg[field_name] = list_class(vals=value)
+            list_class, fld = _list_type_to_class_and_field[t]
+            # Use UintList if any element is of type `uint64`.
+            if t is int and any(_is_uint64(v) for v in value):
+                list_class, fld = UintList, "uint_list"
+            arg[fld] = list_class(vals=value)
             return proto_class(**arg)
     # Invalid types
     raise TypeError(
@@ -413,9 +445,9 @@ def _record_value_from_proto(value_proto: GrpcMessage) -> Any:
 
 def _record_value_dict_to_proto(
     value_dict: TypedDict[str, Any],
-    allowed_types: List[type],
-    value_proto_class: Type[T],
-) -> Dict[str, T]:
+    allowed_types: list[type],
+    value_proto_class: type[T],
+) -> dict[str, T]:
     """Serialize the record value dict to ProtoBuf.
 
     Note: `bool` MUST be put in the front of allowd_types if it exists.
@@ -433,7 +465,7 @@ def _record_value_dict_to_proto(
 
 def _record_value_dict_from_proto(
     value_dict_proto: MutableMapping[str, Any]
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """Deserialize the record value dict from ProtoBuf."""
     return {k: _record_value_from_proto(v) for k, v in value_dict_proto.items()}
 
@@ -484,7 +516,7 @@ def metrics_record_from_proto(record_proto: ProtoMetricsRecord) -> MetricsRecord
     """Deserialize MetricsRecord from ProtoBuf."""
     return MetricsRecord(
         metrics_dict=cast(
-            Dict[str, typing.MetricsRecordValues],
+            dict[str, typing.MetricsRecordValues],
             _record_value_dict_from_proto(record_proto.data),
         ),
         keep_input=False,
@@ -506,7 +538,7 @@ def configs_record_from_proto(record_proto: ProtoConfigsRecord) -> ConfigsRecord
     """Deserialize ConfigsRecord from ProtoBuf."""
     return ConfigsRecord(
         configs_dict=cast(
-            Dict[str, typing.ConfigsRecordValues],
+            dict[str, typing.ConfigsRecordValues],
             _record_value_dict_from_proto(record_proto.data),
         ),
         keep_input=False,
@@ -575,6 +607,7 @@ def message_to_taskins(message: Message) -> TaskIns:
         task=Task(
             producer=Node(node_id=0, anonymous=True),  # Assume driver node
             consumer=Node(node_id=md.dst_node_id, anonymous=False),
+            created_at=md.created_at,
             ttl=md.ttl,
             ancestry=[md.reply_to_message] if md.reply_to_message != "" else [],
             task_type=md.message_type,
@@ -601,7 +634,7 @@ def message_from_taskins(taskins: TaskIns) -> Message:
     )
 
     # Construct Message
-    return Message(
+    message = Message(
         metadata=metadata,
         content=(
             recordset_from_proto(taskins.task.recordset)
@@ -614,6 +647,8 @@ def message_from_taskins(taskins: TaskIns) -> Message:
             else None
         ),
     )
+    message.metadata.created_at = taskins.task.created_at
+    return message
 
 
 def message_to_taskres(message: Message) -> TaskRes:
@@ -626,6 +661,7 @@ def message_to_taskres(message: Message) -> TaskRes:
         task=Task(
             producer=Node(node_id=md.src_node_id, anonymous=False),
             consumer=Node(node_id=0, anonymous=True),  # Assume driver node
+            created_at=md.created_at,
             ttl=md.ttl,
             ancestry=[md.reply_to_message] if md.reply_to_message != "" else [],
             task_type=md.message_type,
@@ -652,7 +688,7 @@ def message_from_taskres(taskres: TaskRes) -> Message:
     )
 
     # Construct the Message
-    return Message(
+    message = Message(
         metadata=metadata,
         content=(
             recordset_from_proto(taskres.task.recordset)
@@ -665,3 +701,246 @@ def message_from_taskres(taskres: TaskRes) -> Message:
             else None
         ),
     )
+    message.metadata.created_at = taskres.task.created_at
+    return message
+
+
+# === FAB ===
+
+
+def fab_to_proto(fab: typing.Fab) -> ProtoFab:
+    """Create a proto Fab object from a Python Fab."""
+    return ProtoFab(hash_str=fab.hash_str, content=fab.content)
+
+
+def fab_from_proto(fab: ProtoFab) -> typing.Fab:
+    """Create a Python Fab object from a proto Fab."""
+    return typing.Fab(fab.hash_str, fab.content)
+
+
+# === User configs ===
+
+
+def user_config_to_proto(user_config: typing.UserConfig) -> Any:
+    """Serialize `UserConfig` to ProtoBuf."""
+    proto = {}
+    for key, value in user_config.items():
+        proto[key] = user_config_value_to_proto(value)
+    return proto
+
+
+def user_config_from_proto(proto: Any) -> typing.UserConfig:
+    """Deserialize `UserConfig` from ProtoBuf."""
+    metrics = {}
+    for key, value in proto.items():
+        metrics[key] = user_config_value_from_proto(value)
+    return metrics
+
+
+def user_config_value_to_proto(user_config_value: typing.UserConfigValue) -> Scalar:
+    """Serialize `UserConfigValue` to ProtoBuf."""
+    if isinstance(user_config_value, bool):
+        return Scalar(bool=user_config_value)
+
+    if isinstance(user_config_value, float):
+        return Scalar(double=user_config_value)
+
+    if isinstance(user_config_value, int):
+        return Scalar(sint64=user_config_value)
+
+    if isinstance(user_config_value, str):
+        return Scalar(string=user_config_value)
+
+    raise ValueError(
+        f"Accepted types: {bool, float, int, str} (but not {type(user_config_value)})"
+    )
+
+
+def user_config_value_from_proto(scalar_msg: Scalar) -> typing.UserConfigValue:
+    """Deserialize `UserConfigValue` from ProtoBuf."""
+    scalar_field = scalar_msg.WhichOneof("scalar")
+    scalar = getattr(scalar_msg, cast(str, scalar_field))
+    return cast(typing.UserConfigValue, scalar)
+
+
+# === Metadata messages ===
+
+
+def metadata_to_proto(metadata: Metadata) -> ProtoMetadata:
+    """Serialize `Metadata` to ProtoBuf."""
+    proto = ProtoMetadata(  # pylint: disable=E1101
+        run_id=metadata.run_id,
+        message_id=metadata.message_id,
+        src_node_id=metadata.src_node_id,
+        dst_node_id=metadata.dst_node_id,
+        reply_to_message=metadata.reply_to_message,
+        group_id=metadata.group_id,
+        ttl=metadata.ttl,
+        message_type=metadata.message_type,
+        created_at=metadata.created_at,
+    )
+    return proto
+
+
+def metadata_from_proto(metadata_proto: ProtoMetadata) -> Metadata:
+    """Deserialize `Metadata` from ProtoBuf."""
+    metadata = Metadata(
+        run_id=metadata_proto.run_id,
+        message_id=metadata_proto.message_id,
+        src_node_id=metadata_proto.src_node_id,
+        dst_node_id=metadata_proto.dst_node_id,
+        reply_to_message=metadata_proto.reply_to_message,
+        group_id=metadata_proto.group_id,
+        ttl=metadata_proto.ttl,
+        message_type=metadata_proto.message_type,
+    )
+    return metadata
+
+
+# === Message messages ===
+
+
+def message_to_proto(message: Message) -> ProtoMessage:
+    """Serialize `Message` to ProtoBuf."""
+    proto = ProtoMessage(
+        metadata=metadata_to_proto(message.metadata),
+        content=(
+            recordset_to_proto(message.content) if message.has_content() else None
+        ),
+        error=error_to_proto(message.error) if message.has_error() else None,
+    )
+    return proto
+
+
+def message_from_proto(message_proto: ProtoMessage) -> Message:
+    """Deserialize `Message` from ProtoBuf."""
+    created_at = message_proto.metadata.created_at
+    message = Message(
+        metadata=metadata_from_proto(message_proto.metadata),
+        content=(
+            recordset_from_proto(message_proto.content)
+            if message_proto.HasField("content")
+            else None
+        ),
+        error=(
+            error_from_proto(message_proto.error)
+            if message_proto.HasField("error")
+            else None
+        ),
+    )
+    # `.created_at` is set upon Message object construction
+    # we need to manually set it to the original value
+    message.metadata.created_at = created_at
+    return message
+
+
+# === Context messages ===
+
+
+def context_to_proto(context: Context) -> ProtoContext:
+    """Serialize `Context` to ProtoBuf."""
+    proto = ProtoContext(
+        run_id=context.run_id,
+        node_id=context.node_id,
+        node_config=user_config_to_proto(context.node_config),
+        state=recordset_to_proto(context.state),
+        run_config=user_config_to_proto(context.run_config),
+    )
+    return proto
+
+
+def context_from_proto(context_proto: ProtoContext) -> Context:
+    """Deserialize `Context` from ProtoBuf."""
+    context = Context(
+        run_id=context_proto.run_id,
+        node_id=context_proto.node_id,
+        node_config=user_config_from_proto(context_proto.node_config),
+        state=recordset_from_proto(context_proto.state),
+        run_config=user_config_from_proto(context_proto.run_config),
+    )
+    return context
+
+
+# === Run messages ===
+
+
+def run_to_proto(run: typing.Run) -> ProtoRun:
+    """Serialize `Run` to ProtoBuf."""
+    proto = ProtoRun(
+        run_id=run.run_id,
+        fab_id=run.fab_id,
+        fab_version=run.fab_version,
+        fab_hash=run.fab_hash,
+        override_config=user_config_to_proto(run.override_config),
+        pending_at=run.pending_at,
+        starting_at=run.starting_at,
+        running_at=run.running_at,
+        finished_at=run.finished_at,
+        status=run_status_to_proto(run.status),
+    )
+    return proto
+
+
+def run_from_proto(run_proto: ProtoRun) -> typing.Run:
+    """Deserialize `Run` from ProtoBuf."""
+    run = typing.Run(
+        run_id=run_proto.run_id,
+        fab_id=run_proto.fab_id,
+        fab_version=run_proto.fab_version,
+        fab_hash=run_proto.fab_hash,
+        override_config=user_config_from_proto(run_proto.override_config),
+        pending_at=run_proto.pending_at,
+        starting_at=run_proto.starting_at,
+        running_at=run_proto.running_at,
+        finished_at=run_proto.finished_at,
+        status=run_status_from_proto(run_proto.status),
+    )
+    return run
+
+
+# === ClientApp status messages ===
+
+
+def clientappstatus_to_proto(
+    status: typing.ClientAppOutputStatus,
+) -> ClientAppOutputStatus:
+    """Serialize `ClientAppOutputStatus` to ProtoBuf."""
+    code = ClientAppOutputCode.SUCCESS
+    if status.code == typing.ClientAppOutputCode.DEADLINE_EXCEEDED:
+        code = ClientAppOutputCode.DEADLINE_EXCEEDED
+    if status.code == typing.ClientAppOutputCode.UNKNOWN_ERROR:
+        code = ClientAppOutputCode.UNKNOWN_ERROR
+    return ClientAppOutputStatus(code=code, message=status.message)
+
+
+def clientappstatus_from_proto(
+    msg: ClientAppOutputStatus,
+) -> typing.ClientAppOutputStatus:
+    """Deserialize `ClientAppOutputStatus` from ProtoBuf."""
+    code = typing.ClientAppOutputCode.SUCCESS
+    if msg.code == ClientAppOutputCode.DEADLINE_EXCEEDED:
+        code = typing.ClientAppOutputCode.DEADLINE_EXCEEDED
+    if msg.code == ClientAppOutputCode.UNKNOWN_ERROR:
+        code = typing.ClientAppOutputCode.UNKNOWN_ERROR
+    return typing.ClientAppOutputStatus(code=code, message=msg.message)
+
+
+# === Run status ===
+
+
+def run_status_to_proto(run_status: typing.RunStatus) -> ProtoRunStatus:
+    """Serialize `RunStatus` to ProtoBuf."""
+    return ProtoRunStatus(
+        status=run_status.status,
+        sub_status=run_status.sub_status,
+        details=run_status.details,
+    )
+
+
+def run_status_from_proto(run_status_proto: ProtoRunStatus) -> typing.RunStatus:
+    """Deserialize `RunStatus` from ProtoBuf."""
+    return typing.RunStatus(
+        status=run_status_proto.status,
+        sub_status=run_status_proto.sub_status,
+        details=run_status_proto.details,
+    )

flwr/common/telemetry.py

@@ -25,8 +25,9 @@ import uuid
 from concurrent.futures import Future, ThreadPoolExecutor
 from enum import Enum, auto
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Optional, Union, cast
 
+from flwr.common.constant import FLWR_DIR
 from flwr.common.version import package_name, package_version
 
 FLWR_TELEMETRY_ENABLED = os.getenv("FLWR_TELEMETRY_ENABLED", "1")
@@ -64,6 +65,18 @@ def _get_home() -> Path:
     return Path().home()
 
 
+def _get_partner_id() -> str:
+    """Get partner ID."""
+    partner_id = os.getenv("FLWR_TELEMETRY_PARTNER_ID")
+    if not partner_id:
+        return "unavailable"
+    try:
+        uuid.UUID(partner_id)
+    except ValueError:
+        partner_id = "invalid"
+    return partner_id
+
+
 def _get_source_id() -> str:
     """Get existing or new source ID."""
     source_id = "unavailable"
@@ -74,7 +87,7 @@ def _get_source_id() -> str:
         # If the home directory can’t be resolved, RuntimeError is raised.
         return source_id
 
-    flwr_dir = home.joinpath(".flwr")
+    flwr_dir = home.joinpath(FLWR_DIR)
     # Create .flwr directory if it does not exist yet.
     try:
         flwr_dir.mkdir(parents=True, exist_ok=True)
@@ -114,67 +127,86 @@ class EventType(str, Enum):
     # The type signature is not compatible with mypy, pylint and flake8
     # so each of those needs to be disabled for this line.
     # pylint: disable-next=no-self-argument,arguments-differ,line-too-long
-    def _generate_next_value_(name: str, start: int, count: int, last_values: List[Any]) -> Any:  # type: ignore # noqa: E501
+    def _generate_next_value_(name: str, start: int, count: int, last_values: list[Any]) -> Any:  # type: ignore # noqa: E501
         return name
 
     # Ping
     PING = auto()
 
-    # Client: start_client
+    # --- LEGACY FUNCTIONS -------------------------------------------------------------
+
+    # Legacy: `start_client` function
     START_CLIENT_ENTER = auto()
     START_CLIENT_LEAVE = auto()
 
-    # Server: start_server
+    # Legacy: `start_server` function
     START_SERVER_ENTER = auto()
     START_SERVER_LEAVE = auto()
 
-    # Driver API
-    RUN_DRIVER_API_ENTER = auto()
-    RUN_DRIVER_API_LEAVE = auto()
+    # Legacy: `start_simulation` function
+    START_SIMULATION_ENTER = auto()
+    START_SIMULATION_LEAVE = auto()
+
+    # --- `flwr` CLI -------------------------------------------------------------------
+
+    # Not yet implemented
+
+    # --- `flwr-*` commands ------------------------------------------------------------
+
+    # CLI: flwr-simulation
+    FLWR_SIMULATION_RUN_ENTER = auto()
+    FLWR_SIMULATION_RUN_LEAVE = auto()
+
+    # CLI: flwr-serverapp
+    FLWR_SERVERAPP_RUN_ENTER = auto()
+    FLWR_SERVERAPP_RUN_LEAVE = auto()
 
-    # Fleet API
-    RUN_FLEET_API_ENTER = auto()
-    RUN_FLEET_API_LEAVE = auto()
+    # --- Simulation Engine ------------------------------------------------------------
 
-    # Driver API and Fleet API
+    # CLI: flower-simulation
+    CLI_FLOWER_SIMULATION_ENTER = auto()
+    CLI_FLOWER_SIMULATION_LEAVE = auto()
+
+    # Python API: `run_simulation`
+    PYTHON_API_RUN_SIMULATION_ENTER = auto()
+    PYTHON_API_RUN_SIMULATION_LEAVE = auto()
+
+    # --- Deployment Engine ------------------------------------------------------------
+
+    # CLI: `flower-superlink`
     RUN_SUPERLINK_ENTER = auto()
     RUN_SUPERLINK_LEAVE = auto()
 
-    # Simulation
-    START_SIMULATION_ENTER = auto()
-    START_SIMULATION_LEAVE = auto()
+    # CLI: `flower-supernode`
+    RUN_SUPERNODE_ENTER = auto()
+    RUN_SUPERNODE_LEAVE = auto()
 
-    # Driver: Driver
-    DRIVER_CONNECT = auto()
-    DRIVER_DISCONNECT = auto()
+    # --- DEPRECATED -------------------------------------------------------------------
 
-    # Driver: start_driver
-    START_DRIVER_ENTER = auto()
-    START_DRIVER_LEAVE = auto()
+    # [DEPRECATED] CLI: `flower-server-app`
+    RUN_SERVER_APP_ENTER = auto()
+    RUN_SERVER_APP_LEAVE = auto()
 
-    # flower-client-app
+    # [DEPRECATED] CLI: `flower-client-app`
     RUN_CLIENT_APP_ENTER = auto()
     RUN_CLIENT_APP_LEAVE = auto()
 
-    # flower-server-app
-    RUN_SERVER_APP_ENTER = auto()
-    RUN_SERVER_APP_LEAVE = auto()
-
 
 # Use the ThreadPoolExecutor with max_workers=1 to have a queue
 # and also ensure that telemetry calls are not blocking.
-state: Dict[str, Union[Optional[str], Optional[ThreadPoolExecutor]]] = {
+state: dict[str, Union[Optional[str], Optional[ThreadPoolExecutor]]] = {
     # Will be assigned ThreadPoolExecutor(max_workers=1)
     # in event() the first time it's required
     "executor": None,
     "source": None,
     "cluster": None,
+    "partner": None,
 }
 
 
 def event(
     event_type: EventType,
-    event_details: Optional[Dict[str, Any]] = None,
+    event_details: Optional[dict[str, Any]] = None,
 ) -> Future:  # type: ignore
     """Submit create_event to ThreadPoolExecutor to avoid blocking."""
     if state["executor"] is None:
@@ -186,7 +218,7 @@ def event(
     return result
 
 
-def create_event(event_type: EventType, event_details: Optional[Dict[str, Any]]) -> str:
+def create_event(event_type: EventType, event_details: Optional[dict[str, Any]]) -> str:
     """Create telemetry event."""
     if state["source"] is None:
         state["source"] = _get_source_id()
@@ -194,11 +226,15 @@ def create_event(event_type: EventType, event_details: Optional[Dict[str, Any]])
     if state["cluster"] is None:
         state["cluster"] = str(uuid.uuid4())
 
+    if state["partner"] is None:
+        state["partner"] = _get_partner_id()
+
     if event_details is None:
         event_details = {}
 
     date = datetime.datetime.now(tz=datetime.timezone.utc).isoformat()
     context = {
+        "partner": state["partner"],
         "source": state["source"],
         "cluster": state["cluster"],
         "date": date,