add pinecone benchmarks and quick fixes for Astra benchmarks

stargate · Feb 20, 2024 · 5ea1ce6 · 5ea1ce6
1 parent 13588e1
commit 5ea1ce6
Show file tree

Hide file tree

Showing 5 changed files with 154 additions and 56 deletions.
diff --git a/benchmarks/README.md b/benchmarks/README.md
@@ -70,4 +70,26 @@ $ env ASTRA_CONNECTION_STRING="https://abc-us-east-2.apps.astra.datastax.com/...
   "name": "benchmark-findone-mongoose-vector-astra",
   "totalTimeMS": 20436
 }
+```
+
+## Pinecone Benchmarks
+
+First, install the Pinecone client:
+
+```
+npm install @pinecone-database/pinecone --no-save
+```
+
+Then set up sample data:
+
+```
+```
+env PINECONE_INDEX=your-index-here env PINECONE_API_KEY=your-api-key-here ./node_modules/.bin/ts-node benchmarks/setup-pinecone-vector-example.ts
+```
+```
+
+Then run Pinecone benchmarks:
+
+```
+env PINECONE_INDEX=your-index-here env PINECONE_API_KEY=your-api-key-here ./node_modules/.bin/ts-node benchmarks/benchmark-findone-pinecone.ts
 ```
diff --git a/benchmarks/benchmark-findone-pinecone.ts b/benchmarks/benchmark-findone-pinecone.ts
@@ -0,0 +1,36 @@
+main().then(
+    () => { process.exit(0); },
+    err => {
+        console.error(err);
+        process.exit(-1);
+    }
+);
+
+async function main() {
+    let Pinecone;
+    try {
+        Pinecone = require('@pinecone-database/pinecone').Pinecone;
+    } catch (err) {
+        console.log('{"name":"benchmark-findone-pinecone"}');
+        return;
+    }
+
+    if (!process.env.PINECONE_API_KEY || !process.env.PINECONE_INDEX) {
+        console.log('{"name":"benchmark-findone-pinecone"}');
+        return;
+    }
+
+    const pinecone = new Pinecone();
+
+    const vector = [1, ...Array(1535).fill(0)];
+    const index = process.env.PINECONE_INDEX;
+    const start = Date.now();
+    for (let i = 0; i < 100; ++i) {
+        await pinecone.index(index).query({ topK: 1, vector, includeValues: true });
+    }
+    const results = {
+        name: 'benchmark-findone-pinecone',
+        totalTimeMS: Date.now() - start
+    };
+    console.log(JSON.stringify(results, null, '  '));
+}
diff --git a/benchmarks/setup-astra-vector-example.ts b/benchmarks/setup-astra-vector-example.ts
@@ -15,57 +15,59 @@ main().then(
 );
 
 async function main() {
-  if (!process.env.ASTRA_CONNECTION_STRING) {
-      throw new Error('Must set ASTRA_CONNECTION_STRING');
-  }
-  if (!process.env.ASTRA_COLLECTION_NAME) {
-    throw new Error('Must set ASTRA_COLLECTION_NAME');
-}
-  await mongoose.connect(process.env.ASTRA_CONNECTION_STRING, {
-      isAstra: true
-  } as mongoose.ConnectOptions);
+    if (!process.env.ASTRA_CONNECTION_STRING) {
+        throw new Error('Must set ASTRA_CONNECTION_STRING');
+    }
+    if (!process.env.ASTRA_COLLECTION_NAME) {
+        throw new Error('Must set ASTRA_COLLECTION_NAME');
+    }
+    await mongoose.connect(process.env.ASTRA_CONNECTION_STRING, {
+        isAstra: true,
+        useHTTP2: false
+    } as mongoose.ConnectOptions);
 
-  const Vector = mongoose.model(
-      'Vector',
-      new mongoose.Schema({
-          $vector: {
-              type: [Number]
-          },
-          prompt: {
-              type: String,
-              required: true
-          }
-      }, {
-          collectionOptions: { vector: { dimension: 1536, metric: 'cosine' } },
-          autoCreate: false
-      }),
-      process.env.ASTRA_COLLECTION_NAME
-  );
+    const Vector = mongoose.model(
+        'Vector',
+        new mongoose.Schema({
+            $vector: {
+                type: [Number]
+            },
+            prompt: {
+                type: String,
+                required: true
+            }
+        }, {
+            collectionOptions: { vector: { dimension: 1536, metric: 'cosine' } },
+            autoCreate: false
+        }),
+        process.env.ASTRA_COLLECTION_NAME
+    );
 
-  console.log('Recreating collection...');
-  await Vector.db.dropCollection(process.env.ASTRA_COLLECTION_NAME);
-  await Vector.createCollection();
+    console.log('Recreating collection...');
+    await Vector.db.dropCollection(process.env.ASTRA_COLLECTION_NAME);
+    await Vector.createCollection();
 
-  console.log('Creating vectors...');
-  const numVectors = 1000;
-  for (let i = 0; i < numVectors; ++i) {
-      console.log(`${i} / ${numVectors}`);
-      const $vector = Array(1536).fill(0);
-      $vector[i] = 1;
-      for (let j = 0; j < 3; ++j) {
-        try {
-          await Vector.create({ $vector, prompt: `Test ${i}` });
-          break;
-        } catch (err) {
-          if (j >= 2) {
-            throw err;
-          }
-          await new Promise(resolve => setTimeout(resolve, 100));
+    console.log('Creating vectors...');
+    const numVectors = 1000;
+    const start = Date.now();
+    for (let i = 0; i < numVectors; ++i) {
+        console.log(`${i} / ${numVectors}`);
+        const $vector = Array(1536).fill(0);
+        $vector[i] = 1;
+        for (let j = 0; j < 3; ++j) {
+            try {
+                await Vector.create({ $vector, prompt: `Test ${i}` });
+                break;
+            } catch (err) {
+                if (j >= 2) {
+                    throw err;
+                }
+                await new Promise(resolve => setTimeout(resolve, 100));
+            }
         }
-      }
 
-  }
+    }
 
-  console.log('Done');
-  process.exit(0);
+    console.log('Done', Date.now() - start);
+    process.exit(0);
 }
diff --git a/benchmarks/setup-pinecone-vector-example.ts b/benchmarks/setup-pinecone-vector-example.ts
@@ -0,0 +1,38 @@
+import { Pinecone } from '@pinecone-database/pinecone';
+
+main().then(
+    () => { process.exit(0); },
+    err => {
+        console.error(err);
+        process.exit(-1);
+    }
+);
+
+async function main() {
+    if (!process.env.PINECONE_API_KEY || !process.env.PINECONE_INDEX) {
+        throw new Error('Cannot run pinecone setup without an API key');
+    }
+
+    const pinecone = new Pinecone();
+
+    console.log('Creating vectors...');
+
+    const numVectors = 1000;
+    const index = process.env.PINECONE_INDEX;
+    const start = Date.now();
+    for (let i = 0; i < numVectors; ++i) {
+        console.log(`${i} / ${numVectors}`);
+        const values: number[] = Array(1536).fill(0);
+        values[i] = 1;
+        await pinecone.index(index).upsert([{
+            id: '' + i,
+            values,
+            metadata: {
+                prompt: `Test ${i}`
+            }
+        }]);
+    }
+
+    console.log('Done', Date.now() - start);
+    process.exit(0);
+}
diff --git a/benchmarks/upload-results.ts b/benchmarks/upload-results.ts
@@ -31,15 +31,15 @@ async function main() {
         'BenchmarkResult',
         new mongoose.Schema({
             githash: {
-              type: String,
-              required: true
+                type: String,
+                required: true
             },
             name: {
-              type: String,
-              required: true
+                type: String,
+                required: true
             },
             totalTimeMS: {
-              type: Number
+                type: Number
             }
         }, { autoCreate: true, timestamps: true })
     );
@@ -53,12 +53,12 @@ async function main() {
             fs.readFileSync(path.join('.', 'benchmarks', file), 'utf8')
         );
         if (totalTimeMS == null) {
-          continue;
+            continue;
         }
         const doc = await BenchmarkResult.findOneAndUpdate(
-          { githash, name },
-          { totalTimeMS },
-          { upsert: true, returnDocument: 'after' }
+            { githash, name },
+            { totalTimeMS },
+            { upsert: true, returnDocument: 'after' }
         );
         console.log(doc);
     }