Skip to main content

SageMaker

Overview​

Sagemaker is a machine learning platform. Learn more in the official Sagemaker documentation.

The DataHub integration for Sagemaker covers ML entities such as models, features, and related lineage metadata. Depending on module capabilities, it can also capture features such as lineage, usage, profiling, ownership, tags, and stateful deletion detection.

Concept Mapping​

While the specific concept mapping is still pending, this shows the generic concept mapping in DataHub.

Source ConceptDataHub ConceptNotes
Platform/account/project scopePlatform Instance, ContainerOrganizes assets within the platform context.
Core technical asset (for example table/view/topic/file)DatasetPrimary ingested technical asset.
Schema fields / columnsSchemaFieldIncluded when schema extraction is supported.
Ownership and collaboration principalsCorpUser, CorpGroupEmitted by modules that support ownership and identity metadata.
Dependencies and processing relationshipsLineage edgesAvailable when lineage extraction is supported and enabled.

Module sagemaker​

Certified

Important Capabilities​

CapabilityStatusNotes
Detect Deleted Entities✅Enabled by default via stateful ingestion.
Table-Level Lineage✅Enabled by default.

Overview​

The sagemaker module ingests metadata from SageMaker into DataHub. It is intended for production ingestion workflows and module-specific capabilities are documented below.

This plugin extracts the following:

  • Feature groups
  • Models, jobs, and lineage between the two (e.g. when jobs output a model or a model is used by a job)

Prerequisites​

Before running ingestion, ensure network connectivity to the source, valid authentication credentials, and read permissions for metadata APIs required by this module.

Install the Plugin​

pip install 'acryl-datahub[sagemaker]'

Starter Recipe​

Check out the following recipe to get started with ingestion! See below for full configuration options.

For general pointers on writing and running a recipe, see our main recipe guide.

source:
  type: sagemaker
  config:
    # Coordinates
    aws_region: "my-aws-region"

sink:
  # sink configs

Config Details​

Note that a . is used to denote nested fields in the YAML recipe.

FieldDescription
aws_access_key_id
One of string, null
AWS access key ID. Can be auto-detected, see the AWS boto3 docs for details.
Default: None
aws_advanced_config
object
Advanced AWS configuration options. These are passed directly to botocore.config.Config.
aws_endpoint_url
One of string, null
The AWS service endpoint. This is normally constructed automatically, but can be overridden here.
Default: None
aws_profile
One of string, null
The named profile to use from AWS credentials. Falls back to default profile if not specified and no access keys provided. Profiles are configured in ~/.aws/credentials or ~/.aws/config.
Default: None
aws_proxy
One of string, null
A set of proxy configs to use with AWS. See the botocore.config docs for details.
Default: None
aws_region
One of string, null
AWS region code.
Default: None
aws_retry_mode
Enum
One of: "legacy", "standard", "adaptive"
Default: standard
aws_retry_num
integer
Number of times to retry failed AWS requests. See the botocore.retry docs for details.
Default: 5
aws_secret_access_key
One of string(password), null
AWS secret access key. Can be auto-detected, see the AWS boto3 docs for details.
Default: None
aws_session_token
One of string(password), null
AWS session token. Can be auto-detected, see the AWS boto3 docs for details.
Default: None
extract_feature_groups
One of boolean, null
Whether to extract feature groups.
Default: True
extract_jobs
One of string, boolean, null
Whether to extract AutoML jobs.
Default: True
extract_models
One of boolean, null
Whether to extract models.
Default: True
read_timeout
number
The timeout for reading from the connection (in seconds).
Default: 60
env
string
The environment that all assets produced by this connector belong to
Default: PROD
aws_role
One of string, array, null
AWS roles to assume. If using the string format, the role ARN can be specified directly. If using the object format, the role can be specified in the RoleArn field and additional available arguments are the same as boto3's STS.Client.assume_role.
Default: None
aws_role.union
One of string, AwsAssumeRoleConfig
aws_role.union.RoleArn ❓
string
ARN of the role to assume.
aws_role.union.ExternalId
One of string, null
External ID to use when assuming the role.
Default: None
database_pattern
AllowDenyPattern
A class to store allow deny regexes
database_pattern.ignoreCase
One of boolean, null
Whether to ignore case sensitivity during pattern matching.
Default: True
table_pattern
AllowDenyPattern
A class to store allow deny regexes
table_pattern.ignoreCase
One of boolean, null
Whether to ignore case sensitivity during pattern matching.
Default: True
stateful_ingestion
One of StatefulStaleMetadataRemovalConfig, null
Default: None
stateful_ingestion.enabled
boolean
Whether or not to enable stateful ingest. Default: True if a pipeline_name is set and either a datahub-rest sink or datahub_api is specified, otherwise False
Default: False
stateful_ingestion.fail_safe_threshold
number
Prevents large amount of soft deletes & the state from committing from accidental changes to the source configuration if the relative change percent in entities compared to the previous state is above the 'fail_safe_threshold'.
Default: 75.0
stateful_ingestion.remove_stale_metadata
boolean
Soft-deletes the entities present in the last successful run but missing in the current run with stateful_ingestion enabled.
Default: True

Capabilities​

Use the Important Capabilities table above as the source of truth for supported features and whether additional configuration is required.

Limitations​

Module behavior is constrained by source APIs, permissions, and metadata exposed by the platform. Refer to capability notes for unsupported or conditional features.

Troubleshooting​

If ingestion fails, validate credentials, permissions, connectivity, and scope filters first. Then review ingestion logs for source-specific errors and adjust configuration accordingly.

Code Coordinates​

  • Class Name: datahub.ingestion.source.aws.sagemaker.SagemakerSource
  • Browse on GitHub
Questions?

If you've got any questions on configuring ingestion for SageMaker, feel free to ping us on our Slack.

💡 Contributing to this documentation

This page is auto-generated from the underlying source code. To make changes, please edit the relevant source files in the metadata-ingestion directory.

Tip: For quick typo fixes or documentation updates, you can click the ✏️ Edit icon directly in the GitHub UI to open a Pull Request. For larger changes and PR naming conventions, please refer to our Contributing Guide.