Implement background transpiler for faster expansion set and sequence…

… generation downstream * Fetches the latest command dictionary, mission model, and expansion logic on a regular basis.Transpiles and caches the results. * Help to pay the upfront cost that we are seeing on Clipper where there expansion logic takes about 15 minutes to generate a expansion set of 130 authoring logic.
NASA-AMMOS · Mar 1, 2024 · 717ac54 · 717ac54
1 parent 18f8f2b
commit 717ac54
Show file tree

Hide file tree

Showing 2 changed files with 238 additions and 1 deletion.
diff --git a/sequencing-server/src/backgroundTranspiler.ts b/sequencing-server/src/backgroundTranspiler.ts
@@ -0,0 +1,143 @@
+import { piscina, graphqlClient, typeCheckingCache, promiseThrottler } from './app.js';
+import { objectCacheKeyFunction } from './lib/batchLoaders/index.js';
+import crypto from 'crypto';
+import { generateTypescriptForGraphQLActivitySchema } from './lib/codegen/ActivityTypescriptCodegen.js';
+import DataLoader from 'dataloader';
+import { activitySchemaBatchLoader } from './lib/batchLoaders/activitySchemaBatchLoader.js';
+import { commandDictionaryTypescriptBatchLoader } from './lib/batchLoaders/commandDictionaryTypescriptBatchLoader.js';
+import type { typecheckExpansion } from './worker';
+import { Result } from '@nasa-jpl/aerie-ts-user-code-runner/build/utils/monads.js';
+import { getLatestCommandDictionary, getLatestMissionModel, getExpansionRule } from './utils/hasura.js';
+
+export async function backgroundTranspiler(numberOfThreads: number = 2) {
+  if (graphqlClient === null) {
+    return;
+  }
+
+  // Fetch latest mission model
+  const { mission_model_aggregate } = await getLatestMissionModel(graphqlClient);
+  if (!mission_model_aggregate) {
+    console.log(
+      '[ Background Transpiler ] Unable to fetch the latest mission model. Aborting background transpiling...',
+    );
+    return;
+  }
+
+  // Fetch latest command dictionary
+  const { command_dictionary_aggregate } = await getLatestCommandDictionary(graphqlClient);
+  if (!command_dictionary_aggregate) {
+    console.log(
+      '[ Background Transpiler ] Unable to fetch the latest command dictionary. Aborting background transpiling...',
+    );
+    return;
+  }
+
+  const commandDictionaryId = command_dictionary_aggregate.aggregate.max.id;
+  const missionModelId = mission_model_aggregate.aggregate.max.id;
+  const { expansion_rule } = await getExpansionRule(graphqlClient, missionModelId, commandDictionaryId);
+
+  if (expansion_rule === null || expansion_rule.length === 0) {
+    console.log(`[ Background Transpiler ] No expansion rules to transpile.`);
+    return;
+  }
+
+  const commandTypescriptDataLoader = new DataLoader(commandDictionaryTypescriptBatchLoader({ graphqlClient }), {
+    cacheKeyFn: objectCacheKeyFunction,
+    name: null,
+  });
+  const activitySchemaDataLoader = new DataLoader(activitySchemaBatchLoader({ graphqlClient }), {
+    cacheKeyFn: objectCacheKeyFunction,
+    name: null,
+  });
+
+  const commandTypes = await commandTypescriptDataLoader.load({
+    dictionaryId: missionModelId,
+  });
+
+  if (commandTypes === null) {
+    console.log(`[ Background Transpiler ] Unable to fetch command ts lib.
+    Aborting transpiling...`);
+    return;
+  }
+
+  // only process 'numberOfThreads' worth at a time ex. transpile 2 logics at a time
+  // This allows for expansion set and sequence expansion to utilize the remaining workers
+  for (let i = 0; i < expansion_rule.length; i += numberOfThreads) {
+    await Promise.all(
+      expansion_rule.slice(i, i + numberOfThreads).map(async expansion => {
+        await promiseThrottler.run(async () => {
+          // Assuming expansion_rule elements have the same type
+          if (expansion instanceof Error) {
+            console.log(`[ Background Transpiler ] Expansion: ${expansion.name} could not be loaded`, expansion);
+            return Promise.reject(`Expansion: ${expansion.name} could not be loaded`);
+          }
+
+          const hash = crypto
+            .createHash('sha256')
+            .update(
+              JSON.stringify({
+                commandDictionaryId,
+                missionModelId,
+                id: expansion.id,
+                expansionLogic: expansion.expansion_logic,
+                activityType: expansion.activity_type,
+              }),
+            )
+            .digest('hex');
+
+          // ignore already transpiled hash info
+          if (typeCheckingCache.has(hash)) {
+            return Promise.resolve();
+          }
+
+          const activitySchema = await activitySchemaDataLoader.load({
+            missionModelId,
+            activityTypeName: expansion.activity_type,
+          });
+
+          // log error
+          if (!activitySchema) {
+            console.log(
+              `[ Background Transpiler ] Activity schema for ${expansion.activity_type} could not be loaded`,
+              activitySchema,
+            );
+            return Promise.reject('Activity schema for ${expansion.activity_type} could not be loaded');
+          }
+
+          const activityTypescript = generateTypescriptForGraphQLActivitySchema(activitySchema);
+
+          // log error
+          if (!activityTypescript) {
+            console.log(
+              `[ Background Transpiler ] Unable to generate typescript for activity ${expansion.activity_type}`,
+              activityTypescript,
+            );
+            return Promise.reject(`Unable to generate typescript for activity ${expansion.activity_type}`);
+          }
+
+          const typecheckingResult = (
+            piscina.run(
+              {
+                expansionLogic: expansion.expansion_logic,
+                commandTypes: commandTypes,
+                activityTypes: activityTypescript,
+                activityTypeName: expansion.activity_type,
+              },
+              { name: 'typecheckExpansion' },
+            ) as ReturnType<typeof typecheckExpansion>
+          ).then(Result.fromJSON);
+
+          //Display any errors
+          typecheckingResult.then(result => {
+            if (result.isErr()) {
+              console.log(`Error transpiling ${expansion.activity_type}:\n ${result.unwrapErr().map(e => e.message)}`);
+            }
+          });
+
+          typeCheckingCache.set(hash, typecheckingResult);
+          return typecheckingResult;
+        });
+      }),
+    );
+  }
+}
diff --git a/sequencing-server/src/utils/hasura.ts b/sequencing-server/src/utils/hasura.ts
@@ -32,7 +32,7 @@ export const ENDPOINTS_WHITELIST = new Set([
   '/seqjson/get-seqjson-for-sequence-standalone',
   '/seqjson/get-seqjson-for-seqid-and-simulation-dataset',
   '/seqjson/bulk-get-edsl-for-seqjson',
-  '/seqjson/get-edsl-for-seqjson'
+  '/seqjson/get-edsl-for-seqjson',
 ]);
 
 /**
@@ -368,3 +368,97 @@ async function getMissionModelId(
 
   return missionModelId;
 }
+
+export async function getLatestMissionModel(graphqlClient: GraphQLClient): Promise<{
+  mission_model_aggregate: {
+    aggregate: {
+      max: {
+        id: number;
+      };
+    };
+  };
+}> {
+  return graphqlClient.request<{
+    mission_model_aggregate: {
+      aggregate: {
+        max: {
+          id: number;
+        };
+      };
+    };
+  }>(
+    gql`
+      query GetLatestMissionModel {
+        mission_model_aggregate(order_by: { uploaded_file: { modified_date: asc } }) {
+          aggregate {
+            max {
+              id
+            }
+          }
+        }
+      }
+    `,
+  );
+}
+
+export async function getLatestCommandDictionary(graphqlClient: GraphQLClient): Promise<{
+  command_dictionary_aggregate: {
+    aggregate: {
+      max: {
+        id: number;
+      };
+    };
+  };
+}> {
+  return graphqlClient.request<{
+    command_dictionary_aggregate: {
+      aggregate: {
+        max: {
+          id: number;
+        };
+      };
+    };
+  }>(
+    gql`
+      query GetLatestCommandDictionary {
+        command_dictionary_aggregate(order_by: { created_at: asc }) {
+          aggregate {
+            max {
+              id
+            }
+          }
+        }
+      }
+    `,
+  );
+}
+
+export async function getExpansionRule(
+  graphqlClient: GraphQLClient,
+  missionModelId: number,
+  commandDictionaryId: number,
+): Promise<{
+  expansion_rule: {
+    id: number;
+    activity_type: string;
+    expansion_logic: string;
+  }[];
+}> {
+  return graphqlClient.request<{
+    expansion_rule: {
+      id: number;
+      activity_type: string;
+      expansion_logic: string;
+    }[];
+  }>(
+    gql`
+        query GetExpansonLogic {
+          expansion_rule(where: {authoring_command_dict_id: {_eq: ${missionModelId}}, authoring_mission_model_id: {_eq: ${commandDictionaryId} }}) {
+            id
+            activity_type
+            expansion_logic
+          }
+        }
+      `,
+  );
+}