flwr 1.24.0__py3-none-any.whl → 1.25.0__py3-none-any.whl

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

@@ -127,7 +127,7 @@ def send_node_heartbeat(
     return SendNodeHeartbeatResponse(success=res)
 
 
-def pull_messages(
+def pull_messages(  # pylint: disable=too-many-locals
     request: PullMessagesRequest,
     state: LinkState,
     store: ObjectStore,
@@ -143,6 +143,8 @@ def pull_messages(
     # Convert to Messages
     msg_proto = []
     trees = []
+    run_id_to_record: int | None = None
+
     for msg in message_list:
         try:
             # Retrieve Message object tree from ObjectStore
@@ -152,12 +154,30 @@ def pull_messages(
             # Add Message and its object tree to the response
             msg_proto.append(message_to_proto(msg))
             trees.append(obj_tree)
+
+            # Track run_id for traffic recording
+            run_id_to_record = msg.metadata.run_id
+
         except NoObjectInStoreError as e:
             log(ERROR, e.message)
             # Delete message ins from state
             state.delete_messages(message_ins_ids={msg_object_id})
 
-    return PullMessagesResponse(messages_list=msg_proto, message_object_trees=trees)
+    response = PullMessagesResponse(messages_list=msg_proto, message_object_trees=trees)
+
+    # Record incoming traffic size
+    bytes_recv = len(request.SerializeToString())
+
+    # Record traffic only if message was successfully processed
+    # All messages in this request are assumed to belong to the same run
+    if run_id_to_record is not None:
+        # Record outgoing traffic size
+        bytes_sent = len(response.SerializeToString())
+        state.store_traffic(
+            run_id_to_record, bytes_sent=bytes_sent, bytes_recv=bytes_recv
+        )
+
+    return response
 
 
 def push_messages(
@@ -170,6 +190,9 @@ def push_messages(
     msg = message_from_proto(message_proto=request.messages_list[0])
     run_id = msg.metadata.run_id
 
+    # Record incoming traffic size
+    bytes_recv = len(request.SerializeToString())
+
     # Abort if the run is not running
     abort_msg = check_abort(
         run_id,
@@ -193,6 +216,16 @@ def push_messages(
         results={str(message_id): 0},
         objects_to_push=objects_to_push,
     )
+
+    # Record outgoing traffic size
+    bytes_sent = len(response.SerializeToString())
+
+    # Record traffic only if message was successfully processed
+    # Only one message is processed per request
+    state.store_traffic(run_id, bytes_sent=bytes_sent, bytes_recv=bytes_recv)
+    if request.clientapp_runtime_list:
+        state.add_clientapp_runtime(run_id, request.clientapp_runtime_list[0])
+
     return response
 
 
@@ -257,6 +290,10 @@ def push_object(
     try:
         store.put(request.object_id, request.object_content)
         stored = True
+        # Record bytes traffic pushed from SuperNode
+        state.store_traffic(
+            request.run_id, bytes_sent=0, bytes_recv=len(request.object_content)
+        )
     except (NoObjectInStoreError, ValueError) as e:
         log(ERROR, str(e))
     except UnexpectedObjectContentError as e:
@@ -283,6 +320,8 @@ def pull_object(
     content = store.get(request.object_id)
     if content is not None:
         object_available = content != b""
+        # Record bytes traffic pulled by SuperNode
+        state.store_traffic(request.run_id, bytes_sent=len(content), bytes_recv=0)
         return PullObjectResponse(
             object_found=True,
             object_available=object_available,

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -576,6 +576,9 @@ class InMemoryLinkState(LinkState, InMemoryCoreState):  # pylint: disable=R0902,
                         ),
                         flwr_aid=flwr_aid if flwr_aid else "",
                         federation=federation,
+                        bytes_sent=0,
+                        bytes_recv=0,
+                        clientapp_runtime=0.0,
                     ),
                 )
                 self.run_ids[run_id] = run_record
@@ -771,3 +774,34 @@ class InMemoryLinkState(LinkState, InMemoryCoreState):  # pylint: disable=R0902,
             index = bisect_right(run.logs, (after_timestamp, ""))
             latest_timestamp = run.logs[-1][0] if index < len(run.logs) else 0.0
             return "".join(log for _, log in run.logs[index:]), latest_timestamp
+
+    def store_traffic(self, run_id: int, *, bytes_sent: int, bytes_recv: int) -> None:
+        """Store traffic data for the specified `run_id`."""
+        # Validate non-negative values
+        if bytes_sent < 0 or bytes_recv < 0:
+            raise ValueError(
+                f"Negative traffic values for run {run_id}: "
+                f"bytes_sent={bytes_sent}, bytes_recv={bytes_recv}"
+            )
+
+        if bytes_sent == 0 and bytes_recv == 0:
+            raise ValueError(
+                f"Both bytes_sent and bytes_recv cannot be zero for run {run_id}"
+            )
+
+        with self.lock:
+            if run_id not in self.run_ids:
+                raise ValueError(f"Run {run_id} not found")
+            run_record = self.run_ids[run_id]
+
+        with run_record.lock:
+            run = run_record.run
+            run.bytes_sent += bytes_sent
+            run.bytes_recv += bytes_recv
+
+    def add_clientapp_runtime(self, run_id: int, runtime: float) -> None:
+        """Add ClientApp runtime to the cumulative total for the specified `run_id`."""
+        with self.lock:
+            if run_id not in self.run_ids:
+                raise ValueError(f"Run {run_id} not found")
+            self.run_ids[run_id].run.clientapp_runtime += runtime

flwr/server/superlink/linkstate/linkstate.py

@@ -480,3 +480,35 @@ class LinkState(CoreState):  # pylint: disable=R0904
             - The timestamp of the latest log entry in the returned logs.
               Returns `0` if no logs are returned.
         """
+
+    @abc.abstractmethod
+    def store_traffic(self, run_id: int, *, bytes_sent: int, bytes_recv: int) -> None:
+        """Store traffic data for the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run for which to store traffic data.
+        bytes_sent : int
+            The number of bytes pulled by SuperNodes from the SuperLink to add to the
+            run's total.
+        bytes_recv : int
+            The number of bytes received by SuperLink from SuperNodes to add to the
+            run's total.
+        """
+
+    @abc.abstractmethod
+    def add_clientapp_runtime(self, run_id: int, runtime: float) -> None:
+        """Add ClientApp runtime to the cumulative total for the specified `run_id`.
+
+        This method accumulates the runtime by adding the provided value to the
+        existing total runtime for the run. Multiple ClientApps can contribute
+        to the same run's total runtime.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run for which to store each ClientApp's runtime.
+        runtime : float
+            The runtime in seconds to add to the `run_id`'s cumulative total.
+        """

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -119,7 +119,10 @@ CREATE TABLE IF NOT EXISTS run(
     details               TEXT,
     federation            TEXT,
     federation_options    BLOB,
-    flwr_aid              TEXT
+    flwr_aid              TEXT,
+    bytes_sent            INTEGER DEFAULT 0,
+    bytes_recv            INTEGER DEFAULT 0,
+    clientapp_runtime     REAL DEFAULT 0.0
 );
 """
 
@@ -905,8 +908,8 @@ class SqliteLinkState(LinkState, SqliteCoreState):  # pylint: disable=R0904
                     (run_id, fab_id, fab_version,
                     fab_hash, override_config, federation, federation_options,
                     pending_at, starting_at, running_at, finished_at, sub_status,
-                    details, flwr_aid)
-                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
+                    details, flwr_aid, bytes_sent, bytes_recv, clientapp_runtime)
+                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
                 """
                 override_config_json = json.dumps(override_config)
                 data = [
@@ -924,6 +927,9 @@ class SqliteLinkState(LinkState, SqliteCoreState):  # pylint: disable=R0904
                     "",  # sub_status
                     "",  # details
                     flwr_aid or "",  # flwr_aid
+                    0,  # bytes_sent
+                    0,  # bytes_recv
+                    0,  # clientapp_runtime
                 ]
                 self.conn.execute(query, tuple(data))
                 return uint64_run_id
@@ -972,6 +978,9 @@ class SqliteLinkState(LinkState, SqliteCoreState):  # pylint: disable=R0904
                 ),
                 flwr_aid=row["flwr_aid"],
                 federation=row["federation"],
+                bytes_sent=row["bytes_sent"],
+                bytes_recv=row["bytes_recv"],
+                clientapp_runtime=row["clientapp_runtime"],
             )
         log(ERROR, "`run_id` does not exist.")
         return None
@@ -1255,6 +1264,54 @@ class SqliteLinkState(LinkState, SqliteCoreState):  # pylint: disable=R0904
 
         return rows[0]
 
+    def store_traffic(self, run_id: int, *, bytes_sent: int, bytes_recv: int) -> None:
+        """Store traffic data for the specified `run_id`."""
+        # Validate non-negative values
+        if bytes_sent < 0 or bytes_recv < 0:
+            raise ValueError(
+                f"Negative traffic values for run {run_id}: "
+                f"bytes_sent={bytes_sent}, bytes_recv={bytes_recv}"
+            )
+
+        if bytes_sent == 0 and bytes_recv == 0:
+            raise ValueError(
+                f"Both bytes_sent and bytes_recv cannot be zero for run {run_id}"
+            )
+
+        sint64_run_id = uint64_to_int64(run_id)
+
+        with self.conn:
+            # Check if run exists, performing the update only if it does
+            update_query = """
+                UPDATE run
+                SET bytes_sent = bytes_sent + ?,
+                    bytes_recv = bytes_recv + ?
+                WHERE run_id = ?
+                RETURNING run_id;
+            """
+            rows = self.conn.execute(
+                update_query, (bytes_sent, bytes_recv, sint64_run_id)
+            ).fetchall()
+
+            if not rows:
+                raise ValueError(f"Run {run_id} not found")
+
+    def add_clientapp_runtime(self, run_id: int, runtime: float) -> None:
+        """Add ClientApp runtime to the cumulative total for the specified `run_id`."""
+        sint64_run_id = uint64_to_int64(run_id)
+        with self.conn:
+            # Check if run exists, performing the update only if it does
+            update_query = """
+                UPDATE run
+                SET clientapp_runtime = clientapp_runtime + ?
+                WHERE run_id = ?
+                RETURNING run_id;
+            """
+            rows = self.conn.execute(update_query, (runtime, sint64_run_id)).fetchall()
+
+            if not rows:
+                raise ValueError(f"Run {run_id} not found")
+
 
 def message_to_dict(message: Message) -> dict[str, Any]:
     """Transform Message to dict."""

flwr/supercore/constant.py

@@ -57,6 +57,9 @@ NOOP_FEDERATION = "default"
 # Constants for exit handling
 FORCE_EXIT_TIMEOUT_SECONDS = 5  # Used in `flwr_exit` function
 
+# Constants for message processing timing
+MESSAGE_TIME_ENTRY_MAX_AGE_SECONDS = 3600
+
 
 class NodeStatus:
     """Event log writer types."""

flwr/supercore/utils.py

@@ -15,6 +15,16 @@
 """Utility functions for the infrastructure."""
 
 
+import json
+import re
+
+import requests
+
+from flwr.common.version import package_version as flwr_version
+
+from .constant import APP_ID_PATTERN, APP_VERSION_PATTERN
+
+
 def mask_string(value: str, head: int = 4, tail: int = 4) -> str:
     """Mask a string by preserving only the head and tail characters.
 
@@ -50,3 +60,183 @@ def int64_to_uint64(signed: int) -> int:
     if signed < 0:
         return signed + (1 << 64)
     return signed
+
+
+def parse_app_spec(app_spec: str) -> tuple[str, str | None]:
+    """Parse app specification string into app ID and version.
+
+    Parameters
+    ----------
+    app_spec : str
+        The app specification string in the format '@account/app' or
+        '@account/app==x.y.z' (digits only).
+
+    Returns
+    -------
+    tuple[str, str | None]
+        A tuple containing the app ID and optional version.
+
+    Raises
+    ------
+    ValueError
+        If the app specification format is invalid.
+    """
+    if "==" in app_spec:
+        app_id, app_version = app_spec.split("==", 1)
+
+        if not re.match(APP_VERSION_PATTERN, app_version):
+            raise ValueError(
+                "Invalid app version. Expected format: x.y.z (digits only)."
+            )
+    else:
+        app_id = app_spec
+        app_version = None
+
+    if not re.match(APP_ID_PATTERN, app_id):
+        raise ValueError(
+            "Invalid remote app ID. Expected format: '@account_name/app_name'."
+        )
+
+    return app_id, app_version
+
+
+def request_download_link(
+    app_id: str, app_version: str | None, in_url: str, out_url: str
+) -> tuple[str, list[dict[str, str]] | None]:
+    """Request a download link for the given app from the Flower Platform API.
+
+    Parameters
+    ----------
+    app_id : str
+        The application identifier in the format '@account/app'.
+    app_version : str | None
+        The application version (e.g., '1.2.3'), or None to request the latest version.
+    in_url : str
+        The Platform API endpoint URL to query.
+    out_url : str
+        The key name in the response that contains the download URL.
+
+    Returns
+    -------
+    tuple[str, list[dict[str, str]] | None]
+        A tuple containing:
+        - The download URL for the application.
+        - A list of verification dictionaries if provided by the API, otherwise None.
+
+    Raises
+    ------
+    ValueError
+        If the API connection fails, the application or version is not found,
+        the API returns a non-200 response, or the response format is invalid.
+    """
+    headers = {
+        "Content-Type": "application/json",
+        "Accept": "application/json",
+    }
+    body = {
+        "app_id": app_id,  # send raw string of app_id
+        "app_version": app_version,
+        "flwr_version": flwr_version,
+    }
+
+    try:
+        resp = requests.post(in_url, headers=headers, data=json.dumps(body), timeout=20)
+    except requests.RequestException as e:
+        raise ValueError(f"Unable to connect to Platform API: {e}") from e
+
+    if resp.status_code == 404:
+        # Expecting a JSON body with a "detail" field
+        try:
+            error_message = resp.json().get("detail")
+        except ValueError:
+            # JSON parsing failed
+            raise ValueError(f"{app_id} not found in Platform API.") from None
+
+        if isinstance(error_message, dict):
+            available_app_versions = error_message.get("available_app_versions", [])
+            available_versions_str = (
+                ", ".join(map(str, available_app_versions))
+                if available_app_versions
+                else "None"
+            )
+            raise ValueError(
+                f"{app_id}=={app_version} not found in Platform API. "
+                f"Available app versions for {app_id}: {available_versions_str}"
+            )
+
+        raise ValueError(f"{app_id} not found in Platform API.")
+
+    if not resp.ok:
+        raise ValueError(
+            f"Platform API request failed with status {resp.status_code}. "
+            f"Details: {resp.text}"
+        )
+
+    data = resp.json()
+    if out_url not in data:
+        raise ValueError("Invalid response from Platform API")
+
+    verifications = data["verifications"] if "verifications" in data else None
+
+    return str(data[out_url]), verifications
+
+
+def humanize_duration(seconds: float) -> str:
+    """Convert a duration in seconds to a human-friendly string.
+
+    Rules:
+      - < 90 seconds: show seconds
+      - < 1 hour: show minutes + seconds
+      - < 1 day: show hours + minutes
+      - >= 1 day: show days + hours
+    """
+    seconds = int(seconds)
+
+    # Under 90 seconds → Seconds only
+    if seconds < 90:
+        return f"{seconds}s"
+
+    # Under 1 hour → Minutes and seconds
+    minutes, sec = divmod(seconds, 60)
+    if minutes < 60:
+        return f"{minutes}m {sec}s"
+
+    # Under 1 day → Hours and minutes
+    hours, minutes = divmod(minutes, 60)
+    if hours < 24:
+        return f"{hours}h {minutes}m"
+
+    # 1+ days → Days and hours
+    days, hours = divmod(hours, 24)
+    return f"{days}d {hours}h"
+
+
+def humanize_bytes(num_bytes: int) -> str:
+    """Convert a number of bytes to a human-friendly string.
+
+    Uses 1024-based units and 0-1 decimal precision.
+    Rules:
+      - < 1 KB: bytes
+      - < 1 MB: KB
+      - < 1 GB: MB
+      - < 1 TB: GB
+    """
+    value = float(num_bytes)
+
+    for suffix in ["B", "KB", "MB", "GB", "TB"]:
+        if value < 1024 or suffix == "TB":
+            # Bytes → no decimals
+            if suffix == "B":
+                return f"{int(value)} B"
+
+            # Decide precision: 1 decimal for <10, otherwise no decimal
+            if value < 10:
+                formatted = f"{value:.1f}"
+            else:
+                formatted = f"{int(value)}"
+
+            return f"{formatted} {suffix}"
+
+        value /= 1024
+
+    raise RuntimeError("Unreachable code")  # Make mypy happy

flwr/superlink/servicer/control/control_grpc.py

@@ -61,6 +61,7 @@ def run_control_api_grpc(
     authz_plugin: ControlAuthzPlugin,
     event_log_plugin: EventLogWriterPlugin | None = None,
     artifact_provider: ArtifactProvider | None = None,
+    fleet_api_type: str | None = None,
 ) -> grpc.Server:
     """Run Control API (gRPC, request-response)."""
     license_plugin: LicensePlugin | None = get_license_plugin()
@@ -74,6 +75,7 @@ def run_control_api_grpc(
         is_simulation=is_simulation,
         authn_plugin=authn_plugin,
         artifact_provider=artifact_provider,
+        fleet_api_type=fleet_api_type,
     )
     interceptors = [ControlAccountAuthInterceptor(authn_plugin, authz_plugin)]
     if license_plugin is not None:

flwr/superlink/servicer/control/control_servicer.py

@@ -16,12 +16,14 @@
 
 
 import hashlib
+import json
 import time
 from collections.abc import Generator, Sequence
 from logging import ERROR, INFO
 from typing import Any, cast
 
 import grpc
+import requests
 
 from flwr.cli.config_utils import get_fab_metadata
 from flwr.common import Context, RecordDict, now
@@ -36,6 +38,7 @@ from flwr.common.constant import (
     PUBLIC_KEY_NOT_VALID,
     PULL_UNFINISHED_RUN_MESSAGE,
     RUN_ID_NOT_FOUND_MESSAGE,
+    TRANSPORT_TYPE_GRPC_ADAPTER,
     Status,
     SubStatus,
 )
@@ -76,9 +79,11 @@ from flwr.proto.control_pb2 import (  # pylint: disable=E0611
 from flwr.proto.federation_pb2 import Federation  # pylint: disable=E0611
 from flwr.proto.node_pb2 import NodeInfo  # pylint: disable=E0611
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
+from flwr.supercore.constant import PLATFORM_API_URL
 from flwr.supercore.ffs import FfsFactory
 from flwr.supercore.object_store import ObjectStore, ObjectStoreFactory
 from flwr.supercore.primitives.asymmetric import bytes_to_public_key, uses_nist_ec_curve
+from flwr.supercore.utils import parse_app_spec, request_download_link
 from flwr.superlink.artifact_provider import ArtifactProvider
 from flwr.superlink.auth_plugin import ControlAuthnPlugin
 
@@ -96,6 +101,7 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
         is_simulation: bool,
         authn_plugin: ControlAuthnPlugin,
         artifact_provider: ArtifactProvider | None = None,
+        fleet_api_type: str | None = None,
     ) -> None:
         self.linkstate_factory = linkstate_factory
         self.ffs_factory = ffs_factory
@@ -103,8 +109,9 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
         self.is_simulation = is_simulation
         self.authn_plugin = authn_plugin
         self.artifact_provider = artifact_provider
+        self.fleet_api_type = fleet_api_type
 
-    def StartRun(  # pylint: disable=too-many-locals
+    def StartRun(  # pylint: disable=too-many-locals, too-many-branches, too-many-statements
         self, request: StartRunRequest, context: grpc.ServicerContext
     ) -> StartRunResponse:
         """Create run ID."""
@@ -112,7 +119,15 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
         state = self.linkstate_factory.state()
         ffs = self.ffs_factory.ffs()
 
-        if len(request.fab.content) > FAB_MAX_SIZE:
+        verification_dict: dict[str, str] = {}
+        if request.app_spec:
+            fab_file, verification_dict = _get_remote_fab(
+                self.fleet_api_type, request.app_spec, context
+            )
+        else:
+            fab_file = request.fab.content
+
+        if len(fab_file) > FAB_MAX_SIZE:
             log(
                 ERROR,
                 "FAB size exceeds maximum allowed size of %d bytes.",
@@ -124,7 +139,6 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
         flwr_aid = _check_flwr_aid_exists(flwr_aid, context)
         override_config = user_config_from_proto(request.override_config)
         federation_options = config_record_from_proto(request.federation_options)
-        fab_file = request.fab.content
 
         try:
             # Check that num-supernodes is set
@@ -150,9 +164,10 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
             fab = Fab(
                 hashlib.sha256(fab_file).hexdigest(),
                 fab_file,
-                dict(request.fab.verifications),
+                verification_dict,
             )
-            fab_hash = ffs.put(fab.content, {})
+            fab_hash = ffs.put(fab.content, fab.verifications)
+
             if fab_hash != fab.hash_str:
                 raise RuntimeError(
                     f"FAB ({fab.hash_str}) hash from request doesn't match contents"
@@ -612,3 +627,71 @@ def _check_flwr_aid_in_run(
             grpc.StatusCode.PERMISSION_DENIED,
             "⛔️ Run ID does not belong to the account",
         )
+
+
+def _format_verification(verifications: list[dict[str, str]]) -> dict[str, str]:
+    """Format verification information for FAB."""
+    # Convert verifications to dict[str, str] type
+    verification_dict = {
+        item["public_key_id"]: json.dumps(
+            {k: v for k, v in item.items() if k != "public_key_id"}
+        )
+        for item in verifications
+    }
+    verification_dict.update({"valid_license": "Valid"})
+
+    return verification_dict
+
+
+def _get_remote_fab(
+    fleet_api_type: str | None,
+    app_spec: str,
+    context: grpc.ServicerContext,
+) -> tuple[bytes, dict[str, str]]:
+    """Get remote FAB from Flower platform API."""
+    if fleet_api_type == TRANSPORT_TYPE_GRPC_ADAPTER:
+        context.abort(
+            grpc.StatusCode.FAILED_PRECONDITION,
+            "The selected SuperLink transport type is not "
+            "supported for connecting to Flower Platform.",
+        )
+
+    # Parse and validate app specification
+    try:
+        app_id, app_version = parse_app_spec(app_spec)
+    except ValueError as e:
+        context.abort(
+            grpc.StatusCode.FAILED_PRECONDITION,
+            f"{e}",
+        )
+
+    # Request download link and verification information
+    url = f"{PLATFORM_API_URL}/hub/fetch-fab"
+    try:
+        presigned_url, verifications = request_download_link(
+            app_id, app_version, url, "fab_url"
+        )
+    except ValueError as e:
+        context.abort(
+            grpc.StatusCode.FAILED_PRECONDITION,
+            f"{e}",
+        )
+
+    # Format verification information
+    verification_dict = (
+        _format_verification(verifications)
+        if verifications is not None
+        else {"valid_license": ""}
+    )
+
+    # Download FAB from Flower platform API
+    try:
+        r = requests.get(presigned_url, timeout=60)
+        r.raise_for_status()
+    except requests.RequestException as e:
+        context.abort(
+            grpc.StatusCode.FAILED_PRECONDITION,
+            f"FAB download failed: {str(e)}",
+        )
+    fab_file = r.content
+    return fab_file, verification_dict