def auth(self, permission=None, debug=False):
headers = cherrypy.request.headers
if 'X-Auth-Token' not in headers:
raise cherrypy.HTTPError(400, "Missing API Token")
token = headers['X-Auth-Token']
publisher = RPCPublisher(self.rabbitmq, "security", "validate")
output = {'token': token}
if permission is not None:
output['permission'] = permission
correlation_id = publisher.publish(output)
if correlation_id is None:
raise cherrypy.HTTPError(500, "Error publishing auth rpc")
data = publisher.get_data(correlation_id)
if data is None:
raise cherrypy.HTTPError(500, "Did not hear back from a manager - security validate")
if data['code'] != 200:
raise cherrypy.HTTPError(data['code'], data['error'])
cherrypy.request.user = {'id': data['user']['id'], 'name': data['user']['name']}
def register(self, rabbitmq, config_path):
self.rabbitmq = rabbitmq
self.config_path = config_path
self.load_config(self.config_path)
RunTaskSubscriber(self.rabbitmq, self).start()
AliveSubscriber(self.rabbitmq, self).start()
tries = 0
while self.id is None:
tries += 1
if tries > 5:
self.logger.error("Did not get a register reply from manager")
return False
publisher = RPCPublisher(self.rabbitmq, "worker", "register")
corr_id = publisher.publish({'target': self.worker_name(), 'framework': self.framework_name,
'datacenter': self.config.datacenter, 'tags': self.get_tags()})
if corr_id is None:
self.logger.warning("RPC Worker Register corr_id is None")
continue
data = publisher.get_data(corr_id)
if data is None:
self.logger.warning("RPC Worker Register data is None")
continue
self.id = data['id']
self.on_register()
return True
def launch_task(self, worker, task):
self.logger.debug("launching task " + task.name)
publisher = RPCPublisher(self.rabbitmq, "task", "launch")
correlation_id = publisher.publish({'task_id': str(task.id), 'worker_id': str(worker.id)})
if correlation_id is None:
raise LaunchTaskException("Error publishing launch task rpc")
data = publisher.get_data(correlation_id, wait=10)
if data is None:
raise LaunchTaskException("Did not receive a reply from a manager - launch task")
if data['code'] != 200:
raise LaunchTaskException(data['error'])
return data['status']
def get_workers(self, datacenter):
self.logger.debug("Getting workers")
publisher = RPCPublisher(self.rabbitmq, "worker", "get")
correlation_id = publisher.publish({'framework': self.name, 'datacenter': datacenter})
if correlation_id is None:
raise GetWorkersException("Error publishing get workers rpc")
data = publisher.get_data(correlation_id)
if data is None:
raise GetWorkersException("Did not receive a reply from a manager - get workers")
workers = []
for worker in data['workers']:
workers.append(Worker(id=worker['id'], target=worker['target'], framework=worker['framework'],
tags=worker['tags']))
return workers
def register_framework(self, rabbitmq, database, config_path):
self.config_path = config_path
self.rabbitmq = rabbitmq
self.database = database
if not self.load_config(self.config_path):
return False
tries = 0
while self.id is None:
tries += 1
if tries >= 5:
self.logger.error("Unable to register framework. Stopping")
self.stop()
return False
self.logger.info("Trying to register framework " + self.name)
publisher = RPCPublisher(self.rabbitmq, "framework", "register")
corr_id = publisher.publish({"name": self.name})
if corr_id is None:
self.logger.warning("RPC Framework Register corr_id is None")
continue
data = publisher.get_data(corr_id)
if data is None:
self.logger.warning("RPC Framework Register data is None")
continue
self.id = uuid.UUID(data['id'])
self._registered()
break
JobProcessSubscriber(self.rabbitmq, self).start()
return True
def message_deliver(self, channel, basic_deliver, properties, body):
data = json.loads(body)
self.logger.info("Received task launch "+str(data['task_id']))
with self.database.session() as session:
task = session.query(Task).filter(Task.id == data['task_id']).first()
publisher = RPCReplyPublisher(self.rabbitmq, properties.reply_to, properties.correlation_id)
if task is None:
publisher.publish({"error": "Unknown task_id "+data['task_id'], "code": 404})
publisher.close()
channel.basic_ack(basic_deliver.delivery_tag)
return
if task.status != TaskStatus.PENDING and task.status != TaskStatus.LOST:
self.logger.warning("Task cannot launch")
publisher.publish({"error": "Task "+data['task_id']+" does not have PENDING or LOST status", "code": 400})
publisher.close()
channel.basic_ack(basic_deliver.delivery_tag)
return
worker = session.query(Worker).filter(Worker.id == data['worker_id']).first()
if worker is None:
publisher.publish({"error": "Invalid worker id supplied", "code": 400})
publisher.close()
channel.basic_ack(basic_deliver.delivery_tag)
return
publisher_worker = RPCPublisher(self.rabbitmq, "worker-"+worker.target, "alive-"+worker.framework)
corr_id = publisher_worker.publish({})
if corr_id is None:
publisher.publish({"error": "Worker Alive RPC corr_id is None", "code": 400})
publisher.close()
channel.basic_ack(basic_deliver.delivery_tag)
return
if publisher_worker.get_data(corr_id) is None:
publisher.publish({"error": "Worker did not reply it is dead", "code": 400})
publisher.close()
channel.basic_ack(basic_deliver.delivery_tag)
return
task.status = TaskStatus.STARTING
session.add(task)
session.commit()
session.refresh(task)
data = {'id': task.id, 'name': task.name, 'actions': []}
for action in task.actions:
action_data = {'processor': action.processor, 'arguments': {}}
if action.arguments is not None:
action_data['arguments'] = action.arguments
data['actions'].append(action_data)
WorkerRunTask(self.rabbitmq, worker).run(data)
publisher.publish({"status": task.status.value, "code": 200})
publisher.close()
channel.basic_ack(basic_deliver.delivery_tag)
