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)