flwr-nightly 1.8.0.dev20240304__tar.gz → 1.8.0.dev20240306__tar.gz

{flwr_nightly-1.8.0.dev20240304 → flwr_nightly-1.8.0.dev20240306}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flwr-nightly
-Version: 1.8.0.dev20240304
+Version: 1.8.0.dev20240306
 Summary: Flower: A Friendly Federated Learning Framework
 Home-page: https://flower.ai
 License: Apache-2.0
@@ -42,6 +42,7 @@ Requires-Dist: pydantic (<2.0.0) ; extra == "simulation"
 Requires-Dist: ray (==2.6.3) ; extra == "simulation"
 Requires-Dist: requests (>=2.31.0,<3.0.0) ; extra == "rest"
 Requires-Dist: starlette (>=0.31.0,<0.32.0) ; extra == "rest"
+Requires-Dist: tomli (>=2.0.1,<3.0.0)
 Requires-Dist: typer[all] (>=0.9.0,<0.10.0)
 Requires-Dist: uvicorn[standard] (>=0.23.0,<0.24.0) ; extra == "rest"
 Project-URL: Documentation, https://flower.ai
@@ -182,7 +183,6 @@ Quickstart examples:
 - [Quickstart (JAX)](https://github.com/adap/flower/tree/main/examples/quickstart-jax)
 - [Quickstart (MONAI)](https://github.com/adap/flower/tree/main/examples/quickstart-monai)
 - [Quickstart (scikit-learn)](https://github.com/adap/flower/tree/main/examples/sklearn-logreg-mnist)
-- [Quickstart (XGBoost)](https://github.com/adap/flower/tree/main/examples/xgboost-quickstart)
 - [Quickstart (Android [TFLite])](https://github.com/adap/flower/tree/main/examples/android)
 - [Quickstart (iOS [CoreML])](https://github.com/adap/flower/tree/main/examples/ios)
 - [Quickstart (MLX)](https://github.com/adap/flower/tree/main/examples/quickstart-mlx)
@@ -194,7 +194,8 @@ Other [examples](https://github.com/adap/flower/tree/main/examples):
 - [PyTorch: From Centralized to Federated](https://github.com/adap/flower/tree/main/examples/pytorch-from-centralized-to-federated)
 - [Vertical FL](https://github.com/adap/flower/tree/main/examples/vertical-fl)
 - [Federated Finetuning of OpenAI's Whisper](https://github.com/adap/flower/tree/main/examples/whisper-federated-finetuning)
-- [Comprehensive XGBoost](https://github.com/adap/flower/tree/main/examples/xgboost-comprehensive)
+- [Federated Finetuning of Large Language Model](https://github.com/adap/flower/tree/main/examples/fedllm-finetune)
+- [Federated Finetuning of a Vision Transformer](https://github.com/adap/flower/tree/main/examples/vit-finetune)
 - [Advanced Flower with TensorFlow/Keras](https://github.com/adap/flower/tree/main/examples/advanced-tensorflow)
 - [Advanced Flower with PyTorch](https://github.com/adap/flower/tree/main/examples/advanced-pytorch)
 - Single-Machine Simulation of Federated Learning Systems ([PyTorch](https://github.com/adap/flower/tree/main/examples/simulation-pytorch)) ([Tensorflow](https://github.com/adap/flower/tree/main/examples/simulation-tensorflow))

{flwr_nightly-1.8.0.dev20240304 → flwr_nightly-1.8.0.dev20240306}/README.md

@@ -132,7 +132,6 @@ Quickstart examples:
 - [Quickstart (JAX)](https://github.com/adap/flower/tree/main/examples/quickstart-jax)
 - [Quickstart (MONAI)](https://github.com/adap/flower/tree/main/examples/quickstart-monai)
 - [Quickstart (scikit-learn)](https://github.com/adap/flower/tree/main/examples/sklearn-logreg-mnist)
-- [Quickstart (XGBoost)](https://github.com/adap/flower/tree/main/examples/xgboost-quickstart)
 - [Quickstart (Android [TFLite])](https://github.com/adap/flower/tree/main/examples/android)
 - [Quickstart (iOS [CoreML])](https://github.com/adap/flower/tree/main/examples/ios)
 - [Quickstart (MLX)](https://github.com/adap/flower/tree/main/examples/quickstart-mlx)
@@ -144,7 +143,8 @@ Other [examples](https://github.com/adap/flower/tree/main/examples):
 - [PyTorch: From Centralized to Federated](https://github.com/adap/flower/tree/main/examples/pytorch-from-centralized-to-federated)
 - [Vertical FL](https://github.com/adap/flower/tree/main/examples/vertical-fl)
 - [Federated Finetuning of OpenAI's Whisper](https://github.com/adap/flower/tree/main/examples/whisper-federated-finetuning)
-- [Comprehensive XGBoost](https://github.com/adap/flower/tree/main/examples/xgboost-comprehensive)
+- [Federated Finetuning of Large Language Model](https://github.com/adap/flower/tree/main/examples/fedllm-finetune)
+- [Federated Finetuning of a Vision Transformer](https://github.com/adap/flower/tree/main/examples/vit-finetune)
 - [Advanced Flower with TensorFlow/Keras](https://github.com/adap/flower/tree/main/examples/advanced-tensorflow)
 - [Advanced Flower with PyTorch](https://github.com/adap/flower/tree/main/examples/advanced-pytorch)
 - Single-Machine Simulation of Federated Learning Systems ([PyTorch](https://github.com/adap/flower/tree/main/examples/simulation-pytorch)) ([Tensorflow](https://github.com/adap/flower/tree/main/examples/simulation-tensorflow))

{flwr_nightly-1.8.0.dev20240304 → flwr_nightly-1.8.0.dev20240306}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "flwr-nightly"
-version = "1.8.0-dev20240304"
+version = "1.8.0-dev20240306"
 description = "Flower: A Friendly Federated Learning Framework"
 license = "Apache-2.0"
 authors = ["The Flower Authors <hello@flower.ai>"]
@@ -58,7 +58,7 @@ flower-fleet-api = "flwr.server:run_fleet_api"
 flower-superlink = "flwr.server:run_superlink"
 flower-client-app = "flwr.client:run_client_app"
 flower-server-app = "flwr.server:run_server_app"
-flower-simulation = "flwr.simulation:run_simulation"
+flower-simulation = "flwr.simulation:run_simulation_from_cli"
 
 [tool.poetry.dependencies]
 python = "^3.8"
@@ -70,7 +70,8 @@ cryptography = "^42.0.4"
 pycryptodome = "^3.18.0"
 iterators = "^0.0.2"
 typer = { version = "^0.9.0", extras=["all"] }
-# Optional dependencies (VCE)
+tomli = "^2.0.1"
+# Optional dependencies (Simulation Engine)
 ray = { version = "==2.6.3", optional = true }
 pydantic = { version = "<2.0.0", optional = true }
 # Optional dependencies (REST transport layer)

{flwr_nightly-1.8.0.dev20240304 → flwr_nightly-1.8.0.dev20240306}/src/py/flwr/cli/app.py

@@ -18,6 +18,7 @@ import typer
 
 from .example import example
 from .new import new
+from .run import run
 
 app = typer.Typer(
     help=typer.style(
@@ -30,6 +31,7 @@ app = typer.Typer(
 
 app.command()(new)
 app.command()(example)
+app.command()(run)
 
 if __name__ == "__main__":
     app()

flwr_nightly-1.8.0.dev20240306/src/py/flwr/cli/flower_toml.py

@@ -0,0 +1,151 @@
+# 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.
+# ==============================================================================
+"""Utility to validate the `flower.toml` file."""
+
+import importlib
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+import tomli
+
+
+def load_flower_toml(path: Optional[str] = None) -> Optional[Dict[str, Any]]:
+    """Load flower.toml and return as dict."""
+    if path is None:
+        cur_dir = os.getcwd()
+        toml_path = os.path.join(cur_dir, "flower.toml")
+    else:
+        toml_path = path
+
+    if not os.path.isfile(toml_path):
+        return None
+
+    with open(toml_path, encoding="utf-8") as toml_file:
+        data = tomli.loads(toml_file.read())
+        return data
+
+
+def validate_flower_toml_fields(
+    config: Dict[str, Any]
+) -> Tuple[bool, List[str], List[str]]:
+    """Validate flower.toml fields."""
+    errors = []
+    warnings = []
+
+    if "project" not in config:
+        errors.append("Missing [project] section")
+    else:
+        if "name" not in config["project"]:
+            errors.append('Property "name" missing in [project]')
+        if "version" not in config["project"]:
+            errors.append('Property "version" missing in [project]')
+        if "description" not in config["project"]:
+            warnings.append('Recommended property "description" missing in [project]')
+        if "license" not in config["project"]:
+            warnings.append('Recommended property "license" missing in [project]')
+        if "authors" not in config["project"]:
+            warnings.append('Recommended property "authors" missing in [project]')
+
+    if "flower" not in config:
+        errors.append("Missing [flower] section")
+    elif "components" not in config["flower"]:
+        errors.append("Missing [flower.components] section")
+    else:
+        if "serverapp" not in config["flower"]["components"]:
+            errors.append('Property "serverapp" missing in [flower.components]')
+        if "clientapp" not in config["flower"]["components"]:
+            errors.append('Property "clientapp" missing in [flower.components]')
+
+    return len(errors) == 0, errors, warnings
+
+
+def validate_object_reference(ref: str) -> Tuple[bool, Optional[str]]:
+    """Validate object reference.
+
+    Returns
+    -------
+    Tuple[bool, Optional[str]]
+        A boolean indicating whether an object reference is valid and
+        the reason why it might not be.
+    """
+    module_str, _, attributes_str = ref.partition(":")
+    if not module_str:
+        return (
+            False,
+            f"Missing module in {ref}",
+        )
+    if not attributes_str:
+        return (
+            False,
+            f"Missing attribute in {ref}",
+        )
+
+    # Load module
+    try:
+        module = importlib.import_module(module_str)
+    except ModuleNotFoundError:
+        return False, f"Unable to load module {module_str}"
+
+    # Recursively load attribute
+    attribute = module
+    try:
+        for attribute_str in attributes_str.split("."):
+            attribute = getattr(attribute, attribute_str)
+    except AttributeError:
+        return (
+            False,
+            f"Unable to load attribute {attributes_str} from module {module_str}",
+        )
+
+    return (True, None)
+
+
+def validate_flower_toml(config: Dict[str, Any]) -> Tuple[bool, List[str], List[str]]:
+    """Validate flower.toml."""
+    is_valid, errors, warnings = validate_flower_toml_fields(config)
+
+    if not is_valid:
+        return False, errors, warnings
+
+    # Validate serverapp
+    is_valid, reason = validate_object_reference(
+        config["flower"]["components"]["serverapp"]
+    )
+    if not is_valid and isinstance(reason, str):
+        return False, [reason], []
+
+    # Validate clientapp
+    is_valid, reason = validate_object_reference(
+        config["flower"]["components"]["clientapp"]
+    )
+
+    if not is_valid and isinstance(reason, str):
+        return False, [reason], []
+
+    return True, [], []
+
+
+def apply_defaults(
+    config: Dict[str, Any],
+    defaults: Dict[str, Any],
+) -> Dict[str, Any]:
+    """Apply defaults to config."""
+    for key in defaults:
+        if key in config:
+            if isinstance(config[key], dict) and isinstance(defaults[key], dict):
+                apply_defaults(config[key], defaults[key])
+        else:
+            config[key] = defaults[key]
+    return config

{flwr_nightly-1.8.0.dev20240304 → flwr_nightly-1.8.0.dev20240306}/src/py/flwr/cli/new/new.py

@@ -28,6 +28,7 @@ from ..utils import prompt_options
 class MlFramework(str, Enum):
     """Available frameworks."""
 
+    NUMPY = "NumPy"
     PYTORCH = "PyTorch"
     TENSORFLOW = "TensorFlow"
 

flwr_nightly-1.8.0.dev20240306/src/py/flwr/cli/new/templates/app/code/client.numpy.py.tpl

@@ -0,0 +1,24 @@
+"""$project_name: A Flower / NumPy app."""
+
+import flwr as fl
+import numpy as np
+
+
+# Flower client, adapted from Pytorch quickstart example
+class FlowerClient(fl.client.NumPyClient):
+    def get_parameters(self, config):
+        return [np.ones((1, 1))]
+
+    def fit(self, parameters, config):
+        return ([np.ones((1, 1))], 1, {})
+
+    def evaluate(self, parameters, config):
+        return float(0.0), 1, {"accuracy": float(1.0)}
+
+
+def client_fn(cid: str):
+    return FlowerClient().to_client()
+
+
+# ClientApp for Flower-Next
+app = fl.client.ClientApp(client_fn=client_fn)

flwr_nightly-1.8.0.dev20240306/src/py/flwr/cli/new/templates/app/code/server.numpy.py.tpl

@@ -0,0 +1,12 @@
+"""$project_name: A Flower / NumPy app."""
+
+import flwr as fl
+
+# Configure the strategy
+strategy = fl.server.strategy.FedAvg()
+
+# Flower ServerApp
+app = fl.server.ServerApp(
+    config=fl.server.ServerConfig(num_rounds=1),
+    strategy=strategy,
+)

{flwr_nightly-1.8.0.dev20240304 → flwr_nightly-1.8.0.dev20240306}/src/py/flwr/cli/new/templates/app/flower.toml.tpl

@@ -1,10 +1,10 @@
-[flower]
+[project]
 name = "$project_name"
 version = "1.0.0"
 description = ""
 license = "Apache-2.0"
 authors = ["The Flower Authors <hello@flower.ai>"]
 
-[components]
+[flower.components]
 serverapp = "$project_name.server:app"
 clientapp = "$project_name.client:app"

flwr_nightly-1.8.0.dev20240306/src/py/flwr/cli/new/templates/app/requirements.numpy.txt.tpl

@@ -0,0 +1,2 @@
+flwr>=1.8, <2.0
+numpy >= 1.21.0

flwr_nightly-1.8.0.dev20240306/src/py/flwr/cli/run/__init__.py

@@ -0,0 +1,21 @@
+# 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 command line interface `run` command."""
+
+from .run import run as run
+
+__all__ = [
+    "run",
+]

flwr_nightly-1.8.0.dev20240306/src/py/flwr/cli/run/run.py

@@ -0,0 +1,102 @@
+# 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 command line interface `run` command."""
+
+import sys
+
+import typer
+
+from flwr.cli.flower_toml import apply_defaults, load_flower_toml, validate_flower_toml
+from flwr.simulation.run_simulation import _run_simulation
+
+
+def run() -> None:
+    """Run Flower project."""
+    print(
+        typer.style("Loading project configuration... ", fg=typer.colors.BLUE),
+        end="",
+    )
+    config = load_flower_toml()
+    if not config:
+        print(
+            typer.style(
+                "Project configuration could not be loaded. "
+                "flower.toml does not exist.",
+                fg=typer.colors.RED,
+                bold=True,
+            )
+        )
+        sys.exit()
+    print(typer.style("Success", fg=typer.colors.GREEN))
+
+    print(
+        typer.style("Validating project configuration... ", fg=typer.colors.BLUE),
+        end="",
+    )
+    is_valid, errors, warnings = validate_flower_toml(config)
+    if warnings:
+        print(
+            typer.style(
+                "Project configuration is missing the following "
+                "recommended properties:\n"
+                + "\n".join([f"- {line}" for line in warnings]),
+                fg=typer.colors.RED,
+                bold=True,
+            )
+        )
+
+    if not is_valid:
+        print(
+            typer.style(
+                "Project configuration could not be loaded.\nflower.toml is invalid:\n"
+                + "\n".join([f"- {line}" for line in errors]),
+                fg=typer.colors.RED,
+                bold=True,
+            )
+        )
+        sys.exit()
+    print(typer.style("Success", fg=typer.colors.GREEN))
+
+    # Apply defaults
+    defaults = {
+        "flower": {
+            "engine": {"name": "simulation", "simulation": {"super-node": {"num": 100}}}
+        }
+    }
+    config = apply_defaults(config, defaults)
+
+    server_app_ref = config["flower"]["components"]["serverapp"]
+    client_app_ref = config["flower"]["components"]["clientapp"]
+    engine = config["flower"]["engine"]["name"]
+
+    if engine == "simulation":
+        num_supernodes = config["flower"]["engine"]["simulation"]["super-node"]["num"]
+
+        print(
+            typer.style("Starting run... ", fg=typer.colors.BLUE),
+        )
+        _run_simulation(
+            server_app_attr=server_app_ref,
+            client_app_attr=client_app_ref,
+            num_supernodes=num_supernodes,
+        )
+    else:
+        print(
+            typer.style(
+                f"Engine '{engine}' is not yet supported in `flwr run`",
+                fg=typer.colors.RED,
+                bold=True,
+            )
+        )

{flwr_nightly-1.8.0.dev20240304 → flwr_nightly-1.8.0.dev20240306}/src/py/flwr/client/app.py

@@ -20,7 +20,9 @@ import sys
 import time
 from logging import DEBUG, INFO, WARN
 from pathlib import Path
-from typing import Callable, ContextManager, Optional, Tuple, Union
+from typing import Callable, ContextManager, Optional, Tuple, Type, Union
+
+from grpc import RpcError
 
 from flwr.client.client import Client
 from flwr.client.client_app import ClientApp
@@ -36,6 +38,7 @@ from flwr.common.constant import (
 )
 from flwr.common.exit_handlers import register_exit_handlers
 from flwr.common.logger import log, warn_deprecated_feature, warn_experimental_feature
+from flwr.common.retry_invoker import RetryInvoker, exponential
 
 from .client_app import load_client_app
 from .grpc_client.connection import grpc_connection
@@ -104,6 +107,8 @@ def run_client_app() -> None:
         transport="rest" if args.rest else "grpc-rere",
         root_certificates=root_certificates,
         insecure=args.insecure,
+        max_retries=args.max_retries,
+        max_wait_time=args.max_wait_time,
     )
     register_exit_handlers(event_type=EventType.RUN_CLIENT_APP_LEAVE)
 
@@ -141,6 +146,22 @@ def _parse_args_run_client_app() -> argparse.ArgumentParser:
         default="0.0.0.0:9092",
         help="Server address",
     )
+    parser.add_argument(
+        "--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,"
+        "it is set to None, meaning there is no limit to the number of tries.",
+    )
+    parser.add_argument(
+        "--max-wait-time",
+        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"
+        "is set to None, meaning there is no limit to the total time.",
+    )
     parser.add_argument(
         "--dir",
         default="",
@@ -180,6 +201,8 @@ def start_client(
     root_certificates: Optional[Union[bytes, str]] = None,
     insecure: Optional[bool] = None,
     transport: Optional[str] = None,
+    max_retries: Optional[int] = None,
+    max_wait_time: Optional[float] = None,
 ) -> None:
     """Start a Flower client node which connects to a Flower server.
 
@@ -213,6 +236,14 @@ def start_client(
         - 'grpc-bidi': gRPC, bidirectional streaming
         - 'grpc-rere': gRPC, request-response (experimental)
         - 'rest': HTTP (experimental)
+    max_retries: Optional[int] (default: None)
+        The maximum number of times the client will try to connect to the
+        server before giving up in case of a connection error. If set to None,
+        there is no limit to the number of tries.
+    max_wait_time: Optional[float] (default: None)
+        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.
 
     Examples
     --------
@@ -254,6 +285,8 @@ def start_client(
         root_certificates=root_certificates,
         insecure=insecure,
         transport=transport,
+        max_retries=max_retries,
+        max_wait_time=max_wait_time,
     )
     event(EventType.START_CLIENT_LEAVE)
 
@@ -272,6 +305,8 @@ def _start_client_internal(
     root_certificates: Optional[Union[bytes, str]] = None,
     insecure: Optional[bool] = None,
     transport: Optional[str] = None,
+    max_retries: Optional[int] = None,
+    max_wait_time: Optional[float] = None,
 ) -> None:
     """Start a Flower client node which connects to a Flower server.
 
@@ -299,7 +334,7 @@ def _start_client_internal(
         The PEM-encoded root certificates as a byte string or a path string.
         If provided, a secure connection using the certificates will be
         established to an SSL-enabled Flower server.
-    insecure : bool (default: True)
+    insecure : Optional[bool] (default: None)
         Starts an insecure gRPC connection when True. Enables HTTPS connection
         when False, using system certificates if `root_certificates` is None.
     transport : Optional[str] (default: None)
@@ -307,6 +342,14 @@ def _start_client_internal(
         - 'grpc-bidi': gRPC, bidirectional streaming
         - 'grpc-rere': gRPC, request-response (experimental)
         - 'rest': HTTP (experimental)
+    max_retries: Optional[int] (default: None)
+        The maximum number of times the client will try to connect to the
+        server before giving up in case of a connection error. If set to None,
+        there is no limit to the number of tries.
+    max_wait_time: Optional[float] (default: None)
+        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.
     """
     if insecure is None:
         insecure = root_certificates is None
@@ -338,7 +381,45 @@ def _start_client_internal(
     # Both `client` and `client_fn` must not be used directly
 
     # Initialize connection context manager
-    connection, address = _init_connection(transport, server_address)
+    connection, address, connection_error_type = _init_connection(
+        transport, server_address
+    )
+
+    retry_invoker = RetryInvoker(
+        wait_factory=exponential,
+        recoverable_exceptions=connection_error_type,
+        max_tries=max_retries,
+        max_time=max_wait_time,
+        on_giveup=lambda retry_state: (
+            log(
+                WARN,
+                "Giving up reconnection after %.2f seconds and %s tries.",
+                retry_state.elapsed_time,
+                retry_state.tries,
+            )
+            if retry_state.tries > 1
+            else None
+        ),
+        on_success=lambda retry_state: (
+            log(
+                INFO,
+                "Connection successful after %.2f seconds and %s tries.",
+                retry_state.elapsed_time,
+                retry_state.tries,
+            )
+            if retry_state.tries > 1
+            else None
+        ),
+        on_backoff=lambda retry_state: (
+            log(WARN, "Connection attempt failed, retrying...")
+            if retry_state.tries == 1
+            else log(
+                DEBUG,
+                "Connection attempt failed, retrying in %.2f seconds",
+                retry_state.actual_wait,
+            )
+        ),
+    )
 
     node_state = NodeState()
 
@@ -347,6 +428,7 @@ def _start_client_internal(
         with connection(
             address,
             insecure,
+            retry_invoker,
             grpc_max_message_length,
             root_certificates,
         ) as conn:
@@ -509,7 +591,7 @@ def start_numpy_client(
 
 def _init_connection(transport: Optional[str], server_address: str) -> Tuple[
     Callable[
-        [str, bool, int, Union[bytes, str, None]],
+        [str, bool, RetryInvoker, int, Union[bytes, str, None]],
         ContextManager[
             Tuple[
                 Callable[[], Optional[Message]],
@@ -520,6 +602,7 @@ def _init_connection(transport: Optional[str], server_address: str) -> Tuple[
         ],
     ],
     str,
+    Type[Exception],
 ]:
     # Parse IP address
     parsed_address = parse_address(server_address)
@@ -535,6 +618,8 @@ def _init_connection(transport: Optional[str], server_address: str) -> Tuple[
     # Use either gRPC bidirectional streaming or REST request/response
     if transport == TRANSPORT_TYPE_REST:
         try:
+            from requests.exceptions import ConnectionError as RequestsConnectionError
+
             from .rest_client.connection import http_request_response
         except ModuleNotFoundError:
             sys.exit(MISSING_EXTRA_REST)
@@ -543,14 +628,14 @@ def _init_connection(transport: Optional[str], server_address: str) -> Tuple[
                 "When using the REST API, please provide `https://` or "
                 "`http://` before the server address (e.g. `http://127.0.0.1:8080`)"
             )
-        connection = http_request_response
+        connection, error_type = http_request_response, RequestsConnectionError
     elif transport == TRANSPORT_TYPE_GRPC_RERE:
-        connection = grpc_request_response
+        connection, error_type = grpc_request_response, RpcError
     elif transport == TRANSPORT_TYPE_GRPC_BIDI:
-        connection = grpc_connection
+        connection, error_type = grpc_connection, RpcError
     else:
         raise ValueError(
             f"Unknown transport type: {transport} (possible: {TRANSPORT_TYPES})"
         )
 
-    return connection, address
+    return connection, address, error_type