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

'use strict';

import * as request from 'request';
import { Deferred } from 'ts-deferred';
import { NNIError, NNIErrorNames } from '../../common/errors';
9
import { getLogger, Logger } from '../../common/log';
10
import { TrialJobStatus } from '../../common/trainingService';
11
12
import { PAIClusterConfig } from './paiConfig';
import { PAITrialJobDetail } from './paiData';
13
14
15
16
17

/**
 * Collector PAI jobs info from PAI cluster, and update pai job status locally
 */
export class PAIJobInfoCollector {
chicm-ms's avatar
chicm-ms committed
18
    private readonly trialJobsMap: Map<string, PAITrialJobDetail>;
19
    private readonly log: Logger = getLogger();
chicm-ms's avatar
chicm-ms committed
20
21
    private readonly statusesNeedToCheck: TrialJobStatus[];
    private readonly finalStatuses: TrialJobStatus[];
22
23
24
25

    constructor(jobMap: Map<string, PAITrialJobDetail>) {
        this.trialJobsMap = jobMap;
        this.statusesNeedToCheck = ['RUNNING', 'UNKNOWN', 'WAITING'];
QuanluZhang's avatar
QuanluZhang committed
26
        this.finalStatuses = ['SUCCEEDED', 'FAILED', 'USER_CANCELED', 'SYS_CANCELED', 'EARLY_STOPPED'];
27
28
    }

chicm-ms's avatar
chicm-ms committed
29
    public async retrieveTrialStatus(paiToken? : string, paiClusterConfig?: PAIClusterConfig): Promise<void> {
30
        if (paiClusterConfig === undefined || paiToken === undefined) {
31
            return Promise.resolve();
32
33
        }

chicm-ms's avatar
chicm-ms committed
34
        const updatePaiTrialJobs: Promise<void>[] = [];
35
36
        for (const [trialJobId, paiTrialJob] of this.trialJobsMap) {
            if (paiTrialJob === undefined) {
37
38
                throw new NNIError(NNIErrorNames.NOT_FOUND, `trial job id ${trialJobId} not found`);
            }
39
            updatePaiTrialJobs.push(this.getSinglePAITrialJobInfo(paiTrialJob, paiToken, paiClusterConfig));
40
41
42
43
44
        }

        await Promise.all(updatePaiTrialJobs);
    }

chicm-ms's avatar
chicm-ms committed
45
46
    private getSinglePAITrialJobInfo(paiTrialJob: PAITrialJobDetail, paiToken: string, paiClusterConfig: PAIClusterConfig): Promise<void> {
        const deferred: Deferred<void> = new Deferred<void>();
47
48
        if (!this.statusesNeedToCheck.includes(paiTrialJob.status)) {
            deferred.resolve();
49

50
51
52
53
54
55
            return deferred.promise;
        }

        // Rest call to get PAI job info and update status
        // Refer https://github.com/Microsoft/pai/blob/master/docs/rest-server/API.md for more detail about PAI Rest API
        const getJobInfoRequest: request.Options = {
56
            uri: `http://${paiClusterConfig.host}/rest-server/api/v1/user/${paiClusterConfig.userName}/jobs/${paiTrialJob.paiJobName}`,
57
58
59
            method: 'GET',
            json: true,
            headers: {
60
61
                'Content-Type': 'application/json',
                Authorization: `Bearer ${paiToken}`
62
63
            }
        };
64

65
        //TODO : pass in request timeout param?
66
        request(getJobInfoRequest, (error: Error, response: request.Response, body: any) => {
67
            if ((error !== undefined && error !== null) || response.statusCode >= 500) {
68
69
                this.log.error(`PAI Training service: get job info for trial ${paiTrialJob.id} from PAI Cluster failed!`);
                // Queried PAI job info failed, set job status to UNKNOWN
70
                if (paiTrialJob.status === 'WAITING' || paiTrialJob.status === 'RUNNING') {
71
72
73
                    paiTrialJob.status = 'UNKNOWN';
                }
            } else {
74
75
                if (response.body.jobStatus && response.body.jobStatus.state) {
                    switch (response.body.jobStatus.state) {
76
                        case 'WAITING':
77
78
79
80
                            paiTrialJob.status = 'WAITING';
                            break;
                        case 'RUNNING':
                            paiTrialJob.status = 'RUNNING';
81
                            if (paiTrialJob.startTime === undefined) {
82
83
                                paiTrialJob.startTime = response.body.jobStatus.appLaunchedTime;
                            }
84
                            if (paiTrialJob.url === undefined) {
85
                                paiTrialJob.url = response.body.jobStatus.appTrackingUrl;
86
87
88
89
90
91
                            }
                            break;
                        case 'SUCCEEDED':
                            paiTrialJob.status = 'SUCCEEDED';
                            break;
                        case 'STOPPED':
92
                            if (paiTrialJob.isEarlyStopped !== undefined) {
93
                                paiTrialJob.status = paiTrialJob.isEarlyStopped === true ?
94
95
                                        'EARLY_STOPPED' : 'USER_CANCELED';
                            } else {
96
97
98
                                /* if paiTrialJob's isEarlyStopped is undefined, that mean we didn't stop it via cancellation,
                                 * mark it as SYS_CANCELLED by PAI
                                 */
99
                                paiTrialJob.status = 'SYS_CANCELED';
QuanluZhang's avatar
QuanluZhang committed
100
                            }
101
102
                            break;
                        case 'FAILED':
103
                            paiTrialJob.status = 'FAILED';
104
105
106
107
108
                            break;
                        default:
                            paiTrialJob.status = 'UNKNOWN';
                    }
                    // For final job statues, update startTime, endTime and url
109
110
                    if (this.finalStatuses.includes(paiTrialJob.status)) {
                        if (paiTrialJob.startTime === undefined) {
111
112
                            paiTrialJob.startTime = response.body.jobStatus.appLaunchedTime;
                        }
113
                        if (paiTrialJob.endTime === undefined) {
114
115
116
                            paiTrialJob.endTime = response.body.jobStatus.completedTime;
                        }
                        // Set pai trial job's url to WebHDFS output path
117
                        if (paiTrialJob.hdfsLogPath !== undefined) {
118
                            paiTrialJob.url += `,${paiTrialJob.hdfsLogPath}`;
119
120
121
122
123
124
125
126
127
                        }
                    }
                }
            }
            deferred.resolve();
        });

        return deferred.promise;
    }
QuanluZhang's avatar
QuanluZhang committed
128
}