add more init options

wandb · Oct 21, 2024 · fe13989 · fe13989
1 parent db06b3d
commit fe13989
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 9 deletions.
diff --git a/sdks/node/src/__tests__/clientApi.test.ts b/sdks/node/src/__tests__/clientApi.test.ts
@@ -17,7 +17,7 @@ describe('Client API', () => {
       defaultEntityName: jest.fn().mockResolvedValue('test-entity'),
     }));
 
-    const client = await init('test-project');
+    const client = await init({ project: 'test-project' });
     const gottenClient = requireGlobalClient();
 
     expect(gottenClient).toBeDefined();

diff --git a/sdks/node/src/__tests__/dataset.test.ts b/sdks/node/src/__tests__/dataset.test.ts
@@ -3,7 +3,7 @@ import { Dataset } from '../dataset';
 
 describe('Dataset', () => {
   test('should save a dataset', async () => {
-    const client = await init('test-project');
+    const client = await init({ project: 'test-project' });
     const data = [
       { id: 1, value: 2 },
       { id: 2, value: 3 },

diff --git a/sdks/node/src/__tests__/table.test.ts b/sdks/node/src/__tests__/table.test.ts
@@ -19,7 +19,7 @@ describe('table', () => {
     }
 
     // Saving the table generates refs for the table and its rows
-    const client = await init('test-project');
+    const client = await init({ project: 'test-project' });
 
     (client as any).saveTable(table); // TODO: Saving a Table is not public... but maybe it should be?
     const ref = await table.__savedRef;

diff --git a/sdks/node/src/__tests__/weaveObject.test.ts b/sdks/node/src/__tests__/weaveObject.test.ts
@@ -28,7 +28,7 @@ describe('weaveObject', () => {
   });
 
   test('class-example', async () => {
-    const client = await init('test-project');
+    const client = await init({ project: 'test-project' });
     const obj = new ExampleObject('test', 1);
 
     // save an object

diff --git a/sdks/node/src/clientApi.ts b/sdks/node/src/clientApi.ts
@@ -5,13 +5,23 @@ import { getApiKey } from './wandb/settings';
 import { WandbServerApi } from './wandb/wandbServerApi';
 import { CallStackEntry, WeaveClient } from './weaveClient';
 
+export interface InitOptions {
+  project: string;
+  entity?: string;
+  projectName?: string;
+  host?: string;
+  apiKey?: string;
+}
+
 // Global client instance
 export let globalClient: WeaveClient | null = null;
 
-export async function init(projectName: string): Promise<WeaveClient> {
-  const host = 'https://api.wandb.ai';
-  const apiKey = getApiKey();
-
+export async function init({
+  project,
+  entity,
+  host = 'https://api.wandb.ai',
+  apiKey = getApiKey(),
+}: InitOptions): Promise<WeaveClient> {
   const headers: Record<string, string> = {
     'User-Agent': `W&B Internal JS Client ${process.env.VERSION || 'unknown'}`,
     Authorization: `Basic ${Buffer.from(`api:${apiKey}`).toString('base64')}`,
@@ -20,7 +30,8 @@ export async function init(projectName: string): Promise<WeaveClient> {
   try {
     const wandbServerApi = new WandbServerApi(host, apiKey);
     const defaultEntityName = await wandbServerApi.defaultEntityName();
-    const projectId = `${defaultEntityName}/${projectName}`;
+    const entityName = entity ?? defaultEntityName;
+    const projectId = `${entityName}/${project}`;
 
     const retryFetch = createFetchWithRetry({
       baseDelay: 1000,