def __manage_keys(config): """ Creates and pushes SSH keys when necessary """ logger.info('Managing SSH keys') nodes_info = config_utils.get_nodes_ip_name_type(config) for hostname, ip, node_type in nodes_info: ssh_client = ssh_utils.ssh_client(ip, 'root') if not ssh_client: logger.debug('Creating and injecting key to %s', ip) password = config_utils.get_node_password(config, hostname) ansible_utils.apply_playbook(consts.MANAGE_KEYS, variables={ 'ip': ip, 'password': password }) else: logger.debug('Key already exists') docker_repo = config_utils.get_docker_repo(config) if docker_repo and isinstance(docker_repo, dict): ip = docker_repo[consts.IP_KEY] ssh_client = ssh_utils.ssh_client(ip, 'root') if not ssh_client: logger.debug('Creating and injecting key to %s', ip) password = docker_repo[consts.PASSWORD_KEY] ansible_utils.apply_playbook(consts.MANAGE_KEYS, variables={ 'ip': ip, 'password': password }) else: logger.debug('Key already exists')
def test_ssh_client_minimal(self, m1): """ Initial test to ensure main code path does not have any syntax or import errors when calling with the minimal parameters :return: """ self.assertIsNotNone(m1) client = ssh_utils.ssh_client('foo', 'user') self.assertIsNotNone(client)
def __block_until_complete(boot_conf, timeout, suspend=0): """ Function that blocks until all nodes have SSH ports opened :param boot_conf: boot configuration dict :param timeout: boot configuration dict :param suspend: the number of seconds to wait before polling :return: """ host_ips = config_utils.get_node_ips_from_config(boot_conf) host_ip_status = dict() for host_ip in host_ips: host_ip_status[host_ip] = False if suspend > 0: logger.info('Waiting %s seconds before polling IPs %s for SSH', suspend, host_ips) time.sleep(suspend) user = config_utils.get_node_user(boot_conf) password = config_utils.get_node_pass(boot_conf) logger.info('Checking nodes for SSH on %s, user - [%s], pass - [%s]', host_ips, user, password) all_completed = True start = time.time() while timeout > time.time() - start: all_completed = True for host_ip in host_ips: if not host_ip_status[host_ip]: logger.debug( 'Attempting to obtain ssh client - IP [%s], user - [%s],' ' pass - [%s]', host_ip, user, password) ssh_client = ssh_utils.ssh_client(host_ip, user, password=password) if ssh_client: logger.info('Obtained ssh client to IP [%s]', host_ip) if __drp_boot_complete(ssh_client): host_ip_status[host_ip] = True for host_ip, status in host_ip_status.items(): if not status: all_completed = False continue if all_completed: break time.sleep(10) if not all_completed: logger.error('Timeout connecting to all nodes - %s', host_ips) raise Exception('Timeout waiting for nodes to finish booting') logger.info('Connected to all nodes')
def get_gpu_nodes(node_ips): """ Returns nodes which has GPU :return list of GPU nodes """ gpuNodes = list() for ip in node_ips: ssh_cli = ssh_utils.ssh_client(ip, 'root', password='') if ssh_cli: stdin, stdout, stderr = ssh_cli.exec_command( 'lspci | grep -i nvidia | wc -l') if (stdout.read()).rstrip() != '0': stdin, stdout, stderr = ssh_cli.exec_command('hostname') gpuNodes.append((stdout.read()).rstrip()) return gpuNodes
def __run(parsed_args): """ Uses ansible_utils for applying Ansible Playbooks to machines with a private key """ logging.basicConfig(level=logging.DEBUG) logger.info('Starting Playbook Runner') proxy_settings = None if parsed_args.http_proxy: tokens = re.split(':', parsed_args.http_proxy) proxy_settings = ProxySettings(host=tokens[0], port=tokens[1], ssh_proxy_cmd=parsed_args.ssh_proxy_cmd) # Ensure can get an SSH client ssh = ssh_utils.ssh_client(parsed_args.ip_addr, parsed_args.host_user, private_key_filepath=parsed_args.priv_key, proxy_settings=proxy_settings) if ssh: ssh.close() env = Environment(loader=FileSystemLoader( searchpath=os.path.dirname(parsed_args.env_file))) template = env.get_template(os.path.basename(parsed_args.env_file)) env_dict = dict() if parsed_args.vars: env_dict = ast.literal_eval(parsed_args.vars) output = template.render(**env_dict) variables = yaml.load(output) if not variables.get('env_file'): variables['env_file'] = parsed_args.env_file ansible_utils.apply_playbook(parsed_args.playbook, [parsed_args.ip_addr], parsed_args.host_user, ssh_priv_key_file_path=parsed_args.priv_key, password=parsed_args.password, variables=variables, proxy_setting=proxy_settings)