Support remote training service use reuse mode

docs/en_US/Tutorial/ExperimentConfig.md

@@ -592,6 +592,16 @@ Specifies the pre-command that will be executed before the remote machine execut
 
 __Note__: Because __preCommand__ will execute before other commands each time, it is strongly not recommended to set __preCommand__ that will make changes to system, i.e. `mkdir` or `touch`.
 
+### remoteConfig
+
+Optional field in remote mode. Users could set per machine information in `machineList` field, and set global configuration for remote mode in this field.
+
+#### reuse
+
+Optional. Bool. default: `false`. It's an experimental feature.
+
+If it's true, NNI will reuse remote jobs to run as many as possible trials. It can save time of creating new jobs. User needs to make sure each trial can run independent in same job, for example, avoid loading checkpoint from previous trials. 
+
 ### kubeflowConfig
 
 #### operator

src/nni_manager/main.ts

@@ -23,9 +23,6 @@ import { KubeflowTrainingService } from './training_service/kubernetes/kubeflow/
 import { LocalTrainingService } from './training_service/local/localTrainingService';
 import { RouterTrainingService } from './training_service/reusable/routerTrainingService';
 import { PAIYarnTrainingService } from './training_service/pai/paiYarn/paiYarnTrainingService';
-import {
-    RemoteMachineTrainingService
-} from './training_service/remote_machine/remoteMachineTrainingService';
 import { DLTSTrainingService } from './training_service/dlts/dltsTrainingService';
 
 function initStartupInfo(
@@ -43,7 +40,7 @@ async function initContainer(foreground: boolean, platformMode: string, logFileN
             .scope(Scope.Singleton);
     } else if (platformMode === 'remote') {
         Container.bind(TrainingService)
-            .to(RemoteMachineTrainingService)
+            .to(RouterTrainingService)
             .scope(Scope.Singleton);
     } else if (platformMode === 'pai') {
         Container.bind(TrainingService)

src/nni_manager/rest_server/restValidationSchemas.ts

@@ -166,6 +166,9 @@ export namespace ValidationSchemas {
             }),
             nni_manager_ip: joi.object({ // eslint-disable-line @typescript-eslint/camelcase
                 nniManagerIp: joi.string().min(1)
+            }),
+            remote_config: joi.object({ // eslint-disable-line @typescript-eslint/camelcase
+                reuse: joi.boolean()
             })
         }
     };

src/nni_manager/training_service/common/trialConfigMetadataKey.ts

@@ -10,6 +10,7 @@ export enum TrialConfigMetadataKey {
     MACHINE_LIST = 'machine_list',
     LOCAL_CONFIG = 'local_config',
     TRIAL_CONFIG = 'trial_config',
+    REMOTE_CONFIG = 'remote_config',
     EXPERIMENT_ID = 'experimentId',
     MULTI_PHASE = 'multiPhase',
     RANDOM_SCHEDULER = 'random_scheduler',

src/nni_manager/training_service/reusable/environment.ts

@@ -125,8 +125,14 @@ export abstract class EnvironmentService {
     public abstract get hasStorageService(): boolean;
     public abstract config(key: string, value: string): Promise<void>;
     public abstract refreshEnvironmentsStatus(environments: EnvironmentInformation[]): Promise<void>;
-    public abstract startEnvironment(environment: EnvironmentInformation): Promise<void>;
     public abstract stopEnvironment(environment: EnvironmentInformation): Promise<void>;
+    public abstract startEnvironment(environment: EnvironmentInformation): Promise<void>;
+
+    // It is used to set prefetched environment count, default value is 0 for OpenPAI and AML mode,
+    // in remote mode, this value is set to the length of machine list.
+    public get prefetchedEnvironmentCount(): number {
+        return 0;
+    }
 
     // It depends on environment pressure and settings
     // for example, OpenPAI relies on API calls, and there is an limitation for frequence, so it need to be bigger.

src/nni_manager/training_service/reusable/environments/remoteEnvironmentService.ts

@@ -0,0 +1,320 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+'use strict';
+
+import * as fs from 'fs';
+import * as path from 'path';
+import { Deferred } from 'ts-deferred';
+import * as component from '../../../common/component';
+import { getExperimentId } from '../../../common/experimentStartupInfo';
+import { getLogger, Logger } from '../../../common/log';
+import { EnvironmentInformation, EnvironmentService } from '../environment';
+import { ObservableTimer } from '../../../common/observableTimer';
+import {
+    getExperimentRootDir,
+} from '../../../common/utils';
+import { TrialConfig } from '../../common/trialConfig';
+import { TrialConfigMetadataKey } from '../../common/trialConfigMetadataKey';
+import { execMkdir, validateCodeDir } from '../../common/util';
+import {
+    ExecutorManager, RemoteMachineMeta,
+} from '../../remote_machine/remoteMachineData';
+import { ShellExecutor } from 'training_service/remote_machine/shellExecutor';
+import { RemoteMachineEnvironmentInformation } from '../remote/remoteConfig';
+
+
+@component.Singleton
+export class RemoteEnvironmentService extends EnvironmentService {
+
+    private readonly initExecutorId = "initConnection";
+    private readonly machineExecutorManagerMap: Map<RemoteMachineMeta, ExecutorManager>;
+    private readonly machineCopyExpCodeDirPromiseMap: Map<RemoteMachineMeta, Promise<void>>;
+    private readonly environmentExecutorManagerMap: Map<string, ExecutorManager>;
+    private readonly remoteMachineMetaOccupiedMap: Map<RemoteMachineMeta, boolean>;
+    private trialConfig: TrialConfig | undefined;
+    private readonly log: Logger;
+    private sshConnectionPromises: any[];
+    private experimentRootDir: string;
+    private experimentId: string;
+
+    constructor() {
+        super();
+        this.experimentId = getExperimentId();
+        this.environmentExecutorManagerMap = new Map<string, ExecutorManager>();
+        this.machineCopyExpCodeDirPromiseMap = new Map<RemoteMachineMeta, Promise<void>>();
+        this.machineExecutorManagerMap = new Map<RemoteMachineMeta, ExecutorManager>();
+        this.remoteMachineMetaOccupiedMap = new Map<RemoteMachineMeta, boolean>();
+        this.sshConnectionPromises = [];
+        this.experimentRootDir = getExperimentRootDir();
+        this.experimentId = getExperimentId();
+        this.log = getLogger();
+    }
+
+    public get prefetchedEnvironmentCount(): number {
+        return this.machineExecutorManagerMap.size;
+    }
+
+    public get environmentMaintenceLoopInterval(): number {
+        return 5000;
+    }
+
+    public get hasMoreEnvironments(): boolean {
+        return false;
+    }
+
+    public get hasStorageService(): boolean {
+        return false;
+    }
+
+    public async config(key: string, value: string): Promise<void> {
+        switch (key) {
+            case TrialConfigMetadataKey.MACHINE_LIST:
+                await this.setupConnections(value);
+                break;
+            case TrialConfigMetadataKey.TRIAL_CONFIG: {
+                const remoteMachineTrailConfig: TrialConfig = <TrialConfig>JSON.parse(value);
+                // Parse trial config failed, throw Error
+                if (remoteMachineTrailConfig === undefined) {
+                    throw new Error('trial config parsed failed');
+                }
+                // codeDir is not a valid directory, throw Error
+                if (!fs.lstatSync(remoteMachineTrailConfig.codeDir)
+                    .isDirectory()) {
+                    throw new Error(`codeDir ${remoteMachineTrailConfig.codeDir} is not a directory`);
+                }
+                try {
+                    // Validate to make sure codeDir doesn't have too many files
+                    await validateCodeDir(remoteMachineTrailConfig.codeDir);
+                } catch (error) {
+                    this.log.error(error);
+                    return Promise.reject(new Error(error));
+                }
+
+                this.trialConfig = remoteMachineTrailConfig;
+                break;
+            }
+            default:
+                this.log.debug(`Remote not support metadata key: '${key}', value: '${value}'`);
+        }
+    }
+
+    private scheduleMachine(): RemoteMachineMeta | undefined {
+        for (const [rmMeta, occupied] of this.remoteMachineMetaOccupiedMap) {
+            if (!occupied) {
+                this.remoteMachineMetaOccupiedMap.set(rmMeta, true);
+                return rmMeta;
+            }
+        }
+        return undefined;
+    }
+
+    private recycleMachineReservation(rmMeta: RemoteMachineMeta): void {
+        if (!this.remoteMachineMetaOccupiedMap.has(rmMeta)) {
+            throw new Error(`${rmMeta} not initialized!`);
+        }
+        this.remoteMachineMetaOccupiedMap.set(rmMeta, false);
+    }
+
+    private async setupConnections(machineList: string): Promise<void> {
+        this.log.debug(`Connecting to remote machines: ${machineList}`);
+        //TO DO: verify if value's format is wrong, and json parse failed, how to handle error
+        const rmMetaList: RemoteMachineMeta[] = <RemoteMachineMeta[]>JSON.parse(machineList);
+
+        for (const rmMeta of rmMetaList) {
+            this.sshConnectionPromises.push(await this.initRemoteMachineOnConnected(rmMeta));
+        }
+    }
+
+    private async initRemoteMachineOnConnected(rmMeta: RemoteMachineMeta): Promise<void> {
+        rmMeta.occupiedGpuIndexMap = new Map<number, number>();
+        const executorManager: ExecutorManager = new ExecutorManager(rmMeta);
+        this.log.info(`connecting to ${rmMeta.username}@${rmMeta.ip}:${rmMeta.port}`);
+        const executor: ShellExecutor = await executorManager.getExecutor(this.initExecutorId);
+        this.log.debug(`reached ${executor.name}`);
+        this.machineExecutorManagerMap.set(rmMeta, executorManager);
+        this.log.debug(`initializing ${executor.name}`);
+
+        // Create root working directory after executor is ready
+        const nniRootDir: string = executor.joinPath(executor.getTempPath(), 'nni');
+        await executor.createFolder(executor.getRemoteExperimentRootDir(getExperimentId()));
+
+        // the directory to store temp scripts in remote machine
+        const remoteGpuScriptCollectorDir: string = executor.getRemoteScriptsPath(getExperimentId());
+
+        // clean up previous result.
+        await executor.createFolder(remoteGpuScriptCollectorDir, true);
+        await executor.allowPermission(true, nniRootDir);
+    }
+
+    public async refreshEnvironmentsStatus(environments: EnvironmentInformation[]): Promise<void> {
+        environments.forEach(async (environment) => {
+            const executor = await this.getExecutor(environment.id);
+            const jobpidPath: string = `${environment.runnerWorkingFolder}/pid`;
+            const trialReturnCodeFilePath: string = `${environment.runnerWorkingFolder}/code`;
+            if (fs.existsSync(jobpidPath)) {
+                /* eslint-disable require-atomic-updates */
+                try {
+                    const isAlive = await executor.isProcessAlive(jobpidPath);
+                    // if the process of jobpid is not alive any more
+                    if (!isAlive) {
+                        this.log.info(`pid in ${jobpidPath} is not alive!`);
+                        if (fs.existsSync(trialReturnCodeFilePath)) {
+                            const trialReturnCode: string = await executor.getRemoteFileContent(trialReturnCodeFilePath);
+                            const match: RegExpMatchArray | null = trialReturnCode.trim()
+                                .match(/^-?(\d+)\s+(\d+)$/);
+                            if (match !== null) {
+                                const { 1: code } = match;
+                                // Update trial job's status based on result code
+                                if (parseInt(code, 10) === 0) {
+                                    environment.setStatus('SUCCEEDED');
+                                } else {
+                                    environment.setStatus('FAILED');
+                                }
+                                this.releaseEnvironmentResource(environment);
+                            }
+                        }
+                    }
+                } catch (error) {
+                    this.releaseEnvironmentResource(environment);
+                    this.log.error(`Update job status exception, error is ${error.message}`);
+                }
+            }
+        });
+    }
+
+    /**
+     * If a environment is finished, release the connection resource
+     * @param environment remote machine environment job detail
+     */
+    public releaseEnvironmentResource(environment: EnvironmentInformation): void {
+        const executorManager = this.environmentExecutorManagerMap.get(environment.id);
+        if (executorManager === undefined) {
+            throw new Error(`ExecutorManager is not assigned for environment ${environment.id}`);
+        }
+
+        // Note, it still keep reference in trialExecutorManagerMap, as there may be following requests from nni manager.
+        executorManager.releaseExecutor(environment.id);
+        const remoteEnvironment: RemoteMachineEnvironmentInformation = environment as RemoteMachineEnvironmentInformation;
+        if (remoteEnvironment.rmMachineMeta === undefined) {
+            throw new Error(`${remoteEnvironment.id} rmMachineMeta not initialized!`);
+        }
+        this.recycleMachineReservation(remoteEnvironment.rmMachineMeta);
+    }
+
+    public async startEnvironment(environment: EnvironmentInformation): Promise<void> {
+        if (this.sshConnectionPromises.length > 0) {
+            await Promise.all(this.sshConnectionPromises);
+            this.log.info('ssh connection initialized!');
+            // set sshConnectionPromises to [] to avoid log information duplicated
+            this.sshConnectionPromises = [];
+            if (this.trialConfig ===  undefined) {
+                throw new Error("trial config not initialized!");
+            }
+            Array.from(this.machineExecutorManagerMap.keys()).forEach(rmMeta => {
+                // initialize remoteMachineMetaOccupiedMap, false means not occupied
+                this.remoteMachineMetaOccupiedMap.set(rmMeta, false);
+            });
+        }
+        const remoteEnvironment: RemoteMachineEnvironmentInformation = environment as RemoteMachineEnvironmentInformation;
+        remoteEnvironment.status = 'WAITING';
+        await this.prepareEnvironment(remoteEnvironment);
+        await this.launchEnvironment(remoteEnvironment);
+    }
+
+    private async prepareEnvironment(environment: RemoteMachineEnvironmentInformation): Promise<boolean> {
+        const deferred: Deferred<boolean> = new Deferred<boolean>();
+        if (this.trialConfig === undefined) {
+            throw new Error('trial config is not initialized');
+        }
+
+        // get an executor from scheduler
+        const rmMachineMeta: RemoteMachineMeta | undefined = this.scheduleMachine();
+        if (rmMachineMeta === undefined) {
+            this.log.warning(`No available machine!`);
+            deferred.resolve(false);
+        } else {
+            environment.rmMachineMeta = rmMachineMeta;
+            const copyExpCodeDirPromise = this.machineCopyExpCodeDirPromiseMap.get(
+                environment.rmMachineMeta);
+            if (copyExpCodeDirPromise !== undefined) {
+                await copyExpCodeDirPromise;
+            }
+            this.allocateExecutorManagerForEnvironment(environment);
+            const executor = await this.getExecutor(environment.id);
+            environment.runnerWorkingFolder = 
+                executor.joinPath(executor.getRemoteExperimentRootDir(getExperimentId()), 
+                'envs', environment.id)
+            environment.command = `cd ${environment.runnerWorkingFolder} && \
+${environment.command} --job_pid_file ${environment.runnerWorkingFolder}/pid \
+&& echo $? \`date +%s%3N\` >${environment.runnerWorkingFolder}/code`;
+
+            await this.launchEnvironment(environment);
+
+            environment.status = 'RUNNING';
+            environment.trackingUrl = `file://${rmMachineMeta.ip}:${environment.runnerWorkingFolder}`;
+            deferred.resolve(true);
+        }
+
+        return deferred.promise;
+    }
+
+    /**
+     * give environment an executor
+     * @param environment RemoteMachineEnvironmentDetail
+     */
+    private allocateExecutorManagerForEnvironment(environment: RemoteMachineEnvironmentInformation): void {
+        if (environment.rmMachineMeta === undefined) {
+            throw new Error(`rmMeta not set in trial ${environment.id}`);
+        }
+        const executorManager: ExecutorManager | undefined = this.machineExecutorManagerMap.get(environment.rmMachineMeta);
+        if (executorManager === undefined) {
+            throw new Error(`executorManager not initialized`);
+        }
+        this.environmentExecutorManagerMap.set(environment.id, executorManager);
+    }
+
+    private async getExecutor(environmentId: string): Promise<ShellExecutor> {
+        const executorManager = this.environmentExecutorManagerMap.get(environmentId);
+        if (executorManager === undefined) {
+            throw new Error(`ExecutorManager is not assigned for environment ${environmentId}`);
+        }
+        return await executorManager.getExecutor(environmentId);
+    }
+
+    public async stopEnvironment(environment: EnvironmentInformation): Promise<void> {
+        const executor = await this.getExecutor(environment.id);
+
+        if (environment.status === 'UNKNOWN') {
+            environment.status = 'USER_CANCELED';
+            this.releaseEnvironmentResource(environment);
+            return
+        }
+
+        const jobpidPath: string = `${environment.runnerWorkingFolder}/pid`;
+        try {
+            await executor.killChildProcesses(jobpidPath);
+            this.releaseEnvironmentResource(environment);
+        } catch (error) {
+            this.log.error(`remoteTrainingService.cancelTrialJob: ${error}`);
+        }
+    }
+
+    private async launchEnvironment(environment: RemoteMachineEnvironmentInformation): Promise<void> {
+        if (this.trialConfig === undefined) {
+            throw new Error('trial config is not initialized');
+        }
+        const executor = await this.getExecutor(environment.id);
+        const environmentLocalTempFolder: string =  
+            path.join(this.experimentRootDir, this.experimentId, "environment-temp")
+        await executor.createFolder(environment.runnerWorkingFolder);
+        await execMkdir(environmentLocalTempFolder);
+        await fs.promises.writeFile(path.join(environmentLocalTempFolder, executor.getScriptName("run")),
+        environment.command, { encoding: 'utf8' });
+        // Copy files in codeDir to remote working directory
+        await executor.copyDirectoryToRemote(environmentLocalTempFolder, environment.runnerWorkingFolder);
+        // Execute command in remote machine
+        executor.executeScript(executor.joinPath(environment.runnerWorkingFolder,
+            executor.getScriptName("run")), true, false);
+    }
+}