feat(worker): Report task runtime state to API

This allows for showing status of tasks and debugging tasks more easily in production.

Author: nellh

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/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,84 @@
 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
+        self.api_token = generate_worker_token()
 
     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 +87,61 @@ 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."""
+        if not self.api_token:
+            logger.warning('DATALAD_WORKER_TOKEN not set, cannot update task status.')
+            return
+
+        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 {self.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),
+        )