Local TrainingService UT (#772)

src/nni_manager/training_service/local/localTrainingService.ts

@@ -255,7 +255,7 @@ class LocalTrainingService implements TrainingService {
         }
         if (trialJob.pid === undefined){
             this.setTrialJobStatus(trialJob, 'USER_CANCELED');
-            return;
+            return Promise.resolve();
         }
         if (trialJob.form.jobType === 'TRIAL') {
             await tkill(trialJob.pid, 'SIGKILL');
@@ -265,6 +265,7 @@ class LocalTrainingService implements TrainingService {
             throw new Error(`Job type not supported: ${trialJob.form.jobType}`);
         }
         this.setTrialJobStatus(trialJob, getJobCancelStatus(isEarlyStopped));
+        return Promise.resolve();
     }
 
     public async setClusterMetadata(key: string, value: string): Promise<void> {

src/nni_manager/training_service/test/localTrainingService.test.ts

@@ -19,14 +19,106 @@
 
 'use strict';
 
-import { TrainingService } from '../../common/trainingService';
-import { LocalTrainingService } from '../local/localTrainingService';
+import * as assert from 'assert';
+import * as chai from 'chai';
+import * as chaiAsPromised from 'chai-as-promised';
+import * as fs from 'fs';
+import * as tmp from 'tmp';
 import * as component from '../../common/component';
+import { TrialJobApplicationForm, TrialJobDetail, TrainingService } from '../../common/trainingService';
+import { cleanupUnitTest, delay, prepareUnitTest } from '../../common/utils';
+import { TrialConfigMetadataKey } from '../common/trialConfigMetadataKey';
+import { LocalTrainingServiceForGPU } from '../local/localTrainingServiceForGPU';
+
+// TODO: copy mockedTrail.py to local folder
+const localCodeDir: string = tmp.dirSync().name
+const mockedTrialPath: string = './training_service/test/mockedTrial.py'
+fs.copyFileSync(mockedTrialPath, localCodeDir + '/mockedTrial.py')
 
 describe('Unit Test for LocalTrainingService', () => {
-    let trainingService: TrainingService
+    let trialConfig: any = `{"command":"sleep 1h && echo hello","codeDir":"${localCodeDir}","gpuNum":1}`
+
+    let localTrainingService: LocalTrainingServiceForGPU;
+
+    before(() => {
+        chai.should();
+        chai.use(chaiAsPromised);
+        prepareUnitTest();
+    });
+
+    after(() => {
+        cleanupUnitTest();
+    });
+
+    beforeEach(() => {
+        localTrainingService = component.get(LocalTrainingServiceForGPU);
+        localTrainingService.run();
+    });
+
+    afterEach(() => {
+        localTrainingService.cleanUp();
+    });
+
+    it('List empty trial jobs', async () => {
+        //trial jobs should be empty, since there are no submitted jobs
+        chai.expect(await localTrainingService.listTrialJobs()).to.be.empty;
+    });
+    
+    it('setClusterMetadata and getClusterMetadata', async () => {
+        await localTrainingService.setClusterMetadata(TrialConfigMetadataKey.TRIAL_CONFIG, trialConfig);
+        localTrainingService.getClusterMetadata(TrialConfigMetadataKey.TRIAL_CONFIG).then((data)=>{
+            chai.expect(data).to.be.equals(trialConfig);
+        });
+    });
+
+    it('Submit job and Cancel job', async () => {
+        await localTrainingService.setClusterMetadata(TrialConfigMetadataKey.TRIAL_CONFIG, trialConfig);
+
+        // submit job
+        const form: TrialJobApplicationForm = {
+            jobType: 'TRIAL',
+            hyperParameters: {
+                value: 'mock hyperparameters',
+                index: 0
+            }
+        };
+        const jobDetail: TrialJobDetail = await localTrainingService.submitTrialJob(form);
+        chai.expect(jobDetail.status).to.be.equals('WAITING');
+        await localTrainingService.cancelTrialJob(jobDetail.id);
+        chai.expect(jobDetail.status).to.be.equals('USER_CANCELED');
+    }).timeout(20000);
+    
+    it('Read metrics, Add listener, and remove listener', async () => {
+        // set meta data
+        const trialConfig: string = `{\"command\":\"python3 mockedTrial.py\", \"codeDir\":\"${localCodeDir}\",\"gpuNum\":0}`
+        await localTrainingService.setClusterMetadata(TrialConfigMetadataKey.TRIAL_CONFIG, trialConfig);
+
+        // submit job
+        const form: TrialJobApplicationForm = {
+            jobType: 'TRIAL',
+            hyperParameters: {
+                value: 'mock hyperparameters',
+                index: 0
+            }
+        };
+        const jobDetail: TrialJobDetail = await localTrainingService.submitTrialJob(form);
+        chai.expect(jobDetail.status).to.be.equals('WAITING');
+        localTrainingService.listTrialJobs().then((jobList)=>{
+            chai.expect(jobList.length).to.be.equals(1);
+        });
+        // Add metrics listeners
+        const listener1 = function f1(metric: any) {
+            chai.expect(metric.id).to.be.equals(jobDetail.id);
+        }
+        localTrainingService.addTrialJobMetricListener(listener1);
+        // Wait to collect metric
+        await delay(1000);
+
+        await localTrainingService.cancelTrialJob(jobDetail.id);
+        localTrainingService.removeTrialJobMetricListener(listener1);
+    }).timeout(20000);
 
-    beforeEach(async () => {
-        trainingService = component.get(LocalTrainingService);
+    it('Test multiphaseSupported', () => {
+        chai.expect(localTrainingService.isMultiPhaseJobSupported).to.be.equals(true)
     })
 });
\ No newline at end of file