Merge pull request #85 from oslabs-beta/stephany/jest

feat: Add jest test for RunManager

Author: yiqunzheng

mlflow/src/utils/interface.ts

@@ -1,30 +1,40 @@
 export interface Run {
-  info: {
-    run_id: string;
-    run_name: string;
-    experiment_id: string;
-    status: string;
-    start_time: number;
-    end_time: number;
-    artifact_uri: string;
-    lifecycle_stage: string;
-  };
-  data: {
-    metrics: Metrics[];
-    params: Params[];
-    tags: Tags[];
+  info: RunInfo;
+  data: RunData;
+  inputs?: RunInputs;
+}
+
+interface RunInfo {
+  run_id: string;
+  run_name: string;
+  experiment_id: string;
+  status: 'RUNNING' | 'SCHEDULED' | 'FINISHED' | 'FAILED' | 'KILLED';
+  start_time: number;
+  end_time: number;
+  artifact_uri: string;
+  lifecycle_stage: string;
+}
+
+interface RunData {
+  metrics: Metrics[];
+  params: Params[];
+  tags: Tags[];
+}
+
+interface RunInputs {
+  dataset_inputs?: DatasetInput[];
+}
+
+interface DatasetInput {
+  tags?: Tags[];
+  dataset: {
+    name: string;
+    digest: string;
+    source_type: string;
+    source: string;
+    schema?: string;
+    profile?: string;
   };
-  inputs: Array<{
-    tags?: Tags[];
-    dataset: {
-      name: string;
-      digest: string;
-      source_type: string;
-      source: string;
-      schema?: string;
-      profile?: string;
-    };
-  }>;
 }
 
 export interface Metrics {
@@ -47,3 +57,20 @@ export interface MetricHistoryResponse {
   metrics: Metrics[];
   next_page_token?: string;
 }
+
+export interface SearchedRuns {
+  runs: Run[];
+  next_page_token?: string;
+}
+
+export interface CleanupRuns {
+  deletedRuns: Run[];
+  total: number;
+  dryRun: boolean;
+}
+
+export interface CopyRun {
+  originalRunId: string;
+  newRunId: string;
+  targetExperimentId: string;
+}

mlflow/src/workflows/RunManager.ts

@@ -1,12 +1,8 @@
 import RunClient from '@tracking/RunClient';
 import ModelVersionClient from '@model-registry/ModelVersionClient';
-import { Run } from '@utils/interface';
+import { Run, SearchedRuns } from '@utils/interface';
 import { ApiError } from '@utils/apiError';
 
-interface keyable {
-  [key: string]: any;
-}
-
 class RunManager {
   private runClient: RunClient;
   private modelVersion: ModelVersionClient;
@@ -34,65 +30,61 @@ class RunManager {
   ): Promise<object> {
     const deletedRuns = [];
     const keepRunIds = new Set();
-    let pageToken = null;
+    const runViewType = undefined;
+    let pageToken = undefined;
     const maxResults = 1000;
 
     try {
       do {
         // get all runs
-        const searchResult: keyable = await this.runClient.searchRuns(
+        const searchResult = (await this.runClient.searchRuns(
           experimentIds,
           '',
-          undefined, // run_view_type
+          runViewType,
           maxResults,
           ['start_time DESC'],
           pageToken
-        );
+        )) as SearchedRuns;
 
         // get runs that match the keep crteria
-        const keepRunsResult: keyable = await this.runClient.searchRuns(
+        const keepRunsResult = (await this.runClient.searchRuns(
           experimentIds,
           query_string,
-          undefined, // run_view_type
+          runViewType,
           maxResults,
           ['start_time DESC'],
           pageToken
-        );
+        )) as SearchedRuns;
 
         // Add runs from keepRunsResult to keepResult
-        keepRunsResult.runs.forEach((run: Run) =>
+        keepRunsResult.runs?.forEach((run: Run) =>
           keepRunIds.add(run.info.run_id)
         );
 
-        // Add runs without the specified metric key to keepRunIds
+        // Process each run
         for (const run of searchResult.runs) {
-          if (Array.isArray(run.data.metrics)) {
-            const hasMetricKey = run.data.metrics.some(
-              (metric: { key: string }) => metric.key === metric_key
-            );
-            if (!hasMetricKey) {
-              keepRunIds.add(run.info.run_id);
-            }
-          } else {
-            // If run.data.metrics is not an array (e.g., undefined), keep the run
-            keepRunIds.add(run.info.run_id);
-          }
-        }
+          const metrics = run.data.metrics;
+          const hasMetricKey = Array.isArray(metrics)
+            ? metrics.some((metric) => metric.key === metric_key)
+            : metric_key in metrics;
 
-        // Delete runs that are not in keepRunIds
-        for (const run of searchResult.runs) {
-          if (!keepRunIds.has(run.info.run_id)) {
+          if (!hasMetricKey || keepRunIds.has(run.info.run_id)) {
+            keepRunIds.add(run.info.run_id);
+          } else {
+            deletedRuns.push(run);
             if (!dryRun) {
               await this.runClient.deleteRun(run.info.run_id);
             }
-            deletedRuns.push(run);
           }
         }
-
-        pageToken = searchResult.page_token;
+        pageToken = searchResult.next_page_token;
       } while (pageToken);
 
-      return { deletedRuns, total: deletedRuns.length, dryRun };
+      return {
+        deletedRuns: deletedRuns,
+        total: deletedRuns.length,
+        dryRun,
+      };
     } catch (error) {
       if (error instanceof ApiError) {
         console.error(`API Error (${error.statusCode}): ${error.message}`);
@@ -119,14 +111,14 @@ class RunManager {
   ): Promise<object> {
     try {
       // get original run
-      const originalRun: keyable = await this.runClient.getRun(runId);
+      const originalRun = (await this.runClient.getRun(runId)) as Run;
 
       // create a new run in the target experiment
-      const newRun: keyable = await this.runClient.createRun(
+      const newRun = (await this.runClient.createRun(
         targetExperimentId,
         undefined,
         originalRun.info.start_time
-      );
+      )) as Run;
 
       const newRunId = newRun.info.run_id;
 
@@ -138,6 +130,7 @@ class RunManager {
         originalRun.info.status,
         endTime
       );
+
       if (originalRun.info.lifecycle_stage !== 'active') {
         await this.runClient.setTag(
           newRunId,

mlflow/tests/RunClient.test.ts

@@ -128,8 +128,12 @@ describe('RunClient', () => {
 
   // POST - Restore a deleted run
   describe('restoreRun', () => {
+    let run: Run;
+    beforeEach(async () => {
+      run = (await runClient.createRun(experimentId)) as Run;
+    });
+
     test('- Should restore a deleted run with run_id', async () => {
-      const run = (await runClient.createRun(experimentId)) as Run;
       await runClient.deleteRun(run.info.run_id);
 
       // check if the run is marked as deleted
@@ -147,8 +151,6 @@ describe('RunClient', () => {
     });
 
     test('- Should not throw error when trying to restore a non-deleted run', async () => {
-      const run = (await runClient.createRun(experimentId)) as Run;
-
       // Attempt to restore a non-deleted run
       await expect(
         runClient.restoreRun(run.info.run_id)
@@ -403,17 +405,19 @@ describe('RunClient', () => {
 
   // POST - Log a batch of metrics, params, and tags for a run
   describe('logBatch', () => {
-    test('- Should not throw error with just run_id', async () => {
-      const run = (await runClient.createRun(experimentId)) as Run;
+    let run: Run;
+
+    beforeEach(async () => {
+      run = (await runClient.createRun(experimentId)) as Run;
+    });
 
+    test('- Should not throw error with just run_id', async () => {
       await expect(
         runClient.logBatch(run.info.run_id)
       ).resolves.toBeUndefined();
     });
 
     test('- Should log batch with optional metrics', async () => {
-      const run = (await runClient.createRun(experimentId)) as Run;
-
       const metrics: Metrics[] = [
         { key: 'accuracy', value: 0.83, timestamp: 1694000700000 },
         { key: 'loss', value: 0.18, timestamp: 1694000700000 },
@@ -448,8 +452,6 @@ describe('RunClient', () => {
     });
 
     test('- Should log batch with optional params', async () => {
-      const run = (await runClient.createRun(experimentId)) as Run;
-
       const params: Params[] = [
         { key: 'learning_rate', value: '0.0001' },
         { key: 'batch_size', value: '256' },
@@ -464,8 +466,6 @@ describe('RunClient', () => {
     });
 
     test('- Should log batch with optional tags', async () => {
-      const run = (await runClient.createRun(experimentId)) as Run;
-
       const tags: Tags[] = [
         { key: 'model_type', value: 'GradientBoosting' },
         { key: 'data_version', value: 'v1.7' },
@@ -480,8 +480,6 @@ describe('RunClient', () => {
     });
 
     test('- Should be able to log up to 1000 metrics', async () => {
-      const run = (await runClient.createRun(experimentId)) as Run;
-
       const metrics = Array.from({ length: 1000 }, (_, index) => ({
         key: `metric${index}`,
         value: index,
@@ -495,8 +493,6 @@ describe('RunClient', () => {
     });
 
     test('- Should throw error when exceeding 1000 metrics', async () => {
-      const run = (await runClient.createRun(experimentId)) as Run;
-
       const metrics = Array.from({ length: 1001 }, (_, index) => ({
         key: `metric${index}`,
         value: index,
@@ -510,8 +506,6 @@ describe('RunClient', () => {
     });
 
     test('- Should be able to log up to 100 params', async () => {
-      const run = (await runClient.createRun(experimentId)) as Run;
-
       const params = Array.from({ length: 100 }, (_, index) => ({
         key: `param${index}`,
         value: `value${index}`,
@@ -523,8 +517,6 @@ describe('RunClient', () => {
     });
 
     test('- Should throw error when exceeding 100 params', async () => {
-      const run = (await runClient.createRun(experimentId)) as Run;
-
       const params = Array.from({ length: 101 }, (_, index) => ({
         key: `param${index}`,
         value: `value${index}`,
@@ -536,8 +528,6 @@ describe('RunClient', () => {
     });
 
     test('- Should be able to log up to 100 tags', async () => {
-      const run = (await runClient.createRun(experimentId)) as Run;
-
       const tags = Array.from({ length: 100 }, (_, index) => ({
         key: `tag${index}`,
         value: `value${index}`,
@@ -549,8 +539,6 @@ describe('RunClient', () => {
     });
 
     test('- Should throw error when exceeding 100 tags', async () => {
-      const run = (await runClient.createRun(experimentId)) as Run;
-
       const tags = Array.from({ length: 101 }, (_, index) => ({
         key: `tag${index}`,
         value: `value${index}`,