def can_pull(self, request, context):
mid = request.id
n_iter = request.n_iter
worker_id = request.worker_id
#print("worker {} ask can_pull >>> id = {}, iter = {}".format(worker_id, mid, n_iter))
if self.model_push_count.__contains__(mid):
return ps_service_pb2.Status(status=self.model_push_count[mid] ==
self.model_parallelism[mid] * n_iter)
else:
print('No model {} in model_push_count on PS'.format(mid))
return ps_service_pb2.Status(status=False)
def push_update(self, request, context):
mid = request.id
length = request.length
data = request.data
n_iter = request.n_iter
worker_id = request.worker_id
print("worker {} push update >>> id = {}, n_iter = {}".format(
worker_id, mid, n_iter))
update_start = time.time()
if self.models.__contains__(mid):
self.w_lock.acquire()
self.models[mid] = np.add(self.models.get(mid), data)
self.model_push_count[mid] = self.model_push_count[mid] + 1
self.w_lock.release()
print("update model cost {} s".format(time.time() - update_start))
return ps_service_pb2.Status(status=True)
else:
print('No model {} on PS'.format(mid))
return ps_service_pb2.Status(status=False)
def register_model(self, request, context):
mid = request.id
length = request.length
parallelism = request.parallelism
self.models[mid] = np.random.rand(length)
self.model_ts[mid] = time.time()
self.model_parallelism[mid] = parallelism
self.model_pull_count[mid] = 0
self.model_push_count[mid] = 0
print(
"register model >>> id = {}, length = {}, parallelism = {}".format(
mid, length, parallelism))
response = ps_service_pb2.Status(status=True)
return response
def exist_model(self, request, context):
mid = request.id
response = ps_service_pb2.Status(status=self.models.__contains__(mid))
return response
def ping(self, request, context):
response = ps_service_pb2.Status(status=True)
return response