wehub-resource-sync
248 lines
7.0 KiB
Python
248 lines
7.0 KiB
Python
"""Manage the roles in different clients.
|
|
|
|
Right now, the clients have different roles. Some clients work as samplers and
|
|
some work as trainers.
|
|
"""
|
|
|
|
import os
|
|
|
|
import numpy as np
|
|
|
|
from . import rpc
|
|
|
|
REGISTER_ROLE = 700001
|
|
REG_ROLE_MSG = "Register_Role"
|
|
|
|
|
|
class RegisterRoleResponse(rpc.Response):
|
|
"""Send a confirmation signal (just a short string message)
|
|
of RegisterRoleRequest to client.
|
|
"""
|
|
|
|
def __init__(self, msg):
|
|
self.msg = msg
|
|
|
|
def __getstate__(self):
|
|
return self.msg
|
|
|
|
def __setstate__(self, state):
|
|
self.msg = state
|
|
|
|
|
|
class RegisterRoleRequest(rpc.Request):
|
|
"""Send client id and role to server
|
|
|
|
Parameters
|
|
----------
|
|
client_id : int
|
|
ID of client
|
|
role : str
|
|
role of client
|
|
"""
|
|
|
|
def __init__(self, client_id, machine_id, role):
|
|
self.client_id = client_id
|
|
self.machine_id = machine_id
|
|
self.role = role
|
|
self.group_id = rpc.get_group_id()
|
|
|
|
def __getstate__(self):
|
|
return self.client_id, self.machine_id, self.role, self.group_id
|
|
|
|
def __setstate__(self, state):
|
|
self.client_id, self.machine_id, self.role, self.group_id = state
|
|
|
|
def process_request(self, server_state):
|
|
kv_store = server_state.kv_store
|
|
role = server_state.roles.setdefault(self.group_id, {})
|
|
if self.role not in role:
|
|
role[self.role] = set()
|
|
if kv_store is not None:
|
|
barrier_count = kv_store.barrier_count.setdefault(
|
|
self.group_id, {}
|
|
)
|
|
barrier_count[self.role] = 0
|
|
role[self.role].add((self.client_id, self.machine_id))
|
|
total_count = 0
|
|
for key in role:
|
|
total_count += len(role[key])
|
|
# Clients are blocked util all clients register their roles.
|
|
if total_count == rpc.get_num_client():
|
|
res_list = []
|
|
for target_id in range(rpc.get_num_client()):
|
|
res_list.append((target_id, RegisterRoleResponse(REG_ROLE_MSG)))
|
|
return res_list
|
|
return None
|
|
|
|
|
|
GET_ROLE = 700002
|
|
GET_ROLE_MSG = "Get_Role"
|
|
|
|
|
|
class GetRoleResponse(rpc.Response):
|
|
"""Send the roles of all client processes"""
|
|
|
|
def __init__(self, role):
|
|
self.role = role
|
|
self.msg = GET_ROLE_MSG
|
|
|
|
def __getstate__(self):
|
|
return self.role, self.msg
|
|
|
|
def __setstate__(self, state):
|
|
self.role, self.msg = state
|
|
|
|
|
|
class GetRoleRequest(rpc.Request):
|
|
"""Send a request to get the roles of all client processes."""
|
|
|
|
def __init__(self):
|
|
self.msg = GET_ROLE_MSG
|
|
self.group_id = rpc.get_group_id()
|
|
|
|
def __getstate__(self):
|
|
return self.msg, self.group_id
|
|
|
|
def __setstate__(self, state):
|
|
self.msg, self.group_id = state
|
|
|
|
def process_request(self, server_state):
|
|
return GetRoleResponse(server_state.roles[self.group_id])
|
|
|
|
|
|
# The key is role, the value is a dict of mapping RPC rank to a rank within the role.
|
|
PER_ROLE_RANK = {}
|
|
|
|
# The global rank of a client process. The client processes of the same role have
|
|
# global ranks that fall in a contiguous range.
|
|
GLOBAL_RANK = {}
|
|
|
|
# The role of the current process
|
|
CUR_ROLE = None
|
|
|
|
IS_STANDALONE = False
|
|
|
|
|
|
def init_role(role):
|
|
"""Initialize the role of the current process.
|
|
|
|
Each process is associated with a role so that we can determine what
|
|
function can be invoked in a process. For example, we do not allow some
|
|
functions in sampler processes.
|
|
|
|
The initialization includes registeration the role of the current process and
|
|
get the roles of all client processes. It also computes the rank of all client
|
|
processes in a deterministic way so that all clients will have the same rank for
|
|
the same client process.
|
|
"""
|
|
global CUR_ROLE
|
|
CUR_ROLE = role
|
|
|
|
global PER_ROLE_RANK
|
|
global GLOBAL_RANK
|
|
global IS_STANDALONE
|
|
|
|
if os.environ.get("DGL_DIST_MODE", "standalone") == "standalone":
|
|
if role == "default":
|
|
GLOBAL_RANK[0] = 0
|
|
PER_ROLE_RANK["default"] = {0: 0}
|
|
IS_STANDALONE = True
|
|
return
|
|
|
|
PER_ROLE_RANK = {}
|
|
GLOBAL_RANK = {}
|
|
|
|
# Register the current role. This blocks until all clients register themselves.
|
|
client_id = rpc.get_rank()
|
|
machine_id = rpc.get_machine_id()
|
|
request = RegisterRoleRequest(client_id, machine_id, role)
|
|
rpc.send_request(0, request)
|
|
response = rpc.recv_response()
|
|
assert response.msg == REG_ROLE_MSG
|
|
|
|
# Get all clients on all machines.
|
|
request = GetRoleRequest()
|
|
rpc.send_request(0, request)
|
|
response = rpc.recv_response()
|
|
assert response.msg == GET_ROLE_MSG
|
|
|
|
# Here we want to compute a new rank for each client.
|
|
# We compute the per-role rank as well as global rank.
|
|
# For per-role rank, we ensure that all ranks within a machine is contiguous.
|
|
# For global rank, we also ensure that all ranks within a machine are contiguous,
|
|
# and all ranks within a role are contiguous.
|
|
global_rank = 0
|
|
|
|
# We want to ensure that the global rank of the trainer process starts from 0.
|
|
role_names = ["default"]
|
|
for role_name in response.role:
|
|
if role_name not in role_names:
|
|
role_names.append(role_name)
|
|
|
|
for role_name in role_names:
|
|
# Let's collect the ranks of this role in all machines.
|
|
machines = {}
|
|
for client_id, machine_id in response.role[role_name]:
|
|
if machine_id not in machines:
|
|
machines[machine_id] = []
|
|
machines[machine_id].append(client_id)
|
|
|
|
num_machines = len(machines)
|
|
PER_ROLE_RANK[role_name] = {}
|
|
per_role_rank = 0
|
|
for i in range(num_machines):
|
|
clients = machines[i]
|
|
clients = np.sort(clients)
|
|
for client_id in clients:
|
|
GLOBAL_RANK[client_id] = global_rank
|
|
global_rank += 1
|
|
PER_ROLE_RANK[role_name][client_id] = per_role_rank
|
|
per_role_rank += 1
|
|
|
|
|
|
def get_global_rank():
|
|
"""Get the global rank
|
|
|
|
The rank can globally identify the client process. For the client processes
|
|
of the same role, their ranks are in a contiguous range.
|
|
"""
|
|
if IS_STANDALONE:
|
|
return 0
|
|
else:
|
|
return GLOBAL_RANK[rpc.get_rank()]
|
|
|
|
|
|
def get_rank(role):
|
|
"""Get the role-specific rank"""
|
|
if IS_STANDALONE:
|
|
return 0
|
|
else:
|
|
return PER_ROLE_RANK[role][rpc.get_rank()]
|
|
|
|
|
|
def get_trainer_rank():
|
|
"""Get the rank of the current trainer process.
|
|
|
|
This function can only be called in the trainer process. It will result in
|
|
an error if it's called in the process of other roles.
|
|
"""
|
|
assert CUR_ROLE == "default"
|
|
if IS_STANDALONE:
|
|
return 0
|
|
else:
|
|
return PER_ROLE_RANK["default"][rpc.get_rank()]
|
|
|
|
|
|
def get_role():
|
|
"""Get the role of the current process"""
|
|
return CUR_ROLE
|
|
|
|
|
|
def get_num_trainers():
|
|
"""Get the number of trainer processes"""
|
|
return len(PER_ROLE_RANK["default"])
|
|
|
|
|
|
rpc.register_service(REGISTER_ROLE, RegisterRoleRequest, RegisterRoleResponse)
|
|
rpc.register_service(GET_ROLE, GetRoleRequest, GetRoleResponse)
|