Update documentation for version 2.0

Author: ivan-khvostishkov

Flows.md

@@ -0,0 +1,18 @@
+## Flow Diagrams
+
+- [Training job with SSM](#training-job-with-ssm)
+- [Training job with SSH](#training-job-with-ssh)
+
+#### Training job with SSM
+
+This flow corresponds to the [Connecting to SageMaker training jobs with SSM](README.md#training) procedure.
+
+![Screenshot](images/Flow_Train_SSM.png)
+
+
+#### Training job with SSH
+
+This flow corresponds to the [Remote code execution with PyCharm / VSCode over SSH](README.md#remote-interpreter) procedure.
+
+![Screenshot](images/Flow_Train_SSH.png)
+

Flows_Diagrams.md

@@ -8,28 +8,27 @@ SageMaker SSH Helper is an "army-knife" library that helps you to securely conne
 realtime inference endpoints, and SageMaker Studio notebook containers for fast interactive experimentation, 
 remote debugging, and advanced troubleshooting.
 
-The three most common scenarios for the library, also known as "SSH into SageMaker", are:
+Three most common asks that motivated to create the library, sometimes referred as "SSH into SageMaker", are:
 1. A terminal session into a container running in SageMaker to diagnose a stuck training job, use CLI commands 
 like nvidia-smi, or iteratively fix and re-execute your training script within seconds. 
 2. Remote debugging of a code running in SageMaker from your local favorite IDE like 
 PyCharm Professional Edition or Visual Studio Code.
 3. Port forwarding to access diagnostic tools running inside SageMaker, e.g., Dask dashboard, TensorBoard or Spark Web UI.
 
-Other scenarios include but not limited to connecting to a remote Jupyter Notebook in SageMaker Studio from your IDE, or start a VNC session to SageMaker Studio to run GUI apps.  
-
-Also see our [Frequently Asked Questions](FAQ.md), especially if you're using Windows on your local machine.
+Other asks include but not limited to connecting to a remote Jupyter Notebook in SageMaker Studio from your IDE, or start a VNC session to SageMaker Studio to run GUI apps.  
 
 ## How it works
-SageMaker SSH helper uses AWS Systems Manager (SSM) Session Manager, to register the SageMaker container in SSM, followed by creating an SSM session between your client machine and the SageMaker container. Then you can "SSH into SageMaker" by creating an SSH (Secure Shell) connection on top of the SSM session, that allows opening a Linux shell, and/or configuring bidirectional SSH port forwarding to enable applications like remote development/debugging/desktop, and others.
+SageMaker SSH helper uses AWS Systems Manager (SSM) Session Manager, to register the SageMaker container in SSM, followed by creating an SSM session between your client machine and the SageMaker container. Then you can "SSH into SageMaker" by creating an SSH (Secure Shell) connection on top of the SSM session, that allows opening a Linux shell and configuring bidirectional SSH port forwarding to run applications like remote development, debugging, desktop GUI, and others.
+
+![Screenshot](images/high-level-architecture.png)
 
-![Screenshot](images/layers.png)
+Once you become familiar with the library, check the [Flow Diagrams](Flows_Diagrams.md) of the common use cases. 
 
-See detailed architecture diagrams of the complete flow of participating components 
-in [Training Diagram](Flows.md), and [IDE integration with SageMaker Studio diagram](Flows_IDE.md).
+Also make sure you looked at our [Frequently Asked Questions](FAQ.md).
 
 ## Getting started
 
-To get started, your AWS system administrator must set up needed IAM and SSM configuration in your AWS account as shown 
+To get started, your AWS system administrator must configure IAM and SSM in your AWS account as shown 
 in [Setting up your AWS account with IAM and SSM configuration](IAM_SSM_Setup.md).
 
 > **Note**: This solution is a sample AWS content. You should not use this content in your production accounts, in a production 

Flows_IDE.md

@@ -0,0 +1,167 @@
+See https://pdf.plantuml.net/PlantUML_Language_Reference_Guide_en.pdf
+
+@startuml
+actor Developer as dev
+participant "SageMaker SSH Helper \n library (local)" as sm_ssh_helper_local
+participant "Amazon SageMaker" as sagemaker
+participant "VPC \n private or managed" as vpc
+participant "SageMaker SSH Helper \n library (remote)" as sm_ssh_helper_remote
+participant "SSH server" as ssh
+participant "SSM agent" as ssm_agent
+participant "AWS Systems Manager \n (SSM)" as ssm
+participant "IAM" as iam
+
+dev -> sm_ssh_helper_local: Get library dir
+
+note right of dev
+.../sagemaker_ssh_helper/
+end note
+
+return
+
+
+note left of dev
+estimator:
+    source_dir/
+    train.py
+
+    .../sagemaker_ssh_helper/
+end note
+
+dev -> sm_ssh_helper_local: Create wrapper \n around estimator
+activate sm_ssh_helper_local
+sm_ssh_helper_local -> sm_ssh_helper_local: Modify estimator \n metadata
+return
+deactivate sm_ssh_helper_local
+
+dev -> sagemaker: Create training job \n (fit estimator)
+note left of dev
+sourcedir.tar.gz
+end note
+
+sagemaker -> vpc: Start containers
+note left vpc
+docker run train
+end note
+
+activate vpc
+
+vpc -> vpc: Start training
+note left vpc
+train.py
+end note
+activate vpc
+
+vpc -> sm_ssh_helper_remote: Setup and start SSH
+note left sm_ssh_helper_remote
+sm-setup-ssh
+end note
+activate vpc
+
+activate sm_ssh_helper_remote
+
+sm_ssh_helper_remote -> sm_ssh_helper_remote: Configure and \n install libs
+sm_ssh_helper_remote -> sm_ssh_helper_remote: Save environment \n for remote shell
+note left sm_ssh_helper_remote
+sm-save-env
+end note
+
+sm_ssh_helper_remote -> ssh: Start SSH server
+activate ssh
+
+sm_ssh_helper_remote -> sm_ssh_helper_remote: Initialize SSM
+activate sm_ssh_helper_remote
+note left sm_ssh_helper_remote
+sm-init-ssm
+end note
+sm_ssh_helper_remote -> ssm: Create activation
+sm_ssh_helper_remote -> ssm_agent: Register
+ssm_agent -> ssm:
+note right ssm_agent
+mi-01234567890abcdef
+end note
+ssm --> ssm_agent:
+deactivate sm_ssh_helper_remote
+
+sm_ssh_helper_remote -> ssm_agent: Start SSM agent
+activate ssm_agent
+ssm_agent -> ssm: Go online
+
+
+sm_ssh_helper_remote -> sm_ssh_helper_remote: Start waiting \n loop
+activate sm_ssh_helper_remote
+note left sm_ssh_helper_remote
+sm-wait
+end note
+
+note right dev
+sm-local-ssh-training connect <training_job_name>
+end note
+dev -> sm_ssh_helper_local: Connect
+sm_ssh_helper_local -> sm_ssh_helper_local: Get instance IDs
+sm_ssh_helper_local -> ssm: List and filter instances
+sm_ssh_helper_local -> sm_ssh_helper_local: Repeat until successful \n or timeout
+
+note right sm_ssh_helper_local
+mi-01234567890abcdef, ...
+end note
+
+activate sm_ssh_helper_local
+note right sm_ssh_helper_local
+sm-connect-ssh-proxy mi-01234567890abcdef
+end note
+sm_ssh_helper_local -> sm_ssh_helper_local: Generate SSH key pair
+sm_ssh_helper_local -> ssm: Copy SSH public key through S3
+
+ssm -> iam: Check SendCommand \n permissions
+ssm -> ssm_agent: Run command
+ssm_agent -> ssm_agent: Copy key from S3
+sm_ssh_helper_local -> ssm: Start SSH session proxy \n over SSM
+
+ssm -> iam: Check StartSession \n permissions
+
+ssm -> ssm_agent: Start SSH session
+ssm_agent -> ssm_agent: Start SSH proxy tunnel
+ssm_agent --> sm_ssh_helper_local:
+sm_ssh_helper_local -> ssm_agent: Start SSH port forwarding \n over SSM proxy tunnel
+ssm_agent -> ssh: Start SSH proxy \n tunnel session
+activate ssh
+note right dev
+ssh sagemaker-training
+end note
+dev -> ssh: Connect with SSH through forwarded SSH port
+ssh --> dev:
+
+dev -> vpc: Run commands inside a container \n (before training)
+
+note right dev
+sm-local-ssh-training stop-waiting
+end note
+
+dev -> vpc: Stop waiting
+note right vpc
+sm-wait stop
+end note
+deactivate sm_ssh_helper_remote
+
+vpc -> vpc: Training begins
+deactivate vpc
+
+deactivate sm_ssh_helper_remote
+
+dev -> vpc: Run commands inside a container \n (during training)
+
+...Training is in progress...
+dev -> vpc: Stop training or wait until finished
+
+deactivate ssh
+
+deactivate ssh
+deactivate ssm_agent
+deactivate vpc
+
+vpc --> sagemaker: Job is finished
+deactivate vpc
+deactivate sm_ssh_helper_local
+
+@enduml