feat(framework) Update fleet connections to include run_id in `get_…

…fab` (#4693)
adap · Dec 13, 2024 · 2619ef0 · 2619ef0
1 parent 331ed9b
commit 2619ef0
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 10 deletions.
diff --git a/src/py/flwr/client/app.py b/src/py/flwr/client/app.py
@@ -474,7 +474,7 @@ def _on_backoff(retry_state: RetryState) -> None:
 
                     run: Run = runs[run_id]
                     if get_fab is not None and run.fab_hash:
-                        fab = get_fab(run.fab_hash)
+                        fab = get_fab(run.fab_hash, run_id)
                         if not isolation:
                             # If `ClientApp` runs in the same process, install the FAB
                             install_from_fab(fab.content, flwr_path, True)
@@ -752,7 +752,7 @@ def _init_connection(transport: Optional[str], server_address: str) -> tuple[
                 Optional[Callable[[], Optional[int]]],
                 Optional[Callable[[], None]],
                 Optional[Callable[[int], Run]],
-                Optional[Callable[[str], Fab]],
+                Optional[Callable[[str, int], Fab]],
             ]
         ],
     ],

diff --git a/src/py/flwr/client/grpc_adapter_client/connection.py b/src/py/flwr/client/grpc_adapter_client/connection.py
@@ -48,7 +48,7 @@ def grpc_adapter(  # pylint: disable=R0913,too-many-positional-arguments
         Optional[Callable[[], Optional[int]]],
         Optional[Callable[[], None]],
         Optional[Callable[[int], Run]],
-        Optional[Callable[[str], Fab]],
+        Optional[Callable[[str, int], Fab]],
     ]
 ]:
     """Primitives for request/response-based interaction with a server via GrpcAdapter.

diff --git a/src/py/flwr/client/grpc_client/connection.py b/src/py/flwr/client/grpc_client/connection.py
@@ -76,7 +76,7 @@ def grpc_connection(  # pylint: disable=R0913,R0915,too-many-positional-argument
         Optional[Callable[[], Optional[int]]],
         Optional[Callable[[], None]],
         Optional[Callable[[int], Run]],
-        Optional[Callable[[str], Fab]],
+        Optional[Callable[[str, int], Fab]],
     ]
 ]:
     """Establish a gRPC connection to a gRPC server.

diff --git a/src/py/flwr/client/grpc_rere_client/connection.py b/src/py/flwr/client/grpc_rere_client/connection.py
@@ -84,7 +84,7 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         Optional[Callable[[], Optional[int]]],
         Optional[Callable[[], None]],
         Optional[Callable[[int], Run]],
-        Optional[Callable[[str], Fab]],
+        Optional[Callable[[str, int], Fab]],
     ]
 ]:
     """Primitives for request/response-based interaction with a server.
@@ -290,9 +290,9 @@ def get_run(run_id: int) -> Run:
         # Return fab_id and fab_version
         return run_from_proto(get_run_response.run)
 
-    def get_fab(fab_hash: str) -> Fab:
+    def get_fab(fab_hash: str, run_id: int) -> Fab:
         # Call FleetAPI
-        get_fab_request = GetFabRequest(node=node, hash_str=fab_hash)
+        get_fab_request = GetFabRequest(node=node, hash_str=fab_hash, run_id=run_id)
         get_fab_response: GetFabResponse = retry_invoker.invoke(
             stub.GetFab,
             request=get_fab_request,

diff --git a/src/py/flwr/client/rest_client/connection.py b/src/py/flwr/client/rest_client/connection.py
@@ -96,7 +96,7 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         Optional[Callable[[], Optional[int]]],
         Optional[Callable[[], None]],
         Optional[Callable[[int], Run]],
-        Optional[Callable[[str], Fab]],
+        Optional[Callable[[str, int], Fab]],
     ]
 ]:
     """Primitives for request/response-based interaction with a server.
@@ -361,9 +361,9 @@ def get_run(run_id: int) -> Run:
 
         return run_from_proto(res.run)
 
-    def get_fab(fab_hash: str) -> Fab:
+    def get_fab(fab_hash: str, run_id: int) -> Fab:
         # Construct the request
-        req = GetFabRequest(node=node, hash_str=fab_hash)
+        req = GetFabRequest(node=node, hash_str=fab_hash, run_id=run_id)
 
         # Send the request
         res = _request(req, GetFabResponse, PATH_GET_FAB)