async def run(autoscale_worker, autoscale_app, autoheal_worker,
autoheal_apps):
async with connect_to_aioredis_pool(config.REDIS_URI) as redis, aiohttp.ClientSession() as session, \
connect_to_aiodocker() as docker_client:
ump = await Umpire.init(docker_client=docker_client,
redis=redis,
session=session,
autoscale_worker=autoscale_worker,
autoscale_app=autoscale_app,
autoheal_worker=autoheal_worker,
autoheal_apps=autoheal_apps)
# Attach our signal handler to cleanly close services we've created
loop = asyncio.get_running_loop()
for signame in {'SIGINT', 'SIGTERM'}:
loop.add_signal_handler(
getattr(signal, signame),
lambda: asyncio.ensure_future(ump.shutdown()))
# logger.info("Bringing up Umpire API...")
# Use --reload when you want to run locally
# os.system("uvicorn umpire.umpire_api:app --host 0.0.0.0 --port 8000 --lifespan on &")
logger.info("Umpire is initialized!")
await asyncio.gather(
asyncio.create_task(ump.workflow_control_listener()),
asyncio.create_task(ump.monitor_queues()))
await ump.shutdown()
async def console_log_generator():
async with connect_to_aioredis_pool(config.REDIS_URI) as conn:
try:
while True:
await asyncio.sleep(1)
event = await conn.rpop(redis_stream)
if event is not None:
event_object = json.loads(event.decode("ascii"))
message = event_object["message"]
for user in USERS:
await user.send_text(message)
logger.info(
f"Sending console message for {exec_id}: {message}"
)
if event_object["close"] == "Done":
await conn.delete(redis_stream)
console_stream_subs.remove(exec_id)
for user in USERS:
await user.close(code=1000)
except Exception as e:
await conn.delete(redis_stream)
console_stream_subs.remove(exec_id)
USERS.remove(websocket)
# Websocket closed so no one else can access.
await websocket.close(code=1000)
logger.info(f"Error: {e}")
async def get_build_status():
async with connect_to_aioredis_pool(config.REDIS_URI) as conn:
ret = []
build_keys = set(await conn.keys(pattern=BUILD_STATUS_GLOB + "*", encoding="utf-8"))
for key in build_keys:
build = await conn.execute('get', key)
build = build.decode('utf-8')
ret.append((key, build))
return ret
async def execute_workflow_helper(request: Request,
workflow_id,
workflow_status_col: AsyncIOMotorCollection,
workflow_col: AsyncIOMotorCollection = None,
execution_id=None,
workflow: WorkflowModel = None):
if not execution_id:
execution_id = str(uuid4())
if not workflow:
workflow = await mongo_helpers.get_item(workflow_col, WorkflowModel,
workflow_id)
# workflow_status_json = {
# "name": workflow.name,
# "status": StatusEnum.PENDING.name,
# "started_at": str(datetime.now().isoformat()),
# "workflow_id": workflow_id,
# "execution_id": execution_id,
# # "completed_at": None,
# "user": (await get_jwt_claims(request)).get('username', None),
# "node_status": [],
# # "app_name": None,
# # "action_name": None,
# # "label": None
# }
workflow_status = WorkflowStatus(
name=workflow.name,
status=StatusEnum.PENDING.name,
started_at=str(datetime.now().isoformat()),
execution_id=execution_id,
workflow_id=workflow_id,
user=(await get_jwt_claims(request)).get('username', None),
node_statuses={})
await mongo_helpers.create_item(workflow_status_col,
WorkflowStatus,
workflow_status,
id_key="execution_id")
# Assign the execution id to the workflow so the worker knows it
workflow.execution_id = execution_id
# ToDo: self.__box.encrypt(message))
async with connect_to_aioredis_pool(config.REDIS_URI) as conn:
await conn.sadd(static.REDIS_PENDING_WORKFLOWS, str(execution_id))
await conn.xadd(static.REDIS_WORKFLOW_QUEUE,
{str(execution_id): workflow.json()})
workflow_status.status = StatusEnum.PENDING
await sio.emit(static.SIO_EVENT_LOG,
json.loads(workflow_status.json()),
namespace=static.SIO_NS_WORKFLOW)
# await push_to_workflow_stream_queue(workflow_status, "PENDING")
logger.info(f"Created Workflow Status {workflow.name} ({execution_id})")
return execution_id
async def create_console_message(body: ConsoleBody, wf_exec_id: UUID = None):
logger.info(f"App console log: {body.message}")
if wf_exec_id in console_stream_subs:
redis_stream = CONSOLE_STREAM_GLOB + "." + str(wf_exec_id)
elif 'all' in console_stream_subs:
redis_stream = CONSOLE_STREAM_GLOB + ".all"
else:
# return body.message
redis_stream = CONSOLE_STREAM_GLOB + "." + str(wf_exec_id)
async with connect_to_aioredis_pool(config.REDIS_URI) as conn:
key = f"{redis_stream}"
value = body.json()
await conn.lpush(key, value)
return str(body.message)
async def update_workflow_status():
async with connect_to_aioredis_pool(config.REDIS_URI) as redis:
wfq_col = mongo.async_client.walkoff_db.workflowqueue
node_id_regex = r"/node_statuses/([0-9a-f]{8}\-[0-9a-f]{4}\-[0-9a-f]{4}\-[0-9a-f]{4}\-[0-9a-f]{12})"
while True:
try:
logger.debug("Waiting for results...")
message = (await redis.brpop(static.REDIS_RESULTS_QUEUE))[1]
message = UpdateMessage(**json.loads(message.decode()))
old_workflow_status = await mongo_helpers.get_item(
wfq_col,
WorkflowStatus,
message.execution_id,
id_key="execution_id")
patch = jsonpatch.JsonPatch.from_string(message.message)
new_workflow_status = WorkflowStatus(
**patch.apply(old_workflow_status.dict()))
update_wfs: WorkflowStatus = await mongo_helpers.update_item(
wfq_col,
WorkflowStatus,
message.execution_id,
new_workflow_status,
id_key="execution_id")
if message.type == "workflow":
update_wfs.to_response()
await sio.emit(static.SIO_EVENT_LOG,
json.loads(update_wfs.json()),
namespace=static.SIO_NS_WORKFLOW)
else:
for patch in json.loads(message.message):
node_id = re.search(node_id_regex, patch["path"],
re.IGNORECASE).group(1)
await sio.emit(
static.SIO_EVENT_LOG,
json.loads(
update_wfs.node_statuses[node_id].json()),
namespace=static.SIO_NS_NODE)
except Exception as e:
traceback.print_exc()
async def run(cls):
""" Connect to Redis and HTTP session, await actions """
async with connect_to_aioredis_pool(config.REDIS_URI) as redis:
with connect_to_socketio(config.SOCKETIO_URI, ["/console"]) as sio:
# TODO: Migrate to the common log config
logging.basicConfig(
format="{asctime} - {name} - {levelname}:{message}",
style='{')
logger = logging.getLogger(f"{cls.__name__}")
logger.setLevel(logging.DEBUG)
handler = logging.StreamHandler(stream=SIOStream(sio))
handler.setFormatter(
logging.Formatter(
fmt="{asctime} - {name} - {levelname}:{message}",
style='{'))
logger.addHandler(handler)
app = cls(redis=redis, logger=logger)
await app.get_actions()
async def control_workflow(
request: Request,
execution,
workflow_to_control: ControlWorkflow,
workflow_status_col: AsyncIOMotorCollection = Depends(get_mongo_c)):
"""
Pause, resume, or abort a workflow currently executing in WALKOFF.
"""
execution = await mongo_helpers.get_item(workflow_status_col,
WorkflowStatus,
execution,
id_key="execution_id")
walkoff_db = get_mongo_d(request)
workflow_col = walkoff_db.workflows
curr_user_id = await get_jwt_identity(request)
workflow_id = execution.workflow_id
data = dict(workflow_to_control)
status = data['status']
workflow: WorkflowModel = await mongo_helpers.get_item(
workflow_col, WorkflowModel, workflow_id)
# workflow = workflow_getter(execution.workflow_id, workflow_status_col)
# The resource factory returns the WorkflowStatus model but we want the string of the execution ID
execution_id = str(execution.execution_id)
to_execute = await auth_check(workflow,
curr_user_id,
"execute",
walkoff_db=walkoff_db)
# TODO: add in pause/resume here. Workers need to store and recover state for this
if to_execute:
if status.lower() == 'abort':
logger.info(
f"User '{(await get_jwt_claims(request)).get('username', None)}' aborting workflow: {execution_id}"
)
message = {
"execution_id": execution_id,
"status": status,
"workflow": dict(workflow)
}
async with connect_to_aioredis_pool(config.REDIS_URI) as conn:
await conn.smove(static.REDIS_PENDING_WORKFLOWS,
static.REDIS_ABORTING_WORKFLOWS, execution_id)
await conn.xadd(static.REDIS_WORKFLOW_CONTROL, message)
return None, HTTPStatus.NO_CONTENT
elif status.lower() == 'trigger':
if execution.status not in (StatusEnum.PENDING,
StatusEnum.EXECUTING,
StatusEnum.AWAITING_DATA):
raise InvalidInputException(
"workflow",
"trigger",
execution_id,
errors=[
"Workflow must be in a running state to accept triggers."
])
trigger_id = data.get('trigger_id')
if not trigger_id:
raise InvalidInputException(
"workflow",
"trigger",
execution_id,
errors=[
"ID of the trigger must be specified in trigger_id."
])
seen = False
for trigger in workflow.triggers:
if str(trigger.id_) == trigger_id:
seen = True
if not seen:
raise InvalidInputException(
"workflow",
"trigger",
execution_id,
errors=[
f"trigger_id {trigger_id} was not found in this workflow."
])
trigger_stream = f"{execution_id}-{trigger_id}:triggers"
try:
async with connect_to_aioredis_pool(config.REDIS_URI) as conn:
info = await conn.xinfo_stream(trigger_stream)
stream_length = info["length"]
except Exception:
stream_length = 0
if stream_length > 0:
return InvalidInputException(
"workflow",
"trigger",
execution_id,
errors=[f"This trigger has already received data."])
trigger_data = data.get('trigger_data')
logger.info(
f"User '{(await get_jwt_claims(request)).get('username', None)}' triggering workflow: {execution_id} at trigger "
f"{trigger_id} with data {trigger_data}")
async with connect_to_aioredis_pool(config.REDIS_URI) as conn:
await conn.xadd(trigger_stream, {
execution_id:
message_dumps({"trigger_data": trigger_data})
})
return ({"trigger_stream": trigger_stream})
else:
return None
async def build_status_from_id(build_id: str):
async with connect_to_aioredis_pool(config.REDIS_URI) as conn:
get = BUILD_STATUS_GLOB + "." + build_id
build_status = await conn.execute('get', get)
build_status = build_status.decode('utf-8')
return build_status
