main.ts 5.19 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
import 'app-module-path/register';
Deshui Yu's avatar
Deshui Yu committed
5
6
import { Container, Scope } from 'typescript-ioc';

7
import * as fs from 'fs';
8
import * as path from 'path';
9
import * as component from './common/component';
Deshui Yu's avatar
Deshui Yu committed
10
11
import { Database, DataStore } from './common/datastore';
import { setExperimentStartupInfo } from './common/experimentStartupInfo';
12
import { getLogger, setLogLevel, startLogging } from './common/log';
SparkSnail's avatar
SparkSnail committed
13
import { Manager, ExperimentStartUpMode } from './common/manager';
14
import { ExperimentManager } from './common/experimentManager';
J-shang's avatar
J-shang committed
15
import { TensorboardManager } from './common/tensorboardManager';
16
import { getLogDir, mkDirP, parseArg } from './common/utils';
Deshui Yu's avatar
Deshui Yu committed
17
18
19
import { NNIDataStore } from './core/nniDataStore';
import { NNIManager } from './core/nnimanager';
import { SqlDB } from './core/sqlDatabase';
20
import { NNIExperimentsManager } from './core/nniExperimentsManager';
J-shang's avatar
J-shang committed
21
import { NNITensorboardManager } from './core/nniTensorboardManager';
22
import { NNIRestServer } from './rest_server/nniRestServer';
Deshui Yu's avatar
Deshui Yu committed
23

24

25
function initStartupInfo(
26
    startExpMode: string, experimentId: string, basePort: number, platform: string,
27
    logDirectory: string, experimentLogLevel: string, readonly: boolean, dispatcherPipe: string, urlprefix: string): void {
SparkSnail's avatar
SparkSnail committed
28
    const createNew: boolean = (startExpMode === ExperimentStartUpMode.NEW);
29
    setExperimentStartupInfo(createNew, experimentId, basePort, platform, logDirectory, experimentLogLevel, readonly, dispatcherPipe, urlprefix);
Deshui Yu's avatar
Deshui Yu committed
30
31
}

32
async function initContainer(foreground: boolean, _platformMode: string, logFileName?: string): Promise<void> {
33
34
35
36
37
38
39
40
41
    Container.bind(Manager)
        .to(NNIManager)
        .scope(Scope.Singleton);
    Container.bind(Database)
        .to(SqlDB)
        .scope(Scope.Singleton);
    Container.bind(DataStore)
        .to(NNIDataStore)
        .scope(Scope.Singleton);
42
43
44
    Container.bind(ExperimentManager)
        .to(NNIExperimentsManager)
        .scope(Scope.Singleton);
J-shang's avatar
J-shang committed
45
46
47
    Container.bind(TensorboardManager)
        .to(NNITensorboardManager)
        .scope(Scope.Singleton);
48
    const DEFAULT_LOGFILE: string = path.join(getLogDir(), 'nnimanager.log');
49
50
51
52
53
54
    if (!foreground) {
        if (logFileName === undefined) {
            startLogging(DEFAULT_LOGFILE);
        } else {
            startLogging(logFileName);
        }
55
    }
J-shang's avatar
J-shang committed
56
57
    // eslint-disable-next-line @typescript-eslint/no-use-before-define
    setLogLevel(logLevel);
Deshui Yu's avatar
Deshui Yu committed
58
59
60
61
62
63
    const ds: DataStore = component.get(DataStore);

    await ds.init();
}

function usage(): void {
64
    console.info('usage: node main.js --port <port> --mode \
65
    <local/remote/pai/kubeflow/frameworkcontroller/aml/adl/hybrid/dlc> --start_mode <new/resume> --experiment_id <id> --foreground <true/false>');
Deshui Yu's avatar
Deshui Yu committed
66
67
68
}

const strPort: string = parseArg(['--port', '-p']);
goooxu's avatar
goooxu committed
69
70
71
if (!strPort || strPort.length === 0) {
    usage();
    process.exit(1);
Deshui Yu's avatar
Deshui Yu committed
72
73
}

74
75
76
77
78
79
80
81
const foregroundArg: string = parseArg(['--foreground', '-f']);
if (!('true' || 'false').includes(foregroundArg.toLowerCase())) {
    console.log(`FATAL: foreground property should only be true or false`);
    usage();
    process.exit(1);
}
const foreground: boolean = foregroundArg.toLowerCase() === 'true' ? true : false;

goooxu's avatar
goooxu committed
82
83
const port: number = parseInt(strPort, 10);

Deshui Yu's avatar
Deshui Yu committed
84
85
86
const mode: string = parseArg(['--mode', '-m']);

const startMode: string = parseArg(['--start_mode', '-s']);
SparkSnail's avatar
SparkSnail committed
87
if (![ExperimentStartUpMode.NEW, ExperimentStartUpMode.RESUME].includes(startMode)) {
88
    console.log(`FATAL: unknown start_mode: ${startMode}`);
Deshui Yu's avatar
Deshui Yu committed
89
90
91
92
93
    usage();
    process.exit(1);
}

const experimentId: string = parseArg(['--experiment_id', '-id']);
94
if (experimentId.trim().length < 1) {
SparkSnail's avatar
SparkSnail committed
95
    console.log(`FATAL: cannot resume the experiment, invalid experiment_id: ${experimentId}`);
Deshui Yu's avatar
Deshui Yu committed
96
97
98
99
    usage();
    process.exit(1);
}

100
101
102
103
104
105
106
107
108
const logDir: string = parseArg(['--log_dir', '-ld']);
if (logDir.length > 0) {
    if (!fs.existsSync(logDir)) {
        console.log(`FATAL: log_dir ${logDir} does not exist`);
    }
}

const logLevel: string = parseArg(['--log_level', '-ll']);

SparkSnail's avatar
SparkSnail committed
109
110
111
112
113
114
115
116
const readonlyArg: string = parseArg(['--readonly', '-r']);
if (!('true' || 'false').includes(readonlyArg.toLowerCase())) {
    console.log(`FATAL: readonly property should only be true or false`);
    usage();
    process.exit(1);
}
const readonly = readonlyArg.toLowerCase() == 'true' ? true : false;

117
118
const dispatcherPipe: string = parseArg(['--dispatcher_pipe']);

119
120
121
const urlPrefix: string = parseArg(['--url_prefix']);

initStartupInfo(startMode, experimentId, port, mode, logDir, logLevel, readonly, dispatcherPipe, urlPrefix);
Deshui Yu's avatar
Deshui Yu committed
122

123
124
mkDirP(getLogDir())
    .then(async () => {
125
126
127
128
        try {
            await initContainer(foreground, mode);
            const restServer: NNIRestServer = component.get(NNIRestServer);
            await restServer.start();
129
            getLogger('main').info(`Rest server listening on: ${restServer.endPoint}`);
130
        } catch (err) {
131
            getLogger('main').error(`${err.stack}`);
132
133
134
135
136
137
            throw err;
        }
    })
    .catch((err: Error) => {
        console.error(`Failed to create log dir: ${err.stack}`);
    });
138

139
140
function cleanUp(): void {
    (component.get(Manager) as Manager).stopExperiment();
141
142
143
144
145
}

process.on('SIGTERM', cleanUp);
process.on('SIGBREAK', cleanUp);
process.on('SIGINT', cleanUp);