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

flwr/simulation/simulationio_connection.py

@@ -0,0 +1,89 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Flower SimulationIo connection."""
+
+
+from logging import DEBUG, WARNING
+from typing import Optional, cast
+
+import grpc
+
+from flwr.common.constant import SIMULATIONIO_API_DEFAULT_CLIENT_ADDRESS
+from flwr.common.grpc import create_channel
+from flwr.common.logger import log
+from flwr.common.retry_invoker import _make_simple_grpc_retry_invoker, _wrap_stub
+from flwr.proto.simulationio_pb2_grpc import SimulationIoStub  # pylint: disable=E0611
+
+
+class SimulationIoConnection:
+    """`SimulationIoConnection` provides an interface to the SimulationIo API.
+
+    Parameters
+    ----------
+    simulationio_service_address : str (default: "[::]:9094")
+        The address (URL, IPv6, IPv4) of the SuperLink SimulationIo API service.
+    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__(  # pylint: disable=too-many-arguments
+        self,
+        simulationio_service_address: str = SIMULATIONIO_API_DEFAULT_CLIENT_ADDRESS,
+        root_certificates: Optional[bytes] = None,
+    ) -> None:
+        self._addr = simulationio_service_address
+        self._cert = root_certificates
+        self._grpc_stub: Optional[SimulationIoStub] = None
+        self._channel: Optional[grpc.Channel] = None
+        self._retry_invoker = _make_simple_grpc_retry_invoker()
+
+    @property
+    def _is_connected(self) -> bool:
+        """Check if connected to the SimulationIo API server."""
+        return self._channel is not None
+
+    @property
+    def _stub(self) -> SimulationIoStub:
+        """SimulationIo stub."""
+        if not self._is_connected:
+            self._connect()
+        return cast(SimulationIoStub, self._grpc_stub)
+
+    def _connect(self) -> None:
+        """Connect to the SimulationIo API."""
+        if self._is_connected:
+            log(WARNING, "Already connected")
+            return
+        self._channel = create_channel(
+            server_address=self._addr,
+            insecure=(self._cert is None),
+            root_certificates=self._cert,
+        )
+        self._grpc_stub = SimulationIoStub(self._channel)
+        _wrap_stub(self._grpc_stub, self._retry_invoker)
+        log(DEBUG, "[SimulationIO] Connected to %s", self._addr)
+
+    def _disconnect(self) -> None:
+        """Disconnect from the SimulationIo API."""
+        if not self._is_connected:
+            log(DEBUG, "Already disconnected")
+            return
+        channel: grpc.Channel = self._channel
+        self._channel = None
+        self._grpc_stub = None
+        channel.close()
+        log(DEBUG, "[SimulationIO] Disconnected")

flwr/superexec/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Flower SuperExec service."""

flwr/superexec/app.py

@@ -0,0 +1,59 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Flower SuperExec app."""
+
+
+import argparse
+import sys
+from logging import INFO
+
+from flwr.common import log
+from flwr.common.object_ref import load_app, validate
+
+from .executor import Executor
+
+
+def run_superexec() -> None:
+    """Run Flower SuperExec."""
+    log(INFO, "Starting Flower SuperExec")
+
+    sys.exit(
+        "Manually launching the SuperExec is deprecated. Since `flwr 1.13.0` "
+        "the executor service runs in the SuperLink. Launching it manually is not "
+        "recommended."
+    )
+
+
+def load_executor(
+    args: argparse.Namespace,
+) -> Executor:
+    """Get the executor plugin."""
+    executor_ref: str = args.executor
+    valid, error_msg = validate(executor_ref, project_dir=args.executor_dir)
+    if not valid and error_msg:
+        raise LoadExecutorError(error_msg) from None
+
+    executor = load_app(executor_ref, LoadExecutorError, args.executor_dir)
+
+    if not isinstance(executor, Executor):
+        raise LoadExecutorError(
+            f"Attribute {executor_ref} is not of type {Executor}",
+        ) from None
+
+    return executor
+
+
+class LoadExecutorError(Exception):
+    """Error when trying to load `Executor`."""

flwr/superexec/deployment.py

@@ -0,0 +1,188 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Deployment engine executor."""
+
+
+import hashlib
+from logging import ERROR, INFO
+from pathlib import Path
+from typing import Optional
+
+from typing_extensions import override
+
+from flwr.cli.config_utils import get_fab_metadata
+from flwr.common import ConfigsRecord, Context, RecordSet
+from flwr.common.constant import (
+    SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS,
+    Status,
+    SubStatus,
+)
+from flwr.common.logger import log
+from flwr.common.typing import Fab, RunStatus, UserConfig
+from flwr.server.superlink.ffs import Ffs
+from flwr.server.superlink.ffs.ffs_factory import FfsFactory
+from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
+
+from .executor import Executor
+
+
+class DeploymentEngine(Executor):
+    """Deployment engine executor.
+
+    Parameters
+    ----------
+    serverappio_api_address: str (default: "127.0.0.1:9091")
+        Address of the SuperLink to connect to.
+    root_certificates: Optional[str] (default: None)
+        Specifies the path to the PEM-encoded root certificate file for
+        establishing secure HTTPS connections.
+    flwr_dir: Optional[str] (default: None)
+        The path containing installed Flower Apps.
+    """
+
+    def __init__(
+        self,
+        serverappio_api_address: str = SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS,
+        root_certificates: Optional[str] = None,
+        flwr_dir: Optional[str] = None,
+    ) -> None:
+        self.serverappio_api_address = serverappio_api_address
+        if root_certificates is None:
+            self.root_certificates = None
+            self.root_certificates_bytes = None
+        else:
+            self.root_certificates = root_certificates
+            self.root_certificates_bytes = Path(root_certificates).read_bytes()
+        self.flwr_dir = flwr_dir
+        self.linkstate_factory: Optional[LinkStateFactory] = None
+        self.ffs_factory: Optional[FfsFactory] = None
+
+    @override
+    def initialize(
+        self, linkstate_factory: LinkStateFactory, ffs_factory: FfsFactory
+    ) -> None:
+        """Initialize the executor with the necessary factories."""
+        self.linkstate_factory = linkstate_factory
+        self.ffs_factory = ffs_factory
+
+    @property
+    def linkstate(self) -> LinkState:
+        """Return the LinkState."""
+        if self.linkstate_factory is None:
+            raise RuntimeError("Executor is not initialized.")
+        return self.linkstate_factory.state()
+
+    @property
+    def ffs(self) -> Ffs:
+        """Return the Flower File Storage (FFS)."""
+        if self.ffs_factory is None:
+            raise RuntimeError("Executor is not initialized.")
+        return self.ffs_factory.ffs()
+
+    @override
+    def set_config(
+        self,
+        config: UserConfig,
+    ) -> None:
+        """Set executor config arguments.
+
+        Parameters
+        ----------
+        config : UserConfig
+            A dictionary for configuration values.
+            Supported configuration key/value pairs:
+            - "superlink": str
+                The address of the SuperLink ServerAppIo API.
+            - "root-certificates": str
+                The path to the root certificates.
+            - "flwr-dir": str
+                The path to the Flower directory.
+        """
+        if not config:
+            return
+        if superlink_address := config.get("superlink"):
+            if not isinstance(superlink_address, str):
+                raise ValueError("The `superlink` value should be of type `str`.")
+            self.serverappio_api_address = superlink_address
+        if root_certificates := config.get("root-certificates"):
+            if not isinstance(root_certificates, str):
+                raise ValueError(
+                    "The `root-certificates` value should be of type `str`."
+                )
+            self.root_certificates = root_certificates
+            self.root_certificates_bytes = Path(str(root_certificates)).read_bytes()
+        if flwr_dir := config.get("flwr-dir"):
+            if not isinstance(flwr_dir, str):
+                raise ValueError("The `flwr-dir` value should be of type `str`.")
+            self.flwr_dir = str(flwr_dir)
+
+    def _create_run(
+        self,
+        fab: Fab,
+        override_config: UserConfig,
+    ) -> int:
+        fab_hash = self.ffs.put(fab.content, {})
+        if fab_hash != fab.hash_str:
+            raise RuntimeError(
+                f"FAB ({fab.hash_str}) hash from request doesn't match contents"
+            )
+        fab_id, fab_version = get_fab_metadata(fab.content)
+
+        run_id = self.linkstate.create_run(
+            fab_id, fab_version, fab_hash, override_config, ConfigsRecord()
+        )
+        return run_id
+
+    def _create_context(self, run_id: int) -> None:
+        """Register a Context for a Run."""
+        # Create an empty context for the Run
+        context = Context(
+            run_id=run_id, node_id=0, node_config={}, state=RecordSet(), run_config={}
+        )
+
+        # Register the context at the LinkState
+        self.linkstate.set_serverapp_context(run_id=run_id, context=context)
+
+    @override
+    def start_run(
+        self,
+        fab_file: bytes,
+        override_config: UserConfig,
+        federation_options: ConfigsRecord,
+    ) -> Optional[int]:
+        """Start run using the Flower Deployment Engine."""
+        run_id = None
+        try:
+
+            # Call SuperLink to create run
+            run_id = self._create_run(
+                Fab(hashlib.sha256(fab_file).hexdigest(), fab_file), override_config
+            )
+
+            # Register context for the Run
+            self._create_context(run_id=run_id)
+            log(INFO, "Created run %s", str(run_id))
+
+            return run_id
+        # pylint: disable-next=broad-except
+        except Exception as e:
+            log(ERROR, "Could not start run: %s", str(e))
+            if run_id:
+                run_status = RunStatus(Status.FINISHED, SubStatus.FAILED, str(e))
+                self.linkstate.update_run_status(run_id, new_status=run_status)
+            return None
+
+
+executor = DeploymentEngine()

flwr/superexec/exec_grpc.py

@@ -0,0 +1,80 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""SuperExec gRPC API."""
+
+
+from collections.abc import Sequence
+from logging import INFO
+from typing import Optional
+
+import grpc
+
+from flwr.common import GRPC_MAX_MESSAGE_LENGTH
+from flwr.common.auth_plugin import ExecAuthPlugin
+from flwr.common.grpc import generic_create_grpc_server
+from flwr.common.logger import log
+from flwr.common.typing import UserConfig
+from flwr.proto.exec_pb2_grpc import add_ExecServicer_to_server
+from flwr.server.superlink.ffs.ffs_factory import FfsFactory
+from flwr.server.superlink.linkstate import LinkStateFactory
+from flwr.superexec.exec_user_auth_interceptor import ExecUserAuthInterceptor
+
+from .exec_servicer import ExecServicer
+from .executor import Executor
+
+
+# pylint: disable-next=too-many-arguments, too-many-positional-arguments
+def run_exec_api_grpc(
+    address: str,
+    executor: Executor,
+    state_factory: LinkStateFactory,
+    ffs_factory: FfsFactory,
+    certificates: Optional[tuple[bytes, bytes, bytes]],
+    config: UserConfig,
+    auth_plugin: Optional[ExecAuthPlugin] = None,
+) -> grpc.Server:
+    """Run Exec API (gRPC, request-response)."""
+    executor.set_config(config)
+
+    exec_servicer: grpc.Server = ExecServicer(
+        linkstate_factory=state_factory,
+        ffs_factory=ffs_factory,
+        executor=executor,
+        auth_plugin=auth_plugin,
+    )
+    interceptors: Optional[Sequence[grpc.ServerInterceptor]] = None
+    if auth_plugin is not None:
+        interceptors = [ExecUserAuthInterceptor(auth_plugin)]
+    exec_add_servicer_to_server_fn = add_ExecServicer_to_server
+    exec_grpc_server = generic_create_grpc_server(
+        servicer_and_add_fn=(exec_servicer, exec_add_servicer_to_server_fn),
+        server_address=address,
+        max_message_length=GRPC_MAX_MESSAGE_LENGTH,
+        certificates=certificates,
+        interceptors=interceptors,
+    )
+
+    if auth_plugin is None:
+        log(INFO, "Flower Deployment Engine: Starting Exec API on %s", address)
+    else:
+        log(
+            INFO,
+            "Flower Deployment Engine: Starting Exec API with user "
+            "authentication on %s",
+            address,
+        )
+    exec_grpc_server.start()
+
+    return exec_grpc_server

flwr/superexec/exec_servicer.py

@@ -0,0 +1,231 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""SuperExec API servicer."""
+
+
+import time
+from collections.abc import Generator
+from logging import ERROR, INFO
+from typing import Any, Optional
+from uuid import UUID
+
+import grpc
+
+from flwr.common import now
+from flwr.common.auth_plugin import ExecAuthPlugin
+from flwr.common.constant import LOG_STREAM_INTERVAL, Status, SubStatus
+from flwr.common.logger import log
+from flwr.common.serde import (
+    configs_record_from_proto,
+    run_to_proto,
+    user_config_from_proto,
+)
+from flwr.common.typing import RunStatus
+from flwr.proto import exec_pb2_grpc  # pylint: disable=E0611
+from flwr.proto.exec_pb2 import (  # pylint: disable=E0611
+    GetAuthTokensRequest,
+    GetAuthTokensResponse,
+    GetLoginDetailsRequest,
+    GetLoginDetailsResponse,
+    ListRunsRequest,
+    ListRunsResponse,
+    StartRunRequest,
+    StartRunResponse,
+    StopRunRequest,
+    StopRunResponse,
+    StreamLogsRequest,
+    StreamLogsResponse,
+)
+from flwr.server.superlink.ffs.ffs_factory import FfsFactory
+from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
+
+from .executor import Executor
+
+
+class ExecServicer(exec_pb2_grpc.ExecServicer):
+    """SuperExec API servicer."""
+
+    def __init__(
+        self,
+        linkstate_factory: LinkStateFactory,
+        ffs_factory: FfsFactory,
+        executor: Executor,
+        auth_plugin: Optional[ExecAuthPlugin] = None,
+    ) -> None:
+        self.linkstate_factory = linkstate_factory
+        self.ffs_factory = ffs_factory
+        self.executor = executor
+        self.executor.initialize(linkstate_factory, ffs_factory)
+        self.auth_plugin = auth_plugin
+
+    def StartRun(
+        self, request: StartRunRequest, context: grpc.ServicerContext
+    ) -> StartRunResponse:
+        """Create run ID."""
+        log(INFO, "ExecServicer.StartRun")
+
+        run_id = self.executor.start_run(
+            request.fab.content,
+            user_config_from_proto(request.override_config),
+            configs_record_from_proto(request.federation_options),
+        )
+
+        if run_id is None:
+            log(ERROR, "Executor failed to start run")
+            return StartRunResponse()
+
+        return StartRunResponse(run_id=run_id)
+
+    def StreamLogs(  # pylint: disable=C0103
+        self, request: StreamLogsRequest, context: grpc.ServicerContext
+    ) -> Generator[StreamLogsResponse, Any, None]:
+        """Get logs."""
+        log(INFO, "ExecServicer.StreamLogs")
+        state = self.linkstate_factory.state()
+
+        # Retrieve run ID
+        run_id = request.run_id
+
+        # Exit if `run_id` not found
+        if not state.get_run(run_id):
+            context.abort(grpc.StatusCode.NOT_FOUND, "Run ID not found")
+
+        after_timestamp = request.after_timestamp + 1e-6
+        while context.is_active():
+            log_msg, latest_timestamp = state.get_serverapp_log(run_id, after_timestamp)
+            if log_msg:
+                yield StreamLogsResponse(
+                    log_output=log_msg,
+                    latest_timestamp=latest_timestamp,
+                )
+                # Add a small epsilon to the latest timestamp to avoid getting
+                # the same log
+                after_timestamp = max(latest_timestamp + 1e-6, after_timestamp)
+
+            # Wait for and continue to yield more log responses only if the
+            # run isn't completed yet. If the run is finished, the entire log
+            # is returned at this point and the server ends the stream.
+            run_status = state.get_run_status({run_id})[run_id]
+            if run_status.status == Status.FINISHED:
+                log(INFO, "All logs for run ID `%s` returned", request.run_id)
+                context.cancel()
+
+            time.sleep(LOG_STREAM_INTERVAL)  # Sleep briefly to avoid busy waiting
+
+    def ListRuns(
+        self, request: ListRunsRequest, context: grpc.ServicerContext
+    ) -> ListRunsResponse:
+        """Handle `flwr ls` command."""
+        log(INFO, "ExecServicer.List")
+        state = self.linkstate_factory.state()
+
+        # Handle `flwr ls --runs`
+        if not request.HasField("run_id"):
+            return _create_list_runs_response(state.get_run_ids(), state)
+        # Handle `flwr ls --run-id <run_id>`
+        return _create_list_runs_response({request.run_id}, state)
+
+    def StopRun(
+        self, request: StopRunRequest, context: grpc.ServicerContext
+    ) -> StopRunResponse:
+        """Stop a given run ID."""
+        log(INFO, "ExecServicer.StopRun")
+        state = self.linkstate_factory.state()
+
+        # Exit if `run_id` not found
+        if not state.get_run(request.run_id):
+            context.abort(
+                grpc.StatusCode.NOT_FOUND, f"Run ID {request.run_id} not found"
+            )
+
+        run_status = state.get_run_status({request.run_id})[request.run_id]
+        if run_status.status == Status.FINISHED:
+            context.abort(
+                grpc.StatusCode.FAILED_PRECONDITION,
+                f"Run ID {request.run_id} is already finished",
+            )
+
+        update_success = state.update_run_status(
+            run_id=request.run_id,
+            new_status=RunStatus(Status.FINISHED, SubStatus.STOPPED, ""),
+        )
+
+        if update_success:
+            task_ids: set[UUID] = state.get_task_ids_from_run_id(request.run_id)
+
+            # Delete TaskIns and TaskRes for the `run_id`
+            state.delete_tasks(task_ids)
+
+        return StopRunResponse(success=update_success)
+
+    def GetLoginDetails(
+        self, request: GetLoginDetailsRequest, context: grpc.ServicerContext
+    ) -> GetLoginDetailsResponse:
+        """Start login."""
+        log(INFO, "ExecServicer.GetLoginDetails")
+        if self.auth_plugin is None:
+            context.abort(
+                grpc.StatusCode.UNIMPLEMENTED,
+                "ExecServicer initialized without user authentication",
+            )
+            raise grpc.RpcError()  # This line is unreachable
+
+        # Get login details
+        details = self.auth_plugin.get_login_details()
+
+        # Return empty response if details is None
+        if details is None:
+            return GetLoginDetailsResponse()
+
+        return GetLoginDetailsResponse(
+            auth_type=details.auth_type,
+            device_code=details.device_code,
+            verification_uri_complete=details.verification_uri_complete,
+            expires_in=details.expires_in,
+            interval=details.interval,
+        )
+
+    def GetAuthTokens(
+        self, request: GetAuthTokensRequest, context: grpc.ServicerContext
+    ) -> GetAuthTokensResponse:
+        """Get auth token."""
+        log(INFO, "ExecServicer.GetAuthTokens")
+        if self.auth_plugin is None:
+            context.abort(
+                grpc.StatusCode.UNIMPLEMENTED,
+                "ExecServicer initialized without user authentication",
+            )
+            raise grpc.RpcError()  # This line is unreachable
+
+        # Get auth tokens
+        credentials = self.auth_plugin.get_auth_tokens(request.device_code)
+
+        # Return empty response if credentials is None
+        if credentials is None:
+            return GetAuthTokensResponse()
+
+        return GetAuthTokensResponse(
+            access_token=credentials.access_token,
+            refresh_token=credentials.refresh_token,
+        )
+
+
+def _create_list_runs_response(run_ids: set[int], state: LinkState) -> ListRunsResponse:
+    """Create response for `flwr ls --runs` and `flwr ls --run-id <run_id>`."""
+    run_dict = {run_id: state.get_run(run_id) for run_id in run_ids}
+    return ListRunsResponse(
+        run_dict={run_id: run_to_proto(run) for run_id, run in run_dict.items() if run},
+        now=now().isoformat(),
+    )