feature node update-done format

Author: AlejandroAvilesSerrano

nebula/addons/reporter.py

@@ -4,6 +4,7 @@
 import logging
 import os
 import sys
+from nebula.controller.federation.utils_requests import NodeUpdateRequest, NodeDoneRequest
 from typing import TYPE_CHECKING
 
 import aiohttp
@@ -171,10 +172,17 @@ async def report_scenario_finished(self):
             - Logs exceptions if the connection attempt to the controller fails.
         """
         url = f"http://{self.config.participant['scenario_args']['controller']}/nodes/{self.config.participant['scenario_args']['federation_id']}/done"
-        data = json.dumps({"idx": self.config.participant["device_args"]["idx"],
-                           "deployment": self.config.participant["scenario_args"]["deployment"],
-                           "name": self.config.participant["scenario_args"]["name"],
-                           "federation_id": self.config.participant["scenario_args"]["federation_id"]})
+        node_done_req = NodeDoneRequest(idx=self.config.participant["device_args"]["idx"],
+                                        deployment=self.config.participant["scenario_args"]["deployment"],
+                                        name=self.config.participant["scenario_args"]["name"],
+                                        federation_id=self.config.participant["scenario_args"]["federation_id"]
+                        )
+        payload = node_done_req.model_dump()
+        data = json.dumps(payload)
+        # data = json.dumps({"idx": self.config.participant["device_args"]["idx"],
+        #                    "deployment": self.config.participant["scenario_args"]["deployment"],
+        #                    "name": self.config.participant["scenario_args"]["name"],
+        #                    "federation_id": self.config.participant["scenario_args"]["federation_id"]})
         headers = {
             "Content-Type": "application/json",
             "User-Agent": f"NEBULA Participant {self.config.participant['device_args']['idx']}",
@@ -266,11 +274,13 @@ async def __report_status_to_controller(self):
             - Delays for 5 seconds upon general exceptions to avoid rapid retry loops.
         """
         try:
+            node_updt_req = NodeUpdateRequest(config=self.config.participant)
+            payload = node_updt_req.model_dump()
             async with (
                 aiohttp.ClientSession() as session,
                 session.post(
                     self.url,
-                    data=json.dumps(self.config.participant),
+                    data=json.dumps(payload),
                     headers={
                         "Content-Type": "application/json",
                         "User-Agent": f"NEBULA Participant {self.config.participant['device_args']['idx']}",

nebula/controller/federation/controllers/docker_federation_controller.py

@@ -8,6 +8,7 @@
 from nebula.controller.federation.federation_controller import FederationController
 from nebula.controller.federation.scenario_builder import ScenarioBuilder
 from nebula.controller.federation.utils_requests import factory_requests_path
+from nebula.controller.federation.utils_requests import NodeUpdateRequest, NodeDoneRequest
 from typing import Dict
 from fastapi import Request
 from nebula.config.config import Config
@@ -200,8 +201,8 @@ async def stop_scenario(self, federation_id: str):
 
         return True #TODO care about cases
 
-    async def update_nodes(self, federation_id: str, request: Request):
-        config = await request.json()
+    async def update_nodes(self, federation_id: str, node_update_request: NodeUpdateRequest):
+        config = node_update_request.config
         scenario_name = config["scenario_args"]["name"]
         fed_id = config["scenario_args"]["federation_id"]
 
@@ -229,27 +230,24 @@ async def update_nodes(self, federation_id: str, request: Request):
                                     nebula_federation.last_index_deployed += 1
                                     #additionals.remove(index)
                                     adds_deployed.add(index)
-            request_body = await request.json()
-            payload = {"scenario_name": scenario_name, "data": request_body}
+            payload = node_update_request.model_dump()
             asyncio.create_task(self._send_to_hub("update", payload, federation_id=fed_id))
             return {"message": "Node updated successfully in Federation Controller"}
         except Exception as e:
             self.logger.info(f"ERROR: federation ID: ({fed_id}), {e}")
             return {"message": "Node updated failed in Federation Controller"} 
 
-    async def node_done(self, federation_id: str, request: Request):
-        request_body = await request.json()
-        scenario_name = request_body["scenario_args"]["name"]
+    async def node_done(self, federation_id: str, node_done_request: NodeDoneRequest):
         nebula_federation = self.nfp[federation_id]
         self.logger.info(f"Node-Done received from node on federation ID: ({federation_id})")
 
         if await nebula_federation.is_experiment_finish():
-            payload = {"federation_id": federation_id, "scenario_name": scenario_name, "data": request_body}
+            payload = node_done_request.model_dump()
             self.logger.info(f"All nodes have finished on federation ID: ({federation_id}), reporting to hub..")
             await self._remove_nebula_federation_from_pool(federation_id)
             asyncio.create_task(self._send_to_hub("finish", payload, federation_id=federation_id))
 
-        payload = {"scenario_name": scenario_name, "data": request_body}
+        payload = node_done_request.model_dump()
         asyncio.create_task(self._send_to_hub("done", payload, federation_id=federation_id))
         return {"message": "Nodes done received successfully"}
 

nebula/controller/federation/controllers/processes_federation_controller.py

@@ -8,6 +8,7 @@
 from nebula.controller.federation.federation_controller import FederationController
 from nebula.controller.federation.scenario_builder import ScenarioBuilder
 from nebula.controller.federation.utils_requests import factory_requests_path
+from nebula.controller.federation.utils_requests import NodeUpdateRequest, NodeDoneRequest
 from typing import Dict
 from fastapi import Request
 from nebula.config.config import Config
@@ -161,8 +162,8 @@ async def stop_scenario(self, federation_id: str = ""):
             self.logger.exception(f"Error while removing current_scenario_commands.sh file: {e}")
         return False
 
-    async def update_nodes(self, federation_id: str, request: Request):
-        config = await request.json()
+    async def update_nodes(self, federation_id: str, node_update_request: NodeUpdateRequest):
+        config = node_update_request.config
         fed_id = config["scenario_args"]["federation_id"]
         scenario_name = config["scenario_args"]["name"]
 
@@ -189,27 +190,24 @@ async def update_nodes(self, federation_id: str, request: Request):
                                     nebula_federation.last_index_deployed += 1
                                     additionals.remove(index)
                                     adds_deployed.add(index)
-            request_body = await request.json()
-            payload = {"scenario_name": scenario_name, "data": request_body}
+            payload = node_update_request.model_dump()
             asyncio.create_task(self._send_to_hub("update", payload, federation_id=fed_id))
             return {"message": "Node updated successfully in Federation Controller"}
         except Exception as e:
             self.logger.info(f"ERROR: federation ID: ({fed_id}) not found on pool..")
             return {"message": "Node updated failed in Federation Controller, ID not found.."}
 
-    async def node_done(self, federation_id: str, request: Request):
-        request_body = await request.json()
-        scenario_name = request_body["scenario_args"]["name"]
+    async def node_done(self, federation_id: str, node_done_request: NodeDoneRequest):
         nebula_federation = self.nfp[federation_id]
         self.logger.info(f"Node-Done received from node on federation ID: ({federation_id})")
 
         if await nebula_federation.is_experiment_finish():
-            payload = {"federation_id": federation_id, "scenario_name": scenario_name, "data": request_body}
+            payload = node_done_request.model_dump()
             self.logger.info(f"All nodes have finished on federation ID: ({federation_id}), reporting to hub..")
             await self._remove_nebula_federation_from_pool(federation_id)
             asyncio.create_task(self._send_to_hub("finish", payload, federation_id=federation_id))
 
-        payload = {"scenario_name": scenario_name, "data": request_body}
+        payload = node_done_request.model_dump()
         asyncio.create_task(self._send_to_hub("done", payload, federation_id=federation_id))
         return {"message": "Nodes done received successfully"}
 

nebula/controller/federation/federation_api.py

@@ -10,7 +10,7 @@
 from nebula.utils import LoggerUtils
 from nebula.controller.federation.federation_controller import FederationController 
 from nebula.controller.federation.factory_federation_controller import federation_controller_factory
-from nebula.controller.federation.utils_requests import RunScenarioRequest, StopScenarioRequest, Routes
+from nebula.controller.federation.utils_requests import RunScenarioRequest, StopScenarioRequest, NodeUpdateRequest, NodeDoneRequest, Routes
 
 fed_controllers: Dict[str, FederationController] = {}
 
@@ -61,7 +61,7 @@ async def run_scenario(run_scenario_request: RunScenarioRequest):
     if controller:
         return await controller.run_scenario(run_scenario_request.federation_id, run_scenario_request.scenario_data, run_scenario_request.user)
     else:
-        return {"message": "Experyment type not allowed"}
+        return {"message": "Experiment type not allowed"}
 
 @app.post(Routes.STOP)
 async def stop_scenario(stop_scenario_request: StopScenarioRequest):
@@ -73,35 +73,33 @@ async def stop_scenario(stop_scenario_request: StopScenarioRequest):
     if controller:
         return await controller.stop_scenario(stop_scenario_request.federation_id)
     else:
-        return {"message": "Experyment type not allowed"}
+        return {"message": "Experiment type not allowed"}
 
 @app.post(Routes.UPDATE)
 async def update_nodes(
     federation_id: str,
-    request: Request,
+    node_update_request: NodeUpdateRequest,
 ):
     global fed_controllers
-    config = await request.json()
-    experiment_type = config["scenario_args"]["deployment"]
+    experiment_type = node_update_request.config["scenario_args"]["deployment"]
     controller = fed_controllers.get(experiment_type, None)
     if controller:
-        return await controller.update_nodes(federation_id, request)
+        return await controller.update_nodes(federation_id, node_update_request)
     else:
-        return {"message": "Experyment type not allowed on response for update message.."}
+        return {"message": "Experiment type not allowed on response for update message.."}
 
 @app.post(Routes.DONE)
-async def update_nodes(
+async def node_done(
     federation_id: str,
-    request: Request,
+    node_done_request: NodeDoneRequest,
 ):
     global fed_controllers
-    config = await request.json()
-    experiment_type = config["deployment"]
+    experiment_type = node_done_request.deployment
     controller = fed_controllers.get(experiment_type, None)
     if controller:
-        return await controller.node_done(federation_id, request)
+        return await controller.node_done(federation_id, node_done_request)
     else:
-        return {"message": "Experyment type not allowed on responde for Node done message.."}
+        return {"message": "Experiment type not allowed on responde for Node done message.."}
 
 if __name__ == "__main__":
     # Parse args from command line

nebula/controller/federation/federation_controller.py

@@ -2,6 +2,7 @@
 from fastapi import Request
 from typing import Dict
 from nebula.controller.federation.scenario_builder import ScenarioBuilder
+from nebula.controller.federation.utils_requests import NodeUpdateRequest, NodeDoneRequest
 import logging 
 
 class NebulaFederation(ABC):
@@ -26,9 +27,9 @@ async def stop_scenario(self, federation_id: str):
         pass
 
     @abstractmethod
-    async def update_nodes(self, federation_id: str, request: Request):
+    async def update_nodes(self, federation_id: str, node_update_request: NodeUpdateRequest):
         pass
 
     abstractmethod
-    async def node_done(self, federation_id: str, request: Request):
+    async def node_done(self, federation_id: str, node_done_request: NodeDoneRequest):
         pass

nebula/controller/federation/utils_requests.py

@@ -10,6 +10,15 @@ class StopScenarioRequest(BaseModel):
     experiment_type: str
     federation_id: str
 
+class NodeUpdateRequest(BaseModel):
+    config: Dict[str, Any] = {}    
+    
+class NodeDoneRequest(BaseModel):
+    idx: int
+    deployment: str
+    name: str
+    federation_id: str
+    
 class Routes:
     INIT = "/init"
     RUN = "/scenarios/run"