execute removeGpuReservation first to release GPUs properly (#3941)

microsoft · Jul 16, 2021 · bccfa9f · bccfa9f
1 parent 30df69b
commit bccfa9f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/ts/nni_manager/training_service/reusable/trialDispatcher.ts b/ts/nni_manager/training_service/reusable/trialDispatcher.ts
@@ -732,6 +732,9 @@ class TrialDispatcher implements TrainingService {
      * @param trial 
      */
     private releaseEnvironment(trial: TrialDetail): void {
+        if (true === this.enableGpuScheduler) {
+            this.gpuScheduler.removeGpuReservation(trial);
+        }
         if (trial.environment !== undefined) {
             if (trial.environment.runningTrialCount <= 0) {
                 throw new Error(`TrialDispatcher: environment ${trial.environment.id} has no counted running trial!`);
@@ -740,9 +743,6 @@ class TrialDispatcher implements TrainingService {
             trial.environment.latestTrialReleasedTime = Date.now();
             trial.environment = undefined;
         }
-        if (true === this.enableGpuScheduler) {
-            this.gpuScheduler.removeGpuReservation(trial);
-        }
     }
 
     private async handleMetricData(trialId: string, data: any): Promise<void> {