From 196916ac34e3521fcf11356469de9722ff74f8c1 Mon Sep 17 00:00:00 2001
From: Sean Smith <seaam@amazon.com>
Date: Sun, 26 May 2024 15:37:38 -0700
Subject: [PATCH] Pcluster easy-ssh.sh (#346)

* Pcluster easy-ssh.sh

Signed-off-by: Sean Smith <seaam@amazon.com>

* Update 1.architectures/2.aws-parallelcluster/easy-ssh.sh

Co-authored-by: Keita Watanabe <mlkeita@amazon.com>

---------

Signed-off-by: Sean Smith <seaam@amazon.com>
Co-authored-by: Keita Watanabe <mlkeita@amazon.com>
---
 .../2.aws-parallelcluster/README.md           | 58 ++++++++++++++++++-
 .../2.aws-parallelcluster/easy-ssh.sh         | 39 +++++++++++++
 2 files changed, 96 insertions(+), 1 deletion(-)
 create mode 100755 1.architectures/2.aws-parallelcluster/easy-ssh.sh

diff --git a/1.architectures/2.aws-parallelcluster/README.md b/1.architectures/2.aws-parallelcluster/README.md
index f3a8648a..f3512b99 100644
--- a/1.architectures/2.aws-parallelcluster/README.md
+++ b/1.architectures/2.aws-parallelcluster/README.md
@@ -28,7 +28,7 @@ pip3 install awscli # install the AWS CLI
 pip3 install aws-parallelcluster # then AWS ParallelCluster
 ```
 
-> **Note**: you can use virtual environments to test different versions of AWS ParallelCluster by setting the version during the installation. For example to use 3.7.1, change the command `pip3 install aws-parallelcluster==3.7.1`.
+> **Note**: you can use virtual environments to test different versions of AWS ParallelCluster by setting the version during the installation. For example to use 3.9.1, change the command `pip3 install aws-parallelcluster==3.9.1`.
 
 ### 2.2. Create your EC2 Keypair (if needed)
 
@@ -56,6 +56,62 @@ aws ec2 create-key-pair --key-name pcluster-workshop-key \
 sudo chmod 600 $KEYPAIR_NAME.pem
 ```
 
+### 2.2 Connect to your cluster
+
+To easily login to your cluster via [AWS Systems Manager](https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-sessions-start.html) we've included a script `easy-ssh.sh` that you can run like so, assuming `ml-cluster` is the name of your cluster:
+
+```bash
+./easy-ssh.sh ml-cluster
+```
+
+You'll need a few pre-requisites for this script:
+* JQ: `brew install jq`
+* aws cli
+* `pcluster` cli
+* [Session Manager Plugin](https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html)
+
+Once you've run the script you'll see the following output:
+
+```
+Instance Id: i-0096542c11ccb02b5
+Os: ubuntu2004
+User: ubuntu
+Add the following to your ~/.ssh/config to easily connect:
+
+cat <<EOF >> ~/.ssh/config
+Host ml-cluster
+  User ubuntu
+  ProxyCommand sh -c "aws ssm start-session --target i-0095542c11ccb02b5 --document-name AWS-StartSSHSession --parameters 'portNumber=%p'"
+EOF
+
+Add your ssh keypair and then you can do:
+
+$ ssh ml-cluster
+
+Connecting to ml-cluster...
+
+Starting session with SessionId: ...
+root@ip-10-0-24-126:~#
+```
+
+1. Add your public key to the file `~/.ssh/authorized_keys`
+
+2. Now paste in the lines from the output of to your terminal, this will add them to your `~/.ssh/config`.
+
+```
+cat <<EOF >> ~/.ssh/config
+Host ml-cluster
+  User ubuntu
+  ProxyCommand sh -c "aws ssm start-session --target i-0095542c11ccb02b5 --document-name AWS-StartSSHSession --parameters 'portNumber=%p'"
+EOF
+```
+3. Now you ssh in, assuming `ml-cluster` is the name of your cluster with:
+
+```
+ssh ml-cluster
+```
+
+
 ## 3. Deploy a Cluster
 
 To create the cluster use the command below and replace `CLUSTER_CONFIG_FILE` by the path to the cluster configuration file (see next section) and `NAME_OF_YOUR_CLUSTER` by the name of your cluster (`realpotato` is a cool name).
diff --git a/1.architectures/2.aws-parallelcluster/easy-ssh.sh b/1.architectures/2.aws-parallelcluster/easy-ssh.sh
new file mode 100755
index 00000000..5c7617df
--- /dev/null
+++ b/1.architectures/2.aws-parallelcluster/easy-ssh.sh
@@ -0,0 +1,39 @@
+#!/bin/bash
+
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+
+
+cluster_name=${1:-ml-cluster}
+
+# grab head node instance id
+instance_id=$(pcluster describe-cluster -n ${cluster_name} | jq '.headNode.instanceId' | tr -d '"')
+os=$(aws ec2 describe-tags --filters "Name=resource-id,Values=$instance_id" "Name=key,Values=parallelcluster:attributes" | jq '.Tags[0].Value' | tr -d '"' | awk -F ',' '{print $1}')
+
+if [ "$os" = 'ubuntu2004' ]; then 
+    user='ubuntu'
+elif [ "$os" = 'ubuntu2204' ]; then 
+    user='ubuntu'
+elif [ "$os" = 'alinux2' ]; then
+    user='ec2-user'
+fi
+
+echo -e "Instance Id: $instance_id"
+echo -e "Os: $os"
+echo -e "User: $user"
+
+echo -e "Add the following to your ~/.ssh/config to easily connect:"
+echo "
+cat <<EOF >> ~/.ssh/config
+Host ${cluster_name}
+  User ${user}
+  ProxyCommand sh -c \"aws ssm start-session --target ${instance_id} --document-name AWS-StartSSHSession --parameters 'portNumber=%p'\"
+EOF
+
+Add your ssh keypair and then you can do:
+
+$ ssh ${cluster_name}
+"
+
+echo "Connecting to $cluster_name..."
+aws ssm start-session --target $instance_id
\ No newline at end of file