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)