Fix dispatcher CUDA_VISIBLE_DEVICES envvar for windows (microsoft#1604)

* Fix dispatcher CUDA_VISIBLE_DEVICES for windows
SparkSnail · Oct 12, 2019 · 5bd994d · 5bd994d
1 parent 1bd2457
commit 5bd994d
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 19 deletions.
diff --git a/src/nni_manager/common/manager.ts b/src/nni_manager/common/manager.ts
@@ -49,8 +49,8 @@ interface ExperimentParams {
         classArgs?: any;
         classFileName?: string;
         checkpointDir: string;
-        gpuNum?: number;
         includeIntermediateResults?: boolean;
+        gpuIndices?: string;
     };
     assessor?: {
         className: string;
@@ -59,7 +59,6 @@ interface ExperimentParams {
         classArgs?: any;
         classFileName?: string;
         checkpointDir: string;
-        gpuNum?: number;
     };
     advisor?: {
         className: string;
@@ -68,7 +67,7 @@ interface ExperimentParams {
         classArgs?: any;
         classFileName?: string;
         checkpointDir: string;
-        gpuNum?: number;
+        gpuIndices?: string;
     };
     clusterMetaData?: {
         key: string;

diff --git a/src/nni_manager/common/utils.ts b/src/nni_manager/common/utils.ts
@@ -219,11 +219,6 @@ function getMsgDispatcherCommand(tuner: any, assessor: any, advisor: any, multiP
         if (advisor.classFileName !== undefined && advisor.classFileName.length > 1) {
             command += ` --advisor_class_filename ${advisor.classFileName}`;
         }
-        if (advisor.gpuIndices !== undefined) {
-            command = `CUDA_VISIBLE_DEVICES=${advisor.gpuIndices} ` + command;
-        } else {
-            command = `CUDA_VISIBLE_DEVICES='' ` + command;
-        }
     } else {
         command += ` --tuner_class_name ${tuner.className}`;
         if (tuner.classArgs !== undefined) {
@@ -248,12 +243,6 @@ function getMsgDispatcherCommand(tuner: any, assessor: any, advisor: any, multiP
                 command += ` --assessor_class_filename ${assessor.classFileName}`;
             }
         }
-
-        if (tuner.gpuIndices !== undefined) {
-            command = `CUDA_VISIBLE_DEVICES=${tuner.gpuIndices} ` + command;
-        } else {
-            command = `CUDA_VISIBLE_DEVICES='' ` + command;
-        }
     }
 
     return command;

diff --git a/src/nni_manager/core/nnimanager.ts b/src/nni_manager/core/nnimanager.ts
@@ -369,7 +369,8 @@ class NNIManager implements Manager {
             NNI_CHECKPOINT_DIRECTORY: dataDirectory,
             NNI_LOG_DIRECTORY: getLogDir(),
             NNI_LOG_LEVEL: getLogLevel(),
-            NNI_INCLUDE_INTERMEDIATE_RESULTS: includeIntermediateResultsEnv
+            NNI_INCLUDE_INTERMEDIATE_RESULTS: includeIntermediateResultsEnv,
+            CUDA_VISIBLE_DEVICES: this.getGpuEnvvarValue()
         };
         let newEnv = Object.assign({}, process.env, nniEnv);
         const tunerProc: ChildProcess = getTunerProc(command,stdio,newCwd,newEnv);
@@ -379,6 +380,22 @@ class NNIManager implements Manager {
         return;
     }
 
+    private getGpuEnvvarValue(): string {
+        let cudaDevices: string | undefined;
+
+        if (this.experimentProfile.params.advisor !== undefined) {
+            cudaDevices = this.experimentProfile.params.advisor.gpuIndices;
+        } else if (this.experimentProfile.params.tuner !== undefined) {
+            cudaDevices = this.experimentProfile.params.tuner.gpuIndices;
+        }
+
+        if (cudaDevices === undefined) {
+            return '';
+        } else {
+            return cudaDevices;
+        }
+    }
+
     private updateTrialConcurrency(trialConcurrency: number): void {
         // we assume trialConcurrency >= 0, which is checked by restserver
         this.trialConcurrencyChange += (trialConcurrency - this.experimentProfile.params.trialConcurrency);

diff --git a/src/nni_manager/core/test/dataStore.test.ts b/src/nni_manager/core/test/dataStore.test.ts
@@ -72,8 +72,7 @@ describe('Unit test for dataStore', () => {
                 }`,
                 tuner: {
                     className: 'testTuner',
-                    checkpointDir: '/tmp/cp',
-                    gpuNum: 0
+                    checkpointDir: '/tmp/cp'
                 }
             },
             id: 'exp123',

diff --git a/src/nni_manager/core/test/sqlDatabase.test.ts b/src/nni_manager/core/test/sqlDatabase.test.ts
@@ -40,8 +40,7 @@ const expParams1: ExperimentParams = {
     searchSpace: 'SS',
     tuner: {
         className: 'testTuner',
-        checkpointDir: '/tmp',
-        gpuNum: 0
+        checkpointDir: '/tmp'
     }
 };