SageMaker

Important Capabilities

Capability	Status	Notes
Detect Deleted Entities	✅	Optionally enabled via stateful_ingestion.remove_stale_metadata
Table-Level Lineage	✅	Enabled by default

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)

CLI based Ingestion

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

Options

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

Field	Description
aws_access_key_id string	AWS access key ID. Can be auto-detected, see the AWS boto3 docs for details.
aws_advanced_config object	Advanced AWS configuration options. These are passed directly to botocore.config.Config.
aws_endpoint_url string	The AWS service endpoint. This is normally constructed automatically, but can be overridden here.
aws_profile string	Named AWS profile to use. Only used if access key / secret are unset. If not set the default will be used
aws_proxy map(str,string)
aws_region string	AWS region code.
aws_secret_access_key string	AWS secret access key. Can be auto-detected, see the AWS boto3 docs for details.
aws_session_token string	AWS session token. Can be auto-detected, see the AWS boto3 docs for details.
extract_feature_groups boolean	Whether to extract feature groups. Default: True
extract_jobs One of string, boolean	Whether to extract AutoML jobs. Default: True
extract_models boolean	Whether to extract models. Default: True
platform_instance string	The instance of the platform that all assets produced by this recipe belong to
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	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 documented at https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts.html?highlight=assume_role#STS.Client.assume_role
aws_role.union One of string, AwsAssumeRoleConfig
aws_role.union.RoleArn ❓ string	ARN of the role to assume.
aws_role.union.ExternalId string	External ID to use when assuming the role.
database_pattern AllowDenyPattern	regex patterns for databases to filter in ingestion. Default: {'allow': ['.*'], 'deny': [], 'ignoreCase': True}
database_pattern.ignoreCase boolean	Whether to ignore case sensitivity during pattern matching. Default: True
database_pattern.allow array	List of regex patterns to include in ingestion Default: ['.*']
database_pattern.allow.string string
database_pattern.deny array	List of regex patterns to exclude from ingestion. Default: []
database_pattern.deny.string string
table_pattern AllowDenyPattern	regex patterns for tables to filter in ingestion. Default: {'allow': ['.*'], 'deny': [], 'ignoreCase': True}
table_pattern.ignoreCase boolean	Whether to ignore case sensitivity during pattern matching. Default: True
table_pattern.allow array	List of regex patterns to include in ingestion Default: ['.*']
table_pattern.allow.string string
table_pattern.deny array	List of regex patterns to exclude from ingestion. Default: []
table_pattern.deny.string string
stateful_ingestion StatefulStaleMetadataRemovalConfig	Base specialized config for Stateful Ingestion with stale metadata removal capability.
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.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

Schema

The JSONSchema for this configuration is inlined below.

{
  "title": "SagemakerSourceConfig",
  "description": "Common AWS credentials config.\n\nCurrently used by:\n    - Glue source\n    - DynamoDB source\n    - SageMaker source",
  "type": "object",
  "properties": {
    "stateful_ingestion": {
      "$ref": "#/definitions/StatefulStaleMetadataRemovalConfig"
    },
    "platform_instance": {
      "title": "Platform Instance",
      "description": "The instance of the platform that all assets produced by this recipe belong to",
      "type": "string"
    },
    "aws_access_key_id": {
      "title": "Aws Access Key Id",
      "description": "AWS access key ID. Can be auto-detected, see [the AWS boto3 docs](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html) for details.",
      "type": "string"
    },
    "aws_secret_access_key": {
      "title": "Aws Secret Access Key",
      "description": "AWS secret access key. Can be auto-detected, see [the AWS boto3 docs](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html) for details.",
      "type": "string"
    },
    "aws_session_token": {
      "title": "Aws Session Token",
      "description": "AWS session token. Can be auto-detected, see [the AWS boto3 docs](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html) for details.",
      "type": "string"
    },
    "aws_role": {
      "title": "Aws Role",
      "description": "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 documented at https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts.html?highlight=assume_role#STS.Client.assume_role",
      "anyOf": [
        {
          "type": "string"
        },
        {
          "type": "array",
          "items": {
            "anyOf": [
              {
                "type": "string"
              },
              {
                "$ref": "#/definitions/AwsAssumeRoleConfig"
              }
            ]
          }
        }
      ]
    },
    "aws_profile": {
      "title": "Aws Profile",
      "description": "Named AWS profile to use. Only used if access key / secret are unset. If not set the default will be used",
      "type": "string"
    },
    "aws_region": {
      "title": "Aws Region",
      "description": "AWS region code.",
      "type": "string"
    },
    "aws_endpoint_url": {
      "title": "Aws Endpoint Url",
      "description": "The AWS service endpoint. This is normally [constructed automatically](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html), but can be overridden here.",
      "type": "string"
    },
    "aws_proxy": {
      "title": "Aws Proxy",
      "description": "A set of proxy configs to use with AWS. See the [botocore.config](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html) docs for details.",
      "type": "object",
      "additionalProperties": {
        "type": "string"
      }
    },
    "read_timeout": {
      "title": "Read Timeout",
      "description": "The timeout for reading from the connection (in seconds).",
      "default": 60,
      "type": "number"
    },
    "aws_advanced_config": {
      "title": "Aws Advanced Config",
      "description": "Advanced AWS configuration options. These are passed directly to [botocore.config.Config](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html).",
      "type": "object"
    },
    "env": {
      "title": "Env",
      "description": "The environment that all assets produced by this connector belong to",
      "default": "PROD",
      "type": "string"
    },
    "database_pattern": {
      "title": "Database Pattern",
      "description": "regex patterns for databases to filter in ingestion.",
      "default": {
        "allow": [
          ".*"
        ],
        "deny": [],
        "ignoreCase": true
      },
      "allOf": [
        {
          "$ref": "#/definitions/AllowDenyPattern"
        }
      ]
    },
    "table_pattern": {
      "title": "Table Pattern",
      "description": "regex patterns for tables to filter in ingestion.",
      "default": {
        "allow": [
          ".*"
        ],
        "deny": [],
        "ignoreCase": true
      },
      "allOf": [
        {
          "$ref": "#/definitions/AllowDenyPattern"
        }
      ]
    },
    "extract_feature_groups": {
      "title": "Extract Feature Groups",
      "description": "Whether to extract feature groups.",
      "default": true,
      "type": "boolean"
    },
    "extract_models": {
      "title": "Extract Models",
      "description": "Whether to extract models.",
      "default": true,
      "type": "boolean"
    },
    "extract_jobs": {
      "title": "Extract Jobs",
      "description": "Whether to extract AutoML jobs.",
      "default": true,
      "anyOf": [
        {
          "type": "object",
          "additionalProperties": {
            "type": "string"
          }
        },
        {
          "type": "boolean"
        }
      ]
    }
  },
  "additionalProperties": false,
  "definitions": {
    "DynamicTypedStateProviderConfig": {
      "title": "DynamicTypedStateProviderConfig",
      "type": "object",
      "properties": {
        "type": {
          "title": "Type",
          "description": "The type of the state provider to use. For DataHub use `datahub`",
          "type": "string"
        },
        "config": {
          "title": "Config",
          "description": "The configuration required for initializing the state provider. Default: The datahub_api config if set at pipeline level. Otherwise, the default DatahubClientConfig. See the defaults (https://github.com/datahub-project/datahub/blob/master/metadata-ingestion/src/datahub/ingestion/graph/client.py#L19).",
          "default": {},
          "type": "object"
        }
      },
      "required": [
        "type"
      ],
      "additionalProperties": false
    },
    "StatefulStaleMetadataRemovalConfig": {
      "title": "StatefulStaleMetadataRemovalConfig",
      "description": "Base specialized config for Stateful Ingestion with stale metadata removal capability.",
      "type": "object",
      "properties": {
        "enabled": {
          "title": "Enabled",
          "description": "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,
          "type": "boolean"
        },
        "remove_stale_metadata": {
          "title": "Remove Stale Metadata",
          "description": "Soft-deletes the entities present in the last successful run but missing in the current run with stateful_ingestion enabled.",
          "default": true,
          "type": "boolean"
        }
      },
      "additionalProperties": false
    },
    "AwsAssumeRoleConfig": {
      "title": "AwsAssumeRoleConfig",
      "type": "object",
      "properties": {
        "RoleArn": {
          "title": "Rolearn",
          "description": "ARN of the role to assume.",
          "type": "string"
        },
        "ExternalId": {
          "title": "Externalid",
          "description": "External ID to use when assuming the role.",
          "type": "string"
        }
      },
      "required": [
        "RoleArn"
      ]
    },
    "AllowDenyPattern": {
      "title": "AllowDenyPattern",
      "description": "A class to store allow deny regexes",
      "type": "object",
      "properties": {
        "allow": {
          "title": "Allow",
          "description": "List of regex patterns to include in ingestion",
          "default": [
            ".*"
          ],
          "type": "array",
          "items": {
            "type": "string"
          }
        },
        "deny": {
          "title": "Deny",
          "description": "List of regex patterns to exclude from ingestion.",
          "default": [],
          "type": "array",
          "items": {
            "type": "string"
          }
        },
        "ignoreCase": {
          "title": "Ignorecase",
          "description": "Whether to ignore case sensitivity during pattern matching.",
          "default": true,
          "type": "boolean"
        }
      },
      "additionalProperties": false
    }
  }
}

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.