Publishing version 1.7.8

Author: ivan-khvostishkov

.gitlab-ci.yml

@@ -0,0 +1,50 @@
+# Official language image. Look for the different tagged releases at:
+# https://hub.docker.com/r/library/python/tags/
+image: python:3.8
+
+stages:
+  - scan_source_code
+  - test
+
+scan_source_code:
+  stage: scan_source_code
+  script:
+    - pip install bandit
+    - bandit -r ./ --skip B603,B404,B101,B607 2>&1 | tee bandit.txt
+  artifacts:
+    paths:
+      - bandit.txt
+
+run_tests:
+  stage: test
+  resource_group: ssh-helper-tests
+  script:
+    # Smoke test - Python version
+    - python --version
+    # Check all running processes on the agent
+    - hostname
+    - ps xfa
+    # Install AWS CLI
+    - curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "/tmp/awscliv2.zip"
+    - unzip -o -q -d /tmp/ /tmp/awscliv2.zip
+    - /tmp/aws/install
+    - aws --version
+    - aws sts get-caller-identity
+    - echo "AWS default region - $AWS_DEFAULT_REGION"
+    - echo "AWS region - $AWS_REGION"
+    - aws configure list
+    # Install SSM plugin
+    - curl "https://s3.amazonaws.com/session-manager-downloads/plugin/latest/ubuntu_64bit/session-manager-plugin.deb" -o "session-manager-plugin.deb"
+    - dpkg -i session-manager-plugin.deb
+    - session-manager-plugin --version
+    # Install the package
+    - pip install '.[test]'
+    # Run tests
+    - echo "SageMaker role - $SAGEMAKER_ROLE"
+    - echo "Kernel Gateway name - $KERNEL_GATEWAY_NAME"
+    - echo "Extra args for pytest - $PYTEST_EXTRA_ARGS"
+    - cd tests; pytest --junitxml=pytest_report.xml -m 'not manual' -o sagemaker_role=$SAGEMAKER_ROLE -o kernel_gateway_name=$KERNEL_GATEWAY_NAME $PYTEST_EXTRA_ARGS
+  artifacts:
+    when: always
+    reports:
+      junit: tests/pytest_report.xml

.idea/.gitignore

@@ -1,4 +1,4 @@
 ## Code of Conduct
-This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct).
-For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact
-[email protected] with any additional questions or comments.
+This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct). 
+For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact 
+[email protected] with any additional questions or comments.

.idea/inspectionProfiles/profiles_settings.xml

@@ -6,6 +6,38 @@ documentation, we greatly value feedback and contributions from our community.
 Please read through this document before submitting any issues or pull requests to ensure we have all the necessary
 information to effectively respond to your bug report or contribution.
 
+## Adding new features to the SageMaker SSH Helper
+
+SageMaker SSH helper uses Test Driven Development (TDD) methodology to implement its features.
+
+Before implementing new features, check [TODO](TODO) list. The contributors are either working already 
+on these features or planning to implement them in the future.
+
+To start development, install the library on your local machine:
+```shell
+pip install '.[test]'
+```
+
+Make sure all tests are working:
+```shell
+cd tests
+pytest -m 'not manual' \
+  -o sagemaker_role=arn:aws:iam::<<YOUR_ACCOUNT_ID>>:role/service-role/<<YOUR_AmazonSageMaker_ExecutionRole>> \
+  -o kernel_gateway_name=<<YOUR_KERNEL_GATEWAY_NAME>>
+```
+*Tip:* You can pass your parameters either in command line or set it in `tests/pytest.ini`.
+
+Check `tests/test_end_to_end.py` for automation of steps 
+required for training / processing and inference:
+
+1. `test_train_e2e()` - a simple training job
+2. `test_inference_e2e()` - a training job followed by deployment
+3. `test_inference_e2e_mms()` - a training job followed by deployment to a multi-model endpoint
+4. `test_processing_e2e()` - a Spark processing job
+5. `test_processing_framework_e2e()` - a PyTorch framework processing job
+
+Then write a failing test, put code to make it pass, and make sure other tests are still working to avoid any regression.
+
 
 ## Reporting Bugs/Feature Requests
 

.idea/misc.xml

@@ -0,0 +1,29 @@
+## Flow Diagrams
+
+#### Training job
+
+Also compare with the similar flow diagram for [IDE integration with SageMaker Studio](Flows_IDE.md).
+
+The following diagram describes the flow of events for the training job use case:
+
+![Screenshot](images/overall-flow.png)
+
+1. Data Scientists (DS) starts a training job, with SSH Helper Lib as a dependency.
+2. The SageMaker Python SDK starts the job, sending the train.py script, and SSH Helper lib as code dependencies.
+3. Amazon SageMaker control plain starts the host and the container.
+4. The user script (train.py) starts running, starting SSH helper, which fetches AWS SSM agent and other packages 
+from the Internet, and installs them.
+5. SSH Helper starts the SSM agent
+6. Through SSM, SSH Helper registers the container as an SSM `managed instance`, and tags it with the DS AWS user/role name.
+7. SSH Helper printouts the `managed instance` ID. The log is streamed to CloudWatch Logs.
+8. The DS manually/automatically tails the training job logs for the `managed instance` ID.
+
+9-12. Optionally: The DS starts a process to copy over his SSH Public key to the container, needed to set up port forwarding via SSH (e.g., for remote debugging)
+
+13. The DS uses the AWS SSM CLI to start a shell providing the managed instance ID as a parameter.
+Optionally, user starts SSM with SSH port forwarding with the helper command 
+`sm-local-ssh-training connnect <<training_job_name>>`.
+14. AWS SSM IAM rules verify that the user is allowed to take this action and that the instance is tagged with the DS's AWS user/role name. Once verified, a session is created with the SSM Agent running in the container.
+15. The SSM agent generates a shell by spinning off a new bash shell process.
+Optionally, SSH port forwarding starts over SSM connection to let user connect to remote processes over TCP
+in both directions.

.idea/modules.xml

@@ -0,0 +1,28 @@
+## Flow Diagrams - IDE
+
+#### IDE integration with SageMaker Studio
+
+Also compare with the similar flow diagram for [Training Job](Flows.md).
+The major difference is that the flow for IDe starts with the step 3 
+(there's no step 1 and 2 to submit a job).
+
+![Screenshot](images/overall-flow-ide.png)
+
+1. (Not applicable)
+2. (Not applicable)
+3. Amazon SageMaker Studio starts the kernel gateway app.
+4. User starts the notebook [SageMaker_SSH_IDE.ipynb](SageMaker_SSH_IDE.ipynb), starting SSH helper, which fetches 
+AWS SSM agent and other packages from the Internet, and installs them.
+5. SSH Helper starts the SSM agent
+6. Through SSM, SSH Helper registers the container as an SSM `managed instance`, and tags it with the DS AWS user/role name.
+7. SSH Helper printouts the `managed instance` ID. The log is streamed to CloudWatch Logs.
+8. The DS manually/automatically tails the training job logs for the `managed instance` ID.
+
+9-12. Optionally: The DS starts a process to copy over his SSH Public key to the container, needed to set up port forwarding via SSH (e.g., for remote debugging)
+
+13. The DS uses the AWS SSM CLI to start a shell providing the managed instance ID as a parameter.
+Optionally, user starts SSM with SSH port forwarding with the helper command `sm-local-ssh-ide <<kernel_gateway_name>>`.
+14. AWS SSM IAM rules verify that the user is allowed to take this action and that the instance is tagged with the DS's AWS user/role name. Once verified, a session is created with the SSM Agent running in the container.
+15. The SSM agent generates a shell by spinning off a new bash shell process.
+Optionally, SSH port forwarding starts over SSM connection to let user connect to remote processes over TCP
+in both directions.

.idea/sagemaker-ssh-helper.iml

@@ -0,0 +1,135 @@
+## <a name="setup"></a>Setting up your AWS account with IAM and SSM configuration
+
+> **NOTE**: This section involves AWS IAM changes and should be completed by an AWS system admin.
+If you plan to use these settings in production, please, carefully review them with your security team.
+
+SageMaker SSH Helper relies on the AWS Systems Manager service to create SSH tunnels between your client and the 
+SageMaker component. To allow that the following setup is required:  
+
+**1. Update your SageMaker IAM role (used with the `role` Estimator parameter)**
+    
+a. Add trust relationship with SSM, so it will look like this:
+
+```json
+{
+    "Version": "2012-10-17",
+    "Statement": [
+        {
+            "Effect": "Allow",
+            "Principal": {
+                "Service": "sagemaker.amazonaws.com"
+            },
+            "Action": "sts:AssumeRole"
+        },
+        {
+            "Effect": "Allow",
+            "Principal": {
+                "Service": "ssm.amazonaws.com"
+            },
+            "Action": "sts:AssumeRole"
+        }
+    ]
+}
+```
+
+b. Add `AmazonSSMManagedInstanceCore` and `CloudWatchLogsFullAccess` managed policies to the role.
+ 
+c. Add to the role a new inline policy named `SageMakerSSMPolicy` as follows, replacing `<<SAGEMAKER_ROLE_ARN>>` 
+     with the arn of the role you're editing:
+
+```json
+{
+   "Version": "2012-10-17",
+   "Statement": [
+       {
+           "Effect": "Allow",
+           "Action": "iam:PassRole",
+           "Resource": "<<SAGEMAKER_ROLE_ARN>>"
+       },
+       {
+           "Effect": "Allow",
+           "Action": "ssm:CreateActivation",
+           "Resource": "*"
+       },
+       {
+        "Effect": "Allow",
+        "Action": "ssm:AddTagsToResource",
+        "Resource": "<<SAGEMAKER_ROLE_ARN>>"
+       }
+   ]
+}
+```
+
+---
+
+**2. Setup Systems Manager (SSM)**
+
+> Note: These steps are needed if you haven't setup SSM before. 
+> Pay attention to the step "2h. Enable advanced instances tier".
+   
+a. Create a new minimal EC2 instance, with amazon linux 2 AMI (you need it for successful setup and for verification).
+
+b. Navigate in AWS Console to Systems Manager > Quick Setup and press 'Get started' button
+   
+c. From Configuration types select 'Host Management' 
+
+d. If you have enabled AWS Organizations before, choose 'Current account'
+
+e. Leave everything else as default and choose 'Create'. It will take several minutes to complete. 
+   Make sure "Configuration association status" shows all associations are successful and not failing or pending.
+
+f. (Optional) Verify that Systems Manager is successfully set up by connecting to the previously created EC2 instance:
+   Systems Manager > Sessions Manager > Start session. In the [Target instances] table filter using the instance id (e.g., Instance ID: "i-0b900fbc0fe61e9e8") check it, and click start session. Verify you we're able to start a session.
+
+g. Terminate the created EC2 instance (not needed anymore).
+
+h. Enable advanced instances tier to use managed instances (such as SageMaker managed containers): 
+   Systems Manager > Fleet Manager > Account management > Instance tier settings > Change account setting > Confirm change from Standard-Tier to Advanced-Tier
+
+---
+
+**3. Make sure that the local user role also has all necessary permissions**
+
+a. For quick setup you may want to try the solution either with `AdministratorAccess` managed policy or 
+    at least with the following policies: `AmazonSSMAutomationRole`, `CloudWatchLogsReadOnlyAccess`,
+    `AmazonSageMakerFullAccess`, `AmazonS3FullAccess`.
+
+> NOTE: you can test the solution with elevated privileges like `AdministratorAccess`,
+> but for stronger security you should always scope permissions to the least privilege principle.
+> For more information see [Apply least-privilege permissions](https://docs.aws.amazon.com/IAM/latest/UserGuide/best-practices.html#grant-least-privilege)
+> section of the IAM documentation.
+
+b. (Optional) To prevent users from accessing each other's instances, 
+    add to their role a new inline policy named `SSMDenyAccessNotOwner`.
+    It helps to make sure that the local user/role, i.e. which fires the training job,
+    is the same one that is connecting via SSH:
+
+```json
+{
+    "Version": "2012-10-17",
+    "Statement": [
+        {
+            "Sid": "VisualEditor0",
+            "Effect": "Deny",
+            "Action": "ssm:StartSession",
+            "Resource": "arn:aws:ssm:*:*:managed-instance/*",
+            "Condition": {
+                "StringNotEquals": {
+                    "ssm:resourceTag/SSHOwner": "${aws:userid}"
+                }
+            }
+        },
+        {
+            "Sid": "VisualEditor1",
+            "Effect": "Deny",
+            "Action": "ssm:SendCommand",
+            "Resource": "arn:aws:ssm:*:*:managed-instance/*",
+            "Condition": {
+                "StringNotEquals": {
+                    "ssm:resourceTag/SSHOwner": "${aws:userid}"
+                }
+            }
+        }
+    ]
+}
+```

.idea/vcs.xml

@@ -1,15 +1,14 @@
 Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 
-Permission is hereby granted, free of charge, to any person obtaining a copy of
-this software and associated documentation files (the "Software"), to deal in
-the Software without restriction, including without limitation the rights to
-use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
-the Software, and to permit persons to whom the Software is furnished to do so.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
-FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
-COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
-IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
-CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+Permission is hereby granted, free of charge, to any person obtaining a copy of this
+software and associated documentation files (the "Software"), to deal in the Software
+without restriction, including without limitation the rights to use, copy, modify,
+merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
+permit persons to whom the Software is furnished to do so.
 
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
+INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
+PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
+HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.