Add support for hardware.list endpoint

Author: mattt

index.d.ts

@@ -14,6 +14,11 @@ declare module 'replicate' {
     models?: Model[];
   }
 
+  export interface Hardware {
+    sku: string;
+    name: string
+  }
+
   export interface Model {
     url: string;
     owner: string;
@@ -115,6 +120,25 @@ declare module 'replicate' {
       get(collection_slug: string): Promise<Collection>;
     };
 
+    deployments: {
+      predictions: {
+        create(
+          deployment_owner: string,
+          deployment_name: string,
+          options: {
+            input: object;
+            stream?: boolean;
+            webhook?: string;
+            webhook_events_filter?: WebhookEventType[];
+          }
+        ): Promise<Prediction>;
+      };
+    };
+
+    hardware: {
+      list(): Promise<Hardware[]>
+    }
+
     models: {
       get(model_owner: string, model_name: string): Promise<Model>;
       list(): Promise<Page<Model>>;
@@ -157,20 +181,5 @@ declare module 'replicate' {
       cancel(training_id: string): Promise<Training>;
       list(): Promise<Page<Training>>;
     };
-
-    deployments: {
-      predictions: {
-        create(
-          deployment_owner: string,
-          deployment_name: string,
-          options: {
-            input: object;
-            stream?: boolean;
-            webhook?: string;
-            webhook_events_filter?: WebhookEventType[];
-          }
-        ): Promise<Prediction>;
-      };
-    };
   }
 }

index.js

@@ -3,6 +3,7 @@ const { withAutomaticRetries } = require('./lib/util');
 
 const collections = require('./lib/collections');
 const deployments = require('./lib/deployments');
+const hardware = require('./lib/hardware');
 const models = require('./lib/models');
 const predictions = require('./lib/predictions');
 const trainings = require('./lib/trainings');
@@ -49,6 +50,16 @@ class Replicate {
       get: collections.get.bind(this),
     };
 
+    this.deployments = {
+      predictions: {
+        create: deployments.predictions.create.bind(this),
+      }
+    };
+
+    this.hardware = {
+      list: hardware.list.bind(this),
+    };
+
     this.models = {
       get: models.get.bind(this),
       list: models.list.bind(this),
@@ -71,12 +82,6 @@ class Replicate {
       cancel: trainings.cancel.bind(this),
       list: trainings.list.bind(this),
     };
-
-    this.deployments = {
-      predictions: {
-        create: deployments.predictions.create.bind(this),
-      }
-    };
   }
 
   /**

index.test.ts

@@ -136,20 +136,20 @@ describe('Replicate client', () => {
       nock(BASE_URL)
         .get('/models')
         .reply(200, {
-          results: [{ url: 'https://replicate.com/some-user/model-1' }],
+          results: [ { url: 'https://replicate.com/some-user/model-1' } ],
           next: 'https://api.replicate.com/v1/models?cursor=cD0yMDIyLTAxLTIxKzIzJTNBMTglM0EyNC41MzAzNTclMkIwMCUzQTAw',
         })
         .get('/models?cursor=cD0yMDIyLTAxLTIxKzIzJTNBMTglM0EyNC41MzAzNTclMkIwMCUzQTAw')
         .reply(200, {
-          results: [{ url: 'https://replicate.com/some-user/model-2' }],
+          results: [ { url: 'https://replicate.com/some-user/model-2' } ],
           next: null,
         });
 
       const results: Model[] = [];
       for await (const batch of client.paginate(client.models.list)) {
         results.push(...batch);
       }
-      expect(results).toEqual([{ url: 'https://replicate.com/some-user/model-1' }, { url: 'https://replicate.com/some-user/model-2' }]);
+      expect(results).toEqual([ { url: 'https://replicate.com/some-user/model-1' }, { url: 'https://replicate.com/some-user/model-2' } ]);
 
       // Add more tests for error handling, edge cases, etc.
     });
@@ -662,6 +662,28 @@ describe('Replicate client', () => {
     // Add more tests for error handling, edge cases, etc.
   });
 
+  describe('hardware.list', () => {
+    test('Calls the correct API route', async () => {
+      nock(BASE_URL)
+        .get('/hardware')
+        .reply(200, [
+          { name: "CPU", sku: "cpu" },
+          { name: "Nvidia T4 GPU", sku: "gpu-t4" },
+          { name: "Nvidia A40 GPU", sku: "gpu-a40-small" },
+          { name: "Nvidia A40 (Large) GPU", sku: "gpu-a40-large" },
+          { name: "Nvidia A40 (Large) GPU (8x)", sku: "gpu-a40-large-8x" },
+          { name: "Nvidia A100 (40GB) GPU", sku: "gpu-a100-small" },
+          { name: "Nvidia A100 (80GB) GPU", sku: "gpu-a100-large" },
+        ]);
+
+      const hardware = await client.hardware.list();
+      expect(hardware.length).toBe(7);
+      expect(hardware[ 0 ].name).toBe('CPU');
+      expect(hardware[ 0 ].sku).toBe('cpu');
+    });
+    // Add more tests for error handling, edge cases, etc.
+  });
+
   describe('run', () => {
     test('Calls the correct API routes', async () => {
       let firstPollingRequest = true;

lib/hardware.js

@@ -0,0 +1,16 @@
+/**
+ * List hardware
+ *
+ * @returns {Promise<object[]>} Resolves with the array of hardware
+ */
+async function listHardware() {
+  const response = await this.request('/hardware', {
+    method: 'GET',
+  });
+
+  return response.json();
+}
+
+module.exports = {
+  list: listHardware,
+};