Allow complex action setups in benches (windmill-labs#869)

Co-authored-by: Ruben Fiszel <[email protected]>

benchmarks/action.ts

@@ -0,0 +1,125 @@
+import {
+  FlowValue,
+  JobService,
+  Preview as ScriptPreview,
+} from "https://deno.land/x/[email protected]/windmill-api/index.ts";
+
+export type Action = FlowAction | ScriptAction | RandomAction;
+
+export type RandomAction = {
+  type: "RANDOM";
+  actions: {
+    action: Action;
+    weight: number;
+  }[];
+};
+
+export type FlowAction = PreviewFlowAction | PathFlowAction;
+export type PreviewFlowAction = {
+  type: "PREVIEW_FLOW";
+  value: FlowValue;
+  args: Record<string, any>;
+  workspace: string;
+};
+export type PathFlowAction = {
+  type: "PATH_FLOW";
+  path: string;
+  workspace: string;
+  args: Record<string, any>;
+};
+
+export type ScriptAction = PreviewScriptAction | HashScriptAction;
+export type PreviewScriptAction = {
+  type: "PREVIEW_SCRIPT";
+  content: string;
+  language: ScriptPreview.language;
+  args: Record<string, any>;
+  workspace: string;
+};
+export type HashScriptAction = {
+  type: "HASH_SCRIPT";
+  hash: string;
+  workspace: string;
+  args: Record<string, any>;
+};
+
+export function evaluate(action: Action): Promise<void> {
+  if (action.type == "HASH_SCRIPT") {
+    return evaluateHashScript(action);
+  } else if (action.type == "PATH_FLOW") {
+    return evaluatePathFlow(action);
+  } else if (action.type == "PREVIEW_FLOW") {
+    return evaluatePreviewFlow(action);
+  } else if (action.type == "PREVIEW_SCRIPT") {
+    return evaluatePreviewScript(action);
+  } else if (action.type == "RANDOM") {
+    return evaluateRandom(action);
+  } else {
+    throw new Error("UNHANDLED ACTION TYPE!!");
+  }
+}
+
+export function evaluateRandom(action: RandomAction): Promise<void> {
+  if (action.actions.length < 1) {
+    console.log("empty random action");
+    return Promise.resolve();
+  } else if (action.actions.length == 1) {
+    return evaluate(action.actions[0].action);
+  } else {
+    const sequentialWeightedActions: { weight: number; action: Action }[] =
+      new Array(action.actions.length);
+    let total_weight = 0;
+    for (let i = 0; i < action.actions.length; i++) {
+      sequentialWeightedActions[i] = {
+        weight: total_weight + action.actions[i].weight,
+        action: action.actions[i].action,
+      };
+      total_weight += action.actions[i].weight;
+    }
+    const r = Math.random() * total_weight;
+    const selectedAction = sequentialWeightedActions.find(
+      (x) => x.weight >= r
+    )!.action;
+
+    return evaluate(selectedAction);
+  }
+}
+
+async function evaluatePreviewScript(
+  action: PreviewScriptAction
+): Promise<void> {
+  await JobService.runScriptPreview({
+    workspace: action.workspace,
+    requestBody: {
+      content: action.content,
+      language: action.language,
+      args: action.args,
+    },
+  });
+}
+
+async function evaluateHashScript(action: HashScriptAction): Promise<void> {
+  await JobService.runScriptByHash({
+    workspace: action.workspace,
+    hash: action.hash,
+    requestBody: action.args,
+  });
+}
+
+async function evaluatePreviewFlow(action: PreviewFlowAction): Promise<void> {
+  await JobService.runFlowPreview({
+    workspace: action.workspace,
+    requestBody: {
+      args: action.args,
+      value: action.value,
+    },
+  });
+}
+
+async function evaluatePathFlow(action: PathFlowAction): Promise<void> {
+  await JobService.runFlowByPath({
+    workspace: action.workspace,
+    path: action.path,
+    requestBody: action.args,
+  });
+}

benchmarks/example_fake_user.json

@@ -0,0 +1,53 @@
+{
+  "type": "RANDOM",
+  "actions": [
+    {
+      "weight": 1,
+      "action": {
+        "type": "RANDOM",
+        "actions": [
+          {
+            "weight": 1,
+            "action": {
+              "type": "PREVIEW_SCRIPT",
+              "workspace": "demo",
+              "language": "deno",
+              "args": {},
+              "content": "export async function main() { return \"Hello World\"; }"
+            }
+          },
+          {
+            "weight": 1,
+            "action": {
+              "type": "PREVIEW_SCRIPT",
+              "workspace": "demo",
+              "language": "python3",
+              "args": {},
+              "content": "def main(): return \"Hello World\";"
+            }
+          },
+          {
+            "weight": 1,
+            "action": {
+              "type": "PREVIEW_SCRIPT",
+              "workspace": "demo",
+              "language": "go",
+              "args": {},
+              "content": "func main() { return \"Hello World\" }"
+            }
+          }
+        ]
+      }
+    },
+    {
+      "weight": 0.5,
+      "action": {
+        "type": "PREVIEW_SCRIPT",
+        "workspace": "demo",
+        "language": "deno",
+        "args": {},
+        "content": "import { delay } from \"https://deno.land/[email protected]/async/delay.ts\"; export async function main() { await delay(1000); return \"Hello World\"; }"
+      }
+    }
+  ]
+}

benchmarks/main.ts

@@ -4,6 +4,7 @@
 import { Command } from "https://deno.land/x/[email protected]/command/mod.ts";
 import { sleep } from "https://deno.land/x/[email protected]/mod.ts";
 import * as windmill from "https://deno.land/x/[email protected]/mod.ts";
+import { Action } from "./action.ts";
 
 async function login(email: string, password: string): Promise<string> {
   return await windmill.UserService.login({
@@ -35,6 +36,7 @@ await new Command()
       default: 30,
     }
   )
+  .option("--max <max:number>", "Maximum number of operations performed.")
   .option("-e --email <email:string>", "The email to use to login.")
   .option("-p --password <password:string>", "The password to use to login.")
   .env(
@@ -81,6 +83,7 @@ await new Command()
     }
   )
   .option("--use-flows", "Run flows instead of jobs.")
+  .option("--custom <custom_path:string>", "Use custom actions during bench")
   .option(
     "--zombie-timeout",
     "The maximum time in ms to wait for jobs to complete.",
@@ -132,9 +135,15 @@ await new Command()
       useFlows,
       zombieTimeout,
       continous,
+      max,
+      custom,
     }) => {
       windmill.setClient("", host);
 
+      const custom_content: Action | undefined = custom
+        ? JSON.parse(await Deno.readTextFile(custom))
+        : undefined;
+
       if (!Array.isArray(histogramBuckets)) {
         histogramBuckets = [];
       }
@@ -209,13 +218,16 @@ await new Command()
       windmill.setClient(final_token, host);
 
       const per_worker_throughput = maximumThroughput / num_workers;
+      const max_per_worker = max ? max / num_workers : undefined;
       const shared_config = {
         server: host,
         token: final_token,
         workspace_id: config.workspace_id,
         per_worker_throughput,
+        max_per_worker,
         useFlows,
         continous,
+        custom: custom_content,
       };
 
       let workers: Worker[] = new Array(num_workers);

benchmarks/worker.ts

@@ -1,16 +1,18 @@
 /// <reference no-default-lib="true" />
 /// <reference lib="deno.worker" />
-import { boolean } from "https://deno.land/x/[email protected]/flags/types/boolean.ts";
 import { sleep } from "https://deno.land/x/[email protected]/sleep.ts";
 import * as windmill from "https://deno.land/x/[email protected]/mod.ts";
 import * as api from "https://deno.land/x/[email protected]/windmill-api/index.ts";
 import { Job } from "https://deno.land/x/[email protected]/windmill-api/index.ts";
+import { Action, evaluate } from "./action.ts";
 
 const promise = new Promise<{
   workspace_id: string;
   per_worker_throughput: number;
   useFlows: boolean;
   continous: boolean;
+  max_per_worker: number;
+  custom: Action | undefined;
 }>((resolve, _reject) => {
   self.onmessage = (evt) => {
     const sharedConfig = evt.data;
@@ -20,6 +22,8 @@ const promise = new Promise<{
       per_worker_throughput: sharedConfig.per_worker_throughput,
       useFlows: sharedConfig.useFlows,
       continous: sharedConfig.continous,
+      max_per_worker: sharedConfig.max_per_worker,
+      custom: sharedConfig.custom,
     };
     self.name = "Worker " + sharedConfig.i;
     resolve(config);
@@ -60,8 +64,15 @@ while (cont) {
     await sleep(0.1);
     continue;
   }
+  total_spawned++;
+  if (total_spawned > config.max_per_worker) {
+    break;
+  }
   let uuid: string;
-  if (config.useFlows) {
+  if (config.custom) {
+    await evaluate(config.custom);
+    continue;
+  } else if (config.useFlows) {
     uuid = await windmill.JobService.runFlowPreview({
       workspace: config.workspace_id,
       requestBody: {
@@ -101,7 +112,6 @@ while (cont) {
     });
   }
   if (!config.continous) outstanding.push(uuid);
-  total_spawned++;
 }
 
 clearInterval(updateStatusInterval);