remoteMachineData.ts 4.83 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

'use strict';

6
import { TrialJobApplicationForm, TrialJobDetail, TrialJobStatus } from '../../common/trainingService';
7
import { GPUInfo, GPUSummary, ScheduleResultType } from '../common/gpuData';
8
import { ShellExecutor } from './shellExecutor';
Deshui Yu's avatar
Deshui Yu committed
9
10
11
12
13

/**
 * Metadata of remote machine for configuration and statuc query
 */
export class RemoteMachineMeta {
chicm-ms's avatar
chicm-ms committed
14
15
16
    public readonly ip: string = '';
    public readonly port: number = 22;
    public readonly username: string = '';
17
    public readonly passwd: string = '';
18
19
    public readonly sshKeyPath?: string;
    public readonly passphrase?: string;
chicm-ms's avatar
chicm-ms committed
20
    public gpuSummary: GPUSummary | undefined;
21
    public readonly gpuIndices?: string;
22
    public readonly maxTrialNumPerGpu?: number;
23
24
    //TODO: initialize varialbe in constructor
    public occupiedGpuIndexMap?: Map<number, number>;
25
    public readonly useActiveGpu?: boolean = false;
26
    public readonly preCommand?: string;
27
28
}

Deshui Yu's avatar
Deshui Yu committed
29
30
31
32
/**
 * The execution result for command executed on remote machine
 */
export class RemoteCommandResult {
chicm-ms's avatar
chicm-ms committed
33
34
35
    public readonly stdout: string;
    public readonly stderr: string;
    public readonly exitCode: number;
Deshui Yu's avatar
Deshui Yu committed
36

chicm-ms's avatar
chicm-ms committed
37
    constructor(stdout: string, stderr: string, exitCode: number) {
Deshui Yu's avatar
Deshui Yu committed
38
39
40
41
42
43
44
45
46
47
48
49
        this.stdout = stdout;
        this.stderr = stderr;
        this.exitCode = exitCode;
    }
}

/**
 * RemoteMachineTrialJobDetail
 */
export class RemoteMachineTrialJobDetail implements TrialJobDetail {
    public id: string;
    public status: TrialJobStatus;
50
51
52
    public submitTime: number;
    public startTime?: number;
    public endTime?: number;
Deshui Yu's avatar
Deshui Yu committed
53
54
55
    public tags?: string[];
    public url?: string;
    public workingDirectory: string;
56
    public form: TrialJobApplicationForm;
Deshui Yu's avatar
Deshui Yu committed
57
    public rmMeta?: RemoteMachineMeta;
58
    public isEarlyStopped?: boolean;
59
    public gpuIndices: GPUInfo[];
Deshui Yu's avatar
Deshui Yu committed
60

61
    constructor(id: string, status: TrialJobStatus, submitTime: number,
62
        workingDirectory: string, form: TrialJobApplicationForm) {
Deshui Yu's avatar
Deshui Yu committed
63
64
65
66
67
68
        this.id = id;
        this.status = status;
        this.submitTime = submitTime;
        this.workingDirectory = workingDirectory;
        this.form = form;
        this.tags = [];
69
        this.gpuIndices = [];
Deshui Yu's avatar
Deshui Yu committed
70
71
72
    }
}

SparkSnail's avatar
SparkSnail committed
73
/**
74
 * The remote machine executor manager
SparkSnail's avatar
SparkSnail committed
75
 */
76
export class ExecutorManager {
77
    private readonly executorMap: Map<string, ShellExecutor> = new Map<string, ShellExecutor>();
SparkSnail's avatar
SparkSnail committed
78
    private readonly rmMeta: RemoteMachineMeta;
79
80
81
82

    private executors: ShellExecutor[] = [];

    constructor(rmMeta: RemoteMachineMeta) {
SparkSnail's avatar
SparkSnail committed
83
84
85
        this.rmMeta = rmMeta;
    }

86
87
88
    public async getExecutor(id: string): Promise<ShellExecutor> {
        let isFound = false;
        let executor: ShellExecutor | undefined;
89

90
91
92
93
94
        // already assigned
        if (this.executorMap.has(id)) {
            executor = this.executorMap.get(id);
            if (executor === undefined) {
                throw new Error("executor shouldn't be undefined before return!");
SparkSnail's avatar
SparkSnail committed
95
            }
96
            return executor;
97
98
        }

99
100
101
102
103
104
105
106
107
108
109
        for (const candidateExecutor of this.executors) {
            if (candidateExecutor.addUsage()) {
                isFound = true;
                executor = candidateExecutor;
                break;
            }
        }
        // init a new executor if no free one.
        if (!isFound) {
            executor = await this.createShellExecutor();
        }
110

111
112
113
114
        if (executor === undefined) {
            throw new Error("executor shouldn't be undefined before set!");
        }
        this.executorMap.set(id, executor);
115

116
        return executor;
SparkSnail's avatar
SparkSnail committed
117
    }
118

SparkSnail's avatar
SparkSnail committed
119
    /**
120
     * close all of executor
SparkSnail's avatar
SparkSnail committed
121
     */
122
123
124
    public releaseAllExecutor(): void {
        this.executorMap.clear();
        for (const executor of this.executors) {
125
            executor.close();
SparkSnail's avatar
SparkSnail committed
126
        }
127
        this.executors = [];
SparkSnail's avatar
SparkSnail committed
128
    }
129

SparkSnail's avatar
SparkSnail committed
130
    /**
131
132
     * retrieve resource, minus a number for given executor
     * @param executor executor
SparkSnail's avatar
SparkSnail committed
133
     */
134
135
    public releaseExecutor(id: string): void {
        const executor = this.executorMap.get(id);
136
        if (executor === undefined) {
137
            throw new Error(`executor for ${id} is not found`);
SparkSnail's avatar
SparkSnail committed
138
        }
139
140
        executor.releaseUsage();
        this.executorMap.delete(id);
SparkSnail's avatar
SparkSnail committed
141
142
    }

143
    /**
144
     * Create a new connection executor and initialize it
145
     */
146
    private async createShellExecutor(): Promise<ShellExecutor> {
147
148
        const executor = new ShellExecutor();
        await executor.initialize(this.rmMeta);
149
150
151
152
        if (!executor.addUsage()) {
            throw new Error("failed to add usage on new created Executor! It's a wired bug!");
        }
        this.executors.push(executor);
153
        return executor;
154
155
    }
}
SparkSnail's avatar
SparkSnail committed
156

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

159
export type RemoteMachineScheduleInfo = { rmMeta: RemoteMachineMeta; cudaVisibleDevice: string };