Fix datastore performance issue

src/nni_manager/common/datastore.ts

@@ -66,7 +66,7 @@ interface TrialJobInfo {
     endTime?: number;
     hyperParameters?: string[];
     logPath?: string;
-    finalMetricData?: string;
+    finalMetricData?: MetricDataRecord;
     stderrPath?: string;
 }
 

src/nni_manager/core/nniDataStore.ts

@@ -156,21 +156,23 @@ class NNIDataStore implements DataStore {
     }
 
     private async queryTrialJobs(status?: TrialJobStatus, trialJobId?: string): Promise<TrialJobInfo[]> {
-        const result: TrialJobInfo[]= [];
+        const result: TrialJobInfo[] = [];
         const trialJobEvents: TrialJobEventRecord[] = await this.db.queryTrialJobEvent(trialJobId);
         if (trialJobEvents === undefined) {
             return result;
         }
         const map: Map<string, TrialJobInfo> = this.getTrialJobsByReplayEvents(trialJobEvents);
 
-        for (let key of map.keys()) {
-            const jobInfo = map.get(key);
+        const finalMetricsMap: Map<string, MetricDataRecord> = await this.getFinalMetricData(trialJobId);
+
+        for (const key of map.keys()) {
+            const jobInfo: TrialJobInfo | undefined = map.get(key);
             if (jobInfo === undefined) {
                 continue;
             }
             if (!(status !== undefined && jobInfo.status !== status)) {
                 if (jobInfo.status === 'SUCCEEDED') {
-                    jobInfo.finalMetricData = await this.getFinalMetricData(jobInfo.id);
+                    jobInfo.finalMetricData = finalMetricsMap.get(jobInfo.id);
                 }
                 result.push(jobInfo);
             }
@@ -179,16 +181,20 @@ class NNIDataStore implements DataStore {
         return result;
     }
 
-    private async getFinalMetricData(trialJobId: string): Promise<any> {
+    private async getFinalMetricData(trialJobId?: string): Promise<Map<string, MetricDataRecord>> {
+        const map: Map<string, MetricDataRecord> = new Map();
         const metrics: MetricDataRecord[] = await this.getMetricData(trialJobId, 'FINAL');
 
         const multiPhase: boolean = await this.isMultiPhase();
 
-        if (metrics.length > 1 && !multiPhase) {
-            this.log.error(`Found multiple FINAL results for trial job ${trialJobId}`);
+        for (const metric of metrics) {
+            if (map.has(metric.trialJobId) && !multiPhase) {
+                this.log.error(`Found multiple FINAL results for trial job ${trialJobId}`);
+            }
+            map.set(metric.trialJobId, metric);
         }
 
-        return metrics[metrics.length - 1];
+        return map;
     }
 
     private async isMultiPhase(): Promise<boolean> {

src/nni_manager/rest_server/test/mockedNNIManager.ts

@@ -158,14 +158,28 @@ export class MockedNNIManager extends Manager {
             status: 'SUCCEEDED',
             startTime: Date.now(),
             endTime: Date.now(),
-            finalMetricData: 'lr: 0.01, val accuracy: 0.89, batch size: 256'
+            finalMetricData: {
+                timestamp: 0,
+                trialJobId: '3456',
+                parameterId: '123',
+                type: 'FINAL',
+                sequence: 0,
+                data: '0.2'
+            }
         };
         const job2: TrialJobInfo = {
             id: '3456',
             status: 'FAILED',
             startTime: Date.now(),
             endTime: Date.now(),
-            finalMetricData: ''
+            finalMetricData: {
+                timestamp: 0,
+                trialJobId: '3456',
+                parameterId: '123',
+                type: 'FINAL',
+                sequence: 0,
+                data: '0.2'
+            }
         };
 
         return Promise.resolve([job1, job2]);