├── .gitignore ├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── LICENSE ├── Makefile ├── Pipfile ├── README.md ├── aws-sso-setup.md ├── cfn-templates ├── iam.yaml ├── sagemaker-domain.yaml └── vpc.yaml ├── design ├── aws-sso-idp-synchronization.drawio.svg ├── iam-roles-setup.drawio.svg ├── network-architecture.drawio.svg ├── solution-architecture.drawio.svg └── solution-flow.drawio.svg ├── functions ├── inline │ ├── get_network_config.py │ └── get_user_profile_metadata.py ├── requirements.txt └── saml-backend │ ├── requirements.txt │ └── saml_backend_function.py ├── img ├── control-panel-profiles.png ├── signing-to-studio-2.png ├── signing-to-studio.png ├── sso-app-mapping.png ├── sso-app-user.png ├── sso-app.png ├── sso-custom-apps-2.png ├── sso-custom-apps.png ├── sso-store-id.png ├── starting-sm-studio.png └── studio-exec-role.png ├── template.yaml ├── test └── cfn-unit-test.sh └── user-profile-metadata.json /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | # Project-specific 132 | deploy 133 | 134 | # Build folder 135 | 136 | */build/* 137 | 138 | # SAM 139 | .aws-sam/* 140 | samconfig.* 141 | .DS_Store 142 | .environment 143 | .not-used-snippets 144 | 145 | .gp2/* 146 | *.pdf 147 | *snapshot.json 148 | .test* 149 | *.zip -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | ## Code of Conduct 2 | This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct). 3 | For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact 4 | opensource-codeofconduct@amazon.com with any additional questions or comments. 5 | 6 | Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 7 | SPDX-License-Identifier: MIT-0 -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing Guidelines 2 | 3 | Thank you for your interest in contributing to our project. Whether it's a bug report, new feature, correction, or additional 4 | documentation, we greatly value feedback and contributions from our community. 5 | 6 | Please read through this document before submitting any issues or pull requests to ensure we have all the necessary 7 | information to effectively respond to your bug report or contribution. 8 | 9 | 10 | ## Reporting Bugs/Feature Requests 11 | 12 | We welcome you to use the GitHub issue tracker to report bugs or suggest features. 13 | 14 | When filing an issue, please check [existing open](https://github.com/aws-samples/amazon-sagemaker-studio-secure-sso/issues), or [recently closed](https://github.com/aws-samples/amazon-sagemaker-studio-secure-sso/issues?utf8=%E2%9C%93&q=is%3Aissue%20is%3Aclosed%20), issues to make sure somebody else hasn't already 15 | reported the issue. Please try to include as much information as you can. Details like these are incredibly useful: 16 | 17 | * A reproducible test case or series of steps 18 | * The version of our code being used 19 | * Any modifications you've made relevant to the bug 20 | * Anything unusual about your environment or deployment 21 | 22 | 23 | ## Contributing via Pull Requests 24 | Contributions via pull requests are much appreciated. Before sending us a pull request, please ensure that: 25 | 26 | 1. You are working against the latest source on the *master* branch. 27 | 2. You check existing open, and recently merged, pull requests to make sure someone else hasn't addressed the problem already. 28 | 3. You open an issue to discuss any significant work - we would hate for your time to be wasted. 29 | 30 | To send us a pull request, please: 31 | 32 | 1. Fork the repository. 33 | 2. Modify the source; please focus on the specific change you are contributing. If you also reformat all the code, it will be hard for us to focus on your change. 34 | 3. Ensure local tests pass. 35 | 4. Commit to your fork using clear commit messages. 36 | 5. Send us a pull request, answering any default questions in the pull request interface. 37 | 6. Pay attention to any automated CI failures reported in the pull request, and stay involved in the conversation. 38 | 39 | GitHub provides additional document on [forking a repository](https://help.github.com/articles/fork-a-repo/) and 40 | [creating a pull request](https://help.github.com/articles/creating-a-pull-request/). 41 | 42 | 43 | ## Finding contributions to work on 44 | Looking at the existing issues is a great way to find something to contribute on. As our projects, by default, use the default GitHub issue labels (enhancement/bug/duplicate/help wanted/invalid/question/wontfix), looking at any ['help wanted'](https://github.com/aws-samples/amazon-sagemaker-studio-secure-sso/labels/help%20wanted) issues is a great place to start. 45 | 46 | 47 | ## Code of Conduct 48 | This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct). 49 | For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact 50 | opensource-codeofconduct@amazon.com with any additional questions or comments. 51 | 52 | 53 | ## Security issue notifications 54 | If you discover a potential security issue in this project we ask that you notify AWS/Amazon Security via our [vulnerability reporting page](http://aws.amazon.com/security/vulnerability-reporting/). Please do **not** create a public github issue. 55 | 56 | 57 | ## Licensing 58 | 59 | See the [LICENSE](./LICENSE) file for our project's licensing. We will ask you to confirm the licensing of your contribution. 60 | 61 | We may ask you to sign a [Contributor License Agreement (CLA)](http://en.wikipedia.org/wiki/Contributor_License_Agreement) for larger changes. 62 | 63 | Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 64 | SPDX-License-Identifier: MIT-0 -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | SPDX-License-Identifier: MIT-0 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy of this 5 | software and associated documentation files (the "Software"), to deal in the Software 6 | without restriction, including without limitation the rights to use, copy, modify, 7 | merge, publish, distribute, sublicense, and/or sell copies of the Software, and to 8 | permit persons to whom the Software is furnished to do so. 9 | 10 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, 11 | INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A 12 | PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT 13 | HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 14 | OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 15 | SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | #SHELL := /bin/sh 5 | PY_VERSION := 3.8 6 | 7 | export PYTHONUNBUFFERED := 1 8 | 9 | SRC_DIR := functions 10 | TEST_DIR := test 11 | SAM_DIR := .aws-sam 12 | TEMPLATE_DIR := . 13 | TESTAPP_DIR := test/integration/testdata/ 14 | 15 | # user can optionally override the following by setting environment variables with the same names before running make 16 | 17 | # Path to system pip 18 | PIP ?= pip 19 | # Region for deployment 20 | AWS_DEPLOY_REGION ?= us-east-1 21 | # Region for publishing 22 | AWS_PUBLISH_REGION ?= us-east-1 23 | # Stack name 24 | APP_STACK_NAME ?= sagemaker-team-mgmt-sso 25 | # S3 bucket used for packaging SAM templates 26 | PACKAGE_BUCKET ?= $(APP_STACK_NAME)-$(AWS_DEPLOY_REGION) 27 | 28 | PYTHON := $(shell /usr/bin/which python$(PY_VERSION)) 29 | 30 | .DEFAULT_GOAL := build 31 | 32 | ifndef PACKAGE_BUCKET 33 | $(error PACKAGE_BUCKET is not set) 34 | endif 35 | 36 | compile: 37 | pipenv run sam build -p -t $(TEMPLATE_DIR)/template.yaml -m $(SRC_DIR)/requirements.txt 38 | 39 | build: compile 40 | 41 | package: compile 42 | pipenv run sam package --template-file $(SAM_DIR)/build/template.yaml --s3-bucket $(PACKAGE_BUCKET) --output-template-file $(SAM_DIR)/packaged.yaml 43 | 44 | publish: package 45 | pipenv run sam publish --template $(SAM_DIR)/packaged.yaml --region $(AWS_PUBLISH_REGION) 46 | 47 | # need to add CAPABILITY_AUTO_EXPAND to be able to update nested stack (aws-lambda-powertools-python-layer) 48 | deploy: package 49 | pipenv run sam deploy --template-file $(SAM_DIR)/packaged.yaml \ 50 | --stack-name $(APP_STACK_NAME) \ 51 | --capabilities CAPABILITY_IAM CAPABILITY_AUTO_EXPAND \ 52 | --region $(AWS_DEPLOY_REGION) \ 53 | --confirm-changeset 54 | 55 | cfn_nag_scan: 56 | cfn_nag_scan --input-path $(TEMPLATE_DIR) 57 | -------------------------------------------------------------------------------- /Pipfile: -------------------------------------------------------------------------------- 1 | [[source]] 2 | url = "https://pypi.org/simple" 3 | verify_ssl = true 4 | name = "pypi" 5 | 6 | [packages] 7 | aws-lambda-powertools = "*" 8 | crhelper = "*" 9 | 10 | [dev-packages] 11 | 12 | [requires] 13 | python_version = "3.8" 14 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Team and user management with Amazon SageMaker and AWS SSO 2 | 3 | This code repository is for the [AWS Machine Learning Blog](https://aws.amazon.com/blogs/machine-learning/) [post](https://aws.amazon.com/blogs/machine-learning/team-and-user-management-with-amazon-sagemaker-and-aws-sso/). 4 | 5 | [Amazon SageMaker Studio](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-ui.html) is a web-based integrated development environment (IDE) for machine learning (ML) that lets you build, train, debug, deploy, and monitor your ML models. Each onboarded user in Studio has their own dedicated set of resources, such as compute instances, a home directory on an [Amazon Elastic File System](https://aws.amazon.com/efs/) (Amazon EFS) volume, and a dedicated [Identity and Access Management](https://aws.amazon.com/iam/) (IAM) execution role. 6 | 7 | One of the most common real-world challenges in setting up user access for Studio is how to manage multiple users, groups, and data science teams for data access and resource isolation. 8 | 9 | Many customers implement user management using federated identities with [AWS Single Sign-On](https://docs.aws.amazon.com/singlesignon/latest/userguide/what-is.html) (AWS SSO) and an external identity provider (IdP), such as Active Directory (AD) or AWS Managed Microsoft AD directory. It's aligned with the AWS [recommended practice](https://wa.aws.amazon.com/wat.question.SEC_2.en.html) of using temporary credentials to access AWS accounts. 10 | 11 | [Amazon SageMaker](https://aws.amazon.com/sagemaker) [domain](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-entity-status.html) supports AWS SSO and can be configured in AWS SSO [authentication mode](https://docs.aws.amazon.com/sagemaker/latest/dg/onboard-sso-users.html). In this case, each entitled AWS SSO user has their own [Studio user profile](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-entity-status.html). Users given access to Studio have a unique sign-in URL that directly opens Studio, and they sign in with their AWS SSO credentials. Organizations manage their users in AWS SSO instead of the SageMaker domain. You can assign multiple users access to the domain at the same time. You can use Studio user profiles for each user to define their security permissions in Studio notebooks via an IAM role attached to the user profile, called an execution role. This role controls permissions for SageMaker operations according to its IAM permission policies. 12 | 13 | In AWS SSO authentication mode, there is always one-to-one mapping between users and user profiles. The SageMaker domain manages the creation of user profiles based on the AWS SSO user ID. You can’t create user profiles via the [AWS Management Console](http://aws.amazon.com/console). This works well in the case when one user is a member of only one data science team or if users have the same or very similar access requirements across their projects and teams. In a more common use case, when a user can participate in multiple ML projects and be a member of multiple teams with slightly different permission requirements, the user requires access to different Studio user profiles with different execution roles and permission policies. Because you can’t manage user profiles independently of AWS SSO in AWS SSO authentication mode, you can’t implement a one-to-many mapping between users and Studio user profiles. 14 | 15 | If you need to establish a strong separation of security contexts, for example for different data categories, or need to entirely prevent the visibility of one group of users’ activity and resources to another, the recommended approach is to create multiple SageMaker domains. At the time of this writing, you can create only one domain per AWS account per Region. To implement the strong separation, you can use multiple AWS accounts with one domain per account as a workaround. 16 | 17 | The second challenge is to [restrict access to the Studio IDE](https://aws.amazon.com/about-aws/whats-new/2020/12/secure-sagemaker-studio-access-using-aws-privatelink-aws-iam-sourceip-restrictions/) to only users from inside a corporate network or a designated VPC. You can achieve this by using [IAM-based access control policies](https://docs.aws.amazon.com/sagemaker/latest/dg/security_iam_id-based-policy-examples.html#api-access-policy). In this case the SageMaker domain must be configured with [IAM authentication mode](https://docs.aws.amazon.com/sagemaker/latest/dg/onboard-iam.html), because the IAM identity-based polices aren’t supported by the sign-in mechanism in AWS SSO mode. The post [Secure access to Amazon SageMaker Studio with AWS SSO and a SAML application](https://aws.amazon.com/blogs/machine-learning/secure-access-to-amazon-sagemaker-studio-with-aws-sso-and-a-saml-application/) solves this challenge and demonstrates how to control network access to a SageMaker domain. 18 | 19 | This solution addresses these challenges of AWS SSO user management for Studio for a common use case of multiple user groups and a many-to-many mapping between users and teams. The solution outlines how to use a [custom SAML 2.0 application](https://docs.aws.amazon.com/singlesignon/latest/userguide/samlapps.html#addconfigcustomapp) as the mechanism to trigger the user authentication for Studio and support multiple Studio user profiles per one AWS SSO user. 20 | 21 | ## Solution overview 22 | The solution implements the following architecture: 23 | 24 | ![](design/solution-architecture.drawio.svg) 25 | 26 | The main high-level architecture components are: 27 | 28 | **1 - Identity provider** 29 | Users and groups are managed in an external identity source, for example in Azure Active Directory. User assignments to AD groups define what permissions a particular user has and which SageMaker Studio "team" they have access to. The identity source must by synchronized with AWS SSO. 30 | 31 | **2 - AWS Single Sign-On** 32 | AWS Single Sign-On service manages SSO users, SSO permission set, and applications. This solution uses custom a SAML 2.0 application to provide access to Studio for entitled SSO users. The solution also uses SAML attribute mapping to populate the SAML assertion with specific access-relevant data, such as user id and user team. 33 | 34 | **3 - custom SAML 2.0 applications** 35 | The solution creates one application per SageMaker Studio team and assigns one or multiple applications to a user or a user group based on entitlements. Users can access these applications from within their SSO user portal based on assigned permissions. Each application is configured with the [Amazon API Gateway](https://aws.amazon.com/api-gateway/) endpoint URL as its SAML backend. 36 | 37 | **4 - Amazon SageMaker domain** 38 | The solution provisions a SageMaker domain in an AWS account and creates a dedicated user profile for each combination of SSO user and Studio team the user is assigned to. The domain must be configured in IAM [authentication mode](https://docs.aws.amazon.com/sagemaker/latest/dg/onboard-iam.html). 39 | 40 | **5 - Studio user profiles** 41 | The solution automatically creates a dedicated user profile for each _user-team_ combination. For example, if a user is a member of two Studio teams and has corresponding permissions, the solution provisions two separate user profiles for this user. Each profile always belongs to one and only one user. 42 | 43 | To demonstrate the configuration, we use two users, _User 1_, _User 2_, and two Studio teams, _Team 1_, _Team 2_. The _User 1_ belongs to both teams, while the _User 2_ belongs to _Team 2_ only. The _User 1_ can access Studio environments for both teams, while the _User 2_ can access only the Studio environment for _Team 2_. 44 | 45 | **6 - Studio execution roles** 46 | Each Studio user profile uses a dedicated execution role with permission polices with required level of access for the specific team the user belongs to. Studio execution roles implement an effective permission isolation between individual users and their team roles. You manage data and resource access for each role and not at an individual user level. 47 | 48 | The solution also implements an attribute-based access control (ABAC) using SAML 2.0 attributes, tags on Studio user profiles, and tags on SageMaker execution roles. 49 | 50 | ❗ In this particular configuration we assume that SSO users don't have permissions to sign into the AWS account and don't have corresponding AWS SSO-controlled IAM roles in the account. Each user signs into the Studio environment via a presigned URL from a SSO portal without the need to go to AWS console in the AWS account. 51 | In a real-world environment you might need to setup [SSO permission sets](https://docs.aws.amazon.com/singlesignon/latest/userguide/permissionsetsconcept.html) for SSO users to allow the authorized users to assume an IAM role and sign into an AWS account. For example, you can provide _Data Scientist_ role permissions for a user to be able to interact with account resources and have the level of access they need to fulfill their role. 52 | 53 | ### How solution works 54 | The following diagram presents the end-to-end sign-on flow for an AWS SSO user: 55 | 56 | ![](design/solution-flow.drawio.svg) 57 | 58 | An AWS SSO user clicks on a corresponding Studio application in their SSO portal. AWS SSO prepares a SAML assertion (**1**) with configured SAML attribute mappings. A custom SAML application is configured with the Amazon API Gateway endpoint URL as its Assertion Consumer Service (ACS), and needs mapping attributes containing the AWS SSO user ID and team ID. We use `ssouserid` and `teamid` custom attributes to send all needed information to the SAML backend. 59 | 60 | The API Gateway calls an SAML backend API. AWS Lambda function (**2**) implements the API, parses the SAML response to extract the user ID and team ID. The function uses these attributes to retrieve a team-specific configuration, such as an execution role and SageMaker domain ID. The backend function checks if a required user profile exists in the domain, and creates a new one with the corresponding configuration settings if no profile exists. Afterwards, the function generates a Studio presigned URL for a specific Studio user profile by calling [`CreatePresignedDomainUrl`](https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreatePresignedDomainUrl.html) API (**3**) via a SageMaker API VPC endpoint. The Lambda function finally returns the presigned URL with HTTP 302 redirection response (**4**) to sign the user in Studio. 61 | 62 | ❗ The solution implements **a non-production sample** version of an SAML backend. The Lambda function parses the SAML assertion and uses only attributes in `` element to construct a `CreatePresignedDomainUrl` API call. 63 | In your production solution you must use a proper SAML backend implementation which must include a validation of an SAML response, a signature, and certificates, replay and redirect prevention, and any other features of an SAML authentication process. For example, you can use a [python3-saml SAML backend implementation](https://python-social-auth.readthedocs.io/en/latest/backends/saml.html) or 64 | [OneLogin open source SAML toolkit](https://developers.onelogin.com/saml/python) to implement a secure SAML backend. 65 | 66 | ### Dynamic creation of Studio user profiles 67 | The solution automatically creates Studio user profiles for each combination _User-Team_, as soon as the SSO sign-in process requests a presigned URL. The creation of a user profile is based on the configured metadata in the [AWS SAM template](./template.yaml): 68 | ```yaml 69 | Metadata: 70 | Team1: 71 | DomainId: !If 72 | - CreateSageMakerDomainCondition 73 | - !GetAtt SageMakerDomain.Outputs.SageMakerDomainId 74 | - !Ref SageMakerDomainId 75 | SessionExpiration: 43200 76 | Tags: 77 | - Key: Team 78 | Value: Team1 79 | UserSettings: 80 | ExecutionRole: !GetAtt IAM.Outputs.SageMakerStudioExecutionRoleTeam1Arn 81 | Team2: 82 | DomainId: !If 83 | - CreateSageMakerDomainCondition 84 | - !GetAtt SageMakerDomain.Outputs.SageMakerDomainId 85 | - !Ref SageMakerDomainId 86 | SessionExpiration: 43200 87 | Tags: 88 | - Key: Team 89 | Value: Team2 90 | UserSettings: 91 | ExecutionRole: !GetAtt IAM.Outputs.SageMakerStudioExecutionRoleTeam2Arn 92 | ``` 93 | 94 | You can configure own teams, custom settings, and tags simply by adding them to the metadata configuration for the CloudFormation resource `GetUserProfileMetadata`. 95 | Refer to the `create_user_profile` boto3 [documentation](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.create_user_profile) for a further configuration elements of `UserSettings`. 96 | 97 | In your productive solution you can use a dedicated microservice to manage and retrieve user profile metadata. For example, you also can use [Amazon DynamoDB](https://aws.amazon.com/dynamodb) to store and manage the metadata. 98 | 99 | ### IAM roles 100 | The following diagram shows the IAM roles in this solution: 101 | 102 | ![](design/iam-roles-setup.drawio.svg) 103 | 104 | **1 - Studio execution role** 105 | A Studio user profile uses a dedicated Studio execution role with data and resource permissions specific for each team or user group. This role can also use tags to implement ABAC for data and resource access. Refer to the [SageMaker Roles](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html) documentation for more details. 106 | 107 | **2 - SAML backend Lambda execution role** 108 | This execution role contains permission to call `CreatePresignedDomainUrl` API. You can configure permission policy to include additional conditional checks using `Condition` keys: for example, allow access to Studio only from a designated range of IP addresses within your private corporate network: 109 | ```json 110 | { 111 | "Version": "2012-10-17", 112 | "Statement": [ 113 | { 114 | "Action": [ 115 | "sagemaker:CreatePresignedDomainUrl" 116 | ], 117 | "Resource": "arn:aws:sagemaker:::user-profile/*/*", 118 | "Effect": "Allow" 119 | }, 120 | { 121 | "Condition": { 122 | "NotIpAddress": { 123 | "aws:VpcSourceIp": "10.100.10.0/24" 124 | } 125 | }, 126 | "Action": [ 127 | "sagemaker:*" 128 | ], 129 | "Resource": "arn:aws:sagemaker:::user-profile/*/*", 130 | "Effect": "Deny" 131 | } 132 | ] 133 | } 134 | ``` 135 | For more examples of how to use conditions in IAM policies, refer to [Control Access to the SageMaker API by Using Identity-based Policies](https://docs.aws.amazon.com/sagemaker/latest/dg/security_iam_id-based-policy-examples.html#api-access-policy) documentation. 136 | 137 | Since the SAML backend Lambda function creates a new Studio user profile, it must have `iam:PassRole` permission for Studio execution roles which are attached to the profile: 138 | 139 | ```json 140 | { 141 | "Version": "2012-10-17", 142 | "Statement": [ 143 | { 144 | "Action": [ 145 | "iam:PassRole" 146 | ], 147 | "Resource": [ 148 | "arn:aws:iam:::role/", 149 | "arn:aws:iam:::role/" 150 | ], 151 | "Effect": "Allow" 152 | } 153 | ] 154 | } 155 | ``` 156 | 157 | **3 - SageMaker service** 158 | SageMaker service assumes the Studio execution role on your behalf, as controlled by a corresponding trust policy on the execution role. This allows the service to access data and resources, and perform actions on your behalf. The Studio execution role must contain a trust policy allowing SageMaker service to assume this role. 159 | 160 | **4 - SSO permission set IAM role** 161 | You can assign your SSO users to AWS accounts in your AWS Organizations via [SSO permission sets](https://docs.aws.amazon.com/singlesignon/latest/userguide/permissionsetsconcept.html). A permission set is a template that defines a collection of user role-specific IAM policies. You manage permission sets in AWS SSO and AWS SSO controls the corresponding IAM roles in each account. 162 | 163 | **5 - AWS Organizations Service Control Policies (SCPs)** 164 | If you use [AWS Organizations](https://aws.amazon.com/organizations/), you can implement [Service Control Policies](https://docs.aws.amazon.com/organizations/latest/userguide/orgs_manage_policies_scps.html) (SCPs) to centrally control the maximum available permissions for all accounts and all IAM roles in your organization. For example, in order to centrally prevent access to Studio via AWS console, you can implement the following SCP and attach it to the accounts with SageMaker domain: 165 | 166 | ```json 167 | { 168 | "Version": "2012-10-17", 169 | "Statement": [ 170 | { 171 | "Action": [ 172 | "sagemaker:*" 173 | ], 174 | "Resource": "*", 175 | "Effect": "Allow" 176 | }, 177 | { 178 | "Condition": { 179 | "NotIpAddress": { 180 | "aws:VpcSourceIp": "" 181 | } 182 | }, 183 | "Action": [ 184 | "sagemaker:CreatePresignedDomainUrl" 185 | ], 186 | "Resource": "*", 187 | "Effect": "Deny" 188 | } 189 | ] 190 | } 191 | ``` 192 | 193 | #### Solution provisioned roles 194 | The solution CFN stack creates three Studio execution roles used in the SageMaker domain: 195 | - `SageMakerStudioExecutionRoleDefault` 196 | - `SageMakerStudioExecutionRoleTeam1` 197 | - `SageMakerStudioExecutionRoleTeam2` 198 | 199 | None of the roles has [`AmazonSageMakerFullAccess`](https://docs.aws.amazon.com/sagemaker/latest/dg/security-iam-awsmanpol.html) policy attached and each has only a limited set of permissions. In your real-world SageMaker environment you need to amend the role's permissions based on your specific requirements. 200 | 201 | `SageMakerStudioExecutionRoleDefault` has only a custom policy `SageMakerReadOnlyPolicy` attached with a restrictive list of allowed actions. 202 | 203 | Both team roles, `SageMakerStudioExecutionRoleTeam1` and `SageMakerStudioExecutionRoleTeam2` additionally have two custom polices `SageMakerAccessSupportingServicesPolicy` and `SageMakerStudioDeveloperAccessPolicy` allowing usage of particular services and one deny-only policy `SageMakerDeniedServicesPolicy` with explicit deny on some SageMaker API calls. 204 | 205 | The Studio developer access policy enforces setting of `Team` tag equal to the same value as user's own execution role for calling any SageMaker `Create*` API: 206 | ```json 207 | { 208 | "Condition": { 209 | "ForAnyValue:StringEquals": { 210 | "aws:TagKeys": [ 211 | "Team" 212 | ] 213 | }, 214 | "StringEqualsIfExists": { 215 | "aws:RequestTag/Team": "${aws:PrincipalTag/Team}" 216 | } 217 | }, 218 | "Action": [ 219 | "sagemaker:Create*" 220 | ], 221 | "Resource": [ 222 | "arn:aws:sagemaker:*::*" 223 | ], 224 | "Effect": "Allow", 225 | "Sid": "AmazonSageMakerCreate" 226 | } 227 | ``` 228 | 229 | Furthermore, it allows using delete, stop, update, and start operations only on resources tagged with the same `Team` tag as user's execution role: 230 | ```json 231 | { 232 | "Condition": { 233 | "StringEquals": { 234 | "aws:PrincipalTag/Team": "${sagemaker:ResourceTag/Team}" 235 | } 236 | }, 237 | "Action": [ 238 | "sagemaker:Delete*", 239 | "sagemaker:Stop*", 240 | "sagemaker:Update*", 241 | "sagemaker:Start*", 242 | "sagemaker:DisassociateTrialComponent", 243 | "sagemaker:AssociateTrialComponent", 244 | "sagemaker:BatchPutMetrics" 245 | ], 246 | "Resource": [ 247 | "arn:aws:sagemaker:*::*" 248 | ], 249 | "Effect": "Allow", 250 | "Sid": "AmazonSageMakerUpdateDeleteExecutePolicy" 251 | } 252 | ``` 253 | 254 | For more information on roles and polices, refer to the post [Configuring Amazon SageMaker Studio for teams and groups with complete resource isolation](https://aws.amazon.com/fr/blogs/machine-learning/configuring-amazon-sagemaker-studio-for-teams-and-groups-with-complete-resource-isolation/). 255 | 256 | ### Network infrastructure 257 | The solution implements a fully isolated SageMaker domain environment with all network traffic going through [AWS PrivateLink](https://aws.amazon.com/privatelink) connections. You may optionally enable internet access from the Studio notebooks. The solution also creates three [VPC security groups](https://docs.aws.amazon.com/vpc/latest/userguide/VPC_SecurityGroups.html) to control traffic between all solution components such as the SAML backend Lambda function, [VPC endpoints](https://docs.aws.amazon.com/vpc/latest/privatelink/concepts.html), and SageMaker Studio notebooks. 258 | 259 | ![](design/network-architecture.drawio.svg) 260 | 261 | This solution provisions all required network infrastructure. The CloudFormation template `./cfn-templates/vpc.yaml` contains the source code. 262 | 263 | ## Deployment 264 | To deploy and test the solution you must complete the following steps: 265 | 1. [Deploy solution's stack via a AWS SAM template](#deploy-sam-template). 266 | 2. [Create AWS SSO users](#create-aws-sso-users), or use your existing AWS SSO users. 267 | 3. [Create custom SAML 2.0 applications](#create-custom-saml-20-applications) and assign AWS SSO users to the applications. 268 | 269 | ### Prerequisites 270 | [AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/cli-chap-getting-started.html), [AWS SAM CLI](https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/serverless-sam-cli-install.html) and [Python3.8 or later](https://www.python.org/downloads/) must be installed. 271 | 272 | The deployment procedure assumes that AWS SSO has been enabled and configured for the [AWS Organization](https://docs.aws.amazon.com/organizations/latest/userguide/orgs_introduction.html) where the solution will be deployed. 273 | 274 | You can follow these [instructions](./aws-sso-setup.md) to setup AWS Single Sign-On. 275 | 276 | ### Solution deployment options 277 | You can choose several solution deployment option to have the best fit for your existing AWS environment. You can select network and SageMaker domain provisioning options. 278 | 279 | #### Network deployment options 280 | There are following network infrastructure deployment options: 281 | - **New VPC**: the solution creates a new VPC with all subnets, one public and one private route tables, NAT and Internet gateways, security groups, and VPC endpoints. 282 | - **Existing VPC**: you can use your **existing** VPC, a public subnet, and NAT and Internet gateways. No one of these resources are created by the solution in this option. If you use an existing VPC you can choose one of the following options: 283 | - **new private subnets**: the solution creates private subnets without internet access, a route table with a local route only, security groups, and VPC endpoints. 284 | - **use existing private subnets**: the solution creates security groups and VPC endpoints only. 285 | 286 | To choose one of these deployment options, provide the following CloudFormation template parameters for SAM deployment process. 287 | 288 | ##### New VPC 289 | - `VPCCIDR` (optional): CIDR block for a new VPC. Default is `10.0.0.0/16` 290 | - `SAMLBackendPrivateSubnetCIDR` (optional): CIDR block for a private subnet for SAML backend. Default is `10.0.0.0/19` 291 | - `SageMakerDomainPrivateSubnetCIDR` (optional): CIDR block for a private subnet for SageMaker domain. Default is `10.0.32.0/19` 292 | - `PublicSubnetCIDR` (optional): CIDR block for a public subnet for Internet and NAT Gateways. Default is `10.0.128.0/20` 293 | 294 | ❗ The provided VPC and subnet CIDR blocks must be compatible with your existing VPC and subnets. Refer to [VPC documentation](https://docs.aws.amazon.com/vpc/latest/userguide/configure-your-vpc.html#vpc-sizing-ipv4) on more details on CIDR block associations. 295 | 296 | ##### Existing VPC and new private subnets 297 | - `ExistingVPCId` (required): Existing VPC id. You can list all VPC in your AWS account by running an AWS CLI command: 298 | ``` 299 | aws ec2 describe-vpcs 300 | ``` 301 | - `CreatePrivateSubnets` (required): Must be set to `YES` 302 | - `SAMLBackendPrivateSubnetCIDR` (required): CIDR block for a **new** private SAML backend subnet. 303 | - `SageMakerDomainPrivateSubnetCIDR` (required): CIDR block for a **new** private subnet for SageMaker domain. 304 | 305 | ❗ The provided private subnet CIDR blocks must be compatible with your VPC and existing subnets. Refer to [VPC documentation](https://docs.aws.amazon.com/vpc/latest/userguide/configure-your-vpc.html#vpc-sizing-ipv4) on more details on CIDR block associations. 306 | ❗ The private subnets are created without internet access. The stack creates a route table with a local route only and associates this table with SAML backend and SageMaker private subnets. You must add a [route to a NAT gateway](https://docs.aws.amazon.com/vpc/latest/userguide/vpc-nat-gateway.html#nat-gateway-create-route) to the route table if you need an internet route for the private subnets. If you don't configure an internet route for a SageMaker private subnet, you won't have internet access in Studio notebooks. 307 | 308 | ##### Existing VPC and existing private subnets 309 | - `ExistingVPCId` (required): Existing VPC id 310 | - `CreatePrivateSubnets` (required): Must be set to `NO` 311 | - `ExistingSAMLBackendPrivateSubnetId` (required): subnet id for an **existing** subnet. The SAML backend will be created in this subnet. 312 | - `ExistingSageMakerDomainPrivateSubnetId` (required): subnet id for an **existing** subnet for a SageMaker domain. To list all SageMaker domain subnets you can run the following AWS CLI commands: 313 | ``` 314 | export DOMAIN_ID=$(aws sagemaker list-domains --output text --query 'Domains[0].DomainId') 315 | aws sagemaker describe-domain --domain-id $DOMAIN_ID --output text --query 'SubnetIds[*]' 316 | ``` 317 | ❗ For this option you must use an existing SageMaker domain private subnet in the Availability Zone `a`, for example in `us-east-1a` for North Virginia AWS Region. The stack creates SageMaker API, Studio, and runtime VPC endpoints in the Availability Zone `a`. 318 | 319 | #### SageMaker domain deployment options 320 | The solution provisions a SageMaker domain in `VpcOnly` network mode. You have an option to use already existing domain and create new Studio user profiles only. 321 | 322 | ##### New domain 323 | To create a new domain, leave `SageMakerDomainId` parameter empty. 324 | 325 | ##### Existing domain 326 | If you want to use your existing domain, you can provide the domain id in `SageMakerDomainId` parameter. 327 | To get the domain id, run the following command in your terminal: 328 | ```sh 329 | export DOMAIN_ID=$(aws sagemaker list-domains --output text --query 'Domains[0].DomainId') 330 | ``` 331 | 332 | ❗ The domain must be in `IAM` authentication mode. Run the following command to check the authentication mode: 333 | ```sh 334 | aws sagemaker describe-domain --domain-id $DOMAIN_ID --output text --query 'AuthMode' 335 | ``` 336 | 337 | ### Deploy the AWS SAM template 338 | 1. Install [AWS SAM CLI](https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/serverless-sam-cli-install.html) if you do not have it 339 | 2. Clone the source code repository to your local environment: 340 | ```sh 341 | git clone https://github.com/aws-samples/users-and-team-management-with-amazon-sagemaker-and-aws-sso.git 342 | ``` 343 | 3. Build AWS SAM application: 344 | ```bash 345 | sam build 346 | ``` 347 | 4. Deploy the application: 348 | ```bash 349 | sam deploy --guided 350 | ``` 351 | 352 | Provide stack parameters according to your existing environment and desired deployment options, such as existing VPC, existing private and public subnets, existing SageMaker domain, as discussed above. 353 | 354 | You can leave **all parameters** at their default values to provision new network resources and a new SageMaker domain. 355 | 356 | The following table summarizes the parameter default values and usage. 357 | Parameter | Default value | Usage 358 | ---|---|--- 359 | `EnvironmentName` | `sagemaker-team-mgmt-sso` | Stack and resource names 360 | `SageMakerDomainId` | Empty | Leave empty for a new domain. Provide a domain id to use an existing domain 361 | `CreatePrivateSubnets` | `YES` | Set to `YES` to create new private subnets, in a new or existing VPC. Set to `NO` if you re-use your existing subnets and VPC 362 | `ExistingVPCId` | Empty | Required only if you use an existing VPC 363 | `SAMLBackendSubnetId` | Empty | Required only if `CreatePrivateSubnets` = `NO`. Provide a subnet id for an existing subnet 364 | `SageMakerDomainSubnetId` | Empty | Required only if `CreatePrivateSubnets` = `NO`. Provide a subnet id for an existing subnet 365 | `VPCCIDR` | `10.0.0.0/16` | Used only for a new VPC. Leave default or use a custom block 366 | `SAMLBackedPrivateSubnetCIDR` | `10.0.0.0/19` | Used only if `CreatePrivateSubnets` = YES. Leave default or use a custom block 367 | `SageMakerDomainPrivateSubnetCIDR` | `10.0.32.0/19` | Used only if `CreatePrivateSubnets` = YES. Leave default or use a custom block 368 | `PublicSubnetCIDR` | `10.0.128.0/20` | Used only for a new VPC. Leave default or use a custom block 369 | `DomainAccessAllowedCIDR` | Empty | Allowed CIDR block for `CreatePresignedDomainURL` API call. Use to restrict access to the Studio to only users from inside a specified network CIDR. 370 | 371 | Wait until the stack deployment is complete. The end-to-end deployment including provisioning all network resources and a SageMaker domain takes about 20 minutes. 372 | 373 | To see the stack output run the following command in the terminal: 374 | ```sh 375 | export STACK_NAME= 376 | 377 | aws cloudformation describe-stacks \ 378 | --stack-name $STACK_NAME \ 379 | --output table \ 380 | --query "Stacks[0].Outputs[*].[OutputKey, OutputValue]" 381 | ``` 382 | 383 | ### Create AWS SSO users 384 | Follow [add AWS SSO user instructions](./aws-sso-setup.md#add-aws-sso-users) to create two users with names `User1` and `User2` or use any two of your existing AWS SSO users to test the solution. Make sure you use AWS SSO in the same AWS Region in which you deployed the solution. 385 | 386 | ### Create custom SAML 2.0 applications 387 | Open the [AWS SSO console](https://console.aws.amazon.com/singlesignon) on the AWS management account of your AWS Organizations. Make sure you select the same Region as the one where you deployed the solution stack. 388 | 389 | Do the following steps for **each** of two required custom SAML 2.0 applications, for _Team 1_ and for _Team 2_. 390 | 1. Choose **Applications** 391 | 2. Choose **Add a new application** 392 | 3. Choose **Add a custom SAML 2.0 application** 393 | 4. For **Display Name**, enter an application name, for example `SageMaker Studio Team 1` 394 | 5. Leave **Application start URL** and **Relay state** empty in **Application Properties** 395 | 6. Choose **If you don't have a metadata file, you can manually type your metadata values.** 396 | 7. For **Application ACS URL**, enter the URL provided in the `SAMLBackendEndpoint` key of the AWS SAM stack output 397 | 8. For **Application SAML audience**, enter the URL provided in the `SAMLAudience` key of the AWS SAM stack output 398 | 9. Choose **Save Changes** 399 | 400 | ![](./img/sso-app.png) 401 | 402 | 10. Navigate to the **Attribute mappings** tab 403 | 11. Set the Subject to **email** and format **emailAddress** 404 | 12. Add the following new attributes: 405 | - `ssouserid` set to `${user:AD_GUID}` 406 | - `teamid` set to `Team1` or `Team2` respectively for each of the two applications 407 | 408 | ![](./img/sso-app-mapping.png) 409 | 410 | 13. Choose **Save Changes** 411 | 14. On the **Assigned users** tab, choose **Assign users** 412 | 16. Choose the _User 1_ for the _Team 1_ application and both _User 1_ and _User 2_ for _Team 2_ application 413 | 17. Choose **Assign users** 414 | 415 | ![](./img/sso-app-user.png) 416 | 417 | ## Test the solution 418 | Go to AWS SSO user portal `https://.awsapps.com/start` and sign as _User 1_. Two SageMaker applications are shown in the portal: 419 | 420 | ![](img/sso-custom-apps.png) 421 | 422 | Choose on **SageMaker Studio Team 1**. 423 | You're redirected to SageMaker Studio instance for _Team 1_ in a new browser window: 424 | 425 | ![](./img/signing-to-studio.png) 426 | 427 | Because the SAML backend Lambda function creates a new user profile first time a user signs in and needs to wait until the profile becomes `InService`, the very first sign-in might timeout. If you get a timeout error, just sign in again. 428 | 429 | The first time you start Studio, SageMaker creates a JupyterServer application. This process takes few minutes: 430 | 431 | ![](./img/starting-sm-studio.png) 432 | 433 | In Studio, on the **File** menu, choose **New** and **Terminal** to start a new terminal. In the terminal command line enter the command: 434 | ```sh 435 | aws sts get-caller-identity 436 | ``` 437 | The command returns the Studio execution role: 438 | 439 | ![](./img/studio-exec-role.png) 440 | 441 | In our setup, this role must be different for each team. You can also check that each user in each instance of Studio has their own home directory on a mounted Amazon EFS volume. 442 | 443 | Return to AWS SSO portal, still logged as _User 1_ and choose on **SageMaker Studio Team 2** application. You're redirected to a _Team 2_ Studio instance: 444 | 445 | ![](./img/signing-to-studio-2.png) 446 | 447 | The start process can again take several minutes, because SageMaker starts a new JupyterServer application for User 2. 448 | 449 | Sign as _User 2_ in AWS SSO portal. _User 2_ has only one application assigned - **SageMaker Studio Team 2**: 450 | 451 | ![](./img/sso-custom-apps-2.png) 452 | 453 | If you start a instance of Studio via this user application, you can verify that it uses the same SageMaker execution role as _User 1's_ _Team 2_ instance. However, each Studio instance is completely isolated. _User 2_ has their own home directory on an Amazon EFS volume and own instance of JupyterServer application. You can verify this by creating a folder and some files for each of the users and see that each user’s home directory is isolated. 454 | 455 | Now you can sign into Amazon SageMaker console and see that there are three user profiles created: 456 | 457 | ![](./img/control-panel-profiles.png) 458 | 459 | You just implemented a PoC solution to manage multiple user and teams with Studio. 460 | 461 | ## Synchronization with identity provider 462 | 463 | ![idp-sso-sync](design/aws-sso-idp-synchronization.drawio.svg) 464 | 465 | _in development_ 466 | 467 | ## Access management with ABAC 468 | _in development_ 469 | 470 | https://docs.aws.amazon.com/singlesignon/latest/userguide/configure-abac.html 471 | 472 | https://docs.aws.amazon.com/singlesignon/latest/userguide/attributemappingsconcept.html#defaultattributemappings 473 | 474 | 475 | ## Clean-up 476 | To avoid charges, you must remove all project-provisioned and generated resources from your AWS account. Use the following AWS SAM CLI command to delete the solution CloudFormation stack: 477 | 478 | ```sh 479 | sam delete delete-stack --stack-name 480 | ``` 481 | 482 | ❗ For security reasons and to prevent data loss, the Amazon EFS mount and the content associated with the Amazon SageMaker Studio Domain deployed in this solution **are not** deleted. The VPC and subnets associated with the SageMaker domain remain in your AWS account. For instructions to delete the file system and VPC, refer to [Deleting an Amazon EFS file system](https://docs.aws.amazon.com/efs/latest/ug/delete-efs-fs.html) and [Work with VPCs](https://docs.aws.amazon.com/vpc/latest/userguide/working-with-vpcs.html), respectively. 483 | 484 | To delete the custom SAML application, complete the following steps: 485 | 1. Open the [AWS SSO console](https://console.aws.amazon.com/singlesignon) in the SSO management account 486 | 2. Choose **Applications**. 487 | 3. Select **SageMaker Studio Team 1**. 488 | 4. On the **Actions**, choose **Remove**. 489 | 5. Repeat these steps for **SageMaker Studio Team 2**. 490 | 491 | # Resources 492 | 493 | ## Documentation 494 | - [AWS Single Sign-On](https://docs.aws.amazon.com/singlesignon/latest/userguide/what-is.html) 495 | - [Attributes for access control](https://docs.aws.amazon.com/singlesignon/latest/userguide/attributesforaccesscontrol.html) 496 | - [Attribute mappings](https://docs.aws.amazon.com/singlesignon/latest/userguide/attributemappingsconcept.html) 497 | 498 | ## Blog posts 499 | - [Onboarding Amazon SageMaker Studio with AWS SSO and Okta Universal Directory](https://aws.amazon.com/fr/blogs/machine-learning/onboarding-amazon-sagemaker-studio-with-aws-sso-and-okta-universal-directory/) 500 | - [Configuring Amazon SageMaker Studio for teams and groups with complete resource isolation](https://aws.amazon.com/fr/blogs/machine-learning/configuring-amazon-sagemaker-studio-for-teams-and-groups-with-complete-resource-isolation/) 501 | - [Secure access to Amazon SageMaker Studio with AWS SSO and a SAML application](https://aws.amazon.com/blogs/machine-learning/secure-access-to-amazon-sagemaker-studio-with-aws-sso-and-a-saml-application/) 502 | - [Access an Amazon SageMaker Studio notebook from a corporate network](https://aws.amazon.com/blogs/machine-learning/access-an-amazon-sagemaker-studio-notebook-from-a-corporate-network/) 503 | - [Secure Amazon SageMaker Studio presigned URLs Part 1: Foundational infrastructure](https://aws.amazon.com/blogs/machine-learning/secure-amazon-sagemaker-studio-presigned-urls-part-1-foundational-infrastructure/) 504 | - [Secure Amazon SageMaker Studio presigned URLs Part 2: Private API with JWT authentication](https://aws.amazon.com/blogs/machine-learning/secure-amazon-sagemaker-studio-presigned-urls-part-2-private-api-with-jwt-authentication/) 505 | 506 | Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 507 | SPDX-License-Identifier: MIT-0 -------------------------------------------------------------------------------- /aws-sso-setup.md: -------------------------------------------------------------------------------- 1 | # Setup AWS Single Sign-On (SSO) 2 | 3 | ## Introduction 4 | 5 | AWS Single Sign-On is a cloud-based single sign-on (SSO) service that makes it easy to centrally manage SSO access to all of your AWS accounts and cloud applications. 6 | AWS SSO helps you manage access and permissions to commonly used third-party software as a service (SaaS) applications as well as custom applications that support Security Assertion Markup Language (SAML) 2.0. 7 | 8 | ## Enable AWS SSO 9 | 10 | 1. Sign in to the AWS Management Console with your AWS Organizations management account credentials. 11 | 2. Open the [AWS SSO console](https://console.aws.amazon.com/singlesignon). 12 | 3. Choose Enable AWS SSO. 13 | 4. If you have not yet set up AWS Organizations, you will be prompted to create an organization. Choose Create AWS organization to complete this process. 14 | 15 | ## Choose your identity source 16 | 17 | By default, AWS SSO offers its identity source which means users's username and password are directly managed within AWS SSO, however you do have a choice to use an alternate identity source such [Active Directory](https://docs.aws.amazon.com/singlesignon/latest/userguide/manage-your-identity-source-ad.html) or other [external identity provider](https://docs.aws.amazon.com/singlesignon/latest/userguide/manage-your-identity-source-idp.html). 18 | 19 | For the purpose of this example we recommend to select the default AWS SSO option to manage identities. 20 | 21 | ## Add AWS SSO users 22 | 23 | Users and groups that you create in your AWS SSO store are available in AWS SSO only. Use the following procedure to add users to your AWS SSO store. 24 | 25 | 1. Open the [AWS SSO console](https://console.aws.amazon.com/singlesignon). 26 | 2. Choose Users. 27 | 3. Choose Add user and provide the following required information: 28 | 1. Username – This user name will be required to sign in to the user portal and cannot be changed later. 29 | 2. Password – Choose from one of the following choices to send the user's password. 30 | * Send an email to the user with password setup instructions – This option automatically sends the user an email addressed from Amazon Web Services. The email invites the user on behalf of your company to access the AWS SSO user portal. 31 | * Generate a one-time password that you can share with the user – This option provides you with the user portal URL and password details that you can manually send to the user from your email address. 32 | 3. Email address – The value you provide here must be unique. 33 | 4. Confirm email address 34 | 5. First name – You must enter a name here for automatic provisioning to work. For more information, see Automatic provisioning. 35 | 6. Last name – You must enter a name here for automatic provisioning to work. 36 | 7. Display name 37 | 38 | 4. Choose Next: Groups. 39 | 5. Select one or more groups that you want the user to be a member of. Then choose Add user. 40 | 41 | ## Retrieve Identity store ID 42 | 43 | 1. Open the [AWS SSO console](https://console.aws.amazon.com/singlesignon). 44 | 2. Choose Settings 45 | 3. Note the Identity store ID under the Identity source, format is ```d-xxxxxx``` 46 | 47 | ![image info](./img/sso-store-id.png) 48 | 49 | Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 50 | SPDX-License-Identifier: MIT-0 -------------------------------------------------------------------------------- /cfn-templates/iam.yaml: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | AWSTemplateFormatVersion: 2010-09-09 5 | Description: Create IAM Roles for SageMaker user profiles and SAML backend Lambda function 6 | 7 | Outputs: 8 | 9 | SageMakerStudioExecutionRoleDefaultArn: 10 | Description: The ARN of the SageMaker default execution role 11 | Value: !GetAtt SageMakerStudioExecutionRoleDefault.Arn 12 | 13 | SageMakerStudioExecutionRoleTeam1Arn: 14 | Description: The ARN of the SageMaker Team1 execution role 15 | Value: !GetAtt SageMakerStudioExecutionRoleTeam1.Arn 16 | 17 | SageMakerStudioExecutionRoleTeam2Arn: 18 | Description: The ARN of the SageMaker Team2 execution role 19 | Value: !GetAtt SageMakerStudioExecutionRoleTeam2.Arn 20 | 21 | SetupLambdaExecutionRoleArn: 22 | Description: Lambda execution role for stack setup Lambda function 23 | Value: !GetAtt SetupLambdaExecutionRole.Arn 24 | 25 | SAMLBackendLambdaExecutionRoleArn: 26 | Description: Lambda execution role for SAML backend Lambda function 27 | Value: !GetAtt SAMLBackendLambdaExecutionRole.Arn 28 | 29 | Parameters: 30 | EnvironmentName: 31 | Type: String 32 | AllowedPattern: '[a-z0-9\-]*' 33 | Description: Please specify your SageMaker environment name. 34 | 35 | AllowedCIDR: 36 | Type: String 37 | Description: Allowed CIDR block for CreatePresignedDomainURL API call 38 | 39 | Conditions: 40 | RestrictSageMakerToCIDRCondition: !Not [ !Equals [ !Ref AllowedCIDR, ''] ] 41 | 42 | Resources: 43 | 44 | SageMakerDeniedServicesPolicy: 45 | Type: AWS::IAM::ManagedPolicy 46 | Properties: 47 | Description: Explicit deny for specific SageMaker services 48 | PolicyDocument: 49 | Version: 2012-10-17 50 | Statement: 51 | - Sid: AmazonSageMakerDeniedServices 52 | Action: 53 | - sagemaker:CreatePresignedNotebookInstanceUrl 54 | - sagemaker:*NotebookInstance 55 | - sagemaker:*NotebookInstanceLifecycleConfig 56 | - sagemaker:CreateUserProfile 57 | - sagemaker:DeleteDomain 58 | - sagemaker:DeleteUserProfile 59 | Resource: 60 | - '*' 61 | Effect: Deny 62 | 63 | SageMakerReadOnlyPolicy: 64 | Type: AWS::IAM::ManagedPolicy 65 | Properties: 66 | Description: Read-only baseline policy for SageMaker execution role 67 | PolicyDocument: 68 | Version: 2012-10-17 69 | Statement: 70 | - Sid: AmazonSageMakerDescribeReadyOnlyPolicy 71 | Action: 72 | - sagemaker:Describe* 73 | - sagemaker:GetSearchSuggestions 74 | Resource: 75 | - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:*' 76 | Effect: Allow 77 | - Sid: AmazonSageMakerListOnlyPolicy 78 | Action: 79 | - 'sagemaker:List*' 80 | Resource: 81 | - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:*' 82 | Effect: Allow 83 | - Sid: AmazonSageMakerUIandMetricsOnlyPolicy 84 | Action: 85 | - sagemaker:*App 86 | - sagemaker:Search 87 | - sagemaker:RenderUiTemplate 88 | - sagemaker:BatchGetMetrics 89 | Resource: 90 | - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:*' 91 | Effect: Allow 92 | - Sid: AmazonSageMakerEC2ReadOnlyPolicy 93 | Action: 94 | - ec2:DescribeDhcpOptions 95 | - ec2:DescribeNetworkInterfaces 96 | - ec2:DescribeRouteTables 97 | - ec2:DescribeSecurityGroups 98 | - ec2:DescribeSubnets 99 | - ec2:DescribeVpcEndpoints 100 | - ec2:DescribeVpcs 101 | Resource: 102 | - !Sub 'arn:aws:ec2:*:${AWS::AccountId}:*' 103 | Effect: Allow 104 | - Sid: AmazonSageMakerIAMReadOnlyPolicy 105 | Action: 106 | - iam:ListRoles 107 | Resource: 108 | - !Sub 'arn:aws:iam::${AWS::AccountId}:*' 109 | Effect: Allow 110 | 111 | SageMakerAccessSupportingServicesPolicy: 112 | Type: AWS::IAM::ManagedPolicy 113 | Properties: 114 | Description: Read-only baseline policy for SageMaker execution role 115 | PolicyDocument: 116 | Version: 2012-10-17 117 | Statement: 118 | - Sid: AmazonSageMakerCRUDAccessS3Policy 119 | Action: 120 | - s3:PutObject 121 | - s3:GetObject 122 | - s3:AbortMultipartUpload 123 | - s3:DeleteObject 124 | - s3:CreateBucket 125 | - s3:ListBucket 126 | - s3:PutBucketCORS 127 | - s3:ListAllMyBuckets 128 | - s3:GetBucketCORS 129 | - s3:GetBucketLocation 130 | Resource: 131 | - arn:aws:s3:::*SageMaker* 132 | - arn:aws:s3:::*Sagemaker* 133 | - arn:aws:s3:::*sagemaker* 134 | Effect: Allow 135 | - Sid: AmazonSageMakerReadOnlyAccessKMSPolicy 136 | Action: 137 | - kms:DescribeKey 138 | - kms:ListAliases 139 | Resource: 140 | - !Sub 'arn:aws:kms:*:${AWS::AccountId}:*' 141 | Effect: Allow 142 | - Sid: AmazonSageMakerCRUDAccessECRPolicy 143 | Action: 144 | - ecr:Set* 145 | - ecr:CompleteLayerUpload 146 | - ecr:Batch* 147 | - ecr:Upload* 148 | - ecr:InitiateLayerUpload 149 | - ecr:Put* 150 | - ecr:Describe* 151 | - ecr:CreateRepository 152 | - ecr:Get* 153 | - ecr:StartImageScan 154 | Resource: 155 | - '*' 156 | Effect: Allow 157 | - Sid: AmazonSageMakerCRUDAccessCloudWatchPolicy 158 | Action: 159 | - cloudwatch:Put* 160 | - cloudwatch:Get* 161 | - cloudwatch:List* 162 | - cloudwatch:DescribeAlarms 163 | - logs:Put* 164 | - logs:Get* 165 | - logs:List* 166 | - logs:CreateLogGroup 167 | - logs:CreateLogStream 168 | - logs:ListLogDeliveries 169 | - logs:Describe* 170 | - logs:CreateLogDelivery 171 | - logs:PutResourcePolicy 172 | - logs:UpdateLogDelivery 173 | Resource: 174 | - '*' 175 | Effect: Allow 176 | 177 | SageMakerStudioDeveloperAccessPolicy: 178 | Type: AWS::IAM::ManagedPolicy 179 | Properties: 180 | Description: Read-only baseline policy for SageMaker execution role 181 | PolicyDocument: 182 | Version: 2012-10-17 183 | Statement: 184 | - Sid: AmazonSageMakerStudioCreateApp 185 | Action: 186 | - sagemaker:CreateApp 187 | Resource: 188 | - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:*' 189 | Effect: Allow 190 | - Sid: AmazonSageMakerStudioIAMPassRole 191 | Action: 192 | - iam:PassRole 193 | Resource: 194 | - !Sub 'arn:aws:iam::${AWS::AccountId}:role/*AmazonSageMaker*' 195 | Effect: Allow 196 | Condition: 197 | StringEquals: 198 | iam:PassedToService: sagemaker.amazonaws.com 199 | - Sid: AmazonSageMakerInvokeEndPointRole 200 | Action: 201 | - sagemaker:InvokeEndpoint 202 | Resource: 203 | - '*' 204 | Effect: Allow 205 | - Sid: AmazonSageMakerTags 206 | Action: 207 | - sagemaker:AddTags 208 | - sagemaker:ListTags 209 | Resource: 210 | - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:*' 211 | Effect: Allow 212 | - Sid: AmazonSageMakerCreate 213 | Action: 214 | - sagemaker:Create* 215 | Resource: 216 | - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:*' 217 | Effect: Allow 218 | Condition: 219 | ForAnyValue:StringEquals: 220 | aws:TagKeys: 221 | - Team 222 | StringEqualsIfExists: 223 | aws:RequestTag/Team: ${aws:PrincipalTag/Team} 224 | - Sid: AmazonSageMakerUpdateDeleteExecutePolicy 225 | Action: 226 | - sagemaker:Delete* 227 | - sagemaker:Stop* 228 | - sagemaker:Update* 229 | - sagemaker:Start* 230 | - sagemaker:DisassociateTrialComponent 231 | - sagemaker:AssociateTrialComponent 232 | - sagemaker:BatchPutMetrics 233 | Resource: 234 | - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:*' 235 | Effect: Allow 236 | Condition: 237 | StringEquals: 238 | aws:PrincipalTag/Team: ${sagemaker:ResourceTag/Team} 239 | 240 | SageMakerStudioExecutionRoleDefault: 241 | Type: 'AWS::IAM::Role' 242 | Properties: 243 | AssumeRolePolicyDocument: 244 | Version: 2012-10-17 245 | Statement: 246 | - Effect: Allow 247 | Principal: 248 | Service: 249 | - sagemaker.amazonaws.com 250 | Action: 251 | - 'sts:AssumeRole' 252 | Path: / 253 | ManagedPolicyArns: 254 | - !Ref SageMakerReadOnlyPolicy 255 | - !Ref SageMakerDeniedServicesPolicy 256 | Tags: 257 | - Key: EnvironmentName 258 | Value: !Ref EnvironmentName 259 | 260 | SageMakerStudioExecutionRoleTeam1: 261 | Type: 'AWS::IAM::Role' 262 | Properties: 263 | AssumeRolePolicyDocument: 264 | Version: 2012-10-17 265 | Statement: 266 | - Effect: Allow 267 | Principal: 268 | Service: 269 | - sagemaker.amazonaws.com 270 | Action: 271 | - 'sts:AssumeRole' 272 | Path: / 273 | ManagedPolicyArns: 274 | - !Ref SageMakerReadOnlyPolicy 275 | - !Ref SageMakerAccessSupportingServicesPolicy 276 | - !Ref SageMakerStudioDeveloperAccessPolicy 277 | - !Ref SageMakerDeniedServicesPolicy 278 | Tags: 279 | - Key: EnvironmentName 280 | Value: !Ref EnvironmentName 281 | - Key: Team 282 | Value: Team1 283 | 284 | SageMakerStudioExecutionRoleTeam2: 285 | Type: 'AWS::IAM::Role' 286 | Properties: 287 | AssumeRolePolicyDocument: 288 | Version: 2012-10-17 289 | Statement: 290 | - Effect: Allow 291 | Principal: 292 | Service: 293 | - sagemaker.amazonaws.com 294 | Action: 295 | - 'sts:AssumeRole' 296 | Path: / 297 | ManagedPolicyArns: 298 | - !Ref SageMakerReadOnlyPolicy 299 | - !Ref SageMakerAccessSupportingServicesPolicy 300 | - !Ref SageMakerStudioDeveloperAccessPolicy 301 | - !Ref SageMakerDeniedServicesPolicy 302 | Tags: 303 | - Key: EnvironmentName 304 | Value: !Ref EnvironmentName 305 | - Key: Team 306 | Value: Team2 307 | 308 | RestrictSageMakerToCIDRPolicy: 309 | Type: AWS::IAM::ManagedPolicy 310 | Properties: 311 | Path: / 312 | PolicyDocument: 313 | Version: 2012-10-17 314 | Statement: 315 | - Effect: Deny 316 | Action: 317 | - 'sagemaker:*' 318 | Resource: !Sub 'arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:user-profile/*/*' 319 | Condition: 320 | NotIpAddress: 321 | aws:VpcSourceIp: !If [RestrictSageMakerToCIDRCondition, !Ref AllowedCIDR, '0.0.0.0/0'] 322 | 323 | SageMakerPermissionsPolicy: 324 | Type: AWS::IAM::ManagedPolicy 325 | Properties: 326 | Path: / 327 | PolicyDocument: 328 | Version: 2012-10-17 329 | Statement: 330 | - Effect: Allow 331 | Action: 332 | - 'sagemaker:CreatePresignedDomainUrl' 333 | - 'sagemaker:ListUserProfiles' 334 | - 'sagemaker:CreateUserProfile' 335 | - 'sagemaker:DescribeUserProfile' 336 | - 'sagemaker:AddTags' 337 | - 'sagemaker:ListTags' 338 | Resource: !Sub 'arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:user-profile/*/*' 339 | 340 | SAMLBackendLambdaExecutionRole: 341 | Type: AWS::IAM::Role 342 | Properties: 343 | ManagedPolicyArns: 344 | - 'arn:aws:iam::aws:policy/CloudWatchLambdaInsightsExecutionRolePolicy' 345 | - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole' 346 | - !Ref SageMakerPermissionsPolicy 347 | - !If [RestrictSageMakerToCIDRCondition, !Ref RestrictSageMakerToCIDRPolicy, !Ref AWS::NoValue] 348 | AssumeRolePolicyDocument: 349 | Version: '2012-10-17' 350 | Statement: 351 | - 352 | Effect: Allow 353 | Principal: 354 | Service: 355 | - 'lambda.amazonaws.com' 356 | Action: 357 | - 'sts:AssumeRole' 358 | Policies: 359 | - PolicyName: PassExecutionRole 360 | PolicyDocument: 361 | Version: '2012-10-17' 362 | Statement: 363 | - Effect: Allow 364 | Action: 365 | - 'iam:PassRole' 366 | Resource: 367 | - !GetAtt SageMakerStudioExecutionRoleTeam1.Arn 368 | - !GetAtt SageMakerStudioExecutionRoleTeam2.Arn 369 | - PolicyName: AccessVPCResources 370 | PolicyDocument: 371 | Version: '2012-10-17' 372 | Statement: 373 | - Effect: Allow 374 | Action: 375 | - 'ec2:CreateNetworkInterface' 376 | - 'ec2:DescribeNetworkInterfaces' 377 | - 'ec2:DeleteNetworkInterface' 378 | - 'ec2:DescribeSecurityGroups' 379 | - 'ec2:DescribeSubnets' 380 | - 'ec2:DescribeVpcs' 381 | Resource: 382 | - '*' 383 | Tags: 384 | - Key: EnvironmentName 385 | Value: !Ref EnvironmentName 386 | 387 | SetupLambdaExecutionPolicy: 388 | Type: AWS::IAM::ManagedPolicy 389 | Properties: 390 | Path: / 391 | PolicyDocument: 392 | Version: 2012-10-17 393 | Statement: 394 | - Sid: SageMakerDomainPermission 395 | Effect: Allow 396 | Action: 397 | - sagemaker:ListDomains 398 | - sagemaker:DescribeDomain 399 | - sagemaker:UpdateDomain 400 | - sagemaker:ListUserProfiles 401 | - sagemaker:DescribeUserProfile 402 | - sagemaker:ListApps 403 | - sagemaker:DescribeApp 404 | - sagemaker:DeleteApp 405 | - sagemaker:UpdateApp 406 | - sagemaker:DeleteUserProfile 407 | Resource: 408 | - !Sub "arn:${AWS::Partition}:sagemaker:*:*:domain/*" 409 | - !Sub "arn:${AWS::Partition}:sagemaker:*:*:user-profile/*" 410 | - !Sub "arn:${AWS::Partition}:sagemaker:*:*:app/*" 411 | - Sid: SCPermissions 412 | Effect: Allow 413 | Action: 414 | - servicecatalog:Associate* 415 | - servicecatalog:Accept* 416 | - servicecatalog:Enable* 417 | - servicecatalog:Get* 418 | - servicecatalog:List* 419 | - servicecatalog:Describe* 420 | Resource: 421 | - !Sub 'arn:aws:servicecatalog:*:${AWS::AccountId}:*' 422 | - !Sub 'arn:aws:catalog:*:${AWS::AccountId}:*' 423 | - # Authorization strategy is ActionOnly for these two operations and require * in resource field 424 | Sid: SageMakerEnableSCPortfolio 425 | Effect: Allow 426 | Action: 427 | - sagemaker:EnableSagemakerServicecatalogPortfolio 428 | - sagemaker:DisableSagemakerServicecatalogPortfolio 429 | Resource: 430 | - '*' 431 | - Sid: SageMakerExecPassRole 432 | Effect: Allow 433 | Action: 434 | - iam:PassRole 435 | - iam:GetRole 436 | Resource: 437 | - !GetAtt SageMakerStudioExecutionRoleDefault.Arn 438 | - !GetAtt SageMakerStudioExecutionRoleTeam1.Arn 439 | - !GetAtt SageMakerStudioExecutionRoleTeam2.Arn 440 | 441 | SetupLambdaExecutionRole: 442 | Type: 'AWS::IAM::Role' 443 | Properties: 444 | AssumeRolePolicyDocument: 445 | Version: 2012-10-17 446 | Statement: 447 | - Effect: Allow 448 | Principal: 449 | Service: 450 | - lambda.amazonaws.com 451 | Action: 452 | - 'sts:AssumeRole' 453 | Path: / 454 | ManagedPolicyArns: 455 | - !Ref SetupLambdaExecutionPolicy 456 | - 'arn:aws:iam::aws:policy/AmazonEC2ReadOnlyAccess' 457 | - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole' 458 | Tags: 459 | - Key: EnvironmentName 460 | Value: !Ref EnvironmentName -------------------------------------------------------------------------------- /cfn-templates/sagemaker-domain.yaml: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | AWSTemplateFormatVersion: 2010-09-09 5 | Description: | 6 | Create a SageMaker Studio domain 7 | 8 | Metadata: 9 | AWS::CloudFormation::Interface: 10 | ParameterGroups: 11 | - Label: 12 | default: Data Science Environment 13 | Parameters: 14 | - EnvironmentName 15 | - Label: 16 | default: Amazon SageMaker Studio 17 | Parameters: 18 | - DomainName 19 | - AuthMode 20 | - Label: 21 | default: Network and Storage Configuration 22 | Parameters: 23 | - VPCId 24 | - SageMakerDomainSubnetIds 25 | - SageMakerDomainSecurityGroupIds 26 | - SageMakerDomainStorageKMSKeyId 27 | - Label: 28 | default: Permissions 29 | Parameters: 30 | - SageMakerDefaultExecutionRoleArn 31 | - SetupLambdaExecutionRoleArn 32 | 33 | ParameterLabels: 34 | EnvironmentName: 35 | default: Environment name 36 | DomainName: 37 | default: Domain name 38 | AuthMode: 39 | default: Authentication mode 40 | VPCId: 41 | default: VPC 42 | SageMakerDomainSubnetIds: 43 | default: Subnet(s) 44 | SageMakerDomainSecurityGroupIds: 45 | default: Security group(s) 46 | SageMakerDomainStorageKMSKeyId: 47 | default: Storage encryption key 48 | NetworkAccessType: 49 | default: Network access for SageMaker Studio 50 | SageMakerDefaultExecutionRoleArn: 51 | default: SageMaker default execution role 52 | SetupLambdaExecutionRoleArn: 53 | default: Execution role for setup Lambda function 54 | 55 | Outputs: 56 | SageMakerDomainId: 57 | Description: SageMaker Domain Id 58 | Value: !GetAtt SageMakerStudioDomain.DomainId 59 | Export: 60 | Name: 'ds-sagemaker-domain-id' 61 | 62 | Parameters: 63 | EnvironmentName: 64 | Type: String 65 | AllowedPattern: '[a-z0-9\-]*' 66 | Description: Please specify your SageMaker environment name. 67 | Default: 'sm-environment' 68 | 69 | DomainName: 70 | Type: String 71 | Description: SageMaker Studio domain name. Leave empty to auto generate. 72 | Default: '' 73 | 74 | VPCId: 75 | Type: AWS::EC2::VPC::Id 76 | Description: Choose a VPC for SageMaker Studio and SageMaker workloads 77 | 78 | SageMakerDomainSubnetIds: 79 | Type: List 80 | Description: Choose subnets or provide a comma-delimited list of subnet ids 81 | 82 | SageMakerDomainSecurityGroupIds: 83 | Type: List 84 | Description: Choose security groups for SageMaker Studio and SageMaker workloads 85 | 86 | SageMakerDomainStorageKMSKeyId: 87 | Type: String 88 | Description: SageMaker uses an AWS managed CMK to encrypt your EFS and EBS file systems by default. To use a customer managed CMK, enter its key Id. 89 | Default: '' 90 | 91 | NetworkAccessType: 92 | Type: String 93 | AllowedValues: 94 | - 'PublicInternetOnly' 95 | - 'VpcOnly' 96 | Description: Choose how SageMaker Studio accesses resources over the Network 97 | Default: 'VpcOnly' 98 | 99 | AuthMode: 100 | Type: String 101 | AllowedValues: 102 | - 'IAM' 103 | Description: The mode of authentication that members use to access the domain. Only IAM is supported for this solution. 104 | Default: 'IAM' 105 | 106 | SageMakerDefaultExecutionRoleArn: 107 | Type: String 108 | Description: The ARN of the SageMaker execution role 109 | 110 | SetupLambdaExecutionRoleArn: 111 | Type: String 112 | Description: The ARN of the execution role for the Lambda function for SageMaker Studio setup 113 | 114 | Conditions: 115 | GenerateDomainNameCondition: !Equals [ !Ref DomainName, '' ] 116 | SageMakerEFSKMSKeyCondition: !Not [ !Equals [ !Ref SageMakerDomainStorageKMSKeyId, ''] ] 117 | 118 | Resources: 119 | SageMakerStudioDomain: 120 | Type: AWS::SageMaker::Domain 121 | Properties: 122 | AppNetworkAccessType: !Ref NetworkAccessType 123 | AuthMode: !Ref AuthMode 124 | DefaultUserSettings: 125 | ExecutionRole: !Ref SageMakerDefaultExecutionRoleArn 126 | SecurityGroups: !Ref SageMakerDomainSecurityGroupIds 127 | DomainName: !If 128 | - GenerateDomainNameCondition 129 | - !Sub '${EnvironmentName}-${AWS::Region}-sagemaker-domain' 130 | - !Ref DomainName 131 | KmsKeyId: !If [ SageMakerEFSKMSKeyCondition, !Ref SageMakerDomainStorageKMSKeyId, !Ref 'AWS::NoValue' ] 132 | SubnetIds: !Ref SageMakerDomainSubnetIds 133 | VpcId: !Ref VPCId 134 | Tags: 135 | - Key: EnvironmentName 136 | Value: !Ref EnvironmentName 137 | 138 | EnableSageMakerProjects: 139 | Type: Custom::ResourceForEnablingSageMakerProjects 140 | DependsOn: SageMakerStudioDomain 141 | Properties: 142 | ServiceToken: !GetAtt EnableSageMakerProjectsLambda.Arn 143 | ExecutionRole: !Ref SageMakerDefaultExecutionRoleArn 144 | 145 | DeleteDomainApps: 146 | Type: Custom::DeleteDomainApps 147 | Properties: 148 | ServiceToken: !GetAtt DeleteDomainAppsLambda.Arn 149 | DomainId: !GetAtt SageMakerStudioDomain.DomainId 150 | 151 | DeleteDomainAppsLambda: 152 | Type: AWS::Lambda::Function 153 | Properties: 154 | ReservedConcurrentExecutions: 1 155 | Code: 156 | ZipFile: | 157 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 158 | # SPDX-License-Identifier: MIT-0 159 | 160 | import time 161 | import boto3 162 | import logging 163 | import json 164 | import cfnresponse 165 | from botocore.exceptions import ClientError 166 | 167 | sm_client = boto3.client('sagemaker') 168 | logger = logging.getLogger(__name__) 169 | 170 | def delete_user_profiles(domain_id): 171 | logger.info(f'Start deleting user profiles for domain id: {domain_id}') 172 | for p in sm_client.get_paginator('list_user_profiles').paginate(DomainIdEquals=domain_id): 173 | for up in p['UserProfiles']: 174 | if up['Status'] not in ('Deleting', 'Pending'): 175 | sm_client.delete_user_profile(DomainId=up['DomainId'], UserProfileName=up['UserProfileName']) 176 | 177 | up = 1 178 | while up: 179 | up = 0 180 | for p in sm_client.get_paginator('list_user_profiles').paginate(DomainIdEquals=domain_id): 181 | up += len([u['UserProfileName'] for u in p['UserProfiles'] if u['Status'] != 'Deleted']) 182 | logger.info(f'Number of active user profiles: {str(up)}') 183 | time.sleep(5) 184 | 185 | def delete_apps(domain_id): 186 | logger.info(f'Start deleting apps for domain id: {domain_id}') 187 | 188 | try: 189 | sm_client.describe_domain(DomainId=domain_id) 190 | except: 191 | logger.info(f'Cannot retrieve {domain_id}') 192 | return 193 | 194 | for p in sm_client.get_paginator('list_apps').paginate(DomainIdEquals=domain_id): 195 | for a in p['Apps']: 196 | if a['Status'] != 'Deleted': 197 | logger.info(f"Deleting {a['AppType']}:{a['AppName']}") 198 | sm_client.delete_app(DomainId=a['DomainId'], UserProfileName=a['UserProfileName'], AppType=a['AppType'], AppName=a['AppName']) 199 | 200 | apps = 1 201 | while apps: 202 | apps = 0 203 | for p in sm_client.get_paginator('list_apps').paginate(DomainIdEquals=domain_id): 204 | apps += len([a['AppName'] for a in p['Apps'] if a['Status'] != 'Deleted']) 205 | logger.info(f'Number of active apps: {str(apps)}') 206 | time.sleep(5) 207 | 208 | logger.info(f'Apps for {domain_id} deleted') 209 | return 210 | 211 | def lambda_handler(event, context): 212 | response_data = {} 213 | try: 214 | physicalResourceId = event.get('PhysicalResourceId') 215 | 216 | logger.info(json.dumps(event)) 217 | 218 | if event['RequestType'] in ['Create', 'Update']: 219 | physicalResourceId = event.get('ResourceProperties')['DomainId'] 220 | 221 | elif event['RequestType'] == 'Delete': 222 | delete_apps(physicalResourceId) 223 | delete_user_profiles(physicalResourceId) 224 | 225 | cfnresponse.send(event, context, cfnresponse.SUCCESS, response_data, physicalResourceId=physicalResourceId) 226 | 227 | except (Exception, ClientError) as exception: 228 | logger.error(exception) 229 | cfnresponse.send(event, context, cfnresponse.FAILED, response_data, physicalResourceId=physicalResourceId, reason=str(exception)) 230 | 231 | Description: Delete SageMaker domain apps to clean up 232 | Handler: index.lambda_handler 233 | MemorySize: 128 234 | Role: !Ref SetupLambdaExecutionRoleArn 235 | Runtime: python3.8 236 | Timeout: 900 237 | Tags: 238 | - Key: EnvironmentName 239 | Value: !Ref EnvironmentName 240 | 241 | EnableSageMakerProjectsLambda: 242 | Type: AWS::Lambda::Function 243 | DependsOn: SageMakerStudioDomain 244 | Properties: 245 | ReservedConcurrentExecutions: 1 246 | Code: 247 | ZipFile: | 248 | # Function: EnableSagemakerProjects 249 | # Purpose: Enables Sagemaker Projects 250 | import json 251 | import boto3 252 | import cfnresponse 253 | from botocore.exceptions import ClientError 254 | 255 | client = boto3.client('sagemaker') 256 | sc_client = boto3.client('servicecatalog') 257 | 258 | def lambda_handler(event, context): 259 | try: 260 | response_status = cfnresponse.SUCCESS 261 | 262 | if 'RequestType' in event and event['RequestType'] == 'Create': 263 | enable_sm_projects(event['ResourceProperties']['ExecutionRole']) 264 | cfnresponse.send(event, context, response_status, {}, '') 265 | except (Exception, ClientError) as exception: 266 | print(exception) 267 | cfnresponse.send(event, context, cfnresponse.FAILED, {}, physicalResourceId=event.get('PhysicalResourceId'), reason=str(exception)) 268 | 269 | def enable_sm_projects(studio_role_arn): 270 | # enable Project on account level (accepts portfolio share) 271 | response = client.enable_sagemaker_servicecatalog_portfolio() 272 | 273 | # associate studio role with portfolio 274 | response = sc_client.list_accepted_portfolio_shares() 275 | 276 | portfolio_id = '' 277 | for portfolio in response['PortfolioDetails']: 278 | if portfolio['ProviderName'] == 'Amazon SageMaker': 279 | portfolio_id = portfolio['Id'] 280 | 281 | response = sc_client.associate_principal_with_portfolio( 282 | PortfolioId=portfolio_id, 283 | PrincipalARN=studio_role_arn, 284 | PrincipalType='IAM' 285 | ) 286 | Description: Enable Sagemaker Projects 287 | Handler: index.lambda_handler 288 | MemorySize: 128 289 | Role: !Ref SetupLambdaExecutionRoleArn 290 | Runtime: python3.8 291 | Timeout: 60 292 | Tags: 293 | - Key: EnvironmentName 294 | Value: !Ref EnvironmentName 295 | 296 | # SSM parameter 297 | SageMakerDomainIdSSM: 298 | Type: 'AWS::SSM::Parameter' 299 | Properties: 300 | Name: !Sub "${EnvironmentName}-sagemaker-domain-id" 301 | Type: String 302 | Value: !GetAtt SageMakerStudioDomain.DomainId 303 | Description: !Sub 'SageMaker Studio domain id for ${SageMakerStudioDomain.DomainArn}' 304 | -------------------------------------------------------------------------------- /cfn-templates/vpc.yaml: -------------------------------------------------------------------------------- 1 | AWSTemplateFormatVersion: '2010-09-09' 2 | Description: Create a secure VPC environment for SageMaker domain and SAML backend Lambda function 3 | 4 | Outputs: 5 | VPCId: 6 | Value: !If [CreateVPCCondition, !Ref SageMakerVPC, !Ref ExistingVPCId] 7 | PublicSubnetId: 8 | Value: !If [CreateVPCCondition, !Ref PublicSubnet, ''] 9 | SageMakerDomainPrivateSubnetId: 10 | Value: !If 11 | - PrivateSubnetsCondition 12 | - !Join 13 | - ',' 14 | - - !Ref SageMakerDomainPrivateSubnet 15 | - !Ref ExistingSageMakerDomainPrivateSubnetId 16 | SAMLBackendPrivateSubnetId: 17 | Value: !If [PrivateSubnetsCondition, !Ref SAMLBackendPrivateSubnet, !Ref ExistingSAMLBackendPrivateSubnetId] 18 | SAMLBackendSecurityGroupId: 19 | Value: !GetAtt SAMLBackendSecurityGroup.GroupId 20 | SageMakerDomainSecurityGroupId: 21 | Value: !Join 22 | - ',' 23 | - - !GetAtt SageMakerDomainSecurityGroup.GroupId 24 | VPCESecurityGroupId: 25 | Value: !GetAtt VPCESecurityGroup.GroupId 26 | APIGatewayVPCE: 27 | Value: !Ref APIGatewayVPCE 28 | 29 | Parameters: 30 | EnvironmentName: 31 | Type: String 32 | 33 | ExistingVPCId: 34 | Type: String 35 | Default: '' 36 | Description: Enter an existing VPC Id for deployment or leave empty to create a new VPC 37 | 38 | VPCCIDR: 39 | AllowedPattern: ^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(\/(1[6-9]|2[0-8]))$ 40 | ConstraintDescription: CIDR block parameter must be in the form x.x.x.x/16-28 41 | Type: String 42 | Default: 10.0.0.0/16 43 | Description: CIDR block for a new or existing VPC, must always be provided 44 | 45 | CreatePrivateSubnets: 46 | AllowedValues: 47 | - 'YES' 48 | - 'NO' 49 | Default: 'YES' 50 | Description: Set to NO when you want to re-use existing subnets in the existing VPC (existing VPC Id must be provided) 51 | Type: String 52 | 53 | ExistingSAMLBackendPrivateSubnetId: 54 | Description: Existing private subnet id for SAML backend. Leave empty if CreatePrivateSubnets = YES 55 | Default: '' 56 | Type: String 57 | 58 | ExistingSageMakerDomainPrivateSubnetId: 59 | Description: Existing private subnet id for SageMaker domain. Leave empty if CreatePrivateSubnets = YES 60 | Default: '' 61 | Type: String 62 | 63 | SAMLBackendPrivateSubnetCIDR: 64 | AllowedPattern: ^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(\/(1[6-9]|2[0-8]))$ 65 | ConstraintDescription: CIDR block parameter must be in the form x.x.x.x/16-28 66 | Type: String 67 | Default: 10.0.0.0/19 68 | Description: CIDR block for a private subnet for SAML backend 69 | 70 | SageMakerDomainPrivateSubnetCIDR: 71 | AllowedPattern: ^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(\/(1[6-9]|2[0-8]))$ 72 | ConstraintDescription: CIDR block parameter must be in the form x.x.x.x/16-28 73 | Type: String 74 | Default: 10.0.32.0/19 75 | Description: CIDR block for a private subnet for SageMaker domain 76 | 77 | PublicSubnetCIDR: 78 | AllowedPattern: ^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(\/(1[6-9]|2[0-8]))$ 79 | ConstraintDescription: CIDR block parameter must be in the form x.x.x.x/16-28 80 | Type: String 81 | Default: 10.0.128.0/20 82 | Description: CIDR block for a public subnet for Internet and NAT Gateways 83 | 84 | Rules: 85 | CreateVPC: 86 | RuleCondition: !Equals [ !Ref ExistingVPCId, ''] 87 | Assertions: 88 | - Assert: !Equals [ !Ref 'CreatePrivateSubnets', 'YES' ] 89 | AssertDescription: Create private subnet must be set to YES if you select create a new VPC 90 | 91 | PrivateSubnets: 92 | RuleCondition: !Equals [ !Ref 'CreatePrivateSubnets', 'NO' ] 93 | Assertions: 94 | - Assert: !Not [!Equals [ !Ref 'ExistingVPCId', '' ]] 95 | AssertDescription: You can set Create private subnets to NO only if you provide an existing VPC id 96 | - Assert: !And 97 | - !Not [ !Equals [ !Ref 'ExistingSAMLBackendPrivateSubnetId', '' ] ] 98 | - !Not [ !Equals [ !Ref 'ExistingSageMakerDomainPrivateSubnetId', '' ] ] 99 | AssertDescription: You must provide existing private subnet ids for both SAML backend and SageMaker domain if you select not to create private subnets 100 | 101 | Conditions: 102 | CreateVPCCondition: !Equals [ !Ref ExistingVPCId, ''] 103 | PrivateSubnetsCondition: !Equals [!Ref 'CreatePrivateSubnets', 'YES'] 104 | 105 | Resources: 106 | 107 | ######## VPC / Subnets ######## 108 | SageMakerVPC: 109 | Type: AWS::EC2::VPC 110 | Condition: CreateVPCCondition 111 | DeletionPolicy: Retain 112 | UpdateReplacePolicy: Delete 113 | Properties: 114 | CidrBlock: !Ref VPCCIDR 115 | EnableDnsSupport: true 116 | EnableDnsHostnames: true 117 | Tags: 118 | - Key: Name 119 | Value: !Sub vpc-${EnvironmentName} 120 | - Key: EnvironmentName 121 | Value: !Ref EnvironmentName 122 | 123 | SAMLBackendPrivateSubnet: 124 | Type: AWS::EC2::Subnet 125 | Condition: PrivateSubnetsCondition 126 | Properties: 127 | VpcId: !If [CreateVPCCondition, !Ref SageMakerVPC, !Ref ExistingVPCId] 128 | CidrBlock: !Ref SAMLBackendPrivateSubnetCIDR 129 | MapPublicIpOnLaunch: false 130 | AvailabilityZone: !Sub "${AWS::Region}a" 131 | Tags: 132 | - Key: Name 133 | Value: !Sub private-sn-1a-${EnvironmentName}-saml-backend 134 | - Key: EnvironmentName 135 | Value: !Sub ${EnvironmentName} 136 | 137 | SageMakerDomainPrivateSubnet: 138 | Type: AWS::EC2::Subnet 139 | Condition: PrivateSubnetsCondition 140 | DeletionPolicy: Retain 141 | UpdateReplacePolicy: Delete 142 | Properties: 143 | VpcId: !If [CreateVPCCondition, !Ref SageMakerVPC, !Ref ExistingVPCId] 144 | CidrBlock: !Ref SageMakerDomainPrivateSubnetCIDR 145 | MapPublicIpOnLaunch: false 146 | AvailabilityZone: !Sub "${AWS::Region}a" 147 | Tags: 148 | - Key: Name 149 | Value: !Sub private-sn-1a-${EnvironmentName}-sm-domain 150 | - Key: EnvironmentName 151 | Value: !Sub ${EnvironmentName} 152 | 153 | PublicSubnet: 154 | Type: AWS::EC2::Subnet 155 | Condition: CreateVPCCondition 156 | Properties: 157 | VpcId: !If [CreateVPCCondition, !Ref SageMakerVPC, !Ref ExistingVPCId] 158 | CidrBlock: !Ref PublicSubnetCIDR 159 | MapPublicIpOnLaunch: true 160 | AvailabilityZone: !Sub "${AWS::Region}a" 161 | Tags: 162 | - Key: Name 163 | Value: !Sub public-sn-1a-${EnvironmentName} 164 | - Key: EnvironmentName 165 | Value: !Sub ${EnvironmentName} 166 | 167 | ####### SecurityGroup SageMaker ######## 168 | SageMakerDomainSecurityGroup: 169 | Type: AWS::EC2::SecurityGroup 170 | Properties: 171 | GroupDescription: SageMaker domain security group 172 | VpcId: !If [CreateVPCCondition, !Ref SageMakerVPC, !Ref ExistingVPCId] 173 | SecurityGroupEgress: 174 | - Description: All traffic is allowed outbound 175 | IpProtocol: '-1' 176 | CidrIp: 0.0.0.0/0 177 | Tags: 178 | - Key: Name 179 | Value: !Sub sg-${EnvironmentName}-sm-domain 180 | - Key: EnvironmentName 181 | Value: !Ref EnvironmentName 182 | 183 | SageMakerDomainSecurityGroupIngress: 184 | Type: AWS::EC2::SecurityGroupIngress 185 | Properties: 186 | Description: !Sub Allow inbound traffic from ${VPCCIDR} 187 | IpProtocol: "-1" 188 | FromPort: 0 189 | ToPort: 65535 190 | CidrIp: !Ref VPCCIDR 191 | GroupId: !Ref SageMakerDomainSecurityGroup 192 | 193 | ####### Security Group SAML backend ######## 194 | SAMLBackendSecurityGroup: 195 | Type: AWS::EC2::SecurityGroup 196 | Properties: 197 | GroupDescription: SAML backend security group 198 | VpcId: !If [CreateVPCCondition, !Ref SageMakerVPC, !Ref ExistingVPCId] 199 | SecurityGroupIngress: 200 | - Description: Allow ingress TCP 443 from SageMaker security group 201 | IpProtocol: tcp 202 | FromPort: 443 203 | ToPort: 443 204 | SourceSecurityGroupId: !GetAtt SageMakerDomainSecurityGroup.GroupId 205 | SecurityGroupEgress: 206 | - Description: All traffic is allowed outbound 207 | IpProtocol: '-1' 208 | CidrIp: 0.0.0.0/0 209 | Tags: 210 | - Key: Name 211 | Value: !Sub sg-${EnvironmentName}-saml-backend 212 | - Key: EnvironmentName 213 | Value: !Ref EnvironmentName 214 | 215 | SAMLBackendSecurityGroupIngress: 216 | Type: AWS::EC2::SecurityGroupIngress 217 | Properties: 218 | IpProtocol: '-1' 219 | CidrIp: !Ref VPCCIDR 220 | Description: !Sub Allow inbound traffic from ${VPCCIDR} 221 | GroupId: !GetAtt SAMLBackendSecurityGroup.GroupId 222 | 223 | ####### SecurityGroup VPCE ######## 224 | VPCESecurityGroup: 225 | Type: AWS::EC2::SecurityGroup 226 | Properties: 227 | GroupDescription: VPC endpoints security group 228 | VpcId: !If [CreateVPCCondition, !Ref SageMakerVPC, !Ref ExistingVPCId] 229 | SecurityGroupIngress: 230 | - Description: Allow ingress TCP 443 from SageMaker security group 231 | IpProtocol: tcp 232 | FromPort: 443 233 | ToPort: 443 234 | SourceSecurityGroupId: !GetAtt SageMakerDomainSecurityGroup.GroupId 235 | SecurityGroupEgress: 236 | - Description: All traffic is allowed outbound 237 | IpProtocol: '-1' 238 | CidrIp: 0.0.0.0/0 239 | Tags: 240 | - Key: Name 241 | Value: !Sub sg-${EnvironmentName}-vpce 242 | - Key: EnvironmentName 243 | Value: !Ref EnvironmentName 244 | 245 | VPCESecurityGroupVPCIngress: 246 | Type: AWS::EC2::SecurityGroupIngress 247 | Properties: 248 | CidrIp: !Ref VPCCIDR 249 | Description: !Sub Allow inbound traffic from from ${VPCCIDR} 250 | IpProtocol: '-1' 251 | GroupId: !GetAtt VPCESecurityGroup.GroupId 252 | 253 | ######## NAT Gateways/ IGW Gateway ######## 254 | InternetGateway: 255 | Type: AWS::EC2::InternetGateway 256 | Condition: CreateVPCCondition 257 | Properties: 258 | Tags: 259 | - Key: Name 260 | Value: !Sub igw-${EnvironmentName} 261 | - Key: EnvironmentName 262 | Value: !Ref EnvironmentName 263 | 264 | InternetGatewayAttachment: 265 | Type: AWS::EC2::VPCGatewayAttachment 266 | Condition: CreateVPCCondition 267 | Properties: 268 | InternetGatewayId: !Ref InternetGateway 269 | VpcId: !If [CreateVPCCondition, !Ref SageMakerVPC, !Ref ExistingVPCId] 270 | 271 | NatGateway: 272 | Type: AWS::EC2::NatGateway 273 | Condition: CreateVPCCondition 274 | Properties: 275 | AllocationId: !GetAtt NatGatewayEIP.AllocationId 276 | SubnetId: !Ref PublicSubnet 277 | Tags: 278 | - Key: Name 279 | Value: !Sub nat-gw-1-${EnvironmentName} 280 | - Key: EnvironmentName 281 | Value: !Ref EnvironmentName 282 | 283 | NatGatewayEIP: 284 | Type: AWS::EC2::EIP 285 | Condition: CreateVPCCondition 286 | DependsOn: InternetGatewayAttachment 287 | Properties: 288 | Domain: vpc 289 | 290 | ######## Route Tables ######## 291 | PublicRouteTable: 292 | Type: AWS::EC2::RouteTable 293 | Condition: CreateVPCCondition 294 | Properties: 295 | VpcId: !If [CreateVPCCondition, !Ref SageMakerVPC, !Ref ExistingVPCId] 296 | Tags: 297 | - Key: Name 298 | Value: !Sub public-rtb-${EnvironmentName} 299 | - Key: EnvironmentName 300 | Value: !Sub ${EnvironmentName} Public Routes 301 | 302 | PublicSubnetRoute: 303 | Type: AWS::EC2::Route 304 | Condition: CreateVPCCondition 305 | DependsOn: InternetGatewayAttachment 306 | Properties: 307 | RouteTableId: !Ref PublicRouteTable 308 | DestinationCidrBlock: 0.0.0.0/0 309 | GatewayId: !Ref InternetGateway 310 | 311 | PublicSubnetRouteTableAssociation: 312 | Type: AWS::EC2::SubnetRouteTableAssociation 313 | Condition: CreateVPCCondition 314 | Properties: 315 | RouteTableId: !Ref PublicRouteTable 316 | SubnetId: !Ref PublicSubnet 317 | 318 | PrivateRouteTable: 319 | Type: AWS::EC2::RouteTable 320 | Condition: PrivateSubnetsCondition 321 | Properties: 322 | VpcId: !If [CreateVPCCondition, !Ref SageMakerVPC, !Ref ExistingVPCId] 323 | Tags: 324 | - Key: Name 325 | Value: !Sub private-rtb-1a-${EnvironmentName} 326 | - Key: EnvironmentName 327 | Value: !Sub ${EnvironmentName} Private Routes 328 | 329 | PrivateSubnetRoute: 330 | Type: AWS::EC2::Route 331 | Condition: CreateVPCCondition 332 | Properties: 333 | RouteTableId: !If [PrivateSubnetsCondition, !Ref PrivateRouteTable, !Ref AWS::NoValue] 334 | DestinationCidrBlock: 0.0.0.0/0 335 | NatGatewayId: !Ref NatGateway 336 | 337 | SAMLBackendPrivateSubnetRouteTableAssociation: 338 | Type: AWS::EC2::SubnetRouteTableAssociation 339 | Condition: PrivateSubnetsCondition 340 | Properties: 341 | RouteTableId: !Ref PrivateRouteTable 342 | SubnetId: !If [PrivateSubnetsCondition, !Ref SAMLBackendPrivateSubnet, !Ref ExistingSAMLBackendPrivateSubnetId] 343 | 344 | SageMakerDomainPrivateSubnetRouteTableAssociation: 345 | Type: AWS::EC2::SubnetRouteTableAssociation 346 | Condition: PrivateSubnetsCondition 347 | Properties: 348 | RouteTableId: !Ref PrivateRouteTable 349 | SubnetId: !If [PrivateSubnetsCondition, !Ref SageMakerDomainPrivateSubnet, !Ref ExistingSageMakerDomainPrivateSubnetId] 350 | 351 | ######## VPC Endpoints ######## 352 | APIGatewayVPCE: 353 | Type: AWS::EC2::VPCEndpoint 354 | Properties: 355 | ServiceName: !Sub "com.amazonaws.${AWS::Region}.execute-api" 356 | PrivateDnsEnabled: true 357 | VpcId: !If [CreateVPCCondition, !Ref SageMakerVPC, !Ref ExistingVPCId] 358 | SubnetIds: 359 | - !If [PrivateSubnetsCondition, !Ref SAMLBackendPrivateSubnet, !Ref ExistingSAMLBackendPrivateSubnetId] 360 | VpcEndpointType: Interface 361 | SecurityGroupIds: 362 | - !GetAtt VPCESecurityGroup.GroupId 363 | 364 | SageMakerStudioVPCE: 365 | Type: AWS::EC2::VPCEndpoint 366 | Properties: 367 | PrivateDnsEnabled: true 368 | ServiceName: !Sub "aws.sagemaker.${AWS::Region}.studio" 369 | VpcId: !If [CreateVPCCondition, !Ref SageMakerVPC, !Ref ExistingVPCId] 370 | SubnetIds: 371 | - !If [PrivateSubnetsCondition, !Ref SageMakerDomainPrivateSubnet, !Ref ExistingSageMakerDomainPrivateSubnetId] 372 | VpcEndpointType: Interface 373 | SecurityGroupIds: 374 | - !GetAtt VPCESecurityGroup.GroupId 375 | 376 | SageMakerAPIVPCE: 377 | Type: AWS::EC2::VPCEndpoint 378 | Properties: 379 | PrivateDnsEnabled: true 380 | ServiceName: !Sub "com.amazonaws.${AWS::Region}.sagemaker.api" 381 | VpcId: !If [CreateVPCCondition, !Ref SageMakerVPC, !Ref ExistingVPCId] 382 | SubnetIds: 383 | - !If [PrivateSubnetsCondition, !Ref SageMakerDomainPrivateSubnet, !Ref ExistingSageMakerDomainPrivateSubnetId] 384 | VpcEndpointType: Interface 385 | SecurityGroupIds: 386 | - !GetAtt VPCESecurityGroup.GroupId 387 | 388 | SageMakerRuntimeVPCE: 389 | Type: AWS::EC2::VPCEndpoint 390 | Properties: 391 | PrivateDnsEnabled: true 392 | ServiceName: !Sub "com.amazonaws.${AWS::Region}.sagemaker.runtime" 393 | VpcId: !If [CreateVPCCondition, !Ref SageMakerVPC, !Ref ExistingVPCId] 394 | SubnetIds: 395 | - !If [PrivateSubnetsCondition, !Ref SageMakerDomainPrivateSubnet, !Ref ExistingSageMakerDomainPrivateSubnetId] 396 | VpcEndpointType: Interface 397 | SecurityGroupIds: 398 | - !GetAtt VPCESecurityGroup.GroupId 399 | 400 | -------------------------------------------------------------------------------- /design/aws-sso-idp-synchronization.drawio.svg: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /design/network-architecture.drawio.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 |
15 |
16 |
17 | SAML backend private subnet 18 |
19 |
20 |
21 | 22 | 23 | SAML backend private subnet 24 | 25 | 26 | 27 | 28 | 29 | 30 |
31 |
32 |
33 | SAML backend Lambda function 34 |
35 |
36 |
37 | 38 | 39 | SAML backend Lambd... 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | SAML backend security group 49 | 50 | 51 | 52 | 53 | 54 | 55 |
56 |
57 |
58 | ``` 59 |
60 |
61 |
62 | 63 | 64 | ``` 65 | 66 | 67 | 68 | 69 | 70 | 71 |
72 |
73 |
74 | API Gateway endpoint 75 |
76 |
77 |
78 | 79 | 80 | API Gatewa... 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 |
90 |
91 |
92 | SageMaker private subnet 93 |
94 |
95 |
96 | 97 | 98 | SageMaker private subnet 99 | 100 | 101 | 102 | 103 | 104 | 105 | 106 |
107 |
108 |
109 | SageMaker Studio Notebook ENI 110 |
111 |
112 |
113 | 114 | 115 | SageMaker Studio Not... 116 | 117 | 118 | 119 | 120 | 121 | 122 | VPCE security group 123 | 124 | 125 | 126 | 127 | 128 | AWS Account 129 | 130 | 131 | 132 | 133 | 134 | 135 | 136 | 137 |
138 |
139 |
140 | VPC 141 |
142 |
143 |
144 | 145 | 146 | VPC 147 | 148 | 149 | 150 | 151 | 152 | 153 |
154 |
155 |
156 | SageMaker Studio endpoint 157 |
158 |
159 |
160 | 161 | 162 | SageMaker... 163 | 164 | 165 | 166 | 167 | 168 | 169 | 170 |
171 |
172 |
173 | SageMaker API 174 |
175 | endpoint 176 |
177 |
178 |
179 | 180 | 181 | SageMaker... 182 | 183 | 184 | 185 | 186 | 187 | 188 | 189 | 190 |
191 |
192 |
193 | SageMaker runtime 194 |
195 | endpoint 196 |
197 |
198 |
199 | 200 | 201 | SageMaker... 202 | 203 | 204 | 205 | 206 | 207 | 208 | SageMaker security group 209 | 210 | 211 | 212 | 213 | 214 | 215 |
216 |
217 |
218 | Internet Gateway 219 |
220 |
221 |
222 | 223 | 224 | Internet G... 225 | 226 | 227 | 228 | 229 | 230 | 231 | 232 | 233 |
234 |
235 |
236 | Public subnet 237 |
238 |
239 |
240 | 241 | 242 | Public subnet 243 | 244 | 245 | 246 | 247 | 248 | 249 | 250 |
251 |
252 |
253 | NAT Gateway 254 |
255 |
256 |
257 | 258 | 259 | NAT Gateway 260 | 261 | 262 | 263 | 264 | 265 | 266 | 267 | 268 | Viewer does not support full SVG 1.1 269 | 270 | 271 | 272 | -------------------------------------------------------------------------------- /design/solution-flow.drawio.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | AWS account 25 | 26 | 27 | 28 | 29 | 30 | AWS Single Sign-On 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 |
39 |
40 |
41 | SAML backend Lambda function 42 |
43 |
44 |
45 | 46 | 47 | SAML backend Lambd... 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 |
57 |
58 |
59 | Amazon SageMaker API 60 |
61 |
62 |
63 | 64 | 65 | Amazon SageMaker A... 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 |
74 |
75 |
76 | AWS IAM policy 77 |
78 |
79 |
80 | 81 | 82 | AWS IAM policy 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 |
92 |
93 |
94 | call CreatePresignedDomainUrl() 95 |
96 |
97 |
98 | 99 | 100 | call CreatePresignedDomainUrl() 101 | 102 | 103 | 104 | 105 | 106 | 107 | 108 | 109 | 110 |
111 |
112 |
113 | Attribute mappings: 114 |
115 | 116 | ssouserid 117 |
118 | teamid 119 | 120 |
121 |
122 |
123 | 124 | 125 | Attribute mappings:... 126 | 127 | 128 | 129 | 130 | 131 | 132 |
133 |
134 |
135 | SAML app 136 |
137 |
138 |
139 | 140 | 141 | SAML app 142 | 143 | 144 | 145 | 146 | 147 | 148 | 149 | 150 | 151 |
152 |
153 |
154 | SAML assertion 155 |
156 |
157 |
158 | 159 | 160 | SAML assertion 161 | 162 | 163 | 164 | 165 | 166 | 167 | 168 | 169 |
170 |
171 |
172 | HTTP 302: https://presigned-url 173 |
174 |
175 |
176 | 177 | 178 | HTTP 302: https://presign... 179 | 180 | 181 | 182 | 183 | 184 | 185 | 186 | 187 | 188 | 189 | 190 | 191 |
192 |
193 |
194 | 1 195 |
196 |
197 |
198 | 199 | 200 | 1 201 | 202 | 203 | 204 | 205 | 206 | 207 | 208 |
209 |
210 |
211 | 2 212 |
213 |
214 |
215 | 216 | 217 | 2 218 | 219 | 220 | 221 | 222 | 223 | 224 | 225 |
226 |
227 |
228 | 3 229 |
230 |
231 |
232 | 233 | 234 | 3 235 | 236 | 237 | 238 | 239 | 240 | 241 | 242 |
243 |
244 |
245 | 4 246 |
247 |
248 |
249 | 250 | 251 | 4 252 | 253 | 254 | 255 | 256 | 257 | 258 | 259 | 260 | 261 | 262 |
263 |
264 |
265 | Amazon API Gateway 266 |
267 |
268 |
269 | 270 | 271 | Amazon API Gateway 272 | 273 | 274 | 275 | 276 | 277 | 278 | 279 |
280 |
281 |
282 | SAML backend API 283 |
284 |
285 |
286 | 287 | 288 | SAML backend API 289 | 290 | 291 | 292 | 293 | 294 | 295 | 296 | 297 | 298 | 299 |
300 |
301 |
302 | VPC 303 |
304 |
305 |
306 | 307 | 308 | VPC 309 | 310 | 311 | 312 | 313 | 314 | 315 |
316 |
317 |
318 | sagemaker.api VPC endpoint 319 |
320 |
321 |
322 | 323 | 324 | sagemaker.api VPC... 325 | 326 | 327 | 328 | 329 | 330 | 331 | 332 | Amazon SageMaker Domain 333 | 334 | 335 | 336 | 337 | 338 |
339 |
340 |
341 | ssouserid-teamid profile 342 |
343 |
344 |
345 | 346 | 347 | ssouserid-teamid p... 348 | 349 | 350 | 351 | 352 | 353 | 354 | 355 | 356 |
357 |
358 |
359 | profiles are dynamically generated 360 |
361 |
362 |
363 | 364 | 365 | profiles are dynamica... 366 | 367 | 368 | 369 | 370 | 371 | 372 | 373 |
374 |
375 |
376 | Lambda function IAM execution role 377 |
378 |
379 |
380 | 381 | 382 | Lambda function IA... 383 | 384 | 385 | 386 | 387 | 388 | 389 | 390 | 391 | 392 | 393 | 394 | 395 | Viewer does not support full SVG 1.1 396 | 397 | 398 | 399 | -------------------------------------------------------------------------------- /functions/inline/get_network_config.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | import json 5 | import boto3 6 | import cfnresponse 7 | from botocore.exceptions import ClientError 8 | 9 | ec2 = boto3.resource("ec2") 10 | 11 | def lambda_handler(event, context): 12 | try: 13 | response_status = cfnresponse.SUCCESS 14 | r = {} 15 | 16 | if 'RequestType' in event and event['RequestType'] == 'Create': 17 | r["VPCCIDR"] = get_vpc_cidr( 18 | event['ResourceProperties']['VPCId'] 19 | ) 20 | 21 | cfnresponse.send(event, context, response_status, r, '') 22 | 23 | except ClientError as exception: 24 | print(exception) 25 | cfnresponse.send(event, context, cfnresponse.FAILED, {}, physicalResourceId=event.get('PhysicalResourceId'), reason=str(exception)) 26 | 27 | def get_vpc_cidr(vpc_id): 28 | print(vpc_id) 29 | 30 | return ec2.Vpc(vpc_id).cidr_block -------------------------------------------------------------------------------- /functions/inline/get_user_profile_metadata.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | import json 5 | import boto3 6 | import cfnresponse 7 | from botocore.exceptions import ClientError 8 | 9 | def lambda_handler(event, context): 10 | try: 11 | response_status = cfnresponse.SUCCESS 12 | r = {} 13 | 14 | if 'RequestType' in event and event['RequestType'] == 'Create': 15 | r["Metadata"] = json.dumps(event['ResourceProperties']['Metadata']) 16 | 17 | cfnresponse.send(event, context, response_status, r, '') 18 | 19 | except ClientError as exception: 20 | print(exception) 21 | cfnresponse.send(event, context, cfnresponse.FAILED, {}, physicalResourceId=event.get('PhysicalResourceId'), reason=str(exception)) -------------------------------------------------------------------------------- /functions/requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/users-and-team-management-with-amazon-sagemaker-and-aws-sso/12ba96a3516e182e9b9eaf34e95221833848cf03/functions/requirements.txt -------------------------------------------------------------------------------- /functions/saml-backend/requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/users-and-team-management-with-amazon-sagemaker-and-aws-sso/12ba96a3516e182e9b9eaf34e95221833848cf03/functions/saml-backend/requirements.txt -------------------------------------------------------------------------------- /functions/saml-backend/saml_backend_function.py: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | """ 5 | This is **a non-production sample** of a SAML backend 6 | The Lambda function parses the SAML assertion and uses only attributes in `` element to construct a `CreatePresignedDomainUrl` API call. 7 | In your production solution you must use a proper SAML backend implementation which must include a validation of a SAML response, a signature, and certificates, replay and redirect prevention, and any other features of a SAML authentication process. 8 | For example, you can use a [python3-saml SAML backend implementation](https://python-social-auth.readthedocs.io/en/latest/backends/saml.html) or 9 | [OneLogin open source SAML toolkit](https://developers.onelogin.com/saml/python) to implement a secure SAML backend. 10 | """ 11 | 12 | import json 13 | import os 14 | import boto3 15 | import logging 16 | import json 17 | import botocore.exceptions 18 | import base64 19 | import urllib 20 | import time 21 | from xml.dom import minidom 22 | 23 | try: 24 | logger = logging.getLogger(__name__) 25 | logging.root.setLevel(os.environ.get("LOG_LEVEL", "INFO")) 26 | sm = boto3.client("sagemaker") 27 | except Exception as e: 28 | print(f"Exception in initializing block: {e}") 29 | 30 | try: 31 | user_profile_metadata = json.loads(os.environ.get("USER_PROFILE_METADATA", "{}")) 32 | except Exception as e: 33 | print(f"Exception in loading user profile metadata: {e}") 34 | 35 | HTTP_REDIRECT = 302 36 | HTTP_EXCEPTION = 400 37 | KEY_NAME_USER_ID = os.environ.get("KEY_NAME_USER_ID", "ssouserid") 38 | KEY_NAME_TEAM_ID = os.environ.get("KEY_NAME_TEAM_ID", "teamid") 39 | 40 | def get_xml(body): 41 | saml_response=urllib.parse.unquote(body.split("&")[0].split("=")[1]) 42 | return base64.b64decode(saml_response) 43 | 44 | def get_saml_attributes(saml_response_xml): 45 | return { 46 | e.attributes["Name"].value:e.getElementsByTagName("saml2:AttributeValue")[0].childNodes[0].nodeValue 47 | for e in minidom.parseString(saml_response_xml).getElementsByTagName("saml2:Attribute") 48 | } 49 | 50 | def get_user_profile_name(user_id, team_id): 51 | return f"{user_id}-{team_id}" 52 | 53 | def get_user_profile_metadata(team): 54 | """ 55 | This function can be implemented as a microservice to return user and team metadata 56 | """ 57 | return user_profile_metadata.get(team) 58 | 59 | def create_presigned_domain_url(user_profile_name, metadata, expires_in_seconds=5): 60 | """ 61 | This function can be implemented as a microservice to manage studio user profiles 62 | """ 63 | user_profiles = sm.list_user_profiles( 64 | DomainIdEquals=metadata["DomainId"], 65 | UserProfileNameContains=user_profile_name)["UserProfiles"] 66 | 67 | if len(user_profiles) > 1: 68 | raise Exception(f"{user_profile_name} contains in more than one user profile for domain {metadata['DomainId']}") 69 | 70 | if not len(user_profiles): 71 | r = sm.create_user_profile( 72 | DomainId=metadata["DomainId"], 73 | UserProfileName=user_profile_name, 74 | Tags=metadata["Tags"], 75 | UserSettings=metadata["UserSettings"] 76 | ) 77 | 78 | while sm.describe_user_profile( 79 | DomainId=metadata["DomainId"], 80 | UserProfileName=user_profile_name)["Status"] != "InService": time.sleep(3) 81 | 82 | return sm.create_presigned_domain_url( 83 | DomainId=metadata["DomainId"], 84 | UserProfileName=user_profile_name, 85 | SessionExpirationDurationInSeconds=int(metadata["SessionExpiration"]), 86 | ExpiresInSeconds=expires_in_seconds 87 | )["AuthorizedUrl"] 88 | 89 | def lambda_handler(event, context): 90 | try: 91 | logger.info(json.dumps(event)) 92 | 93 | body = event.get("body") 94 | if not body: 95 | raise Exception("No body key in the request") 96 | 97 | attr_dict = get_saml_attributes(get_xml(body)) 98 | user_profile_name = get_user_profile_name( 99 | attr_dict[KEY_NAME_USER_ID], 100 | attr_dict[KEY_NAME_TEAM_ID] 101 | ) 102 | 103 | logger.info(f"Got team {attr_dict[KEY_NAME_TEAM_ID]} and constructed user profile name={user_profile_name}") 104 | 105 | try: 106 | metadata = get_user_profile_metadata(attr_dict[KEY_NAME_TEAM_ID]) 107 | if not metadata: 108 | raise Exception(f"no user profile metadata found for team {attr_dict[KEY_NAME_TEAM_ID]}") 109 | 110 | response = { 111 | "statusCode": HTTP_REDIRECT, 112 | "headers": { 113 | "Location": create_presigned_domain_url( 114 | user_profile_name, 115 | metadata, 116 | int(os.environ.get("PRESIGNED_URL_EXPIRATION", 5)) 117 | ) 118 | }, 119 | "isBase64Encoded": False 120 | } 121 | 122 | except botocore.exceptions.ClientError as ce: 123 | logger.error(f"ClientError exception in CreatePresignedDomainUrl: {ce}") 124 | response = { 125 | "statusCode": ce.response["ResponseMetadata"]["HTTPStatusCode"], 126 | "headers": {}, 127 | "body": ce.response["Error"]["Message"], 128 | "isBase64Encoded": False 129 | } 130 | except Exception as e: 131 | logger.error(f"Exception in CreatePresignedDomainUrl: {e}") 132 | response = { 133 | "statusCode": HTTP_EXCEPTION, 134 | "headers": {}, 135 | "body": json.dumps(str(e)), 136 | "isBase64Encoded": False 137 | } 138 | 139 | return response 140 | 141 | except Exception as e: 142 | logger.error(f"Exception in function body: {e}") 143 | raise e -------------------------------------------------------------------------------- /img/control-panel-profiles.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/users-and-team-management-with-amazon-sagemaker-and-aws-sso/12ba96a3516e182e9b9eaf34e95221833848cf03/img/control-panel-profiles.png -------------------------------------------------------------------------------- /img/signing-to-studio-2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/users-and-team-management-with-amazon-sagemaker-and-aws-sso/12ba96a3516e182e9b9eaf34e95221833848cf03/img/signing-to-studio-2.png -------------------------------------------------------------------------------- /img/signing-to-studio.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/users-and-team-management-with-amazon-sagemaker-and-aws-sso/12ba96a3516e182e9b9eaf34e95221833848cf03/img/signing-to-studio.png -------------------------------------------------------------------------------- /img/sso-app-mapping.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/users-and-team-management-with-amazon-sagemaker-and-aws-sso/12ba96a3516e182e9b9eaf34e95221833848cf03/img/sso-app-mapping.png -------------------------------------------------------------------------------- /img/sso-app-user.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/users-and-team-management-with-amazon-sagemaker-and-aws-sso/12ba96a3516e182e9b9eaf34e95221833848cf03/img/sso-app-user.png -------------------------------------------------------------------------------- /img/sso-app.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/users-and-team-management-with-amazon-sagemaker-and-aws-sso/12ba96a3516e182e9b9eaf34e95221833848cf03/img/sso-app.png -------------------------------------------------------------------------------- /img/sso-custom-apps-2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/users-and-team-management-with-amazon-sagemaker-and-aws-sso/12ba96a3516e182e9b9eaf34e95221833848cf03/img/sso-custom-apps-2.png -------------------------------------------------------------------------------- /img/sso-custom-apps.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/users-and-team-management-with-amazon-sagemaker-and-aws-sso/12ba96a3516e182e9b9eaf34e95221833848cf03/img/sso-custom-apps.png -------------------------------------------------------------------------------- /img/sso-store-id.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/users-and-team-management-with-amazon-sagemaker-and-aws-sso/12ba96a3516e182e9b9eaf34e95221833848cf03/img/sso-store-id.png -------------------------------------------------------------------------------- /img/starting-sm-studio.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/users-and-team-management-with-amazon-sagemaker-and-aws-sso/12ba96a3516e182e9b9eaf34e95221833848cf03/img/starting-sm-studio.png -------------------------------------------------------------------------------- /img/studio-exec-role.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/users-and-team-management-with-amazon-sagemaker-and-aws-sso/12ba96a3516e182e9b9eaf34e95221833848cf03/img/studio-exec-role.png -------------------------------------------------------------------------------- /template.yaml: -------------------------------------------------------------------------------- 1 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: MIT-0 3 | 4 | AWSTemplateFormatVersion: '2010-09-09' 5 | Transform: AWS::Serverless-2016-10-31 6 | Description: > 7 | Amazon SageMaker SAML backend. This stack creates a SAML backend API and Lambda functions to suppport a custom SAML application 8 | 9 | Metadata: 10 | AWS::ServerlessRepo::Application: 11 | Name: amazon-sagemaker-saml-backend 12 | Description: Custom SAML application backend for Amazon SageMaker Studio 13 | Author: ilyiny 14 | SpdxLicenseId: MIT 15 | LicenseUrl: ./LICENSE 16 | ReadmeUrl: ./README.md 17 | Labels: ['amazon', 'sagemaker', 'studio', 'saml', 'sso'] 18 | HomePageUrl: https://gitlab.aws.dev/ilyiny/amazon-sagemaker-team-and-user-management-sso 19 | SemanticVersion: 1.0.0 20 | SourceCodeUrl: https://gitlab.aws.dev/ilyiny/amazon-sagemaker-team-and-user-management-sso 21 | 22 | Globals: 23 | Function: 24 | Runtime: python3.8 25 | MemorySize: 128 26 | Timeout: 60 27 | Environment: 28 | Variables: 29 | LOG_LEVEL: INFO 30 | Tags: 31 | Project: amazon-sagemaker-saml-backend 32 | 33 | Outputs: 34 | SageMakerDomainId: 35 | Description: SageMaker Domain Id 36 | Value: !If 37 | - CreateSageMakerDomainCondition 38 | - !GetAtt SageMakerDomain.Outputs.SageMakerDomainId 39 | - !Ref SageMakerDomainId 40 | Export: 41 | Name: 'sagemaker-domain-id' 42 | 43 | SAMLBackendEndpoint: 44 | Description: "API Gateway endpoint URL acting as the Application ACS URL" 45 | Value: !Sub "https://${SageMakerDomainSAMLAPI}.execute-api.${AWS::Region}.amazonaws.com/prod/saml" 46 | 47 | SAMLAudience: 48 | Description: "Application SAML audience" 49 | Value: !Sub "https://${SageMakerDomainSAMLAPI}.execute-api.${AWS::Region}.amazonaws.com/" 50 | 51 | SageMakerStudioExecutionRoleTeam1Arn: 52 | Description: The ARN of the SageMaker Team1 execution role 53 | Value: !GetAtt IAM.Outputs.SageMakerStudioExecutionRoleTeam1Arn 54 | 55 | SageMakerStudioExecutionRoleTeam2Arn: 56 | Description: The ARN of the SageMaker Team2 execution role 57 | Value: !GetAtt IAM.Outputs.SageMakerStudioExecutionRoleTeam2Arn 58 | 59 | Parameters: 60 | EnvironmentName: 61 | Type: String 62 | AllowedPattern: '[a-z0-9\-]*' 63 | Default: 'sagemaker-team-mgmt-sso' 64 | Description: Your Amazon SageMaker environment name 65 | 66 | SageMakerDomainId: 67 | Type: String 68 | AllowedPattern: '[a-z0-9\-]*' 69 | Default: '' 70 | Description: Existing Amazon SageMaker domain id. Leave empty to create a new domain. 71 | 72 | CreatePrivateSubnets: 73 | AllowedValues: 74 | - 'YES' 75 | - 'NO' 76 | Default: 'YES' 77 | Description: Set to NO when you want to re-use existing subnets in the existing VPC (existing VPC Id must be provided in ExistingVPCId). 78 | Type: String 79 | 80 | ExistingVPCId: 81 | Type: String 82 | Description: Choose a VPC for SageMaker Studio. Leave empty to create a new VPC. 83 | Default: '' 84 | 85 | SAMLBackendSubnetId: 86 | Type: String 87 | Description: Existing private subnet id for SAML backend. Leave empty if CreatePrivateSubnets = YES 88 | Default: '' 89 | 90 | SageMakerDomainSubnetId: 91 | Type: String 92 | Description: Choose subnet for SageMaker domain. Leave empty if if CreatePrivateSubnets = YES 93 | Default: '' 94 | 95 | VPCCIDR: 96 | AllowedPattern: ^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(\/(1[6-9]|2[0-8]))$ 97 | ConstraintDescription: CIDR block parameter must be in the form x.x.x.x/16-28 98 | Type: String 99 | Default: 10.0.0.0/16 100 | Description: CIDR block for a new VPC 101 | 102 | SAMLBackendPrivateSubnetCIDR: 103 | AllowedPattern: ^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(\/(1[6-9]|2[0-8]))$ 104 | ConstraintDescription: CIDR block parameter must be in the form x.x.x.x/16-28 105 | Type: String 106 | Default: 10.0.0.0/19 107 | Description: CIDR block for a private subnet for SAML backend 108 | 109 | SageMakerDomainPrivateSubnetCIDR: 110 | AllowedPattern: ^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(\/(1[6-9]|2[0-8]))$ 111 | ConstraintDescription: CIDR block parameter must be in the form x.x.x.x/16-28 112 | Type: String 113 | Default: 10.0.32.0/19 114 | Description: CIDR block for a private subnet for SageMaker domain 115 | 116 | PublicSubnetCIDR: 117 | AllowedPattern: ^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(\/(1[6-9]|2[0-8]))$ 118 | ConstraintDescription: CIDR block parameter must be in the form x.x.x.x/16-28 119 | Type: String 120 | Default: 10.0.128.0/20 121 | Description: CIDR block for a public subnet for Internet and NAT Gateways 122 | 123 | DomainAccessAllowedCIDR: 124 | Type: String 125 | Default: '' 126 | Description: Allowed CIDR block for CreatePresignedDomainURL API call. Leave empty to allow access from public internet 127 | 128 | Conditions: 129 | ExistingVPCCondition: !Not [ !Equals [ !Ref ExistingVPCId, ''] ] 130 | CreateSageMakerDomainCondition: !Equals [ !Ref SageMakerDomainId, '' ] 131 | 132 | Resources: 133 | 134 | ######## CF Stacks ######## 135 | VPC: 136 | Type: AWS::CloudFormation::Stack 137 | Properties: 138 | TemplateURL: cfn-templates/vpc.yaml 139 | Parameters: 140 | EnvironmentName : !Ref EnvironmentName 141 | ExistingVPCId: !Ref ExistingVPCId 142 | VPCCIDR: !If [ ExistingVPCCondition, !GetAtt GetNetworkConfiguration.VPCCIDR, !Ref VPCCIDR ] 143 | CreatePrivateSubnets: !Ref CreatePrivateSubnets 144 | ExistingSAMLBackendPrivateSubnetId: !Ref SAMLBackendSubnetId 145 | ExistingSageMakerDomainPrivateSubnetId: !Ref SageMakerDomainSubnetId 146 | SAMLBackendPrivateSubnetCIDR: !Ref SAMLBackendPrivateSubnetCIDR 147 | SageMakerDomainPrivateSubnetCIDR: !Ref SageMakerDomainPrivateSubnetCIDR 148 | PublicSubnetCIDR: !Ref PublicSubnetCIDR 149 | Tags: 150 | - Key: EnvironmentName 151 | Value: !Ref EnvironmentName 152 | 153 | IAM: 154 | Type: AWS::CloudFormation::Stack 155 | Properties: 156 | TemplateURL: cfn-templates/iam.yaml 157 | Parameters: 158 | EnvironmentName : !Ref EnvironmentName 159 | AllowedCIDR: !Ref DomainAccessAllowedCIDR 160 | Tags: 161 | - Key: EnvironmentName 162 | Value: !Ref EnvironmentName 163 | 164 | SageMakerDomain: 165 | Type: AWS::CloudFormation::Stack 166 | Condition: CreateSageMakerDomainCondition 167 | Properties: 168 | TemplateURL: cfn-templates/sagemaker-domain.yaml 169 | Parameters: 170 | EnvironmentName: !Ref EnvironmentName 171 | VPCId: !GetAtt VPC.Outputs.VPCId 172 | SageMakerDomainSubnetIds: !GetAtt VPC.Outputs.SageMakerDomainPrivateSubnetId 173 | SageMakerDomainSecurityGroupIds: !GetAtt VPC.Outputs.SageMakerDomainSecurityGroupId 174 | SageMakerDefaultExecutionRoleArn: !GetAtt IAM.Outputs.SageMakerStudioExecutionRoleDefaultArn 175 | SetupLambdaExecutionRoleArn: !GetAtt IAM.Outputs.SetupLambdaExecutionRoleArn 176 | Tags: 177 | - Key: EnvironmentName 178 | Value: !Ref EnvironmentName 179 | 180 | GetNetworkConfiguration: 181 | Type: Custom::GetNetworkConfiguration 182 | Condition: ExistingVPCCondition 183 | Properties: 184 | ServiceToken: !GetAtt GetNetworkConfigurationFunction.Arn 185 | VPCId: !Ref ExistingVPCId 186 | 187 | GetUserProfileMetadata: 188 | Type: Custom::GetUserProfileMetadata 189 | Properties: 190 | ServiceToken: !GetAtt GetUserProfileMetadataFunction.Arn 191 | Metadata: 192 | Team1: 193 | DomainId: !If 194 | - CreateSageMakerDomainCondition 195 | - !GetAtt SageMakerDomain.Outputs.SageMakerDomainId 196 | - !Ref SageMakerDomainId 197 | SessionExpiration: 43200 198 | Tags: 199 | - Key: Team 200 | Value: Team1 201 | UserSettings: 202 | ExecutionRole: !GetAtt IAM.Outputs.SageMakerStudioExecutionRoleTeam1Arn 203 | Team2: 204 | DomainId: !If 205 | - CreateSageMakerDomainCondition 206 | - !GetAtt SageMakerDomain.Outputs.SageMakerDomainId 207 | - !Ref SageMakerDomainId 208 | SessionExpiration: 43200 209 | Tags: 210 | - Key: Team 211 | Value: Team2 212 | UserSettings: 213 | ExecutionRole: !GetAtt IAM.Outputs.SageMakerStudioExecutionRoleTeam2Arn 214 | 215 | ######## API Gateway ######## 216 | SageMakerDomainSAMLAPI: 217 | Type: AWS::Serverless::Api 218 | Properties: 219 | StageName: prod 220 | EndpointConfiguration: 221 | Type: REGIONAL 222 | Auth: 223 | ResourcePolicy: 224 | CustomStatements: { 225 | Effect: 'Allow', 226 | Action: 'execute-api:Invoke', 227 | Resource: ['execute-api:/*/*/*'], 228 | Principal: '*' 229 | } 230 | 231 | ######## Lambda Functions ######## 232 | GetNetworkConfigurationFunction: 233 | Type: AWS::Serverless::Function 234 | Properties: 235 | ReservedConcurrentExecutions: 1 236 | InlineCode: | 237 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 238 | # SPDX-License-Identifier: MIT-0 239 | 240 | import json 241 | import boto3 242 | import cfnresponse 243 | from botocore.exceptions import ClientError 244 | 245 | ec2 = boto3.resource("ec2") 246 | 247 | def lambda_handler(event, context): 248 | try: 249 | response_status = cfnresponse.SUCCESS 250 | r = {} 251 | 252 | if 'RequestType' in event and event['RequestType'] == 'Create': 253 | r["VPCCIDR"] = get_vpc_cidr( 254 | event['ResourceProperties']['VPCId'] 255 | ) 256 | 257 | cfnresponse.send(event, context, response_status, r, '') 258 | 259 | except ClientError as exception: 260 | print(exception) 261 | cfnresponse.send(event, context, cfnresponse.FAILED, {}, physicalResourceId=event.get('PhysicalResourceId'), reason=str(exception)) 262 | 263 | def get_vpc_cidr(vpc_id): 264 | print(vpc_id) 265 | 266 | return ec2.Vpc(vpc_id).cidr_block 267 | 268 | Role: !GetAtt IAM.Outputs.SetupLambdaExecutionRoleArn 269 | Handler: index.lambda_handler 270 | 271 | GetUserProfileMetadataFunction: 272 | Type: AWS::Serverless::Function 273 | Properties: 274 | ReservedConcurrentExecutions: 1 275 | InlineCode: | 276 | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 277 | # SPDX-License-Identifier: MIT-0 278 | 279 | import json 280 | import boto3 281 | import cfnresponse 282 | from botocore.exceptions import ClientError 283 | 284 | def lambda_handler(event, context): 285 | try: 286 | response_status = cfnresponse.SUCCESS 287 | r = {} 288 | 289 | if 'RequestType' in event and event['RequestType'] == 'Create': 290 | r["Metadata"] = json.dumps(event['ResourceProperties']['Metadata']) 291 | 292 | cfnresponse.send(event, context, response_status, r, '') 293 | 294 | except ClientError as exception: 295 | print(exception) 296 | cfnresponse.send(event, context, cfnresponse.FAILED, {}, physicalResourceId=event.get('PhysicalResourceId'), reason=str(exception)) 297 | 298 | Handler: index.lambda_handler 299 | Policies: 300 | - AWSLambdaBasicExecutionRole 301 | 302 | SAMLBackEndFunction: 303 | Type: AWS::Serverless::Function 304 | Properties: 305 | ReservedConcurrentExecutions: 1 306 | CodeUri: functions/saml-backend/ 307 | Role: !GetAtt IAM.Outputs.SAMLBackendLambdaExecutionRoleArn 308 | Handler: saml_backend_function.lambda_handler 309 | Environment: 310 | Variables: 311 | PRESIGNED_URL_EXPIRATION: 5 312 | USER_PROFILE_METADATA: !GetAtt GetUserProfileMetadata.Metadata 313 | KEY_NAME_USER_ID: 'ssouserid' 314 | KEY_NAME_TEAM_ID: 'teamid' 315 | Events: 316 | SAMLBackEnd: 317 | Type: Api 318 | Properties: 319 | RestApiId: !Ref SageMakerDomainSAMLAPI 320 | Path: /saml 321 | Method: POST 322 | VpcConfig: 323 | SecurityGroupIds: 324 | - !GetAtt VPC.Outputs.SAMLBackendSecurityGroupId 325 | SubnetIds: 326 | - !GetAtt VPC.Outputs.SAMLBackendPrivateSubnetId 327 | 328 | -------------------------------------------------------------------------------- /test/cfn-unit-test.sh: -------------------------------------------------------------------------------- 1 | # Test CFN template isolated 2 | 3 | # set variables 4 | export ENV_NAME="sagemaker-team-mgmt-sso" 5 | 6 | ########## VPC ########## 7 | # use cases: 8 | # 1. New VPC - the stack creates all network infrastructure 9 | # 2. Existing VPC, new private subnets - the stack creates private subnets, security groups, and VPC endpoints 10 | # 3. Existing VPC, existing private subnets - the stack creates security groups and VPC endpoints 11 | 12 | aws cloudformation validate-template \ 13 | --template-body file://cfn-templates/vpc.yaml 14 | 15 | # 1. New VPC 16 | aws cloudformation deploy \ 17 | --template-file cfn-templates/vpc.yaml \ 18 | --stack-name $ENV_NAME-vpc \ 19 | --parameter-overrides \ 20 | EnvironmentName=$ENV_NAME 21 | 22 | # 2. Existing VPC, new private subnets 23 | export VPC_ID=vpc-c513e9b8 24 | export VPCCIDR=$(aws ec2 describe-vpcs --filters Name="vpc-id",Values=$VPC_ID \ 25 | --output text \ 26 | --query 'Vpcs[0].CidrBlock') 27 | 28 | # CIDR blocks for private subnets 29 | export SAMLBACKEND_SN_CIDR="172.31.96.0/19" 30 | export SMDOMAIN_SN_CIDR="172.31.128.0/19" 31 | 32 | aws cloudformation deploy \ 33 | --template-file cfn-templates/vpc.yaml \ 34 | --stack-name $ENV_NAME-vpc \ 35 | --parameter-overrides \ 36 | EnvironmentName=$ENV_NAME \ 37 | ExistingVPCId=$VPC_ID \ 38 | VPCCIDR=$VPCCIDR \ 39 | SAMLBackendPrivateSubnetCIDR=$SAMLBACKEND_SN_CIDR \ 40 | SageMakerDomainPrivateSubnetCIDR=$SMDOMAIN_SN_CIDR 41 | 42 | # 3. Existing VPC, existing private subnets 43 | export SAMLBACKEND_SN_ID=subnet-0ded16799f1a327d6 44 | export SMDOMAIN_SN_ID=subnet-0cbd39976e7509a90 45 | 46 | aws cloudformation deploy \ 47 | --template-file cfn-templates/vpc.yaml \ 48 | --stack-name $ENV_NAME-vpc \ 49 | --parameter-overrides \ 50 | EnvironmentName=$ENV_NAME \ 51 | ExistingVPCId=$VPC_ID \ 52 | VPCCIDR=$VPCCIDR \ 53 | CreatePrivateSubnets=NO \ 54 | ExistingSAMLBackendPrivateSubnetId=$SAMLBACKEND_SN_ID \ 55 | ExistingSageMakerDomainPrivateSubnetId=$SMDOMAIN_SN_ID 56 | 57 | ########## IAM ########## 58 | export ALLOWED_CIDR=172.31.0.0/16 59 | 60 | aws cloudformation deploy \ 61 | --template-file cfn-templates/iam.yaml \ 62 | --stack-name $ENV_NAME-iam \ 63 | --capabilities CAPABILITY_NAMED_IAM \ 64 | --parameter-overrides \ 65 | EnvironmentName=$ENV_NAME \ 66 | AllowedCIDR=$ALLOWED_CIDR 67 | 68 | ########## SageMaker Domain ########## 69 | 70 | ########## 71 | export VPC_ID= 72 | export VPCCIDR=$(aws ec2 describe-vpcs --filters Name="vpc-id",Values=$VPC_ID \ 73 | --output text \ 74 | --query 'Vpcs[0].CidrBlock') 75 | 76 | 77 | export DOMAIN_ID=$(aws sagemaker list-domains --output text --query 'Domains[0].DomainId') 78 | export SUBNET_IDS=$(aws sagemaker describe-domain --domain-id $DOMAIN_ID --output text --query 'SubnetIds[*]') 79 | 80 | aws ec2 describe-subnets \ 81 | --subnet-ids ${SUBNET_IDS} \ 82 | --filters Name="availability-zone",Values=${AWS_DEFAULT_REGION}a \ 83 | --output text \ 84 | --query 'Subnets[].CidrBlock' 85 | 86 | # SageMaker domain authentication mode 87 | aws sagemaker describe-domain --domain-id $DOMAIN_ID --output text --query 'AuthMode' 88 | 89 | # Get the domain id 90 | export DOMAIN_ID=$(aws sagemaker list-domains --output text --query 'Domains[0].DomainId') 91 | 92 | # Get the execution roles 93 | export STACK_NAME= 94 | export EXEC_ROLE_TEAM1=$(aws cloudformation describe-stacks --stack-name $STACK_NAME | jq -r '.Stacks[].Outputs[] | select(.OutputKey=="SageMakerStudioExecutionRoleTeam1Arn") | .OutputValue') 95 | export EXEC_ROLE_TEAM2=$(aws cloudformation describe-stacks --stack-name $STACK_NAME | jq -r '.Stacks[].Outputs[] | select(.OutputKey=="SageMakerStudioExecutionRoleTeam2Arn") | .OutputValue') 96 | 97 | # Get SSO user id 98 | export SSO_STORE_ID= 99 | export SSO_USER1_NAME= 100 | export SSO_USER2_NAME=g 101 | 102 | export SSO_USER1_ID=$(aws identitystore list-users --identity-store-id $SSO_STORE_ID --filter AttributePath='UserName',AttributeValue=$SSO_USER1_NAME --query 'Users[0].UserId' --output text) 103 | export SSO_USER2_ID=$(aws identitystore list-users --identity-store-id $SSO_STORE_ID --filter AttributePath='UserName',AttributeValue=$SSO_USER2_NAME --query 'Users[0].UserId' --output text) 104 | 105 | # Create Studio user profiles 106 | aws sagemaker create-user-profile \ 107 | --domain-id $DOMAIN_ID \ 108 | --user-profile-name $SSO_USER1_ID-Team1 \ 109 | --tags Key=studiouserid,Value=ilyiny+demo@amazon.com \ 110 | --user-settings ExecutionRole=$EXEC_ROLE_TEAM1 111 | 112 | aws sagemaker create-user-profile \ 113 | --domain-id $DOMAIN_ID \ 114 | --user-profile-name $SSO_USER1_ID-Team2 \ 115 | --tags Key=studiouserid,Value=ilyiny+demo@amazon.com \ 116 | --user-settings ExecutionRole=$EXEC_ROLE_TEAM2 117 | 118 | aws sagemaker create-user-profile \ 119 | --domain-id $DOMAIN_ID \ 120 | --user-profile-name $SSO_USER2_ID-Team2 \ 121 | --tags Key=studiouserid,Value=ilyiny+demo-sm-sso-2@amazon.com \ 122 | --user-settings ExecutionRole=$EXEC_ROLE_TEAM2 123 | 124 | # List the tags assigned to a user profile 125 | export USER_PROFILE_ARN=$(aws sagemaker describe-user-profile \ 126 | --user-profile-name \ 127 | --domain-id $DOMAIN_ID \ 128 | --output text --query 'UserProfileArn') 129 | 130 | aws sagemaker list-tags --resource-arn $USER_PROFILE_ARN 131 | -------------------------------------------------------------------------------- /user-profile-metadata.json: -------------------------------------------------------------------------------- 1 | { 2 | "Team1":{ 3 | "DomainId":"" 13 | } 14 | }, 15 | "Team2":{ 16 | "DomainId":"" 26 | } 27 | } 28 | } --------------------------------------------------------------------------------