def register(self, conn, head_dir):
username, password = head_dir['username'], head_dir['password']
if username not in self.all_clients or password != self.all_clients[
username]:
send_head_dir(conn=conn, head_dir=json.dumps({'msg': "error"}))
else:
send_head_dir(conn=conn, head_dir=json.dumps({'msg': "ok"}))
self.logger.info(username + " successfully registered!")
self.logger.info("Wait a bit longer for other clients to join.")
send_file(conn=conn,
file_path=self.configs["model_path"],
new_file_name=None)
sleep(20)
def recv_model(self, conn, head_dir):
username = head_dir["username"]
status = self.clients_status[username]
if status != 0 or (not self.lock.acquire(blocking=False)):
send_head_dir(conn=conn, head_dir=json.dumps({'msg': "error"}))
else:
try:
send_head_dir(conn=conn, head_dir=json.dumps({'msg': "ok"}))
recv_and_write_file(conn=conn,
file_dir=self.client_weight_dir,
buff_size=self.configs['buff_size'])
self.clients_status[username] = 1
self.logger.info("received model from " + username)
finally:
self.lock.release()
def handle_request(self):
while True:
conn, addr = self.recv_socket.accept()
try:
if self.finish:
send_head_dir(conn=conn,
head_dir=json.dumps({'msg': "finish"}))
head_dir = recv_head_dir(conn=conn)
msg = head_dir["msg"]
if msg == "register":
self.register(conn=conn, head_dir=head_dir)
elif msg == "request_model":
self.send_model(conn=conn, head_dir=head_dir)
elif msg == "send_model":
self.recv_model(conn=conn, head_dir=head_dir)
finally:
conn.close()
def send_model(self, conn, head_dir, _model_path=None):
username = head_dir["username"]
status = self.clients_status[username]
if status == 1 or (not self.lock.acquire(blocking=False)):
send_head_dir(conn=conn, head_dir=json.dumps({'msg': "wait"}))
else:
try:
if status == -1:
send_head_dir(conn=conn,
head_dir=json.dumps({'msg': "ok"}))
model_path = self.model_path if _model_path is None else _model_path
send_file(conn=conn,
file_path=model_path,
new_file_name=None)
self.clients_status[username] = 0
self.logger.info("sent model to " + username)
else:
send_head_dir(conn=conn,
head_dir=json.dumps({'msg': "error"}))
self.clients_status[username] = -1
finally:
self.lock.release()