Fix local gpu issue (#1000)

1d9b0a99 · chicm-ms · SparkSnail · b7ca02c7 · 1d9b0a99
Commit 1d9b0a99 authored Apr 19, 2019 by chicm-ms Committed by SparkSnail Apr 19, 2019
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 7 deletions

src/nni_manager/training_service/local/localTrainingService.ts ...ni_manager/training_service/local/localTrainingService.ts +5 -7

No files found.
--- a/src/nni_manager/training_service/local/localTrainingService.ts
+++ b/src/nni_manager/training_service/local/localTrainingService.ts
@@ -369,7 +369,7 @@ class LocalTrainingService implements TrainingService {
    private getEnvironmentVariables(
        trialJobDetail: TrialJobDetail,
-        resource?: { gpuIndices: number[] }): { key: string; value: string }[] {
+        resource: { gpuIndices: number[] }): { key: string; value: string }[] {
        const envVariables: { key: string; value: string }[] = [
            { key: 'NNI_PLATFORM', value: 'local' },
            { key: 'NNI_SYS_DIR', value: trialJobDetail.workingDirectory },
@@ -379,12 +379,10 @@ class LocalTrainingService implements TrainingService {
            { key: 'MULTI_PHASE', value: this.isMultiPhase.toString() }
        ];
-        if (resource !== undefined && resource.gpuIndices.length > 0) {
+        envVariables.push({
-            envVariables.push({
+            key: 'CUDA_VISIBLE_DEVICES',
-                    key: 'CUDA_VISIBLE_DEVICES',
+            value: this.gpuScheduler === undefined ? '' : resource.gpuIndices.join(',')
-                    value: this.gpuScheduler === undefined ? '' : resource.gpuIndices.join(',')
+        });
-            });
-        }
        return envVariables;
    }