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()