Skip to content

Instantly share code, notes, and snippets.

@sean-smith
Last active November 8, 2023 08:02
Show Gist options
  • Save sean-smith/ffd725ec0d0edfba2135d84630e3dc28 to your computer and use it in GitHub Desktop.
Save sean-smith/ffd725ec0d0edfba2135d84630e3dc28 to your computer and use it in GitHub Desktop.
Easily SSH into your cluster

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()
  1. Add the following line to your ~/.ssh/config:
mkdir -p ~/.ssh/configs
echo "Include configs/*" >> ~/.ssh/config
  1. 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment