| 93 | |
| 94 | # 服务器类 |
| 95 | class MSG_Server: |
| 96 | def __init__(self): |
| 97 | self.clients = {} |
| 98 | self.redis_conn = redis_connection() |
| 99 | self.pubsub = self.redis_conn.pubsub() |
| 100 | self.running = True |
| 101 | # 订阅注册频道 |
| 102 | self.pubsub.subscribe(REGISTER_CHANNEL) |
| 103 | |
| 104 | def register_client(self, client_id): |
| 105 | if client_id not in self.clients: |
| 106 | self.clients[client_id] = { |
| 107 | "status_channel": f"{client_id}_status", |
| 108 | "data_channel": client_id, |
| 109 | "send_channel": f"server_{client_id}" |
| 110 | } |
| 111 | self.pubsub.subscribe(self.clients[client_id]["send_channel"]) |
| 112 | logging.info(f"Server registered client {client_id}") |
| 113 | self.redis_conn.publish(self.clients[client_id]["status_channel"], "REGISTERED") |
| 114 | |
| 115 | def send_message(self, target_client_id, message): |
| 116 | if target_client_id in self.clients: |
| 117 | channel = self.clients[target_client_id]["receive_channel"] |
| 118 | full_message = f"server:{message}" |
| 119 | self.redis_conn.publish(channel, full_message) |
| 120 | print(f"Server sent to {target_client_id}: {full_message}") |
| 121 | |
| 122 | def listen(self): |
| 123 | print(f"Server started, listening for registrations and messages") |
| 124 | while self.running: |
| 125 | message = self.pubsub.get_message(timeout=1.0) |
| 126 | if message and message['type'] == 'message': |
| 127 | channel = message['channel'].decode() |
| 128 | data = message['data'].decode() |
| 129 | if channel == REGISTER_CHANNEL: |
| 130 | # 处理注册消息 |
| 131 | self.register_client(data) |
| 132 | else: |
| 133 | # 处理客户端消息 |
| 134 | client_id = channel.split("_")[-1] |
| 135 | print(f"Server received from {client_id}: {data}") |
| 136 | time.sleep(0.01) |
| 137 | |
| 138 | def send_to_client(self, gpu_info_list, clientID): |
| 139 | # print(f'send to clientID:{clientID}') |
| 140 | logging.info(f'Send to clientID:{clientID}') |
| 141 | client_channel = clientID # 使用 clientID 作为频道 |
| 142 | |
| 143 | # self.msg_server.redis_conn.publish(self.msg_server.clients[client_id]["status_channel"], "ALLOCATED") |
| 144 | self.redis_conn.publish(self.clients[str(client_channel)]["status_channel"], "ALLOCATED") |
| 145 | |
| 146 | if str(clientID) not in self.clients: |
| 147 | # print(f"Client {clientID} not registered yet") |
| 148 | logging.error(f"Client {clientID} not registered yet") |
| 149 | return |
| 150 | print(f'gpu_info_list: {gpu_info_list}') |
| 151 | for gpu_info in gpu_info_list: |
| 152 | if gpu_info is None: |