Task Master (#241)

* clean task queue

* remove task endpoints

* alembic

* add itnernal endpoint dummy

* offer task executions

* information source endpoint

* return payload id

* information source, project id handling

* import wizard org ids

* project id tokenizer

* attribute calculation, embedding

* embedding

* data slice

* gate removal

* remove start gates

* task executions, wizard improvements

* improve information source

* clean

* project id

* model

* task deletion

* replace alembic + model

* PR updates

* blank lines

* clean

* model merge

* test

* model dev

Author: LennartSchmidtKern

alembic/versions/881102ae15f8_rework_task_queue.py

@@ -0,0 +1,61 @@
+"""rework task queue
+
+Revision ID: 881102ae15f8
+Revises: 37d138040614
+Create Date: 2024-08-08 12:57:27.648167
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = '881102ae15f8'
+down_revision = '37d138040614'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('task_queue',
+    sa.Column('id', postgresql.UUID(as_uuid=True), nullable=False),
+    sa.Column('organization_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('task_type', sa.String(), nullable=True),
+    sa.Column('task_info', sa.JSON(), nullable=True),
+    sa.Column('priority', sa.Boolean(), nullable=True),
+    sa.Column('is_active', sa.Boolean(), nullable=True),
+    sa.Column('created_at', sa.DateTime(), nullable=True),
+    sa.Column('created_by', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.ForeignKeyConstraint(['created_by'], ['user.id'], ondelete='SET NULL'),
+    sa.ForeignKeyConstraint(['organization_id'], ['organization.id'], ondelete='CASCADE'),
+    sa.PrimaryKeyConstraint('id'),
+    schema='global'
+    )
+    op.create_index(op.f('ix_global_task_queue_created_by'), 'task_queue', ['created_by'], unique=False, schema='global')
+    op.create_index(op.f('ix_global_task_queue_organization_id'), 'task_queue', ['organization_id'], unique=False, schema='global')
+    op.drop_index('ix_task_queue_created_by', table_name='task_queue')
+    op.drop_table('task_queue')
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('task_queue',
+    sa.Column('id', postgresql.UUID(), autoincrement=False, nullable=False),
+    sa.Column('project_id', postgresql.UUID(), autoincrement=False, nullable=True),
+    sa.Column('task_type', sa.VARCHAR(), autoincrement=False, nullable=True),
+    sa.Column('task_info', postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True),
+    sa.Column('priority', sa.BOOLEAN(), autoincrement=False, nullable=True),
+    sa.Column('is_active', sa.BOOLEAN(), autoincrement=False, nullable=True),
+    sa.Column('created_at', postgresql.TIMESTAMP(), autoincrement=False, nullable=True),
+    sa.Column('created_by', postgresql.UUID(), autoincrement=False, nullable=True),
+    sa.ForeignKeyConstraint(['created_by'], ['user.id'], name='task_queue_created_by_fkey', ondelete='SET NULL'),
+    sa.ForeignKeyConstraint(['project_id'], ['project.id'], name='task_queue_project_id_fkey', ondelete='CASCADE'),
+    sa.PrimaryKeyConstraint('id', name='task_queue_pkey')
+    )
+    op.create_index('ix_task_queue_created_by', 'task_queue', ['created_by'], unique=False)
+    op.drop_index(op.f('ix_global_task_queue_organization_id'), table_name='task_queue', schema='global')
+    op.drop_index(op.f('ix_global_task_queue_created_by'), table_name='task_queue', schema='global')
+    op.drop_table('task_queue', schema='global')
+    # ### end Alembic commands ###

api/project.py

@@ -10,7 +10,7 @@
 from submodules.model import events
 from util import doc_ock, notification, adapter
 
-from controller.task_queue import manager as task_queue_manager
+from controller.task_master import manager as task_master_manager
 from submodules.model.enums import TaskType, RecordTokenizationScope
 
 logging.basicConfig(level=logging.DEBUG)
@@ -71,17 +71,22 @@ async def post(self, request_body) -> JSONResponse:
         adapter.check(data, project.id, user.id)
 
         project_manager.add_workflow_store_data_to_project(
-            user_id=user.id, project_id=project.id, file_name=name, data=data
+            user_id=user.id,
+            project_id=project.id,
+            org_id=project.organization_id,
+            file_name=name,
+            data=data,
         )
 
-        task_queue_manager.add_task(
-            str(project.id),
-            TaskType.TOKENIZATION,
+        task_master_manager.queue_task(
+            str(organization.id),
             str(user.id),
+            TaskType.TOKENIZATION,
             {
                 "scope": RecordTokenizationScope.PROJECT.value,
                 "include_rats": True,
                 "only_uploaded_attributes": False,
+                "project_id": str(project.id),
             },
         )
 

api/transfer.py

@@ -1,7 +1,7 @@
 import logging
 import traceback
 import time
-from typing import Optional, Dict
+from typing import Optional
 from starlette.endpoints import HTTPEndpoint
 from starlette.responses import PlainTextResponse, JSONResponse
 from controller.embedding.manager import recreate_embeddings
@@ -18,12 +18,11 @@
     general,
     organization,
     tokenization,
-    project as refinery_project,
+    project,
 )
 
 from submodules.model.cognition_objects import (
     project as cognition_project,
-    macro as macro_db_bo,
 )
 
 from controller.transfer import manager as transfer_manager
@@ -41,7 +40,7 @@
 from util import daemon, notification
 from controller.transfer.cognition.minio_upload import handle_cognition_file_upload
 
-from controller.task_queue import manager as task_queue_manager
+from controller.task_master import manager as task_master_manager
 from submodules.model.enums import TaskType, RecordTokenizationScope
 
 
@@ -243,86 +242,6 @@ def put(self, request) -> PlainTextResponse:
         return PlainTextResponse("OK")
 
 
-class CognitionParseMarkdownFile(HTTPEndpoint):
-    def post(self, request) -> PlainTextResponse:
-        refinery_project_id = request.path_params["project_id"]
-        refinery_project_item = refinery_project.get(refinery_project_id)
-        if not refinery_project_item:
-            return PlainTextResponse("Bad project id", status_code=400)
-
-        dataset_id = request.path_params["dataset_id"]
-        file_id = request.path_params["file_id"]
-
-        # via thread to ensure the endpoint returns immediately
-
-        daemon.run(
-            CognitionParseMarkdownFile.__add_parse_markdown_file_thread,
-            refinery_project_id,
-            str(refinery_project_item.created_by),
-            {
-                "org_id": str(refinery_project_item.organization_id),
-                "dataset_id": dataset_id,
-                "file_id": file_id,
-            },
-        )
-
-        return PlainTextResponse("OK")
-
-    def __add_parse_markdown_file_thread(
-        project_id: str, user_id: str, task_info: Dict[str, str]
-    ):
-
-        ctx_token = general.get_ctx_token()
-        try:
-            task_queue_manager.add_task(
-                project_id, TaskType.PARSE_MARKDOWN_FILE, user_id, task_info
-            )
-        finally:
-            general.remove_and_refresh_session(ctx_token, False)
-
-
-class CognitionStartMacroExecutionGroup(HTTPEndpoint):
-    def put(self, request) -> PlainTextResponse:
-        macro_id = request.path_params["macro_id"]
-        group_id = request.path_params["group_id"]
-
-        execution_entries = macro_db_bo.get_all_macro_executions(macro_id, group_id)
-
-        if len(execution_entries) == 0:
-            return PlainTextResponse("No executions found", status_code=400)
-        if not (cognition_prj_id := execution_entries[0].meta_info.get("project_id")):
-            return PlainTextResponse("No project id found", status_code=400)
-        cognition_prj = cognition_project.get(cognition_prj_id)
-        refinery_prj_id = str(
-            refinery_project.get_or_create_queue_project(
-                cognition_prj.organization_id, cognition_prj.created_by, True
-            ).id
-        )
-        cached = {str(e.id): str(e.created_by) for e in execution_entries}
-
-        def queue_tasks():
-            token = general.get_ctx_token()
-            try:
-                for exec_id in cached:
-                    task_queue_manager.add_task(
-                        refinery_prj_id,
-                        TaskType.RUN_COGNITION_MACRO,
-                        cached[exec_id],
-                        {
-                            "macro_id": macro_id,
-                            "execution_id": exec_id,
-                            "execution_group_id": group_id,
-                        },
-                    )
-                general.commit()
-            finally:
-                general.remove_and_refresh_session(token, False)
-
-        daemon.run(queue_tasks)
-
-        return PlainTextResponse("OK")
-
-
 class AssociationsImport(HTTPEndpoint):
     async def post(self, request) -> JSONResponse:
         project_id = request.path_params["project_id"]
@@ -404,11 +323,14 @@ def init_file_import(task: UploadTask, project_id: str, is_global_update: bool)
         )
     if task.file_type != "knowledge_base":
         only_usable_attributes = task.file_type == "records_add"
-        task_queue_manager.add_task(
-            project_id,
+        project_item = project.get(project_id)
+        org_id = project_item.organization_id
+        task_master_manager.queue_task(
+            str(org_id),
+            str(task.user_id),
             TaskType.TOKENIZATION,
-            task.user_id,
             {
+                "project_id": str(project_id),
                 "scope": RecordTokenizationScope.PROJECT.value,
                 "include_rats": True,
                 "only_uploaded_attributes": only_usable_attributes,

app.py

@@ -15,8 +15,6 @@
     UploadTaskInfo,
     CognitionImport,
     CognitionPrepareProject,
-    CognitionParseMarkdownFile,
-    CognitionStartMacroExecutionGroup,
 )
 from fast_api.routes.organization import router as org_router
 from fast_api.routes.project import router as project_router
@@ -36,12 +34,12 @@
 from fast_api.routes.record import router as record_router
 from fast_api.routes.weak_supervision import router as weak_supervision_router
 from fast_api.routes.labeling_tasks import router as labeling_tasks_router
+from fast_api.routes.task_execution import router as task_execution_router
 from middleware.database_session import handle_db_session
 from middleware.starlette_tmp_middleware import DatabaseSessionHandler
 from starlette.applications import Starlette
 from starlette.routing import Route, Mount
 
-from controller.task_queue.task_queue import init_task_queues
 from controller.project.manager import check_in_deletion_projects
 from route_prefix import (
     PREFIX_ORGANIZATION,
@@ -62,6 +60,7 @@
     PREFIX_RECORD,
     PREFIX_WEAK_SUPERVISION,
     PREFIX_LABELING_TASKS,
+    PREFIX_TASK_EXECUTION,
 )
 from util import security, clean_up
 from middleware import log_storage
@@ -116,6 +115,10 @@
     labeling_tasks_router, prefix=PREFIX_LABELING_TASKS, tags=["labeling-tasks"]
 )
 
+fastapi_app_internal = FastAPI()
+fastapi_app_internal.include_router(
+    task_execution_router, prefix=PREFIX_TASK_EXECUTION, tags=["task-execution"]
+)
 routes = [
     Route("/notify/{path:path}", Notify),
     Route("/healthcheck", Healthcheck),
@@ -135,19 +138,14 @@
         "/project/{cognition_project_id:str}/cognition/continue/{task_id:str}/finalize",
         CognitionPrepareProject,
     ),
-    Route(
-        "/project/{project_id:str}/cognition/datasets/{dataset_id:str}/files/{file_id:str}/queue",
-        CognitionParseMarkdownFile,
-    ),
     Route("/project/{project_id:str}/import/task/{task_id:str}", UploadTaskInfo),
     Route("/project", ProjectCreationFromWorkflow),
-    Route(
-        "/macro/{macro_id:str}/execution-group/{group_id:str}/queue",
-        CognitionStartMacroExecutionGroup,
-    ),
     Route("/is_managed", IsManagedRest),
     Route("/is_demo", IsDemoRest),
     Mount("/api", app=fastapi_app, name="REST API"),
+    Mount(
+        "/internal/api", app=fastapi_app_internal, name="INTERNAL REST API"
+    ),  # task master requests
 ]
 
 
@@ -156,7 +154,6 @@
 middleware = [Middleware(DatabaseSessionHandler)]
 app = Starlette(routes=routes, middleware=middleware)
 
-init_task_queues()
 check_in_deletion_projects()
 security.check_secret_key()
 clean_up.clean_up_database()