def enqueue(call: APICall, company_id, req_model: EnqueueRequest):
task_id = req_model.task
queue_id = req_model.queue
status_message = req_model.status_message
status_reason = req_model.status_reason
if not queue_id:
# try to get default queue
queue_id = queue_bll.get_default(company_id).id
with translate_errors_context():
query = dict(id=task_id, company=company_id)
task = Task.get_for_writing(
_only=("type", "script", "execution", "status", "project", "id"), **query
)
if not task:
raise errors.bad_request.InvalidTaskId(**query)
res = EnqueueResponse(
**ChangeStatusRequest(
task=task,
new_status=TaskStatus.queued,
status_reason=status_reason,
status_message=status_message,
allow_same_state_transition=False,
).execute()
)
try:
queue_bll.add_task(
company_id=company_id, queue_id=queue_id, task_id=task.id
)
except Exception:
# failed enqueueing, revert to previous state
ChangeStatusRequest(
task=task,
current_status_override=TaskStatus.queued,
new_status=task.status,
force=True,
status_reason="failed enqueueing",
).execute()
raise
# set the current queue ID in the task
if task.execution:
Task.objects(**query).update(execution__queue=queue_id, multi=False)
else:
Task.objects(**query).update(
execution=Execution(queue=queue_id), multi=False
)
res.queued = 1
res.fields.update(**{"execution.queue": queue_id})
call.result.data_model = res
def cleanup_tasks(cls, threshold_sec):
relevant_status = (TaskStatus.in_progress,)
threshold = timedelta(seconds=threshold_sec)
ref_time = datetime.utcnow() - threshold
log.info(
f"Starting cleanup cycle for running tasks last updated before {ref_time}"
)
tasks = list(
Task.objects(status__in=relevant_status, last_update__lt=ref_time).only(
"id", "name", "status", "project", "last_update"
)
)
log.info(f"{len(tasks)} non-responsive tasks found")
if not tasks:
return 0
err_count = 0
for task in tasks:
log.info(
f"Stopping {task.id} ({task.name}), last updated at {task.last_update}"
)
try:
ChangeStatusRequest(
task=task,
new_status=TaskStatus.stopped,
status_reason="Forced stop (non-responsive)",
status_message="Forced stop (non-responsive)",
force=True,
).execute()
except errors.bad_request.FailedChangingTaskStatus:
err_count += 1
return len(tasks) - err_count
def dequeue(call: APICall, company_id, req_model: UpdateRequest):
task = TaskBLL.get_task_with_access(
req_model.task,
company_id=company_id,
only=("id", "execution", "status", "project"),
requires_write_access=True,
)
if task.status not in (TaskStatus.queued,):
raise errors.bad_request.InvalidTaskId(
status=task.status, expected=TaskStatus.queued
)
_dequeue(task, company_id)
status_message = req_model.status_message
status_reason = req_model.status_reason
res = DequeueResponse(
**ChangeStatusRequest(
task=task,
new_status=TaskStatus.created,
status_reason=status_reason,
status_message=status_message,
).execute(unset__execution__queue=1)
)
res.dequeued = 1
call.result.data_model = res
def reset(call: APICall, company_id, req_model: UpdateRequest):
task = TaskBLL.get_task_with_access(req_model.task,
company_id=company_id,
requires_write_access=True)
force = req_model.force
if not force and task.status == TaskStatus.published:
raise errors.bad_request.InvalidTaskStatus(task_id=task.id,
status=task.status)
api_results = {}
updates = {}
try:
dequeued = _dequeue(task, company_id, silent_fail=True)
except APIError:
# dequeue may fail if the task was not enqueued
pass
else:
if dequeued:
api_results.update(dequeued=dequeued)
updates.update(unset__execution__queue=1)
cleaned_up = cleanup_task(task, force)
api_results.update(attr.asdict(cleaned_up))
updates.update(
set__last_iteration=DEFAULT_LAST_ITERATION,
set__last_metrics={},
unset__output__result=1,
unset__output__model=1,
__raw__={"$pull": {
"execution.artifacts": {
"mode": {
"$ne": "input"
}
}
}},
)
res = ResetResponse(**ChangeStatusRequest(
task=task,
new_status=TaskStatus.created,
force=force,
status_reason="reset",
status_message="reset",
).execute(started=None, completed=None, published=None, **updates))
for key, value in api_results.items():
setattr(res, key, value)
call.result.data_model = res
def set_task_status_from_call(request: UpdateRequest,
company_id,
new_status=None,
**kwargs) -> dict:
task = TaskBLL.get_task_with_access(
request.task,
company_id=company_id,
only=("status", "project"),
requires_write_access=True,
)
status_reason = request.status_reason
status_message = request.status_message
force = request.force
return ChangeStatusRequest(
task=task,
new_status=new_status or task.status,
status_reason=status_reason,
status_message=status_message,
force=force,
).execute(**kwargs)
def set_task_status_from_call(
request: UpdateRequest, company_id, new_status=None, **set_fields
) -> dict:
fields_resolver = SetFieldsResolver(set_fields)
task = TaskBLL.get_task_with_access(
request.task,
company_id=company_id,
only=tuple({"status", "project"} | fields_resolver.get_names()),
requires_write_access=True,
)
status_reason = request.status_reason
status_message = request.status_message
force = request.force
return ChangeStatusRequest(
task=task,
new_status=new_status or task.status,
status_reason=status_reason,
status_message=status_message,
force=force,
).execute(**fields_resolver.get_fields(task))
def reset(call: APICall, company_id, req_model: UpdateRequest):
task = TaskBLL.get_task_with_access(req_model.task,
company_id=company_id,
requires_write_access=True)
force = req_model.force
if not force and task.status == TaskStatus.published:
raise errors.bad_request.InvalidTaskStatus(task_id=task.id,
status=task.status)
api_results = {}
updates = {}
cleaned_up = cleanup_task(task, force)
api_results.update(attr.asdict(cleaned_up))
updates.update(
unset__script__requirements=1,
set__last_iteration=DEFAULT_LAST_ITERATION,
set__last_metrics={},
unset__output__result=1,
unset__output__model=1,
)
res = ResetResponse(**ChangeStatusRequest(
task=task,
new_status=TaskStatus.created,
force=force,
status_reason="reset",
status_message="reset",
).execute(started=None, completed=None, published=None, **updates))
for key, value in api_results.items():
setattr(res, key, value)
call.result.data_model = res
def set_task_status_from_call(request: UpdateRequest,
company_id,
new_status=None,
**set_fields) -> dict:
fields_resolver = SetFieldsResolver(set_fields)
task = TaskBLL.get_task_with_access(
request.task,
company_id=company_id,
only=tuple({"status", "project", "started", "duration"}
| fields_resolver.get_names()),
requires_write_access=True,
)
if "duration" not in fields_resolver.get_names():
if new_status == Task.started:
fields_resolver.add_fields(
min__duration=max(0, task.duration or 0))
elif new_status in (
TaskStatus.completed,
TaskStatus.failed,
TaskStatus.stopped,
):
fields_resolver.add_fields(
duration=int((task.started -
datetime.utcnow()).total_seconds()))
status_reason = request.status_reason
status_message = request.status_message
force = request.force
return ChangeStatusRequest(
task=task,
new_status=new_status or task.status,
status_reason=status_reason,
status_message=status_message,
force=force,
).execute(**fields_resolver.get_fields(task))