Bayes_DFL/main.py at production · bayesianinstitute/Bayes_DFL · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
from core.MqttOPS.mqtt_operations import MqttOperations
import argparse

from core.Logs_System.logger import Logger

from core.API.ClientAPI import ApiClient
from core.API.endpoint import *
import json


# Define a class for the Federated Learning Workflow
class DFLWorkflow:
    def __init__(self, broker_service, internal_cluster_topic):

        self.ip=self.get_public_ip()
        # Setup Logger
        self.logger=Logger(name='DFL_logger',api_endpoint=f"{self.ip}:8000/{update_logs}").get_logger()
        # Initialize various attributes and parameters

        self.broker_service = broker_service
        self.internal_cluster_topic = internal_cluster_topic
        self.mqtt_operations = None
        self.apiClient=ApiClient(ip=self.ip)

    def get_public_ip(self):
        import requests

        try:
            response = requests.get("https://api64.ipify.org?format=json")
            if response.status_code == 200:
                public_ip = response.json()["ip"]
                return f"http://{public_ip}"
            else:
                print(f"Failed to retrieve public IP. Status code: {response.status_code}")
        except Exception as e:
            print(f"Error retrieving public IP: {e}")

    def update_role_status(self,data,role):

        try:
            if role =='User':

                connected_status = self.apiClient.put_request(update_user,data)

                if connected_status.status_code == 200:
                    self.logger.info(f"PUT  User Role_status Request Successful: {connected_status.text}")
                    return json.loads(connected_status.text)
                else:
                    self.logger.info(f"PUT  User Request Failed: {connected_status.status_code, connected_status.text}")
                    return None

            elif role =='Admin':

                connected_status = self.apiClient.put_request(update_admin,data)

                if connected_status.status_code == 200:
                    self.logger.info(f"PUT Admin Role_status Request Successful: {connected_status.text}")
                    return json.loads(connected_status.text)
                else:
                    self.logger.info(f"PUT Admin Request Failed: {connected_status.status_code, connected_status.text}")
                    return None
            else:
                import sys
                self.logger.critical(F"unknow role: {role} failed")
                sys.exit(1)

        except Exception as e:
                self.logger.error(f"Error in update_network_status: {str(e)}")
                return None
        else:
            pass

    # Main function to run the federated learning workflow
    def run(self,role='User'):

        self.logger.debug(self.internal_cluster_topic)

        self.logger.info(f"Your IP address is {self.ip}")

        data={
            "role":role
        }


        role_data=self.update_role_status(data,role)


        # Initialize  MQTT operations for communication
        self.mqtt_operations = MqttOperations(self.ip,self.internal_cluster_topic,
                                            self.broker_service,)

        if role_data['role'] == "Admin":
            from core.Role.Admin import Admin

            self.logger.info(f"Role Admin")

            model_type = 'CNN'

            optimizer = "Adam"

            # Fetch from database
            data={
                    "model_name": model_type,
                    "dataset_name": "Mnist",
                    "optimizer": optimizer,
                    "training_name": self.internal_cluster_topic }

            self.logger.warning(data)

            post_response=self.apiClient.post_request(create_training_information,data)

            if post_response.status_code == 201:
                self.logger.info(f"POST Request Successful: {post_response.text}" )
            else:
                self.logger.error(f"POST Request Failed:{ post_response.status_code, post_response.text}")

            admin = Admin(self.internal_cluster_topic , model_type, optimizer,self.mqtt_operations,self.ip,role=role_data['role'])
            admin.admin_logic()
        # User
        elif role_data['role'] == "User":
            from core.Role.User import User

            self.logger.info(f"Role User")
            user = User( self.internal_cluster_topic,self.mqtt_operations,self.ip,role=role_data['role'] )
            user.user_logic()
        else:
             pass

if __name__ == "__main__":
    # Parse command-line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("broker_service", help="Name of broker service", type=str)
    parser.add_argument("cluster_name", help="Name of the cluster", type=str)
    parser.add_argument("role", help="Name of role", type=str)


    args = parser.parse_args()


    internal_cluster_topic=f'{args.cluster_name}'

    workflow = DFLWorkflow(args.broker_service,  internal_cluster_topic)
    workflow.run(args.role)