def _init_function_servers(self):
budget_dir = os.path.join(self.working_dir, 'budget')
budget_cls = BudgetApplyServer if not self.is_multi_process \
else self.manager.budget_server
self.budget_server = budget_cls(budget_dir, self.settings, None,
self.job_name)
if self.rpc_server:
BudgetApplyServer.register_rpc(self.budget_server,
self.rpc_server,
app_name=self.job_name)
self.shutdown_callbacks.append(self.budget_server.shutdown)
counter_dir = os.path.join(self.working_dir, 'counter')
counter_cls = CounterServer if not self.is_multi_process \
else self.manager.counter_server
self.counter_server = counter_cls(counter_dir, self.settings, None,
self.job_name)
if self.rpc_server:
CounterServer.register_rpc(self.counter_server,
self.rpc_server,
app_name=self.job_name)
self.shutdown_callbacks.append(self.counter_server.shutdown)
speed_dir = os.path.join(self.working_dir, 'speed')
speed_cls = SpeedControlServer if not self.is_multi_process \
else self.manager.speed_server
self.speed_server = speed_cls(speed_dir, self.settings, None,
self.job_name, self.counter_server,
self.ctx.ips)
if self.rpc_server:
SpeedControlServer.register_rpc(self.speed_server,
self.rpc_server,
app_name=self.job_name)
self.shutdown_callbacks.append(self.speed_server.shutdown)
def _init_function_servers(self):
budget_dir = os.path.join(self.working_dir, 'budget')
budget_cls = BudgetApplyServer if not self.is_multi_process \
else self.manager.budget_server
self.budget_server = budget_cls(budget_dir, self.settings,
None, self.job_name)
if self.rpc_server:
BudgetApplyServer.register_rpc(self.budget_server, self.rpc_server,
app_name=self.job_name)
self.shutdown_callbacks.append(self.budget_server.shutdown)
counter_dir = os.path.join(self.working_dir, 'counter')
counter_cls = CounterServer if not self.is_multi_process \
else self.manager.counter_server
self.counter_server = counter_cls(counter_dir, self.settings,
None, self.job_name)
if self.rpc_server:
CounterServer.register_rpc(self.counter_server, self.rpc_server,
app_name=self.job_name)
self.shutdown_callbacks.append(self.counter_server.shutdown)
speed_dir = os.path.join(self.working_dir, 'speed')
speed_cls = SpeedControlServer if not self.is_multi_process \
else self.manager.speed_server
self.speed_server = speed_cls(speed_dir, self.settings,
None, self.job_name,
self.counter_server, self.ctx.ips)
if self.rpc_server:
SpeedControlServer.register_rpc(self.speed_server, self.rpc_server,
app_name=self.job_name)
self.shutdown_callbacks.append(self.speed_server.shutdown)
def testBudgetApply(self):
self.serv = BudgetApplyServer(self.dir_, Settings(), rpc_server=self.rpc_server)
self.cli1 = BudgetApplyClient(self.serv)
self.cli2 = BudgetApplyClient('localhost:%s'%self.port)
try:
self.serv.set_budgets(90)
self.assertEqual(self.cli1.apply(50), 50)
self.assertEqual(self.cli2.apply(50), 40)
self.cli1.finish(50)
self.assertEqual(50, self.serv.finished)
self.cli2.finish(50)
self.assertEqual(90, self.serv.finished)
self.cli1.error(10)
self.assertEqual(90, self.serv.applied)
self.serv.finished = 0
self.cli2.error(10)
self.assertEqual(80, self.serv.applied)
finally:
self.serv.shutdown()
def _init_budget_server(self):
budget_dir = os.path.join(self.working_dir, 'budget')
self.budget_server = BudgetApplyServer(budget_dir,
self.settings,
rpc_server=self.rpc_server,
app_name=self.job_name)
class JobMaster(object):
def __init__(self, ctx, job_name, job_desc, workers):
self.working_dir = os.path.join(ctx.working_dir, 'master', 'tracker',
job_name)
if not os.path.exists(self.working_dir):
os.makedirs(self.working_dir)
self.job_name = job_name
self.job_desc = job_desc
self.settings = job_desc.settings
self.rpc_server = ctx.master_rpc_server
self.inited = False
self.lock = threading.Lock()
self.init()
self.workers = workers
def _init_counter_server(self):
counter_dir = os.path.join(self.working_dir, 'counter')
self.counter_server = CounterServer(counter_dir,
self.settings,
rpc_server=self.rpc_server,
app_name=self.job_name)
def _init_budget_server(self):
budget_dir = os.path.join(self.working_dir, 'budget')
self.budget_server = BudgetApplyServer(budget_dir,
self.settings,
rpc_server=self.rpc_server,
app_name=self.job_name)
def _init_speed_server(self):
speed_dir = os.path.join(self.working_dir, 'speed')
self.speed_server = SpeedControlServer(speed_dir,
self.settings,
rpc_server=self.rpc_server,
app_name=self.job_name)
def init(self):
with self.lock:
if self.inited:
return
self._init_counter_server()
self._init_budget_server()
self._init_speed_server()
self.inited = True
def remove_worker(self, worker):
if worker not in self.workers:
return
# rpc call the other workers to remove this worker
self.workers.remove(worker)
for node in self.workers:
client_call(node, 'remove_node', worker)
def add_worker(self, worker):
if worker in self.workers:
return
# rpc call the other workers to add this worker
for node in self.workers:
client_call(node, 'add_node', worker)
self.workers.append(worker)
def has_worker(self, worker):
return worker in self.workers
def shutdown(self):
with self.lock:
if not self.inited:
return
self.counter_server.shutdown()
self.budget_server.shutdown()
self.speed_server.shutdown()
self.inited = False
def _init_budget_server(self):
budget_dir = os.path.join(self.working_dir, 'budget')
self.budget_server = BudgetApplyServer(budget_dir, self.settings,
rpc_server=self.rpc_server,
app_name=self.job_name)
class JobMaster(object):
def __init__(self, ctx, job_name, job_desc, workers):
self.working_dir = os.path.join(ctx.working_dir, 'master',
'tracker', job_name)
if not os.path.exists(self.working_dir):
os.makedirs(self.working_dir)
self.job_name = job_name
self.job_desc = job_desc
self.settings = job_desc.settings
self.rpc_server = ctx.master_rpc_server
self.inited = False
self.lock = threading.Lock()
self.init()
self.workers = workers
def _init_counter_server(self):
counter_dir = os.path.join(self.working_dir, 'counter')
self.counter_server = CounterServer(counter_dir, self.settings,
rpc_server=self.rpc_server,
app_name=self.job_name)
def _init_budget_server(self):
budget_dir = os.path.join(self.working_dir, 'budget')
self.budget_server = BudgetApplyServer(budget_dir, self.settings,
rpc_server=self.rpc_server,
app_name=self.job_name)
def _init_speed_server(self):
speed_dir = os.path.join(self.working_dir, 'speed')
self.speed_server = SpeedControlServer(speed_dir, self.settings,
rpc_server=self.rpc_server,
app_name=self.job_name)
def init(self):
with self.lock:
if self.inited:
return
self._init_counter_server()
self._init_budget_server()
self._init_speed_server()
self.inited = True
def remove_worker(self, worker):
if worker not in self.workers:
return
# rpc call the other workers to remove this worker
self.workers.remove(worker)
for node in self.workers:
client_call(node, 'remove_node', worker)
def add_worker(self, worker):
if worker in self.workers:
return
# rpc call the other workers to add this worker
for node in self.workers:
client_call(node, 'add_node', worker)
self.workers.append(worker)
def has_worker(self, worker):
return worker in self.workers
def shutdown(self):
with self.lock:
if not self.inited:
return
self.counter_server.shutdown()
self.budget_server.shutdown()
self.speed_server.shutdown()
self.inited = False
