Pause and cancel for task master (#255)

* cancel markdown and macro

* return + delete by task id

* cancel macro

* cancel tmp doc

* remove print

* model

* attribute cancel

* attr cancel

* simple payload cancel check

* improve cancel

* model + camel case

* task queue pause info

* model

* directlyy set to failed etl

* model merge

* model update

---------

Co-authored-by: LennartSchmidtKern <[email protected]>

Author: lumburovskalina

controller/attribute/manager.py

@@ -309,7 +309,11 @@ def __calculate_user_attribute_all_records(
     util.add_log_to_attribute_logs(project_id, attribute_id, "Finished writing.")
 
     attribute_item = attribute.get(project_id, attribute_id)
-    if attribute_item.data_type == DataTypes.TEXT.value:
+    if (
+        attribute_item
+        and attribute_item.data_type == DataTypes.TEXT.value
+        and not attribute_item.state == AttributeState.FAILED.value
+    ):
         util.add_log_to_attribute_logs(
             project_id, attribute_id, "Triggering tokenization."
         )
@@ -346,6 +350,15 @@ def __calculate_user_attribute_all_records(
         )
         request_reupload_docbins(project_id)
 
+    attribute_item = attribute.get(project_id, attribute_id)
+    if attribute_item.state == AttributeState.FAILED.value:
+        __notify_attribute_calculation_failed(
+            project_id=project_id,
+            attribute_id=attribute_id,
+            log="Writing to the database failed.",
+        )
+        general.remove_and_refresh_session(session_token)
+        return
     util.set_progress(project_id, attribute_item, 1.0)
     attribute.update(
         project_id=project_id,

controller/attribute/util.py

@@ -20,6 +20,7 @@
 from submodules.s3 import controller as s3
 from util import daemon, notification
 from controller.knowledge_base import util as knowledge_base
+from submodules.model import enums
 
 client = docker.from_env()
 image = os.getenv("AC_EXEC_ENV_IMAGE")
@@ -189,7 +190,11 @@ def read_container_logs_thread(
         c += 1
         if c > 100:
             ctx_token = general.remove_and_refresh_session(ctx_token, True)
-            attribute_item = attribute.get(project_id, attribute_id)
+        attribute_item = attribute.get(project_id, attribute_id)
+        if not attribute_item:
+            break
+        if attribute_item.state == enums.AttributeState.FAILED.value:
+            break
         if not name in __containers_running:
             break
         try:

controller/monitor/manager.py

@@ -1,7 +1,8 @@
-from typing import Any, List
+from typing import Any, List, Dict
 from submodules.model.business_objects import monitor as task_monitor
 from controller.auth import kratos
 from submodules.model.util import sql_alchemy_to_dict
+from submodules.s3 import controller as s3
 
 
 def monitor_all_tasks(page: int, limit: int) -> List[Any]:
@@ -33,36 +34,85 @@ def cancel_upload_task(project_id: str = None, upload_task_id: str = None) -> No
     task_monitor.set_upload_task_to_failed(project_id, upload_task_id, with_commit=True)
 
 
-def cancel_weak_supervision(project_id: str = None, payload_id: str = None) -> None:
-    task_monitor.set_weak_supervision_to_failed(
-        project_id, payload_id, with_commit=True
-    )
+def cancel_weak_supervision(
+    task_info: Dict[str, Any],
+) -> None:
+    project_id = task_info.get("projectId")
+    payload_id = task_info.get("payloadId")
+    if project_id and payload_id:
+        task_monitor.set_weak_supervision_to_failed(
+            project_id, payload_id, with_commit=True
+        )
 
 
 def cancel_attribute_calculation(
-    project_id: str = None, attribute_id: str = None
+    task_info: Dict[str, Any],
 ) -> None:
-    task_monitor.set_attribute_calculation_to_failed(
-        project_id, attribute_id, with_commit=True
-    )
 
+    project_id = task_info.get("projectId")
+    attribute_id = task_info.get("attributeId")
+    if project_id and attribute_id:
+        task_monitor.set_attribute_calculation_to_failed(
+            project_id, attribute_id, with_commit=True
+        )
 
-def cancel_embedding(project_id: str = None, embedding_id: str = None) -> None:
-    task_monitor.set_embedding_to_failed(project_id, embedding_id, with_commit=True)
+
+def cancel_embedding(
+    task_info: Dict[str, Any],
+) -> None:
+    project_id = task_info.get("projectId")
+    embedding_id = task_info.get("embeddingId")
+    if project_id and embedding_id:
+        task_monitor.set_embedding_to_failed(project_id, embedding_id, with_commit=True)
 
 
 def cancel_information_source_payload(
-    project_id: str = None, payload_id: str = None
+    task_info: Dict[str, Any],
 ) -> None:
-    task_monitor.set_information_source_payloads_to_failed(
-        project_id, payload_id, with_commit=True
-    )
+    project_id = task_info.get("projectId")
+    payload_id = task_info.get("payloadId")
+    if project_id and payload_id:
+        task_monitor.set_information_source_payloads_to_failed(
+            project_id, payload_id, with_commit=True
+        )
 
 
 def cancel_record_tokenization_task(
-    project_id: str = None,
-    tokenization_task_id: str = None,
+    task_info: Dict[str, Any],
+) -> None:
+    project_id = task_info.get("projectId")
+    tokenization_task_id = task_info.get("recordTokenizationTaskId")
+    if project_id and tokenization_task_id:
+        task_monitor.set_record_tokenization_task_to_failed(
+            project_id, tokenization_task_id, with_commit=True
+        )
+
+
+def cancel_macro_execution_task(
+    task_info: Dict[str, Any],
 ) -> None:
-    task_monitor.set_record_tokenization_task_to_failed(
-        project_id, tokenization_task_id, with_commit=True
+
+    macro_execution_id = task_info.get("executionId")
+    macro_execution_group_id = task_info.get("groupExecutionId")
+
+    task_monitor.set_macro_execution_task_to_failed(
+        macro_execution_id, macro_execution_group_id, with_commit=True
     )
+
+
+def cancel_markdown_file_task(
+    task_info: Dict[str, Any],
+) -> None:
+    markdown_file_id = task_info.get("fileId")
+    org_id = task_info.get("orgId")
+    task_monitor.set_markdown_file_task_to_failed(
+        markdown_file_id, org_id, with_commit=True
+    )
+
+
+def cancel_tmp_doc_retrieval_task(
+    task_info: Dict[str, Any],
+) -> None:
+    bucket = task_info.get("bucket")
+    minio_path = task_info.get("minioPath")
+    s3.delete_object(bucket, minio_path)

controller/payload/payload_scheduler.py

@@ -533,6 +533,21 @@ def update_records(
 ) -> bool:
     org_id = organization.get_id_by_project_id(project_id)
     tmp_log_store = information_source_payload.logs
+
+    if information_source_payload.state == enums.PayloadState.FAILED.value:
+        berlin_now = datetime.datetime.now(__tz)
+        tmp_log_store.append(
+            " ".join(
+                [
+                    berlin_now.strftime("%Y-%m-%dT%H:%M:%S"),
+                    "Information source task cancelled.",
+                ]
+            )
+        )
+        information_source_payload.logs = tmp_log_store
+        flag_modified(information_source_payload, "logs")
+        general.commit()
+        return True
     try:
         output_data = json.loads(
             s3.get_object(

controller/task_master/manager.py

@@ -35,3 +35,7 @@ def pause_task_queue(task_queue_pause: bool) -> requests.Response:
     return requests.post(
         f"{TASK_MASTER_URL}/task/queue/pause?task_queue_pause={task_queue_pause}"
     )
+
+
+def get_task_queue_pause() -> requests.Response:
+    return requests.get(f"{TASK_MASTER_URL}/task/queue/pause")

fast_api/models.py

@@ -433,8 +433,8 @@ class DeleteUserBody(BaseModel):
 
 
 class CancelTaskBody(BaseModel):
-    project_id: StrictStr
     task_id: StrictStr
+    task_info: Dict[StrictStr, Any]
     task_type: StrictStr
 
 

fast_api/routes/misc.py

@@ -26,6 +26,7 @@
     CustomerButtonType,
     CustomerButtonLocation,
 )
+from submodules.model.business_objects import task_queue as task_queue_bo
 
 router = APIRouter()
 
@@ -114,24 +115,37 @@ def cancel_task(
 
     auth.check_admin_access(request.state.info)
     task_type = body.task_type
-    project_id = body.project_id
+    task_info = body.task_info
     task_id = body.task_id
 
-    if task_type == enums.TaskType.ATTRIBUTE_CALCULATION.value:
-        controller_manager.cancel_attribute_calculation(project_id, task_id)
-    elif task_type == enums.TaskType.EMBEDDING.value:
-        controller_manager.cancel_embedding(project_id, task_id)
-    elif task_type == enums.TaskType.INFORMATION_SOURCE.value:
-        controller_manager.cancel_information_source_payload(project_id, task_id)
-    elif task_type == enums.TaskType.TOKENIZATION.value:
-        controller_manager.cancel_record_tokenization_task(project_id, task_id)
-    elif task_type == enums.TaskType.UPLOAD_TASK.value:
-        controller_manager.cancel_upload_task(project_id, task_id)
-    elif task_type == enums.TaskType.WEAK_SUPERVISION.value:
-        controller_manager.cancel_weak_supervision(project_id, task_id)
-    else:
-        raise ValueError(f"{task_type} is no valid task type")
-
+    task_entity = task_queue_bo.get(task_id)
+    if not task_entity:
+        return pack_json_result({"data": {"cancelTask": {"ok": False}}})
+    if task_entity and (
+        task_entity.is_active or task_type == enums.TaskType.PARSE_MARKDOWN_FILE.value
+    ):
+        if task_type == enums.TaskType.ATTRIBUTE_CALCULATION.value:
+            controller_manager.cancel_attribute_calculation(task_info)
+        elif task_type == enums.TaskType.EMBEDDING.value:
+            controller_manager.cancel_embedding(task_info)
+        elif task_type == enums.TaskType.INFORMATION_SOURCE.value:
+            controller_manager.cancel_information_source_payload(task_info)
+        elif task_type == enums.TaskType.TOKENIZATION.value:
+            controller_manager.cancel_record_tokenization_task(task_info)
+        elif task_type == enums.TaskType.UPLOAD_TASK.value:
+            controller_manager.cancel_upload_task(task_info)
+        elif task_type == enums.TaskType.WEAK_SUPERVISION.value:
+            controller_manager.cancel_weak_supervision(task_info)
+        elif task_type == enums.TaskType.RUN_COGNITION_MACRO.value:
+            controller_manager.cancel_macro_execution_task(task_info)
+        elif task_type == enums.TaskType.PARSE_MARKDOWN_FILE.value:
+            controller_manager.cancel_markdown_file_task(task_info)
+        elif task_type == enums.TaskType.PARSE_COGNITION_TMP_FILE.value:
+            controller_manager.cancel_tmp_doc_retrieval_task(task_info)
+        else:
+            raise ValueError(f"{task_type} is no valid task type")
+
+    task_queue_bo.delete_by_task_id(task_id, True)
     return pack_json_result({"data": {"cancelTask": {"ok": True}}})
 
 
@@ -142,11 +156,30 @@ def cancel_all_running_tasks(request: Request):
     return pack_json_result({"data": {"cancelAllRunningTasks": {"ok": True}}})
 
 
-@router.get("/pause-task-queue")
+@router.post("/pause-task-queue")
 def pause_task_queue(request: Request, task_queue_pause: bool):
     auth.check_admin_access(request.state.info)
-    task_master_manager.pause_task_queue(task_queue_pause)
-    return SILENT_SUCCESS_RESPONSE
+    task_queue_pause_response = task_master_manager.pause_task_queue(task_queue_pause)
+    task_queue_pause = False
+    if task_queue_pause_response.ok:
+        try:
+            task_queue_pause = task_queue_pause_response.json()["task_queue_pause"]
+        except Exception:
+            task_queue_pause = False
+    return pack_json_result({"taskQueuePause": task_queue_pause})
+
+
+@router.get("/pause-task-queue")
+def get_task_queue_pause(request: Request):
+    auth.check_admin_access(request.state.info)
+    task_queue_pause_response = task_master_manager.get_task_queue_pause()
+    task_queue_pause = False
+    if task_queue_pause_response.ok:
+        try:
+            task_queue_pause = task_queue_pause_response.json()["task_queue_pause"]
+        except Exception:
+            task_queue_pause = False
+    return pack_json_result({"taskQueuePause": task_queue_pause})
 
 
 @router.get("/all-users-activity")