def _dispatch(self, method, *args, **kwargs): """ Turns a request into a JSON-RPC call and calls it """ req_method = "zerotask.broker.client_new_task" req_params = dict(client_id=self._client_id, method=method) if args: req_params["params"] = args else: req_params["params"] = kwargs req_obj = jsonrpc.request(req_method, req_params) logging.info("Sending message %s", req_obj) self._req_socket.send_json(req_obj) result = self._req_socket.recv_json() logging.info("Recieved response %s", result) if result.get("error"): raise JSONRPCError(result["error"]["code"], result["error"].get("message")) task_id = result["result"] logging.info("New task id: %s", task_id) result_response = None while True: sub_result = self._sub_socket.recv_json() logging.info("Recieved response %s", sub_result) sub_method = sub_result.get("method") sub_params = sub_result.get("params") if sub_method != "zerotask.client.task_result_ready": logging.info("Method %s not important.", sub_method) continue if sub_params.get("task_id") != task_id: logging.info("%s != %s" % (sub_params.get("task_id"), task_id)) continue result_method = "zerotask.broker.client_task_result" result_params = dict(client_id=self._client_id, task_id=task_id) result_req = jsonrpc.request(result_method, result_params) self._req_socket.send_json(result_req) result_response = self._req_socket.recv_json() break if result_response.has_key("error"): raise JSONRPCError(result_response["error"]["code"], result_response["error"].get("message")) return result_response.get("result")
def request_status(self): """ Calls broker with the current node status """ notify_method = "zerotask.broker.node_status" notify_params = dict(node_id=self.node_id, workers=1) notification = jsonrpc.request(method=notify_method, params=notify_params, id=None) # notification logging.info("Sending status message: %s" % notification) self.broker_req_socket.send_json(notification) self.broker_req_socket.recv_json() return None
def task_ready(self, method, task_id): """ Checks if node support the method, and responds if so. """ if not self.dispatcher.has_handler(method): logging.info("Ignoring -- method %s is not supported.", method) return None # we don't support that method logging.info("Responding for new task %s", method) req_params = dict(node_id=self.node_id, task_id=task_id) req_method = "zerotask.broker.node_task_request" request = jsonrpc.request(req_method, req_params) self.broker_req_socket.send_json(request) result = self.broker_req_socket.recv_json() if result.has_key("error"): raise JSONRPCError(result["error"]["code"], result["error"].get("message")) task_result = result.get("result") if not task_result: return None # we were not elected method = task_result['method'] params = task_result.get("params", []) request = jsonrpc.request(method, params, id=task_id) logging.info("Assigning new request: %s", request) self.push_socket.send_json(request) self.running_tasks += 1
def _get_client_id(self): """ Tries to connect to broker and get client id """ method = "zerotask.broker.client_connect" params = [] request = jsonrpc.request(method, params) self._req_socket.send_json(request) response = self._req_socket.recv_json() if response.has_key("error"): raise JSONRPCError(response["error"]["code"], response["error"].get("message")) client_id = response.get("result") if not client_id: raise JSONRPCError(jsonrpc.INVALID_CLIENT_ID) logging.info("New client id: %s", client_id) return client_id
def node_task_failed(self, node_id, task_id, error): """ Saves task error state and announces to client(s) """ if not self.nodes.has_key(node_id): raise JSONRPCError(jsonrpc.INVALID_NODE_ID) if not self.tasks.has_key(task_id): raise JSONRPCError(jsonrpc.INVALID_TASK_ID) task = self.tasks[task_id] task["_status"] = zerotask.FAILED task["error"] = error announce_method = "zerotask.client.task_result_ready" announce_params = dict(task_id=task_id) pub_message = jsonrpc.request(method=announce_method, params=announce_params, id=None) # notification logging.info("Publishing message %s", pub_message) self.publish_socket.send_json(pub_message) return True
def node_task_finished(self, node_id, task_id, result): """ Run task finished process """ if not self.nodes.has_key(node_id): raise JSONRPCError(jsonrpc.INVALID_NODE_ID) if not self.tasks.has_key(task_id): raise JSONRPCError(jsonrpc.INVALID_TASK_ID) task = self.tasks[task_id] task["result"] = result task["_status"] = zerotask.FINISHED announce_method = "zerotask.client.task_result_ready" announce_params = dict(task_id=task_id) pub_message = jsonrpc.request(method=announce_method, params=announce_params, id=None) # notification logging.info("Publishing message %s", pub_message) self.publish_socket.send_json(pub_message) return True
def worker_task_result(self, result): """ Reports a task result back to broker """ self.running_tasks -= 1 if not result: logging.warning("Why do we have an empty result??") return task_id = result.get("id") result_method = "zerotask.broker.node_task_finished" result_params = dict(node_id=self.node_id, task_id=task_id) if result.has_key("error"): result_method = "zerotask.broker.node_task_failed" result_params["error"] = result["error"] else: result_params["result"] = result["result"] result_req = jsonrpc.request(method=result_method, params=result_params, id=None) # a notification self.broker_req_socket.send_json(result_req) self.broker_req_socket.recv_json()
def add_broker(self, broker_req_uri, broker_sub_uri): """ Sets the (only) broker """ self.broker_req_socket = self.context.socket(zmq.REQ) self.broker_sub_socket = self.context.socket(zmq.SUB) self.broker_req_socket.connect(broker_req_uri) self.broker_sub_socket.connect(broker_sub_uri) self.broker_sub_socket.setsockopt(zmq.SUBSCRIBE, "") # Getting node id from broker connect_method = "zerotask.broker.node_connect" connect_request = jsonrpc.request(connect_method) self.broker_req_socket.send_json(connect_request) connect_result = self.broker_req_socket.recv_json() if connect_result.has_key("error"): connect_error = connect_result["error"] raise JSONRPCError(connect_error["code"], connect_error.get("message")) node_id = connect_result.get("result") if not node_id: raise JSONRPCError(jsonrpc.INVALID_NODE_ID) self.node_id = node_id logging.info("New node id: %s", node_id) self.add_callback(self.broker_sub_socket, self.subscribe)
def client_new_task(self, client_id, method, params, task_id=None): """ Stores a CLIENTID-TASKID task structure and returns the task id to the client. It will create a task id if one is not provided. """ if not self.clients.has_key(client_id): raise JSONRPCError(jsonrpc.INVALID_CLIENT_ID) task_id = task_id or jsonrpc.get_random_id() full_id = "%s-%s" % (client_id, task_id) if self.tasks.has_key(full_id): raise JSONRPCError(jsonrpc.INVALID_TASK_ID) task = self.tasks.get(full_id) announce_params = dict(method=method, task_id=full_id) announce_method = "zerotask.node.task_ready" pub_message = jsonrpc.request(method=announce_method, params=announce_params, id=None) # notification self.publish_socket.send_json(pub_message) task = dict(method=method, params=params, _status=zerotask.QUEUED, id=task_id) self.tasks[full_id] = task return full_id