diff --git a/docs/en_US/Tutorial/Nnictl.md b/docs/en_US/Tutorial/Nnictl.md
index ed5c9761e1..81caf6f047 100644
--- a/docs/en_US/Tutorial/Nnictl.md
+++ b/docs/en_US/Tutorial/Nnictl.md
@@ -444,9 +444,6 @@ Debug mode will disable version check function in Trialkeeper.
|--all| False| |delete all of experiments|
-
-
-
* __nnictl experiment export__
* Description
@@ -531,6 +528,62 @@ Debug mode will disable version check function in Trialkeeper.
nnictl experiment import [experiment_id] -f experiment_data.json
```
+* __nnictl experiment save__
+ * Description
+
+ Save nni experiment metadata and code data.
+
+ * Usage
+
+ ```bash
+ nnictl experiment save [OPTIONS]
+ ```
+
+ * Options
+
+ |Name, shorthand|Required|Default|Description|
+ |------|------|------ |------|
+ |id| True| |The id of the experiment you want to save|
+ |--path, -p| False| |the folder path to store nni experiment data, default current working directory|
+ |--saveCodeDir, -s| False| |save codeDir data of the experiment, default False|
+
+ * Examples
+
+ > save an expeirment
+
+ ```bash
+ nnictl experiment save [experiment_id] --saveCodeDir
+ ```
+
+* __nnictl experiment load__
+ * Description
+
+ Load an nni experiment.
+
+ * Usage
+
+ ```bash
+ nnictl experiment load [OPTIONS]
+ ```
+
+ * Options
+
+ |Name, shorthand|Required|Default|Description|
+ |------|------|------ |------|
+ |--path, -p| True| |the file path of nni package|
+ |--codeDir, -c| True| |the path of codeDir for loaded experiment, this path will also put the code in the loaded experiment package|
+ |--logDir, -l| False| |the path of logDir for loaded experiment|
+
+ * Examples
+
+ > load an expeirment
+
+ ```bash
+ nnictl experiment load --path [path] --codeDir [codeDir]
+ ```
+
+
+
### Manage platform information
diff --git a/tools/nni_cmd/common_utils.py b/tools/nni_cmd/common_utils.py
index 4166bf034c..2edbf667df 100644
--- a/tools/nni_cmd/common_utils.py
+++ b/tools/nni_cmd/common_utils.py
@@ -4,7 +4,10 @@
import os
import sys
import json
+import tempfile
import socket
+import string
+import random
import ruamel.yaml as yaml
import psutil
from colorama import Fore
@@ -83,3 +86,12 @@ def check_tensorboard_version():
print_error('import tensorboard error!')
exit(1)
+def generate_temp_dir():
+ '''generate a temp folder'''
+ def generate_folder_name():
+ return os.path.join(tempfile.gettempdir(), 'nni', ''.join(random.sample(string.ascii_letters + string.digits, 8)))
+ temp_dir = generate_folder_name()
+ while os.path.exists(temp_dir):
+ temp_dir = generate_folder_name()
+ os.makedirs(temp_dir)
+ return temp_dir
diff --git a/tools/nni_cmd/config_utils.py b/tools/nni_cmd/config_utils.py
index 8cc1dc8ada..e6472ee3de 100644
--- a/tools/nni_cmd/config_utils.py
+++ b/tools/nni_cmd/config_utils.py
@@ -54,13 +54,13 @@ def __init__(self):
self.experiment_file = os.path.join(NNICTL_HOME_DIR, '.experiment')
self.experiments = self.read_file()
- def add_experiment(self, expId, port, time, file_name, platform, experiment_name):
+ def add_experiment(self, expId, port, startTime, file_name, platform, experiment_name, endTime='N/A', status='INITIALIZED'):
'''set {key:value} paris to self.experiment'''
self.experiments[expId] = {}
self.experiments[expId]['port'] = port
- self.experiments[expId]['startTime'] = time
- self.experiments[expId]['endTime'] = 'N/A'
- self.experiments[expId]['status'] = 'INITIALIZED'
+ self.experiments[expId]['startTime'] = startTime
+ self.experiments[expId]['endTime'] = endTime
+ self.experiments[expId]['status'] = status
self.experiments[expId]['fileName'] = file_name
self.experiments[expId]['platform'] = platform
self.experiments[expId]['experimentName'] = experiment_name
diff --git a/tools/nni_cmd/constants.py b/tools/nni_cmd/constants.py
index 5a37c3a1f1..0654473ed4 100644
--- a/tools/nni_cmd/constants.py
+++ b/tools/nni_cmd/constants.py
@@ -6,6 +6,8 @@
NNICTL_HOME_DIR = os.path.join(os.path.expanduser('~'), '.local', 'nnictl')
+NNI_HOME_DIR = os.path.join(os.path.expanduser('~'), 'nni-experiments')
+
ERROR_INFO = 'ERROR: '
NORMAL_INFO = 'INFO: '
WARNING_INFO = 'WARNING: '
diff --git a/tools/nni_cmd/nnictl.py b/tools/nni_cmd/nnictl.py
index 6a2991fe50..07afd85fab 100644
--- a/tools/nni_cmd/nnictl.py
+++ b/tools/nni_cmd/nnictl.py
@@ -11,7 +11,8 @@
from .nnictl_utils import stop_experiment, trial_ls, trial_kill, list_experiment, experiment_status,\
log_trial, experiment_clean, platform_clean, experiment_list, \
monitor_experiment, export_trials_data, trial_codegen, webui_url, \
- get_config, log_stdout, log_stderr, search_space_auto_gen, webui_nas
+ get_config, log_stdout, log_stderr, search_space_auto_gen, webui_nas, \
+ save_experiment, load_experiment
from .package_management import package_install, package_uninstall, package_show, package_list
from .constants import DEFAULT_REST_PORT
from .tensorboard_utils import start_tensorboard, stop_tensorboard
@@ -129,15 +130,6 @@ def parse_args():
parser_experiment_clean.add_argument('id', nargs='?', help='the id of experiment')
parser_experiment_clean.add_argument('--all', action='store_true', default=False, help='delete all of experiments')
parser_experiment_clean.set_defaults(func=experiment_clean)
-
- #parse experiment command
- parser_platform = subparsers.add_parser('platform', help='get platform information')
- #add subparsers for parser_experiment
- parser_platform_subparsers = parser_platform.add_subparsers()
- parser_platform_clean = parser_platform_subparsers.add_parser('clean', help='clean up the platform data')
- parser_platform_clean.add_argument('--config', '-c', required=True, dest='config', help='the path of yaml config file')
- parser_platform_clean.set_defaults(func=platform_clean)
-
#import tuning data
parser_import_data = parser_experiment_subparsers.add_parser('import', help='import additional data')
parser_import_data.add_argument('id', nargs='?', help='the id of experiment')
@@ -149,6 +141,29 @@ def parse_args():
parser_trial_export.add_argument('--type', '-t', choices=['json', 'csv'], required=True, dest='type', help='target file type')
parser_trial_export.add_argument('--filename', '-f', required=True, dest='path', help='target file path')
parser_trial_export.set_defaults(func=export_trials_data)
+ #save an NNI experiment
+ parser_save_experiment = parser_experiment_subparsers.add_parser('save', help='save an experiment')
+ parser_save_experiment.add_argument('id', nargs='?', help='the id of experiment')
+ parser_save_experiment.add_argument('--path', '-p', required=False, help='the folder path to store nni experiment data, \
+ default current working directory')
+ parser_save_experiment.add_argument('--saveCodeDir', '-s', action='store_true', default=False, help='save codeDir data \
+ of the experiment')
+ parser_save_experiment.set_defaults(func=save_experiment)
+ #load an NNI experiment
+ parser_load_experiment = parser_experiment_subparsers.add_parser('load', help='load an experiment')
+ parser_load_experiment.add_argument('--path', '-p', required=True, help='the path of nni package file')
+ parser_load_experiment.add_argument('--codeDir', '-c', required=True, help='the path of codeDir for loaded experiment, \
+ this path will also put the code in the loaded experiment package')
+ parser_load_experiment.add_argument('--logDir', '-l', required=False, help='the path of logDir for loaded experiment')
+ parser_load_experiment.set_defaults(func=load_experiment)
+
+ #parse platform command
+ parser_platform = subparsers.add_parser('platform', help='get platform information')
+ #add subparsers for parser_platform
+ parser_platform_subparsers = parser_platform.add_subparsers()
+ parser_platform_clean = parser_platform_subparsers.add_parser('clean', help='clean up the platform data')
+ parser_platform_clean.add_argument('--config', '-c', required=True, dest='config', help='the path of yaml config file')
+ parser_platform_clean.set_defaults(func=platform_clean)
#TODO:finish webui function
#parse board command
diff --git a/tools/nni_cmd/nnictl_utils.py b/tools/nni_cmd/nnictl_utils.py
index bbbf54fcc6..b411cfda77 100644
--- a/tools/nni_cmd/nnictl_utils.py
+++ b/tools/nni_cmd/nnictl_utils.py
@@ -18,9 +18,9 @@
from .rest_utils import rest_get, rest_delete, check_rest_server_quick, check_response
from .url_utils import trial_jobs_url, experiment_url, trial_job_id_url, export_data_url
from .config_utils import Config, Experiments
-from .constants import NNICTL_HOME_DIR, EXPERIMENT_INFORMATION_FORMAT, EXPERIMENT_DETAIL_FORMAT, \
+from .constants import NNICTL_HOME_DIR, NNI_HOME_DIR, EXPERIMENT_INFORMATION_FORMAT, EXPERIMENT_DETAIL_FORMAT, \
EXPERIMENT_MONITOR_INFO, TRIAL_MONITOR_HEAD, TRIAL_MONITOR_CONTENT, TRIAL_MONITOR_TAIL, REST_TIME_OUT
-from .common_utils import print_normal, print_error, print_warning, detect_process, get_yml_content
+from .common_utils import print_normal, print_error, print_warning, detect_process, get_yml_content, generate_temp_dir
from .command_utils import check_output_command, kill_command
from .ssh_utils import create_ssh_sftp_client, remove_remote_directory
@@ -736,3 +736,165 @@ def search_space_auto_gen(args):
print_warning('Expected search space file \'{}\' generated, but not found.'.format(file_path))
else:
print_normal('Generate search space done: \'{}\'.'.format(file_path))
+
+def save_experiment(args):
+ '''save experiment data to a zip file'''
+ experiment_config = Experiments()
+ experiment_dict = experiment_config.get_all_experiments()
+ if args.id is None:
+ print_error('Please set experiment id.')
+ exit(1)
+ if args.id not in experiment_dict:
+ print_error('Cannot find experiment {0}.'.format(args.id))
+ exit(1)
+ if experiment_dict[args.id].get('status') != 'STOPPED':
+ print_error('Can only save stopped experiment!')
+ exit(1)
+ print_normal('Saving...')
+ nni_config = Config(experiment_dict[args.id]['fileName'])
+ logDir = os.path.join(NNI_HOME_DIR, args.id)
+ if nni_config.get_config('logDir'):
+ logDir = os.path.join(nni_config.get_config('logDir'), args.id)
+ temp_root_dir = generate_temp_dir()
+
+ # Step1. Copy logDir to temp folder
+ if not os.path.exists(logDir):
+ print_error('logDir: %s does not exist!' % logDir)
+ exit(1)
+ temp_experiment_dir = os.path.join(temp_root_dir, 'experiment')
+ shutil.copytree(logDir, temp_experiment_dir)
+
+ # Step2. Copy nnictl metadata to temp folder
+ temp_nnictl_dir = os.path.join(temp_root_dir, 'nnictl')
+ os.makedirs(temp_nnictl_dir, exist_ok=True)
+ try:
+ with open(os.path.join(temp_nnictl_dir, '.experiment'), 'w') as file:
+ experiment_dict[args.id]['id'] = args.id
+ json.dump(experiment_dict[args.id], file)
+ except IOError:
+ print_error('Write file to %s failed!' % os.path.join(temp_nnictl_dir, '.experiment'))
+ exit(1)
+ nnictl_config_dir = os.path.join(NNICTL_HOME_DIR, experiment_dict[args.id]['fileName'])
+ shutil.copytree(nnictl_config_dir, os.path.join(temp_nnictl_dir, experiment_dict[args.id]['fileName']))
+
+ # Step3. Copy code dir
+ if args.saveCodeDir:
+ temp_code_dir = os.path.join(temp_root_dir, 'code')
+ shutil.copytree(nni_config.get_config('experimentConfig')['trial']['codeDir'], temp_code_dir)
+
+ # Step4. Archive folder
+ zip_package_name = 'nni_experiment_%s' % args.id
+ if args.path:
+ os.makedirs(args.path, exist_ok=True)
+ zip_package_name = os.path.join(args.path, zip_package_name)
+ shutil.make_archive(zip_package_name, 'zip', temp_root_dir)
+ print_normal('Save to %s.zip success!' % zip_package_name)
+
+ # Step5. Cleanup temp data
+ shutil.rmtree(temp_root_dir)
+
+def load_experiment(args):
+ '''load experiment data'''
+ package_path = os.path.expanduser(args.path)
+ if not os.path.exists(args.path):
+ print_error('file path %s does not exist!' % args.path)
+ exit(1)
+ temp_root_dir = generate_temp_dir()
+ shutil.unpack_archive(package_path, temp_root_dir)
+ print_normal('Loading...')
+ # Step1. Validation
+ if not os.path.exists(args.codeDir):
+ print_error('Invalid: codeDir path does not exist!')
+ exit(1)
+ if args.logDir:
+ if not os.path.exists(args.logDir):
+ print_error('Invalid: logDir path does not exist!')
+ exit(1)
+ experiment_temp_dir = os.path.join(temp_root_dir, 'experiment')
+ if not os.path.exists(os.path.join(experiment_temp_dir, 'db')):
+ print_error('Invalid archive file: db file does not exist!')
+ shutil.rmtree(temp_root_dir)
+ exit(1)
+ nnictl_temp_dir = os.path.join(temp_root_dir, 'nnictl')
+ if not os.path.exists(os.path.join(nnictl_temp_dir, '.experiment')):
+ print_error('Invalid archive file: nnictl metadata file does not exist!')
+ shutil.rmtree(temp_root_dir)
+ exit(1)
+ try:
+ with open(os.path.join(nnictl_temp_dir, '.experiment'), 'r') as file:
+ experiment_metadata = json.load(file)
+ except ValueError as err:
+ print_error('Invalid nnictl metadata file: %s' % err)
+ shutil.rmtree(temp_root_dir)
+ exit(1)
+ experiment_config = Experiments()
+ experiment_dict = experiment_config.get_all_experiments()
+ experiment_id = experiment_metadata.get('id')
+ if experiment_id in experiment_dict:
+ print_error('Invalid: experiment id already exist!')
+ shutil.rmtree(temp_root_dir)
+ exit(1)
+ if not os.path.exists(os.path.join(nnictl_temp_dir, experiment_metadata.get('fileName'))):
+ print_error('Invalid: experiment metadata does not exist!')
+ shutil.rmtree(temp_root_dir)
+ exit(1)
+
+ # Step2. Copy nnictl metadata
+ src_path = os.path.join(nnictl_temp_dir, experiment_metadata.get('fileName'))
+ dest_path = os.path.join(NNICTL_HOME_DIR, experiment_metadata.get('fileName'))
+ if os.path.exists(dest_path):
+ shutil.rmtree(dest_path)
+ shutil.copytree(src_path, dest_path)
+
+ # Step3. Copy experiment data
+ nni_config = Config(experiment_metadata.get('fileName'))
+ nnictl_exp_config = nni_config.get_config('experimentConfig')
+ if args.logDir:
+ logDir = args.logDir
+ nnictl_exp_config['logDir'] = logDir
+ else:
+ if nnictl_exp_config.get('logDir'):
+ logDir = nnictl_exp_config['logDir']
+ else:
+ logDir = NNI_HOME_DIR
+ os.rename(os.path.join(temp_root_dir, 'experiment'), os.path.join(temp_root_dir, experiment_id))
+ src_path = os.path.join(os.path.join(temp_root_dir, experiment_id))
+ dest_path = os.path.join(os.path.join(logDir, experiment_id))
+ if os.path.exists(dest_path):
+ shutil.rmtree(dest_path)
+ shutil.copytree(src_path, dest_path)
+
+ # Step4. Copy code dir
+ codeDir = os.path.expanduser(args.codeDir)
+ if not os.path.isabs(codeDir):
+ codeDir = os.path.join(os.getcwd(), codeDir)
+ print_normal('Expand codeDir to %s' % codeDir)
+ nnictl_exp_config['trial']['codeDir'] = codeDir
+ archive_code_dir = os.path.join(temp_root_dir, 'code')
+ if os.path.exists(archive_code_dir):
+ file_list = os.listdir(archive_code_dir)
+ for file_name in file_list:
+ src_path = os.path.join(archive_code_dir, file_name)
+ target_path = os.path.join(codeDir, file_name)
+ if os.path.exists(target_path):
+ print_error('Copy %s failed, %s exist!' % (file_name, target_path))
+ continue
+ if os.path.isdir(src_path):
+ shutil.copytree(src_path, target_path)
+ else:
+ shutil.copy(src_path, target_path)
+
+ # Step5. Create experiment metadata
+ nni_config.set_config('experimentConfig', nnictl_exp_config)
+ experiment_config.add_experiment(experiment_id,
+ experiment_metadata.get('port'),
+ experiment_metadata.get('startTime'),
+ experiment_metadata.get('fileName'),
+ experiment_metadata.get('platform'),
+ experiment_metadata.get('experimentName'),
+ experiment_metadata.get('endTime'),
+ experiment_metadata.get('status'))
+ print_normal('Load experiment %s succsss!' % experiment_id)
+
+ # Step6. Cleanup temp data
+ shutil.rmtree(temp_root_dir)