gpuScheduler.ts 10.5 KB
Newer Older
liuzhe-lz's avatar
liuzhe-lz committed
1
2
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
Deshui Yu's avatar
Deshui Yu committed
3

4
5
6
7
import assert from 'assert';
import { getLogger, Logger } from 'common/log';
import { randomSelect } from 'common/utils';
import { RemoteMachineConfig } from 'common/experimentConfig';
8
import { GPUInfo, ScheduleResultType } from '../common/gpuData';
9
import { ExecutorManager, RemoteMachineMeta, RemoteMachineScheduleResult, RemoteMachineTrialJobDetail } from './remoteMachineData';
Deshui Yu's avatar
Deshui Yu committed
10

chicm-ms's avatar
chicm-ms committed
11
12
type SCHEDULE_POLICY_NAME = 'random' | 'round-robin';

Deshui Yu's avatar
Deshui Yu committed
13
14
15
16
17
/**
 * A simple GPU scheduler implementation
 */
export class GPUScheduler {

18
    private readonly machineExecutorMap: Map<RemoteMachineConfig, ExecutorManager>;
liuzhe-lz's avatar
liuzhe-lz committed
19
    private readonly log: Logger = getLogger('GPUScheduler');
chicm-ms's avatar
chicm-ms committed
20
21
22
    private readonly policyName: SCHEDULE_POLICY_NAME = 'round-robin';
    private roundRobinIndex: number = 0;
    private configuredRMs: RemoteMachineMeta[] = [];
Deshui Yu's avatar
Deshui Yu committed
23
24
25

    /**
     * Constructor
26
     * @param machineExecutorMap map from remote machine to executor
Deshui Yu's avatar
Deshui Yu committed
27
     */
28
    constructor(machineExecutorMap: Map<RemoteMachineConfig, ExecutorManager>) {
29
30
        assert(machineExecutorMap.size > 0);
        this.machineExecutorMap = machineExecutorMap;
31
        this.configuredRMs = Array.from(machineExecutorMap.values(), manager => manager.rmMeta);
Deshui Yu's avatar
Deshui Yu committed
32
33
34
35
36
37
    }

    /**
     * Schedule a machine according to the constraints (requiredGPUNum)
     * @param requiredGPUNum required GPU number
     */
chicm-ms's avatar
chicm-ms committed
38
    public scheduleMachine(requiredGPUNum: number | undefined, trialJobDetail: RemoteMachineTrialJobDetail): RemoteMachineScheduleResult {
39
        if (requiredGPUNum === undefined) {
SparkSnail's avatar
SparkSnail committed
40
41
            requiredGPUNum = 0;
        }
42
        assert(requiredGPUNum >= 0);
43
        const allRMs: RemoteMachineMeta[] = Array.from(this.machineExecutorMap.values(), manager => manager.rmMeta);
44
45
46
        assert(allRMs.length > 0);

        // Step 1: Check if required GPU number not exceeds the total GPU number in all machines
chicm-ms's avatar
chicm-ms committed
47
        const eligibleRM: RemoteMachineMeta[] = allRMs.filter((rmMeta: RemoteMachineMeta) =>
48
            rmMeta.gpuSummary === undefined || requiredGPUNum === 0 || (requiredGPUNum !== undefined && rmMeta.gpuSummary.gpuCount >= requiredGPUNum));
49
        if (eligibleRM.length === 0) {
Deshui Yu's avatar
Deshui Yu committed
50
51
52
            // If the required gpu number exceeds the upper limit of all machine's GPU number
            // Return REQUIRE_EXCEED_TOTAL directly
            return ({
53
54
                resultType: ScheduleResultType.REQUIRE_EXCEED_TOTAL,
                scheduleInfo: undefined
Deshui Yu's avatar
Deshui Yu committed
55
56
57
            });
        }

58
59
60
61
        // Step 2: Allocate Host/GPU for specified trial job
        // Currenty the requireGPUNum parameter for all trial jobs are identical.
        if (requiredGPUNum > 0) {
            // Trial job requires GPU
62
            const result: RemoteMachineScheduleResult | undefined = this.scheduleGPUHost(requiredGPUNum, trialJobDetail);
63
64
            if (result !== undefined) {
                return result;
Deshui Yu's avatar
Deshui Yu committed
65
            }
66
67
68
69
        } else {
            // Trail job does not need GPU
            const allocatedRm: RemoteMachineMeta = this.selectMachine(allRMs);

70
            return this.allocateHost(requiredGPUNum, allocatedRm, [], trialJobDetail);
71
        }
72
        this.log.warning(`Scheduler: trialJob id ${trialJobDetail.id}, no machine can be scheduled, return TMP_NO_AVAILABLE_GPU `);
73

Deshui Yu's avatar
Deshui Yu committed
74
        return {
75
76
            resultType: ScheduleResultType.TMP_NO_AVAILABLE_GPU,
            scheduleInfo: undefined
Deshui Yu's avatar
Deshui Yu committed
77
78
79
        };
    }

80
81
82
    /**
     * remove the job's gpu reversion
     */
83
    public removeGpuReservation(trialJobId: string, trialJobMap: Map<string, RemoteMachineTrialJobDetail>): void {
84
85
        const trialJobDetail: RemoteMachineTrialJobDetail | undefined = trialJobMap.get(trialJobId);
        if (trialJobDetail === undefined) {
86
            throw new Error(`could not get trialJobDetail by id ${trialJobId}`);
87
88
89
90
        }
        if (trialJobDetail.rmMeta !== undefined &&
            trialJobDetail.rmMeta.occupiedGpuIndexMap !== undefined &&
            trialJobDetail.gpuIndices !== undefined &&
91
92
            trialJobDetail.gpuIndices.length > 0) {
            for (const gpuInfo of trialJobDetail.gpuIndices) {
93
94
95
                const num: number | undefined = trialJobDetail.rmMeta.occupiedGpuIndexMap.get(gpuInfo.index);
                if (num !== undefined) {
                    if (num === 1) {
96
97
                        trialJobDetail.rmMeta.occupiedGpuIndexMap.delete(gpuInfo.index);
                    } else {
98
                        trialJobDetail.rmMeta.occupiedGpuIndexMap.set(gpuInfo.index, num - 1);
99
                    }
100
                }
101
            }
102
        }
103
104
        trialJobDetail.gpuIndices = [];
        trialJobMap.set(trialJobId, trialJobDetail);
105
106
    }

107
    private scheduleGPUHost(requiredGPUNum: number, trialJobDetail: RemoteMachineTrialJobDetail): RemoteMachineScheduleResult | undefined {
108
109
110
111
112
113
114
115
116
117
118
        const totalResourceMap: Map<RemoteMachineMeta, GPUInfo[]> = this.gpuResourceDetection();
        const qualifiedRMs: RemoteMachineMeta[] = [];
        totalResourceMap.forEach((gpuInfos: GPUInfo[], rmMeta: RemoteMachineMeta) => {
            if (gpuInfos !== undefined && gpuInfos.length >= requiredGPUNum) {
                qualifiedRMs.push(rmMeta);
            }
        });
        if (qualifiedRMs.length > 0) {
            const allocatedRm: RemoteMachineMeta = this.selectMachine(qualifiedRMs);
            const gpuInfos: GPUInfo[] | undefined = totalResourceMap.get(allocatedRm);
            if (gpuInfos !== undefined) { // should always true
119
                return this.allocateHost(requiredGPUNum, allocatedRm, gpuInfos, trialJobDetail);
120
121
122
123
            } else {
                assert(false, 'gpuInfos is undefined');
            }
        }
124
        return undefined;
125
126
    }

Deshui Yu's avatar
Deshui Yu committed
127
128
129
130
131
132
133
    /**
     * Detect available GPU resource for a remote machine
     * @param rmMeta Remote machine metadata
     * @param requiredGPUNum required GPU number by application
     * @param availableGPUMap available GPU resource filled by this detection
     * @returns Available GPU number on this remote machine
     */
chicm-ms's avatar
chicm-ms committed
134
135
    private gpuResourceDetection(): Map<RemoteMachineMeta, GPUInfo[]> {
        const totalResourceMap: Map<RemoteMachineMeta, GPUInfo[]> = new Map<RemoteMachineMeta, GPUInfo[]>();
136
137
        this.machineExecutorMap.forEach((executorManager: ExecutorManager, machineConfig: RemoteMachineConfig) => {
            const rmMeta = executorManager.rmMeta;
Deshui Yu's avatar
Deshui Yu committed
138
            // Assgin totoal GPU count as init available GPU number
139
140
            if (rmMeta.gpuSummary !== undefined) {
                const availableGPUs: GPUInfo[] = [];
141
                const designatedGpuIndices: number[] | undefined = machineConfig.gpuIndices;
142
143
144
145
146
147
148
                if (designatedGpuIndices !== undefined) {
                    for (const gpuIndex of designatedGpuIndices) {
                        if (gpuIndex >= rmMeta.gpuSummary.gpuCount) {
                            throw new Error(`Specified GPU index not found: ${gpuIndex}`);
                        }
                    }
                }
149
                this.log.debug(`designated gpu indices: ${designatedGpuIndices}`);
Deshui Yu's avatar
Deshui Yu committed
150
                rmMeta.gpuSummary.gpuInfos.forEach((gpuInfo: GPUInfo) => {
151
                    // if the GPU has active process, OR be reserved by a job,
152
                    // or index not in gpuIndices configuration in machineList,
153
                    // or trial number on a GPU reach max number,
Deshui Yu's avatar
Deshui Yu committed
154
                    // We should NOT allocate this GPU
155
                    // if users set useActiveGpu, use the gpu whether there is another activeProcess
156
                    if (designatedGpuIndices === undefined || designatedGpuIndices.includes(gpuInfo.index)) {
157
158
                        if (rmMeta.occupiedGpuIndexMap !== undefined) {
                            const num: number | undefined = rmMeta.occupiedGpuIndexMap.get(gpuInfo.index);
159
160
                            if ((num === undefined && (!machineConfig.useActiveGpu && gpuInfo.activeProcessNum === 0 || machineConfig.useActiveGpu)) ||
                                (num !== undefined && num < machineConfig.maxTrialNumberPerGpu)) {
161
162
163
164
165
                                availableGPUs.push(gpuInfo);
                            }
                        } else {
                            throw new Error(`occupiedGpuIndexMap initialize error!`);
                        }
Deshui Yu's avatar
Deshui Yu committed
166
167
168
169
170
171
172
173
                    }
                });
                totalResourceMap.set(rmMeta, availableGPUs);
            }
        });

        return totalResourceMap;
    }
174
175
176
177

    private selectMachine(rmMetas: RemoteMachineMeta[]): RemoteMachineMeta {
        assert(rmMetas !== undefined && rmMetas.length > 0);

chicm-ms's avatar
chicm-ms committed
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
        if (this.policyName === 'random') {
            return randomSelect(rmMetas);
        } else if (this.policyName === 'round-robin') {
            return this.roundRobinSelect(rmMetas);
        } else {
            throw new Error(`Unsupported schedule policy: ${this.policyName}`);
        }
    }

    private roundRobinSelect(rmMetas: RemoteMachineMeta[]): RemoteMachineMeta {
        while (!rmMetas.includes(this.configuredRMs[this.roundRobinIndex % this.configuredRMs.length])) {
            this.roundRobinIndex++;
        }

        return this.configuredRMs[this.roundRobinIndex++ % this.configuredRMs.length];
193
194
195
196
197
198
199
200
    }

    private selectGPUsForTrial(gpuInfos: GPUInfo[], requiredGPUNum: number): GPUInfo[] {
        // Sequentially allocate GPUs
        return gpuInfos.slice(0, requiredGPUNum);
    }

    private allocateHost(requiredGPUNum: number, rmMeta: RemoteMachineMeta,
201
        gpuInfos: GPUInfo[], trialJobDetail: RemoteMachineTrialJobDetail): RemoteMachineScheduleResult {
202
203
204
        assert(gpuInfos.length >= requiredGPUNum);
        const allocatedGPUs: GPUInfo[] = this.selectGPUsForTrial(gpuInfos, requiredGPUNum);
        allocatedGPUs.forEach((gpuInfo: GPUInfo) => {
205
206
207
            if (rmMeta.occupiedGpuIndexMap !== undefined) {
                let num: number | undefined = rmMeta.occupiedGpuIndexMap.get(gpuInfo.index);
                if (num === undefined) {
208
209
210
                    num = 0;
                }
                rmMeta.occupiedGpuIndexMap.set(gpuInfo.index, num + 1);
211
            } else {
212
                throw new Error(`Machine ${rmMeta.config.host} occupiedGpuIndexMap initialize error!`);
213
            }
214
        });
215
216
        trialJobDetail.gpuIndices = allocatedGPUs;
        trialJobDetail.rmMeta = rmMeta;
217

218
219
220
221
        return {
            resultType: ScheduleResultType.SUCCEED,
            scheduleInfo: {
                rmMeta: rmMeta,
chicm-ms's avatar
chicm-ms committed
222
                cudaVisibleDevice: allocatedGPUs
223
224
225
226
                    .map((gpuInfo: GPUInfo) => {
                        return gpuInfo.index;
                    })
                    .join(',')
227
228
229
            }
        };
    }
Deshui Yu's avatar
Deshui Yu committed
230
}