def test_receive_workflow_with_env_vars(self): workflow_id = str(uuid4()) execution_id = str(uuid4()) start = str(uuid4()) env_var_id = str(uuid4()) env_var_id2 = str(uuid4()) env_vars = [{"id": env_var_id, "value": "env_var_1"}, {"id": env_var_id2, "value": "env_var_2"}] message = ExecuteWorkflowMessage() message.workflow_id = workflow_id message.workflow_execution_id = execution_id message.resume = True message.start = start env_var = message.environment_variables.add() env_var.id = env_vars[0]['id'] env_var.value = env_vars[0]['value'] env_var = message.environment_variables.add() env_var.id = env_vars[1]['id'] env_var.value = env_vars[1]['value'] receiver = self.get_receiver() encrypted_message = self.box.encrypt(message.SerializeToString()) workflow_generator = receiver.receive_workflows() receiver.cache.lpush('request_queue', encrypted_message) workflow = next(workflow_generator) workflow_env_vars = workflow[5] self.assertEqual(str(workflow_env_vars[0].id), env_vars[0]['id']) self.assertEqual(workflow_env_vars[0].value, env_vars[0]['value']) self.assertEqual(str(workflow_env_vars[1].id), env_vars[1]['id']) self.assertEqual(workflow_env_vars[1].value, env_vars[1]['value'])
def receive_workflows(self): """Receives requests to execute workflows, and sends them off to worker threads""" logger.info('Starting workflow receiver') box = Box(self.key, self.server_key) while not self.exit: received_message = self.cache.rpop("request_queue") if received_message is not None: try: decrypted_msg = box.decrypt(received_message) except CryptoError: logger.error( 'Worker could not decrypt received workflow message') continue try: message = ExecuteWorkflowMessage() message.ParseFromString(decrypted_msg) except DecodeError: logger.error( 'Workflow could not decode received workflow message') else: start = message.start if hasattr(message, 'start') else None start_arguments = [] if hasattr(message, 'arguments'): for arg in message.arguments: start_arguments.append( Argument(**(MessageToDict( arg, preserving_proto_field_name=True)))) yield message.workflow_id, message.workflow_execution_id, start, start_arguments, message.resume else: yield None raise StopIteration
def test_receive_workflow_basic_workflow(self): workflow_id = str(uuid4()) execution_id = str(uuid4()) message = ExecuteWorkflowMessage() message.workflow_id = workflow_id message.workflow_execution_id = execution_id message.resume = True self.check_workflow_message(message, (workflow_id, execution_id, '', [], True, [], ''))
def create_workflow_request_message(workflow_id, workflow_execution_id, start=None, start_arguments=None, resume=False, environment_variables=None, user=None): """Creates a workflow request message to be placed on the redis queue """ message = ExecuteWorkflowMessage() message.workflow_id = str(workflow_id) message.workflow_execution_id = workflow_execution_id message.resume = resume if start: message.start = str(start) if start_arguments: ProtobufWorkflowResultsConverter._add_arguments_to_proto( message, start_arguments) if environment_variables: ProtobufWorkflowResultsConverter.add_env_vars_to_proto( message, environment_variables) if user: message.user = user return message.SerializeToString()
def test_receive_workflow_with_env_vars(self): workflow_id = str(uuid4()) execution_id = str(uuid4()) start = str(uuid4()) env_var_id = str(uuid4()) env_var_id2 = str(uuid4()) env_vars = [{ "id": env_var_id, "value": "env_var_1" }, { "id": env_var_id2, "value": "env_var_2" }] message = ExecuteWorkflowMessage() message.workflow_id = workflow_id message.workflow_execution_id = execution_id message.resume = True message.start = start env_var = message.environment_variables.add() env_var.id = env_vars[0]['id'] env_var.value = env_vars[0]['value'] env_var = message.environment_variables.add() env_var.id = env_vars[1]['id'] env_var.value = env_vars[1]['value'] receiver = self.get_receiver() encrypted_message = self.box.encrypt(message.SerializeToString()) workflow_generator = receiver.receive_workflows() receiver.cache.lpush('request_queue', encrypted_message) workflow = next(workflow_generator) workflow_env_vars = workflow[5] self.assertEqual(str(workflow_env_vars[0].id), env_vars[0]['id']) self.assertEqual(workflow_env_vars[0].value, env_vars[0]['value']) self.assertEqual(str(workflow_env_vars[1].id), env_vars[1]['id']) self.assertEqual(workflow_env_vars[1].value, env_vars[1]['value'])
def test_receive_workflow_with_arguments(self): workflow_id = str(uuid4()) execution_id = str(uuid4()) start = str(uuid4()) ref = str(uuid4()) arguments = [{'name': 'arg1', 'value': 42}, {'name': 'arg2', 'reference': ref, 'selection': ['a', 1]}] message = ExecuteWorkflowMessage() message.workflow_id = workflow_id message.workflow_execution_id = execution_id message.resume = True message.start = start arg = message.arguments.add() arg.name = arguments[0]['name'] arg.value = str(arguments[0]['value']) arg = message.arguments.add() arg.name = arguments[1]['name'] arg.reference = arguments[1]['reference'] arg.selection = str(arguments[1]['selection']) receiver = self.get_receiver() encrypted_message = self.box.encrypt(message.SerializeToString()) workflow_generator = receiver.receive_workflows() receiver.cache.lpush('request_queue', encrypted_message) workflow = next(workflow_generator) workflow_arguments = workflow[3] self.assertEqual(workflow_arguments[0].name, arguments[0]['name']) self.assertEqual(workflow_arguments[0].value, str(arguments[0]['value'])) self.assertEqual(workflow_arguments[1].name, arguments[1]['name']) self.assertEqual(workflow_arguments[1].reference, ref) self.assertEqual(workflow_arguments[1].selection, str(arguments[1]['selection']))
def add_workflow(self, workflow_id, workflow_execution_id, start=None, start_arguments=None, resume=False, environment_variables=None): """Adds a workflow ID to the queue to be executed. Args: workflow_id (UUID): The ID of the workflow to be executed. workflow_execution_id (UUID): The execution ID of the workflow to be executed. start (UUID, optional): The ID of the first, or starting action. Defaults to None. start_arguments (list[Argument], optional): The arguments to the starting action of the workflow. Defaults to None. resume (bool, optional): Optional boolean to resume a previously paused workflow. Defaults to False. environment_variables (list[EnvironmentVariable]): Optional list of environment variables to pass into the workflow. These will not be persistent. """ message = ExecuteWorkflowMessage() message.workflow_id = str(workflow_id) message.workflow_execution_id = workflow_execution_id message.resume = resume if start: message.start = str(start) if start_arguments: self._set_arguments_for_proto(message, start_arguments) if environment_variables: proto_helpers.add_env_vars_to_proto(message, environment_variables) message = message.SerializeToString() encrypted_message = self.box.encrypt(message) self.cache.lpush("request_queue", encrypted_message)
def receive_requests(self): """Receives requests to execute workflows, and sends them off to worker threads""" self.request_sock.send(b"Ready") while True: message_bytes = self.request_sock.recv() message = ExecuteWorkflowMessage() message.ParseFromString(message_bytes) start = message.start if hasattr(message, 'start') else None start_arguments = [] if hasattr(message, 'arguments'): for arg in message.arguments: start_arguments.append( Argument(**(MessageToDict( arg, preserving_proto_field_name=True)))) self.threadpool.submit(self.execute_workflow_worker, message.workflow_id, message.workflow_execution_id, start, start_arguments, message.resume)
def manage_workflows(self): """Manages the workflows to be executed and the workers. It waits for the server to submit a request to execute a workflow, and then passes the workflow off to an available worker, once one becomes available. """ while True: if self.thread_exit: break if len(self.workers) < walkoff.config.config.num_processes: try: worker, message = self.request_socket.recv_multipart( flags=zmq.NOBLOCK) if message == b"Ready": self.workers[ worker] = walkoff.config.config.num_threads_per_process except zmq.ZMQError: pass # There is a worker available and a workflow in the queue, so pop it off and send it to the worker if any(val > 0 for val in self.workers.values() ) and not self.pending_workflows.empty(): workflow_id, workflow_execution_id, start, start_arguments, resume = self.pending_workflows.get( ) executiondb.execution_db.session.expire_all() workflow_status = executiondb.execution_db.session.query( WorkflowStatus).filter_by( execution_id=workflow_execution_id).first() if workflow_status.status == WorkflowStatusEnum.aborted: continue worker = self.__get_available_worker() self.workflow_comms[workflow_execution_id] = worker message = ExecuteWorkflowMessage() message.workflow_id = str(workflow_id) message.workflow_execution_id = workflow_execution_id message.resume = resume if start: message.start = str(start) if start_arguments: self.__set_arguments_for_proto(message, start_arguments) self.request_socket.send_multipart( [worker, message.SerializeToString()]) gevent.sleep(0.1) self.request_socket.close() self.comm_socket.close() return
def test_receive_workflow_with_start(self): workflow_id = str(uuid4()) execution_id = str(uuid4()) start = str(uuid4()) message = ExecuteWorkflowMessage() message.workflow_id = workflow_id message.workflow_execution_id = execution_id message.resume = True message.start = start self.check_workflow_message(message, (workflow_id, execution_id, start, [], True, [], ''))
def create_workflow_request_message(workflow_id, workflow_execution_id, start=None, start_arguments=None, resume=False, environment_variables=None, user=None): """Creates a workflow request message to be placed on the redis queue """ message = ExecuteWorkflowMessage() message.workflow_id = str(workflow_id) message.workflow_execution_id = workflow_execution_id message.resume = resume if start: message.start = str(start) if start_arguments: ProtobufWorkflowResultsConverter._add_arguments_to_proto(message, start_arguments) if environment_variables: ProtobufWorkflowResultsConverter.add_env_vars_to_proto(message, environment_variables) if user: message.user = user return message.SerializeToString()