From 828eb037f86119704c7ecd9c0a63dc0bc42d95d6 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Thu, 13 Jun 2024 14:07:50 +0700
Subject: [PATCH] chore: add benchmark options

---
 .../commanders/benchmark.command.ts           | 34 +++++++++++++++++--
 .../usecases/benchmark.cli.usecases.ts        | 27 ++++++++++++---
 .../src/infrastructure/constants/benchmark.ts |  4 +--
 3 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/cortex-js/src/infrastructure/commanders/benchmark.command.ts b/cortex-js/src/infrastructure/commanders/benchmark.command.ts
index e3e6a69fd..96f94f9fc 100644
--- a/cortex-js/src/infrastructure/commanders/benchmark.command.ts
+++ b/cortex-js/src/infrastructure/commanders/benchmark.command.ts
@@ -1,5 +1,6 @@
-import { CommandRunner, SubCommand } from 'nest-commander';
+import { CommandRunner, SubCommand, Option } from 'nest-commander';
 import { BenchmarkCliUsecases } from './usecases/benchmark.cli.usecases';
+import { BenchmarkConfig } from './types/benchmark-config.interface';
 
 @SubCommand({
   name: 'benchmark',
@@ -12,7 +13,34 @@ export class BenchmarkCommand extends CommandRunner {
     super();
   }
 
-  async run(): Promise<void> {
-    return this.benchmarkUsecases.benchmark();
+  async run(
+    _input: string[],
+    options?: Partial<BenchmarkConfig>,
+  ): Promise<void> {
+    return this.benchmarkUsecases.benchmark(options ?? {});
+  }
+
+  @Option({
+    flags: '-n, --num_rounds <num_rounds>',
+    description: 'Number of rounds to run the benchmark',
+  })
+  parseRounds(value: number) {
+    return value;
+  }
+
+  @Option({
+    flags: '-c, --concurrency <concurrency>',
+    description: 'Number of concurrent requests to run the benchmark',
+  })
+  parseConcurrency(value: number) {
+    return value;
+  }
+
+  @Option({
+    flags: '-o, --output <output>',
+    description: 'Output format for the benchmark results. json or table',
+  })
+  parseOutput(value: string) {
+    return value;
   }
 }
diff --git a/cortex-js/src/infrastructure/commanders/usecases/benchmark.cli.usecases.ts b/cortex-js/src/infrastructure/commanders/usecases/benchmark.cli.usecases.ts
index dfd436cae..19cf6c951 100644
--- a/cortex-js/src/infrastructure/commanders/usecases/benchmark.cli.usecases.ts
+++ b/cortex-js/src/infrastructure/commanders/usecases/benchmark.cli.usecases.ts
@@ -28,9 +28,12 @@ export class BenchmarkCliUsecases {
   /**
    * Benchmark and analyze the performance of a specific AI model using a variety of system resources
    */
-  async benchmark() {
+  async benchmark(options: Partial<BenchmarkConfig>) {
     return this.getBenchmarkConfig().then((config) => {
-      this.config = config;
+      this.config = {
+        ...config,
+        ...options,
+      };
 
       // TODO: Using OpenAI client or Cortex client to benchmark?
       this.openai = new OpenAI({
@@ -41,6 +44,7 @@ export class BenchmarkCliUsecases {
 
       const serveProcess = spawn('cortex', ['serve'], {
         detached: false,
+        shell: process.platform == 'win32',
       });
 
       return this.cortexUsecases
@@ -261,8 +265,21 @@ export class BenchmarkCliUsecases {
     fs.writeFileSync(outputFilePath, JSON.stringify(output, null, 2));
     console.log(`Benchmark results and metrics saved to ${outputFilePath}`);
 
-    console.log(
-      inspect(output, { showHidden: false, depth: null, colors: true }),
-    );
+    if (this.config.output === 'table') {
+      console.log('Results:');
+      output.results.forEach((round) => {
+        console.log('Round ' + round.round + ':');
+        console.table(round.results);
+      });
+      console.log('Metrics:');
+      console.table(output.metrics);
+    } else
+      console.log(
+        inspect(output, {
+          showHidden: false,
+          depth: null,
+          colors: true,
+        }),
+      );
   }
 }
diff --git a/cortex-js/src/infrastructure/constants/benchmark.ts b/cortex-js/src/infrastructure/constants/benchmark.ts
index 6b2edaed6..a55574c73 100644
--- a/cortex-js/src/infrastructure/constants/benchmark.ts
+++ b/cortex-js/src/infrastructure/constants/benchmark.ts
@@ -26,11 +26,11 @@ export const defaultBenchmarkConfiguration: BenchmarkConfig = {
     },
   },
   prompts: {
-    min: 102,
+    min: 1024,
     max: 2048,
     samples: 10,
   },
-  output: 'json',
+  output: 'table',
   hardware: ['cpu', 'gpu', 'psu', 'chassis', 'ram'],
   concurrency: 1,
   num_rounds: 10,