refactor(framework): Improve ContextVar handling

Author: chongshenng

framework/py/flwr/superlink/servicer/control/control_account_auth_interceptor.py

@@ -46,11 +46,19 @@
 )
 
 
-shared_account_info: contextvars.ContextVar[AccountInfo] = contextvars.ContextVar(
-    "account_info", default=AccountInfo(flwr_aid=None, account_name=None)
+shared_account_info: contextvars.ContextVar[AccountInfo | None] = (
+    contextvars.ContextVar("account_info", default=None)
 )
 
 
+def get_current_account_info() -> AccountInfo:
+    """Get the current account info from context, or return a default if not set."""
+    account_info = shared_account_info.get()
+    if account_info is None:
+        return AccountInfo(flwr_aid=None, account_name=None)
+    return account_info
+
+
 class ControlAccountAuthInterceptor(grpc.ServerInterceptor):  # type: ignore
     """Control API interceptor for account authentication."""
 

framework/py/flwr/superlink/servicer/control/control_account_auth_interceptor_test.py

@@ -42,6 +42,7 @@
 
 from .control_account_auth_interceptor import (
     ControlAccountAuthInterceptor,
+    get_current_account_info,
     shared_account_info,
 )
 
@@ -99,7 +100,7 @@ def test_unary_unary_authentication_successful(self, request: GrpcMessage) -> No
         # Assert response is as expected
         self.assertEqual(response, "dummy_response")
         # Assert `shared_account_info` is not set
-        account_info_from_context = shared_account_info.get()
+        account_info_from_context = get_current_account_info()
         self.assertIsNone(account_info_from_context.flwr_aid)
         self.assertIsNone(account_info_from_context.account_name)
 
@@ -206,7 +207,7 @@ def test_unary_validate_tokens_successful(self, request: GrpcMessage) -> None:
             self.assertEqual(response, "dummy_response")
 
         # Assert `shared_account_info` is set
-        account_info_from_context = shared_account_info.get()
+        account_info_from_context = get_current_account_info()
         self.assertEqual(
             account_info_from_context.flwr_aid, self.expected_account_info.flwr_aid
         )

framework/py/flwr/superlink/servicer/control/control_event_log_interceptor.py

@@ -24,7 +24,7 @@
 from flwr.common.event_log_plugin.event_log_plugin import EventLogWriterPlugin
 from flwr.common.typing import LogEntry
 
-from .control_account_auth_interceptor import shared_account_info
+from .control_account_auth_interceptor import get_current_account_info
 
 
 class ControlEventLogInterceptor(grpc.ServerInterceptor):  # type: ignore
@@ -66,7 +66,7 @@ def _generic_method_handler(
             log_entry = self.log_plugin.compose_log_before_event(
                 request=request,
                 context=context,
-                account_info=shared_account_info.get(),
+                account_info=get_current_account_info(),
                 method_name=method_name,
             )
             self.log_plugin.write_log(log_entry)
@@ -85,7 +85,7 @@ def _generic_method_handler(
                     log_entry = self.log_plugin.compose_log_after_event(
                         request=request,
                         context=context,
-                        account_info=shared_account_info.get(),
+                        account_info=get_current_account_info(),
                         method_name=method_name,
                         response=unary_response or error,
                     )
@@ -115,7 +115,7 @@ def response_wrapper() -> Iterator[GrpcMessage]:
                         log_entry = self.log_plugin.compose_log_after_event(
                             request=request,
                             context=context,
-                            account_info=shared_account_info.get(),
+                            account_info=get_current_account_info(),
                             method_name=method_name,
                             response=stream_response or error,
                         )

framework/py/flwr/superlink/servicer/control/control_servicer.py

@@ -80,7 +80,7 @@
 from flwr.superlink.artifact_provider import ArtifactProvider
 from flwr.superlink.auth_plugin import ControlAuthnPlugin
 
-from .control_account_auth_interceptor import shared_account_info
+from .control_account_auth_interceptor import get_current_account_info
 
 
 class ControlServicer(control_pb2_grpc.ControlServicer):
@@ -118,7 +118,7 @@ def StartRun(  # pylint: disable=too-many-locals
             )
             return StartRunResponse()
 
-        flwr_aid = shared_account_info.get().flwr_aid
+        flwr_aid = get_current_account_info().flwr_aid
         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)
@@ -215,7 +215,7 @@ def StreamLogs(  # pylint: disable=C0103
             context.abort(grpc.StatusCode.NOT_FOUND, RUN_ID_NOT_FOUND_MESSAGE)
 
         # Check if `flwr_aid` matches the run's `flwr_aid`
-        flwr_aid = shared_account_info.get().flwr_aid
+        flwr_aid = get_current_account_info().flwr_aid
         _check_flwr_aid_in_run(flwr_aid=flwr_aid, run=cast(Run, run), context=context)
 
         after_timestamp = request.after_timestamp + 1e-6
@@ -254,7 +254,7 @@ def ListRuns(
         if not request.HasField("run_id"):
             # If no `run_id` is specified and account auth is enabled,
             # return run IDs for the authenticated account
-            flwr_aid = shared_account_info.get().flwr_aid
+            flwr_aid = get_current_account_info().flwr_aid
             _check_flwr_aid_exists(flwr_aid, context)
             run_ids = state.get_run_ids(flwr_aid=flwr_aid)
         # Build a set of run IDs for `flwr ls --run-id <run_id>`
@@ -269,7 +269,7 @@ def ListRuns(
                 raise grpc.RpcError()  # This line is unreachable
 
             # Check if `flwr_aid` matches the run's `flwr_aid`
-            flwr_aid = shared_account_info.get().flwr_aid
+            flwr_aid = get_current_account_info().flwr_aid
             _check_flwr_aid_in_run(flwr_aid=flwr_aid, run=run, context=context)
 
             run_ids = {run_id}
@@ -295,7 +295,7 @@ def StopRun(
             raise grpc.RpcError()  # This line is unreachable
 
         # Check if `flwr_aid` matches the run's `flwr_aid`
-        flwr_aid = shared_account_info.get().flwr_aid
+        flwr_aid = get_current_account_info().flwr_aid
         _check_flwr_aid_in_run(flwr_aid=flwr_aid, run=run, context=context)
 
         run_status = state.get_run_status({run_id})[run_id]
@@ -405,7 +405,7 @@ def PullArtifacts(
             )
 
         # Check if `flwr_aid` matches the run's `flwr_aid`
-        flwr_aid = shared_account_info.get().flwr_aid
+        flwr_aid = get_current_account_info().flwr_aid
         _check_flwr_aid_in_run(flwr_aid=flwr_aid, run=run, context=context)
 
         # Call artifact provider
@@ -435,10 +435,10 @@ def RegisterNode(
         state = self.linkstate_factory.state()
         node_id = 0
 
-        flwr_aid = shared_account_info.get().flwr_aid
+        flwr_aid = get_current_account_info().flwr_aid
         flwr_aid = _check_flwr_aid_exists(flwr_aid, context)
         # Account name exists if `flwr_aid` exists
-        account_name = cast(str, shared_account_info.get().account_name)
+        account_name = cast(str, get_current_account_info().account_name)
         try:
             node_id = state.create_node(
                 owner_aid=flwr_aid,
@@ -466,7 +466,7 @@ def UnregisterNode(
         # Init link state
         state = self.linkstate_factory.state()
 
-        flwr_aid = shared_account_info.get().flwr_aid
+        flwr_aid = get_current_account_info().flwr_aid
         flwr_aid = _check_flwr_aid_exists(flwr_aid, context)
         try:
             state.delete_node(owner_aid=flwr_aid, node_id=request.node_id)
@@ -494,7 +494,7 @@ def ListNodes(
         # Init link state
         state = self.linkstate_factory.state()
 
-        flwr_aid = shared_account_info.get().flwr_aid
+        flwr_aid = get_current_account_info().flwr_aid
         flwr_aid = _check_flwr_aid_exists(flwr_aid, context)
         # Retrieve all nodes for the account
         nodes_info = state.get_node_info(owner_aids=[flwr_aid])
@@ -510,7 +510,7 @@ def ListFederations(
         # Init link state
         state = self.linkstate_factory.state()
 
-        flwr_aid = shared_account_info.get().flwr_aid
+        flwr_aid = get_current_account_info().flwr_aid
         flwr_aid = _check_flwr_aid_exists(flwr_aid, context)
 
         # Get federations the account is a member of

framework/py/flwr/superlink/servicer/control/control_servicer_test.py

@@ -283,10 +283,8 @@ def test_list_nodes_cli(self, flwr_aid_retrieving: str, expected: bool) -> None:
 
         # Execute
         with patch(
-            "flwr.superlink.servicer.control.control_servicer.shared_account_info",
-            new=SimpleNamespace(
-                get=lambda: SimpleNamespace(flwr_aid=flwr_aid_retrieving)
-            ),
+            "flwr.superlink.servicer.control.control_servicer.get_current_account_info",
+            return_value=SimpleNamespace(flwr_aid=flwr_aid_retrieving),
         ):
             res: ListNodesResponse = self.servicer.ListNodes(ListNodesRequest(), Mock())
 
@@ -349,8 +347,8 @@ def test_streamlogs_auth_unsucessful(
 
         # Execute & Assert
         with patch(
-            "flwr.superlink.servicer.control.control_servicer.shared_account_info",
-            new=SimpleNamespace(get=lambda: SimpleNamespace(flwr_aid=context_flwr_aid)),
+            "flwr.superlink.servicer.control.control_servicer.get_current_account_info",
+            return_value=SimpleNamespace(flwr_aid=context_flwr_aid),
         ):
             gen = self.servicer.StreamLogs(request, ctx)
             with self.assertRaises(RuntimeError) as cm:
@@ -378,8 +376,8 @@ def test_streamlogs_auth_successful(self) -> None:
                 },
             ),
             patch(
-                "flwr.superlink.servicer.control.control_servicer.shared_account_info",
-                new=SimpleNamespace(get=lambda: SimpleNamespace(flwr_aid="user-123")),
+                "flwr.superlink.servicer.control.control_servicer.get_current_account_info",
+                return_value=SimpleNamespace(flwr_aid="user-123"),
             ),
         ):
             msgs = list(self.servicer.StreamLogs(request, ctx))
@@ -403,8 +401,8 @@ def test_stoprun_auth_unsuccessful(
 
         # Execute & Assert
         with patch(
-            "flwr.superlink.servicer.control.control_servicer.shared_account_info",
-            new=SimpleNamespace(get=lambda: SimpleNamespace(flwr_aid=context_flwr_aid)),
+            "flwr.superlink.servicer.control.control_servicer.get_current_account_info",
+            return_value=SimpleNamespace(flwr_aid=context_flwr_aid),
         ):
             with self.assertRaises(RuntimeError) as cm:
                 self.servicer.StopRun(request, ctx)
@@ -419,8 +417,8 @@ def test_stoprun_auth_successful(self) -> None:
 
         # Execute & Assert
         with patch(
-            "flwr.superlink.servicer.control.control_servicer.shared_account_info",
-            new=SimpleNamespace(get=lambda: SimpleNamespace(flwr_aid="user-123")),
+            "flwr.superlink.servicer.control.control_servicer.get_current_account_info",
+            return_value=SimpleNamespace(flwr_aid="user-123"),
         ):
             response = self.servicer.StopRun(request, ctx)
             self.assertTrue(response.success)
@@ -441,8 +439,8 @@ def test_listruns_auth_unsuccessful(
 
         # Execute & Assert
         with patch(
-            "flwr.superlink.servicer.control.control_servicer.shared_account_info",
-            new=SimpleNamespace(get=lambda: SimpleNamespace(flwr_aid=context_flwr_aid)),
+            "flwr.superlink.servicer.control.control_servicer.get_current_account_info",
+            return_value=SimpleNamespace(flwr_aid=context_flwr_aid),
         ):
             with self.assertRaises(RuntimeError) as cm:
                 self.servicer.ListRuns(request, ctx)
@@ -457,8 +455,8 @@ def test_listruns_auth_run_success(self) -> None:
 
         # Execute & Assert
         with patch(
-            "flwr.superlink.servicer.control.control_servicer.shared_account_info",
-            new=SimpleNamespace(get=lambda: SimpleNamespace(flwr_aid="user-123")),
+            "flwr.superlink.servicer.control.control_servicer.get_current_account_info",
+            return_value=SimpleNamespace(flwr_aid="user-123"),
         ):
             response = self.servicer.ListRuns(request, ctx)
             self.assertEqual(set(response.run_dict.keys()), {run_id})