Tag Archives: Amazon SageMaker Unified Studio

Enforce IAM permissions boundaries for Amazon SageMaker Unified Studio Tooling blueprints

Post Syndicated from Sanjana Sekar original https://aws.amazon.com/blogs/big-data/enforce-iam-permissions-boundaries-for-amazon-sagemaker-unified-studio-tooling-blueprints/

Amazon SageMaker Unified Studio now supports custom permissions boundaries for IAM roles created by the Tooling blueprint. Organizations that enforce Service Control Policies (SCPs) requiring permissions boundaries on all AWS Identity and Access Management (IAM) roles can now adopt Amazon SageMaker Unified Studio without modifying their security posture.

Amazon SageMaker Unified Studio is a unified development environment that brings together data engineering, machine learning, and analytics tools into a single workspace. In Amazon SageMaker Unified Studio, a project is a collaborative workspace that bundles people, tools, and access permissions together. It builds every project from a project profile, which defines a list of blueprints. Blueprints are pre-configured infrastructure templates that provision AWS resources at project creation time or on demand, along with their default parameters. The Tooling blueprint is the only mandatory one. Amazon SageMaker Unified Studio deploys it with every project, creating foundational resources such as the project IAM role and security groups.

In this post, you learn how to create a permissions boundary that restricts AI agent capabilities. You then configure it on the Tooling blueprint using the AWS Command Line Interface (AWS CLI). Finally, you validate that the boundary is enforced on all provisioned roles.

The problem

Enterprises in regulated industries use SCPs to require that every IAM role in an account carries a permissions boundary. A well-scoped boundary prevents privilege escalation and verifies no role exceeds the maximum permissions defined by the organization’s security team. Before this feature, Amazon SageMaker Unified Studio Tooling blueprints created IAM roles without permissions boundaries. When an SCP enforced permissions boundaries, project creation failed with an explicit deny:

User: arn:aws:sts::<account-id>:assumed-role/AmazonSageMakerProvisioning-<account-id>/AmazonDataZoneEnvironmentDeployer-<account-id> is not authorized to perform: iam:CreateRole on resource: arn:aws:iam::<account-id>:role/AmazonBedrockServiceRole-<project-id>-<env-id> with an explicit deny in a service control policy

Amazon SageMaker Unified Studio surfaces the blocked role creation as a Tooling environment provisioning failure, as shown in Figure 1.

SMUS project overview showing the Tooling environment in a failed state from a permissions boundary SCP denial

Figure 1: Project creation fails when the SCP requires a permissions boundary that is not attached

The project is marked as failed because its Tooling environment couldn’t deploy in the US East (N. Virginia) AWS Region (us-east-1). The details show a 403 permissions error, while the preceding IAM message identifies the underlying iam:CreateRole SCP denial. This blocked adoption for any organization with SCP-enforced permissions boundaries. The AWS CloudFormation event for the Tooling stack exposes the IAM failure behind the project-level error, as shown in Figure 2.

CloudFormation stack events showing the BedrockServiceRole in CREATE_FAILED from an iam:CreateRole SCP explicit deny

Figure 2: Detailed error showing the SCP denial in the Tooling blueprint AWS CloudFormation stack

The AmazonBedrockServiceRole resource entered CREATE_FAILED because iam:CreateRole was explicitly denied by the SCP, even though AWS CloudFormation surfaced the wrapper error as UnauthorizedTaggingOperation.

Granular control using a permissions boundary: Example use case

Beyond satisfying SCP requirements, permissions boundaries give administrators granular control over what the Tooling blueprint roles can do. For instance, some organizations have SecOps policies that require disabling Data Agent and Data Notebook capabilities across their accounts. These organizations want project members to access data connections and run SQL queries directly, but must block conversational AI, code generation, and notebook cell execution through the agent.

When PermissionsBoundaryArn is configured on the Tooling blueprint, SageMaker Unified Studio attaches the specified customer-managed permissions boundary to all IAM roles provisioned by that blueprint. If your governance requires a boundary, configure it explicitly and verify the resulting roles.

The following permissions boundary policy scopes the roles to the AWS services that Amazon SageMaker Unified Studio uses and then explicitly denies the Amazon DataZone actions that power the AI agent. This permissions boundary is provided for illustrative purposes only and isn’t a recommendation or reference for environment configuration. You should tailor your permissions boundaries to your specific workloads in accordance with the principle of least privilege.

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "AllowSmusServiceScope",
      "Effect": "Allow",
      "Action": [
        "datazone:*",
        "sagemaker:*",
        "glue:*",
        "s3:*",
        "lakeformation:*",
        "redshift:*",
        "redshift-data:*",
        "redshift-serverless:*",
        "athena:*",
        "q:*",
        "elasticmapreduce:*",
        "bedrock:*",
        "lambda:*",
        "kms:*",
        "secretsmanager:*",
        "codecommit:*",
        "logs:*",
        "cloudwatch:*",
        "sts:AssumeRole",
        "iam:PassRole",
        "ec2:Describe*",
        "ec2:CreateNetworkInterface",
        "ec2:DeleteNetworkInterface",
        "ec2:CreateNetworkInterfacePermission",
        "ec2:DeleteNetworkInterfacePermission"
      ],
      "Resource": "*"
    },
    {
      "Sid": "DenyDataNotebookAndDataAgent",
      "Effect": "Deny",
      "Action": [
        "datazone:*Notebook*",
        "datazone:*Cell*",
        "datazone:*Conversation*",
        "datazone:SendMessage",
        "datazone:GenerateCode",
        "datazone:CancelMessage"
      ],
      "Resource": "*"
    }
  ]
}

Warning: validate before using in production. This example scopes the roles to the service namespaces Amazon SageMaker Unified Studio uses, but it is still coarse (it allows each listed service in full) and is provided only for illustration. Because a permissions boundary is a ceiling, it must remain a superset of everything the three Tooling roles (datazone_usr_role, AmazonBedrockServiceRole, and AmazonBedrockLambdaExecutionRole) actually need. If Amazon SageMaker Unified Studio adds a dependency that isn’t listed, provisioning or in-console actions will fail with an access denied error. Validate in a non-production domain first.

With this boundary attached, the Tooling blueprint provisions normally, project members can access data connections and run SQL queries. However, any attempt to invoke the AI assistant or execute notebook cells through the agent returns an access denied error. The boundary acts as a ceiling that no policy attached to the role can override.

How it works

The custom permissions boundary feature operates at the blueprint configuration level. An administrator sets a PermissionsBoundaryArn in the Tooling blueprint’s regional parameters. When a user creates a new project that includes the Tooling blueprint, Amazon SageMaker Unified Studio provisions an AWS CloudFormation stack that creates three IAM roles and attaches the specified boundary to each:

  • datazone_usr_role – the role that all project members assume to access data and resources in that project.
  • AmazonBedrockServiceRole – for Amazon Bedrock operations.
  • AmazonBedrockLambdaExecutionRole – for Amazon Bedrock-related AWS Lambda functions.

Because the boundary is set at the blueprint level, it applies to every project created under that blueprint. No per-project configuration is needed.

Prerequisites

Before you begin, make sure that you have:

If your organization uses AWS Organizations with SCPs that require permissions boundaries, you will also need an organization with the target account as a member and permissions to create and attach SCPs in the management account.

Setting up the SCP (optional)

This section provides instructions to create an SCP and attach it to your AWS Organizations organizational unit or accounts. If your organization already enforces permissions boundaries through SCPs, skip this section. Otherwise, create an SCP in your AWS Organizations management account that denies IAM role creation unless an approved permissions boundary is attached. This also prevents the boundary from being removed, swapped, or weakened afterward:

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "DenyRoleWithoutApprovedBoundary",
      "Effect": "Deny",
      "Action": [
        "iam:CreateRole",
        "iam:PutRolePermissionsBoundary"
      ],
      "Resource": "*",
      "Condition": {
        "StringNotEquals": {
          "iam:PermissionsBoundary": "arn:aws:iam::${aws:PrincipalAccount}:policy/SMUSToolingBoundary"
        }
      }
    },
    {
      "Sid": "DenyRemovingBoundary",
      "Effect": "Deny",
      "Action": "iam:DeleteRolePermissionsBoundary",
      "Resource": "*"
    },
    {
      "Sid": "ProtectBoundaryPolicy",
      "Effect": "Deny",
      "Action": [
        "iam:DeletePolicy",
        "iam:CreatePolicyVersion",
        "iam:SetDefaultPolicyVersion"
      ],
      "Resource": "arn:aws:iam::${aws:PrincipalAccount}:policy/SMUSToolingBoundary"
    }
  ]
}

This policy does three things:

  • DenyRoleWithoutApprovedBoundary blocks creating a role, or attaching a boundary to an existing role, with anything other than the approved boundary ARN. Denying iam:PutRolePermissionsBoundary in addition to iam:CreateRole stops a privileged principal from swapping in a weaker boundary after the role exists.
  • DenyRemovingBoundary blocks iam:DeleteRolePermissionsBoundary outright, so the boundary cannot be stripped off. (This action doesn’t support the iam:PermissionsBoundary condition key, so it must be denied unconditionally.)
  • ProtectBoundaryPolicy prevents tampering with the boundary policy itself. Deleting it, or publishing and defaulting a new version that quietly widens what it allows.

Note: Scope these denies so you don’t lock yourself out. A broad deny on iam:PutRolePermissionsBoundary and iam:DeleteRolePermissionsBoundary also applies to your own administrators. Add an exception for a break-glass or IAM-admin role (for example, an aws:PrincipalArn StringNotLike condition) so a trusted principal can still manage boundaries.

To create the SCP, sign in to the AWS Organizations console with your management account and go to AWS Organizations → Policies → Service control policies. If SCPs aren’t enabled for your organization yet, choose Enable service control policies first. Choose Create policy, give it a name (for example, test_scp), and replace the default content in the policy editor with the JSON above substituting <account-id> with your account ID. Choose Create policy to save it.

After creating the SCP in the management account, verify its content before attaching it. Figure 3 shows the core create-role control. The full example above adds controls that prevent replacing or removing the boundary and modifying the protected policy.

Figure 3: Service Control Policy defined in the AWS Organizations management account

The AWS Organizations Content tab displays the customer-managed test_scp policy. Its visible statement denies iam:CreateRole unless the request uses the SMUSToolingBoundary policy.

Attach this SCP to the organizational unit or account where your Amazon SageMaker Unified Studio domain and domain-associated accounts reside. To do so, open the test_scp service control policy, choose the Targets tab, and choose Attach. The AWS organization structure appears; select the OU or account where the SCP should apply, then choose Attach policy.

Figure 4 identifies the member account that must inherit the SCP in this example organization. The target member account, datazone-account2, resides under OU2, while datazone-account1 is the organization’s management account. Attaching the SCP to the target account or a parent organizational unit enforces it there.

Figure 4: AWS Organizations account structure showing the management account and the target member account

After attaching the policy, verify the association on the SCP’s Targets tab, as shown in Figure 5.

SCP Targets tab listing datazone-account2 as an account target where test_scp is enforced

Figure 5: Service Control Policy attached to the target account where it should be enforced

The Targets tab lists datazone-account2 as an ACCOUNT target, confirming that test_scp is enforced directly on the intended member account.

Configuring the permissions boundary

In this section you will execute the required steps to create the permissions boundary and enable it in the Tooling blueprint. The example in this walkthrough uses us-east-1. Change it to the Region where your Amazon SageMaker Unified Studio domain is deployed. You must execute the configuration in the account where you plan to create your project. This can be your Amazon SageMaker Unified Studio domain account or accounts associated to your Amazon SageMaker Unified Studio domain.

Step 1: Create the permissions boundary policy

If you haven’t already created the boundary policy, save the following JSON document as a boundary-policy.json file on your workstation:

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "AllowSmusServiceScope",
            "Effect": "Allow",
            "Action": [
                "datazone:*",
                "sagemaker:*",
                "glue:*",
                "s3:*",
                "lakeformation:*",
                "redshift:*",
                "redshift-data:*",
                "redshift-serverless:*",
                "athena:*",
                "q:*",
                "elasticmapreduce:*",
                "bedrock:*",
                "lambda:*",
                "kms:*",
                "secretsmanager:*",
                "codecommit:*",
                "logs:*",
                "cloudwatch:*",
                "sts:AssumeRole",
                "iam:PassRole",
                "ec2:Describe*",
                "ec2:CreateNetworkInterface",
                "ec2:DeleteNetworkInterface",
                "ec2:CreateNetworkInterfacePermission",
                "ec2:DeleteNetworkInterfacePermission",
                "iam:GetRole",
                "sqlworkbench:*"
            ],
            "Resource": "*"
        },
        {
            "Sid": "DenyDataNotebookAndDataAgent",
            "Effect": "Deny",
            "Action": [
                "datazone:*Notebook*",
                "datazone:*Cell*",
                "datazone:*Conversation*",
                "datazone:SendMessage",
                "datazone:GenerateCode",
                "datazone:CancelMessage"
            ],
            "Resource": "*"
        }
    ]
}

As noted previously, this illustrative policy is scoped to the services Amazon SageMaker Unified Studio uses but is still coarse, and its allow list must stay a superset of what all three Tooling roles need.

Then create the policy using the following command:

aws iam create-policy \
--policy-name SMUSToolingBoundary \
--policy-document file://boundary-policy.json \
--description "Permissions boundary for SMUS Tooling roles - denies Data Agent and Data Notebook capabilities"

Note the policy ARN from the output, because it will be used later in the procedure.

Step 2: Retrieve the ID of your domain

Retrieve the ID of your domain by running the following command. Replace <YOUR_DOMAIN_NAME> with the name of your SageMaker Unified Studio domain.

aws datazone list-domains \
--region us-east-1 \
--query "items[?name=='<YOUR_DOMAIN_NAME>'].id | [0]" \
--output text

Note the returned ID, because it will be used later in the procedure.

Step 3: Identify the Tooling blueprint

Retrieve the Tooling blueprint ID by running the following command. Replace <domain-id> with the ID you noted in Step 2.

aws datazone list-environment-blueprints \
  --domain-identifier <domain-id> \
  --managed \
  --region eu-west-1 \
  --query "items[?name=='Tooling'].id" \
  --output json | jq -r '.[0]'

Note the returned ID, because it will be used later in the procedure.

Step 4: Read the current configuration

Retrieve the current Tooling blueprint configuration by executing the following command. Replace <domain-id> with the ID from Step 2 and <tooling-bp-id> with the ID from Step 3.

aws datazone get-environment-blueprint-configuration \
--domain-identifier <domain-id> \
--environment-blueprint-identifier <tooling-bp-id> \
--region us-east-1 | tee tooling-bp-config-backup.json

Important: Back up the output of get-environment-blueprint-configuration before making any changes. The command above pipes the response to tooling-bp-config-backup.json so you have a restore point if you need to revert.

Note the values of provisioningRoleArn, manageAccessRoleArn, enabledRegions, and all fields inside regionalParameters (AZs, S3Location, Subnets, VpcId). You will need all of these in the next step.

Step 5: Set the permissions boundary

Update the blueprint configuration to include PermissionsBoundaryArn in the regional parameters using the following command.

Important: The put-environment-blueprint-configuration API operates in overwrite mode, it replaces the entire configuration with what you provide. You must include all existing values from the previous step’s output. The only new addition is PermissionsBoundaryArn inside the regional parameters. Omitting any existing parameter removes it.

Make sure to replace <domain-id> with the ID you noted in Step 2, <tooling-bp-id> with the ID you noted in Step 3, and all other <placeholder> values with the corresponding values from Step 4’s output.

aws datazone put-environment-blueprint-configuration \
--domain-identifier <domain-id> \
--environment-blueprint-identifier <tooling-bp-id> \
--enabled-regions '<enabledRegions>' \
--provisioning-role-arn "<provisioningRoleArn>" \
--manage-access-role-arn "<manageAccessRoleArn>" \
--regional-parameters '{
  "<region>": {
    "AZs": "<AZs>",
    "S3Location": "<S3Location>",
    "Subnets": "<Subnets>",
    "VpcId": "<VpcId>",
    "PermissionsBoundaryArn": "arn:aws:iam::<account-id>:policy/SMUSToolingBoundary"
  }
}' \
--region <region>

The following anonymized example is based on an existing Tooling blueprint configuration. Its S3Location reflects the bucket naming pattern used in that environment. Copy the exact S3Location returned in Step 4. Don’t use the following illustrative value. Here’s an example:

aws datazone put-environment-blueprint-configuration \
--domain-identifier <domain-id> \
--environment-blueprint-identifier <tooling-bp-id> \
--enabled-regions '["us-east-1"]' \
--provisioning-role-arn "arn:aws:iam::<account-id>:role/service-role/AmazonSageMakerProvisioning-<account-id>" \
--manage-access-role-arn "arn:aws:iam::<account-id>:role/service-role/AmazonSageMakerManageAccess-us-east-1-<domain-id>" \
--regional-parameters '{
  "us-east-1": {
    "AZs": "us-east-1a,us-east-1b,us-east-1c,us-east-1d",
    "S3Location": "s3://amazon-sagemaker-<account-id>-us-east-1-<suffix>",
    "Subnets": "<subnet-1>,<subnet-2>,<subnet-3>,<subnet-4>",
    "VpcId": "<vpc-id>",
    "PermissionsBoundaryArn": "arn:aws:iam::<account-id>:policy/SMUSToolingBoundary"
  }
}' \
--region us-east-1

Step 6: Verify the configuration was applied

Confirm the permissions boundary ARN is now set in the blueprint configuration using the following command. Make sure to replace <domain-id> with the ID you noted in Step 2 and <tooling-bp-id> with the ID you noted in Step 3.

aws datazone get-environment-blueprint-configuration \
--domain-identifier <domain-id> \
--environment-blueprint-identifier <tooling-bp-id> \
--region us-east-1 \
--query "regionalParameters.\"us-east-1\".PermissionsBoundaryArn"

The output should return your boundary policy ARN:

"arn:aws:iam::<account-id>:policy/SMUSToolingBoundary"

Validating the configuration

After configuring the permissions boundary, in this section you will get instructions to create a new project to verify it works end to end and that the IAM roles created with the project actually include the permissions boundary.

Step 1: Select a project profile in enabled state

Use the following command to list project profiles configured in your domain. Make sure to replace <domain-id> with the ID you noted in Step 2 of the “Configuring the permissions boundary” section.

aws datazone list-project-profiles \
--domain-identifier <domain-id> \
--region us-east-1

Choose a project profile that has "status": "ENABLED". Note the id of any project profile returned in the previous command.

Step 2: Create a test project

Create a new project using the following command. Make sure to replace <domain-id> with the ID you noted in Step 2 of the “Configuring the permissions boundary” section and to replace <profile-id> with the project profile ID noted in Step 1 of this section.

aws datazone create-project \
--domain-identifier <domain-id> \
--name "PB-Validation-$(date +%Y%m%d-%H%M%S)" \
--project-profile-id <profile-id> \
--region us-east-1

Note the id (project ID) returned in the response. Wait for the Tooling blueprint to provision. This typically takes a minute or two. After provisioning completes, confirm that the validation project reaches the Active state, as shown in Figure 6.

SMUS Projects list showing the timestamped PB-Validation project in Active status after successful creation

Figure 6: Project created successfully with the permissions boundary configured

The Projects list shows the timestamped PB-Validation-* project with an Active status, confirming that project creation succeeded with the custom boundary configured.

Step 3: Verify the roles have the boundary attached

In this section you check that the IAM roles created with the project have the permissions boundary attached. Use the following commands to get the configuration for the IAM roles created with the project you just created. Replace <domain-id> and <project-id> with the values from the previous steps.

# Get the environment ID
ENV_ID=$(aws datazone list-environments \
--domain-identifier <domain-id> \
--project-identifier <project-id> \
--region us-east-1 \
--query "items[?name=='Tooling'].id" --output text)

# List IAM roles in the AWS CloudFormation stack
aws cloudformation describe-stack-resources \
--stack-name "DataZone-Env-${ENV_ID}" \
--region us-east-1 \
--query "StackResources[?ResourceType=='AWS::IAM::Role'].PhysicalResourceId" \
--output table

# Verify each role has the boundary
aws iam get-role \
--role-name "<role-name>" \
--query 'Role.PermissionsBoundary'

All three roles should return a response showing the permissions boundary ARN:

{
  "PermissionsBoundaryType": "Policy",
  "PermissionsBoundaryArn": "arn:aws:iam::<account-id>:policy/SMUSToolingBoundary"
}

You can also verify each role in the IAM console. Figure 7 shows the permissions boundary for the project user role.

IAM console Permissions tab showing SMUSToolingBoundary as the permissions boundary on datazone_usr_role

Figure 7: IAM console showing the permissions boundary attached to the datazone_usr_role

The datazone_usr_role Permissions tab displays SMUSToolingBoundary as its customer-managed permissions boundary.

Figure 8 confirms that the same boundary is attached to the Amazon Bedrock service role.

IAM console showing SMUSToolingBoundary as the permissions boundary on AmazonBedrockServiceRole

Figure 8: IAM console showing the permissions boundary attached to the AmazonBedrockServiceRole

The AmazonBedrockServiceRole also displays SMUSToolingBoundary as its customer-managed permissions boundary.

Figure 9 verifies the boundary on the third Tooling role, the Bedrock Lambda execution role.

IAM console showing SMUSToolingBoundary as the permissions boundary on AmazonBedrockLambdaExecutionRole

Figure 9: IAM console showing the permissions boundary attached to the AmazonBedrockLambdaExecutionRole

The AmazonBedrockLambdaExecutionRole likewise displays SMUSToolingBoundary, confirming that all three provisioned roles carry the boundary.

Step 4: Verify the boundary denies AI agent actions

In this section you verify the boundary actually denies AI agent actions. If you configured the boundary from the use case section earlier, the boundary blocks Data Notebooks and messages to the Data Agent, such as the Query Editor assistant. Any such attempt returns an access denied error. The project user role has the boundary attached, so even if the role’s identity policies grant the relevant APIs, the boundary’s explicit deny takes precedence.

To confirm, navigate to your project in SageMaker Unified Studio and test the following actions:

  1. Attempt to create a notebook – In the left sidebar, select Notebooks. Select Create notebook. The operation will fail because the permissions boundary prevents the datazone:CreateNotebook action (Figure 10).

Figure 10: Permissions boundary preventing creation of Data Notebooks

After the create action, Amazon SageMaker Unified Studio reports Failed to create notebook and identifies datazone:CreateNotebook as explicitly denied by SMUSToolingBoundary.

  1. Attempt to use Data Agent in the Query Editor – In the left sidebar, select Query Editor, then select the Chat with AI icon. The agent chat will fail to load because the permissions boundary blocks the APIs required by Data Agent (Figure 11).

Figure 11: Permissions boundary preventing using Data Agent on Query Editor

The Query Editor remains available, but the Agent panel reports “You don’t have access to Data Agent“. In this configured test, that message is the user-visible result of denying the Data Agent APIs. The screenshot itself doesn’t display the denied API or boundary ARN.

Important considerations

  • Immutable after project creation – The permissions boundary is set at provisioning time. Changing the boundary ARN on the blueprint configuration only affects new projects. Existing projects retain their original boundary.
  • Applies to all Tooling-provisioned roles – When PermissionsBoundaryArn is configured on the Tooling blueprint, SageMaker Unified Studio attaches the specified customer-managed permissions boundary to all three IAM roles created by that blueprint. It’s applied uniformly — you can’t selectively apply it to individual roles. No boundary is attached unless you configure one, so if your governance requires a boundary, set it explicitly and verify the resulting roles rather than assuming one is present by default.
  • Policy must exist – The IAM policy referenced by PermissionsBoundaryArn must exist in the account before project creation. If the policy is deleted or the ARN is invalid, provisioning will fail.
  • Tooling blueprint only – Among Amazon SageMaker Unified Studio provided blueprints, only the Tooling blueprint supports custom permissions boundaries. Other provided blueprints that create IAM roles (for example, the EmrOnEc2 blueprint) don’t currently support this feature. If your organization requires permissions boundaries on roles created by additional blueprints, you can build custom blueprints that include a permissions boundary configuration so you can extend this security control across your entire project infrastructure.

Clean up

To remove test resources, delete the test project from the SageMaker Unified Studio UI. On the project’s Overview page, choose the ⋮ (more actions) menu in the top-right and choose Delete project.

Figure 12: Deleting the test project from the project Overview page.

In the Delete project dialog, type confirm in the text box to acknowledge that the action is final, then choose Delete project. This permanently deletes the project and its underlying resources, and triggers an asynchronous AWS CloudFormation stack deletion.

Figure 13: Confirming project deletion.

To remove the boundary from future projects, re-run the put-environment-blueprint-configuration command from Step 5: Set the permissions boundary, but omit the PermissionsBoundaryArn field from the regional parameters. Because you backed up the original configuration in Step 4: Read the current configuration (tooling-bp-config-backup.json), you can reuse the exact same provisioningRoleArn, manageAccessRoleArn, enabledRegions, and regionalParameters values (AZs, S3Location, Subnets, VpcId) — just without PermissionsBoundaryArn — so the blueprint returns to provisioning roles with no permissions boundary.

Conclusion

With the custom permissions boundary feature for Amazon SageMaker Unified Studio, organizations can adopt Amazon SageMaker Unified Studio Tooling blueprints without compromising their IAM governance posture. By configuring a single parameter on the Tooling blueprint, all IAM roles provisioned by future projects automatically carry the specified permissions boundary. This satisfies SCPs that mandate a boundary on every role and gives administrators granular control over what the Tooling roles can do, for example disabling AI agent and notebook capabilities. Remember that the example boundary in this post is illustrative, because it scopes to the services SageMaker Unified Studio uses but is still coarse.

“I just updated the EnvironmentBlueprintConfiguration for the Tooling blueprint to include the new PermissionsBoundaryArn param. After that the blueprint provisioned successfully with the required permissions boundary attached to all the IAM roles, in line with our security policies. In the end it was a one-line change.”

— Nat Noordanus, Data Tech Lead at Nexthink

To get started, create your permissions boundary policy, configure it on the Tooling blueprint using the CLI, and create a project to verify the boundary is attached.

For more information, see the documentation for Amazon SageMaker Unified Studio, IAM permissions boundaries, and Service Control Policies.


About the authors

Sanjana Sekar

Sanjana Sekar

Sanjana is a Software Development Engineer on the Amazon SageMaker Unified Studio team. She is focused on improving Data Agent capabilities and the compute blueprints experience within SageMaker Unified Studio. Outside of work, she enjoys hiking and biking.

Luca Perrozzi

Luca Perrozzi

Luca is a Solutions Architect at AWS, based in Switzerland. He focuses on innovation topics at AWS, especially in Artificial Intelligence. Luca holds a PhD in particle physics and has 15 years of hands-on experience as a research scientist and software engineer.

Ganesh Sambandan

Ganesh Sambandan

Ganesh is a Senior Technical Account Manager at AWS, helping organizations adopt best practices for running secure, reliable and well-architected workloads on AWS. He works closely with strategic customers to accelerate the adoption of AI-driven cloud operations, enabling more effective DevOps practices, automation and operational excellence.

Stefano Sandona

Stefano Sandona

Stefano is a Senior Worldwide Specialist Solutions Architect for Big Data at AWS, helping customers build efficient, secure, and scalable data solutions.

Paolo Romagnoli

Paolo Romagnoli

Paolo is a Senior Solutions Architect at AWS who helps global energy organizations design and build data and AI enterprise solutions at scale.

Configure domain-level VPC networking in Amazon SageMaker Unified Studio

Post Syndicated from Prasad Nadig original https://aws.amazon.com/blogs/big-data/configure-domain-level-vpc-networking-in-amazon-sagemaker-unified-studio/

Enterprise operations teams that run domain-level VPC networking in Amazon SageMaker Unified Studio often support dozens of projects spanning data engineering, analytics, and machine learning (ML) teams. Each project requires private connectivity to internal databases, Amazon Simple Storage Service (Amazon S3) buckets, and AWS services. Without a domain-level Amazon Virtual Private Cloud (Amazon VPC) configuration, project owners coordinate with the networking team individually. This piecemeal approach leads to inconsistent subnet choices, missing VPC endpoints, connectivity failures that are hard to troubleshoot, and a network posture that is difficult to audit.

With domain-level VPC networking, you configure the network once, and all new projects get the right network immediately upon creation. In this post, you learn how to:

  • Configure SageMaker Unified Studio domain-level VPC networking.
  • Select subnets and security groups that provide multi-Availability Zone (multi-AZ) resilience.
  • Update projects that have no VPC to inherit the domain VPC, and understand when a project must be recreated instead.
  • Validate network connectivity from within a project.

In this post, you learn how to configure VPC networking for a SageMaker Unified Studio domain that uses AWS Identity and Access Management (IAM)-based authentication. You see how network components map to domain and project resources, and how to plan a configuration that balances security, connectivity, and operational simplicity.

Solution overview

Domain-level VPC networking provides a single network configuration that applies to all new projects in the domain. Projects automatically inherit the VPC settings, including subnets, security groups, and connectivity to AWS services through VPC endpoints. Existing projects are an exception and are handled separately (see Step 3).

The following diagram shows a single VPC with private subnets across two Availability Zones configured at the domain level, with data engineering, analytics, and ML projects all inheriting that configuration.

Architecture diagram showing domain-level VPC configuration in Amazon SageMaker Unified Studio with private subnets across two Availability Zones.

Figure 1: Domain-level VPC configuration in Amazon SageMaker Unified Studio. A single VPC with private subnets across two Availability Zones is configured at the domain level. All projects (data engineering, analytics, ML) inherit this configuration automatically

Key benefits of this approach:

  • Configure once, apply across projects: New projects inherit the domain VPC without manual intervention.
  • Consistent security posture: A single network boundary covers all data, analytics, and ML workloads.
  • Simplified auditing: One VPC to audit rather than one per project. Turn on VPC Flow Logs and review AWS CloudTrail events for network-level auditing.
  • Reduced operational overhead: Project teams start working immediately without submitting networking requests.

The following AWS services are used in this solution:

Prerequisites

Before configuring domain-level VPC networking, verify you have the following:

  • Domain administrator permissions for Amazon SageMaker Unified Studio.
  • An existing VPC with the following requirements:
    • At least two private subnets in different Availability Zones.
    • DNS hostnames and DNS support enabled.
    • At least five available IP addresses per expected Amazon SageMaker Unified Studio project. This is a baseline minimum. Workloads using AWS Glue, Amazon EMR, or Amazon Redshift Serverless consume additional elastic network interfaces (ENIs) per worker or node. We recommend /24 or larger subnets for production domains and forward-looking capacity planning based on your expected users and compute types. For detailed guidance, see How to set up a network-isolated VPC for Amazon SageMaker Unified Studio.
  • VPC endpoints configured for the AWS services your projects access (for example, Amazon S3, AWS Glue, Amazon SageMaker AI).
  • Private DNS enabled on all interface VPC endpoints (you must enable this so that service DNS names resolve to private IPs). If you use centralized VPC endpoints through AWS Resource Access Manager (AWS RAM) or AWS Transit Gateway, configure Amazon Route 53 Resolver inbound endpoints instead.
  • S3 gateway endpoint route table associations configured for all selected private subnets (without this, S3 access fails in subnets whose route table lacks the prefix-list route).
  • A security group (optional), if not provided, SageMaker Unified Studio creates one automatically.
  • The SageMakerStudioAdminIAMConsolePolicy managed policy (or equivalent permissions including ec2:Describe*, ec2:CreateSecurityGroup, and datazone:* actions) attached to the domain administrator IAM role. See SageMakerStudioAdminIAMConsolePolicy in the AWS Managed Policy Reference for the full permission set.

Note: The VPC must be in the same AWS Region as the domain.

For detailed guidance on VPC networking configuration, see Configure VPC networking for IAM-based domains in the SageMaker Unified Studio Administrator Guide.

VPC endpoint requirements

Because your subnets are private (no internet gateway route), compute resources access AWS services through VPC endpoints. At a minimum, configure the following interface and gateway endpoints (add Amazon Athena, AWS Lake Formation, or Amazon Redshift endpoints if you use them in your projects):

Endpoint Type Purpose
com.amazonaws.region.s3 Gateway S3 access for data storage
com.amazonaws.region.glue Interface AWS Glue job connectivity
com.amazonaws.region.sagemaker.api Interface SageMaker API calls
com.amazonaws.region.sagemaker.runtime Interface Model inference
com.amazonaws.region.logs Interface Amazon CloudWatch Logs
com.amazonaws.region.monitoring Interface Amazon CloudWatch metrics
com.amazonaws.region.sts Interface IAM role assumption
com.amazonaws.region.datazone Interface Amazon SageMaker Unified Studio service connectivity
com.amazonaws.region.ecr.api Interface ECR API calls (container image metadata)
com.amazonaws.region.ecr.dkr Interface ECR image layer pulls (Docker registry)
com.amazonaws.region.kms Interface AWS Key Management Service (AWS KMS) encryption/decryption operations

Note: Interface endpoints incur an hourly charge per Availability Zone plus data processing fees. Gateway endpoints (such as S3) have no hourly charge. Factor endpoint count and AZ spread into your cost estimate.

For a comprehensive list of all mandatory and optional VPC endpoints for a fully network-isolated setup, see How to set up a network-isolated VPC for Amazon SageMaker Unified Studio. For current pricing details, see AWS PrivateLink pricing.

Note: Review your account’s service quotas for interface VPC endpoints per VPC (default 50) and ENIs per Region before scaling. Request increases through Service Quotas if needed.

Solution walkthrough

The following steps walk you through configuring the domain VPC and validating it, from signing in to the console through confirming private connectivity from a project.

Step 1: Sign in and navigate to networking settings

  1. Sign in to the AWS Management Console as your Amazon SageMaker Unified Studio domain administrator (the IAM role designated as the domain login role).
  2. Open the Amazon SageMaker console.
  3. Use the Region selector in the top navigation bar to select the Region where your domain exists.
  4. On the Amazon SageMaker Unified Studio landing page, choose Open to launch your IAM-based domain.

The following screenshot shows the Amazon SageMaker Unified Studio landing page, where you choose Open to launch the domain.

Amazon SageMaker Unified Studio landing page with the Open button to launch the IAM-based domain.

Figure 2: Amazon SageMaker Unified Studio landing page with the Open button to launch the IAM-based domain

  1. From the navigation pane, choose Domain management.

The following screenshot shows Domain management in the navigation pane.

Navigation pane showing Domain management link in Amazon SageMaker Unified Studio.

Figure 3: Domain management on navigation pane

Note: Access to the domain administration page is restricted to the IAM role specified as the domain login role during domain creation.

Step 2: Add VPC configuration

  1. In the navigation pane, choose Settings. In the Networking in this account section, choose Add VPC.

The following screenshot shows the Networking in this account section with the Add VPC button.

Domain management Settings page showing the Networking in this account section with Add VPC button.

Figure 4: Domain management Settings page showing the Networking in this account section to add a VPC

  1. For VPC, select the VPC with connectivity to your compute, database, and storage resources. If no VPC exists, choose Create VPC to provision one using AWS CloudFormation.
  2. For Subnets, select a minimum of two private subnets in different Availability Zones.
  3. (Optional) For Security group, select a security group to control inbound and outbound traffic. If you don’t choose one, SageMaker Unified Studio creates one automatically.
  4. Choose Save.
  5. Verify the VPC configuration status shows Ready in the Networking in this account section.

The following screenshots show the Add VPC dialog and the resulting Ready status in the Networking in this account section.

Add VPC dialog with fields for VPC, subnets, and security group selection.

Figure 5: Add VPC dialog with fields for VPC, subnets, and security group selection

VPC configuration status showing Ready in the Networking in this account section.

Figure 6: VPC configuration status showing Ready in the Networking in this account section

Note: IAM-based domains support only one VPC configuration at a time. AWS IAM Identity Center-based domains can have a VPC per Region. For details, see Configure VPC networking for IAM-based domains in the SageMaker Unified Studio Administrator Guide.

New projects created in the domain now automatically use the saved VPC configuration. Existing projects are an exception. See Step 3 to update them.

Step 3: Update existing projects

Existing projects don’t automatically inherit the domain VPC configuration. How you apply the new settings depends on the project’s current state:

Projects with no VPC configured – Update in place to adopt the domain VPC. See the following steps.

Projects that already have a VPC – These can’t be switched to a different VPC configuration. To adopt the domain VPC:

  1. Create a new project (which inherits the domain VPC automatically).
  2. Recreate connections in the new project.
  3. Migrate assets from the old project.
  4. Back up any data you need, then delete the original project.

Because recreation can disrupt in-progress work and doesn’t migrate project data automatically, schedule this as a planned maintenance window.

To update a project that currently has no VPC configured:

  1. From the domain administration page, choose Projects in the navigation pane.
  2. Choose the project you want to update.
  3. On the project detail page, a banner appears: “Configurations have changed. Please update this project to access the latest configuration.”
  4. In the banner, choose Update.
  5. Confirm the update when prompted.

Repeat this process for each existing project that should use the domain VPC. The following screenshot shows the project detail page with the configuration update banner.

Project detail page showing the update banner for VPC configuration changes.

Figure 7: Project detail page showing the configuration update banner

Step 4: Validate connectivity

After configuring the domain VPC and updating your projects, verify connectivity. Compute resources should have private connectivity to AWS services through the VPC, without any additional project-level network configuration.

Create a notebook in one of your projects as shown in the following figure and run the following code:

Creating a notebook in a SageMaker Unified Studio project to validate VPC connectivity.

Figure 8: Creating a notebook in a SageMaker Unified Studio project to validate VPC connectivity

Requirements: Python 3.8+, Boto3 1.26 or later. Run in a notebook within your SageMaker Unified Studio project.

import boto3
import socket
import ipaddress

def validate_vpc_connectivity():
    """Validate that the project has private connectivity to AWS services
    through the domain-level VPC configuration."""

    results = {}
    region = boto3.session.Session().region_name
    if not region:
        raise RuntimeError('Could not determine AWS Region. Run this notebook inside a SageMaker Unified Studio project.')

    # Test Amazon S3 access via VPC endpoint
    try:
        s3 = boto3.client('s3')
        response = s3.list_buckets()
        results['S3'] = f"[PASS] Accessible ({len(response['Buckets'])} buckets)"
    except Exception as e:
        results['S3'] = f"[FAIL] Failed: {e}"

    # Test AWS Glue access via VPC endpoint
    try:
        glue = boto3.client('glue')
        dbs = glue.get_databases()
        results['Glue'] = f"[PASS] Accessible ({len(dbs['DatabaseList'])} databases)"
    except Exception as e:
        results['Glue'] = f"[FAIL] Failed: {e}"

    # Test STS (role assumption through VPC endpoint)
    try:
        sts = boto3.client('sts')
        identity = sts.get_caller_identity()
        results['STS'] = f"[PASS] Accessible (Account: {identity['Account']})"
    except Exception as e:
        results['STS'] = f"[FAIL] Failed: {e}"

    # Verify interface endpoint resolves to private IP
    try:
        sts_endpoint = f"sts.{region}.amazonaws.com"
        addr_info = socket.getaddrinfo(sts_endpoint, 443, family=socket.AF_INET)
        ip = addr_info[0][4][0]
        is_private = ipaddress.ip_address(ip).is_private
        if is_private:
            results['DNS Resolution'] = f"[PASS] Private IP ({ip}) (traffic stays on AWS network)"
        else:
            results['DNS Resolution'] = f"[WARN] Public IP ({ip}) - check VPC endpoint config"
    except Exception as e:
        results['DNS Resolution'] = f"[FAIL] Failed: {e}"

    # Print results
    print("-" * 40)
    print("Domain VPC Connectivity Validation")
    print("-" * 40)
    for service, status in results.items():
        print(f" {service}: {status}")
    print("-" * 40)
    print(f"\n Region: {region}")

    # Check if all tests passed
    all_passed = all("[PASS]" in status for status in results.values())
    has_warn = any("[WARN]" in status for status in results.values())
    if all_passed:
        print(f"\n [PASS] All services accessible via private VPC endpoints.")
        print(f" This project inherited its network configuration")
        print(f" from the domain without per-project setup.")
    elif has_warn and all("[PASS]" in s or "[WARN]" in s for s in results.values()):
        print(f"\n [WARN] Services are reachable, but DNS resolves to public IPs.")
        print(f" Verify that Private DNS is enabled on your interface VPC endpoints.")
    else:
        print(f"\n [FAIL] Some services are not reachable.")
        print(f" Check that VPC endpoints are configured and security")
        print(f" groups allow outbound traffic on port 443.")

validate_vpc_connectivity()

Expected output when VPC is correctly configured:

Successful validation output showing all services accessible through private VPC endpoints.

Figure 9: Successful validation output showing all services accessible through private VPC endpoints

If any service shows a failure, one common cause is security groups preventing traffic on port 443 to the VPC endpoint. Other causes include missing VPC endpoints, incorrect route table entries, or DNS resolution issues. For more information, see Configure VPC networking for IAM-based domains in the SageMaker Unified Studio Administrator Guide.

Note: An AccessDenied error indicates the request reached the service. Connectivity is working, but IAM permissions need adjustment (for example, the S3 test requires s3:ListAllMyBuckets, which some project roles lack). A timeout or connection error points to a networking problem (missing endpoint, route, or security group rule). The following screenshot shows the validation output when VPC endpoints are missing, where the affected services report timeout errors.

Validation output when VPC endpoints are not configured showing timeout errors.

Figure 10: Validation output when VPC endpoints are not configured. Timeout errors indicate missing endpoints

The security group applied at the domain level controls network access for all projects. To review or tighten the rules:

  1. Navigate to the Amazon VPC console.
  2. Choose Security groups and choose the security group shown in your domain’s Networking settings.
  3. Review the Inbound rules and Outbound rules tabs.

By default, the auto-created security group allows all outbound traffic on port 443 (HTTPS) to reach AWS services through VPC endpoints. Consider restricting outbound rules to only the specific VPC endpoint security groups for least-privilege access. Additionally, make sure your VPC endpoint security groups allow inbound TCP 443 from the domain security group or subnet CIDRs. For distributed compute services (AWS Glue, Amazon EMR), add a self-referencing inbound rule to allow worker-to-worker communication.

Updating VPC configuration

After the initial setup, you can modify the VPC configuration to change the VPC, subnets, or security group:

  1. From the domain administration page, choose Settings in the navigation pane.
  2. In the Networking in this account section, under the Actions column, choose Update.
  3. Update the VPC, subnets, or security group as needed.
  4. Choose Update.

The following screenshot shows the Update VPC dialog, where you modify the VPC, subnets, or security group.

Settings page with the Actions menu showing Update and Remove options for VPC configuration.

Figure 11: Update VPC dialog showing the option to modify VPC, subnets, or security group for the domain

Important: Updating the VPC does not affect already provisioned resources. Newly created resources in projects use the updated VPC. Existing projects that already have a VPC keep their original settings and must be recreated to adopt the change. Projects with no VPC can be updated in place (see Step 3).

Clean up

To remove the VPC configuration from your domain:

  1. From the domain administration page, choose Settings in the navigation pane.
  2. In the Networking in this account section, choose the Actions menu (⋮) and choose Remove.

The following screenshot shows the Actions menu with the Remove option.

Actions menu in the Networking in this account section showing the Remove option.

Figure 12: Actions menu in the Networking in this account section showing the Remove option

If you created a dedicated VPC for this walkthrough and no longer need it:

  • Delete the VPC and associated resources (subnets, VPC endpoints, security groups) from the Amazon VPC console. Before deleting, remove the domain VPC configuration and make sure all project resources are terminated. Active projects create ENIs that block VPC and subnet deletion.
  • If you used an AWS CloudFormation template to create the VPC, delete the stack to remove all resources cleanly. Open the AWS CloudFormation console and delete the stack.

Note: Removing the domain VPC configuration does not retroactively change projects that already have VPC applied. Those projects retain their existing network configuration. New projects created after removal do not have a VPC configured.

Conclusion

In this post, we showed how to configure domain-level VPC networking in Amazon SageMaker Unified Studio. A single domain-level VPC eliminates per-project networking overhead, enforces a consistent security posture, and simplifies compliance auditing.

Key takeaways:

  • Domain-level VPC is a one-time configuration that automatically applies to all new projects.
  • Projects with no VPC can be updated in place. Projects that already have a VPC must be recreated to adopt a changed configuration.
  • Private subnets with VPC endpoints provide secure, private connectivity to AWS services without traversing the public internet.

As next steps, consider:

  • Reviewing your auto-created security group rules and tightening them for least-privilege access.
  • Adding VPC endpoints for additional AWS services as your projects’ needs evolve.
  • Monitoring subnet IP address utilization to plan capacity as you add more projects. Use the AvailableIpAddressCount Amazon CloudWatch metric for your subnets to track utilization and set alarms.

For more information, see Configure VPC networking for IAM-based domains in the Amazon SageMaker Unified Studio Administrator Guide.

 


About the authors

Prasad Nadig

Prasad Nadig

Prasad is a Senior Analytics Specialist Solutions Architect at Amazon Web Services (AWS), specializing in large-scale data analytics and AI. Prasad partners with customers to design, migrate, and modernize their analytics platforms on AWS into scalable, cost-effective solutions, with deep expertise in data lakes, data warehousing, distributed processing, and performance tuning at petabyte scale.

Amit Shyam Jaisinghani

Amit Shyam Jaisinghani

Amit is a Software Engineer on the SageMaker Studio team at Amazon Web Services, and he earned his Master’s degree in Computer Science from Rochester Institute of Technology. Since joining Amazon in 2019, he has built and enhanced several AWS services, including Amazon WorkSpaces and Amazon SageMaker Studio. Outside of work, he explores hiking trails, plays with his two cats, Missy and Minnie, and enjoys playing Age of Empire.

Arun Shanmugam

Arun Shanmugam

Arun is a Senior Analytics Solutions Architect at AWS, with a focus on building modern data architecture. He has been successfully delivering scalable data analytics solutions for customers across diverse industries. Outside of work, Arun is an avid outdoor enthusiast who actively engages in CrossFit, road biking, and cricket.

Query unstructured data in Amazon SageMaker Catalog using generative AI

Post Syndicated from Nishchai JM original https://aws.amazon.com/blogs/big-data/query-unstructured-data-in-amazon-sagemaker-catalog-using-generative-ai/

Each day, businesses generate massive amounts of unstructured data, such as PDFs, images, email, customer feedback, and medical reports. But knowing data exists isn’t enough. You need to find it, access it, and extract answers from it fast. In Part 1 of this series, you saw how to set up the producer side of the pipeline: using Amazon Textract and Anthropic Claude on Amazon Bedrock to extract and enrich metadata, and then publish those enriched assets to Amazon SageMaker Catalog so your organization can discover them.

In this post, you take the next step: the consumer side. You sign in as a data consumer, search for and subscribe to the enriched unstructured data assets, and then query them using two approaches. The first is a no-code chat agent for natural language queries. The second is Amazon Bedrock model inference for programmatic access. By the end of this post, you will know how to unlock the business knowledge inside your unstructured data and make it available to analysts and application engineers alike.

Solution overview

This post continues the two-part series architecture, where Amazon SageMaker Catalog acts as the central hub connecting data producers and consumers through a publish-subscribe model.

The consumer workflow picks up after the producer has enriched and published the unstructured data assets. As a consumer, you will:

  • Sign in to your SageMaker Unified Studio consumer project and search the catalog using keywords from the enriched metadata README.
  • Subscribe to the published Amazon Simple Storage Service (Amazon S3) asset and get the subscription approved by the producer.
  • Interact with the subscribed data through two options:
    • Option 1 – A no-code chat agent for natural language queries (NLQs), ideal for data analysts and business users.
    • Option 2 – Amazon Bedrock model inference for programmatic NLQ integration, suited for application engineers building data-driven applications.

The following diagram illustrates the consumer workflow in this solution. The consumer (1) signs in to SageMaker Unified Studio, (2) searches the Amazon SageMaker Catalog for enriched unstructured data assets using keywords from the AI-generated metadata, (3) subscribes to the S3 data asset and receives approval from the producer, and then (4) queries the data using either the Amazon Bedrock chat agent app (Option 1) or Amazon Bedrock model inference through a Jupyter notebook (Option 2).

With both a no-code and a programmatic path, consumers across different roles, from analysts to engineers, can query data in the way that fits their workflow, while the SageMaker Catalog approval workflow maintains governed access throughout.

Consumer workflow architecture: sign in to SageMaker Unified Studio, search the SageMaker Catalog, subscribe to the S3 asset with producer approval, then query with the Amazon Bedrock chat agent or model inference

Figure 1: Consumer workflow for the publish-subscribe solution

Prerequisites

Before you begin, make sure you have completed all steps in Part 1 of this series, including:

Consume published data from the consumer project

In this section, you sign in as a consumer user in the SageMaker Unified Studio consumer project. You then subscribe to the S3 bucket by searching for a keyword that is part of the README published in Part 1.

  1. Sign in to the consumer project and search for the keyword emergency, which was added to the README file during publishing. The search returns the enriched asset that the producer published in Part 1.

    SageMaker Unified Studio catalog search for the emergency keyword, returning the enriched asset published in Part 1

    Figure 2: Catalog search results for the emergency keyword

  2. Choose the asset from the results to view its details, including the AI-generated business metadata, glossary terms, and README content. Then choose Subscribe.

    Asset details page showing AI-generated business metadata, glossary terms, and README content, with the Subscribe button

    Figure 3: Asset details with AI-generated metadata and the Subscribe option

  3. Enter analysis as the Reason for request in the Comment section, then choose Request.

    Subscription request dialog with analysis entered as the reason for request in the Comment box

    Figure 4: Subscription request with the reason for request entered

  4. Sign back in to the producer project (unstructured-producer-project) to approve the subscription request.
  5. After approval, return to the consumer project and confirm that the subscribed asset now appears under Manage, Assets, Subscribed assets.

    Consumer project Subscribed assets list confirming the approved subscription

    Figure 5: Approved subscription under the Subscribed assets tab

With the subscription approved, you can now access the enriched unstructured data through two approaches.

Option 1: As a data or business analyst, you can use the Amazon Bedrock chat agent app for natural language queries.

Option 2: As an application engineer, you can use Amazon Bedrock model inference for programmatic natural language queries.

Let’s explore both options.

Option 1: Amazon Bedrock chat agent app

The Amazon Bedrock chat agent app gives you a no-code, conversational interface to query your enriched unstructured data using natural language. As a data analyst or business user, you can ask questions in plain English. You get answers grounded in the documents your organization has ingested, without writing any code. For production workloads, especially in sensitive domains such as healthcare, you can apply Amazon Bedrock Guardrails to add content filtering and grounding validation to your model responses.

Data scientists and application engineers can also extend these capabilities by integrating the chat agent app APIs into custom applications, so users can interact with unstructured Amazon S3 data programmatically.

To set up the Amazon Bedrock chat agent app on your subscribed dataset, complete the following steps.

Prerequisite: Add the S3 data location.

Before creating the chat agent app, you need to add the S3 location of your subscribed data as a registered location in your project.

  1. Choose the Data tab in Overview.
  2. Choose the S3 bucket, and then choose Add to add the S3 location.

    Data tab in the project Overview with the S3 bucket selected and the Add button to register the S3 location

    Figure 6: Adding the S3 location from the Data tab

  3. On the S3 location page, provide the following details:
    • Add a name: producerprojectdata.
    • Add the producer’s S3 path as a new S3 location: s3://amzn-sagemaker-bucket-<domain-id>-<project-id>/medical/.

    Note: You can get the S3 location details from the technical name of your subscribed asset.

    • Choose the AWS Region, and then choose Add data to add this as a new location.

    Note: Make sure the AWS Region you select supports the Amazon Bedrock foundation models used later in this post. For a list of available models by Region, see Supported Regions and models for Amazon Bedrock.

    S3 location page with the location name, producer S3 path, and AWS Region entered before choosing Add data

    Figure 7: S3 location details and AWS Region selection

    Note: Make sure to select only the PDF files within the S3 path for the data source.

    Data source selection showing only the PDF files within the S3 path selected

    Figure 8: Selecting the PDF files as the data source

After the location is added, it appears as a selectable S3 location when creating a knowledge base in AI Apps.

Complete the following steps to configure the chat agent app:

  1. In the left navigation pane, under Generative AI, choose AI Apps.
  2. In the Build section of the page, choose Chat agent.

    AI Apps Build section with Chat agent selected in the left navigation under Generative AI

    Figure 9: Choosing Chat agent in the AI Apps Build section

  3. Expand the Data tab to create a knowledge base with your S3 bucket. On the Create a new knowledge base page, enter the following:
    • Add a name: MedicalKB.
    • Add a description: Knowledge base built from subscribed medical S3 data assets. Contains medical documents used to provide grounded, context-aware responses to medical domain queries.
    • Choose the data source. You will see the S3 bucket that you added in the previous step.
    Create a new knowledge base page with the MedicalKB name, description, and the added S3 bucket as the data source

    Figure 10: Creating the MedicalKB knowledge base from the S3 data source

  4. Choose your embedding model. You can leave the default settings and choose Create. It might take 10–15 minutes to create the knowledge base, depending on file sizes.
  5. After the knowledge base is created, on the Chat agent page:
    • Choose your preferred model from the Model menu (you can switch between different large language models as needed).
    • Under Data, choose your published S3 bucket as the knowledge base.
    • Begin interacting with the agent by entering questions in the Enter prompt field.
    Chat agent page with a model selected and the MedicalKB knowledge base chosen, ready to enter a prompt

    Figure 11: Chat agent page with the model and knowledge base selected

For example, entering “Which age groups had the highest rates of emergency department visits for tooth disorders?” returns an answer grounded in the enriched dental dataset published in Part 1.

The chat agent uses the enriched README metadata along with the underlying documents to surface contextually relevant answers. Analysts can explore unstructured content without needing to know where the data lives or how it’s structured.

Option 2: Natural language queries using Amazon Bedrock model inference

This option demonstrates how to use Amazon Bedrock model inference to query subscribed data using natural language. You can integrate this capability with external chat applications so users can run natural language queries through Amazon Bedrock.

  1. In your consumer project, choose Manage, Assets from the bottom of the left navigation pane. On the Subscribed tab, choose your subscribed S3 asset. Under Actions, choose Open JupyterLab notebook.

    Subscribed S3 asset Actions menu with Open JupyterLab notebook selected in the consumer project

    Figure 12: Opening the JupyterLab notebook from the subscribed asset

  2. This opens the JupyterLab notebook environment. Upload the s3_document_consumer_v2.ipynb notebook and run all the cells. You can download the notebook from s3_document_consumer_v2.ipynb.Note: The project role requires permissions for Amazon S3, Amazon Textract, and Amazon Bedrock. If you followed Part 1, you might already have these policies attached. For details on the required policies and guidance, see the prerequisites in Part 1.
  3. Review the notebook cells.
    JupyterLab notebook cells with the final cell showing a sample question answered by Amazon Bedrock

    Figure 13: Sample question answered by Amazon Bedrock in the notebook

    In the final cell, you find a sample question that Amazon Bedrock answers: “Which primary payer types (Medicare, Medicaid, private insurance, and so on) account for the highest proportion of dental-related emergency department visits?”

    Amazon Bedrock processes the question against the enriched content in the S3 bucket and returns a grounded answer. You can replace this sample question with any query relevant to your documents.

The Amazon Bedrock model inference approach gives you programmatic control, making it possible to embed natural language query capabilities directly into your existing data applications and business intelligence tools.

Clean up

To avoid ongoing charges, make sure to delete the resources used in this solution immediately after completing the walkthrough. The primary cost drivers are SageMaker Unified Studio notebook instances, Amazon Bedrock model inference calls, and Amazon S3 storage.

  1. Stop SageMaker Unified Studio resources:
    • Close running notebooks.
    • Stop running notebook instances.
    • Shut down unused kernels.

    Note: Running notebook instances continue to incur charges even when not in use.

  2. Clean Amazon S3 storage:
    • Delete temporary files created during processing.
    • Remove uploaded test documents that are no longer needed.

    Note: Although Amazon S3 costs are minimal, large volumes of data can accumulate significant charges, so it’s best to remove unneeded data.

Conclusion

In this post, you saw how to consume and query the enriched unstructured data assets published in Part 1 of this series. By subscribing to assets through the Amazon SageMaker Catalog publish-subscribe model, you can discover, access, and interact with your organization’s unstructured data, whether through the no-code chat agent or Amazon Bedrock model inference.

Together, both parts of this series show you how to build a comprehensive pipeline that transforms raw unstructured documents into governed, queryable knowledge assets. The combination of Amazon Textract for extraction, Amazon Bedrock for intelligent summarization and NLQ, and Amazon SageMaker Catalog for governance and discoverability means your teams can focus on extracting business insights rather than managing infrastructure.

To continue your Amazon SageMaker journey, see the following resources:


About the authors

Nishchai JM

Nishchai JM

Nishchai is an Analytics and generative AI Specialist Solutions Architect at Amazon Web Services. He specializes in building larger scale distributed applications and helps customers modernize their workloads on AWS. He thinks Data is new oil and spends most of his time deriving insights from data.

KiKi Nwangwu

KiKi Nwangwu

KiKi is an Analytics and generative AI Specialist Solutions Architect at AWS. She specializes in helping customers architect, build, and modernize scalable data analytics and generative AI solutions. She enjoys traveling and exploring new cultures.

Narendra Gupta

Narendra Gupta

Narendra is a Sr. Specialist Solutions Architect for Data & AI (Analytics) at AWS. He works with customers to design data-driven solutions and has deep expertise in data governance and cataloging.

Aditya Edara

Aditya Edara

Aditya is a Support Engineer at AWS. He serves as a Subject Matter Expert in AWS Analytics services, specializing in Amazon EMR and AWS Glue. Aditya provides expert guidance and technical support to enterprise and strategic customers, helping them optimize data analytics solutions.

Discover and govern Snowflake data using SageMaker Unified Studio

Post Syndicated from Marco Duarte original https://aws.amazon.com/blogs/big-data/discover-and-govern-snowflake-data-using-sagemaker-unified-studio/

Many organizations operate in hybrid data environments where critical assets live in Snowflake while analytics workloads run on AWS, which can create governance gaps, discovery friction, and duplicated efforts when the two aren’t connected.

With Amazon SageMaker Unified Studio, you can govern data across Snowflake and AWS through its integrated catalog and AWS Glue Data Quality, a capability of AWS Glue. You connect directly to Snowflake tables without moving data, apply quality rules using AWS Glue Visual ETL, and publish validated assets to Amazon SageMaker Catalog, maintaining consistent governance across your entire distributed data estate.

Without this integration, cataloging Snowflake data requires building extraction pipelines, often taking days. With SageMaker Unified Studio connected to Snowflake, you can query, catalog, and validate the quality of federated data in 5–15 minutes. No data replication or custom ETL code required.

In this post, we show you how to connect Snowflake to Amazon SageMaker Unified Studio, register data assets in Amazon SageMaker Catalog, configure data quality validation using AWS Glue Visual ETL, and publish assets for unified collaboration. By following these steps, you enrich federated assets with data quality scores so that consumers across your organization can discover and trust the data, all while keeping it in Snowflake.

Solution overview

This solution integrates Snowflake with Amazon SageMaker Unified Studio for centralized data cataloging and quality validation.

The architecture uses an AWS Glue connection to federate the Snowflake catalog into Amazon SageMaker Unified Studio. Tables become available in the project catalog without complex storage configurations. You can query data directly using SQL analytics, publish datasets to Amazon SageMaker Catalog for organization-wide discovery, and apply data quality rules through AWS Glue Visual ETL pipelines.

The workflow consists of the following steps:

Architecture diagram: Snowflake federated into SageMaker Unified Studio through AWS Glue, with data quality validation and publishing to SageMaker Catalog

Figure 1: Architecture for federating Snowflake into SageMaker Unified Studio and validating data quality

  1. Snowflake connection creation on Amazon SageMaker Unified Studio — Amazon SageMaker Unified Studio uses an AWS Glue connection to federate Snowflake tables and views into its open data lakehouse architecture. The federated catalog entry is registered in AWS Glue Data Catalog and governed by AWS Lake Formation for centralized access control, without moving data out of Snowflake.
  2. Federate Snowflake tables into the Amazon SageMaker publisher project — The Amazon SageMaker publisher project discovers the federated Snowflake tables through the AWS Glue Data Catalog integration.
  3. Publish the dataset to Amazon SageMaker Catalog — The publisher project publishes the dataset as a governed asset to the Amazon SageMaker Catalog, making it discoverable for data consumers across the organization.
  4. Validate data quality — AWS Glue Data Quality runs validation rules against the federated Snowflake data and publishes the data quality results directly to the corresponding asset in Amazon SageMaker Catalog.
  5. Consume data — Users access Snowflake data through two paths:
    1. Publisher project users — Query data with SQL Analytics — Users in the publisher project can query the Snowflake data directly using Amazon SageMaker Unified Studio SQL Analytics for interactive exploration and analysis, without copying or moving data.
    2. Consumer project users — Discovery and subscription through SageMaker Catalog — Other Amazon SageMaker consumer projects discover the published asset in the Amazon SageMaker Catalog, subscribe to it, and consume the data for their analytics and machine learning workloads.

Prerequisites

To follow along, you need:

  • An active Snowflake account with administrator access.
  • Tables or views created within a schema inside a Snowflake database.
  • An Amazon SageMaker Unified Studio and project created.
  • An Amazon Simple Storage Service (Amazon S3) bucket for AWS Glue assets.
  • Appropriate AWS Identity and Access Management (IAM) permissions configured (Amazon SageMaker Catalog is built on Amazon DataZone, so the IAM actions use the datazone: prefix.)

Your AWS Glue job execution role requires specific permissions to interact with Amazon SageMaker Catalog.

Required IAM policies for the AWS Glue job role

1. Amazon SageMaker Catalog search and listing permissions: Attach a policy that allows the AWS Glue job to search and list assets in Amazon SageMaker Catalog.

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "datazone:SearchListings",
        "datazone:GetListing",
        "datazone:ListDomains",
        "datazone:GetDomain"
      ],
      "Resource": "arn:aws:datazone:<REGION>:<ACCOUNT_ID>:domain/<DOMAIN_ID>"
    }
  ]
}

2. Amazon SageMaker Catalog time series data posting permissions: Add permissions to post data quality metrics:

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "datazone:PostTimeSeriesDataPoints",
        "datazone:GetAsset",
        "datazone:ListAssetRevisions"
      ],
      "Resource": "arn:aws:datazone:<REGION>:<ACCOUNT_ID>:domain/<DOMAIN_ID>"
    }
  ]
}

Configure the AWS Glue job role as an Amazon SageMaker domain user

Configure the IAM role used by your AWS Glue job as a domain user. In the Amazon SageMaker console, navigate to your domain, choose Access management, and add the AWS Glue job execution IAM role as a domain user.

Project-level permissions

Add the AWS Glue job execution role as a project member with Owner permissions. Navigate to your project, go to Project settings > Members, and add the role.

For more information about IAM roles for AWS Glue, see the AWS Glue security documentation. For Amazon SageMaker Unified Studio permissions, refer to the Amazon SageMaker Unified Studio administrator guide.

Querying Snowflake datasets from Amazon SageMaker Unified Studio

The following sections walk you through connecting Snowflake to Amazon SageMaker Unified Studio and running data quality validation with results displayed in Amazon SageMaker Catalog.

Identifying information in Snowflake

First, gather your Snowflake connection details. You need a Snowflake account with tables or views created at the schema level within a database.

To obtain Snowflake connection information:

  1. Navigate to your Snowflake environment and sign in with administrator credentials.

    Snowflake sign-in screen for administrator credentials
  2. Choose your user account and choose Connect a tool to Snowflake.

  3. Note the Account/Server URL displayed on the screen.
  4. Choose the Config File tab, select values for Warehouse, Database, and Schema, and copy these values for use in the next section.

Creating the connection in Amazon SageMaker Unified Studio

The Add Connection feature stores Snowflake connectivity details including credentials, server, and database information. Amazon SageMaker Unified Studio uses this connection to federate the Snowflake catalog through AWS Glue, so you can query data within minutes of setup.

You need an Amazon SageMaker Unified Studio domain and a project, which acts as a data producer project.

To create the Snowflake connection:

  1. In your Amazon SageMaker Unified Studio project, go to Overview.

    SageMaker Unified Studio project Overview page
  2. Choose Data.

    Data option in the SageMaker Unified Studio project navigation
  3. Choose + Add, then choose Add Connection.

    Add menu in SageMaker Unified Studio with the Add Connection option
    Add Connection panel in SageMaker Unified Studio
  4. Choose Next.
  5. Select Snowflake and choose Next.

    Connection type selection showing Snowflake in SageMaker Unified Studio
  6. Complete the connection details:
    • Name: snowflake-connection.
    • Description (Optional): Enter a description for your connection.
    • Host: Your Snowflake account URL (for example, XXXXXXXXX-XXX000000.snowflakecomputing.com).
    • Port: 443.
    • Database: Your database name (for example, sm_demo).
    • Warehouse: Your warehouse name (for example, COMPUTE_WH).
    • Schema: Your schema name (for example, demo).
    • Additional Properties:
      • Register in AWS Glue Data Catalog: Turn on checkbox.
      • Case conflict handling: Select the option based on Snowflake naming syntax.
    • Authentication:
      • Username: Your Snowflake username.
      • Password: Your Snowflake password.
    Snowflake connection details form with name, host, port, database, warehouse, and schema fields
    Connection form showing authentication and AWS Glue Data Catalog registration options
  7. Choose Add Data.

After creating the connection, wait a few minutes for the federated connection to be established. Search within Amazon SageMaker Unified Studio for the database and created objects.

Federated Snowflake database and objects appearing in SageMaker Unified Studio search

Federated Snowflake tables registered in the AWS Glue Data Catalog

With the Snowflake connection established and the federated tables registered in AWS Glue Catalog, you’re now ready to query Snowflake data directly from Amazon SageMaker Unified Studio, without moving or replicating any data.

Query results from a federated Snowflake table in the SageMaker Unified Studio query editor

How federated queries work

When you run a query in the Amazon SageMaker Unified Studio query editor against a federated Snowflake table, Amazon Athena runs the request. Athena is the underlying query engine integrated into Amazon SageMaker Unified Studio. Athena reads the table definition from AWS Glue Catalog, connects to Snowflake through the established connection, and pushes the query down for execution. Athena returns results directly to the query editor while Snowflake processes the data in place, and only the query results travel across the connection. Amazon SageMaker Unified Studio doesn’t copy data to S3 or any intermediate storage.

After you’ve validated that queries return the expected results, the next step is to publish this dataset to Amazon SageMaker Catalog, making it discoverable and shareable across your organization.

Publishing Snowflake datasets to the SageMaker Catalog

Now that your Snowflake connection is configured, you can publish your datasets to the Amazon SageMaker Catalog, making them discoverable and shareable across your organization.

Creating data assets in SageMaker Catalog

Data assets in Amazon SageMaker Catalog are the cataloged representation of your data resources. They help teams discover, govern, and share data across your organization.

In this section, you create a data asset associated with a Snowflake table. This process transforms a technical Snowflake table into a cataloged resource enriched with business metadata.

To create a data source:

  1. In your Amazon SageMaker Unified Studio project, go to Manage.

    Manage tab in the SageMaker Unified Studio project
  2. Choose Data Sources.
  3. Choose Create Data Source.

  4. Select the AWS Glue option.

    Data source type selection showing the AWS Glue option
  5. Turn on the Import data lineage checkbox and select the connection: project.default_lakehouse.

    Data source configuration with Import data lineage and the project.default_lakehouse connection selected
  6. Complete the form and choose Next:
    • Catalog: Select Enter the catalog name and enter snowflake-connection.
    • Database name: Enter your database name (for example, movies).
    • Table selection criteria: Enter * for all tables in the database, or enter a specific table name.
    Data source form showing catalog name, database name, and table selection criteria
  7. Keep the default options and choose Next until you reach the summary screen.

    SageMaker Unified Studio data source configuration summary screen
    Data source review screen before creation
  8. Review your settings and choose Create.

To extract metadata and publish assets:

  1. Choose Run to start extracting metadata from AWS Glue Data Catalog.

    Data source detail page with the Run option to extract metadata from the AWS Glue Data Catalog
  2. Wait for the run to complete.
  3. Go to Assets to view the Asset Inventory.

    Asset inventory in SageMaker Catalog after the data source run completes

The following screenshot shows the asset inventory after the data source run completes.

  1. Choose an asset to view its details.

    Asset detail page in SageMaker Catalog showing the Snowflake table metadata

At this point, you can enrich the business context by choosing Generate Descriptions. Amazon SageMaker Catalog analyzes the asset’s technical structure and generate:

  • Business descriptions in natural language for the asset.
  • Contextual definitions for each field/column.
  • Suggested glossary terms that could be applied.
  1. After your asset has been enriched with the necessary business metadata, you can publish it to the Amazon SageMaker Catalog by choosing Publish Asset.

Publish Asset option on the enriched Snowflake asset in SageMaker Catalog

The Snowflake enriched asset is now available to data consumers across your organization. Other users can discover it, subscribe to it, and consume it without data replication.

Implementing data quality rules with AWS Glue Data Quality

This section explains how to apply data quality validations to Snowflake data using AWS Glue Data Quality and visualize results in Amazon SageMaker Catalog.

Setting up the custom transform

Upload two files to an Amazon S3 bucket in the same AWS account where you run AWS Glue:

Copy both files to your AWS Glue assets S3 bucket in the transforms folder (s3://aws-glue-assets-<account-id>-<region>/transforms). AWS Glue Studio reads all JSON files from this folder to register custom visual transforms.

Custom transform files uploaded to the transforms folder in the AWS Glue assets S3 bucket

In the following sections, we walk you through the steps of building an ETL pipeline for data quality validation using AWS Glue Studio.

Creating the AWS Glue Visual ETL job

AWS Glue for Spark provides built-in support for reading from Snowflake data sources.

To create a new visual ETL job:

  1. Open the AWS Glue console at https://console.aws.amazon.com/glue/. Choose ETL jobs, then Visual ETL.

    AWS Glue console showing ETL jobs and the Visual ETL option

Establishing the Snowflake connection

To add a Snowflake source:

  1. In the job pane, choose Snowflake as your source. For Snowflake connection, select the connection that you created earlier. Specify the relevant schema and table for data quality checks.

    Snowflake source node configured in the AWS Glue visual ETL job

The visual editor displays the Data source properties panel where you select your connection, database, and enter a custom query targeting your Snowflake table.

Applying data quality rules

After establishing the Snowflake connection, configure the data quality evaluation step using the Data Quality Definition Language (DQDL).

To add data quality validation:

  1. Choose Transform and choose Evaluate Data Quality.
  2. Define domain-specific data quality rules using DQDL. For more information, see the AWS DQDL documentation.

    Evaluate Data Quality transform with DQDL rules in AWS Glue Studio
  3. Choose to output the data quality results. Optionally, store outcomes in Amazon S3 or publish to Amazon CloudWatch with alert notifications.

The preview of the data quality results from the ruleOutcomes node shows the outcomes of each rule.

Preview of the data quality rule outcomes from the ruleOutcomes node

Post the data quality results to Amazon SageMaker Catalog

To configure the custom transform:

  1. Add the Datazone DQ Result Sink transform to your job.
  2. Connect the ruleOutcomes node output to this transform.
  3. Complete the parameters:
    • Role to assume (Optional): Only needed for associated accounts.
    • Domain ID: Your Amazon SageMaker Unified Studio domain ID (found in the Amazon SageMaker Unified Studio portal).
    • Table name and Schema name: Same values used when creating the Snowflake source transform.
    • Data quality ruleset name: The name you want to give to the ruleset in Amazon SageMaker Catalog.
    • Max results: Maximum number of assets to return in case of multiple matches.

The following image shows the complete job graph with the Datazone DQ Result Sink transform configured.

AWS Glue visual ETL job graph with Snowflake source, Evaluate Data Quality, ruleOutcomes, and Datazone DQ Result Sink nodes

The visual editor displays four nodes connected sequentially: the Snowflake data source, the Evaluate Data Quality transform, the ruleOutcomes SelectFromCollection transform, and the Datazone DQ Result Sink transform.

To configure job parameters:

  1. Choose Job details.
  2. In Job parameters, add the following key-value pair:
    • --additional-python-modules
    • boto3>=1.34.105
  3. Save and run the job.

AWS Glue job parameters with the additional-python-modules key set to boto3

Visualizing data quality results in the SageMaker Catalog

After the AWS Glue ETL job completes, you can view the data quality information directly in Amazon SageMaker Catalog. This is the key outcome of running data quality on a federated source: the asset gains quality scores and metadata without ever leaving Snowflake. This makes it trustworthy and ready for other teams across your organization to use. Data consumers can now discover this asset in Amazon SageMaker Catalog and evaluate its quality before subscribing, without needing direct access to Snowflake or running their own validation.

To view data quality results:

  1. Open the Amazon SageMaker Unified Studio console.
  2. Navigate to your project.
  3. Go to Assets.
  4. Choose the Snowflake data asset.
  5. View the data quality information displayed on the asset page.

The following image shows the asset page in Amazon SageMaker Catalog with the data quality score populated.

SageMaker Catalog asset page showing a populated data quality score for the Snowflake asset

Data Quality tab in SageMaker Catalog showing an overall score of 100 with the movies rule set passed

The Data Quality tab shows an overall score of 100 and lists the rule set movies with a Passed result (1/1). This confirms that the data quality checks from AWS Glue posted successfully to Amazon SageMaker Catalog.

Clean up

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

  1. Delete the AWS Glue ETL job — Open the AWS Glue console, choose ETL jobs, select your job, and then choose Delete.
  2. Remove the AWS Glue connection — In the AWS Glue console, go to Connections, select the Snowflake connection, and then choose Delete.
  3. Delete the data source in SageMaker Catalog — In your Amazon SageMaker Unified Studio project, go to Data Sources, select the data source you created, and then choose Delete.
  4. Remove S3 assets — Delete the custom transform files from your s3://aws-glue-assets-<account-id>-<region>/transforms/ bucket.
  5. Remove IAM policies — Detach and delete the IAM policies you attached to the AWS Glue job execution role. Remove the role as a domain user and project member.

Conclusion

In this post, we showed you how to connect Snowflake to Amazon SageMaker Unified Studio for centralized data cataloging and quality validation. This approach maintains consistent governance without replicating data. Key benefits include:

  • Query without data movement: Access Snowflake data directly from Amazon SageMaker Unified Studio through federated queries, using the interoperable data architecture of AWS and eliminating time-consuming data replication.
  • Centralized governance: Maintain a single source of truth for data discovery, quality metrics, and governance policies across your distributed data estate.
  • Automated quality validation: Apply consistent data quality rules using AWS Glue Data Quality and visualize results directly in Amazon SageMaker Catalog.
  • Unified collaboration: Support data discovery and sharing across your organization through the publishing capabilities of Amazon SageMaker Catalog.

To get started, open the Amazon SageMaker Unified Studio console. To learn more about related topics, see Cross-account lakehouse governance with Amazon S3 Tables and SageMaker Catalog and Get started with AWS Glue Data Quality dynamic rules for ETL pipelines.


About the authors

Marco Duarte López

Marco Duarte López

Marco is a Data Specialist Solutions Architect at AWS, based in Santiago, Chile. He works with organizations across the region to design modern data architectures and governance frameworks that enable trusted, scalable data consumption. He is a member of the AWS Technical Field Community (TFC) for Analytics, where he specializes in Data & AI Governance, and has led data transformation programs for some of the largest enterprises in the region.

Diego Ortiz

Diego Ortiz

Diego is a Senior Data Strategy Solutions Architect for Latin America based in San Juan, Puerto Rico, with 14+ years of experience in technology roles. He supports organizations across countries and industries to develop data and AI strategies aligned with their business objectives, combining strategic vision with deep technical expertise in data and AI technologies. He is a core member of the Data Governance global community at AWS and leads the analytics technical community in the Spanish-speaking countries of Latin America.

Connect Amazon SageMaker Unified Studio to Microsoft Power BI – Part 1: IAM Identity Center (IDC)-based domains

Post Syndicated from Ramesh H Singh original https://aws.amazon.com/blogs/big-data/connect-amazon-sagemaker-unified-studio-to-microsoft-power-bi-part-1-iam-identity-center-idc-based-domains/

Connecting Power BI to your Amazon SageMaker Unified Studio data catalogs typically required third-party bridges. These bridges added complexity and licensing costs. In this post, you create a direct connection using new authentication modes in the Amazon Athena ODBC driver, removing those dependencies entirely. If your organization uses Power BI as its business intelligence (BI) tool, your analysts can configure access to governed data in Amazon SageMaker Unified Studio without changing their tools or workflows. As an AWS alternative, Amazon Quick Sight provides serverless BI integration with Amazon SageMaker Unified Studio at pay-per-session pricing.

A previous post showed the connection method using a third-party ODBC-JDBC bridge. The Amazon Athena ODBC driver (version 2.2.0 and later) now supports Amazon SageMaker Unified Studio authentication directly, eliminating the need for customers to configure third-party bridge components previously required for this connection. This bridge also created additional components and required ongoing maintenance. The native connection simplifies the architecture by reducing these requirements.

UC Irvine, a top-ten U.S. public research university, consolidates student data from systems across multiple departments into a single governed repository that supports reporting, research, and analytics for decision-making at the strategic, tactical, and operational levels. Many of their analysts rely on Power BI to explore and visualize this governed data.

“Our users rely on Power BI for data visualization and reporting, but connecting to governed data in AWS previously required workarounds. The ODBC connection feature gives a direct path from Power BI into our SageMaker Unified Studio projects—no bridge software, no extra licensing, just a connection string and we’re ready to go.”

— Bernadette Theologidy, Manager, Student Analytics, UC Irvine

The Athena ODBC driver introduces two new authentication modes for SageMaker Unified Studio:

  1. SageMakerBrowserIdc (for IDC-based domains): The driver opens a browser window and authenticates through AWS IAM Identity Center (and your external identity provider, if configured). No local AWS credentials are needed.
  2. SageMakerIam (for AWS Identity and Access Management (IAM)-based and IDC-based domains): The driver uses AWS credentials from the default credential provider chain. For this walkthrough, we use AWS IAM Identity Center to provide those credentials.

You connect Microsoft Power BI to Amazon SageMaker Unified Studio through Athena. The Athena ODBC driver supports using two connection methods that use these authentication modes:

Method 1: DSN-based (Athena Power BI connector): You configure an ODBC Data Source Name (DSN) and use the Athena connector in Power BI. This method supports DirectQuery and Import mode with both SageMakerBrowserIdc and SageMakerIam authentication.

Method 2: DSN-less (Power BI ODBC connector): You use the Power BI ODBC connector with a connection string, requiring no DSN configuration. This method supports Import mode only with SageMakerIam authentication. DirectQuery isn’t available because the Power BI ODBC connector doesn’t support it. The connection string in Power BI Desktop must match exactly the one on Power BI Service. Because the gateway runs as a Windows service without interactive browser access, both ends must use SageMakerIam.

Feature Method 1: DSN-based Method 2: DSN-less
Power BI Connector Amazon Athena connector ODBC connector
Data connectivity mode DirectQuery and Import Import only
Requires DSN configuration Yes No
Data freshness Real-time (DirectQuery) or scheduled (Import) Scheduled refresh only
Authentication types SageMakerIam and SageMakerBrowserIdc SageMakerIam only
Domain types supported IAM-based and IDC-based IAM-based and IDC-based
Best for Dashboards requiring live data Scenarios where DSN management is not possible or scheduled refresh is acceptable

This is Part 1 of a two-part series. This post covers IDC-based domains using both connection methods. Part 2 covers IAM-based domains.

Solution overview

In this walkthrough, you take the role of a data analyst at an energy company. You need to understand the current state and future direction of the U.S. power generation fleet using the Public Utility Data Liberation Project, available on the Registry of Open Data on AWS. Our goal is to analyze generation capacity and identify where new investment is flowing. We connect Power BI to Athena through Amazon SageMaker Unified Studio and query the EIA-860 generators dataset directly from our data catalog. The result is a single visualization that reveals the energy transition.

The following diagram illustrates the solution architecture for connecting Power BI to Amazon SageMaker Unified Studio through Amazon Athena.

Architecture diagram showing Power BI connecting to Amazon Athena through Amazon SageMaker Unified Studio, with a Microsoft on-premises data gateway on Amazon EC2

Figure 1: Architecture diagram

The following architecture demonstrates a six-step workflow.

  1. Data engineers and analysts connect Power BI Desktop to Athena as a data source.
  2. They build their reports locally.
  3. They then publish them to the Power BI Service.
  4. Microsoft On-Premises Data Gateway on an Amazon Elastic Compute Cloud (Amazon EC2) instance connects to Athena using the instance’s attached IAM role.
  5. The Power BI Service then uses this gateway connection.
  6. Report viewers access the published reports through Power BI Service to make data-driven decisions.

On the AWS side, Athena queries the data catalog managed by AWS Glue Data Catalog. The catalog references data stored in Amazon Simple Storage Service (Amazon S3). An Amazon SageMaker Unified Studio project governs all access.

In an IDC-based domain (covered in this post), Power BI Desktop uses SageMakerBrowserIdc for Method 1 and SageMakerIam for Method 2. Power BI Desktop can run on-premises or on an EC2 instance. The gateway always uses SageMakerIam (it runs as a Windows service without browser access) and authenticates using instance profile credentials, which rotate automatically. The gateway can only query data within projects where its IAM role has been added as a member. For IAM-based domains, see Part 2.

Prerequisites

Before connecting Power BI to Amazon SageMaker Unified Studio, verify that your environment meets these requirements:

  • Athena ODBC driver – The latest Amazon Athena ODBC driver (version 2.2.0 or more recent) for Windows 64-bit.
  • Microsoft Power BI Desktop – The latest version installed on your Windows machine.
  • Microsoft Power BI Pro License – Required for publishing reports and configuring the on-premises data gateway.
  • Microsoft Power BI on-premises data gateway – The latest version installed on the EC2 instance.
  • Amazon SageMaker Unified Studio – An Amazon SageMaker Unified Studio IDC-based domain.

You need an Amazon SageMaker Unified Studio project with data assets. For detailed instructions, refer to the Amazon SageMaker Unified Studio User Guide.

The following screenshot shows the Amazon SageMaker Unified Studio project Query Editor interface, which runs a preview query against the EIA-860 generators dataset.

SageMaker Unified Studio Query Editor previewing the EIA-860 generators dataset

Figure 2: SageMaker Unified Studio project with the EIA-860 generators dataset available in the data catalog

Method 1: DSN-based connection (Athena Power BI connector)

This method uses the Amazon Athena Power BI connector with an ODBC Data Source Name (DSN), supporting DirectQuery and Import mode.

You configure Power BI Desktop to connect to your data assets in Amazon SageMaker Unified Studio using the SageMakerBrowserIdc authentication mode. The driver opens a browser window and authenticates through IAM Identity Center (and your external identity provider, if configured).

Add your SSO user as a member of your SageMaker Unified Studio project

Your single sign-on (SSO) user needs project-level access to query data with Athena. Verify your user is listed as a project member or add it by following Add project members in the Amazon SageMaker Unified Studio User Guide.

The following screenshot shows the SageMaker Unified Studio project user management page, where project owners can add or remove project users and roles.

SageMaker Unified Studio project members page listing users and roles

Figure 3: Members of a SageMaker Unified Studio project

Gather configuration values to configure your Amazon Athena ODBC DSN

Gather the following values from your Amazon SageMaker Unified Studio project:

  1. Open your Amazon SageMaker Unified Studio project.
  2. In the top right, select the three dots.
  3. Choose Project details.
  4. Select JDBC and ODBC details.
  5. Under ODBC connection details copy the following information: IDC issuer URL, domain ID, project ID, Athena workgroup name and AWS Region.

The following screenshot shows the Amazon SageMaker Unified Studio project overview page, where you can copy these details.

SageMaker Unified Studio project overview showing ODBC connection details

Figure 4: ODBC connection details

Configure the ODBC DSN

Create a System DSN using the Amazon Athena ODBC driver. For the general DSN creation steps, see Configuring a data source name on Windows in the Amazon Athena User Guide.

Enter the following values:

Field Value
Data Source Name Name your datasource (for example, pbi-idcdomain)
Region The AWS Region where your Amazon SageMaker domain is provisioned (for example, us-east-1)
Catalog AwsDataCatalog
Database default
Workgroup Your Athena workgroup name (for example, workgroup-abcdefghij-klmexample)

In the Authentication Options, configure the following values:

Field Value
Authentication Type SageMakerBrowserIdc
SSO Start URL IAM Identity Center entry point (for example, https://identitycenter.amazonaws.com/ssoins-0example)
SSO Region Region of IAM Identity Center (for example, us-east-1)
SageMaker Domain ID dzd-123456example
SageMaker Project ID abcd12example
SageMaker Domain Region Region of your Amazon SageMaker Unified Studio project (for example, us-east-1)

Choose OK, then Test to verify the connection. Choose Allow Access when prompted by the browser.

The following screenshot shows the consent prompt.

Browser consent prompt requesting access approval during authentication

Figure 5: Browser consent prompt

The following screenshot shows the successful connection test.

ODBC DSN configuration showing a successful connection test with SageMakerBrowserIdc

Figure 6: Successful connection test in the ODBC DSN configuration with SageMakerBrowserIdc authentication

Connect Power BI Desktop to your data

With the DSN configured, you can connect Power BI Desktop to your data catalog and load the generators dataset.

  1. Open Power BI Desktop.
  2. Open the Get Data menu and select More.
  3. Search for and select Amazon Athena and choose Connect.
  4. For Data Source Name (DSN), enter pbi-idcdomain.
  5. Select DirectQuery.
  6. Choose OK.
  7. Choose Use Data Source Configuration and then Connect.
  8. In the AwsDataCatalog folder, navigate to your database.
  9. Select the core_eia860__scd_generators table.
  10. Choose Load.

The following screenshot shows Power BI Desktop successfully connected to the AWS data catalog.

Power BI Desktop connected to the data catalog with the generators table loaded

Figure 7: Power BI Desktop connected to the data catalog with the generators table loaded using SageMakerBrowserIdc authentication

Create your dashboard and publish it

You can create a dashboard to visualize U.S. power generation data. To create a visualization, complete the following steps:

  1. In the Visualizations pane, choose the Stacked bar chart.
  2. Assign the Y-Axis: Drag technology_description to the Y-Axis.
  3. Assign the X-Axis (Values): Drag capacity_mw to the X-Axis (automatically summed).
  4. Assign the Legend (Stack): Drag operational_status to the Legend field.
  5. Choose Publish.
  6. Give your report a name (for example, generation-idcdomain) and choose Save.
  7. Sign in and choose a destination workspace.
Power BI Desktop stacked bar chart of generation capacity by technology and operational status

Figure 8: Power BI Desktop report using the EIA-860 generators dataset

After publishing, the report structure is available on Power BI Service.

Method 2: DSN-less connection (Power BI ODBC connector)

In this method, you use the Power BI ODBC connector with a connection string (no DSN required). This method supports Import mode only and SageMakerIam authentication. Because the gateway cannot perform browser authentication, both Desktop and gateway must use SageMakerIam. If your workflow requires SageMakerBrowserIdc, use Method 1.

If your machine already has AWS credentials through another method in the default credential provider chain, skip the following setup.

Administrator setup

Create a custom permission set named SageMakerDataAnalyst in IAM Identity Center with the following inline policy. For detailed steps, see Create a permission set in the AWS IAM Identity Center User Guide.

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "SageMakerAccess",
            "Effect": "Allow",
            "Action": [
                "datazone:GetConnection",
                "datazone:ListConnections",
                "datazone:GetDomain",
                "datazone:GetProject"
            ],
            "Resource": "*"
        },
        {
            "Sid": "STSForDriver",
            "Effect": "Allow",
            "Action": [
                "sts:GetCallerIdentity"
            ],
            "Resource": "*"
        }
    ]
}

Assign your user to this permission set for the AWS account containing your SageMaker Unified Studio domain. Then configure your AWS Command Line Interface (AWS CLI) SSO profile by running aws configure sso. For the full CLI configuration walkthrough with detailed steps, see Part 2. After your profile is configured, run aws sso login to authenticate.

Add the IAM identity as a member of SageMaker Unified Studio project

The IAM identity providing credentials needs both domain-level and project-level access to query data through Athena.

  1. Add AWSReservedSSO_SageMakerDataAnalyst_1234example as a domain IAM user: see Managing users in the Amazon SageMaker Unified Studio Admin Guide. Choose Current account.
SageMaker Unified Studio domain users list including the IAM identity

Figure 9: List of users of your SageMaker Unified Studio domain including the IAM identity

  1. Add AWSReservedSSO_SageMakerDataAnalyst_1234example as a project member: see Add project members in the Amazon SageMaker Unified Studio User Guide.
SageMaker Unified Studio project members list including the IAM identity

Figure 10: Members of a SageMaker Unified Studio project including the IAM identity

Gather configuration values

Gather the following connection values from your Amazon SageMaker Unified Studio project:

  1. Open your Amazon SageMaker Unified Studio Project.
  2. On the navigation pane, choose Overview.
  3. Select JDBC and ODBC details.
  4. Select the Using IAM auth toggle.
  5. Copy the ODBC connection string.
SageMaker Unified Studio project overview showing the ODBC connection string for IAM auth

Figure 11: ODBC connection string on the SageMaker Unified Studio project overview

Connect Power BI Desktop to your data and publish

With the configuration parameters of your project, you can connect Power BI Desktop to your data catalog and load the generators dataset.

  1. Open Power BI Desktop.
  2. Open the Get Data menu and select More.
  3. Search for and select ODBC and choose Connect.
  4. For Data Source Name (DSN), select (None).
  5. Expand Advanced Options.
  6. In the Connection string field, enter your connection string. For example, Driver={Amazon Athena ODBC (x64)};AwsRegion=us-east-1;Catalog=AwsDataCatalog;Schema=default;Workgroup=workgroup-abcdefghij-klmexample;SageMakerDomainId= dzd-123456example;SageMakerProjectId= abcd12example;SageMakerDomainRegion=us-east-1;AuthenticationType=SageMakerIam;
  7. Choose OK.
  8. Choose Default or Custom and then Connect.
  9. In the AwsDataCatalog folder, navigate to your database.
  10. Select the core_eia860__scd_generators table.
  11. Choose Load.

When publishing, name your report generation-idcdomain-dsnless.

Configure the on-premises data gateway and view your report on Power BI Service

After creating your reports in Power BI Desktop, configure the on-premises data gateway to view your report on Power BI Service.

You can configure the gateway using either a DSN or a DSN-less connection string, matching the method you used in Power BI Desktop.

Create and attach an IAM role to the Power BI Gateway EC2 instance

Create an IAM role for the EC2 instance that will host your Power BI gateway. Name the role pbi-gateway-role (or a name of your choice). The role must use EC2 as the trusted entity and include the following inline policy:

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "SageMakerAccess",
            "Effect": "Allow",
            "Action": [
                "datazone:GetConnection",
                "datazone:ListConnections",
                "datazone:GetDomain",
                "datazone:GetProject"
            ],
            "Resource": "*"
        },
        {
            "Sid": "STSForDriver",
            "Effect": "Allow",
            "Action": [
                "sts:GetCallerIdentity"
            ],
            "Resource": "*"
        }
    ]
}

Attach this role to your Power BI Gateway EC2 instance. For detailed steps on creating and attaching an IAM role to an EC2 instance, refer to IAM roles for Amazon EC2 in the Amazon EC2 User Guide.

Add the Power BI Gateway IAM role as a member of SageMaker Unified Studio project

The gateway IAM role needs project-level access to query data through Athena.

  1. Add the IAM pbi-gateway-role role as a domain IAM user: see Managing users in the Amazon SageMaker Unified Studio Admin Guide. Choose Current account (or Associated account if your gateway is deployed in a different account).

The following screenshot, from the Amazon SageMaker page of the AWS Management Console, shows the list of users of your Amazon SageMaker Unified Studio domain, including the IAM gateway role.

SageMaker Unified Studio domain users list including the Power BI gateway IAM role

Figure 12: List of users of your SageMaker Unified Studio domain including the IAM gateway role

Add the IAM pbi-gateway-role role as a project member: see Add project members in the Amazon SageMaker Unified Studio User Guide.

The following screenshot shows the Amazon SageMaker Unified Studio project user management page listing the project members.

SageMaker Unified Studio project members list including the Power BI gateway IAM role

Figure 13: Members of a SageMaker Unified Studio project including the IAM gateway role

Configure the data source on Power BI Gateway

How you configure the data source depends on the method you used in Power BI Desktop.

Method 1 (DSN-based)

Configure a System DSN on the gateway EC2 instance following the same ODBC DSN steps described in Method 1. When configuring, make sure that:

  • You use the System DSN tab (not User DSN) because the gateway runs as a Windows service under a separate account.
  • The authentication type is set to SageMakerIam regardless of what you used on Desktop.
  • The DSN name matches exactly the one configured on Power BI Desktop (for example, pbi-idcdomain)

Method 2 (DSN-less)

No configuration is needed on the gateway machine itself. You configure the data source directly in Power BI Service.

Configure the data source and view your report on Power BI Service

To view your report, complete the following steps:

  1. Open the workspace where you saved your report.
  2. Search the Semantic Model which has the same name as your report (for example, generation-idcdomain) and choose the More options icon (three dots).
  3. Choose Settings.
  4. Expand Gateway and Cloud Connection.
  5. Choose View Datasources (play icon) on your gateway.
  6. Choose Manually add to gateway.
  7. Add a connection name (for example, pbi-idcdomain).

The next step depends on the method that you chose:

Method 1 (DSN-based)

  1. Add the DSN (for example, pbi-idcdomain) that matches exactly the one configured on Power BI Desktop.

Method 2 (DSN-less)

  1. In the Connection string field, enter the connection string that matches exactly the one used in Power BI Desktop.

Next, continue with the configuration:

  1. Select Anonymous as Authentication Method.
  2. Choose Create.
  3. Expand again Gateway and Cloud Connection.
  4. For Maps to, choose the connection that you created (for example, pbi-idcdomain).
  5. Choose Apply.
  6. Return to the workspace where you saved your report.
  7. On the Content section, choose your report (for example, generation-idcdomain).

The following screenshot shows a Power BI report on Power BI Service.

Published Power BI report rendering on Power BI Service

Figure 14: Power BI report on Power BI Service

You can now see your report online with the data from your Amazon SageMaker Unified Studio project.

Clean up

To avoid additional charges after testing, delete the Amazon SageMaker Unified Studio domain and EC2 instances. Refer to Delete domains and Terminate Instances for instructions.

Conclusion

In this post, you connected Microsoft Power BI to Amazon SageMaker Unified Studio using an IDC-based domain with both DSN-based and DSN-less methods. This provides a direct connection, with no third-party licensing, that maintains data governance. In Part 2, we cover IAM-based domains.

You can automate many steps of this process. For information about automating DSN creation on the Power BI Gateway or Service, refer to How ENGIE automates the deployment of Amazon Athena data sources on Microsoft Power BI. If you don’t want users adding the gateway IAM role directly, you can create a custom blueprint as a self-service tool for gateway role addition. The blueprint uses a ProjectMembership resource with a configurable parameter that project owners can activate at project creation, automatically adding the gateway role as a project contributor.

For additional best practices, refer to the Using Microsoft Power BI with the AWS Cloud Whitepaper. To learn more, visit Amazon SageMaker Unified Studio and Amazon Athena.


About the authors

Ramesh Singh

Ramesh Singh

Ramesh is a Senior Product Manager Technical (External Services) at AWS in Seattle, Washington, currently with the Amazon SageMaker team. He is passionate about building high-performance ML/AI and analytics products that help enterprise customers achieve their critical goals.

Armando Segnini

Armando Segnini

Armando is a Senior Analytics Specialist Solutions Architect at AWS, partnering with enterprise customers to architect scalable data, analytics, and AI platforms. He helps organizations turn complex data challenges into business value through expertise in streaming, BI integration, and generative AI. Outside of work, Armando enjoys traveling with his family, exploring new cultures, photography, and functional fitness competitions.

Gaurav Sharma

Gaurav is a Specialist Solutions Architect (Analytics) at AWS, supporting US public sector customers on their cloud journey. Outside of work, Gaurav enjoys spending time with his family and reading books.

Krishna Atluru

Krishna Atluru

Krishna is an Enterprise Support Lead TAM at AWS. He provides customers with in-depth guidance on improving security posture and operational excellence for their workloads, helping them build secure, resilient, and cost-effective solutions. His areas of expertise include building serverless architectures, and data and analytics solutions. Outside of work, Krishna enjoys cooking, swimming, and traveling.

Saushthav Saxena

Saushthav Saxena

Saushthav is a Software Development Engineer at AWS on the Amazon Athena team, where he has spent the past few years working on distributed systems and data analytics at scale. Based in the San Francisco Bay Area, his background spans full-stack development, high performance computing, and large-scale infrastructure. Outside of work, he enjoys reading sci-fi novels, swimming, and traveling with family and friends.

Connect Amazon SageMaker Unified Studio to Microsoft Power BI – Part 2: IAM-based domains

Post Syndicated from Ramesh H Singh original https://aws.amazon.com/blogs/big-data/connect-amazon-sagemaker-unified-studio-to-microsoft-power-bi-part-2-iam-based-domains/

In Part 1 of this series, we connected Microsoft Power BI to Amazon SageMaker Unified Studio using an IAM Identity Center (IDC)-based domain. The Amazon Athena ODBC driver (version 2.2.0 and later) supports Amazon SageMaker Unified Studio authentication natively, removing the third-party ODBC-JDBC bridge previously required. We walked through both the DSN-based connection and the DSN-less connection, from Power BI Desktop through the on-premises data gateway to Power BI Service, where report viewers access published dashboards.

In this post, you create the same direct connection using an AWS Identity and Access Management (IAM)-based domain. The walkthrough covers the same two connection methods. The differences are the Amazon SageMaker Unified Studio console navigation paths, the configuration values, and an additional administrator setup that provides AWS credentials through AWS IAM Identity Center. This is Part 2 of a two-part series. For a detailed comparison of the two connection methods, see Part 1.

Solution overview

The architecture is the same as the previous post (see the architecture diagram and walkthrough scenario in Part 1). Power BI Desktop connects to Amazon Athena through the ODBC driver and the Amazon SageMaker Unified Studio project governs all data access. At the same time, the on-premises data gateway on an Amazon Elastic Compute Cloud (Amazon EC2) instance bridges the connection to Power BI Service so report viewers can access published dashboards.

The difference is in authentication: An IAM-based domain uses SageMakerIam authentication for both connection methods. The driver retrieves credentials from the AWS default credential provider chain. For this walkthrough, AWS IAM Identity Center provides those credentials through a custom permission set. Power BI Desktop can run on-premises or on an EC2 instance in the AWS Cloud. The gateway EC2 instance authenticates using its attached IAM role.

Prerequisites

Complete the prerequisites from Part 1. Additionally, you need:

  • AWS Command Line Interface (AWS CLI) – The latest version of the AWS CLI installed on your Windows machine. In this post series, the ODBC driver uses the AWS IAM Identity Center profile configured through the CLI for authentication.
  • Amazon SageMaker Unified Studio – An Amazon SageMaker Unified Studio IAM-based domain with AWS IAM Identity Center single sign-on (SSO) enabled.

The following screenshot shows the Amazon SageMaker Unified Studio (IAM-based domain) project Query Editor interface. It runs a preview query on the EIA-860 generators dataset.

SageMaker Unified Studio Query Editor previewing the EIA-860 generators dataset in an IAM-based domain

Figure 1: SageMaker Unified Studio (IAM-based domain) project with the EIA-860 generators dataset available in the data catalog

Administrator setup

This section configures AWS IAM Identity Center to provide credentials for the SageMakerIam authentication mode. It applies to Method 1 (IAM-based domain) and Method 2 (both domain types). If your machine already has AWS credentials available through another method in the default credential provider chain, you can skip this section and proceed directly to the method of your choice. For the full list of credential sources, refer to Credential providers in the AWS SDKs and Tools Reference Guide.

Create a permission set in IAM Identity Center

Create a custom permission set named SageMakerDataAnalyst in IAM Identity Center with the following inline policy. For detailed steps, see Create a permission set in the AWS IAM Identity Center User Guide.

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "SageMakerAccess",
            "Effect": "Allow",
            "Action": [
                "datazone:GetConnection",
                "datazone:ListConnections",
                "datazone:GetDomain",
                "datazone:GetProject"
            ],
            "Resource": "*"
        },
        {
            "Sid": "STSForDriver",
            "Effect": "Allow",
            "Action": [
                "sts:GetCallerIdentity"
            ],
            "Resource": "*"
        }
    ]
}

The "Resource": "*" is required because these API actions do not support resource-level permissions. For more information, see Actions, resources, and condition keys for Amazon DataZone.

This doesn’t grant broad access to your data. These are read-only metadata actions that allow the ODBC driver to discover connection details and retrieve temporary Athena credentials. The actual data access is governed by Amazon SageMaker Unified Studio project membership: Users can only query data within projects where they have been explicitly added as members. The Amazon SageMaker Unified Studio project IAM role provides Athena and Amazon S3 permissions separately.

Assign users to the permission set

To assign users or groups to the target AWS account, complete the following steps:

  1. In the IAM Identity Center console, choose AWS accounts.
  2. Select the target account where your Amazon SageMaker Unified Studio IAM-based domain is deployed.
  3. Choose Assign users or groups.
  4. Select the SSO users or groups that need access.
  5. Select the SageMakerDataAnalyst permission set.
  6. Choose Submit.

Configure AWS IAM Identity Center profile

To configure the AWS IAM Identity Center profile, run the following command in your terminal on Windows:

aws configure sso

When prompted, enter the following values:

Prompt Value
SSO session name For example, smus
SSO start URL The IDC issuer URL. For example, https://identitycenter.amazonaws.com/ssoins-0example
SSO region The SSO Region. For example, us-east-1
SSO registration scopes sso:account:access

A browser window opens for authentication. After authentication, select your account and the SageMakerDataAnalyst role.

The following screenshots show the consent window and the successful authentication message.

Browser consent prompt requesting access approval during AWS CLI SSO authentication

Figure 2: Browser consent prompt

Browser page confirming successful AWS CLI SSO authentication

Figure 3: Browser authentication successful message

When prompted, enter the following values:

Prompt Value
Default client Region None
CLI default output format None
Profile Name Change value by default

The resulting ~/.aws/config file should look like the following:

[default]
sso_session = smus
sso_account_id = 1234example
sso_role_name = SageMakerDataAnalyst

[sso-session smus]
sso_start_url = https://identitycenter.amazonaws.com/ssoins-0example
sso_region = us-east-1
sso_registration_scopes = sso:account:access

Verify authentication and daily use

To verify that your SSO profile is working correctly, run the following command:

aws sts get-caller-identity

You should receive a response like the following:

{
    "UserId": "AROARHJJNFBQD6EXAMPLE:[email protected]",
    "Account": "111122223333",
    "Arn": "arn:aws:sts::111122223333:assumed-role/AWSReservedSSO_SageMakerDataAnalyst_1234example/[email protected]"
}

For daily use, no passwords or EC2 instance roles are required. When your SSO session expires, run the following command to quickly refresh it:

aws sso login

Add your IAM identity as a member of your Amazon SageMaker Unified Studio project

The IAM identity providing credentials to the ODBC driver needs project-level access to query data through Athena. If you completed the administrator setup, this is the SSO role associated with your permission set (for example, AWSReservedSSO_SageMakerDataAnalyst_1234example). If you’re using another credential source, add the IAM role or user that provides those credentials. For detailed steps, see Managing users for IAM-based domains in the Amazon SageMaker Unified Studio Administrator Guide.

The following screenshot shows the Amazon SageMaker Unified Studio domain management page, which lists the members in a project.

SageMaker Unified Studio project members list

Figure 4: List of members of your SageMaker Unified Studio project

Gather the information to authenticate

To get the parameters that you need to authenticate, complete these steps:

  1. Open your Amazon SageMaker Unified Studio Project.
  2. Open Domain Management.
  3. Choose Users.
  4. Choose View SSO connection.
  5. Copy the end of the Instance ARN, so we can build the Instance URL like https://identitycenter.amazonaws.com/ssoins-0example

The following screenshot shows the Amazon SageMaker Unified Studio domain management page with SSO connection details.

SageMaker Unified Studio domain SSO connection details showing the IAM Identity Center instance ARN

Figure 5: AWS IAM Identity Center information

  1. Choose the user icon and copy the Region as shown in the following screenshot.
SageMaker Unified Studio user menu showing the Region

Figure 6: User icon with the Region information

Method 1: DSN-based connection (Athena Power BI connector)

In this method, you configure an ODBC Data Source Name (DSN) and use the Amazon Athena connector in Power BI. This method uses SageMakerIam authentication mode and supports both DirectQuery and Import mode.

This section covers IAM-based domains. For IDC-based domains, see Part 1.

Gather configuration values to configure your Amazon Athena ODBC DSN

Before configuring the ODBC DSN, gather the following connection values from your Amazon SageMaker Unified Studio project:

  1. Open your Amazon SageMaker Unified Studio Project.
  2. Top right, select the three dots.
  3. Choose Project details.
  4. Select JDBC and ODBC details.
  5. Copy the following values: domain ID, Amazon SageMaker project ID, AWS Region, and Athena workgroup.

The following screenshot shows the Amazon SageMaker Unified Studio project overview page, which provides the project details to copy.

SageMaker Unified Studio project details showing domain ID, project ID, Region, and Athena workgroup

Figure 7: Project details with SageMaker domain ID, SageMaker project ID, Region, and Athena workgroup

Configure the ODBC DSN

Create a System DSN using the Amazon Athena ODBC driver. For the general DSN creation steps, see Configuring a data source name on Windows in the Amazon Athena User Guide. Enter the following values:

Field Value
Data Source Name Name your datasource (for example, pbi-iamdomain)
Region The AWS Region where your Amazon SageMaker domain is provisioned (for example, us-east-1)
Catalog AwsDataCatalog
Database default
Workgroup Your Athena workgroup name (for example, workgroup-abcdefghij-klmexample)

In the Authentication Options, configure the following values:

Field Value
Authentication Type SageMakerIam
SageMaker Domain ID dzd-123456example
SageMaker Project ID abcd12example
SageMaker Region Region of your SageMaker Unified Studio project (for example, us-east-1)

Choose OK, then Test to verify the connection. Choose Allow Access when prompted by the browser.

The following screenshot shows the successful connection test.

ODBC DSN configuration showing a successful connection test with SageMakerIam

Figure 8: Successful connection test in the ODBC DSN configuration with SageMakerIam authentication

Connect Power BI Desktop to your data

With the DSN configured, you can connect Power BI Desktop to your data catalog and load the generators dataset.

  1. Open Microsoft Power BI Desktop.
  2. Open the Get Data menu and select More.
  3. Search for and select Amazon Athena and choose Connect.
  4. For Data Source Name (DSN), enter pbi-iamdomain.
  5. Select DirectQuery.
  6. Choose OK.
  7. Choose Use Data Source Configuration and then Connect.
  8. In the AwsDataCatalog folder, navigate to your database.
  9. Select the core_eia860__scd_generators table.
  10. Choose Load.

The following screenshot shows Power BI Desktop successfully connected to the data catalog.

Power BI Desktop connected to the data catalog with the generators table loaded

Figure 9: Power BI Desktop connected to the data catalog with the generators table loaded using SageMakerIam authentication

Create your dashboard and publish it

You can create a dashboard to visualize U.S. power generation data. To create a visualization, complete the following steps:

  1. In the Visualizations pane, choose the Stacked bar chart.
  2. Assign the Y-Axis: Drag technology_description to the Y-Axis.
  3. Assign the X-Axis (Values): Drag capacity_mw to the X-Axis (automatically summed).
  4. Assign the Legend (Stack): Drag operational_status to the Legend field.
  5. Choose Publish.
  6. Give your report a name (for example, generation-iamdomain) and choose Save.
  7. Sign in and choose a destination workspace.

The following screenshot shows the Power BI dashboard with U.S. power generation data.

Power BI stacked bar chart of U.S. generation capacity by technology and operational status

Figure 10: Power BI dashboard with U.S. power generation data

After you publish, the report structure becomes available on Microsoft Power BI Service.

Method 2: DSN-less connection (Power BI ODBC connector)

In this method, you use the Power BI ODBC connector with a connection string (no DSN required). This method supports Import mode only and SageMakerIam authentication. Because the gateway can’t perform browser authentication and connection strings need to match, both Desktop and gateway must use SageMakerIam.

This section covers IAM-based domains. For IDC-based domains, see Part 1.

Gather configuration values to configure your DSN-less connection

Gather the following connection values from your Amazon SageMaker Unified Studio project:

  1. Open your Amazon SageMaker Unified Studio Project.
  2. Top right, select the three dots.
  3. Choose Project details.
  4. Select JDBC and ODBC details.
  5. Copy the ODBC connection string.

The following screenshot shows the Amazon SageMaker Unified Studio project overview page with the ODBC connection string to copy.

SageMaker Unified Studio project overview showing the ODBC connection string

Figure 11: Project details with ODBC connection string

Connect Power BI Desktop to your data and publish

With the configuration parameters of your project, you can connect Power BI Desktop to your data catalog and load the generators dataset.

  1. Open Power BI Desktop.
  2. Open the Get Data menu and select More.
  3. Search for and select ODBC and choose Connect.
  4. For Data Source Name (DSN), select (None).
  5. Expand Advanced Options.
  6. In the Connection string field, enter your connection string. For example, Driver={Amazon Athena ODBC (x64)};AwsRegion=us-east-1;Catalog=AwsDataCatalog;Schema=default;Workgroup=workgroup-abcdefghij-klmexample;SageMakerDomainId= dzd-123456example;SageMakerProjectId= abcd12example;SageMakerDomainRegion=us-east-1;AuthenticationType=SageMakerIam;
  7. Choose OK.
  8. Choose Default or Custom and then Connect.
  9. In the AwsDataCatalog folder, navigate to your database.
  10. Select the core_eia860__scd_generators table.
  11. Choose Load.

When publishing, name your report generation-iamdomain-dsnless.

Configure the gateway and view your report on Power BI Service

After creating your reports in Power BI Desktop, configure the on-premises data gateway to view your report on Power BI Service.

You can configure the gateway using either a DSN or a DSN-less connection string, matching the method you used in Power BI Desktop.

Create and attach an IAM role to the Power BI Gateway EC2 instance

Create an IAM role for the EC2 instance that will host your Power BI gateway. Name the role pbi-gateway-role (or a name of your choice). The role must use EC2 as the trusted entity and include the following inline policy:

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "SageMakerAccess",
            "Effect": "Allow",
            "Action": [
                "datazone:GetConnection",
                "datazone:ListConnections",
                "datazone:GetDomain",
                "datazone:GetProject"
            ],
            "Resource": "*"
        },
        {
            "Sid": "STSForDriver",
            "Effect": "Allow",
            "Action": [
                "sts:GetCallerIdentity"
            ],
            "Resource": "*"
        }
    ]
}

Attach this role to your Power BI Gateway EC2 instance. For detailed steps on creating and attaching an IAM role to an EC2 instance, refer to IAM roles for Amazon EC2 in the Amazon EC2 User Guide.

Add the Power BI Gateway IAM role as a member of SageMaker Unified Studio project

The gateway IAM role needs project-level access to query data through Athena. The steps to add the role differ depending on your domain type.

IAM-based domain

  1. Open your Amazon SageMaker Unified Studio Project.
  2. Open Domain Management.
  3. Choose your Project Name.
  4. Choose Members.
  5. Choose Add members.
  6. Select the IAM role of your Power BI gateway (for example, pbi-gateway-role).
  7. Choose Add.

The following screenshot shows the Amazon SageMaker Unified Studio project domain management page with options to add members to a project.

SageMaker Unified Studio project members list including the Power BI gateway IAM role

Figure 12: List of members of a SageMaker Unified Studio project with the IAM gateway role

Configure the data source on Power BI Gateway

How you configure the data source depends on the method you used in Power BI Desktop.

Method 1 (DSN-based)

Configure a System DSN on the gateway EC2 instance following the same ODBC DSN steps described in Method 1. When configuring, make sure that:

  • You use the System DSN tab (not User DSN) because the gateway runs as a Windows service under a separate account.
  • The authentication type is set to SageMakerIam.
  • The DSN name matches exactly the one configured on Power BI Desktop (for example, pbi-iamdomain).

Method 2 (DSN-less)

No configuration is needed on the gateway machine itself. You configure the data source directly in Power BI Service.

Configure the data source and view your report on Power BI Service

To view your report, complete the following steps:

  1. Open the workspace where you saved your report.
  2. Search the Semantic Model which has the same name as your report (for example, generation-iamdomain) and choose the More options icon (three dots).
  3. Choose Settings.
  4. Expand Gateway and Cloud Connection.
  5. Choose View Datasources (play icon) on your gateway.
  6. Choose Manually add to gateway.
  7. Add a connection name (for example, pbi-iamdomain).

The next step depends on the method that you chose:

Method 1 (DSN-based)

  1. Add the DSN (for example, pbi-iamdomain) that matches exactly the one configured on Power BI Desktop.

Method 2 (DSN-less)

  1. In the Connection string field, enter the connection string that matches exactly the one used in Power BI Desktop.

Next, continue with the configuration:

  1. Select Anonymous as Authentication Method.
  2. Choose Create.
  3. Expand again Gateway and Cloud Connection.
  4. For Maps to, choose the connection that you created (for example, pbi-iamdomain).
  5. Choose Apply.
  6. Return to the workspace where you saved your report.
  7. On the Content section, choose your report (for example, generation-iamdomain).

The following screenshot shows a report on Power BI Service.

Published Power BI report rendering on Power BI Service

Figure 13: Power BI report on Power BI Service

You can now see your report online with the data from your Amazon SageMaker Unified Studio project.

Clean up

To avoid additional charges after testing, delete the Amazon SageMaker Unified Studio domain and EC2 instances. Refer to Delete domains and Terminate Instances for instructions.

Conclusion

In this two-part series, you connected Power BI to Amazon SageMaker Unified Studio through Amazon Athena. Part 1 covered IDC-based domains. This post covered IAM-based domains using SageMakerIam authentication. This provides a direct connection path, with no third-party licensing, while maintaining data governance and security.

You can automate many steps of this process. For information about automating DSN creation on the Power BI Gateway or Service, refer to How ENGIE automates the deployment of Amazon Athena data sources on Microsoft Power BI. If you don’t want users adding the gateway IAM role directly, you can create a custom blueprint as a self-service tool for gateway role addition. The blueprint uses a ProjectMembership resource with a configurable parameter that project owners can activate at project creation, automatically adding the gateway role as a project contributor.

For additional best practices, refer to the Using Microsoft Power BI with the AWS Cloud Whitepaper. To learn more, visit Amazon SageMaker Unified Studio and Amazon Athena.


About the authors

Ramesh H Singh

Ramesh H Singh

Ramesh is a Senior Product Manager Technical at AWS in Seattle, focused on Amazon SageMaker. He’s passionate about building analytics and AI products that help enterprise customers unlock real value from their data. Away from work, he spends his time hiking with family and exploring spirituality. Connect with him on LinkedIn.

Armando Segnini

Armando Segnini

Armando is a Senior Analytics Specialist Solutions Architect at AWS, partnering with enterprise customers to architect scalable data, analytics, and AI platforms. He helps organizations turn complex data challenges into business value through expertise in streaming, BI integration, and generative AI. Outside of work, Armando enjoys traveling with his family, exploring new cultures, photography, and functional fitness competitions.

Gaurav Sharma

Gaurav is a Specialist Solutions Architect (Analytics) at AWS, supporting US public sector customers on their cloud journey. Outside of work, Gaurav enjoys spending time with his family and reading books.

Krishna Atluru

Krishna Atluru

Krishna is an Enterprise Support Lead TAM at AWS. He provides customers with in-depth guidance on improving security posture and operational excellence for their workloads, helping them build secure, resilient, and cost-effective solutions. His areas of expertise include building serverless architectures, and data and analytics solutions. Outside of work, Krishna enjoys cooking, swimming, and traveling.

Saushthav Saxena

Saushthav Saxena

Saushthav is a Software Development Engineer at AWS on the Amazon Athena team, where he has spent the past few years working on distributed systems and data analytics at scale. Based in the San Francisco Bay Area, his background spans full-stack development, high-performance computing, and large-scale infrastructure. Outside of work, he enjoys reading sci-fi novels, swimming, and traveling with family and friends.

Track SageMaker Unified Studio project costs with custom tags and AWS CUR

Post Syndicated from Nisha Gambhir original https://aws.amazon.com/blogs/big-data/track-sagemaker-unified-studio-project-costs-with-custom-tags-and-aws-cur/

Organizations running machine learning (ML), analytics, and generative AI workloads on Amazon SageMaker Unified Studio domains and projects face a common cost governance challenge. System tags (AmazonDataZoneDomainId and AmazonDataZoneProject) are automatically propagated to all underlying project resources. However, custom tags such as CostCenter, Team, or Environment are not propagated to dynamic resources created through the Studio UI. This creates a gap when you need to report project costs grouped by custom tags.

In this post, we walk through a serverless solution that bridges this gap by enriching AWS Cost and Usage Report (CUR) data with custom project tags. By the end of this post, you can build an Amazon Quick Sight dashboard to filter and analyze Amazon SageMaker Unified Studio project costs by any custom tag dimension that you define. This gives your team the visibility to make informed spending decisions.

Solution overview

The solution consists of three automated subsystems:

  1. Event-driven tag lookup management – An Amazon EventBridge rule captures Amazon DataZone project lifecycle events (Create, Update, Delete) and triggers an AWS Lambda function. The function maintains an Amazon DynamoDB lookup table that maps each project’s DomainId and ProjectId to its custom tags.
  2. CUR enrichment pipeline – An AWS Glue extract, transform, and load (ETL) job reads CUR 2.0 Parquet data from Amazon Simple Storage Service (Amazon S3). The job joins each billing line item with the DynamoDB lookup table using the system tags (DomainId, ProjectId), appends the custom tag values as new columns, and writes the enriched data back to Amazon S3.
  3. Cost visualization – An Amazon Quick Sight dashboard backed by a custom SQL dataset over Amazon Athena provides interactive cost and consumption analytics filtered by custom tags.

Architecture

The following diagram shows the end-to-end architecture:

Figure 1: SageMaker Unified Studio project custom tag cost reporting

The workflow is as follows:

  • An Amazon SageMaker Unified Studio administrator creates or updates a project with custom tags.
  • AWS CloudTrail captures the API call.
  • Amazon EventBridge matches the event.
  • The Lambda orchestrator writes the tag mapping to DynamoDB.
  • Separately, AWS Data Exports delivers CUR data to Amazon S3.
  • The AWS Glue ETL job enriches CUR line items with custom tags from DynamoDB.
  • The AWS Glue Crawler catalogs the enriched data.
  • Amazon Quick Sight visualizes costs by custom tags.

Prerequisites

Before deploying this solution, you need:

  • An Amazon SageMaker Unified Studio domain (you create projects after deployment).
  • AWS Cloud Development Kit (AWS CDK) CLI installed.
  • Python 3.12+.
  • Amazon Quick Sight Enterprise edition enabled in your account.
  • An AWS Identity and Access Management (IAM) user or role with permissions to deploy AWS CloudFormation stacks.

Step 1: Configure custom tags on your project profile

You configure custom tags on project profiles through the Amazon DataZone API. First, enable custom tags on your project profile:

aws datazone update-project-profile \
  --domain-identifier $DOMAIN_ID \
  --identifier $PROJECT_PROFILE_ID \
  --region $REGION \
  --allow-custom-project-resource-tags \
  --project-resource-tags '[
  {"key": "CostCenter", "value": "default", "isValueEditable": true},
  {"key": "Team", "value": "default", "isValueEditable": true},
  {"key": "Environment", "value": "default", "isValueEditable": true}
]'

When creating or updating a project, set the tag values:

aws datazone update-project \
  --domain-identifier $DOMAIN_ID \
  --identifier $PROJECT_ID \
  --project-profile-version latest \
  --region $REGION \
  --resource-tags '{"CostCenter": "CC-100", "Team": "ML-Platform", "Environment": "Production"}'

Important: The AmazonSageMakerProvisioning-<domainAccountId> role needs an inline policy that permits your custom tag keys. Without this, project environment deployment fails.

The following is the inline policy that’s used for the custom tags shared in this post:

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "AllowCustomTagKeys",
      "Effect": "Allow",
      "Action": [
        "sagemaker:AddTags",
        "sagemaker:DeleteTags",
        "cloudformation:TagResource",
        "cloudformation:CreateStack",
        "cloudformation:UpdateStack"
      ],
      "Resource": "*",
      "Condition": {
        "ForAnyValue:StringLike": {
          "aws:TagKeys": [
            "AmazonDataZone*",
            "CostCenter",
            "Team",
            "Environment"
          ]
        }
      }
    }
  ]
}

Step 2: Activate cost allocation tags

Activate the SageMaker Unified Studio system tags as cost allocation tags so they appear in CUR data:

aws ce update-cost-allocation-tags-status \
  --cost-allocation-tags-status '[
  {"TagKey": "AWSDataZoneProject", "Status": "Active"},
  {"TagKey": "AmazonDataZoneDomainId", "Status": "Active"}
]'

These tags take up to 24 hours to start appearing in CUR reports after activation.

Step 3: Deploy the infrastructure

The solution is packaged as a CDK application. Clone the GitHub repository and deploy:

# Install dependencies
pip install -r requirements.txt

# Bootstrap CDK (first time only)
cdk bootstrap aws://$ACCOUNT_ID/$REGION

# Deploy
cdk deploy

This creates the following resources:

  • DynamoDB table (smus-project-tag-lookup) – stores project-to-tag mappings.
  • Lambda function (smus-orchestrator) – processes project lifecycle events.
  • Amazon EventBridge rule – matches Amazon DataZone CreateProject/UpdateProject/DeleteProject events.
  • S3 buckets – for raw CUR and enriched CUR data.
  • AWS Glue ETL job (smus-cur-enrichment) – enriches CUR with custom tags.
  • AWS Glue Crawler – catalogs enriched data.
  • Amazon Simple Notification Service (Amazon SNS) topic – pipeline failure alerts.

Note: The solution uses serverless components (Lambda, DynamoDB on-demand, AWS Glue, Amazon Quick Sight), so you only pay for what you use. The primary cost drivers are AWS Glue ETL job execution time and Amazon Quick Sight SPICE storage.

Step 4: Configure CUR delivery

Create a CUR 2.0 export through AWS Data Exports that delivers Parquet files to the CUR S3 bucket created by the stack. The following screenshots show the complete configuration process in the AWS Billing and Cost Management console.

To create the export, follow these steps:

  1. Go to AWS Billing and Cost Management and then choose Data Exports.
  2. Choose Create in the upper right corner of the Exports and dashboards page. The Data Exports console shows any existing exports, their status, export type, data table, and last refresh date.
  3. On the Create export page, under Export details, select Standard data export and enter an export name. Under Data table content settings, select CUR 2.0.
  4. Under Data table configurations, set Time granularity to Hourly. The configuration page also lets you choose additional export content options such as including resource IDs, split cost allocation data, caller identity allocation data, and capacity reservation columns.
  5. Under Data export delivery options, set Compression type and file format to Parquet. Under Data export storage settings, configure the S3 bucket to: smus-cur-report-{account-id}-{region} and set the S3 path prefix as needed. Choose Create to finish.
Data Exports console listing existing exports with status, type, and last refresh date

Figure 2: Data Exports page listing existing exports

Create export page with Standard data export selected and CUR 2.0 chosen

Figure 3: Create export page with Standard data export and CUR 2.0 selected

Data table configurations with time granularity set to Hourly

Figure 4: Data table configurations with time granularity set to Hourly

Data export delivery options with Parquet format and the S3 storage destination configured

Figure 5: Data export delivery options with Parquet format and S3 storage settings

Step 5: How the event-driven tag capture works

When a project is created or updated in Amazon SageMaker Unified Studio (through the Studio UI or API), the following happens automatically:

  1. CloudTrail logs the Amazon DataZone API call.
  2. Amazon EventBridge matches the event.
  3. Amazon EventBridge invokes the Lambda function.
  4. The Lambda extracts custom tags from the CloudTrail event payload.
  5. The Lambda writes a record to DynamoDB with the DomainId, ProjectId, and all custom tag key-value pairs.

The Lambda function reads tags directly from the responseElements.resourceTags field of the CloudTrail event rather than making a separate GetProject API call. This avoids a race condition where GetProject might return empty tags while the project is in the UPDATING state.

def _extract_tags_from_event(detail):
    tags = {}
    response_elements = detail.get("responseElements") or {}
    for tag_entry in response_elements.get("resourceTags", []):
        if isinstance(tag_entry, dict) and "key" in tag_entry:
            tags[tag_entry["key"]] = tag_entry["value"]
    request_params = detail.get("requestParameters") or {}
    req_tags = request_params.get("resourceTags", {})
    if isinstance(req_tags, dict):
        tags.update(req_tags)
    return tags

Step 6: How the CUR enrichment works

The AWS Glue ETL job runs on a schedule (after each CUR delivery):

  1. Reads CUR Parquet files from the CUR S3 bucket.
  2. Reads all records from the DynamoDB lookup table.
  3. Performs a left outer join on DomainId and ProjectId.
  4. Appends custom tag columns (CostCenter, Team, Environment, and so on) to each CUR line item.
  5. Writes enriched Parquet to the enriched S3 bucket.

Line items without a matching project in the lookup table retain all original columns with NULL custom tag values. No data is dropped.

joined_df = cur_df.join(
    lookup_df,
    on=(
        (cur_df[DOMAIN_COL] == lookup_df["domainId"])
        & (cur_df[PROJECT_COL] == lookup_df["projectId"])
    ),
    how="left_outer",
)

Step 7: Set up the Amazon Quick Sight dashboard

After the first ETL run and crawler execution, set up the Amazon Quick Sight dashboard:

python scripts/setup_quicksight.py \
  --account-id $ACCOUNT_ID \
  --region $REGION \
  --quicksight-user $QUICKSIGHT_USER_ARN

This creates a dashboard with five visuals:

  • Cost by Custom Tag (CostCenter) – horizontal bar chart.
  • Cost by Project – horizontal bar chart.
  • Daily Cost Trend – line chart.
  • Cost by Service per Project – stacked bar chart.
  • Usage by Project & Service – summary table.

And six interactive list filters: Domain, Project, CostCenter, Team, Environment, Service.

The custom SQL includes a CASE statement for service categorization:

SELECT
  line_item_usage_start_date,
  line_item_product_code,
  line_item_usage_amount,
  line_item_unblended_cost,
  resource_tags_user_amazondatazone_domain_id AS domain_id,
  resource_tags_user_amazondatazone_project AS project_id,
  costcenter, team, environment,
  CASE
    WHEN line_item_product_code = 'AmazonSageMaker' THEN 'SageMaker'
    WHEN line_item_product_code = 'AmazonS3' THEN 'S3'
    WHEN line_item_product_code = 'AWSGlue' THEN 'Glue'
    ELSE line_item_product_code
  END AS service_category
FROM "smus_cost_reporting"."enriched_cur"
WHERE line_item_unblended_cost > 0

Step 8: Verifying the solution

After deploying the infrastructure and setting up the dashboard, verify that each component of the pipeline is functioning correctly.

8.1 Verify Amazon EventBridge is capturing project events

  1. Open the Amazon EventBridge console.
  2. In the navigation pane, choose Rules.
  3. Select the rule created by the CDK stack (for example, SmusCostReporting-ProjectTagRule).
  4. Choose the Monitoring tab.
  5. Confirm that the invocations are being recorded in the metrics.
  6. Create or update an Amazon SageMaker Unified Studio project with custom tags using the following command:
    aws datazone update-project \
      --domain-identifier <domain-id> \
      --identifier <project-id> \
      --custom-tags CostCenter=Engineering Team=DataPlatform Environment=Production

  7. Within a few seconds, the Amazon EventBridge rule should show a new invocation in its metrics.

8.2 Verify DynamoDB schema and tag mappings

The DynamoDB lookup table uses a simple key schema:

Attribute Type Role
domainId String Partition Key
projectId String Sort Key
CostCenter String Custom tag
Team String Custom tag
Environment String Custom tag

Custom tags are stored as dynamic attributes. Any tag key set on a project becomes a column in the table.

8.2.1 Verify DynamoDB table contains tag mappings

  1. Open the DynamoDB console.
  2. Navigate to the table created by the stack (for example, SmusCostReporting-ProjectTagsTable).
  3. Choose Explore table items.
  4. Scan for your project with the following keys:
    Partition key (domainId): <your-domain-id>
    Sort key (projectId): <your-project-id>

  5. Confirm the item contains the expected custom tag attributes (CostCenter, Team, Environment) with the values you assigned.
  6. Alternatively, use the AWS CLI:
    aws dynamodb get-item \
      --table-name SmusCostReporting-ProjectTagsTable \
      --key '{"domainId": {"S": "<domain-id>"}, "projectId": {"S": "<project-id>"}}'

8.3 Verify the AWS Glue ETL job enriches CUR data

  1. Wait for the next CUR delivery (hourly if configured as described in Step 4).
  2. Wait for the subsequent AWS Glue job execution.
  3. Open the AWS Glue console.
  4. In the navigation pane, choose ETL Jobs.
  5. Confirm the job completed successfully (status: Succeeded).
  6. Query the enriched data in Amazon Athena to confirm custom tag columns are populated:
    SELECT
      line_item_usage_start_date,
      line_item_product_code,
      line_item_unblended_cost,
      costcenter,
      team,
      environment
    FROM "smus_cost_reporting"."enriched_cur"
    WHERE costcenter IS NOT NULL
    LIMIT 10;

You should see rows with your custom tag values populated in the costcenter, team, and environment columns.

8.4 Verify the Amazon Quick Sight dashboard displays enriched data

  1. Open the Amazon Quick Sight console and navigate to the dashboard created by the setup script.
  2. Confirm that:
    • The Cost by Custom Tag (CostCenter) bar chart displays cost data grouped by your CostCenter values.
    • The list filters for CostCenter, Team, and Environment contain selectable values.
    • Selecting a filter value correctly narrows the displayed data.
  3. If the dashboard shows no data, verify that:
    • The AWS Glue Crawler has run after the ETL job (check the crawler’s last run status in the AWS Glue console).
    • The SPICE dataset has been refreshed. In the Amazon Quick Sight console, navigate to Datasets, select the dataset, and then choose Refresh now.

Figure 6 shows the Amazon Quick Sight dashboard with two side-by-side horizontal bar charts: Cost by Cost Center and Cost by Project. Domain Name and Project Name list filters appear at the top.

Amazon Quick Sight dashboard with Cost by Cost Center and Cost by Project bar charts and Domain and Project filters

Figure 6: Amazon Quick Sight dashboard showing cost data by custom tags, including Cost by Cost Center and Cost by Project bar charts with Domain Name and Project Name filters

Note: The first end-to-end cycle can take up to 48 hours depending on CUR delivery timing. After the initial cycle completes, subsequent updates will flow automatically on the configured schedule.

Operational considerations

Monitoring: The Amazon SNS topic smus-cost-reporting-alerts receives notifications when the AWS Glue ETL job fails or the Lambda orchestrator encounters repeated errors. Subscribe an email address or Slack webhook to stay informed. For instructions on how to create a subscription, see Subscribing to an Amazon SNS topic.

Cost: The solution uses serverless components (Lambda, DynamoDB on-demand, AWS Glue, Amazon Quick Sight, SPICE) so you only pay for what you use. The primary cost drivers are AWS Glue ETL job execution time and Amazon Quick Sight SPICE storage.

Scaling: The DynamoDB table uses on-demand capacity and can scale to accommodate your projects. You can scale the AWS Glue ETL job by increasing the number of workers for larger CUR datasets. For more information, see Managing throughput capacity automatically with DynamoDB auto scaling.

New tag keys: When you add new custom tag keys to projects, the ETL automatically picks them up as new columns. The AWS Glue Crawler’s UPDATE_IN_DATABASE policy adds new columns to the catalog table without manual intervention.

Cleanup

Warning: The following cleanup steps will permanently delete all CUR data, project tag mappings, and Amazon Quick Sight dashboards.

To remove all resources:

# Delete Amazon Quick Sight resources
python scripts/setup_quicksight.py --account-id $ACCOUNT_ID --region $REGION --quicksight-user $QS_USER --clean

# Delete CDK stack
cdk destroy

Go to AWS Billing and Cost Management, and then choose Data Exports and delete the CUR 2.0 export created in Step 4.

Deactivate the cost allocation tags that were activated in Step 2:

aws ce update-cost-allocation-tags-status \
  --cost-allocation-tags-status '[
  {"TagKey": "AWSDataZoneProject", "Status": "Inactive"},
  {"TagKey": "AmazonDataZoneDomainId", "Status": "Inactive"}
]'

Conclusion

In this post, we showed how to build an end-to-end cost reporting solution for Amazon SageMaker Unified Studio projects using custom tags. This solution combines tag capture driven by Amazon EventBridge, CUR enrichment through AWS Glue ETL, and visualization in Amazon Quick Sight. With it, organizations can track and attribute costs by CostCenter, Team, Environment, or any custom dimension. This works even for resources created through the Studio UI that don’t receive custom tag propagation.

This solution serves as an extension to the custom tag propagation feature and reports cost for all project resources. The architecture is fully serverless, automated, and can be deployed to any AWS account using the provided CDK application.

To start building your custom tag cost reporting pipeline, visit the GitHub repository. To learn more about the underlying services, visit the Amazon SageMaker Unified Studio service page. For a related approach to custom tag governance, see Use Amazon SageMaker custom tags for project resource governance and cost tracking

References


About the authors

Nisha Gambhir

Nisha Gambhir

Nisha is a Senior AI/ML & Cloud Architect based out of India. She is passionate about helping customers design, architect and develop secure, scalable and reliable applications using AI/ML and Agentic AI. She loves working on latest technologies, providing simple and scalable solutions that drive positive business outcomes.

Dr Anil Giri

Dr Anil Giri

Anil is a Solutions Architect at AWS, based in London, UK, where he helps ISV customers design and deploy agentic AI systems in production. He specializes in multi-agent orchestration, retrieval-augmented generation, and event-driven serverless architectures on Amazon Bedrock, with a focus on building reliable, secure, and scalable solutions that deliver measurable business outcomes.

Satish Sarapuri

Satish Sarapuri

Satish is a Sr. Data Architect, Data Mesh / Data Lake/Gen AI at AWS. He helps enterprise-level customers build high-performance, highly available, cost-effective, resilient, and secure generative AI, data mesh, data lake, and analytics platform solutions on AWS, through which customers can make data-driven decisions to gain impactful outcomes for their business and help them on their digital and data transformation journey. In his spare time, he enjoys trail running and spending quality time with his family.

Ram Vittal

Ram Vittal

Ram is a Principal GenAI/ML Specialist at AWS. He has over 3 decades of experience building distributed, hybrid, and cloud applications. He is passionate about building secure, scalable, reliable AI/ML and big data solutions to help customers with their cloud adoption and optimization journey. In his spare time, he rides motorcycle and enjoys the nature with his family.

Secure SageMaker Unified Studio access with SAML and conditional policies

Post Syndicated from Manos Samatas original https://aws.amazon.com/blogs/big-data/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 Applications → Create 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 providers → Add 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 Roles → Create 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 General → SAML Settings → Edit.
  3. Choose Next until you reach Attribute Statements.
  4. Add the following mappings:
    • https://aws.amazon.com/SAML/Attributes/PrincipalTag:Email → user.email.
    • https://aws.amazon.com/SAML/Attributes/Role → {IAMROLEARN,IdentityProviderARN}.
    • https://aws.amazon.com/SAML/Attributes/RoleSessionName → user.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.

Scaling fine-grained access control for enterprise lakehouse using SageMaker Unified Studio and AWS Lake Formation

Post Syndicated from Chintan Agrawal original https://aws.amazon.com/blogs/big-data/scaling-fine-grained-access-control-for-enterprise-lakehouse-using-sagemaker-unified-studio-and-aws-lake-formation/

As enterprise lakehouses grow to thousands of tables across multiple business domains and regions, scaling fine-grained access control becomes a critical governance challenge. Data governance teams spend significant time manually granting table-level permissions, only to face permission drift, inconsistent enforcement, and limited auditability. Without a scalable approach, each new dataset requires manual policy updates, increasing the risk of unauthorized access and slowing time-to-insight for analysts and data scientists.

In this post, we show you how to solve this problem by combining AWS IAM Identity Center, AWS Lake Formation tag-based access control (TBAC), and trusted identity propagation in Amazon SageMaker Unified Studio. You deploy a complete governance architecture using AWS Cloud Development Kit (AWS CDK) that classifies data with LF-Tags, maps IAM Identity Center groups to tag-based policies, and enforces permissions at query time across analytics engines. The solution uses Apache Iceberg tables stored in Amazon Simple Storage Service (Amazon S3) and registered in the AWS Glue Data Catalog.

The core governance challenge

As organizations mature their lakehouse environments, governance complexity increases with each new dataset. Several challenges commonly emerge:

  • Explosive dataset growth: Iceberg-based lakehouses often contain thousands of tables distributed across raw, curated, and conformed zones. Each new dataset introduces additional governance requirements, making table-level permission grants operationally expensive.
  • Multi-domain data ownership: Enterprise lakehouses typically serve multiple business domains such as commercial analytics, clinical research, and regulatory reporting. These domains require strict isolation while still supporting controlled data sharing.
  • Regional data sovereignty: Organizations operating globally must enforce geographic boundaries for sensitive datasets. EU clinical trial data might be restricted by GDPR regulations, whereas US commercial datasets follow different compliance frameworks.
  • Sensitivity-based access controls: Within each domain, datasets vary in sensitivity. Pricing strategies, drug discovery research, and patient-related datasets require stricter access controls than standard operational data.
  • Role explosion: Pure RBAC approaches attempt to encode these dimensions into roles, leading to role proliferation. Manual Lake Formation grants at the table level create permission drift and limited scalability.

To address these challenges, enterprise lakehouse governance must satisfy several criteria:

  • Least-privilege access.
  • Dynamic scalability as new datasets are onboarded.
  • Multi-dimensional enforcement across domain, region, and sensitivity.
  • Auditability traceable to individual users.
  • Automation-ready, configuration-driven workflows.

TBAC addresses each of these challenges directly. Instead of granting permissions on individual tables, you define tag-based policies that automatically apply to any resource matching the tag expression. New datasets inherit access rules through tag inheritance, eliminating manual policy updates (solving explosive dataset growth). Domain and region tags enforce strict isolation between business units (solving multi-domain ownership and regional sovereignty). Sensitivity tags control access within domains without role proliferation (solving sensitivity-based controls and role explosion). The following sections describe the architecture that implements this model and walk you through deploying it end to end.

Reference architecture overview

The governance model integrates identity, metadata, and lakehouse services into a unified access architecture that enforces fine-grained permissions consistently across analytics and machine learning (ML) workloads. The architecture consists of five layers, each handling a distinct responsibility in the access control flow.

The following diagram illustrates the end-to-end architecture, showing how user identity flows from IAM Identity Center through SageMaker Unified Studio to Lake Formation for tag-based policy evaluation against the AWS Glue Data Catalog and Amazon S3 storage layer.

Architecture linking IAM Identity Center, SageMaker Unified Studio, Lake Formation, the Glue Data Catalog, and Amazon S3

Figure 1: End-to-end governance architecture for the enterprise lakehouse

1. Identity and authentication layer: IAM Identity Center manages user identities and group memberships, integrates with corporate identity providers, and provides centralized lifecycle management for enterprise users. IAM Identity Center groups represent business roles and serve as the principals that receive Lake Formation permissions.

2. Unified analytics and ML access layer: Amazon SageMaker Unified Studio serves as the primary interface where analysts, data scientists, and ML engineers discover datasets, run queries, and build ML workflows. Because SageMaker Unified Studio integrates with multiple compute engines, including Amazon Athena, AWS Glue, Amazon EMR, and Amazon Redshift, users can access data using their preferred analytics tools while maintaining consistent governance.

3. Governance and authorization layer: AWS Lake Formation provides fine-grained access control across AWS Glue catalog resources using LF-Tags. Instead of granting permissions directly on databases and tables, Lake Formation evaluates LF-Tag policies dynamically and grants or denies access at query time. Governance teams define access rules once, and Lake Formation automatically applies them to new datasets as they are onboarded.

4. Governance automation layer: Two AWS Lambda functions automate tag assignment and permission provisioning. JSON metadata configuration files drive both pipelines, so governance teams manage access control through configuration rather than manual console operations.

5. Metadata and storage layer: Apache Iceberg tables stored in Amazon S3 form the foundation of the lakehouse. You register these tables in the AWS Glue Data Catalog, which provides centralized metadata management and interoperability across analytics services. Lake Formation evaluates governance decisions at the catalog level rather than independently by each analytics engine.

End-to-end access flow

When a user queries a dataset from SageMaker Unified Studio, the following sequence occurs:

  1. The user authenticates through IAM Identity Center and accesses SageMaker Unified Studio.
  2. SageMaker passes the user’s identity context to downstream analytics services using trusted identity propagation.
  3. The analytics engine requests data access from Lake Formation.
  4. Lake Formation evaluates LF-Tag policies against the user’s IAM Identity Center group membership.
  5. Access is granted or denied dynamically at query time.

Because authorization decisions are centralized in Lake Formation, governance remains consistent regardless of which analytics engine the user employs.

Hybrid RBAC + ABAC governance model

The governance model combines identity context from IAM Identity Center with metadata-driven classification using LF-Tags. The following table summarizes how each layer contributes to the overall governance workflow.

Governance capability IAM Identity Center contribution Lake Formation LF-Tag contribution Governance outcome
Identity context Organizes users into groups aligned with business roles Evaluates permissions using group membership Role-aligned access boundaries
Data classification Provides role eligibility for data access Classifies datasets by domain, region, sensitivity, and layer Attribute-aware authorization
Scalability Simplifies user lifecycle management Automatically applies policies to newly tagged datasets Governance that scales with dataset growth
Operational model Centralizes role lifecycle operations Enables metadata-driven policy automation Reduced administrative overhead

IAM Identity Center defines who can request access, LF-Tags define what datasets are eligible, and Lake Formation enforces policies dynamically at query time.

Enterprise LF-Tag data model

A structured tagging strategy is the foundation of scalable Lake Formation governance. In this solution, the solution classifies datasets across four governance dimensions.

Tag Key Tag Values Purpose Example Usage
region us, eu, global Geographic data location Enforce GDPR compliance for EU data
domain commercial, clinical_research, regulatory Business domain Separate commercial from clinical data
data_class standard, sensitive, regulated Data sensitivity level Restrict access to sensitive pricing data
layer raw, curated, conformed Data processing stage Grant analysts access to curated data only

Together, these dimensions enable multi-dimensional authorization policies that reflect both organizational structure and regulatory requirements.

Tag inheritance and evaluation

LF-Tags can be applied at three resource levels within the Glue Data Catalog: database, table, and column. In this implementation, database-level tags define broad governance attributes (domain, region, layer), table-level tags capture dataset-specific sensitivity (data_class), and column-level tags can further restrict access to individual fields. Lake Formation evaluates the effective tag set at query time by combining inherited and explicitly assigned tags.

For example, a database tagged domain=commercial, region=us, layer=raw automatically applies those tags to all tables within it. A table-level data_class=sensitive tag supplements the inherited tags to distinguish sensitive pricing data from standard sales data. This inheritance model means new tables automatically receive governance coverage without manual tag assignment. To learn more, refer to Lake Formation tag-based access control best practices.

Prerequisites

Before deploying the solution, complete the following setup in the us-east-1 Region. Use the same AWS Region throughout all steps.

  1. AWS account and IAM Identity Center: Enable IAM Identity Center and create test users. Note your Identity Store ID from the IAM Identity Center console under Settings. For setup guidance, see Getting started with IAM Identity Center.
  2. Lake Formation configuration: Complete the following setup in the Lake Formation console:2.1. Change Data Catalog default permissions. In the navigation pane under Administration, choose Data Catalog settings. Uncheck Use only IAM access control for new databases and uncheck Use only IAM access control for new tables in new databases. Choose Save. This makes sure Lake Formation permissions govern access to databases and tables created by the CDK stacks.
    Lake Formation Data Catalog settings with both IAM-only access control checkboxes cleared

    Figure 2: Lake Formation Data Catalog settings with both IAM-only access control checkboxes unchecked

    2.2. Integrate with IAM Identity Center. Complete the prerequisites for IAM Identity Center integration with Lake Formation, including enabling trusted identity propagation.You don’t need to manually create a Lake Formation administrator. The CDK deployment in Step 2: Deploy all stacks automatically registers the required administrators via the LfAdminStack (see lf-admin-stack.ts). S3 data location registration is a post-deployment console step covered after the CDK creates the buckets.

  3. SageMaker Unified Studio: Create a SageMaker Unified Studio domain, select your IAM Identity Center instance for authentication, and enable trusted identity propagation. For a detailed walkthrough, see Accelerate your analytics with Amazon S3 Tables and Amazon SageMaker Lakehouse and enable trusted identity propagation for the domain.
  4. Local tooling: Install AWS Command Line Interface (AWS CLI), Python 3.x, Node.js 18+, AWS CDK CLI (npm install -g aws-cdk), and Git.

Solution overview

Now that you understand the governance model and tag taxonomy, the following section walks you through deploying the complete infrastructure and configuring access control.

The deployment uses AWS CDK (TypeScript) and consists of seven stacks that create the complete governance infrastructure. The CDK app manages stack dependencies automatically, so a single cdk deploy --all command deploys everything in the correct order.

The architecture uses a two-layer data lake pattern. The raw layer stores data as CSV files in Amazon S3, registered as external tables in the AWS Glue Data Catalog. The curated layer uses Apache Iceberg v2 tables for ACID transactions and schema evolution. Three business domains (US Commercial, EU Clinical Research, and Global Regulatory) each have one representative table per layer, giving six tables total.

Lake Formation tag-based access control (TBAC) governs all access using four tag dimensions:

Tag Key Values Purpose
domain commercial, clinical_research, regulatory Business domain isolation
region us, eu Geographic data boundary
data_class standard, sensitive, regulated Sensitivity classification
layer raw, curated Data layer identification

Step 1: Clone the repository and install dependencies

Clone the accompanying repository and install the CDK project dependencies:

git clone https://github.com/aws-samples/sample-aws-smus-governance-automation
cd aws-smus-governance-automation/cdk
npm install

The CDK project is written in TypeScript and uses aws-cdk-lib v2. The lib/ directory contains seven stack definitions, and bin/app.ts wires them together with explicit dependency ordering.

If this is your first CDK deployment in this account and Region, bootstrap the CDK environment. Bootstrapping provisions an S3 bucket and IAM roles that CDK uses to deploy assets:

cdk bootstrap aws://<ACCOUNT_ID>/us-east-1

Step 2: Deploy all stacks

Deploy the entire infrastructure with a single command. Pass your IAM Identity Center Identity Store ID as a CDK context variable:

cdk deploy --all -c identityStoreId=d-xxxxxxxxxx --require-approval never --region us-east-1

CDK will prompt for IAM permission changes on each stack. The --require-approval never flag auto-approves these so the deployment runs unattended.

CDK deploys the seven stacks in dependency order:

  1. LfSetupStack: Lake Formation admin registration + LF-Tags (domain, region, data_class, layer)
  2. GlueRawTablesStack: S3 bucket + three Glue databases + three CSV-backed tables.
  3. GlueCuratedTablesStack: S3 bucket + three Glue databases + three Iceberg v2 tables.
  4. SsoGroupsStack: three IAM Identity Center groups (DataLake-US-Commercial, DataLake-EU-Clinical-Research-Sensitive, DataLake-Regulatory)The three groups map to specific tag combinations that control data access:
    • DataLake-US-Commercial: domain=commercial, region=us, data_class=standard.
    • DataLake-EU-Clinical-Research-Sensitive: domain=clinical_research, region=eu, data_class=sensitive,regulated.
    • DataLake-Regulatory: domain=regulatory (all regions, all data classes within regulatory).

    The following table summarizes the user personas, their group assignments, and the data access each group provides:

  5. AssetTaggingAutomationStack: Tag automation Lambda.
  6. SsoPermissionAutomationStack: Permission automation Lambda.
  7. LfAdminStack: Registers CDK + Lambda roles as Lake Formation admins.

After deployment completes, review the CloudFormation stack outputs. They include S3 bucket names, database names, SSO group IDs, and Lambda function ARNs.

The following figure shows all seven CDK stacks deployed successfully in the CloudFormation console.

CloudFormation console showing all seven CDK stacks in CREATE_COMPLETE status

Figure 3: CloudFormation console showing all seven CDK stacks in CREATE_COMPLETE status

Register S3 data locations with Lake Formation: Now that the S3 buckets exist, register them with Lake Formation. In the Lake Formation console, under Administration, choose Data lake locations, then choose Register location. Register both buckets from the stack outputs (for example, s3://datalake-raw-data-<ACCOUNT_ID>-us-east-1 and s3://datalake-curated-data-<ACCOUNT_ID>-us-east-1). For IAM role, use the default AWSServiceRoleForLakeFormationDataAccess and choose Lake Formation as the permission mode. See Registering an Amazon S3 location for step-by-step instructions.

The following figure shows both data lake S3 locations registered in the Lake Formation console.

Lake Formation Data lake locations page listing the registered raw and curated S3 buckets

Figure 4: Lake Formation Data lake locations page with raw and curated S3 buckets registered

Step 3: Populate sample datasets

The scripts use Amazon Athena to insert sample data. Athena stores query results under the athena-results/ prefix in the shared governance metadata bucket (lf-governance-metadata-<ACCOUNT_ID>-<REGION>) created by the CDK deployment.

Populate the raw and curated tables:

cd ../scripts
python3 populate_raw_layer.py
python3 populate_curated_layer.py

Each script executes INSERT INTO statements through the Athena StartQueryExecution API and waits for completion. You should see success messages for all six tables (three raw, three curated).

After populating the tables, you can verify the data in the Glue Data Catalog. The following figure shows the six tables across the three raw and three curated databases.

AWS Glue Data Catalog showing the six databases and tables created by the deployment

Figure 5: AWS Glue Data Catalog showing the six databases and tables created by the CDK deployment

You can also preview the data by querying a table. The following figure shows sample data from the us_sales_summary table.

Athena query results showing sample commercial rows from the us_sales_summary table

Figure 6: Query results for the us_sales_summary table with sample commercial data

Step 4: Apply LF-Tags to data assets

The following diagram illustrates the governance automation flow, showing how metadata JSON configuration files drive the two Lambda pipelines for asset tagging and SSO permission management.

Governance automation flow with the asset tagging and SSO permission Lambda pipelines

Figure 7: Governance automation flow showing the asset tagging and SSO permission Lambda pipelines

The diagram shows two parallel pipelines, each following three steps:

Asset tagging pipeline (left):

  1. Metadata upload – A data governance administrator uploads metadata JSON files (metadata-raw-tables.json and metadata-curated-tables.json) to the asset-tagging/ prefix in the shared S3 governance metadata bucket. These files define which LF-Tags to assign to each AWS Glue database and table.
  2. Lambda processing – The S3 upload triggers the LakeFormationTagAutomation Lambda function, which reads the metadata and calls the Lake Formation API.
  3. Tag operations – The Lambda creates or updates LF-Tags, then assigns them to the target databases and tables in the AWS Glue Data Catalog.

SSO permission pipeline (right):

  1. Permission upload – Three permission JSON files (one per IAM Identity Center group) are uploaded to the sso-permissions/ prefix. These files define the LF-Tag policy expressions that control data access.
  2. Lambda processing – The upload triggers the LakeFormationSSOPermissionAutomation Lambda function.
  3. Permission operations – The Lambda grants tag-based permissions to the corresponding IAM Identity Center groups through the Lake Formation API.

Both pipelines log execution details to Amazon CloudWatch for monitoring and troubleshooting.

Two metadata JSON configuration files drive the asset tagging Lambda that declaratively define which LF-Tags to apply to each AWS Glue resource:

  • metadata-raw-tables.json: Tag definitions for the three raw layer databases and tables.
  • metadata-curated-tables.json: Tag definitions for the three curated layer databases and tables.

Each entry in these files specifies the following fields:

Field Description Example
catalog_id Your AWS account ID (Glue Data Catalog ID) 123456789012
resource_type DATABASE or TABLE DATABASE
database_name AWS Glue database name raw_us_commercial_db
table_name AWS Glue table name (only for TABLE entries) us_sales_summary
lf_tags Array of LF-Tag key/value pairs to assign [{“TagKey”:“domain”,“TagValues”:[“commercial”]}]
access_type Action to perform (GRANT) GRANT

Parameters you must update before invoking: Replace the catalog_id value in every entry of both files with your own AWS account ID. The database and table names match the resources created by the CDK stacks, so those should not be changed unless you customized the stack parameters.

The following snippet from metadata-raw-tables.json shows a database-level entry and a table-level entry:

[
  {
    "comment": "DATABASE LEVEL TAGS - US Commercial RAW Domain",
    "access_type": "GRANT",
    "resource_type": "DATABASE",
    "catalog_id": "<YOUR_ACCOUNT_ID>",
    "database_name": "raw_us_commercial_db",
    "lf_tags": [
      { "TagKey": "region", "TagValues": ["us"] },
      { "TagKey": "domain", "TagValues": ["commercial"] },
      { "TagKey": "layer", "TagValues": ["raw"] }
    ]
  },
  {
    "comment": "TABLE LEVEL TAGS - US Commercial RAW Table (Standard Access)",
    "access_type": "GRANT",
    "resource_type": "TABLE",
    "catalog_id": "<YOUR_ACCOUNT_ID>",
    "database_name": "raw_us_commercial_db",
    "table_name": "us_sales_summary",
    "lf_tags": [
      { "TagKey": "data_class", "TagValues": ["standard"] }
    ]
  }
]

The Lambda applies tags at two levels: database-level entries assign domain, region, and layer tags, while table-level entries assign the data_class tag (standard, sensitive, or regulated). Because of two-level tagging, new tables added to a tagged database automatically inherit the database-level tags. Only the table-specific data_class tag needs explicit assignment. To learn more about this pattern, refer to Lake Formation tag-based access control best practices.

Invoke the Lambda for both layers:

cd ../lf-asset-tagging-automation
aws lambda invoke \
    --function-name LakeFormationTagAutomation \
    --payload fileb://metadata-raw-tables.json \
    --cli-binary-format raw-in-base64-out \
    response.json

aws lambda invoke \
    --function-name LakeFormationTagAutomation \
    --payload fileb://metadata-curated-tables.json \
    --cli-binary-format raw-in-base64-out \
    response.json

Verify tag assignment using the GetResourceLFTags API:

aws lakeformation get-resource-lf-tags \
    --resource '{"Table":{"DatabaseName":"raw_us_commercial_db","Name":"us_sales_summary"}}' \
    --region us-east-1

You should see domain=commercial, region=us, layer=raw, and data_class=standard in the response.

The following figure shows the LF-Tags assigned to the us_sales_summary table in the Lake Formation console, confirming that both database-level inherited tags and table-level tags are applied correctly.

Lake Formation console showing inherited and table-level LF-Tags on the us_sales_summary table

Figure 8: LF-Tags on the us_sales_summary table showing inherited and table-level tags

Step 5: Provision SSO group permissions

Three permission JSON files (one per IAM Identity Center group) define the LF-Tag policy expressions. Update sso_group with the group UUID from the SsoGroupsStack outputs and identity_center_account_id with your AWS account ID. For detailed configuration, see the repository README.

[
  {
    "sso_name": "DataLake-US-Commercial",
    "sso_group": "<GROUP_UUID_FROM_CDK_OUTPUT>",
    "identity_center_account_id": "<YOUR_ACCOUNT_ID>",
    "resources": [
      {
        "resource_type": "DATABASE",
        "permissions": ["DESCRIBE"],
        "lf_tag_expression": [
          { "TagKey": "domain", "TagValues": ["commercial"] },
          { "TagKey": "region", "TagValues": ["us"] },
          { "TagKey": "layer", "TagValues": ["curated", "raw"] }
        ]
      },
      {
        "resource_type": "TABLE",
        "permissions": ["SELECT", "DESCRIBE"],
        "lf_tag_expression": [
          { "TagKey": "domain", "TagValues": ["commercial"] },
          { "TagKey": "region", "TagValues": ["us"] },
          { "TagKey": "data_class", "TagValues": ["standard"] },
          { "TagKey": "layer", "TagValues": ["curated", "raw"] }
        ]
      }
    ]
  }
]

Apply permissions for each group:

cd ../lf-sso-permission-automation
python3 lambda_function.py us-commercial-permissions.json
python3 lambda_function.py eu-clinical-research-sensitive-permissions.json
python3 lambda_function.py regulatory-permissions.json
aws lakeformation list-permissions \
    --principal '{"DataLakePrincipalIdentifier":"arn:aws:identitystore:::group/<GROUP_ID>"}' \
    --region us-east-1

Step 6: Validate fine-grained access control

With all permissions in place, validate that Lake Formation TBAC enforces the correct access boundaries by signing in to SageMaker Unified Studio as different IAM Identity Center users.

Test as Sarah (US Commercial Analyst) — Sarah belongs to DataLake-US-Commercial, which grants access to standard commercial data only.

SELECT * FROM raw_us_commercial_db.us_sales_summary LIMIT 10;

Sarah sees all rows and columns successfully:

SageMaker Unified Studio results: Sarah’s successful query on us_sales_summary

Figure 9: Sarah’s successful query on us_sales_summary in SageMaker Unified Studio

Querying outside her authorized domain returns an access denied error:

SELECT * FROM raw_eu_clinical_research_db.eu_drug_discovery LIMIT 10;
Access denied error when Sarah queries eu_drug_discovery outside her domain

Figure 10: Access denied when Sarah queries eu_drug_discovery, confirming TBAC enforcement

Test as Dr. Chen (EU Clinical Research Lead) — Dr. Chen can access sensitive and regulated EU clinical research data (eu_drug_discovery) but is denied access to US commercial data (us_sales_summary), confirming regional and domain isolation.

Query results showing Dr. Chen’s successful query on eu_drug_discovery

Figure 11: Dr. Chen’s successful query on eu_drug_discovery

Access denied error when Dr. Chen queries us_sales_summary

Figure 12: Access denied when Dr. Chen queries us_sales_summary

Test as Alex (Regulatory Affairs Specialist) — Alex’s tag expression uses only domain=regulatory without a region constraint, granting cross-regional access to regulatory data while maintaining strict isolation from commercial and clinical research domains.

Query results showing Alex’s successful query on fda_submissions

Figure 13: Alex’s successful query on fda_submissions

Access denied error when Alex queries us_sales_summary

Figure 14: Access denied when Alex queries us_sales_summary

These tests demonstrate that TBAC enforces fine-grained permissions based on user identity, data classification, regional boundaries, and domain separation, without per-table permission grants. As new tables are added and tagged, existing groups automatically gain or are denied access based on their tag expressions. This is the core advantage of TBAC over named resource permissions.

Audit user access with CloudTrail

A key benefit of integrating Lake Formation with IAM Identity Center is the detailed audit trail available through AWS CloudTrail. Filter Event history by Event name GetDataAccess to see every data access event. Each record includes the IAM Identity Center user UUID (userIdentity.onBehalfOf.userId), the specific table accessed (requestParameters.tableArn), and confirmation that trusted identity propagation was used (additionalEventData.LakeFormationTrustedCallerInvocation: true).

CloudTrail GetDataAccess event showing Identity Center user identity and table access details

Figure 15: CloudTrail GetDataAccess event showing Identity Center user identity and table access details

To resolve the user UUID to a human-readable name, query the Identity Store:

aws identitystore describe-user \
    --identity-store-id d-xxxxxxxxxx \
    --user-id <USER_UUID_FROM_EVENT> \
    --region us-east-1

This audit capability provides the detailed access logs required for HIPAA, GDPR, and FDA compliance, showing exactly which users accessed which data and when. Learn about configuring CloudTrail for Lake Formation in Logging Lake Formation API calls with CloudTrail.

Cleanup

Run cdk destroy --all to remove all stacks. Manually delete the retained S3 data buckets (datalake-raw-data-* and datalake-curated-data-*) and revoke any remaining Lake Formation permissions. For detailed cleanup steps, see the repository README.

Conclusion

In this post, we showed you how to implement scalable fine-grained access control for an enterprise lakehouse by combining AWS Lake Formation tag-based access control, IAM Identity Center, and trusted identity propagation in SageMaker Unified Studio. The four-dimension LF-Tag taxonomy, hybrid RBAC + ABAC governance model, and metadata-driven Lambda automation together create a governance architecture where new datasets automatically inherit access policies through tag inheritance, permissions scale without per-table grants, and every data access event is auditable to the individual user through CloudTrail.

To extend this solution, consider adding new business domains, implementing column-level security with LF-Tags, scaling to multi-account architectures with Lake Formation cross-account sharing, or integrating additional analytics services such as Amazon Redshift Spectrum or Amazon EMR.

Get started by deploying the CDK stacks from the accompanying repository. To learn more:


About the authors

Chintan Agrawal

Chintan Agrawal

Chintan is a Solutions Architect with over 7 years of experience, with a specialization in Analytics and Healthcare domain. He possesses a strong enthusiasm for assisting clients in discovering valuable insights from their data. Through his expertise, he constructs innovative solutions that empower businesses to arrive at informed, data-driven choices.

Chaitanya Vejendla

Chaitanya Vejendla

Chaitanya is a Senior Solutions Architect and part of Global Healthcare and Life Sciences industry division at AWS. He focuses on developing strategic plans for building an end-to-end analytical strategy for large biopharma, healthcare, and life sciences organizations. His expertise spans across data analytics, data governance, AI, ML, big data, and healthcare-related technologies.

Govern Amazon Redshift Data Warehouses Data Across Accounts using Amazon SageMaker Unified Studio

Post Syndicated from Bandana Das original https://aws.amazon.com/blogs/big-data/govern-amazon-redshift-data-across-accounts-with-sagemaker-unified-studio/

Managing data governance across multiple Amazon Redshift clusters in different AWS accounts presents significant challenges. Organizations operating multiple Amazon Redshift clusters across AWS accounts often rely on manual processes for secure data sharing, which increases operational overhead and governance requirements. In this post, we show you how to use Amazon SageMaker Unified Studio to implement cross-account data sharing in Amazon Redshift using data mesh principles. We demonstrate how to build a scalable data mesh architecture that supports secure, auditable data sharing across AWS accounts while reducing operational burden.

Amazon SageMaker Unified Studio as the backbone of our data mesh

Amazon SageMaker Unified Studio is a data and AI development service which brings together functionality and tools from existing AWS Analytics and AI and machine learning (ML) services, including Amazon EMR, AWS Glue, Amazon Athena, Amazon Redshift, Amazon Bedrock and Amazon SageMaker AI. With the service, organizations can catalog, discover, share, and govern data stored across Amazon Web Services (AWS) without relying on manual coordination between AWS accounts.

A data mesh is an architectural approach that treats data as a product, with decentralized ownership by data producers while maintaining centralized governance. This architecture separates source systems, data producers (data publishers), data consumers (data subscribers), and central governance. The solution we present is tailored for cross-AWS account usage, creating a foundation for data governance so you can share data across Amazon Redshift clusters in different AWS accounts.

Our proposed solution addresses the following common challenges that organizations face when sharing data across AWS accounts:

  • Manual, ad-hoc data sharing processes are replaced with automated, event-driven data publishing to the SageMaker Unified Studio catalog.
  • Inconsistent governance across different use cases is resolved through a consistent governance framework with proper access controls.
  • High load on producer Amazon Redshift clusters is reduced through decoupled publishing that lowers the operational burden on data producers.
  • Complex credential management is simplified using AWS Secrets Manager and AWS KMS encryption.
  • Lack of auditable data publishing is addressed with full traceability of access and permissions supported by the SageMaker Unified Studio service.

With this approach, you can help reduce the time and effort required for cross-account data sharing while maintaining security and governance standards.

Architectural overview

The architecture spans three AWS accounts, each with a distinct role in the data mesh:

Central Data Governance Account (Account A) hosts the Amazon SageMaker Unified Studio domain, which serves as the unified catalog and governance layer for data discovery, access control, and subscription management across accounts.

Data Producer (Account B) hosts the source of data and processing workflows. Raw data lands in an Amazon Simple Storage Service (Amazon S3) source bucket and is processed through AWS Glue extract, transform, and load (ETL) jobs or Amazon Redshift auto copy into the Amazon Redshift source database. Amazon Redshift credentials are securely stored in AWS Secrets Manager.

Data Consumer (Account C) hosts the target Amazon Redshift database and analytics workflows. After access is granted, consumers can query shared data and connect downstream visualization tools.

While this diagram shows a single producer and consumer for simplicity, in a real-world deployment there might be hundreds of producer and consumer accounts connecting through the central governance layer. Amazon SageMaker Unified Studio scales to support this by providing a single place for managing data products regardless of the number of participating accounts.

The data sharing workflow is driven by Amazon SageMaker Unified Studio. The data owner publishes data to the catalog, where it becomes discoverable by consumers across accounts. Consumers browse the catalog, subscribe to data products, and the data owner approves the request. After approval, Amazon SageMaker Unified Studio handles the cross-account sharing, granting the consumer access without requiring direct connectivity between producer and consumer Amazon Redshift clusters.

Publishing Amazon Redshift data assets to the data mesh

In a data mesh architecture, data producers need to make their data products discoverable and accessible across the organization. Amazon SageMaker Unified Studio provides a centralized catalog where data assets can be published for consumer subscription.

In practice, this means registering your data sources with the catalog so they can be discovered, governed, and subscribed to by consuming teams. This section walks through the steps required to register Amazon Redshift data sources with SageMaker Unified Studio.

Before you can publish data assets from your producer account, you need to complete several configuration steps across your Amazon Redshift cluster, AWS Secrets Manager, and Amazon SageMaker Unified Studio.

Prerequisites

  • Install the AWS Command Line Interface (AWS CLI) (v2.15+ recommended).
  • Obtain temporary credentials with permissions to administer each account (producer, consumer, and domain account)
  • IAM permissions required: redshift:* on the relevant clusters, secretsmanager:CreateSecret / PutResourcePolicy / TagResource, kms:CreateKey / PutKeyPolicy / TagResource, datazone:* for subscription-target creation, and iam:PassRole for the Amazon Redshift cluster role.
  • Amazon Redshift clusters must use RA3 node types (ra3.xlplus, ra3.4xlarge, or ra3.16xlarge). Data sharing is not supported on other node types.
  • Amazon SageMaker Unified Studio domain must already be created in Account A with the Tooling and LakeHouseCatalog blueprints available.
  • All resources must be in an AWS Region where Amazon SageMaker Unified Studio is available.

Step 1: Account association and blueprint enablement

To implement the data mesh architecture described in the previous section, you need to set up the following accounts and enable the required blueprints. This ensures that the central governance layer can discover and manage data assets across your producer and consumer accounts.

This post uses three separate AWS accounts to illustrate the cross-account data sharing pattern. However, Amazon SageMaker Unified Studio also supports publishing and subscribing to data within a single account or across any number of accounts depending on your organizational setup. Additionally, this walkthrough uses a provisioned Amazon Redshift cluster, but Amazon SageMaker Unified Studio also supports Amazon Redshift Serverless for both publishing and subscribing to data assets.

Step 2: Configure your Amazon Redshift cluster and credentials

  • In the producer account (Account B), the data to be shared resides in an Amazon Redshift cluster.
  • Verify that your Amazon Redshift cluster uses node types from the RA3 family.
  • Add the following tags to your Amazon Redshift cluster.

Amazon Redshift console showing tags added to the cluster

  • Create a superuser in Amazon Redshift for Amazon SageMaker Unified Studio. For the Amazon Redshift cluster, the database user you provide in AWS Secrets Manager must have superuser permissions. With superuser permission, your Amazon Redshift cluster can publish data and subscribe from the data mesh created with Amazon SageMaker Unified Studio, and it manages the subscriptions (access) on your behalf. For reference, see the note section in this QuickStart guide with sample Amazon Redshift data.

Tag key-value pairs configured on the Amazon Redshift cluster

  • Store the user’s credentials in Secrets Manager. Select the credential type, enter the credential values, and choose the AWS Key Management Service (AWS KMS) key with which to encrypt the secret

QuickStart guide note about providing superuser credentials for Amazon SageMaker Unified Studio

Tags on the AWS Secrets Manager secret including the Amazon Redshift cluster ARN

Resource policy added to the AWS Secrets Manager secret for Amazon SageMaker Unified Studio access

  • If your secret is encrypted with a customer managed AWS KMS key, append the key policy with the following statement and add a tag to the key: AmazonDataZoneEnvironment = All. You can skip this step if you’re using an AWS managed KMS key.
{
    "Sid": "AllowSMUSRolesSecretsAccess",
    "Effect": "Allow",
    "Principal": {
        "AWS": "*"
    },
    "Action": [
        "kms:Decrypt",
        "kms:DescribeKey",
        "kms:GenerateDataKey"
    ],
    "Resource": "*",
    "Condition": {
        "StringEquals": {
            "kms:ViaService": "secretsmanager.<<AWS_Region>>.amazonaws.com"
        },
        "StringLike": {
            "aws:PrincipalArn": [
                "arn:aws:iam::<<Data_Producer_Acct_Id(Account B)>>:role/aws-service-role/redshift.amazonaws.com/AWSServiceRoleForRedshift",
                "arn:aws:iam::<<Data_Producer_Acct_Id(Account B)>>:role/<<Redshift_Cluster_IAM_Role_Name>>",
                "arn:aws:iam::<<Data_Producer_Acct_Id(Account B)>>:role/datazone*",
                "arn:aws:iam::<<Data_Producer_Acct_Id(Account B)>>:role/service-role/AmazonSageMaker*"
            ]
        }
    }
},
{
    "Sid": "AllowSMUSRolesCreateGrant",
    "Effect": "Allow",
    "Principal": {
        "AWS": "*"
    },
    "Action": "kms:CreateGrant",
    "Resource": "*",
    "Condition": {
        "Bool": {
            "kms:GrantIsForAWSResource": "true"
        },
        "StringEquals": {
            "kms:ViaService": "secretsmanager.<<AWS_Region>>.amazonaws.com"
        },
        "StringLike": {
            "aws:PrincipalArn": [
                "arn:aws:iam::<<Data_Producer_Acct_Id(Account B)>>:role/aws-service-role/redshift.amazonaws.com/AWSServiceRoleForRedshift",
                "arn:aws:iam::<<Data_Producer_Acct_Id(Account B)>>:role/<<Redshift_Cluster_IAM_Role_Name>>"
            ]
        }
    }
}

Note: Enable automatic rotation. Configure Secrets Manager automatic rotation for this secret with a rotation interval appropriate to your security policy (for example, every 30 days). When implementing rotation, verify that the rotation Lambda function updates the credentials in both Secrets Manager and Amazon Redshift database users simultaneously. Note that Amazon SageMaker Unified Studio retrieves the secret at connection time, so rotation must produce credentials that are valid immediately upon storage: use the alternating-users rotation strategy if you need to avoid downtime during rotation. See the Secrets Manager rotation documentation for setup instructions.

Using Amazon Redshift Serverless?

  • Add the following Tags to the Amazon Redshift Serverless namespace and workgroup.

Tags added to the Amazon Redshift Serverless namespace and workgroup

  • In the Secrets Manager secret, verify the host points to your Serverless endpoint.

AWS Secrets Manager secret showing the host pointing to the Serverless endpoint

  • Add the following tags to the AWS Secrets Manager secret.

Tags added to the AWS Secrets Manager secret for the Redshift Serverless credentials

Publish Amazon Redshift data to the data mesh

With prerequisites complete, you can now register your Amazon Redshift cluster as a data source in Amazon SageMaker Unified Studio.

Step 1: Create an Amazon Redshift type connection

  • Sign in to Account B, navigate to your Amazon SageMaker Unified Studio associated domain, and open the Amazon SageMaker Unified Studio URL.

Amazon SageMaker Unified Studio associated domain sign-in page

Add an Amazon Redshift connection form in Amazon SageMaker Unified Studio

  • The newly created Amazon Redshift connection appears here.

Newly created Amazon Redshift connection listed in Amazon SageMaker Unified Studio

Step 2: Create the data source for your Amazon Redshift data warehouse

Add an Amazon Redshift data source form in Amazon SageMaker Unified Studio

Amazon Redshift data source configuration in Amazon SageMaker Unified Studio

  • For Publishing settings, choose whether assets are immediately discoverable in Amazon SageMaker Catalog.

Publishing settings controlling asset discoverability in the Amazon SageMaker catalog

Using Amazon Redshift Serverless?

When creating the connection and data source, use your workgroupName instead of clusterName. The rest of the data source configuration remains the same.

Step 3: Run the data source and publish the data asset to the data mesh

Data source run configuration in Amazon SageMaker Unified Studio

Data source run results in Amazon SageMaker Unified Studio

  • During creation of data source if you choose Publishing settings such as assets are immediately discoverable, the Amazon Redshift tables and views appear in the catalog as Published, ready for discovery and subscription by data consumers.

Published Amazon Redshift tables and views in the Amazon SageMaker catalog

Data discovery view in the Amazon SageMaker Unified Studio portal

Subscribe Amazon Redshift data through the data mesh

To complete the end-to-end test, you need to set up a consumer Amazon Redshift cluster in Account C.

Step 1: Setting up the consumer cluster

  • Follow the prerequisites from Steps 1 and 2 in the previous section, make sure the cluster and secret are properly tagged as in the following screenshots:
  • Amazon Redshift cluster tags:

Tags applied to the consumer Amazon Redshift cluster

  • Tags for the AWS Secrets Manager secret that stores the user credentials for the Amazon Redshift cluster:

Tags on the AWS Secrets Manager secret storing the consumer cluster credentials

Step 2: Connect the consumer cluster to the data mesh

  • Log into Amazon SageMaker Unified Studio and navigate to your consumer project.
  • In the Compute section of your project, choose Add compute, then choose Connect to existing compute resources.
  • Choose Amazon Redshift Provisioned.
  • Select your consumer Amazon Redshift cluster from the dropdown list and enter the Secrets Manager name.
  • Choose Add compute.
  • Your newly added Amazon Redshift cluster should now show as available.

Consumer Amazon Redshift cluster added as compute in Amazon SageMaker Unified Studio

  • The newly added Amazon Redshift cluster shows an Available state.

Consumer Amazon Redshift cluster showing an Available state

  • In the Data section you can see that objects (table/views) from Amazon Redshift cluster are visible and you can query them.

Data section showing Amazon Redshift tables and views available to query

Step 3: Creating a subscription target

  • Find the tooling environment ID: in your local terminal after obtaining correct credentials as a project member, run this command to find the tooling environment ID.
export REGION='<your-region>'
export SUBSCRIBER_PROJECT_ID='<your-project-id>'
export DOMAIN_ID='dzd-xxxxxxx'

aws datazone list-environments \
  --domain-identifier $DOMAIN_ID \
  --project-identifier $SUBSCRIBER_PROJECT_ID \
  --region $REGION
  • In the response, find and copy the tooling environment ID as shown in the following example.
{
    "items": [
        {
            "projectId": "<PROJECT_ID>",
            "id": "<ENVIRONMENT_ID>",
            "createdBy": "SYSTEM",
            "createdAt": "<TIMESTAMP>",
            "updatedAt": "<TIMESTAMP>",
            "name": "Tooling",
            "awsAccountId": "<AWS_ACCOUNT_ID>",
            "awsAccountRegion": "eu-west-1",
            "provider": "Amazon SageMaker",
            "status": "ACTIVE",
            "environmentConfigurationId": "<ENVIRONMENT_CONFIGURATION_ID>"
        }
    ]
}
  • Locate the Manage Access Role: In Account C, navigate to SageMaker Unified Studio and find the Tooling blueprint. In the Provisioning Tab you will find the Manage Access role and copy the value, as it is needed for the next CLI call.

Provisioning tab showing the Manage Access role for the Tooling blueprint

  • Create the Subscription Target.

With all the information collected, you can create the subscription target for the Amazon Redshift cluster as shown by the CLI call.

export TOOLING_ENV_ID='<tooling-environment-id>'
export AUTHORIZED_PRINCIPAL='datazone_env_<tooling-env-id>'
export MANAGE_ACCESS_ROLE='arn:aws:iam::<account-id>:role/service-role/AmazonSageMakerManageAccess-<domain-id>'

aws datazone create-subscription-target \
  --domain-identifier $DOMAIN_ID \
  --environment-identifier $TOOLING_ENV_ID \
  --name "RedshiftCluster-default-target" \
  --subscription-target-config '[{
    "formName": "RedshiftSubscriptionTargetConfigForm",
    "content": "{\"databaseName\":\"<db-name>\",\"secretManagerArn\":\"arn:aws:secretsmanager:<region>:<account>:secret:<secret-name>\",\"clusterIdentifier\":\"<cluster-id>\",\"schemaName\":\"<schema-name>"}"
  }]' \
  --applicable-asset-types RedshiftViewAssetType RedshiftTableAssetType \
  --manage-access-role $MANAGE_ACCESS_ROLE \
  --provider "Amazon SageMaker" \
  --type RedshiftSubscriptionTargetType \
  --authorized-principals $AUTHORIZED_PRINCIPAL

Using Amazon Redshift Serverless?

Use RedshiftServerlessSubscriptionTargetType as the --type and RedshiftServerlessSubscriptionTargetConfigForm as the formName in the subscription target config. Replace clusterIdentifier with workgroupName in the content JSON.

  • Verify the Subscription Target.

To verify that the subscription target was created successfully, make a last CLI call. You should find in the return a new subscription target with the name RedshiftCluster-default-target.

aws datazone list-subscription-targets \
  --environment-identifier $TOOLING_ENV_ID \
  --domain-identifier $DOMAIN_ID \
  --region $REGION

Step 4: Subscribing to data assets

  • Open the data catalog inside SageMaker Unified Studio and search for the assets you want to subscribe to.

Data catalog search for assets to subscribe to in Amazon SageMaker Unified Studio

Adding multiple databases and schemas

To publish assets from entirely different databases on the same Amazon Redshift cluster, you need to create a separate data source for each database, meaning repeating the steps mentioned in the section before. Each data source points to the same cluster connection but specifies a different database name. This approach gives you independent control over scheduling, publishing settings, and metadata generation for each database’s assets.

On the consumer side, each subscription target is bound to a specific database and schema combination. This is the target location where SageMaker Unified Studio will create views that give the consumer access to subscribed assets. To receive subscribed data in multiple databases or schemas, you create one subscription target per database-schema combination. For example, different teams within the consumer account might want the data materialized in their own schema. The following example shows this pattern:

# Subscription target for the sales schema
aws datazone create-subscription-target \
  --domain-identifier $DOMAIN_ID \
  --environment-identifier $TOOLING_ENV_ID \
  --name "RedshiftCluster-sales-target" \
  --subscription-target-config '[{ "formName": "RedshiftSubscriptionTargetConfigForm", "content": "{\"databaseName\":\"consumer_db\",\"secretManagerArn\":\"arn:aws:secretsmanager:<region>:<account>:secret:<secret-name>\",\"host\":\"<endpoint>\",\"port\":\"5439\",\"schemaName\":\"sales\"}" }]' \
  --applicable-asset-types RedshiftViewAssetType RedshiftTableAssetType \
  --manage-access-role $MANAGE_ACCESS_ROLE \
  --provider "Amazon SageMaker" \
  --type RedshiftSubscriptionTargetType \
  --authorized-principals $AUTHORIZED_PRINCIPAL

# Subscription target for the marketing schema
aws datazone create-subscription-target \
  --domain-identifier $DOMAIN_ID \
  --environment-identifier $TOOLING_ENV_ID \
  --name "RedshiftCluster-marketing-target" \
  --subscription-target-config '[{ "formName": "RedshiftSubscriptionTargetConfigForm", "content": "{\"databaseName\":\"consumer_db\",\"secretManagerArn\":\"arn:aws:secretsmanager:<region>:<account>:secret:<secret-name>\",\"host\":\"<endpoint>\",\"port\":\"5439\",\"schemaName\":\"marketing\"}" }]' \
  --applicable-asset-types RedshiftViewAssetType RedshiftTableAssetType \
  --manage-access-role $MANAGE_ACCESS_ROLE \
  --provider "Amazon SageMaker" \
  --type RedshiftSubscriptionTargetType \
  --authorized-principals $AUTHORIZED_PRINCIPAL

Verifying the audit trail

To substantiate the governance and traceability claims in this architecture, enable AWS CloudTrail in all three accounts with data events for Secrets Manager and KMS. Enable Amazon Redshift audit logging on clusters to capture connection and query activity through STL_CONNECTION_LOG and STL_QUERY. Subscription approvals and rejections are recorded by SageMaker Unified Studio and emitted to CloudTrail under the datazone.amazonaws.com event source. Look for CreateSubscriptionRequest, AcceptSubscriptionRequest, and RejectSubscriptionRequest events.

Clean up

If you deployed this solution for testing or evaluation purposes and no longer need the resources, we recommend cleaning up to avoid unnecessary costs. Amazon Redshift clusters, Secrets Manager secrets, and SageMaker Unified Studio projects all incur charges when left running. The following steps guide you through a structured teardown in the correct order: subscriptions first, then data assets, and finally the infrastructure itself. This order verifies that no orphaned resources remain.

  • Remove all subscriptions
  • Delete your data assets
  • Delete the projects
    • Delete the project within your SageMaker Unified Studio Domain after all subscriptions are removed. Make sure to delete both the consumer and producer projects.
  • Delete the SageMaker Unified Studio Domain in Account A.

Conclusion

In this post, we demonstrated how Amazon SageMaker Unified Studio simplifies cross-account data governance for Amazon Redshift. By implementing this solution, organizations can move away from ad-hoc, non-auditable data sharing processes to a secure, scalable, and fully governed approach. Amazon SageMaker Unified Studio serves as the central governance layer that data producers and consumers use to publish, discover, and subscribe to data products across AWS accounts. This turns a fragmented data landscape into a well-governed data mesh without the need for custom tooling or manual coordination.

With cross-account data sharing and governance in place, the natural next step is to use this well-governed data for machine learning and generative AI workloads. Because Amazon SageMaker Unified Studio brings together data, analytics, and AI capabilities in a single environment, teams can more efficiently transition from discovering and subscribing to data products to building ML models and generative AI applications, all within the same environment. This reduces the traditional friction between data engineering and data science, accelerating time to value. To get started with establishing your organization’s data mesh using Amazon SageMaker Unified Studio, follow the guidance for Setting up Amazon SageMaker Unified Studio.


About the authors

Bandana Das

Bandana Das is a senior Data Architect in Amazon Web Services and specializes in Data and Analytics. She builds event-driven data architectures to support customers in Data management and data-driven decision making. She is also passionate about enabling customers on their Data management journey to the cloud.

Sindi Cali

Sindi Cali is a ProServe Consultant with AWS Professional Services. She supports customers in building data driven applications in AWS.

Anirban Saha

Anirban Saha is a DevOps Architect at AWS, specializing in architecting and implementation of solutions for customer challenges. He is passionate about well-architected infrastructures, automation, data-driven solutions and helping make the customer’s cloud journey as smooth as possible.

Stoyan Stoyanov

Stoyan Stoyanov works for AWS as a DevOps Engineer. He has more than 10 years of experience in software engineering, cloud technologies, DevOps, data engineering, and security.

Viral Thakkar

Viral Thakkar is a Software Engineer at AWS, working on Amazon DataZone and Amazon SageMaker Unified Studio with a primary focus on distributed systems and data governance with deep expertise in building large-scale data analytics and pipelining solutions. He is passionate about tackling complex distributed systems challenges while also creating tools and automated scripts that simplify day-to-day workflows and improve productivity.

Introducing Apache Spark Connect support in AWS Glue interactive sessions

Post Syndicated from Zach Mitchell original https://aws.amazon.com/blogs/big-data/introducing-apache-spark-connect-support-in-aws-glue-interactive-sessions/

When we built AWS Glue interactive sessions, our goal was to make AWS Glue as interactive as running local Python from a notebook. We mostly succeeded. With a straightforward Python package and a Jupyter notebook, you could execute remotely against the AWS Glue ephemeral Spark backend. The Livy-based approach was ahead of its time, but it had limitations from its REST-based protocol. Running local PySpark unlocked powerful integrated development environment (IDE) features such as debugging and linting, so your environment could understand the code and help you develop Spark applications more quickly. Customers would often split their development work. They used local Spark (or Docker containers) to develop in an IDE on a small amount of data, then switched to AWS Glue interactive sessions to validate scaling and tuning against the full dataset.

With modern PySpark releases came a new protocol: Apache Spark Connect. Spark Connect bridges the gap between these two worlds: you develop in local Python, but execute on AWS Glue against actual data. Today, AWS Glue interactive sessions support Spark Connect natively. You can connect from any environment that supports the PySpark remote() API, including VS Code, PyCharm, Amazon SageMaker Unified Studio notebooks, and standalone Python applications. You don’t need to install specialized kernels or manage cluster infrastructure.

What Spark Connect changes

Spark Connect, introduced in Spark 3.4, decouples the Spark client from the server through a lightweight gRPC protocol. Instead of running your driver program on the cluster, your IDE communicates with a remote Spark server through a thin client layer. This architecture unlocks the key workflow improvement: you develop locally and execute remotely.

Spark Connect architecture diagram showing a thin client communicating with a remote Apache Spark server

Spark Connect architecture — thin client with the full power of Apache Spark

With Spark Connect support in AWS Glue interactive sessions, you get:

  • IDE freedom – Use VS Code, PyCharm, JupyterLab, or any Python environment. No kernel installation required.
  • Programmatic access – Build Spark into your Python applications and automation scripts with a standard SparkSession.builder.remote() call.
  • Serverless execution – AWS Glue provisions and manages the Spark cluster. You pay only for the data processing units (DPUs) consumed while your session is active.
  • Spark Connect monitoring – The Spark Live UI now includes a dedicated Connect tab showing active Spark Connect sessions and operations alongside the existing Jobs, Stages, and Executors views.

Getting started with SageMaker Unified Studio

Amazon SageMaker Unified Studio provides the most direct path to Spark Connect on AWS Glue. The notebook environment handles session creation, endpoint retrieval, and token refresh automatically, so no connection boilerplate is required.

Prerequisite: You need an Amazon SageMaker Unified Studio project to use this workflow. If you don’t have one, create a project in your SageMaker Unified Studio domain first.

To connect to an AWS Glue Spark Connect session:

  1. Sign in to SageMaker Unified Studio, choose your project, and create or open a Notebook.

A notebook open in SageMaker Unified Studio

A notebook open in SageMaker Unified Studio

  1. Choose the compute icon in the left toolbar to open the Compute environment panel. Expand the Spark section.

Compute environment panel in SageMaker Unified Studio with the Spark section expanded

The Compute environment panel with the Spark dropdown list

  1. Select a Glue Spark connection. Depending on your SageMaker domain configuration, you will see either default.spark or named connections such as project.spark.compatibility. Select the appropriate Glue (Spark) connection and choose Apply.

Notebook cell showing spark.version returns 3.5.6-amzn-1 after connecting to Glue Spark Connect

Connected to Glue Spark Connect — running spark.version returns ‘3.5.6-amzn-1’

After you make your selection, you’re connected. The spark session object is available natively. No imports or configuration are needed. Start running PySpark immediately:

spark.sql("SHOW DATABASES").show()

The session manages itself in the background, including automatic token refresh.

Using the sagemaker_studio SDK

The sagemaker-studio Python package extends the Spark Connect experience beyond SageMaker Unified Studio notebooks into local IDEs, continuous integration and continuous delivery (CI/CD) pipelines, and any Python environment. The sparkutils module handles session initialization and connection configuration in a single call. You get the same streamlined experience as in the notebook, anywhere you run Python:

from sagemaker_studio import sparkutils

# Initialize a Glue Spark Connect session using your project connection
spark = sparkutils.init(connection_name="default.spark")

# Run queries immediately
spark.sql("SHOW DATABASES").show()

You can also use sparkutils.get_spark_options() to retrieve pre-configured Java Database Connectivity (JDBC) options for reading and writing to data sources through your project connections. Supported sources include Amazon Redshift, Amazon Aurora, and Amazon DocumentDB (with MongoDB compatibility):

# Get connection options for a Redshift connection in your project
options = sparkutils.get_spark_options("my_redshift_connection")

# Read from Redshift via Spark Connect
df = spark.read.format("jdbc").options(**options).option("dbtable", "analytics.orders").load()
df.show()

Within SageMaker Unified Studio, the sagemaker-studio SDK is native to the environment. The spark session and sparkutils are available without installation. For local IDE use, install it with pip install sagemaker-studio and configure credentials through an AWS named profile or boto3 session.

How it works

Spark Connect sessions in AWS Glue use a three-step workflow:

  1. Create a session – Call the CreateSession API with SessionType set to SPARK_CONNECT. The session provisions in approximately 30 seconds.
  2. Retrieve the endpoint – Call GetSessionEndpoint to receive a sc:// gRPC endpoint URL and a time-limited authentication token.
  3. Connect with PySpark – Pass the endpoint and token to SparkSession.builder.remote() and start running Spark operations.

Spark Connect protocol flow from the DataFrame API to a logical plan, sent over gRPC and protobuf, with results streamed back over gRPC and Arrow

Spark Connect protocol flow — DataFrame API translated to logical plan, sent via gRPC/protobuf, results streamed back via gRPC/Arrow

Connecting with the low-level API

Some environments don’t have the sagemaker-studio SDK, such as custom containers, AWS Lambda functions, or non-Python toolchains. In these environments, or if you’re not using SageMaker Unified Studio, you can use the AWS SDK (Boto3) to manage sessions directly. The following example demonstrates the full workflow:

import time, boto3, urllib.parse
from pyspark.sql import SparkSession

glue = boto3.client("glue", region_name="us-east-1")

# 1. Create a Spark Connect session
session_id = "my-spark-connect-session"
glue.create_session(
    Id=session_id,
    Role="arn:aws:iam::123456789012:role/GlueServiceRole",
    Command={"Name": "glueetl"},
    GlueVersion="5.1",
    SessionType="SPARK_CONNECT",
    DefaultArguments={"--enable-spark-live-ui": "true"},
)

# 2. Wait for the session to reach READY
while True:
    status = glue.get_session(Id=session_id)["Session"]["Status"]
    if status == "READY":
        break
    time.sleep(5)

# 3. Get the Spark Connect endpoint
sc = glue.get_session_endpoint(SessionId=session_id)["SparkConnect"]
endpoint_url = sc["Url"]
auth_token = sc["AuthToken"]

# 4. Connect with PySpark
encoded_token = urllib.parse.quote(auth_token, safe="")
connection_string = f"{endpoint_url}:443/;use_ssl=true;x-aws-proxy-auth={encoded_token}"
spark = SparkSession.builder.remote(connection_string).getOrCreate()
spark.sql("SELECT 1 + 1 AS result").show()

Monitoring with Spark Live UI

When you enable the Spark Live UI at session creation, you gain access to a real-time dashboard showing:

  • Jobs and Stages – Track active, completed, and failed jobs with stage-level metrics.
  • Executors – Monitor memory usage, shuffle data, and executor health.
  • SQL – Inspect query plans and execution details.
  • Connect tab – View active Spark Connect sessions and operations (specific to Spark Connect).

Access the dashboard through the GetDashboardUrl API or directly from the AWS Glue console.

import boto3, webbrowser

glue = boto3.client("glue", region_name="us-east-1")
dashboard = glue.get_dashboard_url(
    ResourceId="my-spark-connect-session",
    ResourceType="SESSION",
)
webbrowser.open(dashboard["Url"])

In SageMaker Unified Studio, no API call is needed. Choose Ready in the notebook status bar to open the kernel info popover. From there, open the Spark UI link for the live dashboard or Spark Driver Logs for real-time log output.

Notebook status bar Ready button that opens the Spark UI and Spark Driver Logs links

Image showing “Ready” in the status bar to access Spark UI and Driver Logs directly from the notebook

Token refresh

Authentication tokens expire after 30 minutes. In SageMaker Unified Studio, this is handled automatically. For programmatic use, you can use a background thread to keep the connection alive. The following helper reconnects transparently before the token expires:

import threading, time, boto3, urllib.parse
from pyspark.sql import SparkSession

class GlueSparkConnect:
    """Maintains a SparkSession with automatic token refresh."""

    def __init__(self, session_id, region="us-east-1", refresh_margin=300):
        self.session_id = session_id
        self.glue = boto3.client("glue", region_name=region)
        self.refresh_margin = refresh_margin  # seconds before expiry to refresh
        self._lock = threading.Lock()
        self.spark = self._connect()
        self._start_refresh_loop()

    def _connect(self):
        sc = self.glue.get_session_endpoint(SessionId=self.session_id)["SparkConnect"]
        encoded_token = urllib.parse.quote(sc["AuthToken"], safe="")
        remote_url = f"{sc['Url']}:443/;use_ssl=true;x-aws-proxy-auth={encoded_token}"
        self._token_expiry = sc["AuthTokenExpirationTime"].timestamp()
        return SparkSession.builder.remote(remote_url).getOrCreate()

    def _start_refresh_loop(self):
        def _loop():
            while True:
                sleep_for = max(self._token_expiry - time.time() - self.refresh_margin, 30)
                time.sleep(sleep_for)
                with self._lock:
                    self.spark = self._connect()
        t = threading.Thread(target=_loop, daemon=True)
        t.start()

# Usage
session = GlueSparkConnect("my-spark-connect-session")
session.spark.sql("SELECT 1 + 1 AS result").show()

The background thread sleeps until 5 minutes before token expiry, then transparently reconnects. Because the daemon thread exits when your script ends, there is no cleanup required.

Getting started

To start using Spark Connect with AWS Glue interactive sessions:

  1. Use AWS Glue version 5.1 (Apache Spark 3.5.6).
  2. Install PySpark 3.5.6 locally: pip install pyspark==3.5.6.
  3. Grant your AWS Identity and Access Management (IAM) identity permissions for glue:CreateSession, glue:GetSession, and glue:GetSessionEndpoint.
  4. Create a session with --session-type SPARK_CONNECT and connect from your preferred environment.

VPC note: If you connect to AWS Glue interactive sessions through a virtual private cloud (VPC) endpoint, add the new Spark Connect endpoint (com.amazonaws.{region}.glue.sessions) to your VPC configuration. Existing AWS Glue VPC endpoints don’t cover Spark Connect traffic.

For detailed instructions, see Connecting to a Spark Connect session in the AWS Glue Developer Guide.


About the authors

Zach Mitchell

Zach Mitchell

Zach is a Senior Big Data Architect at AWS Worldwide Specialist Organization for Analytics. He works with customers to design and build data applications on AWS, with a focus on SageMaker Unified Studio, AWS Glue, and AWS Lake Formation. Outside of work, he enjoys building things with code and occasionally writing about it.

Shrey Malpani

Shrey Malpani

Shrey is a Senior Technical Product Manager at AWS Analytics. He is focused on building and scaling data processing, data integration, and data management capabilities across services like AWS Glue, Amazon EMR, and Amazon Redshift that help customers build AI-ready data platforms for their analytics or machine learning workflows.

Vaibhav Naik

Vaibhav Naik

Vaibhav is a Software Engineer at AWS Glue, where he leads the development of enterprise Generative AI managed services and Agentic data systems. He has over a decade of experience designing massive-scale cloud infrastructure and distributed computing platforms.

Tom Olson

Tom Olson

Tom is a Software Development Engineer on the AWS Glue team, focused on Interactive Sessions and operational excellence. He brings over 20 years of software development experience, including government contracting and EC2 Networking at AWS. Outside of work, he enjoys running and playing board games.

Gaurav Krishnan

Gaurav Krishnan

Gaurav is a Software Development Engineer at AWS Glue. He has a deep interest in distributed systems and creating low-friction developer experiences for interactive data workloads on Apache Spark. In his spare time, he enjoys running and trying new restaurants.

Modernizing financial analytics with Amazon SageMaker Unified Studio

Post Syndicated from Umang Aggarwal original https://aws.amazon.com/blogs/architecture/modernizing-financial-analytics-with-amazon-sagemaker-unified-studio/

Avanse Financial Services is one of India’s leading education loan providers. Their Data Engineering Team had built a data lake on AWS using Amazon Simple Storage Service (Amazon S3), Amazon Athena, and AWS Glue for data ingestion and processing. However, their analytics and reporting layer ran on an external analytics application that wasn’t integrated with AWS. Data had to be copied from Amazon S3 into this external application before analysts could run any report, its license consumed a significant portion of their budget despite low utilization, and every integration with AWS services required custom-built pipelines.

After evaluating their options, Avanse migrated to a cloud-native lakehouse architecture using Amazon SageMaker Unified Studio, which unified their data engineering, analytics, and artificial intelligence (AI) workflows in a single governed environment on AWS. In this post, we walk through their migration journey so you can adapt their approach to your own environment.

Why Avanse chose to modernize

The separation between their AWS data lake and their external analytics application created five problems:

  1. Daily data synchronization bottleneck. Every report required a 4-hour batch copy from Amazon S3 into the external analytics application before analysts could query it. Business decisions were based on data that was at least a day old.
  2. Fixed licensing costs disconnected from usage. The external analytics application charged an annual fee regardless of how many queries analysts ran. Avanse needed usage-based pricing that matched what they actually consumed, not a fixed fee for capacity they weren’t using.
  3. Limited auditability. The external analytics application ran on a shared server where different business units (risk, collections, portfolio management) shared the same resources. It lacked granular audit trails, making it difficult to trace who accessed what data and when, or to allocate costs per team.
  4. No centralized data discovery. Although AWS Glue Data Catalog managed schema metadata for the data lake, the external analytics application couldn’t access it. Analysts working in that application relied on folder structures and manual documentation to find the right datasets, slowing onboarding and increasing the risk of using outdated data.
  5. Disconnected from AWS services. The external analytics application couldn’t query data in Amazon S3 or use AWS Glue catalogs natively. Every data flow required connectors and custom-built pipelines, adding maintenance overhead.

Additionally, some datasets were stored on Network File System (NFS) storage outside of Amazon S3, creating another data silo that needed to be consolidated.

Avanse chose Amazon SageMaker Unified Studio because it addressed all five challenges: direct querying of data in Amazon S3 avoiding synchronization, usage-based compute through Amazon Athena and Amazon EMR Serverless, project-based isolation with per-project billing, lineage tracking with AWS IAM Identity Center, and native integration with their existing AWS services.

Solution overview

The core architectural change was moving from a two-application model to a single integrated stack:

Previous architecture
Avanse’s data ingestion and processing ran on AWS (Amazon S3, AWS Glue, Athena), but analytics and reporting ran on an external analytics application. Data had to be batch-copied from Amazon S3 into this external application daily before analysts could query it. Each system had its own access controls, and there was no shared catalog or lineage tracking between them.
New architecture
Analytics now run directly against data in Amazon S3 through Amazon SageMaker Unified Studio. There’s no data copy step. Analysts query the same data that the ingestion pipelines produce, using Athena for SQL and EMR Serverless for large-scale processing. Governance, access control, and lineage are centralized through IAM Identity Center and SageMaker Catalog.

The following diagram illustrates the target architecture. It follows a lakehouse pattern, storing data in open formats on Amazon S3 while maintaining ACID transaction support for the consistency financial regulators expect.

Three-layer lakehouse architecture for Avanse on AWS, showing the data layer with Amazon S3 and AWS Glue Data Catalog, the compute layer with Amazon SageMaker Unified Studio, AWS Glue ETL, AWS Lambda, Amazon EMR Serverless, Amazon SageMaker AI, and Amazon Bedrock, and the governance layer with AWS IAM Identity Center, SageMaker Catalog, and Amazon DataZone

The architecture has three layers:

  1. Data layer – Amazon S3 stores data in open formats (Parquet, Delta Lake) with S3 Intelligent-Tiering for automatic cost optimization. AWS Glue Data Catalog maintains schema metadata, making data discoverable across tools.
  2. Compute layer – Amazon SageMaker Unified Studio provides project-based workspaces organized by business function. Collections uses the built-in SQL Query Editor powered by Athena, Risk Reporting uses JupyterLab for interactive analysis, and MIS runs large-scale Spark jobs through Amazon EMR Serverless. AWS Glue ETL handles data transformations and AWS Lambda provides event-driven triggers for report generation. For machine learning (ML) workloads, Amazon SageMaker AI supports model training and deployment, with Amazon Bedrock available for generative AI capabilities such as enhancing risk narratives.
  3. Governance layer – IAM Identity Center provides SSO and audit logging across workspaces. SageMaker Catalog serves as the business glossary with data lineage tracking and access controls. Amazon DataZone connects components through a common metadata layer.

Migration journey

Avanse followed a five-phase approach. The timelines can be adapted to your environment, but the systematic progression from validation through production deployment is key.

Phase 1: Technical validation (72-hour workshop)

Avanse started with a focused 72-hour workshop using isolated SageMaker environments where developers could experiment without impacting production. Their team tested SQL analytics against existing Athena tables and validated that Python and PySpark could replicate their existing analytics workflows.

The team confirmed that querying data directly in Amazon S3 addressed their synchronization bottleneck entirely. The 4-hour daily data copy was no longer necessary, which validated the migration approach.

Phase 2: Data migration and storage optimization

Avanse migrated datasets from NFS storage and legacy analytics formats into Amazon S3, consolidating the data into a single location. They implemented S3 Intelligent-Tiering, which automatically moves data between access tiers based on usage patterns, optimizing costs without impacting retrieval performance.

They replaced legacy analytics connectors with native Athena workgroups within SageMaker Unified Studio, avoiding data synchronization entirely. Source data remained in Amazon S3, queryable by both Athena SQL and SageMaker notebooks, establishing a single source of truth.

Phase 3: Compute modernization

Avanse moved from a shared analytics server to project-based isolation in SageMaker Unified Studio. Each business function (Risk Reporting, Collections, MIS) received its own project with dedicated compute spaces running JupyterLab. Project-specific IAM execution roles provided access controls and cost allocation per business unit.

A single browser-based URL with multi-factor authentication (MFA) now provides access to SQL analytics using the built-in query editor, ML development in JupyterLab notebooks, and big data processing through Amazon EMR Serverless. This replaced the need for local analytics client installations.

Phase 4: Governance implementation

Avanse deployed SageMaker Catalog as their central business data catalog. Analysts now discover approved datasets through semantic search rather than navigating folder structures or relying on manual documentation. They mapped technical Athena table names to business terms. For example, analysts search for “collection efficiency” and find the relevant tables with descriptions, schemas, and lineage.

Lineage capture traces each metric in risk reports back to source tables, transformations, and intermediate datasets. Every action (notebook execution, SQL query, data access) is tied to IAM Identity Center users, creating the comprehensive audit trail their compliance team needed.

Phase 5: Use case migration

Rather than attempting a big-bang migration, Avanse moved critical workflows one at a time:

Portfolio MIS (Monthly/Fortnightly)
Previously required the daily 4-hour data copy from Amazon S3 into the external analytics application before report generation could begin. Avanse avoided the data synchronization step entirely and now generates MIS reports by querying existing Athena tables directly in Amazon S3. Because the source data was already on AWS, there was no need to involve the external application for this activity. Report generation dropped from hours to under 30 minutes.
Collection Efficiency and Bounce Calculation
Ported complex legacy analytics procedures for calculating metrics like collection efficiency and bounce rates to event-driven processing using AWS Glue ETL, AWS Lambda, and PySpark jobs for high-volume data aggregation. The serverless execution model charges only for compute time consumed.
EDW Risk Reporting
Large-scale regulatory joins of Enterprise Data Warehouse assets previously ran as legacy scheduled procedures. These now run as SQL queries in the SageMaker Unified Studio query editor, where analysts execute them on-demand or schedule them through Athena workgroups. The distributed query engine handles complex multi-table joins spanning millions of rows.
Scorecard Generation
Model building shifted from the external analytics application to SageMaker AI workflows. Data scientists use JupyterLab with Python libraries and deploy models directly to SageMaker endpoints, avoiding data movement between separate environments.

Overcoming technical challenges

One technical challenge was code migration. Avanse’s analytics code base contained years of accumulated proprietary scripts and procedures. Direct line-by-line translation was not practical. Instead, they took a pragmatic approach: basic data transformations moved to SQL in Athena, complex business logic was rewritten in PySpark for scalability, and statistical procedures were replaced with Python libraries like pandas and scikit-learn. The approach was to focus on what the code accomplishes, then implement it using cloud-native patterns.

The other technical challenge was performance validation. The team needed to confirm that querying data in Amazon S3 would deliver acceptable performance compared to the external analytics application’s in-memory processing. Queries against Parquet-formatted data in Amazon S3 using Athena delivered comparable performance for standard reporting workloads, while avoiding the 4-hour daily data synchronization step entirely. For large-scale regulatory joins spanning millions of rows, Amazon EMR Serverless provided distributed Spark processing that completed in minutes rather than the hours required in the external application.

Key outcomes

Area Result
Licensing costs Avoided external analytics application fees entirely
Storage costs Reduced through S3 Intelligent-Tiering, which automatically moves data between access tiers based on usage patterns
Report generation From over 4 hours (including data synchronization from Amazon S3 to the external analytics application) to under 30 minutes with direct Amazon S3 querying
Compliance audits From weeks of manual investigation to days with automated lineage reports
Compute costs Usage-based serverless model replaced always-on external analytics infrastructure
Collaboration Unified browser-based environment for data scientists, analysts, and engineers

“By adopting SageMaker Unified Studio, we as the Data Team eliminated legacy licensing costs, reduced storage and compute expenses with a serverless, usage-based model, and accelerated our periodic report generation. At the same time, we transformed compliance and collaboration by cutting audit timelines while unifying our teams in a single, efficient data environment.” – Komal Thakkar, AVP – Lead, Data Engineering, Avanse Financial Services

Best practices

Based on their experience, Avanse recommends:

  • Start with a workshop. Validate your specific use cases in a 72-hour technical validation before committing to full migration.
  • Migrate use cases, not code. Focus on what your analytics accomplish, then implement using cloud-native patterns rather than translating legacy scripts line by line.
  • Invest in governance early. Implement the data catalog and lineage tracking from day one.
  • Embrace project-based isolation. Organize around business functions for clear cost allocation and security boundaries.
  • Document business logic. Use migration as an opportunity to capture undocumented knowledge in the business glossary and dataset descriptions.

Conclusion

Avanse’s migration from an external analytics application to Amazon SageMaker Unified Studio consolidated their analytics stack into a single integrated environment on AWS. By querying data directly in Amazon S3 instead of copying it into the external application, they alleviated their biggest operational bottleneck. Project-based isolation replaced a shared server model, giving each business unit independent compute and clear cost visibility. And centralized governance through SageMaker Catalog and IAM Identity Center gave their compliance team the audit trails they had been missing.

The serverless, usage-based model means Avanse no longer pays for idle capacity. The lakehouse architecture supports new analytics patterns as they emerge, and native integration with AWS services, including generative AI through Amazon Bedrock, positions them to adopt new capabilities as their needs evolve.

Next steps

Start your analytics modernization journey by scheduling a 72-hour technical validation workshop. Contact your AWS account team to discuss your migration approach.

For more information, see:

Detecting fraud patterns across Snowflake and AWS using SageMaker Data Agent

Post Syndicated from Akash Gupta original https://aws.amazon.com/blogs/big-data/detecting-fraud-patterns-across-snowflake-and-aws-using-sagemaker-data-agent/

Financial services organizations increasingly run analytical workloads across multiple systems. For example, customers typically store transaction records in Snowflake for its concurrency handling during peak volumes, while they store risk scores, customer profiles, and behavioral signals on AWS. To bridge that divide, practitioners have had to stitch together manual exports, custom extract, transform, and load (ETL) code, and external business intelligence (BI) tools to query both sources, cache expensive aggregations, and visualize results.

Amazon SageMaker Data Agent now closes these gaps with three new capabilities in Amazon SageMaker Unified Studio notebooks: SQL analytics on Snowflake data sources, materialized view management, and interactive charting. Practitioners can use them together to query Snowflake alongside AWS data, pre-compute and schedule repeated aggregations, and create interactive visualizations from natural language prompts in a single notebook, without writing boilerplate code or switching tools.

In this post, we describe the challenges these capabilities address, introduce each one, and walk through a fraud analytics scenario that demonstrates them working together in an end-to-end investigation workflow.

Challenges with fraud detection

Fraud analytics teams working in SageMaker Unified Studio notebooks encounter several recurring friction points that slow their path from alert to insight:

  • Querying across AWS and third-party warehouses. Customers store transaction data in Snowflake and maintain risk scores and customer profiles on AWS. SageMaker Data Agent supported SQL generation for AWS-native engines: Amazon Athena, Amazon Redshift, Apache Spark, and DuckDB. However, it didn’t yet generate Snowflake-dialect SQL. This created a gap for customers working with data distributed across both AWS services and Snowflake. Analysts had to write Snowflake SQL manually and export results as CSV files to join with AWS data. The process consumed 1–2 hours before any actual investigation could begin.
  • Rich visualization requires coding expertise. When analysts want to plot query results, they must write Python code using packages like matplotlib, seaborn, or plotly. They must choose the right chart type, format axes, handle data transformations, and debug rendering issues. For fraud teams whose expertise is in investigation rather than data visualization code, each chart becomes a detour: either learn the package interface, ask an engineer for help, or export to an external BI tool. This slows the exploratory cycle that fraud investigations depend on, where every new angle (time-of-day patterns, category breakdowns, geographic clusters) ideally takes seconds, not minutes of code iteration.
  • Expensive repeated queries with no caching. Fraud signal queries flag transactions that exceed a customer’s historical average and compute risk-score distributions by merchant category. These queries re-scan entire tables on each execution. A team running the same aggregation every morning over millions of rows pays the full compute cost each time, with no mechanism to pre-compute results or schedule automatic refreshes. For fraud teams, this means investigations start with a 30-minute wait for queries that ran identically yesterday.

These three friction points (accessing data across platforms, visualizing it interactively, and operationalizing repeated analyses) are what the new Data Agent capabilities address together.

What’s new in Data Agent

Snowflake connectivity

SageMaker Data Agent can now connect to Snowflake data warehouses through connections registered in Amazon SageMaker Unified Studio. The agent discovers available Snowflake databases, browses schemas progressively (databases → schemas → tables → columns), and generates Snowflake-dialect SQL, including Snowflake-specific syntax like FLATTEN, VARIANT column access, and semi-structured data handling. Analysts query Snowflake tables alongside AWS data sources from a single notebook conversation, and the agent handles dialect differences automatically: Snowflake SQL for extraction, Spark SQL for Amazon Simple Storage Service (Amazon S3) Tables operations, with no manual translation required.

Materialized view management

Data Agent now creates and manages materialized views through natural language prompts. Analysts describe the aggregation they want, for example, “create a materialized view that flags transactions where risk_score is above 0.7, refreshed every 6 hours,” and the agent generates the Spark SQL DDL, including SCHEDULE REFRESH syntax. Materialized views store pre-computed results in Apache Iceberg format for fast repeated access, turning expensive full-table scans into sub-second queries. Supported operations include create, refresh, drop, describe, and scheduled refresh. When asked, Data Agent can also analyze notebook query patterns and recommend which queries would benefit from materialization.

Interactive charting

Instead of generating matplotlib code that produces static images, Data Agent now creates native interactive chart cells powered by Vega-Lite. Supported chart types include bar, line, scatter, pie, area, heatmap, and more. Charts render inline in the notebook with hover tooltips, zoom, and filtering. Analysts can reconfigure them through the sidebar or by typing inline instructions like “change this to a heatmap showing volume by hour and category.” This removes the cycle of modifying Python plotting code or exporting to an external BI tool every time the analysis needs a different view.

Detecting fraud patterns across Snowflake and AWS: a walkthrough

Solution overview

In this section, we walk through how these three capabilities work together in a realistic fraud investigation. A fraud analytics lead at a mid-size fintech processes a high volume of card transactions daily. Customers store transaction data in Snowflake and maintain customer risk profiles on AWS.

This morning, the real-time alerting system flagged an unusual spike in declined transactions from a cluster of new accounts, all purchasing high-value electronics. The analyst suspects a fraud ring using synthetic identities, fabricated customer profiles that pass initial verification but share telltale patterns like similar device fingerprints or overlapping IP ranges. The analyst has three goals:

  • Confirm the fraud ring hypothesis. Determine whether the flagged accounts share device fingerprints, IP ranges, or behavioral patterns indicating coordinated fraud.
  • Quantify the exposure. Calculate total fraudulent transaction volume and identify all affected accounts, not only the ones that triggered today’s alert.
  • Set up ongoing monitoring. Create a reusable, auto-refreshing query so the team catches the next ring faster.

The analyst wants to do all of this without leaving the SageMaker notebook, without writing boilerplate data-engineering code, and within a single morning standup cycle so the investigations team can be briefed by noon.

How Data Agent approaches this analysis

Data Agent is context-aware. It discovers your actual table names, column schemas, and data source connections through Amazon SageMaker Unified Studio rather than requiring you to specify them manually. It generates SQL in the correct dialect for each source (Snowflake SQL for Snowflake, Spark SQL for S3 Tables) and operates within your existing AWS Identity and Access Management (IAM) permissions boundaries.

You interact with Data Agent in two modes: the Agent Panel for multi-step investigations like the example walkthrough that follows, where each prompt builds on previous context, and inline interactions for quick adjustments like “change this to a heatmap” directly on a chart cell.

Prerequisites

Before starting this walkthrough, verify that you have:

  • An Amazon SageMaker Unified Studio domain with a project configured.
  • A Snowflake account with a warehouse and USAGE grants on the database and schemas you want to query.
  • A Snowflake connection registered in your SageMaker Unified Studio project.
  • An S3 Tables catalog in your project containing customer data (or equivalent AWS-hosted tables for joining with Snowflake data).
  • A notebook open in SageMaker Unified Studio with Data Agent available in the chat panel.

Step 1: Explore Snowflake transaction data

What the analyst wants: Before investigating the fraud ring, the analyst must understand what data is available in Snowflake and verify recent transactions are accessible. The schema isn’t memorized (the payments team manages these tables), so Data Agent needs to discover the structure.

In the SageMaker notebook Agent Panel, the analyst types:

“Show me a preview of transactions over $500 for the last 24 hours. I’m looking for repeated high-value purchases that might indicate synthetic identity fraud.”

What Data Agent does for you: Data Agent discovers the Snowflake connection through SageMaker Unified Studio, browses the available databases, and locates PAYMENTS_DB → CARD_TRANSACTIONS schema → transactions table. It surfaces the column structure (transaction_id, customer_id, amount, merchant_category, transaction_timestamp, device_fingerprint, ip_address) so the analyst can confirm the right data is available without writing a single DESCRIBE TABLE statement.

Data Agent then generates a Snowflake-dialect SQL query to preview the last 24 hours of high-value transactions (amount > $500), returning hundreds of results. The preview immediately reveals what was suspected: alongside legitimate high-value purchases (mortgage payments, business supplies), there are clusters of electronics purchases at similar price points from different customer_id values but the same device_fingerprint, a classic synthetic identity pattern.

Data Agent querying Snowflake transaction data and generating equivalent code in the cell

Figure 1: Data Agent querying Snowflake transaction data and generating equivalent code in the cell.

Notebook cell results showing high-value Snowflake transactions

Figure 2: Displaying results when the notebook cell runs.

Step 2: Land Snowflake data into S3 Tables and join with risk profiles

What the analyst wants: Pulling historical high-value transactions into S3 Tables makes this data available for downstream analysis, including the materialized view that will cross-reference risk profiles automatically.

“Load the last 90 days of transactions where amount is greater than 500 into S3 Tables.”

What Data Agent does for you: Data Agent queries Snowflake to extract a large volume of high-value transactions from the last 90 days, converts the result to a PySpark DataFrame, creates an Apache Iceberg table at payments.fraud_analytics.high_value_transactions, and writes all the rows. Data Agent stores the transaction data (transaction_id, customer_id, amount, merchant_category, transaction_timestamp, device_fingerprint, ip_address) as Iceberg in S3 Tables, allowing you to query it entirely on AWS.

Data Agent handles the cross-source complexity: Snowflake-dialect SQL for extraction, automatic schema inference for the Iceberg table, and PySpark for the write. The analyst didn’t write a single line of ETL code.

Prompt sent to Data Agent to land Snowflake transactions into an S3 Tables

Figure 3: Sending a prompt to land Snowflake transactions into an S3 Tables catalog.

Generated PySpark code that reads transaction data from Snowflake

Figure 4: Reading data from Snowflake using code Data Agent generated.

Generated cell creating an S3 Tables Iceberg table populated with Snowflake data

Figure 5: Data Agent creating a new cell to create an S3 Tables Iceberg table and populate it with the Snowflake data.

Step 3: Create a materialized view for ongoing fraud monitoring

What the analyst wants: The pattern is confirmed, but re-running this expensive join across two tables every morning isn’t sustainable. A pre-computed view that automatically refreshes and surfaces transactions from high-risk customers means tomorrow’s investigation starts with answers instead of queries (goal #3, ongoing monitoring).

“Create a materialized view called mv_fraud_signals that joins high_value_transactions with customer_risk_profiles, flagging transactions where risk_score is above 0.7. Refresh it every 6 hours.”

What Data Agent does for you: Data Agent browses the S3 Tables catalog to discover both tables and their schemas, generates the Spark SQL DDL with SCHEDULE REFRESH EVERY 6 HOURS, and creates an INNER JOIN on customer_id with a risk_score > 0.7 filter. The resulting materialized view contains only the high-risk subset of transactions, and subsequent queries against it return significantly faster compared to a full table scan.

Data Agent can also recommend materialized views when asked. If the analyst prompts “analyze my notebook and suggest which queries would benefit from materialized views,” Data Agent examines query patterns and suggests candidates. This is useful when a team runs the same expensive aggregations repeatedly without realizing a materialized view would help.

New cell created by Data Agent to create the mv_fraud_signals materialized view

Figure 6: Data Agent creates a new cell to create the materialized view.

Generated query against the newly created materialized view

Figure 7: Data Agent adds code to query the newly created materialized view.

Step 4: Visualize fraud patterns with interactive charting

What the analyst wants: The data is ready, but the investigations team needs a clear visual story by noon to see which merchant categories are targeted and what time of day the fraud occurs, so they can build detection rules. The team needs interactive charts that can be explored on the fly, not static matplotlib images that need regenerating every time someone asks “what about category X?”

“Show me a scatter plot of flagged transactions: amount vs risk_score, colored by merchant_category.”

What Data Agent does for you: Data Agent queries the materialized view, generates a Vega-Lite specification, and renders an interactive scatter plot directly in the notebook cell, with no matplotlib code and no BI tool export. Hovering over any point reveals the transaction details. A dense cluster immediately stands out: Electronics & Computers transactions with risk scores between 0.75–0.95, all in the $950–$1,000 range.

Generated scatter plot of flagged transactions colored by merchant category

Detail view of the scatter plot highlighting the Electronics cluster

Figures 8, 9, and 10: Data Agent creates a scatter plot showing a dense cluster of Electronics transactions in the $950–$1,000 range with risk scores between 0.75–.95.

The analyst follows up with a second prompt to explore temporal patterns:

“Change this to a heatmap showing transaction volume by hour of day and merchant category.”

What Data Agent does for you: Data Agent generates a new heatmap visualization from the same materialized view. The heatmap reveals that Business Supplies and Mortgage Payments maintain steady transaction volumes throughout the day. However, Electronics shows a distinctly uneven temporal distribution, with noticeable volume dips during early morning hours (midnight to 5 AM) and late evening. This variability, absent in legitimate purchase categories, is a signal the detection rules team can act on immediately.

Heatmap of transaction volume by hour and merchant category

Detail view of the heatmap showing off-hours dips in the Electronics row

Figures 11 and 12: Data Agent creates a heat map to show transaction volume by hour of day and merchant category, revealing uneven temporal distribution in high-risk categories.

From insight to action

This investigation, from Snowflake connection to visual evidence, streamlined a workflow that previously required significant time across multiple tools. The analyst shares the notebook link with the investigations team, who confirm a fraud ring of dozens of synthetic identities responsible for significant fraudulent purchases. The temporal pattern, uneven Electronics transaction distribution with off-hours variability, is added to the company’s real-time detection rules that same afternoon.

The materialized view continues refreshing every 6 hours. The next morning, it flags three new accounts matching the same pattern, caught within hours of their first transaction instead of days.

Why SageMaker Data Agent for fraud analytics

This walkthrough demonstrates three new capabilities working together:

  • SQL analytics on Snowflake data sources removed the CSV export and manual ETL that consumed half of the investigation time.
  • Materialized view management turned a one-time query into persistent, auto-refreshing monitoring, transforming reactive investigations into proactive detection.
  • Interactive charting kept the entire analysis in the notebook, removing the BI tool context switch and making the inline exploration that revealed the Electronics temporal anomaly possible.

For the team, the combined effect is a reduction in time-to-insight, allowing faster fraud pattern analysis. This means daily fraud pattern reviews instead of weekly, and an investigation workflow that’s reproducible. The notebook itself serves as documentation for compliance and audit purposes.

Cleanup

The walkthrough creates notebook cells, SQL queries, and materialized views in your SageMaker Unified Studio session. To remove the generated cells, delete them from your notebook or delete the notebook itself.

If you created resources specifically for this walkthrough, remove the following to avoid ongoing charges:

  • Materialized view. In the notebook Agent Panel, prompt: “Drop the materialized view mv_fraud_signals.” This removes the Iceberg table from S3 Tables and cancels the scheduled refresh. Alternatively, run the Spark SQL statement DROP MATERIALIZED VIEW payments.fraud_analytics.mv_fraud_signals directly.
  • Landed Iceberg tables. Drop any tables created during the data landing step (for example, payments.fraud_analytics.high_value_transactions) by prompting Data Agent or running DROP TABLE in a Spark SQL cell. This removes the data from S3 Tables and the underlying Amazon Simple Storage Service (Amazon S3) storage.
  • SageMaker Unified Studio domain. If you created a domain solely for this walkthrough, delete it to stop incurring charges. Refer to the SageMaker Unified Studio administration guide for deletion steps.
  • Amazon S3 storage. Verify that dropping the materialized view and Iceberg tables removed the associated S3 objects. If residual Iceberg metadata files remain in your S3 Tables bucket, delete them manually.
  • Snowflake compute. No persistent Snowflake resources are created. Queries use your existing warehouse. Review your Snowflake query history to estimate the compute credits consumed during the walkthrough.

Conclusion

In this post, we walked through three new capabilities in Amazon SageMaker Data Agent for notebooks: Snowflake connectivity, materialized views, and native interactive charting. Using a fraud analytics scenario, we demonstrated how these features work together. We connected to a Snowflake warehouse to explore transaction data, landed results into S3 Tables and joined them with AWS-hosted risk profiles, created a materialized view for ongoing fraud monitoring, and visualized patterns with interactive charts that revealed temporal anomalies in Electronics transactions linked to dozens of synthetic identities.

These capabilities are available now in Amazon SageMaker Unified Studio. To get started, open a notebook in your SageMaker Unified Studio domain and begin a conversation with Data Agent in the chat panel.

To learn more, see the following resources:


About the authors

Akash Gupta

Akash Gupta

Akash is a Software Development Engineer on the Amazon SageMaker Unified Studio team, where he builds integrated tools and agentic experiences. An alumnus of Santa Clara University, he is passionate about building scalable solutions that simplify how customers interact with their data. In his spare time, he enjoys singing and cooking.

Mukesh Sahay

Mukesh Sahay

Mukesh Sahay is a Software Development Engineer at Amazon SageMaker, focused on building the SageMaker Data Agent. The agent provides intelligent assistance for code generation, error diagnosis, and data analysis recommendations for data engineers, analysts, and scientists. His work spans agentic AI architectures that transform natural language prompts into executable code and analysis plans across diverse data sources. An alumnus of San Jose State University, Mukesh brings over a decade and a half of experience in building scalable, intelligent data systems.

Eason Ma

Eason Ma

Eason is a Software Development Engineer within SageMaker’s Agentic AI Experiences. His focus is on building agentic infrastructure and intelligent data experiences that help users seamlessly interact with their data across multiple sources. He holds a Master’s in Computer Science from the University of Illinois at Urbana-Champaign and a Bachelor’s in Computer Science from the University of Tennessee, Knoxville. A proud Vol, he brings that same volunteer energy to everything he builds.

Anagha Barve

Anagha Barve

Anagha is a Software Development Manager on the Amazon SageMaker Unified Studio team. Her team is focused on building tools and integrated experiences for the developers using Amazon SageMaker Unified Studio. In her spare time, she enjoys cooking, gardening and traveling.

Siddharth Gupta

Siddharth Gupta

Siddharth is heading Generative AI within SageMaker’s Unified Experiences. His focus is on driving agentic experiences, where AI systems act autonomously on behalf of users to accomplish complex tasks. An alumnus of the University of Illinois at Urbana-Champaign, he brings extensive experience from his roles at Yahoo, Glassdoor, and Twitch.

AI-assisted data development with Kiro and SageMaker Unified Studio

Post Syndicated from Zach Mitchell original https://aws.amazon.com/blogs/big-data/ai-assisted-data-development-with-kiro-and-sagemaker-unified-studio/

AI coding assistants are transforming software development, but data engineering presents unique challenges: governed data access, shared compute environments, and compliance controls that are designed to remain in place. How do you bring the power of agentic AI development into a governed data environment? With the AWS Toolkit for Visual Studio Code, you can connect Kiro, VS Code, or Cursor directly to Amazon SageMaker Unified Studio.

When you connect your editor to a SageMaker Unified Studio Space (a cloud-based compute environment inside your project), you get AI-assisted development with your preferred tools while your data governance, project permissions, and compute are managed by SageMaker Unified Studio. Additionally, SageMaker Unified Studio automatically generates steering files (like AGENTS.md) that provide your AI assistant with context about your project environment, so it understands your data and project configuration from the first prompt.

This post demonstrates the integration using Kiro. The same Remote Access connection works with VS Code and Cursor. The post starts by showing what you can do with this integration: using natural language to explore and analyze data in a governed environment. We then walk through the setup so you can try it yourself.

What’s new

With the AWS Toolkit, you can connect Kiro, VS Code, and Cursor to your SageMaker Space over a secure SSH tunnel. No additional extensions or SSH key management required. After the connection is established, your IDE has full access to your Space’s file system, compute, and data services.

Two capabilities make this especially powerful for data work:

  • Automatic AI steering – When connecting Kiro to SageMaker Unified Studio,  Kiro generates AGENTS.md and smus-context.md files that provide your AI assistant with context about your environment, including project configuration, environment details, and utilities for discovering your data catalog and project structure. Kiro detects these files automatically; other editors can use them as context for their own AI features.
  • MCP server support – have Kiro discover and configure itself for the Model Context Protocol servers on your remote SageMaker space ( like smus_local and aws-dataprocessing) to give your agent direct access to your AWS Glue Data Catalog, Amazon Athena queries, and SageMaker Unified Studio project metadata.

The following diagram shows how the components connect:

Architecture diagram: How the components connect

See it in action: AI-assisted development with governed data

Before walking through the setup, we explain what you can do with this integration. This walkthrough uses Kiro as the editor. With Kiro connected to a SageMaker Unified Studio Space, MCP servers configured, and steering documents in place, we can use natural language to explore data and build analytics. The AI assistant has all the context it needs to do this well.

Note: Agentic AI output is nondeterministic. The exact code, tool choices, and responses Kiro produces will vary between sessions, even with the same prompt. The following walkthrough shows one representative session. Your experience will differ in the specifics, but the patterns and capabilities demonstrated here are consistent.

Step 1: Explore the data

Start with a simple prompt:

show my databases and the tables I have access to

Even with native MCP tools available, Kiro often prefers the AWS Command Line Interface (AWS CLI) and bash to retrieve information. This is expected and typically does not affect the outcome. If you prefer MCP tools for every operation, you can add that preference to a steering document.

Kiro used the sagemaker_studio SDK to discover the catalog:

python3 -c "
from sagemaker_studio import Project
project = Project()
conn = project.connection()
catalog = conn.catalog()
print('Databases:')
for db in catalog.databases:
    print(f'  - {db.name}')
"
Databases:
  - default
  - sagemaker_sample_db

Then it drilled into the table schema:

python3 -c "
from sagemaker_studio import Project
project = Project()
conn = project.connection()
catalog = conn.catalog()
db = catalog.database('sagemaker_sample_db')
print('Tables in sagemaker_sample_db:')
for t in db.tables:
    print(f'  - {t.name}')
    print(f'    Columns: {[(c.name, c.type) for c in t.columns]}')
"
Tables in sagemaker_sample_db:
  - churn
    Columns: [('state', 'string'), ('account_length', 'int'),
              ('area_code', 'string'), ('phone', 'string'),
              ('intl_plan', 'string'), ('vmail_plan', 'string'),
              ('vmail_message', 'int'), ('day_mins', 'double'),
              ('day_calls', 'int'), ('day_charge', 'double'),
              ('eve_mins', 'double'), ('eve_calls', 'int'),
              ('eve_charge', 'double'), ('night_mins', 'double'),
              ('night_calls', 'int'), ('night_charge', 'double'),
              ('intl_mins', 'double'), ('intl_calls', 'int'),
              ('intl_charge', 'double'), ('custserv_calls', 'int'),
              ('churn', 'boolean')]

Kiro discovered the sagemaker_sample_db.churn dataset, a sample dataset that ships with SageMaker Unified Studio containing 10,000 rows and 21 columns of customer churn data (state, account length, call minutes, service calls, churn flag, and more). Notice that we did not write any of this code. We asked a question in natural language, and Kiro chose the right SDK calls, explored the catalog, and surfaced the results.

Another, more natural way to get the same answer is to ask directly. Prompting “Let us sample the churn table.” yields the same catalog paths and schema output, along with additional metrics like row count and a data sample, all from a single conversational prompt:

SageMaker Unified Studio console showing the sagemaker_sample_db.churn dataset listed in the catalog

Figure 1 — The sagemaker_sample_db.churn dataset in the catalog

Schema view showing the 21 columns of the churn table including state, account_length, call minutes, and the churn boolean

Figure 2 — Churn dataset schema with 21 columns

from sagemaker_studio import sqlutils
result = sqlutils.sql(
    'SELECT COUNT(*) AS total_rows FROM sagemaker_sample_db.churn',
    connection_name='default.sql'
)
print('=== Total Row Count ===')
print(result)
=== Total Row Count ===
   total_rows
0       10000

With the schema and row count in hand, Kiro sampled the data to round out its understanding of the dataset:

Comprehensive data sample showing 10 rows from the churn table with all 21 columns populated

Figure 3 — Comprehensive data sample after Kiro catalog exploration

Step 2: Run analytics with full context

With the data explored, ask Kiro to run a data quality evaluation:

Can we run basic statistical evaluations for data quality?

Because Kiro had already explored the catalog and sampled the data, it made smart choices about how to run the analysis. Instead of using PySpark for this 10,000-row table, Kiro used Athena using sqlutils to run the evaluation directly. It produced a thorough data quality report:

  • 10,000 rows, 21 columns, zero nulls across all columns. Clean on that front.
  • 5,000 duplicate rows (50 percent). Significant, worth investigating before modeling.
  • Outliers minimal. Most columns have less than 1 percent outlier rate by IQR.
  • Churn is nearly 50/50 split (50.04 percent False, 49.96 percent True). Unusually balanced, indicating synthetic data.
  • Clear signal in key features. Churners and non-churners show differences in day_mins (7.52 vs. 3.52), eve_mins (5.95 vs. 4.11), and vmail_message (175 vs. 278).
  • State distribution roughly uniform (~2% each), intl_plan and vmail_plan near 50/50.

The key insight here is what Kiro did not do. It did not default to PySpark because the environment supports Spark. Having explored the data first, understanding the table size, column types, and that churn is a proper Boolean (not a string), Kiro independently chose the right engine for the workload and produced correct analytics on the first pass.

Best practice: Explore first, code second

Start every AI-assisted development session with data exploration. Ask your AI assistant to discover your catalog, sample your tables, and understand the schema before asking it to build anything. This single step helps reduce a common source of errors in AI-assisted data work: the LLM making assumptions about data it has not seen.

Exploring your data gives the large language model (LLM) the context it needs to properly help with your project. It saves hallucinations and rework, results in faster development time, and reduces token costs.

Ready to try it yourself? The following sections walk through the full setup: prerequisites, connecting your editor to your SageMaker Space, configuring MCP servers, and working with notebooks.

Prerequisites

Before you begin, make sure you have the following:

  • A SageMaker Unified Studio domain and project with at least one project that has a compute environment provisioned (Tooling or ToolingLight). These should come standard with every SageMaker project except those provisioned with the SQL & Gen AI blueprints. If you need to set up SageMaker Unified Studio, see Getting started with Amazon SageMaker Unified Studio.
  • A Space with Remote Access enabled. Either a JupyterLab or Code Editor Space works. The instance must have at least 8 GiB of memory (for example, ml.t3.large or larger). The default ml.t3.medium (4 GiB) can’t enable Remote Access. You must upgrade the instance type first, then toggle Remote Access to Enabled in the Configure Space dialog.
  • A VS Code-compatible editor. Kiro, VS Code, Cursor, or another VS Code-based IDE installed on your local machine. This walkthrough uses Kiro, but the Remote Access connection has been tested with VS Code and Cursor as well.
  • AWS Toolkit v4.1.0 or later. Kiro ships with the AWS Toolkit pre-installed. For VS Code and Cursor, install the AWS Toolkit extension and verify your version is 4.1.0 or later (Cmd+Shift+X and search for “AWS Toolkit”).
  • AWS credentials. You must be authenticated in the SageMaker Unified Studio panel of the AWS Toolkit with the same identity (AWS IAM Identity Center or AWS Identity and Access Management (IAM)) that you use to access SageMaker Unified Studio in the browser.
  • Network connectivity. Your Space must have internet access (PublicInternetOnly mode, or virtual private cloud (VPC) with a NAT gateway or HTTP proxy that allows VS Code and Open VSX endpoints).

The following screenshots show the SageMaker Unified Studio portal and the Configure Space dialog. Navigate to your project, select your Space, and verify the configuration. Remote Access is disabled when the instance has less than 8 GiB of memory. Select an instance with at least 8 GiB, such as ml.t3.large, then enable Remote Access. This is a one-time configuration per Space.

SageMaker Unified Studio portal showing the Spaces list for a project

Figure 4 — SMUS project Spaces overview in the portal

Configure Space dialog with the instance type selector open and ml.t3.large highlighted

Figure 5 — Configure Space dialog showing instance type selection

Configure Space dialog with the Remote Access toggle set to Enabled on an 8 GiB instance

Figure 6 — Enabling Remote Access on a Space with 8 GiB or more

Connecting your editor to your SageMaker Space

There are two ways to connect: directly from the SageMaker Unified Studio portal, or from your local IDE using the AWS Toolkit.

Method 1: Connect from the SageMaker Unified Studio portal

To launch your IDE directly from the portal, navigate to your project’s Code Spaces page, find your Space, and choose Open in to select your editor (Kiro, VS Code, or Cursor):

Code Spaces list with the Open in menu showing options for Kiro, VS Code, and Cursor

Figure 7 — Open in Local IDE from the Code Spaces list

You can also launch from within a Space’s details page:

Space details page with the Open in menu expanded

Figure 8 — Open in Local IDE from the Space details page

Or from within the JupyterLab or Code Editor browser environment:

JupyterLab toolbar with the Open in Local IDE option visible

Figure 9 — Open in Local IDE from JupyterLab

Your browser will prompt you to allow opening the IDE. Confirm, and the editor launches with an SSH connection to your Space already established via the AWS Toolkit. No additional configuration is typically required.

Method 2: Connect from your IDE via the AWS Toolkit

  1. Open your editor on your local machine. Then, in the AWS Toolkit panel, choose Sign in. Authenticate with your IAM Identity Center or IAM credentials, the same identity you use to access SageMaker Unified Studio in the browser. The following screenshots show Kiro, but the steps are the same in VS Code and Cursor.Figure 10 — AWS Toolkit button in Kiro
    Figure 10 — AWS Toolkit button in KiroAWS Toolkit panel expanded in Kiro showing the Sign in option

    Figure 11 — AWS Toolkit panel expanded

    AWS Toolkit Sign in dialog with profile selection

    Figure 12 — AWS Toolkit Sign in dialog

  2. Choose your AWS profile. You must have a profile configured in the AWS CLI with the correct account and AWS Region set.
  3. In the Toolkit panel, browse your SageMaker Unified Studio domains and projects. Select the project that you want to work in.

Kiro AWS Toolkit panel showing SageMaker Unified Studio domains and projects in a tree view

Figure 13 — Browsing SMUS domains and projects in Kiro

Important: The credentials that you use in the AWS Toolkit must match the identity that you use in the SageMaker Unified Studio portal. The Toolkit validates that your identity has access to the Space.

AI steering: How SageMaker Unified Studio pre-seeds AI context

The real value of the feature comes from what you don’t need to do. When connected to Kiro SageMaker Unified Studio automatically generates steering files that guide your AI assistant with project context, so you can focus on building analytics rather than configuring connections. When you open a SageMaker Unified Studio project, SageMaker Unified Studio presents a prompt to create steering files: an AGENTS.md file that references a newly created smus-context.md. These files provide context about your project environment, such as project configuration, environment details, and utilities for discovering your data catalog and project structure. Kiro detects and applies these files automatically; in other editors, you can reference them as context for your AI features.

SageMaker Unified Studio popup offering to create AGENTS.md and smus-context.md steering files

Figure 14 — SMUS popup offering to create steering files

Kiro file explorer showing the generated AGENTS.md and smus-context.md files at the project root

Figure 15 — Generated AGENTS.md and smus-context.md steering files

Without these steering files, your AI assistant would need several back-and-forth prompts to discover what data you have and how to access it. With them, the assistant understands your project from the first prompt: how to discover your databases, how your environment is configured, and what tools are available. The steering files also help properly configure MCP servers, which you set up in the next section.

Exploring your project

After you’re connected, the project structure expands into Data and Compute sections in the sidebar, as it would in the SageMaker Unified Studio portal.

Kiro sidebar showing the Data and Compute sections expanded under a SageMaker Unified Studio project

Figure 16 — Project Data and Compute sections in the Kiro sidebar

You can explore your data catalog and S3 buckets directly from the sidebar:

Kiro sidebar with the data catalog tree and S3 buckets expanded under the project

Figure 17 — Exploring the data catalog and S3 buckets from the sidebar

You can also remote into a compatible Space for direct development. Hover over a Space and select the remote icon on the right:

Kiro sidebar showing the remote connection icon next to a compatible Space

Figure 18 — Remote connection icon on a compatible Space

After a moment, the Space opens in a new Kiro window:

New Kiro window opened with a remote connection to the SageMaker Unified Studio Space

Figure 19 — Space opened in a new Kiro window

You must sign in again, and then trust the authors of the files in the Space:

Trust authors dialog asking to confirm trust for files in the remote Space

Figure 20 — Trust authors dialog for the Space files

You’re now connected to your Space. The Toolkit works on the Space the way it does locally, except the resources are scoped to the project’s permissions.

Kiro window connected to a SageMaker Unified Studio Space with the AWS Toolkit panel active

Figure 21 — Connected to the SMUS Space with the Toolkit active

Setting up MCP servers

Before you can use AI-assisted development effectively, you must give Kiro access to your data services through Model Context Protocol (MCP) servers. MCP servers extend the Kiro agent with tools: the ability to query catalogs, run SQL, manage credentials, and more.

Out of the box, Kiro has no MCP servers configured:

Kiro MCP servers panel with no servers configured

Figure 22 — Kiro MCP servers panel with no servers configured

Prompt Kiro to find and configure the MCP servers that ship pre-installed on your SageMaker Space. Using the steering file context, Kiro located the servers and generated the configuration. If a server fails to connect, select the failed entry and Kiro will suggest fixes. You might need additional prompts to get the smus_spark_upgrade server (a pre-installed MCP server for managing Spark session upgrades) working correctly.

Kiro chat panel showing the agent discovering and configuring SageMaker Unified Studio MCP servers

Figure 23 — Kiro discovering and configuring SMUS MCP servers

MCP servers panel after iterating on configuration fixes, showing servers connected

Figure 24 — MCP servers after iterating on configuration fixes

For more deterministic results, you can also configure the MCP servers manually. Here is a sample configuration:

{
    "mcpServers": {
        "smus_local": {
            "command": "python3",
            "args": ["-m", "sagemaker_studio.mcp_server"],
            "env": {}
        },
        "aws-dataprocessing": {
            "command": "uvx",
            "args": ["awslabs.aws-dataprocessing-mcp-server@latest"],
            "env": {
                "AWS_REGION": "us-east-1",
                "FASTMCP_LOG_LEVEL": "ERROR"
            },
            "disabled": ["emr_*"]
        }
    }
}

Note: Your MCP configuration might vary depending on your SageMaker Unified Studio environment. Use the preceding configuration as a starting point and let your editor adjust if a server fails to connect.

Next, add the AWS Data Processing MCP server to get catalog information and Athena query capabilities. This isn’t strictly required (Kiro can use Python or AWS CLI for the same tasks), but it gives the agent native tools for catalog and query operations.

AWS Data Processing MCP server tools listed in Kiro with the Amazon EMR tool group disabled

Figure 25 — AWS Data Processing MCP server tools with Amazon EMR tools disabled

You can list the tools that each MCP server provides. Because the AWS Data Processing MCP server includes tools for many services, we recommend disabling tools that you don’t need for a given project to save model context. For this walkthrough, disable the Amazon EMR tools to focus on AWS Glue and Amazon Athena.

Exploring data with notebooks

Kiro supports Jupyter notebooks in your SageMaker Space with the same language and connection selectors that you would find in SageMaker JupyterLab or Code Editor. Open the command palette (Cmd+Shift+P) and create a new Jupyter notebook:

Kiro command palette filtered to the Create New Jupyter Notebook command

Figure 26 — Command palette to create a new Jupyter notebook

New Jupyter notebook open in Kiro showing language and connection selectors at the bottom-right of a cell

Figure 27 — New Jupyter notebook opened in Kiro with language and connection selectors in a notebook cell

As in SageMaker JupyterLab, you get language and connection selectors in the bottom right of each cell. Choose the connection selector to see your available connections:

SageMaker connection selector dropdown showing the available connections for the project

Figure 28 — SageMaker connection selector

Select PySpark to fill in the magic commands for your cell. Write your code (in this case, enter spark and press Shift+Enter) to verify the session starts:

Notebook cell prefilled with the PySpark magic command and a spark verification statement

Figure 29 — PySpark magic command and spark verification code

PySpark cell running in the Kiro notebook

Figure 30 — Running the PySpark cell

If this is your first time using Jupyter with Kiro, you’re prompted to install the Jupyter extension. After it’s installed, select the kernel from Python Environments → Base:

Jupyter kernel selection prompt in Kiro after installing the Jupyter extension

Figure 31 — Jupyter kernel selection prompt

Kernel picker showing the Python kernel selected from the Base environment

Figure 32 — Selecting the Python kernel from the Base environment

Re-run your cell. After a few moments, AWS Glue provisions a PySpark session:

AWS Glue provisioning a PySpark session in a Jupyter notebook in Kiro

Figure 33 — AWS Glue provisioning a PySpark session in a Jupyter notebook in Kiro

You see results the way you would in JupyterLab in the SageMaker Unified Studio portal:

PySpark code running in a Jupyter notebook in Kiro with output cells populated

Figure 34 — PySpark code running in a Jupyter notebook in Kiro

The notebook generate button

You will notice a Generate button underneath notebook cells. Let’s test it with a simple prompt:

looking at the above cell for reference, show me the accounts where state = california
using pyspark prefixing the cell with `%%pyspark default.spark` and sorting by
account_length

Notebook cell showing the Generate button populated with a natural language prompt

Figure 35 — Using the Generate button with a natural language prompt

Generated PySpark code populating a notebook cell after using the Generate button

Figure 36 — Generated PySpark code from the prompt

This prompt builder, like other notebook generation features, doesn’t have good context on the surrounding cells. You must be explicit about what you want because it won’t read other code or cells as input.

While the Kiro notebook generate button works for straightforward edits, for serious code generation, we recommend that you use Kiro agent mode. This mode has full project and SageMaker context, as demonstrated in the “See it in action” walkthrough earlier in this post.

What’s happening under the hood

When you connect your editor to a SageMaker Unified Studio Space, the AWS Toolkit extension establishes a secure SSH tunnel between your local IDE and your cloud-based Space.

Key details:

  • SSH tunnel. The connection is managed entirely by the AWS Toolkit (v4.1.0+) or VS Code’s built-in SSH extension. No separate Remote SSH extension is needed; the capability is built in.
  • File system access. Your editor sees the Space’s persistent storage at /home/sagemaker-user/, including shared project files and notebooks or scripts you create.
  • SageMaker Unified Studio steering context. The integration generates AGENTS.md and smus-context.md files that provide your AI assistant with context about your project environment and utilities for understanding your data. This is what makes the assistant effective from the first prompt.
  • MCP server integration. MCP servers like smus_local (for project metadata and environment utilities) and aws-dataprocessing (for AWS Glue Data Catalog and Amazon Athena) extend your editor’s AI with direct access to your data services. Your own MCP servers will be equally valuable here.
  • Credential flow. The Toolkit uses your existing AWS identity (IAM Identity Center or IAM) to authenticate to the Space. No separate SSH keys to manage. The aws_context_provider tool from the smus_local MCP server handles credential discovery for agent operations.

Best practices

To work effectively with your IDE and SageMaker Unified Studio:

  • Explore your data before building. Start every session by asking your AI assistant to discover your catalog, sample your data, and understand the schema. This single step helps reduce the most common source of errors in AI-assisted data work: the LLM making assumptions about data it has not seen. See the “See it in action” walkthrough earlier in this post for a concrete example of the difference this makes.
  • Use the SageMaker Unified Studio steering files. When prompted to create AGENTS.md and smus-context.md, accept. These files are the foundation that makes everything else work: environment context, MCP server configuration, and project understanding. Without them, your AI assistant starts from zero on every prompt. Kiro detects these automatically; in other editors, add them as context.
  • Disable unused MCP tools. The AWS Data Processing MCP server includes tools for AWS Glue, Amazon EMR, Amazon Athena, and more. Disable the services that you’re not using for a given project to save model context and reduce noise.
  • Be specific in your prompts. The more detail you give your AI (column names, query patterns you prefer, output formats), the closer the first pass will be. “Run data quality evaluation using Athena SQL” gets you better code than “check my data.”
  • Always test interactively first. Whether in notebooks or the terminal, validate code before deploying it. AI agents can iterate quickly, but catching issues in an interactive session is faster than debugging a failed AWS Glue job. Athena PySpark and the SageMaker sqlutils and sparkutils packages are great for this.
  • Stop your Space when idle. Your Space runs on compute (the same instance types as Code Editor and JupyterLab). If idle, the Space will terminate after 60 minutes and close your remote connection. Close the remote window and reconnect to continue.

Things to know

  • Notebook agent mode. For notebook-heavy analytics workflows where you want agentic AI to generate and run cells directly, SageMaker Notebooks with Data Agent in SageMaker Unified Studio is the recommended option today. Current notebook support in local editors covers editing, running, and generating code in individual cells.
  • MCP setup takes iteration. Configuring MCP servers may require iteration, especially for servers with complex authentication. Many AI-enabled editors can self-correct when a server fails. For more deterministic results, use the preceding MCP configuration JSON as a starting point rather than relying solely on auto-discovery.
  • CLI preference. AI agents often prefer the AWS CLI and bash even when MCP tools are available. This doesn’t affect outcomes, but you can steer your assistant toward MCP tools using a steering document if you prefer consistency.

Security and governance boundaries

A core benefit of this integration is that your existing security and governance controls remain enforced. Your editor connects to your SageMaker Space through a secure SSH tunnel managed by the AWS Toolkit. It does not bypass your organization’s access controls. Data access is governed by the same AWS Lake Formation permissions and IAM Identity Center authentication that apply when you work in the SageMaker Unified Studio portal directly. Your project-level permissions, database grants, and column-level security policies apply consistently whether a query originates from an AI agent, a notebook cell, or the SageMaker console. Data access is governed by the boundaries you define in your SageMaker Unified Studio domain and project configuration.

Clean up

To avoid ongoing charges from billable resources (SageMaker Space compute charges per hour, AWS Glue sessions charge per DPU-hour, Amazon Athena queries charge per TB scanned):

  1. Stop your Space – In the SageMaker Unified Studio portal, navigate to your project’s Spaces and stop the Space you used for this walkthrough.
  2. Disconnect: Close the remote connection in your editor (File → Close Remote Connection).
  3. Verify AWS Glue sessions are terminated – If you ran PySpark queries during this walkthrough, verify that the sessions are stopped. In the SageMaker Unified Studio portal, navigate to Data processing and confirm no active AWS Glue sessions remain. Sessions auto-terminate when the Space stops, but verify to avoid unexpected charges.
  4. Delete demo resources (optional) – File deletion is permanent and cannot be undone. Back up any work that you want to retain before proceeding. If you created scripts or files during this walkthrough that you no longer need, delete them from /home/sagemaker-user/. For example, delete any test notebooks, Python scripts, or generated data files. The sample sagemaker_sample_db.churn dataset is read-only and doesn’t need cleanup.

Conclusion

This post showed what happens when agentic AI meets governed data, and walked through how to set it up yourself.

Three key insights emerged from this hands-on experience:

  1. SageMaker Unified Studio steering files transform the developer experience. Your AI assistant is project-aware from the first prompt, understanding your environment and available data without manual setup.
  2. MCP servers bridge “AI that writes code” with “AI that queries your data”. The smus_local and aws-dataprocessing servers are essential for effective agentic data work.
  3. The “explore first” pattern pays immediate dividends. When your AI assistant understands your data before writing code, it makes smarter engine choices and produces correct analytics on the first pass.

This integration brings together two capabilities that are stronger together: your IDE handles the AI-assisted coding and iteration, while SageMaker Unified Studio handles data governance, access control, and compute management. You get the productivity of an agentic AI coding assistant without compromising on the controls your organization requires.

To get started, download Kiro, install VS Code or Cursor, and add the AWS Toolkit for Visual Studio Code (v4.1.0 or later). Then visit the Amazon SageMaker Unified Studio documentation and the AWS Data Processing MCP Server to set up your first Space. For related reading, see Speed up delivery of ML workloads using Code Editor in Amazon SageMaker Unified Studio.


About the authors

Zach Mitchell

Zach Mitchell

Zach is a Senior Big Data Architect in AWS Worldwide Specialist Organization for Analytics. He works with customers to design and build data applications on AWS, with a focus on SageMaker Unified Studio, AWS Glue, and AWS Lake Formation. Outside of work, he enjoys building things with code and occasionally writing about it.

Anchit Gupta

Anchit Gupta

Anchit is a Senior Product Manager on the Amazon SageMaker Unified Studio team at AWS.

Leah Wagner

Leah Wagner

Leah is a Senior Solutions Architect in AWS Worldwide Specialist Organization for Analytics.

Bhargava Varadharajan

Bhargava Varadharajan

Bhargava is a Senior Software Engineer on the Amazon SageMaker Unified Studio team at AWS.

Majisha Namath Parambath

Majisha Namath Parambath

Majisha is a Software Development Engineer on the Amazon SageMaker Unified Studio team at AWS.

Build governance dashboards for Amazon SageMaker Catalog with Amazon Quick

Post Syndicated from Steve Phillips original https://aws.amazon.com/blogs/big-data/build-governance-dashboards-for-amazon-sagemaker-catalog-with-amazon-quick/

Maintaining visibility into your data catalog’s health requires more than ad-hoc queries. Data stewards and compliance teams need automated dashboards that surface governance metrics and alert them when issues arise. These issues include undocumented assets, missing ownership, and stale metadata.

In a previous post, we showed you how to query Amazon SageMaker Catalog metadata using SQL by using the metadata export feature. This post builds on that foundation by demonstrating how to create governance dashboards with Amazon Quick.

Amazon Quick is an agentic AI-powered digital workspace that provides integrated analytics, automation, and research capabilities. With Amazon Quick Sight, a component of Amazon Quick, you can create interactive dashboards and visualizations with automatic chart suggestions and machine learning (ML) insights.

We walk through how to connect Amazon Quick Sight to your Amazon SageMaker Catalog metadata and build governance dashboards using natural language prompts.

Solution overview

This solution extends the metadata export architecture by adding a visualization layer:

  1. Amazon SageMaker Catalog exports asset metadata daily to Amazon Simple Storage Service (Amazon S3) Tables
  2. Amazon Athena queries the metadata using standard SQL
  3. Amazon Quick Sight connects to Athena for interactive dashboards
  4. Amazon Quick uses natural language to build visualizations

AWS Cloud architecture diagram showing the data flow for SageMaker Catalog metadata visualization. Amazon SageMaker Catalog exports asset metadata and daily exports to Amazon S3 Tables in a bucket named aws-sagemaker-catalog. Amazon Athena queries the S3 Tables data using SQL queries. Amazon Quick connects to Athena to provide interactive dashboards, natural language queries, and executive summaries. Arrows indicate the left-to-right data flow from SageMaker Catalog to S3 Tables to Athena, with Athena connecting down to Amazon Quick.

Figure 1 – Amazon SageMaker Catalog governance dashboard architecture

Prerequisites

Before you begin, complete the following steps from Analyzing your data catalog: Query SageMaker Catalog metadata with SQL. You must also have the following:

  • Amazon SageMaker Catalog metadata export enabled
  • Amazon Athena configured with query results S3 bucket
  • AWS Lake Formation permissions configured for AWS Identity and Access Management (IAM)-based access
  • Verified that the asset_metadata.asset table contains data

Additionally, you need:

Building a governance dashboard with Amazon Quick Sight

To visualize catalog health metrics, connect Amazon Quick Sight to your Athena metadata tables.

Configure Amazon Quick Sight permissions

  1. Grant permissions to the Amazon Quick Sight service role.

The Amazon Quick Sight service role (default name: aws-quicksight-service-role-v0) needs permissions to access Amazon S3 Tables and AWS Glue catalog:

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "s3tables:GetTableBucket",
        "s3tables:GetTable",
        "s3tables:GetTableMetadataLocation"
      ],
      "Resource": "arn:aws:s3tables:REGION:ACCOUNT_ID:bucket/aws-sagemaker-catalog/*"
    },
    {
      "Effect": "Allow",
      "Action": "glue:GetCatalog",
      "Resource": "arn:aws:glue:REGION:ACCOUNT_ID:catalog"
    }
  ]
}

Add this as an inline policy to the Amazon Quick Sight service role in the IAM console.

  1. Grant AWS Lake Formation permissions:

Both the Amazon Quick Sight service role and your Amazon Quick Sight admin user need AWS Lake Formation permissions on the S3 Tables catalog. First, find your Amazon Quick Sight admin user ARN by running this AWS Command Line Interface (AWS CLI) command:

aws quicksight list-users \
  --aws-account-id ACCOUNT_ID \
  --namespace default \
  --region us-east-1

Amazon Quick Sight users are managed in the Amazon Quick Sight home AWS Region (us-east-1).To grant permissions, use the Lake Formation console.

  1. Navigate to AWS Lake Formation in the AWS Management Console.
  2. Select Data permissions and Grant.
  3. For Principals, choose SAML users and groups.
  4. Enter your Amazon Quick Sight admin user ARN (from the preceding command).
  5. Under LF-Tags or catalog resources, choose Named Data Catalog resources.
  6. For Catalogs, choose the S3 Tables catalog: ACCOUNT_ID:s3tablescatalog/aws-sagemaker-catalog.
  7. For Databases, choose asset_metadata.
  8. Under Tables, choose asset.
  9. For Table permissions, choose Select and Describe.
  10. Select Grant.

Screenshot of AWS Lake Formation Grant permissions page showing the complete permission configuration workflow. At the top, the resource selection shows the 'asset_metadata' database and 'asset' table from the s3tablescatalog/aws-sagemaker-catalog catalog. Below that are optional sections for Views and Data filters, both unselected. The main content area displays three permission configuration sections. First, the 'Table permissions' section shows two subsections: 'Table permissions' with checkboxes for Select (checked, highlighted with orange box), Describe (checked, highlighted with orange box), Insert, Alter, Delete, Drop, and Super; and 'Grantable permissions' with the same permission options all unchecked. The Super permission includes explanatory text stating it is the union of all individual permissions and supersedes them. The Grantable permissions section explains that this allows the principal to grant any of the permissions to others and supersedes grantable permissions. At the bottom, the 'Data permissions' section displays two radio button options: 'All data access' (selected) which grants access to all data without restrictions, and 'Column-based access' which grants data access to specific columns only. An orange arrow points from the right side down to the bottom right corner where 'Cancel' and 'Grant' buttons are located, with the Grant button highlighted in orange.

Figure 2 – Grant access to Amazon SageMaker Catalog resources

  1. Repeat steps 1–9 for the Amazon Quick Sight service role, but in step 2 choose IAM users and roles instead.

When choosing the catalog in the Lake Formation console, you must choose the full S3 Tables catalog identifier (ACCOUNT_ID:s3tablescatalog/aws-sagemaker-catalog) to see the asset_metadata database.

Create an Amazon Quick Sight dataset.

Access S3 Tables data by creating a Quick Sight dataset using an Amazon Athena data source and the custom SQL option. An S3 Tables data source is also available but requires additional permissions. See Introducing new data source with S3 Tables in Amazon Quick for using S3 Tables as an Amazon Quick data source.

  1. Open Amazon Quick Sight in the AWS Management Console.
  2. Select Analyses and Create analysis.

Amazon QuickSight Analyses page showing the left navigation menu with Analyses selected under the Quick Sight section. The main content area displays a promotional banner for creating insightful and interactive visualizations with sample chart previews. Below the banner, an orange arrow points to the Create analysis button in the upper right. A table lists an existing analysis named New custom SQL analysis owned by Me and last updated a month ago.

Figure 3 – Create Amazon Quick Sight analysis

  1. Choose Create dataset and Create data source.

Amazon QuickSight Create Analysis dialog prompting the user to choose a dataset. A search field for datasets is shown at the top left. An orange arrow points to the Create dataset button in the upper right. A table below lists one available dataset named New custom SQL with a data source of New custom SQL, owned by Me, and last modified on March 5, 2026.

Figure 4 – Create dataset

  1. Select Amazon Athena as the data source and select Next.
  2. Enter a Data source name (for example, “SageMaker Catalog Metadata”) and choose Create data source.

Amazon QuickSight New Amazon Athena data source configuration dialog. The Data source name field is highlighted with an orange box and contains the value SageMaker Catalog Metadata. The Athena workgroup dropdown is set to primary. A Validate connection button and SSL is enabled label appear at the bottom left. An orange box highlights the Create data source button at the bottom right.

Figure 5 – Create data source

  1. Select Use custom SQL and enter a custom SQL query that references the S3 Tables catalog using the full three-part name.

Amazon QuickSight Choose your table dialog for the SageMaker Catalog Metadata data source. The Catalog dropdown is set to AwsDataCatalog and the Database dropdown shows a Select prompt. An instructional message explains to choose Prepare data to create a SQL query or choose Select table. An orange arrow points down to the Use custom SQL button highlighted with a blue box at the bottom center. The Select button is highlighted with an orange box at the bottom right.

Figure 6 – Use custom SQL

Amazon QuickSight Enter custom SQL query dialog. The query name field shows New custom SQL. The SQL editor contains a query reading SELECT FROM s3tablescatalog/aws-sagemaker-catalog with the query text underlined in orange. An orange box highlights the Confirm query button at the bottom right. An Edit/Preview data button appears at the bottom left.

Figure 7 – Enter custom SQL

SELECT * FROM "s3tablescatalog/aws-sagemaker-catalog".asset_metadata.asset

  1. Select Confirm query.
  2. Choose Directly query your data (SPICE import may fail with S3 Tables catalogs)
    Amazon QuickSight Finish dataset creation dialog showing the custom SQL dataset named New custom SQL with the SageMaker Catalog Metadata data source. Two radio button options are displayed: Import to SPICE for quicker analytics with 100 GB available shown in green, and Directly query your data which is selected and highlighted with an orange box. An orange box highlights the Visualize button at the bottom right. Edit/Preview data and Augment with SageMaker buttons appear at the bottom left and center.

Figure 8 – Directly query your data

  1. Choose Visualize and Create to start building your dashboard.

Create visualizations with Amazon Quick.

With Amazon Quick, you can build governance dashboards using natural language prompts. This removes the need for manual field configuration. This approach is faster and more intuitive than traditional dashboard building.The Amazon Quick Sight user must have AdminPro or AuthorPro subscription (the Build feature isn’t available for Reader users).Start building your dashboard with the following steps:

  1. Select Build in the top toolbar to open the natural language builder.

Amazon QuickSight analysis editor for New custom SQL analysis. The left Data panel shows the dataset fields including accountid, asset_created_time, asset_id, asset_name, asset_updated_time, business_description, catalog, extended_metadata, namespace, region, resource_description, resource_id, resource_name, resource_type_enum, and snapshot_time. The center Visuals panel shows the Build button highlighted with an orange box and a grid of available chart types. The right canvas area displays an empty AutoGraph placeholder with the message Add 1 or more fields to build a visual. An Add Data section with a dashed border prompts to add a dimension or measure.

Figure 9 – Amazon Quick build dashboard

  1. You will see a text box where you can describe the visualization that you want to create.

Amazon QuickSight analysis editor with the Build a visual panel open on the right side. An orange arrow points to the natural language input field where the user has typed a prompt requesting asset distribution by resource type as a pie chart, with a Build button next to it. Below the input field, a tooltip explains to describe the visual you would like to build with examples including map showing the top 5 cities by sales, MoM profit in 2026, and average revenue by quarter. The left Data panel shows dataset fields and the center Visuals panel displays available chart types.

Create each visualization using natural language. For each of the six recommended visualizations, enter the corresponding natural language prompt, select Build, then choose ADD TO ANALYSIS.

Amazon QuickSight analysis editor with the Build a visual panel open on the right side. The natural language prompt reads Show asset distribution by resource type as a pie chart with a Build button. Below, the system shows the interpretation as Unique number of Asset Id by Resource Type Enum using the New custom SQL dataset. A pie chart preview is displayed showing the distribution with a large segment labeled GlueTable. An orange arrow points to the Add to Analysis button at the bottom of the panel.

Figure 11 – Add to analysis

Visualization 1: Asset inventory by type

Show count of asset_id by resource_type_enum as a pie chart

After the pie chart is created, choose ADD TO ANALYSIS.

Visualization 2: Documentation completeness

Show count of asset_id where business_description is not null asa KPI

After the KPI is created, choose ADD TO ANALYSIS.

Visualization 3: Monthly registration trends

Show count of asset_id by asset_created_time month as a line chart

After the line chart is created, choose ADD TO ANALYSIS.

Visualization 4: Asset count by account

Show count of asset_id by account_id as a bar chart

After the bar chart is created, choose ADD TO ANALYSIS.

Visualization 5: Namespace distribution

Show count of asset_id by namespace as a treemap

After the treemap is created, choose ADD TO ANALYSIS.

Visualization 6: Resource type by namespace

Show count of asset_id by resource_type_enum and namespace as a heat map

Choose ADD TO ANALYSIS

  1. Arrange and publish your governance dashboard with the following steps:
  2. Delete any empty or unwanted visualizations by choosing the three dots menu and choosing Delete.
  3. Arrange visualizations by dragging them into your preferred layout.
  4. Resize visualizations to emphasize key metrics.
  5. Add titles to each visualization for clarity.
  6. Choose PUBLISH in the top right corner.
  7. Enter a dashboard name: “SageMaker Catalog Governance Dashboard”.
  8. Verify these options are selected:
    1. Allow executive summary.
    2. Allow sharing stories.
    3. Allow sharing scenarios.
  9. Choose Publish dashboard.

Amazon QuickSight SageMaker Catalog Governance Dashboard showing five visualizations. Top left is a pie chart titled Unique number of Asset Id by Resource Type showing all assets as GlueTable type. Top center is a key performance indicator displaying a total of 500 unique assets. Top right is a horizontal bar chart titled Unique number of Asset Id by Account Id showing five AWS account IDs with values of 109, 105, 104, 103, and 79 assets respectively. Middle left is a stacked bar chart titled Unique number of Asset Id by Resource Type Enum and Namespace showing GlueTable assets distributed across namespaces with values ranging from 33 to 52. Middle right is a treemap titled Unique number of Asset Id by Namespace with trading_analytics at 52, compliance_reporting at 51, treasury_ops at 50, market_data at 44, fraud_detection at 42, customer_analytics at 40, credit_scoring at 40, risk_management at 39, portfolio_mgmt at 37, regulatory at 37, loan_origination at 35, and payments at 33. Bottom is a line chart titled Unique number of Asset Id by Asset Created Time month showing asset creation trends from April 2025 through March 2026 with values fluctuating between approximately 30 and 50 assets per month.

Figure 12 – Amazon SageMaker Catalog governance dashboard

    1. Analyze your dashboard with natural language.

After you publish, you can ask questions about your governance data:

    1. On the dashboard, choose Analyze this dashboard in a Scenario in the top center.
    2. In the Data to Insights panel, enter natural language questions such as:
      1. “Which resource types have the lowest documentation rates?”
      2. “How many assets were registered last month compared to this month?”
      3. “What percentage of assets lack ownership information?”
    3. Choose Submit to generate AI-powered insights.

Amazon Quick analyzes your data and provides insights with supporting visualizations.

    1. Generate executive summaries

Create automated governance reports for data stewards and compliance teams:

    1. Choose the Amazon Quick logo in the top left to return to the home page
    2. Select Dashboards from the left panel
    3. Choose your “SageMaker Catalog Governance Dashboard”
    4. Choose the Create dropdown menu in the top right
    5. Select Executive Summary

Amazon Quick will automatically generate a summary with key governance insights, including Total asset counts and growth trends, Documentation completeness metrics, Ownership coverage statistics, and Classification distribution analysis.

    1. Create governance stories.

Build governance reports that combine multiple dashboards:

    1. From the Create dropdown, select Story.
    2. Enter a prompt: “Write a summary of catalog governance metrics and data quality trends”.
    3. Choose Add to select dashboards to include in the report.
    4. Choose Build (this might take a few minutes to complete).

Amazon Quick will generate a narrative report combining your visualizations with AI-generated insights. Share the reports with leadership or compliance teams.

Governance dashboards contain metadata such as ownership and classification details. Restrict access to users who need it. In the Amazon Quick Sight console, open the dashboard, choose Share, and grant access to named users or a dedicated Quick Sight group (for example, data-stewards) instead of selecting Everyone in this account. Review the dashboard’s permissions periodically and remove entries that are no longer needed.

Cleaning up

To avoid ongoing charges, clean up the resources created in this walkthrough. Delete Amazon Quick Sight resources including the dashboard, analyses, and dataset.

Conclusion

In this post, you connected Amazon Quick Sight to your Amazon SageMaker Catalog metadata export, built governance dashboards using the Amazon Quick natural language prompts. This approach gives data stewards and compliance teams visibility into catalog health through six key visualizations covering asset inventory, documentation completeness, registration trends, account distribution, classification coverage, and stale asset detection.

Together with the metadata export and SQL query capabilities covered in the Analyzing your data catalog: Query SageMaker Catalog metadata with SQL post, this solution provides a complete, low-overhead governance monitoring pipeline from raw catalog metadata to executive-ready.

To learn more about Amazon SageMaker Catalogs, see Amazon SageMaker Catalog documentation. To expand the work done with Amazon Quick, review Amazon Quick Sight documentation.


About the authors

Steve Phillips

Steve is a Principal Technical Account Manager and Analytics specialist at AWS in the North America region. Steve currently focuses on data warehouse architectural design, data lakes, data ingestion pipelines, and cloud distributed architectures.

Ramesh Singh

Ramesh is a Senior Product Manager Technical (External Services) at AWS in Seattle, Washington, currently with the Amazon SageMaker team. He is passionate about building high-performance ML/AI and analytics products that help enterprise customers achieve their critical goals using cutting-edge technology.

Pradeep Misra

Pradeep is a Principal Analytics and Applied AI Solutions Architect at AWS. He is passionate about solving customer challenges using data, analytics, and Applied AI. Outside of work, he likes exploring new places and playing badminton with his family. He also likes doing science experiments, building LEGOs, and watching anime with his daughters.

Rohith Kayathi

Rohith is a Senior Software Engineer at Amazon Web Services (AWS) working with Amazon SageMaker team. He leads business data catalog, generative AI–powered metadata curation, and lineage solutions. He is passionate about building large-scale distributed systems, solving complex problems, and setting the bar for engineering excellence for his team.

Accelerate SQL development with SageMaker Data Agent in Query Editor

Post Syndicated from Jason Ramos original https://aws.amazon.com/blogs/big-data/accelerate-sql-development-with-sagemaker-data-agent-in-query-editor/

When you develop SQL against Amazon Redshift and Amazon Athena, you spend time finding the right tables across hundreds of databases, writing complex joins and aggregations, debugging failed queries without context from previous attempts, and re-specifying filters for every new question. Amazon SageMaker Data Agent in Query Editor takes a different approach. You describe what you need in natural language, and the Data Agent generates the SQL. It references your actual tables through AWS Glue Data Catalog, proposes step-by-step plans for complex questions, retains context across your session, and offers one-click error recovery with Fix with AI. In this post, you learn how to use Data Agent in Query Editor to explore data, build multi-step analyses, recover from errors, and summarize results using a public education dataset.

Solution overview

You can go from a natural language question to executable SQL in seconds. Data Agent in Query Editor provides a conversational interface with direct access to your AWS data environment, so you spend less time on query mechanics and more time on analysis. Data Agent in Query Editor focuses specifically on SQL development against Amazon Redshift and Amazon Athena. (For Python, SQL, and PySpark across broader analytical and machine learning (ML) workloads, use Data Agent in notebooks.)

Data Agent provides four key capabilities:

  • Catalog-aware SQL generation. You don’t need to browse catalog structures or memorize schema details. Data Agent reads your table metadata directly.
  • Querybook and session context. You build on previous work. Data Agent uses context from your earlier queries and results.
  • Step-by-step planning. You review and approve a structured plan before Data Agent generates SQL.
  • Fix with AI. You recover from failed queries with one click.

Data Agent integrates with AWS Glue Data Catalog and reads your actual table names, column types, descriptions, and relationships, so generated SQL references your real tables. Each follow-up question builds on your current Query Editor session—the SQL cells in your querybook, the active connection, your selected cell, and execution results from previously run cells. For complex requests, Data Agent produces a structured plan that specifies which data to retrieve, how to aggregate it, and what filters to apply. You review and approve each step before Data Agent proceeds. When a query fails, choose Fix with AI to get a corrected query based on the error and the failed cell’s context.

Query Editor Fix with AI panel showing a corrected SQL query ready for review

[Figure 1: The Query Editor Fix with AI panel, showing a corrected SQL query ready for your review.]

Walkthrough: Education data analysis

In this section, you use Data Agent in Query Editor to analyze California schools data and identify where SAT improvement investment has the most impact. The walkthrough covers four tasks:

  • Explore available data.
  • Build a multi-step analysis plan.
  • Summarize insights from your queries.
  • Recover from a failed query.

The same workflow applies to your own data, whether you are analyzing sales figures, operational metrics, or financial records.

The California schools dataset contains SAT score results, school demographic information, and county-level data for public schools across California. The dataset includes tables that organize SAT scores by subject (reading, writing, math), school details (name, address, county, district), and enrollment figures. After you upload the data into your project database, you directly access the tables from Query Editor through your Amazon Athena or Amazon Redshift Lakehouse connection.

Prerequisites

To complete this walkthrough, you need intermediate SQL knowledge and basic familiarity with the AWS Management Console. You don’t need prior AWS Glue experience, but familiarity with data catalogs (centralized metadata repositories) helps.

You can choose one of two setup paths:

  • Quick start (5 minutes). SageMaker Unified Studio provides a sample database (sagemaker_sample_db) with pre-loaded data. To explore it, choose Data in the navigation pane, expand AwsDataCatalog, and select sagemaker_sample_db.
  • Full setup (30–45 minutes). Upload the California schools dataset into your project’s Lakehouse database. This dataset is publicly available from the California Department of Education. Download the SAT scores, school information, and county-level data files, then upload them through the SageMaker Unified Studio UI. In your project, go to Build, choose Query editor, right-click your project database in the Data explorer, and choose Create table. Drag and drop each CSV file to create the tables. SageMaker Unified Studio stores the data in the project-managed Amazon Simple Storage Service (Amazon S3) location, registers it in AWS Glue Data Catalog, and applies AWS Lake Formation governance automatically.

Running queries against Amazon Athena or Amazon Redshift might incur costs. For pricing details, refer to Amazon Athena pricing and Amazon Redshift pricing. For detailed setup instructions, refer to AWS Identity and Access Management (IAM)-based domains and projects. Before starting the walkthrough, you must have a SageMaker Unified Studio IAM-based domain with a project using the SQL analytics or All Capabilities project profile. The project automatically provisions an AWS Glue database, the required IAM role, and Athena or Redshift Lakehouse connections.

Data Explorer panel in Query Editor showing the california_schools_db and sagemaker_sample_db tables

[Figure 2: The Data Explorer panel in Query Editor, showing the california_schools_db and sagemaker_sample_db tables.]

Explore available data. To start, enter the following prompt in the Data Agent panel:

Query my SAT scores from my california_schools_db

Data Agent searches AWS Glue Data Catalog, locates the relevant tables, and generates an initial exploratory query that retrieves SAT score records. It adds a SQL cell directly to your querybook.

  • Review the generated SQL in the comparison view, which highlights the proposed code.
  • Choose Accept, Reject, or Accept and run.
  • After you run the cell, the results appear inline, giving you a view of the data (column names, score ranges, and the number of records) before you write SQL.

Data Agent returns an exploratory query for the california_schools_db tables, ready for review

[Figure 3: Data Agent returns an exploratory query for the california_schools_db tables, ready for your review.]

SQL query results appear beneath the cell after choosing Accept and run

[Figure 4: The SQL query results appear beneath the cell after you choose Accept and run.]

Build a multi-step analysis plan. With the data explored, enter a more complex analytical question:

Identify which subjects need investment to improve SAT scores in the lowest-performing counties. Include school-level details with addresses.

Data Agent proposes a step-by-step plan before generating SQL. For this request, Data Agent breaks the question into three steps:

  1. Aggregate SAT scores by county and subject to find performance patterns.
  2. Filter to counties with a sufficient number of schools and rank the lowest performers.
  3. Join school address data to produce a final detailed list.

Review the plan in the Data Agent panel and choose Run step-by-step to proceed.

Data Agent proposes a multi-step plan with Cancel plan and Run step-by-step options

[Figure 5: Data Agent proposes a multi-step plan with options to Cancel plan or Run step-by-step.]

Data Agent generates SQL for each step and adds it as a separate querybook cell. Review each cell’s SQL in the comparison view, then choose Accept and run to execute it. The results from each step are visible inline, so you can verify the intermediate output (county-level aggregations, the filtered ranking, and the final school list) before moving to the next step. When the steps are complete, your querybook contains the full analytical progression from raw scores to a detailed investment list.

Each plan step produces a separate querybook cell that can be reviewed and run independently

[Figure 6: Each plan step produces a querybook cell that you can review and run independently.]

Summarize insights from your queries. After running the analysis, enter the following prompt:

Summarize the insights from my queries

Data Agent has context on your querybook, including the SQL and the query results from each cell. It generates a natural language summary: which counties are underperforming, which subjects (reading, writing, or math) need the most attention in each county, and how many schools appear on the investment list. This summary provides a starting point for a report or presentation.

Data Agent summarizes insights from the accumulated query results in the querybook

[Figure 7: Data Agent summarizes insights from the accumulated query results in the querybook.]

Recover from a failed query. During the analysis, a generated query might produce an error, for example, referencing a column name that doesn’t match the schema or a join condition that returns unexpected results. When a cell fails, Query Editor displays the error message and a Fix with AI option.

Choose Fix with AI, and Data Agent reads the error in the context of the failed cell, then generates corrected SQL and updates the querybook cell. Run the corrected cell to verify the fix.

After choosing Fix with AI, Data Agent generates a corrected query for the failed cell

[Figure 8: After you choose Fix with AI, Data Agent is prompted to generate a corrected query for the failed cell.]

Data Agent returns corrected SQL for review

[Figure 9: Data Agent returns corrected SQL for you to review.]

Security and governance

Data Agent operates within your AWS environment and only accesses data that your IAM policies explicitly permit. Your existing IAM access controls and AWS Lake Formation permissions determine what data Data Agent can reach. To use Data Agent, your project role must have permissions to invoke specific Amazon DataZone APIs. For more information, refer to Actions, resources, and condition keys for Amazon DataZone.

Data Agent includes content filtering that prevents it from responding to off-topic requests, requests to reveal its system prompt, and requests for internal technical implementation details. Data Agent is restricted to AWS-related topics and English-language output.

Amazon SageMaker stores your natural language prompts and generated SQL in the AWS Region where you created your SageMaker Unified Studio domain. Data Agent doesn’t store your data, querybook context, or catalog metadata.

To opt out of data usage for service improvement, configure an AI services opt-out policy for Amazon DataZone in AWS Organizations. For more information, refer to Data storage in the SageMaker Data Agent, Service improvement, and AI services opt-out policies.

Clean up

The walkthrough creates querybook cells in your Query Editor session but doesn’t provision standalone infrastructure. To remove the generated SQL cells, delete them from your querybook or delete the querybook itself.

If you uploaded the California schools dataset specifically for this walkthrough, remove the following resources to avoid ongoing charges:

  • SageMaker Unified Studio domain. If you created a domain solely for this walkthrough, delete it to stop incurring charges. Refer to the SageMaker Unified Studio administration guide for deletion steps.
  • Uploaded tables. In the Data explorer, right-click each table you created and choose Delete table to remove the data from your project database and the underlying S3 storage.
  • Amazon Athena query results. Amazon Athena stores query results in an S3 output location. Delete the query result files from that bucket, or delete the bucket if you created it solely for this walkthrough.
  • Amazon CloudWatch logs. If Amazon Athena queries generated CloudWatch log groups, delete those log groups to avoid storage charges.

Conclusion

Data Agent in Query Editor brings conversational, catalog-aware SQL development to your Amazon Redshift and Amazon Athena workloads. In this post, you explored unfamiliar data, built a multi-step investment analysis, recovered from query errors, and summarized findings through natural language prompts.

Data Agent works within your existing IAM and AWS Lake Formation security controls, keeps your data within your AWS environment, and retains context across your analytical workflow so each question builds on the last.

Get started with these next steps:

  1. Run your first prompt. Open Query Editor in your SageMaker Unified Studio domain and enter Show me the top 10 tables in my catalog with the most columns. For setup, refer to the SageMaker Unified Studio getting started guide.
  2. Add descriptions to your AWS Glue Data Catalog. Table descriptions and column-level business metadata improve the quality of generated SQL. For best practices, refer to Populating the AWS Glue Data Catalog.
  3. Try a multi-step analysis. Enter Which product categories had declining revenue quarter-over-quarter, and which regions drove the decline? and review Data Agent’s plan step by step.

For more information, refer to the Amazon SageMaker Data Agent documentation, the What’s New blog post, Amazon Redshift documentation, and Amazon Athena documentation. To learn how Data Agent works in notebooks, refer to Accelerate context-aware data analysis and ML workflows with Amazon SageMaker Data Agent.


About the authors

Jason Ramos

Jason Ramos

Jason is a Front-End Engineer on the Amazon SageMaker Unified Studio team. He builds the scalable frontend experiences that power SageMaker Data Agent, bringing conversational AI capabilities to data scientists, analysts, and engineers across SageMaker Unified Studio. Outside of work, he enjoys playing piano and exploring the Bay Area food scene.

Olena Mursalova

Olena Mursalova

Olena is a Software Development Engineer on the Amazon SageMaker Unified Studio team, where she develops the SageMaker Data Agent — an intelligent assistant that turns natural language prompts into code, visualizations, and data insights for data engineers and analysts.

Jessica Cheng

Jessica Cheng

Jessica is a Front-End Engineer on the Amazon SageMaker Unified Studio team based in the Bay Area, where she builds intelligent data agent experiences. At work, she is passionate about creating accessible, easy-to-use experiences at scale. Outside of work, her passions lie in finding the best swimming hole in California.

Sanjana Sekar

Sanjana Sekar

Sanjana is a Software Development Engineer on the Amazon SageMaker Unified Studio team. She was one of the engineers who built the SageMaker Data Agent, bringing conversational AI-powered SQL generation and debugging to Query Editor. She is focused on improving data agent capabilities and the compute blueprints experience within SageMaker Unified Studio. Outside of work, she enjoys hiking and biking.

Siddharth Gupta

Siddharth Gupta

Siddharth is heading Generative AI within SageMaker’s Unified Experiences. His focus is on driving agentic experiences, where AI systems act autonomously on behalf of users to accomplish complex tasks. An alumnus of the University of Illinois at Urbana-Champaign, he brings extensive experience from his roles at Yahoo, Glassdoor, and Twitch.

Schedule notebook runs in Amazon SageMaker Unified Studio

Post Syndicated from Shivani Mehendarge original https://aws.amazon.com/blogs/big-data/schedule-notebook-runs-in-amazon-sagemaker-unified-studio/

If you build notebooks for recurring tasks such as daily customer analysis, weekly report generation, or data quality checks in Amazon SageMaker Unified Studio, you’ve likely wanted to run them automatically on a schedule. Until now, there wasn’t a native way to do this. Teams had to manage orchestration separately, even though the interactive notebook experience was already in place. Now, notebook scheduling is available, so you can configure your production workloads to run automatically with minimal manual intervention.

In this post, we walk you through the new scheduling and orchestrating capabilities for notebooks in Amazon SageMaker Unified Studio. You will learn how to:

  • Trigger on-demand background runs, such as a model re-training job, without waiting at your desk.
  • Create recurring schedules for tasks such as nightly data freshness checks or weekly business reviews.
  • Parameterize notebooks so a single template can generate reports across different AWS Regions or customer segments.
  • Orchestrate multi-notebook workflows where one notebook’s output feeds into the next. For example, an extract, transform, and load (ETL) pipeline followed by a summary dashboard refresh.
  • Debug failed runs with AI-assisted troubleshooting.

Sample use case overview

In this walkthrough, you will take on the role of a logistics analyst who monitors shipping performance across carriers. The notebook loads shipping data from the ShippingLogs.csv dataset, identifies late deliveries, and generates a performance summary. You want to run this notebook every morning without manual intervention, reuse it across different carriers, and know when something goes wrong.

You will start by running a notebook in the background and viewing the results. Next, you will create a recurring schedule for daily runs, then parameterize the notebook to generate reports for different carriers. You will also orchestrate the notebook in a multi-step workflow and debug a failed run using AI-assisted troubleshooting.

Prerequisites

Before you begin, you need:

  • An Amazon SageMaker Unified Studio project with Notebooks enabled. See Set up IAM-based domains for permission requirements.
  • A sample dataset. We use the ShippingLogs.csv dataset, which contains shipping data including estimated and actual delivery times, carriers, and origins. You can download it from the Workshop Studio (the file is named ShippingLogs.csv on the linked page).

Setting up the notebook

Start by creating a new notebook in your SageMaker Unified Studio project. If you haven’t already, upload the ShippingLogs.csv file under the Shared tab in the Files panel.

SageMaker Unified Studio Notebook Files panel showing the Shared tab with the ShippingLogs.csv dataset uploaded

In the first cell, we load and explore the dataset. To reference the file in code, select the file in the Shared tab and copy the Amazon Simple Storage Service (Amazon S3) URI shown in the file details. Alternatively, you can reference it with this code:

import pandas as pd
from sagemaker_studio import Project

# Initialize the project
proj = Project()

# Get the S3 root path
s3_root = proj.s3.root

df = pd.read_csv(s3_root + '/ShippingLogs.csv')
df.head()

The dataset contains columns including Carrier, ActualShippingDays, ExpectedShippingDays, ShippingOrigin, ShippingPriority, and OnTimeDelivery. Add a second cell to analyze shipping performance for a single carrier:

import matplotlib.pyplot as plt

carrier_data = df[df['Carrier'] == 'GlobalFreight']
# Flag late deliveries
carrier_data['is_late'] = carrier_data['ActualShippingDays'] > carrier_data['ExpectedShippingDays']
late_pct = carrier_data['is_late'].mean() * 100
# Visualize actual vs expected shipping days
plt.figure(figsize=(12, 4))
plt.hist(carrier_data['ActualShippingDays'] - carrier_data['ExpectedShippingDays'], bins=20, edgecolor='black')
plt.axvline(x=0, color='red', linestyle='--', label='On time')
plt.title(f'Shipping Delay Distribution - GlobalFreight ({late_pct:.1f}% late)')
plt.xlabel('Days Over Expected')
plt.ylabel('Number of Shipments')
plt.legend()
plt.show()

With the notebook working interactively, you’re ready to automate it.

Running a notebook asynchronously

To trigger an asynchronous run, open your notebook. In the notebook header, choose the menu on the Run all button, and then choose Run in background.

Notebook header with the Run all menu expanded, showing the Run in background option

This captures a snapshot of the notebook in its current state and starts a run on a separate dedicated compute. You can continue working on other tasks or close the browser entirely. Your interactive session isn’t affected.

You will see a notification at the bottom of your screen confirming that the run started. To check the status of your run, choose View Run in the notification. This opens a view showing every background and scheduled run with its status, duration, and a link to view the full output.

Run history view showing background and scheduled runs with status, duration, and output links

You can choose to view the run details at any point to view results as cells run. The run details include three tabs:

  • Output: The notebook in read-only mode with cell results rendered, including dataframe outputs, visualizations, and print statements.
  • Parameters: The parameter values used for this run.
  • Logs: Run logs for debugging.

Run details view showing the Output, Parameters, and Logs tabs with rendered cell output

You can also access past runs by selecting the View Runs option in the notebook header.

Notebook header with the View Runs option highlighted

Stopping an in-progress run

If you need to cancel a run, open the run, and choose Stop. The run terminates, and its status updates to reflect the cancellation.

Run detail view with the Stop button selected to terminate an in-progress run

What to know about background runs

Compute: Each background run uses its own dedicated compute, separate from your interactive session. Your interactive work isn’t interrupted.

Packages: The packages that you install through the notebook’s package manager will be available in your background runs. When you use !pip install in code cells, the asynchronous run installs those packages as well.

Local files: Background runs can’t access files stored locally in your notebook environment. Reference data from your project’s shared storage (Amazon S3) or connected data sources instead.

Startup time: Expect a few minutes of startup time while compute is provisioned and your environment is prepared.

Creating a recurring schedule

Now that you’ve confirmed asynchronous runs work correctly, you can automate the notebook on a schedule. Choose the schedule icon in the notebook header to open the schedule creation form.

Schedule creation form opened from the notebook header schedule icon

Configure the following settings:

  • Schedule name: Enter a descriptive name, such as Daily Shipping Report.
  • Schedule type: Choose Recurring for repeated runs or One-time for a single future run.
  • Frequency: Define how often the notebook runs using a rate (for example, every one day) or a cron expression. Set the time zone and the start and end dates for the schedule. For example, set the schedule to run every day at 7:00 AM UTC starting tomorrow.
  • Flexible time window (optional): The number of minutes after the scheduled start time within which the run can be invoked. For example, with a 5-minute window, the notebook runs within 5 minutes of the start time.
  • Advanced settings:
    • Compute Instance: Keep the current settings or override with a different instance type for the asynchronous run to use.
    • Timeout: Set a maximum run duration to help prevent notebooks from running indefinitely. If left blank, it defaults to 60 minutes.

Choose Create.

Configured schedule form with name, recurring type, daily frequency, and advanced settings populated

The schedule appears in the Schedules tab of the activity panel. SageMaker Unified Studio creates an Amazon EventBridge Scheduler schedule for each schedule you configure.

Schedules tab in the activity panel listing the newly created Daily Shipping Report schedule

Viewing schedule run history

To view past runs for a schedule, choose the schedule name in the Schedules activity panel. This opens the schedule details view, where you can see the list of runs triggered by that schedule, the duration of each run, and a link to open the notebook output for an individual run.

Schedule details view showing the list of past runs with status, duration, and output links

Editing and deleting schedules

To modify a schedule, choose Edit next to it in the Schedules panel. You can change the frequency, instance type, timeout, and other configuration fields. To pause or resume a schedule, choose Pause or Resume from the same menu. To remove a schedule, choose Delete from that menu. Deleting a schedule stops future runs but preserves historical run outputs in Amazon S3 for auditing purposes.

Schedules panel with the Edit, Pause, Resume, and Delete options for a schedule

Parameterizing notebooks

With parameters, you can reuse a single notebook across different inputs without duplicating code. For example, you can run the same shipping performance report for each carrier by passing a different carrier name to each run.

Defining parameters

Open the Parameters activity panel and choose Add. Set the parameter name to carrier and the default value to GlobalFreight.

Parameters activity panel with the carrier parameter and GlobalFreight default value configured

Using parameters in code

In your notebook, replace the second cell with the following code. This retrieves the carrier parameter value using the SageMaker Unified Studio Python SDK instead of the hardcoded value:

import sagemaker_studio
import matplotlib.pyplot as plt

carrier = sagemaker_studio.nbutils.parameters.get("carrier")

carrier_data = df[df['Carrier'] == carrier].copy()
carrier_data['is_late'] = carrier_data['ActualShippingDays'] > carrier_data['ExpectedShippingDays']
late_pct = carrier_data['is_late'].mean() * 100

plt.figure(figsize=(12, 4))
plt.hist(carrier_data['ActualShippingDays'] - carrier_data['ExpectedShippingDays'], bins=20, edgecolor='black')
plt.axvline(x=0, color='red', linestyle='--', label='On time')
plt.title(f'Shipping Delay Distribution - {carrier} ({late_pct:.1f}% late)')
plt.xlabel('Days Over Expected')
plt.ylabel('Number of Shipments')
plt.legend()
plt.show()

Creating schedules with different parameter values

Now create three schedules for the same notebook, each targeting a different carrier:

  • “daily-shipping-gf” with carrier = GlobalFreight.
  • “daily-shipping-mc” with carrier = MicroCarrier.
  • “daily-shipping-shipper” with carrier = Shipper.

When you view a historical run, a separate Parameters tab in the run output displays the parameter values that were active for that run.

You can also override parameter values when triggering an on-demand background run. Choose the menu on the Run all button, then choose Run with settings. You can keep the defaults or provide custom values for that run.

Orchestrating with Workflows

To combine notebooks into a multi-step pipeline, such as running a data calculation notebook before the shipping log notebook, you can use the Notebook Operator in the Workflows tool to orchestrate them.

To do this, choose the Add to workflows button under the options menu of the notebook header.

Notebook header options menu with the Add to workflows button highlighted

This takes you to the Workflows tool, adding a new Notebook Operator task with prefilled properties from your notebook. When configuring the Operator task:

  • Select the target notebook from the notebook menu.
  • Use the Parameters widget to pass notebook parameters into the run of the notebook.
  • Specify optional arguments such as the compute instance and timeout configuration for the run.

Workflows canvas with a Notebook Operator task configured with notebook, parameters, and compute settings

Workflows also supports polling for the status of a notebook run for a particular notebook using Notebook Sensor. In Workflows, you can add a new Sensor task by hovering on the edge of the existing Operator task, where a plus (+) button is displayed.

Workflows canvas showing the plus button on the edge of an Operator task for adding a Sensor

You can then search for and add the Notebook Sensor to the canvas.

Task picker dialog with Notebook Sensor selected for adding to the workflow canvas

When configuring the Sensor task, specify the notebook run ID within the text field. The Operator’s form field contains Jinja templating to retrieve the notebook run. If the Sensor is used within the same workflow as the Operator, this template can be copied to use within a Sensor to poll the notebook run. Select the target notebook from the notebook menu.

Notebook Sensor configuration panel with the notebook run ID field populated using Jinja templating

Within Workflows, you can configure notebook runs to emit outputs and use those outputs as inputs for subsequent notebook runs.

Building off of the previous shipping log notebook example, we will pass the carrier parameter from an upstream notebook’s output. Your shipping-logs-analysis notebook should be already set up.

Because the notebook depends on the carrier parameter, you can specify it in the Parameters panel.

Parameters panel for the shipping-logs-analysis Operator with the carrier parameter dependency configured

Now, define a second notebook, calculate-best-carrier, which performs a calculation to determine our best carrier to use for shipping:

import pandas as pd
from sagemaker_studio import Project

# Initialize the project
proj = Project()

# Get the S3 root path
s3_root = proj.s3.root

df = pd.read_csv(s3_root + '/ShippingLogs.csv')
df.head()

carrier_stats = df.groupby('Carrier').agg(
    total=('OrderID', 'count'),
    late=('OnTimeDelivery', lambda x: (x == 'Late').sum())
).reset_index()
carrier_stats['late_pct'] = carrier_stats['late'] / carrier_stats['total'] * 100

best = carrier_stats.sort_values('late_pct', ascending=True).iloc[0]
best_carrier = best['Carrier']

print("Late % by carrier:")
print(carrier_stats.to_string(index=False))
print(f"\nBest carrier: {best_carrier} ({best['late_pct']:.1f}% late)")

To configure the calculate-best-carrier notebook’s outputs, you can choose the Variables panel. A new selector is available at the bottom of this panel which allows you to select variables to mark as outputs.

Variables panel with the selector at the bottom for marking notebook variables as outputs

We want this notebook to emit the best_carrier variable.

Variables panel showing best_carrier marked as an output variable for the calculate-best-carrier notebook

Now, use the Add to workflows button as previously demonstrated to quickly add this notebook within a workflow. Chain a second Notebook Operator that points to our shipping-logs-analysis notebook. Because we specified a parameter dependency on carrier for this notebook, it’s available as an option in the Parameters widget menu.

Parameters widget menu of a Notebook Operator showing carrier as a configurable parameter dependency

When they’re chained, the notebook tasks detect the outputs set in upstream notebook runs. These outputs can be selected as keys within the Parameters widget of the Operator to pass into the run. This can be done recursively for an arbitrary number of Operator tasks. We can select the emitted best_carrier output from the calculate-best-carrier notebook.

Parameters widget displaying best_carrier as a selectable upstream output to pass into the next Operator

You can now choose the Save button on the top left of the visual canvas and the Run button to start the workflow. When the workflow is completed, the specified notebook outputs are available in the Task Output panel and the notebook run result can be viewed in the Notebooks tool.

Task Output panel showing the emitted notebook outputs after a successful workflow run

Notebook run result rendered in the Notebooks tool after the chained workflow completes

In a similar manner, the Notebook Sensor will also emit the notebook outputs from a particular notebook’s run which can be used within other tasks. This is useful when you want to retrieve outputs from a notebook run in another workflow.

Debugging a failed run with AI assistance

When viewing your past runs, you notice that a run from earlier today has a Failed status. Choose the failed run to open the notebook output in read-only mode.

In this example, suppose you incorrectly referred to column name ActualShippingDays as DeliveryDays. The run would fail with a KeyError: 'DeliveryDays' in the cell that computes late deliveries.

At the top of the failed run output, choose Troubleshoot with AI. Choosing the Troubleshoot with AI button lands you in the notebook with the Agent chat panel open.

Failed run output with the Troubleshoot with AI button highlighted at the top of the page

The data agent analyzes the cell outputs, identifies the cell that errored, explains the root cause, and suggests a fix. In this case, it identifies that the column DeliveryDays doesn’t exist in the dataframe and suggests updating the code reference. You can review the change, then verify the fix by choosing Run in background from the Run all menu to trigger a test run before the next scheduled run.

Note: You can also use the Data Agent to create schedules and start notebook runs using natural language, without having to navigate.

Cleaning up

To avoid incurring future charges, delete the resources that you created in this walkthrough:

  • Delete any schedules that you created from the Schedules panel in your notebook.
  • Delete test notebooks if you don’t need them.
  • Navigate to the Workflows page and delete any workflows that you created during this walkthrough.
  • Your project’s Amazon S3 storage retains historical run outputs until you manually remove them.

Conclusion

In this post, we showed how to run notebooks in the background in Amazon SageMaker Unified Studio using background runs, schedules, parameterization, workflow orchestration, and AI-assisted debugging. Using a shipping logistics dataset, we demonstrated how a single notebook can be parameterized to generate performance reports for different carriers on independent schedules, all without duplicating code or managing extensive infrastructure.

To get started, open a notebook in your SageMaker Unified Studio project, choose the menu on the Run all button in the notebook header, and choose Run in background. For more advanced use cases, explore workflows in Amazon SageMaker Unified Studio to build multi-step data pipelines, or review the Amazon SageMaker Unified Studio User Guide for additional configuration options.

Learn more:

If you have feedback or questions, reach out on AWS re:Post for Amazon SageMaker Unified Studio.


About the authors

Shivani Mehendarge

Shivani Mehendarge

Shivani is a Software Development Engineer at Amazon Web Services, where she builds scalable infrastructure that helps data teams run and automate their workloads in Amazon SageMaker Unified Studio. She is passionate about solving complex distributed systems challenges and building reliable cloud services.

Regan Perk

Regan Perk

Regan is a Senior Software Development Engineer on the Amazon SageMaker Unified Studio team. She designs, implements, and maintains features that enable customers to manage schedules and workflows in SageMaker Unified Studio.

Qazi Ashikin

Qazi Ashikin

Qazi is a Software Development Engineer at Amazon Web Services, where he works on developing features that allow customers to orchestrate workflows and schedules in SageMaker Unified Studio. He also works on AWS Glue Studio, where he builds agentic systems and maintains services that enable data analytics.

How Amazon is moving to integrate catalogs to improve data discovery with Amazon SageMaker

Post Syndicated from Pradeep Misra original https://aws.amazon.com/blogs/big-data/how-amazon-is-moving-to-integrate-catalogs-to-improve-data-discovery-with-amazon-sagemaker/

Enterprises face challenges when teams create data assets outside of central data catalogs. It adds overhead for discovery, and limits collaboration. Amazon’s Business Data Technologies (BDT) team has built an enterprise data catalog (Andes) for sharing datasets under well-defined policies. However, teams created catalog of local datasets and other non-tabular assets such as dashboards and metrics, outside Andes. This made it difficult to discover all assets in a consolidated way.

In this post, we share how Amazon.com is working to integrate catalogs by extending enterprise data catalog Andes with Amazon SageMaker.

Need for expanding catalog and governance from datasets to data assets

Without a single solution, users had to search multiple catalogs depending upon the asset type. Teams spent considerable time indexing the different catalogs and identifying the right one for their task. This slowed them down and took time away from solving the business problems.

To address these challenges, BDT team identified four critical capabilities needed:

  1. Multimodal catalog – Data consumers required the ability to blend enterprise data with local datasets and use them together for specific use cases. Teams sought to discover not only datasets, but also assets such as metrics, dashboards, and business files, to obtain a complete view of available resources. This necessitated a catalog that consolidates datasets and data assets in one location.
  2. Uniform governance and enforcement – To maintain best data protection practices and support business goals, teams need consistent enterprise-wide data governance where they request access once and the system enforces that access uniformly across all compute engines, alleviating fragmented or redundant access management. For internal systems, there was need for trusted identity propagation so user identity is preserved and used across AWS and internal systems for consistent enforcing.
  3. Multi-approval workflows – The solution supports multiple approval workflows within a single system, using Andes for dataset approvals and a custom workflow for dashboard approvals to maintain total governance and visibility across data assets.
  4. Delegated ownership – While enterprise teams retain overarching governance responsibility, business-specific data stewards required the ability to modify select attributes and apply appropriate tags to assets produced by their respective producers and consumers.

Solution: Unify datasets and data assets with Amazon SageMaker

Amazon chose to extend Andes with Amazon SageMaker to enhance the discovery experience. SageMaker offers native support for multimodal catalogs, and integrated with enterprise identity management, making it the ideal foundation for extending Andes’ governance model.

Rather than broadcasting assets across multiple domains, a single enterprise-wide domain standardizes and synchronizes data assets in one place. This domain is associated with AWS IAM Identity Center, which is connected to Amazon’s corporate identity system to maintain best data protection practices by limiting direct permissions and using corporate identity and group-based permissions.

Architecture diagram showing how Amazon SageMaker integrates with enterprise data catalog Andes and AWS IAM Identity Center

This integrated architecture directly addresses the identified challenges:

  • Single-pane asset discovery – Datasets and data assets are accessible through a single, consolidated view, avoiding the need to navigate across disparate systems or domains. This simplifies discovery and reduces the time to insight for teams across the organization.
  • Extended governance – Governance of both enterprise-wide and local datasets is orchestrated through a single system.
  • Extended observability – Trusted Identity Propagation (TIP) through AWS IAM Identity Center allows human users to access data interactively using their corporate identities. This provides audit-trail visibility into who is accessing what data for audits and organization’s observability requirements.
  • Amazon tool integration – Integration with Git and other internal systems automates management of accounts, permissions, and approvals. This reduces manual overhead and helps maintain that access controls remain tightly aligned with existing business workflows.

Design overview

This section describes the key features and design of the Amazon SageMaker integration. The technical implementation consists of three core components:

1) Catalog connectors

Amazon built connectors and ingestion paths to bring data assets into Amazon SageMaker while maintaining business continuity and preserving existing governance:

  • Andes integration: SageMaker provides APIs to synchronize assets from external catalogs. BDT extended this to bring Andes datasets (with their sophisticated metadata, business context) into the integrated experience. The integration preserves Andes’ permission model and governance workflows, to maintain existing security standards and best practices intact.
  • Account onboarding: Teams self-serve onboard their AWS accounts through an AWS Lambda-based integration. When creating projects, SageMaker queries this service to determine which accounts a user’s identity can access.

2) Delegated ownership

When data systems scale across business units, centralized governance teams need to delegate permissions for catalog enrichment, policy enforcement, and metadata management.

  • Catalog enhancement allows business teams to define and publish their own business glossaries, curated vocabularies of domain-specific terms, definitions, and relationships, directly within the catalog. Allowing business owners to author and maintain these glossaries increased accuracy and discoverability of catalog assets. Data consumers across the enterprise benefit from clearer, more consistent terminology.

3) Integration with consumption and access tooling

Teams discover data in SageMaker Unified Studio and consume it through both SageMaker Unified Studio and internal tooling:

  • Data discovery: SageMaker Unified Studio integrates with Amazon-wide Identity Center allowing almost all Amazon users to authenticate and search for cataloged assets. This integration addresses the data discovery problem by providing enterprise-wide visibility into available data resources.
  • Integrated development environment: SageMaker Unified Studio provides built-in tooling out of the box including a Query Editor for SQL analytics and Amazon SageMaker AI for machine learning (ML), which helps teams access data, build models, and collaborate across organizational boundaries.
  • Code repository integration: Manage code with full Git operations supported from SageMaker Unified Studio. Query code and notebook code persist to GitFarm (Amazon’s internal Git system), allowing teams to view and manage their work through Amazon’s standard version control system.
  • Native analytics integration: Projects directly connect to AWS analytics engines including Amazon Athena for SQL, AWS Glue and Amazon EMR for Apache Spark, and Amazon Redshift for data warehousing. User-authored jobs use Andes governance and permissions across engines for consistent access control.

SageMaker implementation results

SageMaker catalog now encompasses various types of data assets from across the organization, representing an expansion from datasets alone to a complete inventory of data, dashboards, metrics, models, and other data assets, all while maintaining best practices and appropriate access and use guardrails.

“SageMaker provides a unified catalog that makes discovery and sharing of data assets, metrics and dashboards across teams straightforward, with direct integration to Andes datasets. SageMaker delivers deep integration through Git repository connections and enterprise identity management that aligns with existing Amazon workflows.”

– Gerry Moses, Sr. Principal TPM, Amazon

  • Faster data discovery – Data consumers can go to one place to locate trusted, high-quality assets with significantly less friction, which reduces the time from question to insight. By surfacing well-documented, governed assets through an enriched catalog, teams can confidently identify the right data for their use cases without navigating sprawling, inconsistent inventories or relying on tribal knowledge.
  • Improved collaboration – Breaks down data silos by making curated assets discoverable and reusable across Amazon. When teams can build on shared, authoritative datasets rather than creating redundant copies, data proliferation is reduced.

Conclusion

By integrating their existing governance tooling with Amazon SageMaker to build a centralized data catalog, BDT is creating a foundation for faster, more efficient data discovery across teams. Amazon SageMaker helped unify diverse data types with their existing catalog and enabled collaboration across teams to help them find the right data. By integrating with existing governance frameworks, BDT demonstrates how organizations can expand their catalog capabilities while preserving existing enterprise investments.

To learn more and get started with Amazon SageMaker Unified Studio, visit aws.amazon.com/sagemaker/unified-studio or the AWS console.


About the authors

Matt David

Matt David

Matt is a Sr PMM, specializing in helping data teams with AI-powered analytics. His areas of interest include self-service analytics, data democratization, and preparing organizations for the age of AI agents. He brings extensive experience from his roles at Atlassian, Hex, and DataCamp.

Gerry Moses

Gerry Moses

Gerry is a Senior Principal Technical Program Manager in Business Data Technologies where he leads joint Amazon/AWS programs. His work improved data governance for Amazon’s Andes data lake, enabled broader AWS technology adoption by data lake users, and influenced product improvements that benefited all AWS customers.

Ramesh Singh

Ramesh Singh

Ramesh is a Senior Product Manager Technical at AWS in Seattle, Washington, currently with the Amazon SageMaker team. He is passionate about building high-performance ML/AI and analytics products that help enterprise customers achieve their critical goals using cutting-edge technology.

Pradeep Misra

Pradeep Misra

Pradeep is a Principal Analytics and Applied AI leader at AWS. He is passionate about solving customer challenges using data, analytics, and AI/ML. Outside of work, he likes exploring new places, trying new cuisines, and playing badminton with his family. He also likes doing science experiments, building LEGOs, and watching movies with his daughters.

Eunji Kang

Eunji Kang

Eunji is a Principal Product Manager Technical focusing on democratizing data across Amazon teams for fast data-driven business decisions without compromising security and compliance.

Trevor Gasdaska

Trevor Gasdaska

Trevor is a Principal Engineer focusing on data compliance and agentic AI workflows for Big Data Technologies at Amazon. He builds tools that help teams govern and use data at scale.

Brad Porter

Brad Porter

Brad is a Principal Business Development Manager at Amazon Web Services. He works with Amazon.com and enterprise customers to define and accelerate go-to-market strategies across Data Analytics, AI/ML, and Generative AI. He has over 20 years of experience in cloud strategy, enterprise infrastructure, and technology leadership.

Automate deployment of data and AI applications with Amazon SageMaker Unified Studio CI/CD CLI

Post Syndicated from Saurabh Bhutyani original https://aws.amazon.com/blogs/big-data/automate-deployment-of-data-and-ai-applications-with-amazon-sagemaker-unified-studio-ci-cd-cli/

Organizations building data and AI applications in Amazon SageMaker Unified Studio combine multiple AWS services, including AWS Glue, Amazon Athena, Amazon Managed Workflows for Apache Airflow (Amazon MWAA), Amazon SageMaker AI, and Amazon Quick Sight, into single applications. Promoting these applications from development to test and production stages requires substituting service-specific configurations for each stage and provisioning resources in the correct order.

Data teams understand which services their applications need but lack continuous integration and continuous delivery (CI/CD) expertise, while DevOps teams understand deployment automation but must learn each AWS service’s provisioning requirements.

The CI/CD CLI for Amazon SageMaker Unified Studio (aws-smus-cicd-cli) is an open source command line tool that automates deployment of multi-service data and AI applications across pipeline stages. Data teams define their application once in a YAML manifest, DevOps teams deploy with a single command, and the CLI handles configuration substitution, dependency ordering, and resource provisioning automatically. For details, see the CI/CD CLI documentation.

In this post, we walk through how the CI/CD CLI works, show you how to deploy a real application across environments, and demonstrate how it fits into your existing CI/CD workflows.

Customer spotlight

Bureau Veritas, a global leader in testing, inspection, and certification, operates across multiple SageMaker Unified Studio environments to support its data and AI teams. With their data and DevOps teams working on different parts of the application lifecycle, Bureau Veritas needed a controlled way to promote workloads from development through test to production while preserving clear ownership boundaries between the two teams.

“We need to promote data and AI applications across SageMaker Unified Studio environments in a controlled way that respects the boundaries between our data teams and our DevOps teams. The CI/CD CLI does exactly that — a single manifest from the data team, a single deploy command from DevOps, and full control over what goes to production.”

— Gilles Kempf, Architecture Manager, Bureau Veritas

How the CI/CD CLI works

The CI/CD CLI introduces a clean separation of concerns between data teams and DevOps teams.

Data teams define what to deploy in a declarative YAML manifest (manifest.yaml). The manifest describes the application’s resources, including AWS Glue extract, transform, and load (ETL) jobs, Athena queries, Airflow directed acyclic graphs (DAGs), Quick Sight dashboards, and SageMaker training jobs, along with stage-specific configurations for each environment.

DevOps teams define how and when to deploy using their existing CI/CD systems. They retain full control over their deployment methodology. They choose whether to promote content through git branches, a bundle artifactory, or both; they decide the shape of the pipeline, including which stages to include (dev, staging, pre-prod, prod) and which manual approvals or security gates are required. They run aws-smus-cicd-cli deploy inside GitHub Actions, Jenkins, or GitLab CI workflows without needing to understand which AWS services the application uses or how SageMaker Unified Studio projects are structured. The CLI is a utility for AWS analytics service deployment, not a CI/CD methodology. Your team’s existing conventions for branches, approvals, and pipeline shape stay exactly as they are.

The CLI is the abstraction layer between the two. It reads the manifest, substitutes stage-specific configurations (S3 paths, AWS Identity and Access Management (IAM) roles, account IDs, and connection strings), provisions resources in dependency order, and handles all AWS service interactions.The following diagram illustrates this separation:

SageMaker CI/CD

Key concepts

Application manifest

Each stage maps to a dedicated SageMaker Unified Studio project. This one-stage-to-one-project mapping is the foundation of CI/CD isolation: each project has its own domain, IAM boundaries, connections, and data, so changes in dev can never affect prod. For stronger isolation, projects can span different AWS accounts and AWS Regions. For example, dev in a sandbox account and prod in a production account in a different Region. Because each stage is a real SageMaker Unified Studio project, teams can open it in the console at any time to observe workflows, inspect resources, and troubleshoot deployments. Project membership is managed per project, so you control exactly who has access to each stage. For example, developers in dev and a release team in prod.The manifest file is the single source of truth for your application. It declares:

  • Content: application code from git repositories, data files from S3, Quick Sight dashboards, and workflow definitions.
  • Stages: environment-specific project mappings (dev, test, prod, etc.), each isolated as described earlier.
  • Configuration: stage-specific settings that are substituted automatically at deploy time.

Here is an example manifest for an analytics application with AWS Glue ETL and Quick Sight:
applicationName: SalesAnalyticsDashboard

content: 
  storage: 
    - name: etl-code 
      include: ["*.py"] 
    - name: workflows 
      include: ["*.yaml"] 
  quicksight: 
    - name: SalesDashboard 
      type: dashboard 
  workflows: 
    - workflowName: sales_etl_pipeline 
      connectionName: default.workflow_serverless 
 
stages: 
  dev: 
    domain: 
      region: us-east-1 
    project: 
      name: analytics-dev 
    deployment_configuration: 
      storage: 
        - name: etl-code 
          connectionName: default.s3_shared 
          targetDirectory: sales/bundle/etl 
        - name: workflows 
          connectionName: default.s3_shared 
          targetDirectory: sales/bundle/workflows 
 
  prod: 
    domain: 
      region: us-west-2 
    project: 
      name: analytics-prod 
    deployment_configuration: 
      storage: 
        - name: etl-code 
          connectionName: default.s3_shared 
          targetDirectory: sales/bundle/etl 
        - name: workflows 
          connectionName: default.s3_shared 
          targetDirectory: sales/bundle/workflows 
      quicksight: 
        assets: 
          - name: SalesDashboard 
            owners: 
              - arn:aws:quicksight:${AWS_REGION}:${AWS_ACCOUNT_ID}:user/default/Admin/* 

Each stage must map to a separate SageMaker Unified Studio project, providing full isolation between environments. The CLI substitutes variables like ${AWS_ACCOUNT_ID} and ${AWS_REGION} at deploy time based on the target environment.

Bundles

A bundle is an immutable, versioned archive of your application. The bundle command reads from a source stage (typically dev) and packages the application code, workflow definitions, and resolved configurations into a self-contained artifact. The deploy command then applies that artifact to one or more target stages (test or prod).

This stage-to-bundle-to-stage promotion model supports controlled rollout through quality gates:

# Package from dev 
aws-smus-cicd-cli bundle --manifest manifest.yaml 
 
# Deploy to test 
aws-smus-cicd-cli deploy --manifest app.tar.gz --targets test 
 
# Validate the test deployment 
aws-smus-cicd-cli test --manifest manifest.yaml --targets test 
 
# Promote the same bundle to prod 
aws-smus-cicd-cli deploy --manifest app.tar.gz --targets prod 

The same artifact is deployed at every stage without rebuilding, providing audit trails and reproducible deployments for regulated industries.

SageMaker Catalog integration

The CLI manages Amazon SageMaker Catalog resources as part of the deployment process. You can define catalog assets, glossaries, glossary terms, form types, asset types, and metadata forms, in your manifest. During deployment, the CLI searches for assets in the catalog, creates subscription requests for required data access, and waits for approval before proceeding. This automates the data governance workflow that teams previously handled manually.

CLI commands

The CI/CD CLI provides commands that cover the full deployment lifecycle:

Command Description
describe Validates the manifest, checks that target projects exist, and confirms the execution role has required permissions. Use –connect to validate against live AWS environments.
bundle Reads from a source stage and packages application code, workflow definitions, and configurations into an immutable, versioned archive.
deploy Applies bundle contents to one or more target stages. Provisions resources in dependency order.
test Runs post-deployment validation to confirm services are running and ready for workloads.
create Generates a starter manifest from an existing SageMaker Unified Studio project.
run Triggers Airflow workflow execution on MWAA or Airflow Serverless connections.
monitor Monitors workflow execution status in real time.
logs Fetches and streams workflow execution logs.
destroy Removes deployed resources and projects for cleanup or failure recovery.

Walkthrough: deploying a Quick Sight dashboard with AWS Glue ETL

In this section, we walk through deploying an analytics application that uses AWS Glue for ETL, Athena for queries, and Quick Sight for dashboards. This example is available in the GitHub repository.

Use case

An analytics team owns a Sales Analytics Dashboard built on AWS Glue ETL, Athena, and Quick Sight. They want to promote changes from a development environment to production with reproducible builds, automated validation, and a clear approval gate between stages, without writing custom deployment scripts or exposing data engineers to AWS provisioning details.

Solution overview

We use a sample application from the CI/CD CLI GitHub repository that includes AWS Glue ETL scripts, an Airflow workflow definition, a Quick Sight dashboard bundle, and integration tests. A single manifest.yaml describes the application and its dev and prod stages. The CLI handles the full lifecycle: bundle the app from dev, deploy it to test, run validation, and promote the same immutable artifact to prod.

Prerequisites

Before you begin, make sure you have the following:

Solution architecture

Each stage in the manifest maps to a dedicated SageMaker Unified Studio project (see the separation-of-concerns diagram in “How the CI/CD CLI works” earlier in this post). At deploy time, the CLI uploads ETL scripts and workflow definitions to the project’s S3 storage connection, provisions the Airflow workflow in MWAA Serverless, runs the workflow to create AWS Glue jobs and databases, and imports the Quick Sight dashboard. The same bundle artifact is applied to every downstream stage, ensuring dev, test, and prod stay in sync while remaining fully isolated.

Solution implementation

Step 1: Install the CLI

Install the CLI from PyPI:

pip install aws-smus-cicd-cli

Step 2: Create or customize a manifest

Clone the repository and start from the analytics example:

git clone https://github.com/aws/CICD-for-SageMakerUnifiedStudio.gitcd CICD-for-SageMakerUnifiedStudio/examples/analytic-workflow/dashboard-glue-quick

The example includes AWS Glue ETL scripts, an Airflow workflow definition, a Quick Sight dashboard bundle, and integration tests. Open manifest.yaml and update the project, domain, and deployment_configuration values under each stage so they match your own SageMaker Unified Studio projects and connection names.Alternatively, generate a manifest from an existing project: aws-smus-cicd-cli create --domain-id <your-domain-id> --dev-project-id <your-project-id>

Step 3: Validate your configuration

Run the describe command with --connect to verify your environment is ready. This connects to your AWS environment and validates that target projects exist, the execution role has the required permissions, and connections are reachable. Fix any issues before deploying.

aws-smus-cicd-cli describe --manifest manifest.yaml --connect

Step 4: Deploy

Run the deployment:

aws-smus-cicd-cli deploy --targets test --manifest manifest
During deployment, the CLI:
  1. Uploads ETL scripts and workflow definitions to S3 using the project’s storage connection.
  2. Creates the Airflow workflow in MWAA Serverless.
  3. Runs the workflow, which provisions AWS Glue jobs, creates databases, and runs ETL transformations.
  4. Imports the Quick Sight dashboard and refreshes datasets with the latest data.
  5. Processes any catalog asset subscriptions defined in the manifest.

Step 5: Validate

Run post-deployment validation to confirm services are running and ready for workloads:

aws-smus-cicd-cli test --manifest manifest.yaml --targets test

Step 6: Promote to production

Promote the same bundle artifact that was validated in the test stage to production. This guarantees the exact same artifact runs in prod:

# Promote the same bundle that was validated in test to prod

aws-smus-cicd-cli deploy --manifest app.tar.gz --targets prod

Integrating with GitHub Actions

The CLI works with existing CI/CD solutions. The GitHub repository includes reusable workflow templates that DevOps teams can adopt directly.The following is an example of a GitHub Actions workflow that implements a full bundle-based deployment pipeline:

name: Deploy Analytics Application 
on: 
  push: 
    branches: [main] 
 
jobs: 
  deploy-test: 
    runs-on: ubuntu-latest 
    steps: 
      - uses: actions/checkout@v4 
 
      - name: Install CLI 
        run: pip install aws-smus-cicd-cli 
 
      - name: Configure AWS credentials 
        uses: aws-actions/configure-aws-credentials@v4 
        with: 
          role-to-assume: ${{ secrets.AWS_ROLE_ARN }} 
          aws-region: us-east-1 
 
      - name: Validate 
        run: aws-smus-cicd-cli describe --manifest manifest.yaml --connect 
 
      - name: Bundle 
        run: aws-smus-cicd-cli bundle --manifest manifest.yaml 
 
      - name: Deploy to test 
        run: aws-smus-cicd-cli deploy --targets test --manifest manifest.yaml 
 
      - name: Run tests 
        run: aws-smus-cicd-cli test --manifest manifest.yaml --targets test 
 
  deploy-prod: 
    needs: deploy-test 
    runs-on: ubuntu-latest 
    environment: production 
    steps: 
      - uses: actions/checkout@v4 
 
      - name: Install CLI 
        run: pip install aws-smus-cicd-cli 
 
      - name: Configure AWS credentials 
        uses: aws-actions/configure-aws-credentials@v4 
        with: 
          role-to-assume: ${{ secrets.AWS_PROD_ROLE_ARN }} 
          aws-region: us-west-2 
 
      - name: Deploy to production 
        run: aws-smus-cicd-cli deploy --targets prod --manifest manifest.yaml

The CLI also works with Jenkins, GitLab CI, and Azure DevOps. See the CI/CD integration guide for additional examples.

In the next section, we cover which AWS services and workload types the CLI supports.

Supported workloads

The CLI deploys applications that span the following AWS services through Airflow workflow definitions:

  • Analytics and BI: AWS Glue ETL jobs and crawlers, Amazon Athena queries, Amazon Quick Sight dashboards, Amazon EMR jobs, Amazon Redshift queries.
  • Machine learning: SageMaker training jobs, ML model endpoints, SageMaker AI Pipelines.
  • Code and workflows: Jupyter notebooks, Python scripts, Airflow DAGs (MWAA and MWAA Serverless).
  • Data and storage: S3 data files, Git repositories, SageMaker Catalog resources (glossaries, glossary terms, form types, asset types, assets, data products, metadata forms).

The examples directory includes working applications for each of these patterns, with manifests, workflow definitions, and integration tests.

Failure recovery

If a deployment fails, the CLI stops at the point of failure and reports the error with a detailed stack trace. To recover:

  1. Run aws-smus-cicd-cli describe --connect to check which resources exist and which permissions are missing.
  2. Fix the issue and rerun aws-smus-cicd-cli deploy.
  3. For bundle-based deployments, redeploy a previous bundle version.
  4. Use aws-smus-cicd-cli destroy --targets <target> --force to clean up a failed deployment.

For detailed rollback procedures, see the Rollback Guide.

Conclusion

In this post, you learned how the Amazon SageMaker Unified Studio CI/CD CLI gives data and DevOps teams a clean separation of concerns: data teams describe their application once in a YAML manifest, and DevOps teams deploy it with a single command through their existing CI/CD pipelines. You saw how stages map to isolated SageMaker Unified Studio projects (optionally spanning AWS accounts and Regions), how bundles provide immutable, reproducible promotion through test and production, and how the CLI integrates with GitHub Actions, Jenkins, GitLab CI, and Azure DevOps. You also walked through deploying a Glue-and-Quick-Sight analytics application from dev through to prod.

Get started

The CI/CD CLI is available at no additional cost in all AWS Regions where Amazon SageMaker Unified Studio is available. You pay only for the underlying AWS resources provisioned during deployment.

Use the following steps to try it out:

  1. Install the CLI:
    pip install aws-smus-cicd-cli
  2. Browse the example applications for analytics and ML patterns.
  3. Follow the CI/CD CLI documentation to deploy your first application in 10 minutes.
  4. Review the Admin Guide for infrastructure setup.

For feedback and bug reports, open an issue on the GitHub repository.


About the authors

Ramesh H Singh

Ramesh H Singh

Ramesh H Singh is a Senior Product Manager Technical (External Services) at AWS in Seattle, Washington, currently with the Amazon SageMaker team. He is passionate about building high-performance ML/AI and analytics products that help enterprise customers achieve their critical goals using cutting-edge technology.

Vasudevan Venkataramanan

Vasudevan Venkataramanan

Vasudevan Venkataramanan is a Senior Software Engineer on the Amazon SageMaker Unified Studio team. He is responsible for technical direction of scheduling and orchestration within SageMaker Unified Studio. Outside of his professional work, he enjoys spending time with his kid, and playing pickleball and cricket.

Amir Bar Or

Amir Bar Or

Amir Bar Or is a Senior Software Engineer on the Amazon SageMaker Unified Studio team. He is responsible for technical direction of scheduling and orchestration within SageMaker Unified Studio. Outside of his professional work, he enjoys spending time with his kid, and playing pickleball and cricket.

Nikita Arbuzov

Nikita Arbuzov

Nikita is Software Engineer on the Amazon SageMaker Unified Studio team. He is responsible for building support for CI/CD features within SageMaker Unified Studio.

Saurabh Bhutyani

Saurabh Bhutyani

Saurabh Bhutyani is a Principal Analytics Specialist Solutions Architect at AWS. He is passionate about new technologies. He joined AWS in 2019 and works with customers to provide architectural guidance for running generative AI use cases, scalable analytics solutions and data mesh architectures using AWS services like Amazon Bedrock, Amazon SageMaker Unified Studio, Amazon EMR, Amazon Athena, AWS Glue, AWS Lake Formation, and Amazon DataZone.

Capture data lineage of Amazon EMR spark jobs into Amazon SageMaker Unified Studio

Post Syndicated from Jose Romero original https://aws.amazon.com/blogs/big-data/capture-data-lineage-of-amazon-emr-spark-jobs-into-amazon-sagemaker-unified-studio/

Data engineers running Apache Spark jobs on Amazon EMR face a persistent challenge: understanding how data moves through Spark pipelines as it’s transformed, joined, and written to downstream tables . Tracking these transformations manually requires examining job logs, reviewing code, and piecing together transformation logic across multiple sources. As pipelines scale, this process becomes complex. The visibility gap affects key business activities: troubleshooting data quality issues takes longer – impact analysis for schema changes requires more effort – and compliance audits need extensive documentation of data provenance.

Amazon SageMaker is the center for all your data and analytics where you can find and access all the data in your organization and act on it using tools across various use case. This unified platform addresses the data visibility challenge by bringing together data governance, collaboration, and discovery into a single interface. At the heart of this platform is Amazon SageMaker Catalog, a centralized hub that enables organizations to catalog, govern, and discover all their data assets with complete visibility into lineage. By capturing data lineage across your entire data ecosystem from raw sources through transformations to final outputs, SageMaker Catalog enables you to track data provenance across your entire platform, enable collaboration with clear visibility into data ownership and quality metrics, build trust through comprehensive data lineage that supports compliance and confident decision-making, and accelerate discovery of trustworthy, governance-ready data assets. You can access and visualize this lineage directly in Amazon SageMaker Unified Studio, which serves as the unified interface to explore data relationships and collaborate across your analytics workflows.

Amazon EMR, starting from version 7.11, now includes native OpenLineage support that automates lineage capture. OpenLineage is an open-source framework for data lineage that automatically emits lineage metadata from your data transformation jobs directly into Amazon SageMaker Catalog, or other data governance solutions, without requiring customizations.

This EMR native support of OpenLineage is part of a growing set of integrations across AWS analytics services including AWS Glue, Amazon EMR Serverless, and Amazon Redshift. The complete list of services with native OpenLineage integration can be found in the data lineage support matrix.

In this post, you’ll walk through a practical, step-by-step example that shows how to capture and track data lineage from Spark jobs running on Amazon EMR directly into Amazon SageMaker Catalog using OpenLineage. You’ll see how lineage metadata flows automatically and explore data relationships and dependencies across your workflows in Amazon SageMaker Unified Studio.

Solution overview

Imagine you’re part of a large enterprise that relies on HR analytics to optimize workforce planning, compensation strategies, and talent retention practices. Your data engineering team owns the delivery of these analytical products by processing raw HR datasets (including employee records, attendance logs, and compensation details), with Spark jobs running on your Amazon EMR infrastructure.

With time, Spark jobs have grown in complexity. Your team now struggles to maintain visibility into how data moves through pipelines, who modified it, and how to map dependencies between datasets and final analytical products.

The following solution demonstrates how you can address these challenges by automatically capturing data lineage end-to-end from Spark jobs running on your EMR infrastructure and visualizing it in Amazon SageMaker Unified Studio so that you and the business understand data provenance of the final analytical products.

AWS cloud data pipeline architecture diagram showing data flowing from Amazon S3 CSV files (employees.csv, attendance.csv) through Amazon EMR with Apache Spark processing, AWS Glue Data Catalog metadata management, and Amazon SageMaker Catalog integration, producing salary_adjustments.csv and bonus_payments.csv output files stored in Amazon S3.

The architecture includes a Data Layer with CSV files containing employee, attendance, salary, and bonus data stored in Amazon S3 (Simple Storage Service), representing typical HR and payroll source systems.

The Processing Layer uses Amazon EMR cluster running Apache Spark jobs that transform raw data into analytical tables. The first Spark job joins employee and attendance data while the second Spark job combines attendance with compensation data. Both jobs use Apache Iceberg table format to provide ACID (Atomic, Consistent, Isolated, and Durable) transactions and time travel capabilities.

The Metadata Layer uses AWS Glue Data Catalog to store Iceberg table metadata, making tables discoverable and accessible across AWS analytics services. A Lineage Layer uses the OpenLineage integration in EMR to automatically track input/output datasets (CSV files and Iceberg tables), transformation logic at column level (joins, filters, aggregations), and job execution metadata.

Finally, the Data Governance Layer uses Amazon SageMaker Catalog to capture and process OpenLineage events posted by the EMR Spark jobs and automatically build a comprehensive lineage graph that shows complete data provenance from CSV source files through Spark transformations to Iceberg analytical tables.

Before you deploy this solution, make sure you have the following resources in place.

Prerequisites

For this walkthrough, you should have the following prerequisites:

  • An AWS account.
  • Your assumed role should have full access to Amazon EMR serverless, Amazon S3, Amazon Identity and Access Management (IAM) and AWS Lambda. Note that for production workloads, minimum permissions are recommended.
  • A Amazon VPC (Virtual Private Cloud) with at least one subnet with internet access. You can provision this VPC as you create the Amazon SageMaker domain next.
  • An existing Amazon SageMaker Unified Studio domain and project. To get started, use the quick setup option as explained here. To create a project, follow the instructions here.
  • An S3 bucket with the sample data files and Spark scripts uploaded (see Prepare Your Source Data below)
  • Default EMR service roles — if this is your first time using EMR in this account, run `aws emr create-default-roles` from the AWS CLI or CloudShell to create them.

With these prerequisites in place, let’s examine what the AWS CloudFormation template will deploy to your AWS environment.

Architecture components

The deployment creates several interconnected components that work together to capture and visualize lineage:

  • An S3 bucket to store all data and artifacts for the solution.
  • An EMR cluster (v 7.12.0) with Apache Iceberg support enabled and OpenLineage integration pre-installed, ready to run Spark jobs with lineage tracking.
  • A set of IAM policies that grant the necessary permissions to the EMR cluster to post lineage events to your SageMaker Unified Studio domain.
  • A set of AWS Lake Formation permissions that grant the EMR cluster to create, alter, and drop Iceberg tables in your specified Glue database.

With an understanding of what will be deployed, you’re ready to launch the CloudFormation stack.

Deploy the solution

Note: While this walkthrough uses the AWS EMR console and AWS CLI to verify the cluster and run Spark jobs, you can also perform these steps directly from Amazon SageMaker Unified Studio. SMUS provides a unified interface to create and manage EMR clusters, submit Spark jobs, and monitor execution — all within the same environment where you’ll later explore the lineage captured in Amazon SageMaker Catalog.

Prepare your source data

Before deploying the CloudFormation stack, clone or download the following git repository.PutHereGitRepo

Upload the CSV files downloaded from git to the input/ prefix and the spark scripts in scripts/ prefix. You can run the following command to upload the files:

aws s3 cp employees.csv s3://YOUR-BUCKET/input/
aws s3 cp attendance.csv s3://YOUR-BUCKET/input/
aws s3 cp salary_adjustments.csv s3://YOUR-BUCKET/input/
aws s3 cp bonus_payments.csv s3://YOUR-BUCKET/input/
aws s3 cp emr-lineage-spark-job.py s3://YOUR-BUCKET/scripts/
aws s3 cp emr-lineage-compensation-job.py s3://YOUR-BUCKET/scripts/

To deploy the solution, complete the following steps in CloudFormation console:

  1. Create new stack by specifying the CloudFormation yaml file previously download from git repository PutHereThe YMLFileName
  2. Enter a stack name (e.g., emr-lineage-demo) and provide the following parameters:
    • SourceS3BucketName: S3 bucket containing your CSV files and Spark scripts
    • SourceCSVPrefix: S3 prefix where CSV files are located
    • SourceScriptsPrefix: S3 prefix where Spark scripts are located
    • GlueDatabaseName: The name of the Glue database associated to your Amazon SageMaker Unified Studio project.
    • DataZoneDomainId: Your SageMaker Unified Studio domain ID.
    • VpcId: The id of the VPC that was deployed as part of the prerequisites.
    • For EMRReleaseLabel, MasterInstanceType, CoreInstanceType and CoreInstanceCount, keep the default values.
  3. Acknowledge IAM resource creation, choose Next and then Submit. The CloudFormation stack takes approximately 10 to 15 minutes to complete.
  4. In the EMR console, wait for the cluster status to show as WAITING before moving to the next step.

Screenshot of the Amazon EMR on EC2 Clusters management console showing a list of 14 clusters, with the cluster "EMR-Lineage-Demo-emr-ec2-lineage-demo-stack" (ID: j-3APWOTUDNYO2T) highlighted in a "Waiting – Ready to run steps" status with a green badge.

Now that the EMR cluster is running with OpenLineage enabled, let’s examine how the Spark jobs are configured to capture lineage metadata.

Explore data lineage configuration in EMR

When submitting Spark jobs to EMR, specific configurations enable OpenLineage to create and post lineage events to SageMaker Unified Studio as the job runs:

  • spark.hadoop.hive.metastore.client.factory.class – Configures Spark to use AWS Glue as the Hive metastore.
  • spark.jars – Path to the pre-installed OpenLineage library (available on EMR 7.11+).
  • spark.extraListeners – Registers an OpenLineage listener to capture metadata of input / output datasets and transformations.
  • spark.openlineage.transport.type – Uses the OpenLineage DataZone transport option to send lineage events directly into SageMaker Catalog.
  • spark.openlineage.transport.domainId – The ID of your SageMaker Unified Studio domain, that serves as the target for lineage events.
  • spark.glue.accountId – Your AWS account ID for Glue data catalog operations.

Now that you understand the configuration that enables automatic lineage capture, you’re ready to run the data pipeline.

When running this two-step pipeline, you will calculate the total employee compensation by combining salary adjustments, bonuses, and attendance data. The final analytical asset will serve payroll processing and budgeting.

Run employee attendance analysis job

The first job reads employee details (in employees.csv dataset) and attendance records (in attendance.csv dataset), joins the datasets on EmployeeID and creates a unified dataset (employee_attendance Iceberg table) in your Glue database.

Follow the steps below to run this first job:

  1. In the CloudFormation console, navigate to the stack’s Outputs tab
  2. Copy the value of the Job1SubmitCommand output key. Note that this is the command you’ll use to submit the first job in EMR with the right configuration.

AWS CloudFormation console screenshot showing the Outputs tab for the "emr-ec2-lineage-demo-stack" stack, displaying 9 outputs including the Job1SubmitCommand — an AWS EMR add-steps command with Apache Spark configuration for the EMR Lineage Demo Job targeting cluster j-3APWOTUDNYO2T.

  1. Run the command in your terminal or AWS CloudShell.
  2. Monitor the job in the Amazon EMR console under Steps.

Screenshot of the Amazon EMR console Steps tab for the cluster "EMR-Lineage-Demo-emr-ec2-lineage-demo-stack," showing one completed step named "EMR-Lineage-Demo-Job" with Step ID s-0270631D8DHBCJZKBAZ and a green "Completed" status checkmark.

Run employee compensation analysis job

Now, you will calculate the total employee compensation (Iceberg table) by combining salary adjustments (salary_adjustments.csv dataset), bonuses (bonus_payments.csv dataset), and attendance (calculated in the last step):

  1. Repeat the steps 1 to 4 to run Job 2.
  2. After completion, open the AWS Glue console.
  3. Navigate to Data Catalog, then Tables and select your SageMaker project’s database.
  4. Confirm that employee_attendance and employee_compensation tables are listed.

With both Spark jobs complete, you can now visualize the complete data lineage graph in Amazon SageMaker Unified Studio.

Visualizing lineage in SageMaker Unified Studio

SageMaker Unified Studio provides a graph-based data lineage visualization that helps data engineers, analysts, and data scientists clearly understand which source datasets (files or tables) feed into each dataset, what transformations and logic are applied at every step, which downstream analytics assets consume the data, and how changes to upstream data or transformations may impact the rest of the data pipeline.

Now that the data pipeline run successfully, let’s review the captured lineage for the HR data in SageMaker Unified Studio:

  1. Navigate to the SageMaker Unified Studio console, sign in to your domain.
  2. Open your project and go to Data Sources
  3. Find your AWS Glue Data Catalog source

Screenshot of the Amazon SageMaker project catalog Data Sources page listing three configured data sources: a Redshift Serverless source, an AWS Glue Lakehouse source named "AwsDataCatalog-emr_ec2_lineage_blogpost_glue_db-default-datasource" (highlighted), and a Tooling SageMaker model package group source — all scheduled MTWTFSS and in Ready or Running status.

  1. Click RUN. Two new assets will be created.

Screenshot of the AWS Glue Data Catalog interface showing run activities for the data source "AwsDataCatalog-emr_ec2_lineage_blogpost_glue_db-default-datasource," with two completed on-demand runs and a highlighted asset table showing employee_attendance and employee_compensation successfully created in the emr_ec2_lineage_blogpost_glue_db database.

  1. Navigate to Assets and Click on employee_compensation. Under the LINEAGE tab you’ll find the lineage graph view that SageMaker builds based on the OpenLineage metadata captured from the EMR Spark jobs as they run.

AWS Glue data lineage visualization showing the flow of the employee_compensation dataset from an Apache Spark job (default.emr_lineage_compensa, COMPLETE, Dec 22 2025 11:42:47 AM) through an AWS Glue Iceberg table (20 columns) to an AWS Glue Inventory destination table, with a right sidebar displaying lineage metadata including the dataset ARN, OpenLineage producer URL, Iceberg snapshot ID, and projected field names EmployeeID, Name, and Department.

    • You’ll first see three lineage nodes from left to right: one representing the EMR Spark job that created the final Iceberg table, a second one representing the actual Iceberg table in the Glue catalog, and a third one representing the data asset in the SageMaker Catalog inventory that maps to the Glue table.
    • Click on any lineage node to view its underlying metadata in the details pane, including dataset names, S3 locations, schema, data types, job execution details and more.
  1. Expand the lineage to the left by clicking on the double arrow next to the first lineage node. Keep expanding until you hit the originating datasets.

Data pipeline lineage diagram showing the complete ETL flow from Amazon S3 source files (input/attendance.csv with 6 columns, input/employees.csv with 5 columns) through two Apache Spark jobs to intermediate tables (input/salary_adjustments.csv, iceberg/employee.csv, AWS Glue employee_attendance with 14 columns) and final destination tables (AWS Glue iceberg/employee_compensation with 29 columns, AWS Glue Inventory employee_compensation_hive with 30 columns), all timestamped Dec 22, 2025.

    • Expanding the graph to the left reveals the complete data pipeline back to original CSV source files. You can see how compensation data depends on upstream attendance analytics.
    • Note how each lineage node represents an element in the data pipeline you run, including both Spark jobs and even the intermediate employee_attendance Iceberg table that connects them.
  1. You can expand column-level lineage by clicking on the column section of a lineage node of a dataset or data asset. This allows you to understand how data changes at a column level as it goes downstream your data pipeline.

Data lineage diagram showing the employee compensation ETL pipeline with four Amazon S3 source tables (employee.csv with 5 columns, input/attendance.csv with 6 columns, input/salary_adjustments.csv with 4 columns, output/employee_attendance.csv with 14 columns) processed by two Apache Spark jobs to produce a final s3://employee_compensation table with 20 columns, all dated Dec 22, 2025.

Cleanup

To avoid ongoing charges, clean up the resources:

  1. First, empty the destination bucket by running the following command in your terminal or with AWS CloudShell.

aws s3 rm s3://${DEST_BUCKET}/ --recursive

  1. Delete the CloudFormation stack.
    • On the AWS CloudFormation console, choose Stacks in the navigation pane.
    • Choose the stack you created, then choose Delete and then Delete stack when prompted.

Conclusion

In this post, you explore how to capture data lineage from Spark jobs in Amazon EMR (v7.11+) directly into Amazon SageMaker Unified Studio. You learned how to set up an Amazon EMR cluster with native OpenLineage support to automatically track lineage metadata from Spark jobs processing your data. You also configured the integration between EMR and Amazon SageMaker Catalog to ensure lineage information flows seamlessly into your governance platform. Finally, you explored the resulting lineage graph in SageMaker Unified Studio and saw how it provides comprehensive visibility into data transformations, from source CSV files through Spark processing jobs to final analytical tables using Apache Iceberg format.

We encourage you to now test these capabilities with your own data pipelines running on EMR. By implementing automated lineage tracking, many customers have strengthened their governance frameworks while gaining valuable insights into data dependencies, impact analysis, and compliance requirements. This approach enables data teams to build trust in their analytics outputs while maintaining the agility needed to derive business value from their data assets.


About the authors

Yanick Houngbedji is a Solutions Architect for Independent Software Vendors (ISV) at Amazon Web Services (AWS), based in Montréal, Canada. He specializes in helping customers architect and implement highly scalable, performant, and secure cloud solutions on AWS. Before joining AWS, he spent over 8 years providing technical leadership in data engineering, big data analytics, business intelligence, and data science solutions.

Jose Romero is a Senior Solutions Architect for Startups at Amazon Web Services (AWS) based in Austin, TX, US. He is passionate about helping customers architect modern platforms at scale for data, AI, and ML. As a former senior architect in AWS Professional Services, he enjoys building and sharing solutions for common complex problems so that customers can accelerate their cloud journey and adopt best practices. Connect with him on LinkedIn.