flwr-nightly 1.19.0.dev20250611__py3-none-any.whl → 1.19.0.dev20250612__py3-none-any.whl

flwr/proto/run_pb2.py

@@ -18,7 +18,7 @@ from flwr.proto import recorddict_pb2 as flwr_dot_proto_dot_recorddict__pb2
 from flwr.proto import transport_pb2 as flwr_dot_proto_dot_transport__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x14\x66lwr/proto/run.proto\x12\nflwr.proto\x1a\x14\x66lwr/proto/fab.proto\x1a\x15\x66lwr/proto/node.proto\x1a\x1b\x66lwr/proto/recorddict.proto\x1a\x1a\x66lwr/proto/transport.proto\"\xce\x02\n\x03Run\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\x12\x0e\n\x06\x66\x61\x62_id\x18\x02 \x01(\t\x12\x13\n\x0b\x66\x61\x62_version\x18\x03 \x01(\t\x12<\n\x0foverride_config\x18\x04 \x03(\x0b\x32#.flwr.proto.Run.OverrideConfigEntry\x12\x10\n\x08\x66\x61\x62_hash\x18\x05 \x01(\t\x12\x12\n\npending_at\x18\x06 \x01(\t\x12\x13\n\x0bstarting_at\x18\x07 \x01(\t\x12\x12\n\nrunning_at\x18\x08 \x01(\t\x12\x13\n\x0b\x66inished_at\x18\t \x01(\t\x12%\n\x06status\x18\n \x01(\x0b\x32\x15.flwr.proto.RunStatus\x1aI\n\x13OverrideConfigEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12!\n\x05value\x18\x02 \x01(\x0b\x32\x12.flwr.proto.Scalar:\x02\x38\x01\"@\n\tRunStatus\x12\x0e\n\x06status\x18\x01 \x01(\t\x12\x12\n\nsub_status\x18\x02 \x01(\t\x12\x0f\n\x07\x64\x65tails\x18\x03 \x01(\t\"?\n\rGetRunRequest\x12\x1e\n\x04node\x18\x01 \x01(\x0b\x32\x10.flwr.proto.Node\x12\x0e\n\x06run_id\x18\x02 \x01(\x04\".\n\x0eGetRunResponse\x12\x1c\n\x03run\x18\x01 \x01(\x0b\x32\x0f.flwr.proto.Run\"S\n\x16UpdateRunStatusRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\x12)\n\nrun_status\x18\x02 \x01(\x0b\x32\x15.flwr.proto.RunStatus\"\x19\n\x17UpdateRunStatusResponse\"F\n\x13GetRunStatusRequest\x12\x1e\n\x04node\x18\x01 \x01(\x0b\x32\x10.flwr.proto.Node\x12\x0f\n\x07run_ids\x18\x02 \x03(\x04\"\xb1\x01\n\x14GetRunStatusResponse\x12L\n\x0frun_status_dict\x18\x01 \x03(\x0b\x32\x33.flwr.proto.GetRunStatusResponse.RunStatusDictEntry\x1aK\n\x12RunStatusDictEntry\x12\x0b\n\x03key\x18\x01 \x01(\x04\x12$\n\x05value\x18\x02 \x01(\x0b\x32\x15.flwr.proto.RunStatus:\x02\x38\x01\"-\n\x1bGetFederationOptionsRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\"T\n\x1cGetFederationOptionsResponse\x12\x34\n\x12\x66\x65\x64\x65ration_options\x18\x01 \x01(\x0b\x32\x18.flwr.proto.ConfigRecordb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x14\x66lwr/proto/run.proto\x12\nflwr.proto\x1a\x14\x66lwr/proto/fab.proto\x1a\x15\x66lwr/proto/node.proto\x1a\x1b\x66lwr/proto/recorddict.proto\x1a\x1a\x66lwr/proto/transport.proto\"\xe0\x02\n\x03Run\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\x12\x0e\n\x06\x66\x61\x62_id\x18\x02 \x01(\t\x12\x13\n\x0b\x66\x61\x62_version\x18\x03 \x01(\t\x12<\n\x0foverride_config\x18\x04 \x03(\x0b\x32#.flwr.proto.Run.OverrideConfigEntry\x12\x10\n\x08\x66\x61\x62_hash\x18\x05 \x01(\t\x12\x12\n\npending_at\x18\x06 \x01(\t\x12\x13\n\x0bstarting_at\x18\x07 \x01(\t\x12\x12\n\nrunning_at\x18\x08 \x01(\t\x12\x13\n\x0b\x66inished_at\x18\t \x01(\t\x12%\n\x06status\x18\n \x01(\x0b\x32\x15.flwr.proto.RunStatus\x12\x10\n\x08\x66lwr_aid\x18\x0b \x01(\t\x1aI\n\x13OverrideConfigEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12!\n\x05value\x18\x02 \x01(\x0b\x32\x12.flwr.proto.Scalar:\x02\x38\x01\"@\n\tRunStatus\x12\x0e\n\x06status\x18\x01 \x01(\t\x12\x12\n\nsub_status\x18\x02 \x01(\t\x12\x0f\n\x07\x64\x65tails\x18\x03 \x01(\t\"?\n\rGetRunRequest\x12\x1e\n\x04node\x18\x01 \x01(\x0b\x32\x10.flwr.proto.Node\x12\x0e\n\x06run_id\x18\x02 \x01(\x04\".\n\x0eGetRunResponse\x12\x1c\n\x03run\x18\x01 \x01(\x0b\x32\x0f.flwr.proto.Run\"S\n\x16UpdateRunStatusRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\x12)\n\nrun_status\x18\x02 \x01(\x0b\x32\x15.flwr.proto.RunStatus\"\x19\n\x17UpdateRunStatusResponse\"F\n\x13GetRunStatusRequest\x12\x1e\n\x04node\x18\x01 \x01(\x0b\x32\x10.flwr.proto.Node\x12\x0f\n\x07run_ids\x18\x02 \x03(\x04\"\xb1\x01\n\x14GetRunStatusResponse\x12L\n\x0frun_status_dict\x18\x01 \x03(\x0b\x32\x33.flwr.proto.GetRunStatusResponse.RunStatusDictEntry\x1aK\n\x12RunStatusDictEntry\x12\x0b\n\x03key\x18\x01 \x01(\x04\x12$\n\x05value\x18\x02 \x01(\x0b\x32\x15.flwr.proto.RunStatus:\x02\x38\x01\"-\n\x1bGetFederationOptionsRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\"T\n\x1cGetFederationOptionsResponse\x12\x34\n\x12\x66\x65\x64\x65ration_options\x18\x01 \x01(\x0b\x32\x18.flwr.proto.ConfigRecordb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -30,27 +30,27 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _globals['_GETRUNSTATUSRESPONSE_RUNSTATUSDICTENTRY']._options = None
   _globals['_GETRUNSTATUSRESPONSE_RUNSTATUSDICTENTRY']._serialized_options = b'8\001'
   _globals['_RUN']._serialized_start=139
-  _globals['_RUN']._serialized_end=473
-  _globals['_RUN_OVERRIDECONFIGENTRY']._serialized_start=400
-  _globals['_RUN_OVERRIDECONFIGENTRY']._serialized_end=473
-  _globals['_RUNSTATUS']._serialized_start=475
-  _globals['_RUNSTATUS']._serialized_end=539
-  _globals['_GETRUNREQUEST']._serialized_start=541
-  _globals['_GETRUNREQUEST']._serialized_end=604
-  _globals['_GETRUNRESPONSE']._serialized_start=606
-  _globals['_GETRUNRESPONSE']._serialized_end=652
-  _globals['_UPDATERUNSTATUSREQUEST']._serialized_start=654
-  _globals['_UPDATERUNSTATUSREQUEST']._serialized_end=737
-  _globals['_UPDATERUNSTATUSRESPONSE']._serialized_start=739
-  _globals['_UPDATERUNSTATUSRESPONSE']._serialized_end=764
-  _globals['_GETRUNSTATUSREQUEST']._serialized_start=766
-  _globals['_GETRUNSTATUSREQUEST']._serialized_end=836
-  _globals['_GETRUNSTATUSRESPONSE']._serialized_start=839
-  _globals['_GETRUNSTATUSRESPONSE']._serialized_end=1016
-  _globals['_GETRUNSTATUSRESPONSE_RUNSTATUSDICTENTRY']._serialized_start=941
-  _globals['_GETRUNSTATUSRESPONSE_RUNSTATUSDICTENTRY']._serialized_end=1016
-  _globals['_GETFEDERATIONOPTIONSREQUEST']._serialized_start=1018
-  _globals['_GETFEDERATIONOPTIONSREQUEST']._serialized_end=1063
-  _globals['_GETFEDERATIONOPTIONSRESPONSE']._serialized_start=1065
-  _globals['_GETFEDERATIONOPTIONSRESPONSE']._serialized_end=1149
+  _globals['_RUN']._serialized_end=491
+  _globals['_RUN_OVERRIDECONFIGENTRY']._serialized_start=418
+  _globals['_RUN_OVERRIDECONFIGENTRY']._serialized_end=491
+  _globals['_RUNSTATUS']._serialized_start=493
+  _globals['_RUNSTATUS']._serialized_end=557
+  _globals['_GETRUNREQUEST']._serialized_start=559
+  _globals['_GETRUNREQUEST']._serialized_end=622
+  _globals['_GETRUNRESPONSE']._serialized_start=624
+  _globals['_GETRUNRESPONSE']._serialized_end=670
+  _globals['_UPDATERUNSTATUSREQUEST']._serialized_start=672
+  _globals['_UPDATERUNSTATUSREQUEST']._serialized_end=755
+  _globals['_UPDATERUNSTATUSRESPONSE']._serialized_start=757
+  _globals['_UPDATERUNSTATUSRESPONSE']._serialized_end=782
+  _globals['_GETRUNSTATUSREQUEST']._serialized_start=784
+  _globals['_GETRUNSTATUSREQUEST']._serialized_end=854
+  _globals['_GETRUNSTATUSRESPONSE']._serialized_start=857
+  _globals['_GETRUNSTATUSRESPONSE']._serialized_end=1034
+  _globals['_GETRUNSTATUSRESPONSE_RUNSTATUSDICTENTRY']._serialized_start=959
+  _globals['_GETRUNSTATUSRESPONSE_RUNSTATUSDICTENTRY']._serialized_end=1034
+  _globals['_GETFEDERATIONOPTIONSREQUEST']._serialized_start=1036
+  _globals['_GETFEDERATIONOPTIONSREQUEST']._serialized_end=1081
+  _globals['_GETFEDERATIONOPTIONSRESPONSE']._serialized_start=1083
+  _globals['_GETFEDERATIONOPTIONSRESPONSE']._serialized_end=1167
 # @@protoc_insertion_point(module_scope)

flwr/proto/run_pb2.pyi

@@ -41,6 +41,7 @@ class Run(google.protobuf.message.Message):
     RUNNING_AT_FIELD_NUMBER: builtins.int
     FINISHED_AT_FIELD_NUMBER: builtins.int
     STATUS_FIELD_NUMBER: builtins.int
+    FLWR_AID_FIELD_NUMBER: builtins.int
     run_id: builtins.int
     fab_id: typing.Text
     fab_version: typing.Text
@@ -53,6 +54,7 @@ class Run(google.protobuf.message.Message):
     finished_at: typing.Text
     @property
     def status(self) -> global___RunStatus: ...
+    flwr_aid: typing.Text
     def __init__(self,
         *,
         run_id: builtins.int = ...,
@@ -65,9 +67,10 @@ class Run(google.protobuf.message.Message):
         running_at: typing.Text = ...,
         finished_at: typing.Text = ...,
         status: typing.Optional[global___RunStatus] = ...,
+        flwr_aid: typing.Text = ...,
         ) -> None: ...
     def HasField(self, field_name: typing_extensions.Literal["status",b"status"]) -> builtins.bool: ...
-    def ClearField(self, field_name: typing_extensions.Literal["fab_hash",b"fab_hash","fab_id",b"fab_id","fab_version",b"fab_version","finished_at",b"finished_at","override_config",b"override_config","pending_at",b"pending_at","run_id",b"run_id","running_at",b"running_at","starting_at",b"starting_at","status",b"status"]) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["fab_hash",b"fab_hash","fab_id",b"fab_id","fab_version",b"fab_version","finished_at",b"finished_at","flwr_aid",b"flwr_aid","override_config",b"override_config","pending_at",b"pending_at","run_id",b"run_id","running_at",b"running_at","starting_at",b"starting_at","status",b"status"]) -> None: ...
 global___Run = Run
 
 class RunStatus(google.protobuf.message.Message):

flwr/proto/serverappio_pb2.py

@@ -20,15 +20,13 @@ from flwr.proto import run_pb2 as flwr_dot_proto_dot_run__pb2
 from flwr.proto import fab_pb2 as flwr_dot_proto_dot_fab__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1c\x66lwr/proto/serverappio.proto\x12\nflwr.proto\x1a\x1a\x66lwr/proto/heartbeat.proto\x1a\x14\x66lwr/proto/log.proto\x1a\x15\x66lwr/proto/node.proto\x1a\x18\x66lwr/proto/message.proto\x1a\x14\x66lwr/proto/run.proto\x1a\x14\x66lwr/proto/fab.proto\"!\n\x0fGetNodesRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\"3\n\x10GetNodesResponse\x12\x1f\n\x05nodes\x18\x01 \x03(\x0b\x32\x10.flwr.proto.Node\"\x8d\x02\n\x16PushInsMessagesRequest\x12*\n\rmessages_list\x18\x01 \x03(\x0b\x32\x13.flwr.proto.Message\x12\x0e\n\x06run_id\x18\x02 \x01(\x04\x12\x61\n\x19msg_to_descendant_mapping\x18\x03 \x03(\x0b\x32>.flwr.proto.PushInsMessagesRequest.MsgToDescendantMappingEntry\x1aT\n\x1bMsgToDescendantMappingEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12$\n\x05value\x18\x02 \x01(\x0b\x32\x15.flwr.proto.ObjectIDs:\x02\x38\x01\"\xcc\x01\n\x17PushInsMessagesResponse\x12\x13\n\x0bmessage_ids\x18\x01 \x03(\t\x12O\n\x0fobjects_to_push\x18\x02 \x03(\x0b\x32\x36.flwr.proto.PushInsMessagesResponse.ObjectsToPushEntry\x1aK\n\x12ObjectsToPushEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12$\n\x05value\x18\x02 \x01(\x0b\x32\x15.flwr.proto.ObjectIDs:\x02\x38\x01\"=\n\x16PullResMessagesRequest\x12\x13\n\x0bmessage_ids\x18\x01 \x03(\t\x12\x0e\n\x06run_id\x18\x02 \x01(\x04\"\xe3\x01\n\x17PullResMessagesResponse\x12*\n\rmessages_list\x18\x01 \x03(\x0b\x32\x13.flwr.proto.Message\x12O\n\x0fobjects_to_pull\x18\x02 \x03(\x0b\x32\x36.flwr.proto.PullResMessagesResponse.ObjectsToPullEntry\x1aK\n\x12ObjectsToPullEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12$\n\x05value\x18\x02 \x01(\x0b\x32\x15.flwr.proto.ObjectIDs:\x02\x38\x01\"\x1c\n\x1aPullServerAppInputsRequest\"\x7f\n\x1bPullServerAppInputsResponse\x12$\n\x07\x63ontext\x18\x01 \x01(\x0b\x32\x13.flwr.proto.Context\x12\x1c\n\x03run\x18\x02 \x01(\x0b\x32\x0f.flwr.proto.Run\x12\x1c\n\x03\x66\x61\x62\x18\x03 \x01(\x0b\x32\x0f.flwr.proto.Fab\"S\n\x1bPushServerAppOutputsRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\x12$\n\x07\x63ontext\x18\x02 \x01(\x0b\x32\x13.flwr.proto.Context\"\x1e\n\x1cPushServerAppOutputsResponse2\xe4\x08\n\x0bServerAppIo\x12G\n\x08GetNodes\x12\x1b.flwr.proto.GetNodesRequest\x1a\x1c.flwr.proto.GetNodesResponse\"\x00\x12Y\n\x0cPushMessages\x12\".flwr.proto.PushInsMessagesRequest\x1a#.flwr.proto.PushInsMessagesResponse\"\x00\x12Y\n\x0cPullMessages\x12\".flwr.proto.PullResMessagesRequest\x1a#.flwr.proto.PullResMessagesResponse\"\x00\x12\x41\n\x06GetRun\x12\x19.flwr.proto.GetRunRequest\x1a\x1a.flwr.proto.GetRunResponse\"\x00\x12\x41\n\x06GetFab\x12\x19.flwr.proto.GetFabRequest\x1a\x1a.flwr.proto.GetFabResponse\"\x00\x12h\n\x13PullServerAppInputs\x12&.flwr.proto.PullServerAppInputsRequest\x1a\'.flwr.proto.PullServerAppInputsResponse\"\x00\x12k\n\x14PushServerAppOutputs\x12\'.flwr.proto.PushServerAppOutputsRequest\x1a(.flwr.proto.PushServerAppOutputsResponse\"\x00\x12\\\n\x0fUpdateRunStatus\x12\".flwr.proto.UpdateRunStatusRequest\x1a#.flwr.proto.UpdateRunStatusResponse\"\x00\x12S\n\x0cGetRunStatus\x12\x1f.flwr.proto.GetRunStatusRequest\x1a .flwr.proto.GetRunStatusResponse\"\x00\x12G\n\x08PushLogs\x12\x1b.flwr.proto.PushLogsRequest\x1a\x1c.flwr.proto.PushLogsResponse\"\x00\x12_\n\x10SendAppHeartbeat\x12#.flwr.proto.SendAppHeartbeatRequest\x1a$.flwr.proto.SendAppHeartbeatResponse\"\x00\x12M\n\nPushObject\x12\x1d.flwr.proto.PushObjectRequest\x1a\x1e.flwr.proto.PushObjectResponse\"\x00\x12M\n\nPullObject\x12\x1d.flwr.proto.PullObjectRequest\x1a\x1e.flwr.proto.PullObjectResponse\"\x00\x62\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1c\x66lwr/proto/serverappio.proto\x12\nflwr.proto\x1a\x1a\x66lwr/proto/heartbeat.proto\x1a\x14\x66lwr/proto/log.proto\x1a\x15\x66lwr/proto/node.proto\x1a\x18\x66lwr/proto/message.proto\x1a\x14\x66lwr/proto/run.proto\x1a\x14\x66lwr/proto/fab.proto\"!\n\x0fGetNodesRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\"3\n\x10GetNodesResponse\x12\x1f\n\x05nodes\x18\x01 \x03(\x0b\x32\x10.flwr.proto.Node\"\x8a\x01\n\x16PushInsMessagesRequest\x12*\n\rmessages_list\x18\x01 \x03(\x0b\x32\x13.flwr.proto.Message\x12\x0e\n\x06run_id\x18\x02 \x01(\x04\x12\x34\n\x14message_object_trees\x18\x03 \x03(\x0b\x32\x16.flwr.proto.ObjectTree\"\xcc\x01\n\x17PushInsMessagesResponse\x12\x13\n\x0bmessage_ids\x18\x01 \x03(\t\x12O\n\x0fobjects_to_push\x18\x02 \x03(\x0b\x32\x36.flwr.proto.PushInsMessagesResponse.ObjectsToPushEntry\x1aK\n\x12ObjectsToPushEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12$\n\x05value\x18\x02 \x01(\x0b\x32\x15.flwr.proto.ObjectIDs:\x02\x38\x01\"=\n\x16PullResMessagesRequest\x12\x13\n\x0bmessage_ids\x18\x01 \x03(\t\x12\x0e\n\x06run_id\x18\x02 \x01(\x04\"\xe3\x01\n\x17PullResMessagesResponse\x12*\n\rmessages_list\x18\x01 \x03(\x0b\x32\x13.flwr.proto.Message\x12O\n\x0fobjects_to_pull\x18\x02 \x03(\x0b\x32\x36.flwr.proto.PullResMessagesResponse.ObjectsToPullEntry\x1aK\n\x12ObjectsToPullEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12$\n\x05value\x18\x02 \x01(\x0b\x32\x15.flwr.proto.ObjectIDs:\x02\x38\x01\"\x1c\n\x1aPullServerAppInputsRequest\"\x7f\n\x1bPullServerAppInputsResponse\x12$\n\x07\x63ontext\x18\x01 \x01(\x0b\x32\x13.flwr.proto.Context\x12\x1c\n\x03run\x18\x02 \x01(\x0b\x32\x0f.flwr.proto.Run\x12\x1c\n\x03\x66\x61\x62\x18\x03 \x01(\x0b\x32\x0f.flwr.proto.Fab\"S\n\x1bPushServerAppOutputsRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\x12$\n\x07\x63ontext\x18\x02 \x01(\x0b\x32\x13.flwr.proto.Context\"\x1e\n\x1cPushServerAppOutputsResponse2\xd7\t\n\x0bServerAppIo\x12G\n\x08GetNodes\x12\x1b.flwr.proto.GetNodesRequest\x1a\x1c.flwr.proto.GetNodesResponse\"\x00\x12Y\n\x0cPushMessages\x12\".flwr.proto.PushInsMessagesRequest\x1a#.flwr.proto.PushInsMessagesResponse\"\x00\x12Y\n\x0cPullMessages\x12\".flwr.proto.PullResMessagesRequest\x1a#.flwr.proto.PullResMessagesResponse\"\x00\x12\x41\n\x06GetRun\x12\x19.flwr.proto.GetRunRequest\x1a\x1a.flwr.proto.GetRunResponse\"\x00\x12\x41\n\x06GetFab\x12\x19.flwr.proto.GetFabRequest\x1a\x1a.flwr.proto.GetFabResponse\"\x00\x12h\n\x13PullServerAppInputs\x12&.flwr.proto.PullServerAppInputsRequest\x1a\'.flwr.proto.PullServerAppInputsResponse\"\x00\x12k\n\x14PushServerAppOutputs\x12\'.flwr.proto.PushServerAppOutputsRequest\x1a(.flwr.proto.PushServerAppOutputsResponse\"\x00\x12\\\n\x0fUpdateRunStatus\x12\".flwr.proto.UpdateRunStatusRequest\x1a#.flwr.proto.UpdateRunStatusResponse\"\x00\x12S\n\x0cGetRunStatus\x12\x1f.flwr.proto.GetRunStatusRequest\x1a .flwr.proto.GetRunStatusResponse\"\x00\x12G\n\x08PushLogs\x12\x1b.flwr.proto.PushLogsRequest\x1a\x1c.flwr.proto.PushLogsResponse\"\x00\x12_\n\x10SendAppHeartbeat\x12#.flwr.proto.SendAppHeartbeatRequest\x1a$.flwr.proto.SendAppHeartbeatResponse\"\x00\x12M\n\nPushObject\x12\x1d.flwr.proto.PushObjectRequest\x1a\x1e.flwr.proto.PushObjectResponse\"\x00\x12M\n\nPullObject\x12\x1d.flwr.proto.PullObjectRequest\x1a\x1e.flwr.proto.PullObjectResponse\"\x00\x12q\n\x16\x43onfirmMessageReceived\x12).flwr.proto.ConfirmMessageReceivedRequest\x1a*.flwr.proto.ConfirmMessageReceivedResponse\"\x00\x62\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
 _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'flwr.proto.serverappio_pb2', _globals)
 if _descriptor._USE_C_DESCRIPTORS == False:
   DESCRIPTOR._options = None
-  _globals['_PUSHINSMESSAGESREQUEST_MSGTODESCENDANTMAPPINGENTRY']._options = None
-  _globals['_PUSHINSMESSAGESREQUEST_MSGTODESCENDANTMAPPINGENTRY']._serialized_options = b'8\001'
   _globals['_PUSHINSMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._options = None
   _globals['_PUSHINSMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._serialized_options = b'8\001'
   _globals['_PULLRESMESSAGESRESPONSE_OBJECTSTOPULLENTRY']._options = None
@@ -38,27 +36,25 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _globals['_GETNODESRESPONSE']._serialized_start=222
   _globals['_GETNODESRESPONSE']._serialized_end=273
   _globals['_PUSHINSMESSAGESREQUEST']._serialized_start=276
-  _globals['_PUSHINSMESSAGESREQUEST']._serialized_end=545
-  _globals['_PUSHINSMESSAGESREQUEST_MSGTODESCENDANTMAPPINGENTRY']._serialized_start=461
-  _globals['_PUSHINSMESSAGESREQUEST_MSGTODESCENDANTMAPPINGENTRY']._serialized_end=545
-  _globals['_PUSHINSMESSAGESRESPONSE']._serialized_start=548
-  _globals['_PUSHINSMESSAGESRESPONSE']._serialized_end=752
-  _globals['_PUSHINSMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._serialized_start=677
-  _globals['_PUSHINSMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._serialized_end=752
-  _globals['_PULLRESMESSAGESREQUEST']._serialized_start=754
-  _globals['_PULLRESMESSAGESREQUEST']._serialized_end=815
-  _globals['_PULLRESMESSAGESRESPONSE']._serialized_start=818
-  _globals['_PULLRESMESSAGESRESPONSE']._serialized_end=1045
-  _globals['_PULLRESMESSAGESRESPONSE_OBJECTSTOPULLENTRY']._serialized_start=970
-  _globals['_PULLRESMESSAGESRESPONSE_OBJECTSTOPULLENTRY']._serialized_end=1045
-  _globals['_PULLSERVERAPPINPUTSREQUEST']._serialized_start=1047
-  _globals['_PULLSERVERAPPINPUTSREQUEST']._serialized_end=1075
-  _globals['_PULLSERVERAPPINPUTSRESPONSE']._serialized_start=1077
-  _globals['_PULLSERVERAPPINPUTSRESPONSE']._serialized_end=1204
-  _globals['_PUSHSERVERAPPOUTPUTSREQUEST']._serialized_start=1206
-  _globals['_PUSHSERVERAPPOUTPUTSREQUEST']._serialized_end=1289
-  _globals['_PUSHSERVERAPPOUTPUTSRESPONSE']._serialized_start=1291
-  _globals['_PUSHSERVERAPPOUTPUTSRESPONSE']._serialized_end=1321
-  _globals['_SERVERAPPIO']._serialized_start=1324
-  _globals['_SERVERAPPIO']._serialized_end=2448
+  _globals['_PUSHINSMESSAGESREQUEST']._serialized_end=414
+  _globals['_PUSHINSMESSAGESRESPONSE']._serialized_start=417
+  _globals['_PUSHINSMESSAGESRESPONSE']._serialized_end=621
+  _globals['_PUSHINSMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._serialized_start=546
+  _globals['_PUSHINSMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._serialized_end=621
+  _globals['_PULLRESMESSAGESREQUEST']._serialized_start=623
+  _globals['_PULLRESMESSAGESREQUEST']._serialized_end=684
+  _globals['_PULLRESMESSAGESRESPONSE']._serialized_start=687
+  _globals['_PULLRESMESSAGESRESPONSE']._serialized_end=914
+  _globals['_PULLRESMESSAGESRESPONSE_OBJECTSTOPULLENTRY']._serialized_start=839
+  _globals['_PULLRESMESSAGESRESPONSE_OBJECTSTOPULLENTRY']._serialized_end=914
+  _globals['_PULLSERVERAPPINPUTSREQUEST']._serialized_start=916
+  _globals['_PULLSERVERAPPINPUTSREQUEST']._serialized_end=944
+  _globals['_PULLSERVERAPPINPUTSRESPONSE']._serialized_start=946
+  _globals['_PULLSERVERAPPINPUTSRESPONSE']._serialized_end=1073
+  _globals['_PUSHSERVERAPPOUTPUTSREQUEST']._serialized_start=1075
+  _globals['_PUSHSERVERAPPOUTPUTSREQUEST']._serialized_end=1158
+  _globals['_PUSHSERVERAPPOUTPUTSRESPONSE']._serialized_start=1160
+  _globals['_PUSHSERVERAPPOUTPUTSRESPONSE']._serialized_end=1190
+  _globals['_SERVERAPPIO']._serialized_start=1193
+  _globals['_SERVERAPPIO']._serialized_end=2432
 # @@protoc_insertion_point(module_scope)

flwr/proto/serverappio_pb2.pyi

@@ -42,36 +42,21 @@ global___GetNodesResponse = GetNodesResponse
 class PushInsMessagesRequest(google.protobuf.message.Message):
     """PushMessages messages"""
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
-    class MsgToDescendantMappingEntry(google.protobuf.message.Message):
-        DESCRIPTOR: google.protobuf.descriptor.Descriptor
-        KEY_FIELD_NUMBER: builtins.int
-        VALUE_FIELD_NUMBER: builtins.int
-        key: typing.Text
-        @property
-        def value(self) -> flwr.proto.message_pb2.ObjectIDs: ...
-        def __init__(self,
-            *,
-            key: typing.Text = ...,
-            value: typing.Optional[flwr.proto.message_pb2.ObjectIDs] = ...,
-            ) -> None: ...
-        def HasField(self, field_name: typing_extensions.Literal["value",b"value"]) -> builtins.bool: ...
-        def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
-
     MESSAGES_LIST_FIELD_NUMBER: builtins.int
     RUN_ID_FIELD_NUMBER: builtins.int
-    MSG_TO_DESCENDANT_MAPPING_FIELD_NUMBER: builtins.int
+    MESSAGE_OBJECT_TREES_FIELD_NUMBER: builtins.int
     @property
     def messages_list(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[flwr.proto.message_pb2.Message]: ...
     run_id: builtins.int
     @property
-    def msg_to_descendant_mapping(self) -> google.protobuf.internal.containers.MessageMap[typing.Text, flwr.proto.message_pb2.ObjectIDs]: ...
+    def message_object_trees(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[flwr.proto.message_pb2.ObjectTree]: ...
     def __init__(self,
         *,
         messages_list: typing.Optional[typing.Iterable[flwr.proto.message_pb2.Message]] = ...,
         run_id: builtins.int = ...,
-        msg_to_descendant_mapping: typing.Optional[typing.Mapping[typing.Text, flwr.proto.message_pb2.ObjectIDs]] = ...,
+        message_object_trees: typing.Optional[typing.Iterable[flwr.proto.message_pb2.ObjectTree]] = ...,
         ) -> None: ...
-    def ClearField(self, field_name: typing_extensions.Literal["messages_list",b"messages_list","msg_to_descendant_mapping",b"msg_to_descendant_mapping","run_id",b"run_id"]) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["message_object_trees",b"message_object_trees","messages_list",b"messages_list","run_id",b"run_id"]) -> None: ...
 global___PushInsMessagesRequest = PushInsMessagesRequest
 
 class PushInsMessagesResponse(google.protobuf.message.Message):

flwr/proto/serverappio_pb2_grpc.py

@@ -84,6 +84,11 @@ class ServerAppIoStub(object):
                 request_serializer=flwr_dot_proto_dot_message__pb2.PullObjectRequest.SerializeToString,
                 response_deserializer=flwr_dot_proto_dot_message__pb2.PullObjectResponse.FromString,
                 )
+        self.ConfirmMessageReceived = channel.unary_unary(
+                '/flwr.proto.ServerAppIo/ConfirmMessageReceived',
+                request_serializer=flwr_dot_proto_dot_message__pb2.ConfirmMessageReceivedRequest.SerializeToString,
+                response_deserializer=flwr_dot_proto_dot_message__pb2.ConfirmMessageReceivedResponse.FromString,
+                )
 
 
 class ServerAppIoServicer(object):
@@ -180,6 +185,13 @@ class ServerAppIoServicer(object):
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
+    def ConfirmMessageReceived(self, request, context):
+        """Confirm Message Received
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
 
 def add_ServerAppIoServicer_to_server(servicer, server):
     rpc_method_handlers = {
@@ -248,6 +260,11 @@ def add_ServerAppIoServicer_to_server(servicer, server):
                     request_deserializer=flwr_dot_proto_dot_message__pb2.PullObjectRequest.FromString,
                     response_serializer=flwr_dot_proto_dot_message__pb2.PullObjectResponse.SerializeToString,
             ),
+            'ConfirmMessageReceived': grpc.unary_unary_rpc_method_handler(
+                    servicer.ConfirmMessageReceived,
+                    request_deserializer=flwr_dot_proto_dot_message__pb2.ConfirmMessageReceivedRequest.FromString,
+                    response_serializer=flwr_dot_proto_dot_message__pb2.ConfirmMessageReceivedResponse.SerializeToString,
+            ),
     }
     generic_handler = grpc.method_handlers_generic_handler(
             'flwr.proto.ServerAppIo', rpc_method_handlers)
@@ -478,3 +495,20 @@ class ServerAppIo(object):
             flwr_dot_proto_dot_message__pb2.PullObjectResponse.FromString,
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def ConfirmMessageReceived(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/flwr.proto.ServerAppIo/ConfirmMessageReceived',
+            flwr_dot_proto_dot_message__pb2.ConfirmMessageReceivedRequest.SerializeToString,
+            flwr_dot_proto_dot_message__pb2.ConfirmMessageReceivedResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)

flwr/proto/serverappio_pb2_grpc.pyi

@@ -78,6 +78,11 @@ class ServerAppIoStub:
         flwr.proto.message_pb2.PullObjectResponse]
     """Pull Object"""
 
+    ConfirmMessageReceived: grpc.UnaryUnaryMultiCallable[
+        flwr.proto.message_pb2.ConfirmMessageReceivedRequest,
+        flwr.proto.message_pb2.ConfirmMessageReceivedResponse]
+    """Confirm Message Received"""
+
 
 class ServerAppIoServicer(metaclass=abc.ABCMeta):
     @abc.abstractmethod
@@ -184,5 +189,13 @@ class ServerAppIoServicer(metaclass=abc.ABCMeta):
         """Pull Object"""
         pass
 
+    @abc.abstractmethod
+    def ConfirmMessageReceived(self,
+        request: flwr.proto.message_pb2.ConfirmMessageReceivedRequest,
+        context: grpc.ServicerContext,
+    ) -> flwr.proto.message_pb2.ConfirmMessageReceivedResponse:
+        """Confirm Message Received"""
+        pass
+
 
 def add_ServerAppIoServicer_to_server(servicer: ServerAppIoServicer, server: grpc.Server) -> None: ...

flwr/server/grid/grpc_grid.py

@@ -28,7 +28,11 @@ from flwr.common.constant import (
     SUPERLINK_NODE_ID,
 )
 from flwr.common.grpc import create_channel, on_channel_state_change
-from flwr.common.inflatable import get_all_nested_objects
+from flwr.common.inflatable import (
+    get_all_nested_objects,
+    get_object_tree,
+    no_object_id_recompute,
+)
 from flwr.common.inflatable_grpc_utils import (
     make_pull_object_fn_grpc,
     make_push_object_fn_grpc,
@@ -43,7 +47,9 @@ from flwr.common.message import remove_content_from_message
 from flwr.common.retry_invoker import _make_simple_grpc_retry_invoker, _wrap_stub
 from flwr.common.serde import message_to_proto, run_from_proto
 from flwr.common.typing import Run
-from flwr.proto.message_pb2 import ObjectIDs  # pylint: disable=E0611
+from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ConfirmMessageReceivedRequest,
+)
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
@@ -213,18 +219,15 @@ class GrpcGrid(Grid):
         """Push one message and its associated objects."""
         # Compute mapping of message descendants
         all_objects = get_all_nested_objects(message)
-        all_object_ids = list(all_objects.keys())
-        msg_id = all_object_ids[-1]  # Last object is the message itself
-        descendant_ids = all_object_ids[:-1]  # All but the last object are descendants
+        msg_id = message.object_id
+        object_tree = get_object_tree(message)
 
         # Call GrpcServerAppIoStub method
         res: PushInsMessagesResponse = self._stub.PushMessages(
             PushInsMessagesRequest(
                 messages_list=[message_to_proto(remove_content_from_message(message))],
                 run_id=run_id,
-                msg_to_descendant_mapping={
-                    msg_id: ObjectIDs(object_ids=descendant_ids)
-                },
+                message_object_trees=[object_tree],
             )
         )
 
@@ -262,7 +265,8 @@ class GrpcGrid(Grid):
                 # Check message
                 self._check_message(msg)
                 # Try pushing message and its objects
-                message_ids.append(self._try_push_message(run_id, msg))
+                with no_object_id_recompute():
+                    message_ids.append(self._try_push_message(run_id, msg))
 
         except grpc.RpcError as e:
             if e.code() == grpc.StatusCode.RESOURCE_EXHAUSTED:  # pylint: disable=E1101
@@ -308,6 +312,13 @@ class GrpcGrid(Grid):
                         run_id=run_id,
                     ),
                 )
+
+                # Confirm that the message has been received
+                self._stub.ConfirmMessageReceived(
+                    ConfirmMessageReceivedRequest(
+                        node=self.node, run_id=run_id, message_object_id=msg_id
+                    )
+                )
                 message = cast(
                     Message, inflate_object_from_contents(msg_id, all_object_contents)
                 )

flwr/server/superlink/fleet/grpc_rere/fleet_servicer.py

@@ -40,6 +40,8 @@ from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatResponse,
 )
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ConfirmMessageReceivedRequest,
+    ConfirmMessageReceivedResponse,
     PullObjectRequest,
     PullObjectResponse,
     PushObjectRequest,
@@ -151,6 +153,7 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             res = message_handler.get_run(
                 request=request,
                 state=self.state_factory.state(),
+                store=self.objectstore_factory.store(),
             )
         except InvalidRunStatusException as e:
             abort_grpc_context(e.message, context)
@@ -167,6 +170,7 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
                 request=request,
                 ffs=self.ffs_factory.ffs(),
                 state=self.state_factory.state(),
+                store=self.objectstore_factory.store(),
             )
         except InvalidRunStatusException as e:
             abort_grpc_context(e.message, context)
@@ -219,3 +223,24 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             abort_grpc_context(e.message, context)
 
         return res
+
+    def ConfirmMessageReceived(
+        self, request: ConfirmMessageReceivedRequest, context: grpc.ServicerContext
+    ) -> ConfirmMessageReceivedResponse:
+        """Confirm message received."""
+        log(
+            DEBUG,
+            "[Fleet.ConfirmMessageReceived] Message with ID '%s' has been received",
+            request.message_object_id,
+        )
+
+        try:
+            res = message_handler.confirm_message_received(
+                request=request,
+                state=self.state_factory.state(),
+                store=self.objectstore_factory.store(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res

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

@@ -44,6 +44,8 @@ from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatResponse,
 )
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ConfirmMessageReceivedRequest,
+    ConfirmMessageReceivedResponse,
     ObjectIDs,
     PullObjectRequest,
     PullObjectResponse,
@@ -146,6 +148,7 @@ def push_messages(
         msg.metadata.run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
+        store,
     )
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
@@ -165,7 +168,9 @@ def push_messages(
     return response
 
 
-def get_run(request: GetRunRequest, state: LinkState) -> GetRunResponse:
+def get_run(
+    request: GetRunRequest, state: LinkState, store: ObjectStore
+) -> GetRunResponse:
     """Get run information."""
     run = state.get_run(request.run_id)
 
@@ -177,6 +182,7 @@ def get_run(request: GetRunRequest, state: LinkState) -> GetRunResponse:
         request.run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
+        store,
     )
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
@@ -193,7 +199,7 @@ def get_run(request: GetRunRequest, state: LinkState) -> GetRunResponse:
 
 
 def get_fab(
-    request: GetFabRequest, ffs: Ffs, state: LinkState  # pylint: disable=W0613
+    request: GetFabRequest, ffs: Ffs, state: LinkState, store: ObjectStore
 ) -> GetFabResponse:
     """Get FAB."""
     # Abort if the run is not running
@@ -201,6 +207,7 @@ def get_fab(
         request.run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
+        store,
     )
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
@@ -220,6 +227,7 @@ def push_object(
         request.run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
+        store,
     )
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
@@ -245,6 +253,7 @@ def pull_object(
         request.run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
+        store,
     )
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
@@ -259,3 +268,25 @@ def pull_object(
             object_content=content,
         )
     return PullObjectResponse(object_found=False, object_available=False)
+
+
+def confirm_message_received(
+    request: ConfirmMessageReceivedRequest,
+    state: LinkState,
+    store: ObjectStore,
+) -> ConfirmMessageReceivedResponse:
+    """Confirm message received handler."""
+    abort_msg = check_abort(
+        request.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+        store,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
+    # Delete the message object
+    store.delete(request.message_object_id)
+    store.delete_message_descendant_ids(request.message_object_id)
+
+    return ConfirmMessageReceivedResponse()

flwr/server/superlink/fleet/rest_rere/rest_api.py

@@ -39,6 +39,8 @@ from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatResponse,
 )
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ConfirmMessageReceivedRequest,
+    ConfirmMessageReceivedResponse,
     PullObjectRequest,
     PullObjectResponse,
     PushObjectRequest,
@@ -176,9 +178,10 @@ async def get_run(request: GetRunRequest) -> GetRunResponse:
     """GetRun."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
 
     # Handle message
-    return message_handler.get_run(request=request, state=state)
+    return message_handler.get_run(request=request, state=state, store=store)
 
 
 @rest_request_response(GetFabRequest)
@@ -189,9 +192,25 @@ async def get_fab(request: GetFabRequest) -> GetFabResponse:
 
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
+
+    # Handle message
+    return message_handler.get_fab(request=request, ffs=ffs, state=state, store=store)
+
+
+@rest_request_response(ConfirmMessageReceivedRequest)
+async def confirm_message_received(
+    request: ConfirmMessageReceivedRequest,
+) -> ConfirmMessageReceivedResponse:
+    """Confirm message received."""
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
 
     # Handle message
-    return message_handler.get_fab(request=request, ffs=ffs, state=state)
+    return message_handler.confirm_message_received(
+        request=request, state=state, store=store
+    )
 
 
 routes = [
@@ -204,6 +223,11 @@ routes = [
     Route("/api/v0/fleet/send-node-heartbeat", send_node_heartbeat, methods=["POST"]),
     Route("/api/v0/fleet/get-run", get_run, methods=["POST"]),
     Route("/api/v0/fleet/get-fab", get_fab, methods=["POST"]),
+    Route(
+        "/api/v0/fleet/confirm-message-received",
+        confirm_message_received,
+        methods=["POST"],
+    ),
 ]
 
 app: Starlette = Starlette(

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -18,6 +18,7 @@
 import threading
 import time
 from bisect import bisect_right
+from collections import defaultdict
 from dataclasses import dataclass, field
 from logging import ERROR, WARNING
 from typing import Optional
@@ -79,6 +80,9 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         self.message_res_store: dict[str, Message] = {}
         self.message_ins_id_to_message_res_id: dict[str, str] = {}
 
+        # Map flwr_aid to run_ids for O(1) reverse index lookup
+        self.flwr_aid_to_run_ids: dict[str, set[int]] = defaultdict(set)
+
         self.node_public_keys: set[bytes] = set()
 
         self.lock = threading.RLock()
@@ -398,6 +402,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         fab_hash: Optional[str],
         override_config: UserConfig,
         federation_options: ConfigRecord,
+        flwr_aid: Optional[str],
     ) -> int:
         """Create a new run for the specified `fab_hash`."""
         # Sample a random int64 as run_id
@@ -421,9 +426,13 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                             sub_status="",
                             details="",
                         ),
+                        flwr_aid=flwr_aid if flwr_aid else "",
                     ),
                 )
                 self.run_ids[run_id] = run_record
+                # Add run_id to the flwr_aid_to_run_ids mapping if flwr_aid is provided
+                if flwr_aid:
+                    self.flwr_aid_to_run_ids[flwr_aid].add(run_id)
 
                 # Record federation options. Leave empty if not passed
                 self.federation_options[run_id] = federation_options
@@ -451,9 +460,15 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         with self.lock:
             return self.node_public_keys.copy()
 
-    def get_run_ids(self) -> set[int]:
-        """Retrieve all run IDs."""
+    def get_run_ids(self, flwr_aid: Optional[str]) -> set[int]:
+        """Retrieve all run IDs if `flwr_aid` is not specified.
+
+        Otherwise, retrieve all run IDs for the specified `flwr_aid`.
+        """
         with self.lock:
+            if flwr_aid is not None:
+                # Return run IDs for the specified flwr_aid
+                return set(self.flwr_aid_to_run_ids.get(flwr_aid, ()))
             return set(self.run_ids.keys())
 
     def _check_and_tag_inactive_run(self, run_ids: set[int]) -> None:

flwr/server/superlink/linkstate/linkstate.py

@@ -164,12 +164,16 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         fab_hash: Optional[str],
         override_config: UserConfig,
         federation_options: ConfigRecord,
+        flwr_aid: Optional[str],
     ) -> int:
         """Create a new run for the specified `fab_hash`."""
 
     @abc.abstractmethod
-    def get_run_ids(self) -> set[int]:
-        """Retrieve all run IDs."""
+    def get_run_ids(self, flwr_aid: Optional[str]) -> set[int]:
+        """Retrieve all run IDs if `flwr_aid` is not specified.
+
+        Otherwise, retrieve all run IDs for the specified `flwr_aid`.
+        """
 
     @abc.abstractmethod
     def get_run(self, run_id: int) -> Optional[Run]: