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,9 @@ 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>;
+    public abstract getInitializeEnvironmentNumber(): number;
 
     // 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/amlEnvironmentService.ts

@@ -37,6 +37,10 @@ export class AMLEnvironmentService extends EnvironmentService {
         this.experimentRootDir = getExperimentRootDir();
     }
 
+    public getInitializeEnvironmentNumber(): number {
+        return 0;
+    }
+
     public get hasStorageService(): boolean {
         return false;
     }

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

@@ -37,6 +37,10 @@ export class OpenPaiEnvironmentService extends EnvironmentService {
         this.experimentId = getExperimentId();
     }
 
+    public getInitializeEnvironmentNumber(): number {
+        return 0;
+    }
+
     public get environmentMaintenceLoopInterval(): number {
         return 5000;
     }

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

@@ -0,0 +1,309 @@
+// 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 { MachineScheduler } from '../remote/machineScheduler';
+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 environmentJobsMap: Map<string, RemoteMachineEnvironmentInformation>;
+    private trialConfig: TrialConfig | undefined;
+    private machineScheduler?: MachineScheduler;
+    private readonly log: Logger;
+    private sshConnectionPromises: any[];
+    private experimentRootDir: string;
+    private experimentId: string;
+
+    constructor() {
+        super();
+        this.experimentId = getExperimentId();
+        this.environmentJobsMap = new Map<string, RemoteMachineEnvironmentInformation>();
+        this.environmentExecutorManagerMap = new Map<string, ExecutorManager>();
+        this.machineCopyExpCodeDirPromiseMap = new Map<RemoteMachineMeta, Promise<void>>();
+        this.machineExecutorManagerMap = new Map<RemoteMachineMeta, ExecutorManager>();
+        this.sshConnectionPromises = [];
+        this.experimentRootDir = getExperimentRootDir();
+        this.experimentId = getExperimentId();
+        this.log = getLogger();
+        this.log.info('Construct remote machine training service.');
+    }
+
+    public getInitializeEnvironmentNumber(): 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 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.releaseTrialResource(environment);
+                            }
+                        }
+                    }
+                } catch (error) {
+                    this.releaseTrialResource(environment);
+                    this.log.error(`Update job status exception, error is ${error.message}`);
+                }
+            }
+        });
+    }
+
+    /**
+     * If a trial is finished, release the connection resource
+     * @param trial remote machine trial job detail
+     */
+    public releaseTrialResource(environment: EnvironmentInformation): void {
+        const executorManager = this.environmentExecutorManagerMap.get(environment.id);
+        if (executorManager === undefined) {
+            throw new Error(`ExecutorManager is not assigned for environment ${environment.id}`);
+        }
+        if (this.machineScheduler === undefined) {
+            throw new Error(`machineScheduler is not initialized!`);
+        }
+        // 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.machineScheduler.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!");
+            }
+            this.machineScheduler = new MachineScheduler(Array.from(this.machineExecutorManagerMap.keys()));
+        }
+        const remoteEnvironment: RemoteMachineEnvironmentInformation = environment as RemoteMachineEnvironmentInformation;
+        remoteEnvironment.status = 'WAITING';
+        await this.prepareEnvironment(remoteEnvironment);
+        await this.launchEnvironmentOnScheduledMachine(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');
+        }
+        if (this.machineScheduler === undefined) {
+            throw new Error('machineScheduler is not initialized');
+        }
+
+        // get an executor from scheduler
+        const rmMachineMeta: RemoteMachineMeta | undefined = this.machineScheduler.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.launchEnvironmentOnScheduledMachine(environment);
+
+            environment.status = 'RUNNING';
+            environment.trackingUrl = `file://${rmMachineMeta.ip}:${environment.runnerWorkingFolder}`;
+            this.environmentJobsMap.set(environment.id, environment);
+            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.releaseTrialResource(environment);
+            return
+        }
+
+        const jobpidPath: string = `${environment.runnerWorkingFolder}/pid`;
+        try {
+            await executor.killChildProcesses(jobpidPath);
+            this.releaseTrialResource(environment);
+        } catch (error) {
+            this.log.error(`remoteTrainingService.cancelTrialJob: ${error}`);
+        }
+    }
+
+    private async launchEnvironmentOnScheduledMachine(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);
+    }
+}