main.ts 2.79 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';
liuzhe-lz's avatar
liuzhe-lz committed
22
import { RestServer } from './rest_server';
23
import { parseArgs } from 'common/globals/arguments';
24

25
const args = parseArgs(process.argv.slice(2));
Deshui Yu's avatar
Deshui Yu committed
26

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

    await ds.init();
}

54
55
56
57
58
59
60
61
62
63
64
setExperimentStartupInfo(
    args.action === 'create',
    args.experimentId,
    args.port,
    args.mode,
    args.experimentsDirectory,
    args.logLevel,
    args.action === 'view',
    args.dispatcherPipe ?? '',
    args.urlPrefix
);
Deshui Yu's avatar
Deshui Yu committed
65

66
67
mkDirP(getLogDir())
    .then(async () => {
68
        try {
69
            await initContainer();
liuzhe-lz's avatar
liuzhe-lz committed
70
            const restServer: RestServer = component.get(RestServer);
71
72
            await restServer.start();
        } catch (err) {
73
            getLogger('main').error(`${err.stack}`);
74
75
76
77
78
79
            throw err;
        }
    })
    .catch((err: Error) => {
        console.error(`Failed to create log dir: ${err.stack}`);
    });
80

81
82
function cleanUp(): void {
    (component.get(Manager) as Manager).stopExperiment();
83
84
85
86
87
}

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