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

0
1
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.

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"
          ]
        }
      }
    }
  ]
}

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.