Easily SSH into your cluster

pcluster-ssh.md

ParallelCluster Easy SSH

All credit to @tpbrown for this solution.

Usage:

ssh clustername

Setup

1. Create a file easy-ssh.py:

#!/usr/bin/env python3
import argparse
import os
import pcluster.lib as pc
import boto3

def get_user(info):
    instance_id = info.get('headNode').get('instanceId')
    ec2 = boto3.client('ec2')
    ami_id = ec2.describe_instances(InstanceIds=[instance_id])['Reservations'][0]['Instances'][0]['ImageId']
    description = ec2.describe_images(ImageIds=[ami_id])['Images'][0]['Description']
    if 'alinux2' in description :
        user = 'ec2-user'
    elif 'centos' in description:
        user = 'centos'
    elif 'ubuntu' in description:
        user = 'ubuntu'
    return user

def add(name):
    home = os.getenv("HOME")
    region = os.getenv("AWS_DEFAULT_REGION")
    info = pc.describe_cluster(cluster_name=name, region=region)
    user = get_user(info)
    ip = info.get('headNode').get('publicIpAddress')
    with open(f"{home}/.ssh/configs/{name}", "w") as file:
        file.write(f"Host {name}\n  Hostname {ip}\n")
        print(f"Host {name}\n  Hostname {ip}\n  User {user}\n")

def remove(name):
    home = os.getenv("HOME")
    os.remove(f"{home}/.ssh/configs/{name}")

def all():
   home = os.getenv("HOME")
   region = os.getenv("AWS_DEFAULT_REGION")
   print(f"Adding all clusters in {region} to {home}/.ssh/configs")
   clusters = pc.list_clusters()
   for cluster in clusters.get('clusters'):
        cluster_name = cluster.get('clusterName')
        if not os.path.exists(f"{home}/.ssh/configs/{cluster_name}"):
            print(f"Adding ... {cluster_name}")
            add(cluster_name)
        else:
            print(f"Skipping ... {cluster_name} ... file exists.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--all", "--all", action="store_true")
    parser.add_argument("-a", "--add", action="store")
    parser.add_argument("-r", "--remove", action="store")
    parser.add_argument("--name")
    args = parser.parse_args()
    if args.all:
        all()
    elif args.add:
        add(args.name)
    elif args.remove:
        remove(args.name)
    else:
        parser.print_help()

2. Add the following line to your ~/.ssh/config:

mkdir -p ~/.ssh/configs
echo "Include configs/*" >> ~/.ssh/config

3. Execute the script:

python easy-ssh.py --all

Test

Now you'll be able to ssh by only specifying ssh [clustername]. It also supports tab to complete.