Support nnictl tensorboard (#268)

support tensorboard

docs/NNICTLDOC.md

@@ -285,3 +285,45 @@ nnictl webui
        | Name, shorthand | Required|Default | Description |
        | ------ | ------ | ------ |------ |
      | id|  False| |ID of the experiment you want to set|
+
+
+### Manage tensorboard
+* __nnictl tensorboard start__
+   * Description
+     
+	     Start the tensorboard process. 
+   
+   * Usage
+         
+		    nnictl tensorboard start
+        
+    	Options:
+    	
+       | Name, shorthand | Required|Default | Description |
+       | ------ | ------ | ------ |------ |
+     | id|  False| |ID of the experiment you want to set|
+     | --trialid|  False| |ID of the trial|
+     | --port|  False| 6006|The port of the tensorboard process|
+
+   * Detail
+     
+	     1. NNICTL support tensorboard function in local and remote platform for the moment, other platforms will be supported later.   
+         2. If you want to use tensorboard, you need to write your tensorboard log data to environment variable [NNI_OUTPUT_DIR] path.  
+         3. In local mode, nnictl will set --logdir=[NNI_OUTPUT_DIR] directly and start a tensorboard process.
+         4. In remote mode, nnictl will create a ssh client to copy log data from remote machine to local temp directory firstly, and then start a tensorboard process in your local machine. You need to notice that nnictl only copy the log data one time when you use the command, if you want to see the later result of tensorboard, you should execute nnictl tensorboard command again.
+         5. If there is only one trial job, you don't need to set trialid. If there are multiple trial jobs running, you should set the trialid, or you could use [nnictl tensorboard start --trialid all] to map --logdir to all trial log paths.
+
+* __nnictl tensorboard stop__
+   * Description
+     
+	     Stop all of the tensorboard process. 
+   
+   * Usage
+         
+		    nnictl tensorboard stop
+        
+    	Options:
+    	
+       | Name, shorthand | Required|Default | Description |
+       | ------ | ------ | ------ |------ |
+     | id|  False| |ID of the experiment you want to set|
\ No newline at end of file

setup.py

@@ -62,7 +62,8 @@ setup(
         'requests',
         'scipy',
         'schema',
-        'pyhdfs'
+        'pyhdfs',
+        'paramiko'
     ],
 
     cmdclass={

tools/nnicmd/launcher.py

@@ -28,11 +28,10 @@ import tempfile
 from nni_annotation import *
 from .launcher_utils import validate_all_content
 from .rest_utils import rest_put, rest_post, check_rest_server, check_rest_server_quick, check_response
-from .url_utils import cluster_metadata_url, experiment_url
+from .url_utils import cluster_metadata_url, experiment_url, get_local_urls
 from .config_utils import Config, Experiments
 from .common_utils import get_yml_content, get_json_content, print_error, print_normal, print_warning, detect_process, detect_port
 from .constants import *
-from .webui_utils import *
 import time
 import random
 import string
@@ -288,7 +287,8 @@ def launch_experiment(args, experiment_config, mode, config_file_name, experimen
         except Exception:
             raise Exception(ERROR_INFO % 'Restful server stopped!')
         exit(1)
-    web_ui_url_list = get_web_ui_urls(args.port, config_file_name)
+    web_ui_url_list = get_local_urls(args.port)
+    nni_config.set_config('webuiUrl', web_ui_url_list)
     
     #save experiment information
     experiment_config = Experiments()

tools/nnicmd/nnictl.py

@@ -25,6 +25,7 @@ from .updater import update_searchspace, update_concurrency, update_duration, up
 from .nnictl_utils import *
 from .package_management import *
 from .constants import *
+from .tensorboard_utils import *
 
 def nni_help_info(*args):
     print('please run "nnictl {positional argument} --help" to see nnictl guidance')
@@ -148,6 +149,18 @@ def parse_args():
     parser_package_show = parser_package_subparsers.add_parser('show', help='show the information of packages')
     parser_package_show.set_defaults(func=package_show)
 
+    #parse tensorboard command
+    parser_tensorboard = subparsers.add_parser('tensorboard', help='manage tensorboard')
+    parser_tensorboard_subparsers = parser_tensorboard.add_subparsers()
+    parser_tensorboard_start = parser_tensorboard_subparsers.add_parser('start', help='start tensorboard')
+    parser_tensorboard_start.add_argument('id', nargs='?', help='the id of experiment')
+    parser_tensorboard_start.add_argument('--trialid', dest='trialid', help='the id of trial')
+    parser_tensorboard_start.add_argument('--port', dest='port', default=6006, help='the port to start tensorboard')
+    parser_tensorboard_start.set_defaults(func=start_tensorboard)
+    parser_tensorboard_start = parser_tensorboard_subparsers.add_parser('stop', help='stop tensorboard')
+    parser_tensorboard_start.add_argument('id', nargs='?', help='the id of experiment')
+    parser_tensorboard_start.set_defaults(func=stop_tensorboard)
+
     args = parser.parse_args()
     args.func(args)
 

tools/nnicmd/nnictl_utils.py

@@ -174,8 +174,17 @@ def stop_experiment(args):
             time.sleep(3)
             rest_pid = nni_config.get_config('restServerPid')
             if rest_pid:
-                cmds = ['pkill', '-P', str(rest_pid)]
+                stop_rest_cmds = ['pkill', '-P', str(rest_pid)]
+                call(stop_rest_cmds)
+                tensorboard_pid_list = nni_config.get_config('tensorboardPidList')
+                if tensorboard_pid_list:
+                    for tensorboard_pid in tensorboard_pid_list:
+                        try:
+                            cmds = ['kill', '-9', str(tensorboard_pid)]
                             call(cmds)
+                        except Exception as exception:
+                            print_error(exception)
+                    nni_config.set_config('tensorboardPidList', [])
             if stop_rest_result:
                 print_normal('Stop experiment success!')
             experiment_config.update_experiment(experiment_id, 'status', 'stopped')
@@ -343,3 +352,4 @@ def experiment_list(args):
         experiment_information += (EXPERIMENT_DETAIL_FORMAT % (key, experiment_dict[key]['status'], \
         experiment_dict[key]['startTime'], experiment_dict[key]['endTime']))
     print(EXPERIMENT_INFORMATION_FORMAT % experiment_information)
+

tools/nnicmd/webui_utils.py → tools/nnicmd/ssh_utils.py

@@ -18,16 +18,32 @@
 # DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
-import psutil
-from socket import AddressFamily
-from .config_utils import Config
+import paramiko
+import os
+from .common_utils import print_error
 
-def get_web_ui_urls(port, CONFIG_FILE_NAME):
-    webui_url_list = []
-    for name, info in psutil.net_if_addrs().items():
-        for addr in info:
-            if AddressFamily.AF_INET == addr.family:
-                webui_url_list.append('http://{}:{}'.format(addr.address, port))
-    nni_config = Config(CONFIG_FILE_NAME)
-    nni_config.set_config('webuiUrl', webui_url_list)
-    return webui_url_list
+def copy_remote_directory_to_local(sftp, remote_path, local_path):
+    '''copy remote directory to local machine'''
+    try:
+        os.makedirs(local_path, exist_ok=True)
+        files = sftp.listdir(remote_path)
+        for file in files:
+            remote_full_path = os.path.join(remote_path, file)
+            local_full_path = os.path.join(local_path, file)
+            try:
+                if sftp.listdir(remote_full_path):
+                    copy_remote_directory_to_local(sftp, remote_full_path, local_full_path)
+            except:
+                sftp.get(remote_full_path, local_full_path)
+    except Exception:
+        pass
+
+def create_ssh_sftp_client(host_ip, port, username, password):
+    '''create ssh client'''
+    try:
+        conn = paramiko.Transport(host_ip, port)
+        conn.connect(username=username, password=password)
+        sftp = paramiko.SFTPClient.from_transport(conn)
+        return sftp
+    except Exception as exception:
+        print_error('Create ssh client error %s\n' % exception)
\ No newline at end of file

tools/nnicmd/tensorboard_utils.py

+# Copyright (c) Microsoft Corporation
+# All rights reserved.
+#
+# MIT License
+#
+# Permission is hereby granted, free of charge,
+# to any person obtaining a copy of this software and associated
+# documentation files (the "Software"), to deal in the Software without restriction,
+# including without limitation the rights to use, copy, modify, merge, publish,
+# distribute, sublicense, and/or sell copies of the Software, and
+# to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+# The above copyright notice and this permission notice shall be included
+# in all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
+# BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+import os
+import psutil
+import json
+import datetime
+import time
+from subprocess import call, check_output, Popen, PIPE
+from .rest_utils import rest_get, rest_delete, check_rest_server_quick, check_response
+from .config_utils import Config, Experiments
+from .url_utils import trial_jobs_url, experiment_url, trial_job_id_url, get_local_urls
+from .constants import NNICTL_HOME_DIR, EXPERIMENT_INFORMATION_FORMAT, EXPERIMENT_DETAIL_FORMAT, COLOR_GREEN_FORMAT
+import time
+from .common_utils import print_normal, print_error, print_warning, detect_process, detect_port
+from .nnictl_utils import *
+import re
+from .ssh_utils import create_ssh_sftp_client, copy_remote_directory_to_local
+import tempfile
+
+def parse_log_path(args, trial_content):
+    '''parse log path'''
+    path_list = []
+    host_list = []
+    for trial in trial_content:
+        if args.trialid and args.trialid != 'all' and trial.get('id') != args.trialid:
+            continue
+        pattern = r'(?P<head>.+)://(?P<host>.+):(?P<path>.*)'
+        match = re.search(pattern,trial['logPath'])
+        if match:
+            path_list.append(match.group('path'))
+            host_list.append(match.group('host'))
+    if not path_list:
+        print_error('Trial id %s error!' % args.trialid)
+        exit(1)
+    return path_list, host_list
+
+def copy_data_from_remote(args, nni_config, trial_content, path_list, host_list, temp_nni_path):
+    '''use ssh client to copy data from remote machine to local machien'''
+    machine_list = nni_config.get_config('experimentConfig').get('machineList')
+    machine_dict = {}
+    local_path_list = []
+    for machine in machine_list:
+        machine_dict[machine['ip']] = {'port': machine['port'], 'passwd': machine['passwd'], 'username': machine['username']}
+    for index, host in enumerate(host_list):
+        local_path = os.path.join(temp_nni_path, trial_content[index].get('id'))
+        local_path_list.append(local_path)
+        print_normal('Copying log data from %s to %s' % (host + ':' + path_list[index], local_path))
+        sftp = create_ssh_sftp_client(host, machine_dict[host]['port'], machine_dict[host]['username'], machine_dict[host]['passwd'])
+        copy_remote_directory_to_local(sftp, path_list[index], local_path)
+    print_normal('Copy done!')
+    return local_path_list
+
+def get_path_list(args, nni_config, trial_content, temp_nni_path):
+    '''get path list according to different platform'''
+    path_list, host_list = parse_log_path(args, trial_content)
+    platform = nni_config.get_config('experimentConfig').get('trainingServicePlatform')
+    if platform == 'local':
+        print_normal('Log path: %s' % ' '.join(path_list))
+        return path_list
+    elif platform == 'remote':
+        path_list = copy_data_from_remote(args, nni_config, trial_content, path_list, host_list, temp_nni_path)
+        print_normal('Log path: %s' % ' '.join(path_list))
+        return path_list
+    else:
+        print_error('Not supported platform!')
+        exit(1)
+
+def format_tensorboard_log_path(path_list):
+    new_path_list = []
+    for index, value in enumerate(path_list):
+        new_path_list.append('name%d:%s' % (index + 1, value))
+    return ','.join(new_path_list)
+
+def start_tensorboard_process(args, nni_config, path_list, temp_nni_path):
+    '''call cmds to start tensorboard process in local machine'''
+    if detect_port(args.port):
+        print_error('Port %s is used by another process, please reset port!' % str(args.port))
+        exit(1)
+    
+    stdout_file = open(os.path.join(temp_nni_path, 'tensorboard_stdout'), 'a+')
+    stderr_file = open(os.path.join(temp_nni_path, 'tensorboard_stderr'), 'a+')
+    cmds = ['tensorboard', '--logdir', format_tensorboard_log_path(path_list), '--port', str(args.port)]
+    tensorboard_process = Popen(cmds, stdout=stdout_file, stderr=stderr_file)
+    url_list = get_local_urls(args.port)
+    print_normal(COLOR_GREEN_FORMAT % 'Start tensorboard success!\n' + 'Tensorboard urls: ' + '     '.join(url_list))
+    tensorboard_process_pid_list = nni_config.get_config('tensorboardPidList')
+    if tensorboard_process_pid_list is None:
+        tensorboard_process_pid_list = [tensorboard_process.pid]
+    else:
+        tensorboard_process_pid_list.append(tensorboard_process.pid)
+    nni_config.set_config('tensorboardPidList', tensorboard_process_pid_list)
+
+def stop_tensorboard(args):
+    '''stop tensorboard'''
+    experiment_id = check_experiment_id(args)
+    experiment_config = Experiments()
+    experiment_dict = experiment_config.get_all_experiments()
+    config_file_name = experiment_dict[experiment_id]['fileName']
+    nni_config = Config(config_file_name)
+    tensorboard_pid_list = nni_config.get_config('tensorboardPidList')
+    if tensorboard_pid_list:
+        for tensorboard_pid in tensorboard_pid_list:
+            try:
+                cmds = ['kill', '-9', str(tensorboard_pid)]
+                call(cmds)
+            except Exception as exception:
+                print_error(exception)
+        nni_config.set_config('tensorboardPidList', [])
+        print_normal('Stop tensorboard success!')
+    else:
+        print_error('No tensorboard configuration!')
+
+
+def start_tensorboard(args):
+    '''start tensorboard'''
+    experiment_id = check_experiment_id(args)
+    experiment_config = Experiments()
+    experiment_dict = experiment_config.get_all_experiments()
+    config_file_name = experiment_dict[experiment_id]['fileName']
+    nni_config = Config(config_file_name)
+    rest_port = nni_config.get_config('restServerPort')
+    rest_pid = nni_config.get_config('restServerPid')
+    if not detect_process(rest_pid):
+        print_error('Experiment is not running...')
+        return
+    running, response = check_rest_server_quick(rest_port)
+    trial_content = None
+    if running:
+        response = rest_get(trial_jobs_url(rest_port), 20)
+        if response and check_response(response):
+            trial_content = json.loads(response.text)
+        else:
+            print_error('List trial failed...')
+    else:
+        print_error('Restful server is not running...')
+    if not trial_content:
+        print_error('No trial information!')
+        exit(1)
+    if len(trial_content) > 1 and not args.trialid:
+        print_error('There are multiple trials, please set trial id!')
+        exit(1)
+    experiment_id = nni_config.get_config('experimentId')
+    temp_nni_path = os.path.join(tempfile.gettempdir(), 'nni', experiment_id)
+    os.makedirs(temp_nni_path, exist_ok=True)
+
+    path_list = get_path_list(args, nni_config, trial_content, temp_nni_path)
+    start_tensorboard_process(args, nni_config, path_list, temp_nni_path)

tools/nnicmd/url_utils.py

@@ -18,6 +18,8 @@
 # DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+import psutil
+from socket import AddressFamily
 
 BASE_URL = 'http://localhost'
 
@@ -53,6 +55,7 @@ def trial_jobs_url(port):
     '''get trial_jobs url'''
     return '{0}:{1}{2}{3}'.format(BASE_URL, port, API_ROOT_URL, TRIAL_JOBS_API)
 
+
 def trial_job_id_url(port, job_id):
     '''get trial_jobs with id url'''
     return '{0}:{1}{2}{3}/:{4}'.format(BASE_URL, port, API_ROOT_URL, TRIAL_JOBS_API, job_id)
@@ -61,3 +64,13 @@ def trial_job_id_url(port, job_id):
 def tensorboard_url(port):
     '''get tensorboard url'''
     return '{0}:{1}{2}{3}'.format(BASE_URL, port, API_ROOT_URL, TENSORBOARD_API)
+
+
+def get_local_urls(port):
+    '''get urls of local machine'''
+    url_list = []
+    for name, info in psutil.net_if_addrs().items():
+        for addr in info:
+            if AddressFamily.AF_INET == addr.family:
+                url_list.append('http://{}:{}'.format(addr.address, port))
+    return url_list
\ No newline at end of file

tools/setup.py

@@ -12,7 +12,8 @@ setuptools.setup(
         'psutil',
         'astor',
         'schema',
-        'pyhdfs'
+        'pyhdfs',
+        'paramiko'
     ],
 
     author = 'Microsoft NNI Team',