Secure SageMaker Unified Studio access with SAML and conditional policies

0
1
Secure SageMaker Unified Studio access with SAML and conditional policies


Amazon SageMaker Unified Studio is a single data and AI development environment that brings together data preparation, analytics, and machine learning (ML) development in one place. By unifying these workflows, it saves teams from managing multiple tools and makes it straightforward for data scientists, analysts, and developers to build, train, and deploy ML models while collaborating. In Amazon SageMaker Unified Studio, a domain is the organizing entity for connecting your assets, users, and their projects. With Amazon SageMaker unified domains, you have the flexibility to reflect the data and analytics needs of your organizational structure. You can create a single unified domain for your enterprise or multiple domains for different business units.

Some enterprises, especially those in regulated industries, might require limiting access to trusted networks (such as VPN CIDRs) or to managed devices that meet compliance standards through device attestation.

In this post, we demonstrate how to integrate SageMaker Unified Studio as a custom SAML application and apply conditional access policies for enforcing device compliance, IP-based restrictions, or multi-factor authentication (MFA). For this post, we use Okta as the identity provider (IdP).

Solution overview

This solution demonstrates how to integrate Amazon SageMaker Unified Studio (SMUS) with external SAML identity providers such as Okta. The integration enforces enterprise security controls, including trusted network access, device compliance, and multi-factor authentication. With this integration, organizations in regulated industries can maintain strict access controls while providing single sign-on for their data science and AI development teams. By using SAML 2.0 federation with conditional access policies, you can help make sure that only authenticated users on compliant devices from trusted networks gain access. This access applies to your SageMaker Unified Studio domains and the associated data and AI workloads.

SAML authentication flow from a corporate device through the identity provider and AWS STS to Amazon SageMaker Unified Studio

Authentication flow for accessing SageMaker Unified Studio through SAML

The architecture diagram illustrates the secure authentication flow for accessing SageMaker Unified Studio through SAML integration:

  1. Users typically initiate access from corporate-managed devices through VPN or trusted network connections.
  2. The IdP authenticates the user and evaluates conditional access policies defined by your organization. Based on these policies, it checks for trusted devices, approved source IP ranges, and MFA completion. If any policy fails, the login is rejected. Otherwise, authentication proceeds.
  3. Upon successful authentication and policy validation, the IdP generates a digitally signed SAML assertion containing user attributes and group memberships, securely delivering it to the user’s browser through HTTP POST binding.
  4. The client browser automatically posts the SAML assertion to the AWS Security Token Service (AWS STS) sign-in endpoint. There, the AWS IAM Identity Provider validates the trust relationship with your corporate IdP through pre-configured SAML federation settings.
  5. AWS STS validates the SAML assertion signature and authenticity. It then maps the user attributes to a specifically configured IAM role with SageMaker Unified Studio permissions, including the datazone:GetIamPortalLoginUrl permission required for domain access.
  6. AWS STS confirms successful role assumption and generates temporary AWS credentials with a defined session duration. It then issues an HTTP redirect that returns the browser to the SageMaker Unified Studio domain with authenticated session tokens.
  7. Users gain access to the unified environment for data preparation, analytics, and machine learning development. All activities are governed by the assumed IAM role permissions and logged for comprehensive audit trails.

Walkthrough

In this walkthrough, you create a SAML application in Okta, connect it to AWS, and configure a SageMaker Unified Studio domain to use it for authentication.

Prerequisites

Before you get started, make sure you have the following:

  1. Familiarity with Amazon SageMaker Unified Studio.
  2. A basic understanding of SAML 2.0.
  3. AWS Identity and Access Management (IAM) permissions to create a domain in Amazon SageMaker Unified Studio.
  4. Access to your SAML IdP (such as Okta or Entra ID) to create and configure a SAML application.

Step 1: Create an application in Okta

The first step is to set up a new SAML application in Okta that manages authentication for SMUS.

  1. In Okta, go to ApplicationsCreate App Integration, and choose SAML 2.0.
  2. Provide an App name.
  3. Set the Single sign-on URL to https://signin.aws.amazon.com/saml.
  4. Set Name ID format to Persistent.
  5. Set the Audience URI (SP Entity ID) to https://signin.aws.amazon.com/saml.
  6. Choose Next, and finish creating the application.
  7. Once created, copy the Metadata URL and Sign On URL. You need these in later steps.

Step 2: Create an identity provider in IAM

Now, let’s connect Okta to AWS by creating an IAM identity provider. This allows AWS to trust authentication responses from Okta.

  1. Open the IAM console.
  2. Go to Identity providersAdd provider.
  3. Select SAML as the provider type.
  4. Provide a Provider name.
  5. In Okta, go to your application’s Sign On tab, choose Identity Provider metadata, and save the XML file. Upload it here.
  6. Choose Add provider.
  7. Copy the ARN of this provider. You need it when you create the role.

Step 3: Create an IAM role for Okta

Next, create an IAM role that Okta can assume. This role defines what access users have when they sign in through Okta.

  1. In IAM, go to RolesCreate role.
  2. Use the following trust policy (replace both instances of “{Replace with Identity provider ARN}” with the ARN you copied in Step 2):
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Federated": "{Replace with Identity provider ARN}"
            },
            "Action": "sts:AssumeRoleWithSAML",
            "Condition": {
                "StringEquals": {
                    "SAML:aud": "https://signin.aws.amazon.com/saml"
                }
            }
        },
        {
            "Effect": "Allow",
            "Principal": {
                "Federated": "{Replace with Identity provider ARN}"
            },
            "Action": "sts:TagSession",
            "Condition": {
                "StringLike": {
                    "aws:RequestTag/Email": "*"
                }
            }
        }
    ]
}

  1. Attach a permission policy. For example:
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "VisualEditor0",
            "Effect": "Allow",
            "Action": "datazone:GetIamPortalLoginUrl",
            "Resource": "arn:aws:datazone:<REGION>:<ACCOUNT-ID>:domain/<DOMAIN-ID>"
        }
    ]
}

Replace <REGION>, <ACCOUNT-ID>, and <DOMAIN-ID> with the corresponding values from your SageMaker Unified Studio domain ARN (arn:aws:sagemaker:<REGION>:<ACCOUNT-ID>:domain/<DOMAIN-ID>). You can find the domain ARN in the SageMaker console under Domains.

Step 4: Configure SAML assertions

To make sure AWS understands who is signing in, configure the SAML assertions in Okta.

  1. Open your application in Okta.
  2. Go to GeneralSAML SettingsEdit.
  3. Choose Next until you reach Attribute Statements.
  4. Add the following mappings:
    • https://aws.amazon.com/SAML/Attributes/PrincipalTag:Emailuser.email.
    • https://aws.amazon.com/SAML/Attributes/Role{IAMROLEARN,IdentityProviderARN}.
    • https://aws.amazon.com/SAML/Attributes/RoleSessionNameuser.email.

Step 5: Create an SMUS domain

Finally, let’s set up the SMUS domain and tie it all together.

Note: Creating a SageMaker Unified Studio domain incurs charges. For pricing details, see the Amazon SageMaker pricing page.

  1. Open the Amazon SageMaker console.
  2. Choose Create domain.
  3. Choose Manual setup (this allows for SAML integration).
  4. Enter a domain name, then choose Create.
  5. In Configure SSO user access, select SAML, then choose Next.
  6. Set the IdP SSO URL to the Sign On URL from Step 1.
  7. Select Do not require assignments. (Access is instead managed by your IdP team through Okta or Entra.)
  8. Choose Next, then choose Save.

To verify the integration works, open your SMUS domain and choose Sign in with SSO. You are redirected to Okta, and conditional access policies such as VPN, device attestation, or MFA apply automatically.

  1. Open your SMUS domain URL in a browser.
  2. Choose Sign in with SSO.
  3. Confirm that you are redirected to Okta for authentication.
  4. Sign in with your Okta credentials.
  5. Verify that you are redirected back to the SMUS domain with access to your projects.

Step 6: Assign users to the Okta application

Before users can authenticate through Okta to access SMUS, you must assign them to the application.

  1. In Okta, navigate to your SAML application.
  2. Go to the Assignments tab.
  3. Choose Assign, and select Assign to People or Assign to Groups.
  4. Select the users or groups who need access to SMUS.
  5. Choose Save and Go Back, then choose Done.

Step 7: Apply conditional access policies

Up to Step 5, we configured SMUS with an external SAML IdP. At this point, anyone assigned to the new application in your IdP can sign in and access the SMUS domain.

This is where conditional access policies come into play. Based on your organization’s governance model, you can add policies in your IdP to further control how and when users gain access. For example:

  • Restricting access to specific corporate IP address ranges (for example, only through VPN).
  • Enforcing device compliance so that only managed or secure devices can connect.
  • Adding MFA requirements for sensitive actions.
  • Applying device attestation to help assess whether the endpoint conforms to security baselines.

Most major IdPs, including Okta and Entra ID, support conditional access. You can find more details in their documentation:

These policies allow you to enforce the right level of protection, from something as simple as requiring users to connect through corporate networks to something as advanced as verifying device attestation across your fleet.

Clean up

To avoid incurring ongoing charges, delete the resources you created during this walkthrough:

  1. Delete the Amazon SageMaker Unified Studio domain from the SageMaker console.
  2. Delete the IAM role you created for Okta.
  3. Delete the IAM identity provider.
  4. Delete the SAML application in Okta.

Important: Deleting the SMUS domain permanently removes all projects, assets, and data within it. Back up any important work before proceeding.

Conclusion

By integrating SMUS with an external IdP through SAML, you can help enforce modern access controls based on your organization’s security requirements. This post walked through how to configure SMUS with a custom SAML application and pointed you toward resources for setting up conditional access policies.

With conditional access in place, you can decide, based on your organization’s needs, whether access should be limited to trusted users on trusted networks, trusted devices, or both. This approach can help provide a more secure and compliant login experience that aligns SMUS access with your company’s broader identity and security strategy.


About the authors

Amit Samal

Amit Samal

Amit is a Sr. Delivery Consultant in World Wide Public Sector, Professional Services at AWS working with UKGI Customers. Amit has been with AWS for about 4 years and has been helping customers across the UKGI to design & implement secure, resilient and cost-effective workloads on AWS. Amit is passionate about all areas of technology, but has focus areas in Networking, Migrations, and Application Modernizations.

Manos Samatas

Manos Samatas

Manos is a Principal Solutions Architect in Data and AI with Amazon Web Services. He works with government, non-profit, education and healthcare customers in the UK on data and AI projects, helping build solutions using AWS. Manos lives and works in London. In his spare time, he enjoys reading, watching sports, playing video games and socialising with friends.