TaskMan

serverless.yml

@@ -39,6 +39,9 @@ provider:
             - Fn::GetAtt:
                 - exportQueue
                 - Arn
+            - Fn::GetAtt:
+                - taskQueue
+                - Arn
         - Effect: Allow
           Action:
             - sagemaker:InvokeEndpoint
@@ -181,6 +184,19 @@ functions:
               - Arn
           batchSize: 1
     timeout: 120
+  task:
+    image:
+      name: base
+      command: src/task/handler.handler
+    reservedConcurrency: 10
+    events:
+      - sqs:
+          arn:
+            Fn::GetAtt:
+              - taskQueue
+              - Arn
+          batchSize: 1
+    timeout: 900
 
 resources:
   Resources:
@@ -215,7 +231,14 @@ resources:
         QueueName: inferenceDLQ-${self:provider.stage}
         MessageRetentionPeriod: 1209600 # 14 days in seconds
 
-    # Export SQS queue
+    # Task SQS queue
+    taskQueue:
+      Type: AWS::SQS::Queue
+      Properties:
+        QueueName: taskQueue-${self:provider.stage}
+        VisibilityTimeout: 900
+
+    # Export SQS queue - will eventually be replaced by unified tasks queue
     exportQueue:
       Type: AWS::SQS::Queue
       Properties:
@@ -240,6 +263,15 @@ resources:
         Type: String
         Value: !Ref inferenceQueue
 
+    # SSM Param - task queue url
+    SSMParameterTaskqueue:
+      Type: AWS::SSM::Parameter
+      Properties:
+        Description: Task queue url
+        Name: /tasks/task-queue-url-${opt:stage, self:provider.stage, 'dev'}
+        Type: String
+        Value: !Ref taskQueue
+
     # SSM Param - export queue url
     SSMParameterExportqueue:
       Type: AWS::SSM::Parameter

src/api/auth/roles.js

@@ -6,6 +6,7 @@ const EXPORT_DATA_ROLES =                [MANAGER, MEMBER];
 const WRITE_OBJECTS_ROLES =              [MANAGER, MEMBER];
 const WRITE_VIEWS_ROLES =                [MANAGER, MEMBER];
 const WRITE_COMMENTS_ROLES =             [MANAGER, MEMBER];
+const READ_TASKS_ROLES =                 [MANAGER, MEMBER];
 const WRITE_PROJECT_ROLES =              [MANAGER];
 const WRITE_IMAGES_ROLES =               [MANAGER];
 const DELETE_IMAGES_ROLES =              [MANAGER];
@@ -15,6 +16,7 @@ const WRITE_AUTOMATION_RULES_ROLES =     [MANAGER];
 const WRITE_CAMERA_REGISTRATION_ROLES =  [MANAGER];
 
 export {
+  READ_TASKS_ROLES,
   WRITE_COMMENTS_ROLES,
   DELETE_IMAGES_ROLES,
   EXPORT_DATA_ROLES,

src/api/db/models/Image.js

@@ -6,6 +6,7 @@ import GraphQLError, { InternalServerError, ForbiddenError, DuplicateImageError,
 import crypto from 'node:crypto';
 import mongoose from 'mongoose';
 import MongoPaging from 'mongo-cursor-pagination';
+import { TaskModel } from './Task.js';
 import Image from '../schemas/Image.js';
 import Project from '../schemas/Project.js';
 import ImageError from '../schemas/ImageError.js';
@@ -972,7 +973,12 @@ export default class AuthedImageModel {
   }
 
   async getStats(input, context) {
-    return await ImageModel.getStats(input, context);
+    return await TaskModel.create({
+      type: 'GetStats',
+      projectId: context.user['curr_project'],
+      user: context.user['cognito:username'],
+      config: input
+    }, context);
   }
 
   async export(input, context) {

src/api/db/models/Project.js

@@ -488,7 +488,9 @@ export class ProjectModel {
       const label = (project.labels || []).filter((p) => { return p._id.toString() === input._id.toString(); })[0];
       if (!label) throw new DeleteLabelError('Label not found on project');
 
-      const count = await ImageModel.countImagesByLabel([input._id], context);
+      const count = await ImageModel.countImagesByLabel([input._id], {
+        limit: MAX_LABEL_DELETE
+      }, context);
 
       if (count > MAX_LABEL_DELETE) {
         const msg = `This label is already in extensive use (>${MAX_LABEL_DELETE} images) and cannot be ` +

src/api/db/models/Task.js

@@ -0,0 +1,101 @@
+import { NotFoundError, ForbiddenError, AuthenticationError } from '../../errors.js';
+import SQS from '@aws-sdk/client-sqs';
+import MongoPaging from 'mongo-cursor-pagination';
+import Task from '../schemas/Task.js';
+import { hasRole } from './utils.js';
+import {
+  READ_TASKS_ROLES
+} from '../../auth/roles.js';
+
+/**
+ * Tasks manage the state of async events (except for batch uploads) on the platform
+ * @class
+ */
+export class TaskModel {
+  /**
+   * Query Tasks by Filter, returning a paged list
+   *
+   * @param {Object} input
+   * @param {String} input.limit
+   * @param {String} input.paginatedField
+   * @param {String} input.sortAscending
+   * @param {String} input.next
+   * @param {String} input.previous
+   * @param {Object} context
+   */
+  static async queryByFilter(input, context) {
+    return await MongoPaging.aggregate(Task.collection, {
+      aggregation: [
+        { '$match': { 'projectId': context.user['curr_project'] } },
+        { '$match': { 'user': context.user.sub } }
+      ],
+      limit: input.limit,
+      paginatedField: input.paginatedField,
+      sortAscending: input.sortAscending,
+      next: input.next,
+      previous: input.previous
+    });
+  }
+
+  static async queryById(_id, context) {
+    const query = { _id };
+    const task = await Task.findOne(query);
+    if (!task) throw new NotFoundError('Task not found');
+
+    if (task.projectId !== context.user['curr_project']) {
+      throw new NotFoundError('Task does not belong to current project');
+    }
+
+    return task;
+  }
+
+  static async create(input, context) {
+    const task = new Task({
+      user: input.user,
+      projectId: input.projectId,
+      type: input.type
+    });
+
+    const sqs = new SQS.SQSClient({ region: process.env.AWS_DEFAULT_REGION });
+
+    await task.save();
+
+    await sqs.send(new SQS.SendMessageCommand({
+      QueueUrl: context.config['/TASKS/TASK_QUEUE_URL'],
+      MessageBody: JSON.stringify({
+        config: input.config,
+        ...task.toJSON()
+      })
+    }));
+    return task;
+  }
+
+  static async update(input, context) {
+    const task = await this.queryById(input._id, context);
+
+    input.updated = new Date();
+
+    Object.assign(task, input);
+    await task.save();
+    return task;
+  }
+}
+
+export default class AuthedTaskModel {
+  constructor(user) {
+    if (!user) throw new AuthenticationError('Authentication failed');
+    this.user = user;
+  }
+
+  async queryById(input, context) {
+    if (!hasRole(this.user, READ_TASKS_ROLES)) throw new ForbiddenError();
+
+    return await TaskModel.queryById(input, context);
+  }
+
+  async queryByFilter(input, context) {
+    if (!hasRole(this.user, READ_TASKS_ROLES)) throw new ForbiddenError();
+
+    return await TaskModel.queryByFilter(input, context);
+  }
+}

src/api/db/schemas/Task.js

@@ -0,0 +1,27 @@
+import mongoose from 'mongoose';
+import MongoPaging from 'mongo-cursor-pagination';
+
+const Schema = mongoose.Schema;
+
+const TaskSchema = new Schema({
+  user: { type: String, required: true },
+  projectId: { type: String, required: true, ref: 'Project' },
+  type: {
+    type: String,
+    required: true,
+    enum : ['GetStats']
+  },
+  status: {
+    type: String,
+    required: true,
+    enum : ['SUBMITTED', 'RUNNING', 'FAIL', 'COMPLETE'],
+    default: 'SUBMITTED'
+  },
+  created: { type: Date, default: Date.now, required: true },
+  updated: { type: Date, default: Date.now, required: true },
+  output: { type: Object }
+});
+
+TaskSchema.plugin(MongoPaging.mongoosePlugin);
+
+export default mongoose.model('Task', TaskSchema);

src/api/handler.js

@@ -8,6 +8,7 @@ import AuthedMLModelModel from './db/models/MLModel.js';
 import AuthedBatchModel from './db/models/Batch.js';
 import AuthedBatchErrorModel from './db/models/BatchError.js';
 import AuthedImageErrorModel from './db/models/ImageError.js';
+import AuthedTaskModel from './db/models/Task.js';
 import Query from './resolvers/Query.js';
 import Mutation from './resolvers/Mutation.js';
 import Scalars from './resolvers/Scalars.js';
@@ -52,6 +53,7 @@ const context = async ({ event, context }) => {
     config,
     models: {
       User: new AuthedUserModel(user),
+      Task: new AuthedTaskModel(user),
       Project: new AuthedProjectModel(user),
       Image: new AuthedImageModel(user),
       ImageError: new AuthedImageErrorModel(user),

src/api/resolvers/Query.js

@@ -7,6 +7,24 @@ const Query = {
     return await context.models.User.listUsers(input, context);
   },
 
+  tasks: async (_, { input }, context) => {
+    const response = await context.models.Task.queryByFilter(input, context);
+    const { previous, hasPrevious, next, hasNext, results } = response;
+    return {
+      pageInfo: {
+        previous,
+        hasPrevious,
+        next,
+        hasNext
+      },
+      tasks: results
+    };
+  },
+
+  task: async (_, { input }, context) => {
+    return await context.models.Task.queryById(input, context);
+  },
+
   batches: async (_, { input }, context) => {
     const response = await context.models.Batch.queryByFilter(input, context);
     const { previous, hasPrevious, next, hasNext, results } = response;

src/api/type-defs/inputs/QueryTasksInput.js

@@ -0,0 +1,8 @@
+export default `
+  input QueryTasksInput {
+    paginatedField: String
+    sortAscending: Boolean
+    limit: Int
+    next: String
+    previous: String
+  }`;

src/api/type-defs/objects/Task.js

@@ -0,0 +1,11 @@
+export default `
+  type Task {
+    _id: ID!
+    user: String!
+    projectId: String!
+    type: String!
+    status: String!
+    created: String!
+    updated: String!
+    output: JSONObject
+  }`;

src/api/type-defs/payloads/TasksPayload.js

@@ -0,0 +1,5 @@
+export default `
+  type TasksPayload {
+    tasks: [Task]!
+  }
+`;

src/api/type-defs/root/Query.js

@@ -1,6 +1,8 @@
 export default `
   type Query {
     users(input: QueryUsersInput): UsersPayload
+    tasks(input: QueryTasksInput): TasksPayload
+    task(input: String!): Task
     projects(input: QueryProjectsInput): [Project]
     image(input: QueryImageInput!): Image
     images(input: QueryImagesInput!): ImagesConnection
@@ -9,10 +11,10 @@ export default `
     wirelessCameras(input: QueryWirelessCamerasInput): [WirelessCamera]
     mlModels(input: QueryMLModelsInput): [MLModel]
     batches(input: QueryBatchesInput!): BatchesConnection
-    stats(input: QueryStatsInput!): ImagesStats
     export(input: ExportInput!): ExportPayload
     exportErrors(input: ExportErrorsInput!): ExportPayload
     exportStatus(input: ExportStatusInput!): ExportStatusPayload
+    stats(input: QueryStatsInput!): Task
   }
 `;
 

src/config/config.js

@@ -52,6 +52,7 @@ const ssmNames = [
   `/ml/inference-queue-url-${process.env.STAGE}`,
   `/exports/exported-data-bucket-${process.env.STAGE}`,
   `/exports/export-queue-url-${process.env.STAGE}`,
+  `/tasks/task-queue-url-${process.env.STAGE}`,
   `/ml/megadetector-v5a-realtime-endpoint-${process.env.STAGE}`,
   `/ml/megadetector-v5a-batch-endpoint-${process.env.STAGE}`,
   `/ml/megadetector-v5b-realtime-endpoint-${process.env.STAGE}`,