support create ssh connection by using sshkey (issue #1950) (#1957)

tools/nni_cmd/nnictl_utils.py

@@ -403,11 +403,13 @@ def remote_clean(machine_list, experiment_id=None):
         userName = machine.get('username')
         host = machine.get('ip')
         port = machine.get('port')
+        sshKeyPath = machine.get('sshKeyPath')
+        passphrase = machine.get('passphrase')
         if experiment_id:
             remote_dir = '/' + '/'.join(['tmp', 'nni', 'experiments', experiment_id])
         else:
             remote_dir = '/' + '/'.join(['tmp', 'nni', 'experiments'])
-        sftp = create_ssh_sftp_client(host, port, userName, passwd)
+        sftp = create_ssh_sftp_client(host, port, userName, passwd, sshKeyPath, passphrase)
         print_normal('removing folder {0}'.format(host + ':' + str(port) + remote_dir))
         remove_remote_directory(sftp, remote_dir)
 

tools/nni_cmd/ssh_utils.py

@@ -30,12 +30,16 @@ def copy_remote_directory_to_local(sftp, remote_path, local_path):
     except Exception:
         pass
 
-def create_ssh_sftp_client(host_ip, port, username, password):
+def create_ssh_sftp_client(host_ip, port, username, password, ssh_key_path, passphrase):
     '''create ssh client'''
     try:
         paramiko = check_environment()
         conn = paramiko.Transport(host_ip, port)
-        conn.connect(username=username, password=password)
+        if ssh_key_path is not None:
+            ssh_key = paramiko.RSAKey.from_private_key_file(ssh_key_path, password=passphrase)
+            conn.connect(username=username, pkey=ssh_key)
+        else:
+            conn.connect(username=username, password=password)
         sftp = paramiko.SFTPClient.from_transport(conn)
         return sftp
     except Exception as exception:

tools/nni_cmd/tensorboard_utils.py

@@ -37,12 +37,14 @@ def copy_data_from_remote(args, nni_config, trial_content, path_list, host_list,
     machine_dict = {}
     local_path_list = []
     for machine in machine_list:
-        machine_dict[machine['ip']] = {'port': machine['port'], 'passwd': machine['passwd'], 'username': machine['username']}
+        machine_dict[machine['ip']] = {'port': machine['port'], 'passwd': machine['passwd'], 'username': machine['username'],
+                                       'sshKeyPath': machine.get('sshKeyPath'), 'passphrase': machine.get('passphrase')}
     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'])
+        sftp = create_ssh_sftp_client(host, machine_dict[host]['port'], machine_dict[host]['username'], machine_dict[host]['passwd'],
+                                      machine_dict[host]['sshKeyPath'], machine_dict[host]['passphrase'])
         copy_remote_directory_to_local(sftp, path_list[index], local_path)
     print_normal('Copy done!')
     return local_path_list