Merge pull request #3565 from OpenNeuroOrg/feat/taskiq-dashboard

Log task execution for workers to MongoDB

Author: nellh

packages/openneuro-server/src/graphql/permissions.ts

@@ -169,3 +169,12 @@ export const checkAdmin = (userId, userInfo) =>
   userId && userInfo.admin
     ? Promise.resolve(true)
     : Promise.reject(states.ADMIN.errorMessage)
+
+/**
+ * Check if the user is a worker
+ * @param userInfo User context
+ */
+export const checkWorker = (userInfo) => {
+  if (userInfo?.worker) return true
+  else throw new Error("You must be a worker to make this request.")
+}

packages/openneuro-server/src/graphql/resolvers/mutation.ts

@@ -46,6 +46,7 @@ import {
 import { saveAdminNote } from "./datasetEvents"
 import { createGitEvent } from "./gitEvents"
 import { updateFileCheck } from "./fileCheck"
+import { updateWorkerTask } from "./worker"
 
 const Mutation = {
   createDataset,
@@ -95,6 +96,7 @@ const Mutation = {
   saveAdminNote,
   createGitEvent,
   updateFileCheck,
+  updateWorkerTask,
 }
 
 export default Mutation

packages/openneuro-server/src/graphql/resolvers/worker.ts

@@ -0,0 +1,23 @@
+import WorkerTask from "../../models/worker-task"
+import { checkWorker } from "../permissions"
+
+/**
+ * Update a worker task record
+ *
+ * This can be called for new tasks, or to update existing tasks.
+ */
+export const updateWorkerTask = async (obj, args, { userInfo }) => {
+  checkWorker(userInfo)
+  const { id, ...updateData } = args
+
+  // Don't allow null values to unset fields
+  const update = Object.fromEntries(
+    Object.entries(updateData).filter(([, value]) => value != null),
+  )
+
+  const task = await WorkerTask.findOneAndUpdate({ id }, { $set: update }, {
+    new: true,
+    upsert: true,
+  }).exec()
+  return task
+}

packages/openneuro-server/src/graphql/schema.ts

@@ -214,6 +214,19 @@ export const typeDefs = `
       annexFsck: [AnnexFsckInput!]!
       remote: String
     ): FileCheck
+    # Update worker task queue status
+    updateWorkerTask(
+      id: ID!,
+      args: JSON,
+      kwargs: JSON,
+      taskName: String,
+      worker: String,
+      queuedAt: DateTime,
+      startedAt: DateTime,
+      finishedAt: DateTime,
+      error: String,
+      executionTime: Int,
+    ): WorkerTask
   }
 
   # Anonymous dataset reviewer
@@ -969,6 +982,18 @@ export const typeDefs = `
     input: [String]
   }
 
+  type WorkerTask {
+    id: ID!
+    args: JSON
+    kwargs: JSON
+    taskName: String
+    worker: String
+    queuedAt: DateTime
+    startedAt: DateTime
+    finishedAt: DateTime
+    error: String
+    executionTime: Int
+  }
 `
 
 schemaComposer.addTypeDefs(typeDefs)

packages/openneuro-server/src/models/worker-task.ts

@@ -0,0 +1,38 @@
+import mongoose from "mongoose"
+import type { Document } from "mongoose"
+const { Schema, model } = mongoose
+
+export interface WorkerTaskDocument extends Document {
+  id: string
+  args?: Record<string, unknown>
+  kwargs?: Record<string, unknown>
+  taskName?: string
+  worker?: string
+  queuedAt?: Date
+  startedAt?: Date
+  finishedAt?: Date
+  error?: string
+  executionTime?: number
+}
+
+const workerTaskSchema = new Schema({
+  id: { type: String, required: true, unique: true },
+  args: { type: Object },
+  kwargs: { type: Object },
+  taskName: { type: String },
+  worker: { type: String },
+  queuedAt: { type: Date },
+  startedAt: { type: Date },
+  finishedAt: { type: Date },
+  error: { type: String },
+  executionTime: { type: Number },
+})
+
+workerTaskSchema.index({ id: 1 })
+
+const WorkerTaskModel = model<WorkerTaskDocument>(
+  "WorkerTask",
+  workerTaskSchema,
+)
+
+export default WorkerTaskModel

services/datalad/datalad_service/broker/worker_middleware.py

@@ -1,19 +1,83 @@
 import random
+from datetime import datetime, timedelta, timezone, UTC
+from typing import Any
+import httpx
+import logging
+import jwt
+
+from taskiq import TaskiqMessage, TaskiqMiddleware, TaskiqResult
 
-from taskiq import TaskiqMessage, TaskiqMiddleware
 
 from datalad_service import config
 
+logger = logging.getLogger('datalad_service.' + __name__)
+
+_UPDATE_WORKER_TASK_MUTATION = """
+mutation UpdateWorkerTask(
+  $id: ID!
+  $args: JSON
+  $kwargs: JSON
+  $taskName: String
+  $worker: String
+  $queuedAt: DateTime
+  $startedAt: DateTime
+  $finishedAt: DateTime
+  $error: String
+  $executionTime: Int
+) {
+  updateWorkerTask(
+    id: $id
+    args: $args
+    kwargs: $kwargs
+    taskName: $taskName
+    worker: $worker
+    queuedAt: $queuedAt
+    startedAt: $startedAt
+    finishedAt: $finishedAt
+    error: $error
+    executionTime: $executionTime
+  ) {
+    id
+  }
+}
+"""
+
+
+def _update_worker_task_body(**kwargs):
+    """Create a GraphQL mutation body for updateWorkerTask."""
+    return {
+        'query': _UPDATE_WORKER_TASK_MUTATION,
+        'variables': kwargs,
+        'operationName': 'UpdateWorkerTask',
+    }
+
+
+def generate_worker_token():
+    utc_now = datetime.now(timezone.utc)
+    one_day_ahead = utc_now + timedelta(hours=24)
+    return jwt.encode(
+        {
+            'sub': 'dataset-worker',
+            'iat': int(utc_now.timestamp()),
+            'exp': int(one_day_ahead.timestamp()),
+            'scopes': ['dataset:worker'],
+        },
+        config.JWT_SECRET,
+        algorithm='HS256',
+    )
+
 
 class WorkerMiddleware(TaskiqMiddleware):
     """
-    This middleware adds a custom worker label to outgoing tasks scheduled by workers.
+    This middleware adds a custom worker label to outgoing tasks scheduled by workers
+    and reports task status back to the OpenNeuro API.
     """
 
     def __init__(self, worker_id=None):
         self.worker_id = worker_id
 
     async def pre_send(self, message: TaskiqMessage) -> TaskiqMessage:
+        """Assign new tasks to the correct worker."""
         if self.worker_id:
             message.labels['queue_name'] = f'worker-{self.worker_id}'
         else:
@@ -22,3 +86,58 @@ async def pre_send(self, message: TaskiqMessage) -> TaskiqMessage:
                 f'worker-{random.randint(0, config.DATALAD_WORKERS - 1)}'
             )
         return message
+
+    async def _update_task_status(self, **kwargs):
+        """Helper to send updates to the GraphQL API."""
+        api_token = generate_worker_token()
+        body = _update_worker_task_body(**kwargs)
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    url=config.GRAPHQL_ENDPOINT,
+                    json=body,
+                    headers={'Authorization': f'Bearer {api_token}'},
+                )
+                response.raise_for_status()
+                response_json = response.json()
+                if 'errors' in response_json:
+                    logger.error(
+                        f'GraphQL error updating task status: {response_json["errors"]}',
+                    )
+        except httpx.HTTPError as e:
+            logger.error(f'HTTP error updating task status: {e}')
+            logger.error(f'Response: {e.response.text}')
+        except Exception as e:
+            logger.error(f'Unexpected error updating task status: {e}')
+
+    async def post_send(self, message: TaskiqMessage):
+        """Called after a task is sent to the broker."""
+        now = datetime.now(UTC).isoformat()
+        await self._update_task_status(
+            id=message.task_id,
+            args=message.args,
+            kwargs=message.kwargs,
+            taskName=message.task_name,
+            worker=message.labels.get('queue_name'),
+            queuedAt=now,
+        )
+
+    async def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage:
+        """Called before a worker executes a task."""
+        now = datetime.now(UTC).isoformat()
+        await self._update_task_status(id=message.task_id, startedAt=now)
+        return message
+
+    async def post_execute(
+        self,
+        message: TaskiqMessage,
+        result: TaskiqResult[Any],
+    ):
+        """Called after a worker executes a task."""
+        now = datetime.now(UTC).isoformat()
+        await self._update_task_status(
+            id=message.task_id,
+            finishedAt=now,
+            error=None if result.error is None else repr(result.error),
+            executionTime=round(result.execution_time * 1000),
+        )

services/datalad/pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
     "taskiq[reload]>=0.11.18",
     "taskiq-redis>=1.0.9",
     "uvicorn[standard]>=0.34.3",
+    "httpx>=0.28.1",
 ]
 dynamic = ["version"]