remoteMachineData.ts 9.49 KB
Newer Older
Deshui Yu's avatar
Deshui Yu committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
/**
 * Copyright (c) Microsoft Corporation
 * All rights reserved.
 *
 * MIT License
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
 * documentation files (the "Software"), to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and
 * to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 * The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
 * BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

'use strict';

22
import * as fs from 'fs';
SparkSnail's avatar
SparkSnail committed
23
24
import { Client, ConnectConfig } from 'ssh2';
import { Deferred } from 'ts-deferred';
25
import { JobApplicationForm, TrialJobDetail, TrialJobStatus  } from '../../common/trainingService';
26
import { GPUInfo, GPUSummary } from '../common/gpuData';
Deshui Yu's avatar
Deshui Yu committed
27
28
29
30
31

/**
 * Metadata of remote machine for configuration and statuc query
 */
export class RemoteMachineMeta {
32
33
34
35
    public readonly ip : string = '';
    public readonly port : number = 22;
    public readonly username : string = '';
    public readonly passwd: string = '';
36
37
    public readonly sshKeyPath?: string;
    public readonly passphrase?: string;
Deshui Yu's avatar
Deshui Yu committed
38
    public gpuSummary : GPUSummary | undefined;
39
    public readonly gpuIndices?: string;
40
    public readonly maxTrialNumPerGpu?: number;
41
42
    //TODO: initialize varialbe in constructor
    public occupiedGpuIndexMap?: Map<number, number>;
43
    public readonly useActiveGpu?: boolean = false;
44
45
46
47
48
49
50
51
52
53
54
}

export function parseGpuIndices(gpuIndices?: string): Set<number> | undefined {
    if (gpuIndices !== undefined) {
        const indices: number[] = gpuIndices.split(',')
            .map((x: string) => parseInt(x, 10));
        if (indices.length > 0) {
            return new Set(indices);
        } else {
            throw new Error('gpuIndices can not be empty if specified.');
        }
Deshui Yu's avatar
Deshui Yu committed
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
    }
}

/**
 * The execution result for command executed on remote machine
 */
export class RemoteCommandResult {
    public readonly stdout : string;
    public readonly stderr : string;
    public readonly exitCode : number;

    constructor(stdout : string, stderr : string, exitCode : number) {
        this.stdout = stdout;
        this.stderr = stderr;
        this.exitCode = exitCode;
    }
}

/**
 * RemoteMachineTrialJobDetail
 */
export class RemoteMachineTrialJobDetail implements TrialJobDetail {
    public id: string;
    public status: TrialJobStatus;
79
80
81
    public submitTime: number;
    public startTime?: number;
    public endTime?: number;
Deshui Yu's avatar
Deshui Yu committed
82
83
84
85
    public tags?: string[];
    public url?: string;
    public workingDirectory: string;
    public form: JobApplicationForm;
86
    public sequenceId: number;
Deshui Yu's avatar
Deshui Yu committed
87
    public rmMeta?: RemoteMachineMeta;
88
    public isEarlyStopped?: boolean;
89
    public gpuIndices: GPUInfo[];
Deshui Yu's avatar
Deshui Yu committed
90

91
92
    constructor(id: string, status: TrialJobStatus, submitTime: number,
                workingDirectory: string, form: JobApplicationForm, sequenceId: number) {
Deshui Yu's avatar
Deshui Yu committed
93
94
95
96
97
        this.id = id;
        this.status = status;
        this.submitTime = submitTime;
        this.workingDirectory = workingDirectory;
        this.form = form;
98
        this.sequenceId = sequenceId;
Deshui Yu's avatar
Deshui Yu committed
99
        this.tags = [];
100
        this.gpuIndices = [];
Deshui Yu's avatar
Deshui Yu committed
101
102
103
    }
}

SparkSnail's avatar
SparkSnail committed
104
105
106
107
108
109
110
111
112
113
/**
 * The remote machine ssh client used for trial and gpu detector
 */
export class SSHClient {
    private readonly sshClient: Client;
    private usedConnectionNumber: number; //count the connection number of every client
    constructor(sshClient: Client, usedConnectionNumber: number) {
        this.sshClient = sshClient;
        this.usedConnectionNumber = usedConnectionNumber;
    }
114

SparkSnail's avatar
SparkSnail committed
115
116
117
118
119
120
121
122
    public get getSSHClientInstance(): Client {
        return this.sshClient;
    }

    public get getUsedConnectionNumber(): number {
        return this.usedConnectionNumber;
    }

123
    public addUsedConnectionNumber(): void {
SparkSnail's avatar
SparkSnail committed
124
125
126
        this.usedConnectionNumber += 1;
    }

127
    public minusUsedConnectionNumber(): void {
SparkSnail's avatar
SparkSnail committed
128
129
130
131
        this.usedConnectionNumber -= 1;
    }
}

132
133
134
/**
 * The remote machine ssh client manager
 */
SparkSnail's avatar
SparkSnail committed
135
export class SSHClientManager {
136
    private readonly sshClientArray: SSHClient[];
SparkSnail's avatar
SparkSnail committed
137
138
139
140
141
142
143
144
145
146
147
148
149
    private readonly maxTrialNumberPerConnection: number;
    private readonly rmMeta: RemoteMachineMeta;
    constructor(sshClientArray: SSHClient[], maxTrialNumberPerConnection: number, rmMeta: RemoteMachineMeta) {
        this.rmMeta = rmMeta;
        this.sshClientArray = sshClientArray;
        this.maxTrialNumberPerConnection = maxTrialNumberPerConnection;
    }

    /**
     * find a available ssh client in ssh array, if no ssh client available, return undefined
     */
    public async getAvailableSSHClient(): Promise<Client> {
        const deferred: Deferred<Client> = new Deferred<Client>();
150
151
152
        for (const index of this.sshClientArray.keys()) {
            const connectionNumber: number = this.sshClientArray[index].getUsedConnectionNumber;
            if (connectionNumber < this.maxTrialNumberPerConnection) {
SparkSnail's avatar
SparkSnail committed
153
154
                this.sshClientArray[index].addUsedConnectionNumber();
                deferred.resolve(this.sshClientArray[index].getSSHClientInstance);
155

SparkSnail's avatar
SparkSnail committed
156
157
                return deferred.promise;
            }
158
159
        }

SparkSnail's avatar
SparkSnail committed
160
        //init a new ssh client if could not get an available one
161
        return this.initNewSSHClient();
SparkSnail's avatar
SparkSnail committed
162
    }
163

SparkSnail's avatar
SparkSnail committed
164
165
    /**
     * add a new ssh client to sshClientArray
166
     * @param sshClient SSH Client
SparkSnail's avatar
SparkSnail committed
167
     */
168
    public addNewSSHClient(client: Client): void {
SparkSnail's avatar
SparkSnail committed
169
170
        this.sshClientArray.push(new SSHClient(client, 1));
    }
171

SparkSnail's avatar
SparkSnail committed
172
    /**
173
     * first ssh client instance is used for gpu collector and host job
SparkSnail's avatar
SparkSnail committed
174
     */
175
    public getFirstSSHClient(): Client {
SparkSnail's avatar
SparkSnail committed
176
177
        return this.sshClientArray[0].getSSHClientInstance;
    }
178

SparkSnail's avatar
SparkSnail committed
179
180
181
    /**
     * close all of ssh client
     */
182
183
    public closeAllSSHClient(): void {
        for (const sshClient of this.sshClientArray) {
SparkSnail's avatar
SparkSnail committed
184
185
186
            sshClient.getSSHClientInstance.end();
        }
    }
187

SparkSnail's avatar
SparkSnail committed
188
189
    /**
     * retrieve resource, minus a number for given ssh client
190
     * @param client SSH Client
SparkSnail's avatar
SparkSnail committed
191
     */
192
193
    public releaseConnection(client: Client | undefined): void {
        if (client === undefined) {
SparkSnail's avatar
SparkSnail committed
194
195
            throw new Error(`could not release a undefined ssh client`);
        }
196
197
        for (const index of this.sshClientArray.keys()) {
            if (this.sshClientArray[index].getSSHClientInstance === client) {
SparkSnail's avatar
SparkSnail committed
198
199
200
201
202
203
                this.sshClientArray[index].minusUsedConnectionNumber();
                break;
            }
        }
    }

204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
    /**
     * Create a new ssh connection client and initialize it
     */
    // tslint:disable:non-literal-fs-path
    private initNewSSHClient(): Promise<Client> {
        const deferred: Deferred<Client> = new Deferred<Client>();
        const conn: Client = new Client();
        const connectConfig: ConnectConfig = {
            host: this.rmMeta.ip,
            port: this.rmMeta.port,
            username: this.rmMeta.username };
        if (this.rmMeta.passwd !== undefined) {
            connectConfig.password = this.rmMeta.passwd;
        } else if (this.rmMeta.sshKeyPath !== undefined) {
            if (!fs.existsSync(this.rmMeta.sshKeyPath)) {
                //SSh key path is not a valid file, reject
                deferred.reject(new Error(`${this.rmMeta.sshKeyPath} does not exist.`));
            }
            const privateKey: string = fs.readFileSync(this.rmMeta.sshKeyPath, 'utf8');

            connectConfig.privateKey = privateKey;
            connectConfig.passphrase = this.rmMeta.passphrase;
        } else {
            deferred.reject(new Error(`No valid passwd or sshKeyPath is configed.`));
        }
        conn.on('ready', () => {
            this.addNewSSHClient(conn);
            deferred.resolve(conn);
        })
          .on('error', (err: Error) => {
            // SSH connection error, reject with error message
            deferred.reject(new Error(err.message));
        })
          .connect(connectConfig);

        return deferred.promise;
    }
}
SparkSnail's avatar
SparkSnail committed
242

243
export type RemoteMachineScheduleResult = { scheduleInfo : RemoteMachineScheduleInfo | undefined; resultType : ScheduleResultType};
Deshui Yu's avatar
Deshui Yu committed
244

245
export type RemoteMachineScheduleInfo = { rmMeta : RemoteMachineMeta; cuda_visible_device : string};
Deshui Yu's avatar
Deshui Yu committed
246
247

export enum ScheduleResultType {
248
    // Schedule succeeded
Deshui Yu's avatar
Deshui Yu committed
249
250
    SUCCEED,

251
    // Temporarily, no enough available GPU right now
Deshui Yu's avatar
Deshui Yu committed
252
253
    TMP_NO_AVAILABLE_GPU,

254
    // Cannot match requirement even if all GPU are a
Deshui Yu's avatar
Deshui Yu committed
255
256
257
    REQUIRE_EXCEED_TOTAL
}

SparkSnail's avatar
SparkSnail committed
258
export const REMOTEMACHINE_TRIAL_COMMAND_FORMAT: string =
Deshui Yu's avatar
Deshui Yu committed
259
`#!/bin/bash
260
261
export NNI_PLATFORM=remote NNI_SYS_DIR={0} NNI_OUTPUT_DIR={1} NNI_TRIAL_JOB_ID={2} NNI_EXP_ID={3} \
NNI_TRIAL_SEQ_ID={4} export MULTI_PHASE={5}
Deshui Yu's avatar
Deshui Yu committed
262
cd $NNI_SYS_DIR
SparkSnail's avatar
SparkSnail committed
263
264
sh install_nni.sh
echo $$ >{6}
265
266
python3 -m nni_trial_tool.trial_keeper --trial_command '{7}' --nnimanager_ip '{8}' --nnimanager_port '{9}' \
--nni_manager_version '{10}' --log_collection '{11}' 1>$NNI_OUTPUT_DIR/trialkeeper_stdout 2>$NNI_OUTPUT_DIR/trialkeeper_stderr
SparkSnail's avatar
SparkSnail committed
267
echo $? \`date +%s%3N\` >{12}`;
Deshui Yu's avatar
Deshui Yu committed
268

269
export const HOST_JOB_SHELL_FORMAT: string =
Deshui Yu's avatar
Deshui Yu committed
270
271
272
273
274
`#!/bin/bash
cd {0}
echo $$ >{1}
eval {2} >stdout 2>stderr
echo $? \`date +%s%3N\` >{3}`;