forked from mindspore-Ecosystem/mindspore
279 lines
11 KiB
Python
279 lines
11 KiB
Python
# Copyright 2021 Huawei Technologies Co., Ltd
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ============================================================================
|
|
|
|
import argparse
|
|
import time
|
|
import datetime
|
|
import random
|
|
import sys
|
|
import requests
|
|
import flatbuffers
|
|
import numpy as np
|
|
from mindspore.schema import (RequestFLJob, ResponseFLJob, ResponseCode,
|
|
RequestUpdateModel, ResponseUpdateModel,
|
|
FeatureMap, RequestGetModel, ResponseGetModel)
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--pid", type=int, default=0)
|
|
parser.add_argument("--http_ip", type=str, default="10.113.216.106")
|
|
parser.add_argument("--http_port", type=int, default=6666)
|
|
parser.add_argument("--use_elb", type=bool, default=False)
|
|
parser.add_argument("--server_num", type=int, default=1)
|
|
|
|
args, _ = parser.parse_known_args()
|
|
pid = args.pid
|
|
http_ip = args.http_ip
|
|
http_port = args.http_port
|
|
use_elb = args.use_elb
|
|
server_num = args.server_num
|
|
|
|
str_fl_id = 'fl_lenet_' + str(pid)
|
|
|
|
server_not_available_rsp = ["The cluster is in safemode.",
|
|
"The server's training job is disabled or finished."]
|
|
|
|
def generate_port():
|
|
if not use_elb:
|
|
return http_port
|
|
port = random.randint(0, 100000) % server_num + http_port
|
|
return port
|
|
|
|
|
|
def build_start_fl_job():
|
|
start_fl_job_builder = flatbuffers.Builder(1024)
|
|
|
|
fl_name = start_fl_job_builder.CreateString('fl_test_job')
|
|
fl_id = start_fl_job_builder.CreateString(str_fl_id)
|
|
data_size = 32
|
|
timestamp = start_fl_job_builder.CreateString('2020/11/16/19/18')
|
|
|
|
RequestFLJob.RequestFLJobStart(start_fl_job_builder)
|
|
RequestFLJob.RequestFLJobAddFlName(start_fl_job_builder, fl_name)
|
|
RequestFLJob.RequestFLJobAddFlId(start_fl_job_builder, fl_id)
|
|
RequestFLJob.RequestFLJobAddDataSize(start_fl_job_builder, data_size)
|
|
RequestFLJob.RequestFLJobAddTimestamp(start_fl_job_builder, timestamp)
|
|
fl_job_req = RequestFLJob.RequestFLJobEnd(start_fl_job_builder)
|
|
|
|
start_fl_job_builder.Finish(fl_job_req)
|
|
buf = start_fl_job_builder.Output()
|
|
return buf
|
|
|
|
def build_feature_map(builder, names, lengths):
|
|
if len(names) != len(lengths):
|
|
return None
|
|
feature_maps = []
|
|
np_data = []
|
|
for j, _ in enumerate(names):
|
|
name = names[j]
|
|
length = lengths[j]
|
|
weight_full_name = builder.CreateString(name)
|
|
FeatureMap.FeatureMapStartDataVector(builder, length)
|
|
weight = np.random.rand(length) * 32
|
|
np_data.append(weight)
|
|
for idx in range(length - 1, -1, -1):
|
|
builder.PrependFloat32(weight[idx])
|
|
data = builder.EndVector(length)
|
|
FeatureMap.FeatureMapStart(builder)
|
|
FeatureMap.FeatureMapAddData(builder, data)
|
|
FeatureMap.FeatureMapAddWeightFullname(builder, weight_full_name)
|
|
feature_map = FeatureMap.FeatureMapEnd(builder)
|
|
feature_maps.append(feature_map)
|
|
return feature_maps, np_data
|
|
|
|
def build_update_model(iteration):
|
|
builder_update_model = flatbuffers.Builder(1)
|
|
fl_name = builder_update_model.CreateString('fl_test_job')
|
|
fl_id = builder_update_model.CreateString(str_fl_id)
|
|
timestamp = builder_update_model.CreateString('2020/11/16/19/18')
|
|
|
|
feature_maps, np_data = build_feature_map(builder_update_model,
|
|
["conv1.weight", "conv2.weight", "fc1.weight",
|
|
"fc2.weight", "fc3.weight", "fc1.bias", "fc2.bias", "fc3.bias"],
|
|
[450, 2400, 48000, 10080, 5208, 120, 84, 62])
|
|
|
|
RequestUpdateModel.RequestUpdateModelStartFeatureMapVector(builder_update_model, 1)
|
|
for single_feature_map in feature_maps:
|
|
builder_update_model.PrependUOffsetTRelative(single_feature_map)
|
|
feature_map = builder_update_model.EndVector(len(feature_maps))
|
|
|
|
RequestUpdateModel.RequestUpdateModelStart(builder_update_model)
|
|
RequestUpdateModel.RequestUpdateModelAddFlName(builder_update_model, fl_name)
|
|
RequestUpdateModel.RequestUpdateModelAddFlId(builder_update_model, fl_id)
|
|
RequestUpdateModel.RequestUpdateModelAddIteration(builder_update_model, iteration)
|
|
RequestUpdateModel.RequestUpdateModelAddFeatureMap(builder_update_model, feature_map)
|
|
RequestUpdateModel.RequestUpdateModelAddTimestamp(builder_update_model, timestamp)
|
|
req_update_model = RequestUpdateModel.RequestUpdateModelEnd(builder_update_model)
|
|
builder_update_model.Finish(req_update_model)
|
|
buf = builder_update_model.Output()
|
|
return buf, np_data
|
|
|
|
def build_get_model(iteration):
|
|
builder_get_model = flatbuffers.Builder(1)
|
|
fl_name = builder_get_model.CreateString('fl_test_job')
|
|
timestamp = builder_get_model.CreateString('2020/12/16/19/18')
|
|
|
|
RequestGetModel.RequestGetModelStart(builder_get_model)
|
|
RequestGetModel.RequestGetModelAddFlName(builder_get_model, fl_name)
|
|
RequestGetModel.RequestGetModelAddIteration(builder_get_model, iteration)
|
|
RequestGetModel.RequestGetModelAddTimestamp(builder_get_model, timestamp)
|
|
req_get_model = RequestGetModel.RequestGetModelEnd(builder_get_model)
|
|
builder_get_model.Finish(req_get_model)
|
|
buf = builder_get_model.Output()
|
|
return buf
|
|
|
|
def datetime_to_timestamp(datetime_obj):
|
|
local_timestamp = time.mktime(datetime_obj.timetuple()) * 1000.0 + datetime_obj.microsecond // 1000.0
|
|
return local_timestamp
|
|
|
|
weight_to_idx = {
|
|
"conv1.weight": 0,
|
|
"conv2.weight": 1,
|
|
"fc1.weight": 2,
|
|
"fc2.weight": 3,
|
|
"fc3.weight": 4,
|
|
"fc1.bias": 5,
|
|
"fc2.bias": 6,
|
|
"fc3.bias": 7
|
|
}
|
|
|
|
session = requests.Session()
|
|
current_iteration = 1
|
|
np.random.seed(0)
|
|
|
|
def start_fl_job():
|
|
start_fl_job_result = {}
|
|
iteration = 0
|
|
url = "http://" + http_ip + ":" + str(generate_port()) + '/startFLJob'
|
|
print("Start fl job url is ", url)
|
|
|
|
x = session.post(url, data=build_start_fl_job())
|
|
if x.text in server_not_available_rsp:
|
|
start_fl_job_result['reason'] = "Restart iteration."
|
|
start_fl_job_result['next_ts'] = datetime_to_timestamp(datetime.datetime.now()) + 500
|
|
print("Start fl job when safemode.")
|
|
return start_fl_job_result, iteration
|
|
|
|
rsp_fl_job = ResponseFLJob.ResponseFLJob.GetRootAsResponseFLJob(x.content, 0)
|
|
iteration = rsp_fl_job.Iteration()
|
|
if rsp_fl_job.Retcode() != ResponseCode.ResponseCode.SUCCEED:
|
|
if rsp_fl_job.Retcode() == ResponseCode.ResponseCode.OutOfTime:
|
|
start_fl_job_result['reason'] = "Restart iteration."
|
|
start_fl_job_result['next_ts'] = int(rsp_fl_job.NextReqTime().decode('utf-8'))
|
|
print("Start fl job out of time. Next request at ",
|
|
start_fl_job_result['next_ts'], "reason:", rsp_fl_job.Reason())
|
|
else:
|
|
print("Start fl job failed, return code is ", rsp_fl_job.Retcode())
|
|
sys.exit()
|
|
else:
|
|
start_fl_job_result['reason'] = "Success"
|
|
start_fl_job_result['next_ts'] = 0
|
|
return start_fl_job_result, iteration
|
|
|
|
def update_model(iteration):
|
|
update_model_result = {}
|
|
|
|
url = "http://" + http_ip + ":" + str(generate_port()) + '/updateModel'
|
|
print("Update model url:", url, ", iteration:", iteration)
|
|
update_model_buf, update_model_np_data = build_update_model(iteration)
|
|
x = session.post(url, data=update_model_buf)
|
|
if x.text in server_not_available_rsp:
|
|
update_model_result['reason'] = "Restart iteration."
|
|
update_model_result['next_ts'] = datetime_to_timestamp(datetime.datetime.now()) + 500
|
|
print("Update model when safemode.")
|
|
return update_model_result, update_model_np_data
|
|
|
|
rsp_update_model = ResponseUpdateModel.ResponseUpdateModel.GetRootAsResponseUpdateModel(x.content, 0)
|
|
if rsp_update_model.Retcode() != ResponseCode.ResponseCode.SUCCEED:
|
|
if rsp_update_model.Retcode() == ResponseCode.ResponseCode.OutOfTime:
|
|
update_model_result['reason'] = "Restart iteration."
|
|
update_model_result['next_ts'] = int(rsp_update_model.NextReqTime().decode('utf-8'))
|
|
print("Update model out of time. Next request at ",
|
|
update_model_result['next_ts'], "reason:", rsp_update_model.Reason())
|
|
else:
|
|
print("Update model failed, return code is ", rsp_update_model.Retcode())
|
|
sys.exit()
|
|
else:
|
|
update_model_result['reason'] = "Success"
|
|
update_model_result['next_ts'] = 0
|
|
return update_model_result, update_model_np_data
|
|
|
|
def get_model(iteration, update_model_data):
|
|
get_model_result = {}
|
|
|
|
url = "http://" + http_ip + ":" + str(generate_port()) + '/getModel'
|
|
print("Get model url:", url, ", iteration:", iteration)
|
|
|
|
while True:
|
|
x = session.post(url, data=build_get_model(iteration))
|
|
if x.text in server_not_available_rsp:
|
|
print("Get model when safemode.")
|
|
time.sleep(0.5)
|
|
continue
|
|
|
|
rsp_get_model = ResponseGetModel.ResponseGetModel.GetRootAsResponseGetModel(x.content, 0)
|
|
ret_code = rsp_get_model.Retcode()
|
|
if ret_code == ResponseCode.ResponseCode.SUCCEED:
|
|
break
|
|
elif ret_code == ResponseCode.ResponseCode.SucNotReady:
|
|
time.sleep(0.5)
|
|
continue
|
|
else:
|
|
print("Get model failed, return code is ", rsp_get_model.Retcode())
|
|
sys.exit()
|
|
|
|
for i in range(0, 1):
|
|
print(rsp_get_model.FeatureMap(i).WeightFullname())
|
|
origin = update_model_data[weight_to_idx[rsp_get_model.FeatureMap(i).WeightFullname().decode('utf-8')]]
|
|
after = rsp_get_model.FeatureMap(i).DataAsNumpy() * 32
|
|
print("Before update model", args.pid, origin[0:10])
|
|
print("After get model", args.pid, after[0:10])
|
|
sys.stdout.flush()
|
|
|
|
get_model_result['reason'] = "Success"
|
|
get_model_result['next_ts'] = 0
|
|
return get_model_result
|
|
|
|
|
|
while True:
|
|
result, current_iteration = start_fl_job()
|
|
sys.stdout.flush()
|
|
if result['reason'] == "Restart iteration.":
|
|
current_ts = datetime_to_timestamp(datetime.datetime.now())
|
|
duration = result['next_ts'] - current_ts
|
|
if duration >= 0:
|
|
time.sleep(duration / 1000)
|
|
continue
|
|
|
|
result, update_data = update_model(current_iteration)
|
|
sys.stdout.flush()
|
|
if result['reason'] == "Restart iteration.":
|
|
current_ts = datetime_to_timestamp(datetime.datetime.now())
|
|
duration = result['next_ts'] - current_ts
|
|
if duration >= 0:
|
|
time.sleep(duration / 1000)
|
|
continue
|
|
|
|
result = get_model(current_iteration, update_data)
|
|
sys.stdout.flush()
|
|
if result['reason'] == "Restart iteration.":
|
|
current_ts = datetime_to_timestamp(datetime.datetime.now())
|
|
duration = result['next_ts'] - current_ts
|
|
if duration >= 0:
|
|
time.sleep(duration / 1000)
|
|
continue
|
|
|
|
print("")
|
|
sys.stdout.flush()
|