paiConfig.ts 5.05 KB
Newer Older
liuzhe-lz's avatar
liuzhe-lz committed
1
2
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
3
4
5

'use strict';

6
import {TrialConfig} from '../common/trialConfig';
7

8
9
10
/**
 * Task role for PAI
 */
11
12
13
14
15
16
17
18
19
20
21
22
23
export class PAITaskRole {
    // Name for the task role
    public readonly name: string;
    // Number of tasks for the task role, no less than 1
    public readonly taskNumber: number;
    // CPU number for one task in the task role, no less than 1
    public readonly cpuNumber: number;
    // Memory for one task in the task role, no less than 100
    public readonly memoryMB: number;
    // GPU number for one task in the task role, no less than 0
    public readonly gpuNumber: number;
    // Executable command for tasks in the task role, can not be empty
    public readonly command: string;
24
25
    //Shared memory for one task in the task role
    public readonly shmMB?: number;
26
    //portList to specify the port used in container
chicm-ms's avatar
chicm-ms committed
27
    public portList?: PortListMetaData[];
28

29
30
31
32
33
34
35
36
37
    /**
     * Constructor
     * @param name Name for the task role
     * @param taskNumber Number of tasks for the task role, no less than 1
     * @param cpuNumber CPU number for one task in the task role, no less than 1
     * @param memoryMB Memory for one task in the task role, no less than 100
     * @param gpuNumber GPU number for one task in the task role, no less than 0
     * @param command Executable command for tasks in the task role, can not be empty
     */
chicm-ms's avatar
chicm-ms committed
38
39
    constructor(name: string, taskNumber: number, cpuNumber: number, memoryMB: number, gpuNumber: number,
                command: string, shmMB?: number, portList?: PortListMetaData[]) {
40
41
42
43
44
        this.name = name;
        this.taskNumber = taskNumber;
        this.cpuNumber = cpuNumber;
        this.memoryMB = memoryMB;
        this.gpuNumber = gpuNumber;
45
        this.command = command;
46
        this.shmMB = shmMB;
47
        this.portList = portList;
48
49
50
    }
}

51
52
53
54
/**
 * Trial job configuration submitted to PAI
 */
export class PAIJobConfig {
55
56
57
58
59
60
    // Name for the job, need to be unique
    public readonly jobName: string;
    // URL pointing to the Docker image for all tasks in the job
    public readonly image: string;
    // Code directory on HDFS
    public readonly codeDir: string;
61
62
    //authentication file used for private Docker registry 
    public readonly authFile?: string;
63
64
65
66

    // List of taskRole, one task role at least
    public taskRoles: PAITaskRole[];

67
68
69
    //The virtual cluster job runs on.
    public readonly virtualCluster: string;

70
71
72
73
74
75
76
77
    /**
     * Constructor
     * @param jobName Name for the job, need to be unique
     * @param image URL pointing to the Docker image for all tasks in the job
     * @param dataDir Data directory existing on HDFS
     * @param outputDir Output directory on HDFS
     * @param taskRoles List of taskRole, one task role at least
     */
chicm-ms's avatar
chicm-ms committed
78
79
    constructor(jobName: string, image: string, codeDir: string,
                taskRoles: PAITaskRole[], virtualCluster: string, authFile?: string) {
80
81
82
83
        this.jobName = jobName;
        this.image = image;
        this.codeDir = codeDir;
        this.taskRoles = taskRoles;
84
        this.virtualCluster = virtualCluster;
85
        this.authFile = authFile;
86
87
88
    }
}

89
90
91
/**
 * PAI cluster configuration
 */
92
93
export class PAIClusterConfig {
    public readonly userName: string;
94
    public readonly passWord?: string;
95
    public readonly host: string;
96
    public readonly token?: string;
97
98
99
100
101
102

    /**
     * Constructor
     * @param userName User name of PAI Cluster
     * @param passWord password of PAI Cluster
     * @param host Host IP of PAI Cluster
103
     * @param token PAI token of PAI Cluster
104
     */
chicm-ms's avatar
chicm-ms committed
105
    constructor(userName: string, host: string, passWord?: string, token?: string) {
106
107
108
        this.userName = userName;
        this.passWord = passWord;
        this.host = host;
109
        this.token = token;
110
111
112
    }
}

113
114
115
/**
 * portList data structure used in PAI taskRole
 */
chicm-ms's avatar
chicm-ms committed
116
117
export class PortListMetaData {
    public readonly label: string = '';
118
119
120
121
122
    public readonly beginAt: number = 0;
    public readonly portNumber: number = 0;
}
  

123
124
125
126
/**
 * PAI trial configuration
 */
export class NNIPAITrialConfig extends TrialConfig {
127
128
129
130
    public readonly cpuNum: number;
    public readonly memoryMB: number;
    public readonly image: string;

131
132
    //The virtual cluster job runs on. If omitted, the job will run on default virtual cluster
    public virtualCluster?: string;
133
134
    //Shared memory for one task in the task role
    public shmMB?: number;
135
136
    //authentication file used for private Docker registry 
    public authFile?: string;
137
    //portList to specify the port used in container
chicm-ms's avatar
chicm-ms committed
138
    public portList?: PortListMetaData[];
139

chicm-ms's avatar
chicm-ms committed
140
141
    constructor(command: string, codeDir: string, gpuNum: number, cpuNum: number, memoryMB: number,
                image: string, virtualCluster?: string, shmMB?: number, authFile?: string, portList?: PortListMetaData[]) {
142
143
144
145
        super(command, codeDir, gpuNum);
        this.cpuNum = cpuNum;
        this.memoryMB = memoryMB;
        this.image = image;
146
        this.virtualCluster = virtualCluster;
147
        this.shmMB = shmMB;
148
        this.authFile = authFile;
149
        this.portList = portList;
150
    }
151
}