feat(ingestion/SageMaker): Remove deprecated apis and add stateful in…

…gestion capability (datahub-project#10573)

metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import DefaultDict, Dict, Iterable
+from typing import DefaultDict, Dict, Iterable, List, Optional
 
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
@@ -10,7 +10,7 @@
  platform_name,
  support_status,
 )
-from datahub.ingestion.api.source import Source
+from datahub.ingestion.api.source import MetadataWorkUnitProcessor
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.aws.sagemaker_processors.common import (
  SagemakerSourceConfig,
@@ -26,13 +26,19 @@
 )
 from datahub.ingestion.source.aws.sagemaker_processors.lineage import LineageProcessor
 from datahub.ingestion.source.aws.sagemaker_processors.models import ModelProcessor
+from datahub.ingestion.source.state.stale_entity_removal_handler import (
+ StaleEntityRemovalHandler,
+)
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+ StatefulIngestionSourceBase,
+)
 
 
 @platform_name("SageMaker")
 @config_class(SagemakerSourceConfig)
 @support_status(SupportStatus.CERTIFIED)
 @capability(SourceCapability.LINEAGE_COARSE, "Enabled by default")
-class SagemakerSource(Source):
+class SagemakerSource(StatefulIngestionSourceBase):
  """
  This plugin extracts the following:
 
@@ -45,7 +51,7 @@ class SagemakerSource(Source):
  report = SagemakerSourceReport()
 
  def __init__(self, config: SagemakerSourceConfig, ctx: PipelineContext):
- super().__init__(ctx)
+ super().__init__(config, ctx)
  self.source_config = config
  self.report = SagemakerSourceReport()
  self.sagemaker_client = config.sagemaker_client
@@ -56,6 +62,14 @@ def create(cls, config_dict, ctx):
  config = SagemakerSourceConfig.parse_obj(config_dict)
  return cls(config, ctx)
 
+ def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
+ return [
+ *super().get_workunit_processors(),
+ StaleEntityRemovalHandler.create(
+ self, self.source_config, self.ctx
+ ).workunit_processor,
+ ]
+
  def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
  # get common lineage graph
  lineage_processor = LineageProcessor(

metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/common.py

@@ -1,13 +1,22 @@
-from dataclasses import dataclass
-from typing import Dict, Optional, Union
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
 
 from pydantic.fields import Field
 
-from datahub.ingestion.api.source import SourceReport
+from datahub.configuration.source_common import PlatformInstanceConfigMixin
 from datahub.ingestion.source.aws.aws_common import AwsSourceConfig
+from datahub.ingestion.source.state.stale_entity_removal_handler import (
+ StaleEntityRemovalSourceReport,
+ StatefulIngestionConfigBase,
+ StatefulStaleMetadataRemovalConfig,
+)
 
 
-class SagemakerSourceConfig(AwsSourceConfig):
+class SagemakerSourceConfig(
+ AwsSourceConfig,
+ PlatformInstanceConfigMixin,
+ StatefulIngestionConfigBase,
+):
  extract_feature_groups: Optional[bool] = Field(
  default=True, description="Whether to extract feature groups."
  )
@@ -17,21 +26,24 @@ class SagemakerSourceConfig(AwsSourceConfig):
  extract_jobs: Optional[Union[Dict[str, str], bool]] = Field(
  default=True, description="Whether to extract AutoML jobs."
  )
+ # Custom Stateful Ingestion settings
+ stateful_ingestion: Optional[StatefulStaleMetadataRemovalConfig] = None
 
  @property
  def sagemaker_client(self):
  return self.get_sagemaker_client()
 
 
 @dataclass
-class SagemakerSourceReport(SourceReport):
+class SagemakerSourceReport(StaleEntityRemovalSourceReport):
  feature_groups_scanned = 0
  features_scanned = 0
  endpoints_scanned = 0
  groups_scanned = 0
  models_scanned = 0
  jobs_scanned = 0
  datasets_scanned = 0
+ filtered: List[str] = field(default_factory=list)
 
  def report_feature_group_scanned(self) -> None:
  self.feature_groups_scanned += 1
@@ -53,3 +65,6 @@ def report_job_scanned(self) -> None:
 
  def report_dataset_scanned(self) -> None:
  self.datasets_scanned += 1
+
+ def report_dropped(self, name: str) -> None:
+ self.filtered.append(name)

metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/job_classes.py

@@ -1,4 +1,6 @@
-from typing import Dict, Final
+from typing import Dict
+
+from typing_extensions import Final
 
 from datahub.metadata.schema_classes import JobStatusClass
 
@@ -63,8 +65,9 @@ class AutoMlJobInfo(SageMakerJobInfo):
  list_key: Final = "AutoMLJobSummaries"
  list_name_key: Final = "AutoMLJobName"
  list_arn_key: Final = "AutoMLJobArn"
- # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_auto_ml_job
- describe_command: Final = "describe_auto_ml_job"
+ # DescribeAutoMLJobV2 are new versions of CreateAutoMLJob and DescribeAutoMLJob which offer backward compatibility.
+ # https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_DescribeAutoMLJobV2.html
+ describe_command: Final = "describe_auto_ml_job_v2"
  describe_name_key: Final = "AutoMLJobName"
  describe_arn_key: Final = "AutoMLJobArn"
  describe_status_key: Final = "AutoMLJobStatus"
@@ -101,28 +104,6 @@ class CompilationJobInfo(SageMakerJobInfo):
  processor = "process_compilation_job"
 
 
-class EdgePackagingJobInfo(SageMakerJobInfo):
- # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_edge_packaging_jobs
- list_command: Final = "list_edge_packaging_jobs"
- list_key: Final = "EdgePackagingJobSummaries"
- list_name_key: Final = "EdgePackagingJobName"
- list_arn_key: Final = "EdgePackagingJobArn"
- # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_edge_packaging_job
- describe_command: Final = "describe_edge_packaging_job"
- describe_name_key: Final = "EdgePackagingJobName"
- describe_arn_key: Final = "EdgePackagingJobArn"
- describe_status_key: Final = "EdgePackagingJobStatus"
- status_map = {
- "INPROGRESS": JobStatusClass.IN_PROGRESS,
- "COMPLETED": JobStatusClass.COMPLETED,
- "FAILED": JobStatusClass.FAILED,
- "STARTING": JobStatusClass.STARTING,
- "STOPPING": JobStatusClass.STOPPING,
- "STOPPED": JobStatusClass.STOPPED,
- }
- processor = "process_edge_packaging_job"
-
-
 class HyperParameterTuningJobInfo(SageMakerJobInfo):
  # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_hyper_parameter_tuning_jobs
  list_command: Final = "list_hyper_parameter_tuning_jobs"

metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/jobs.py

@@ -26,7 +26,6 @@
 from datahub.ingestion.source.aws.sagemaker_processors.job_classes import (
  AutoMlJobInfo,
  CompilationJobInfo,
- EdgePackagingJobInfo,
  HyperParameterTuningJobInfo,
  LabelingJobInfo,
  ProcessingJobInfo,
@@ -53,7 +52,6 @@
  "JobInfo",
  AutoMlJobInfo,
  CompilationJobInfo,
- EdgePackagingJobInfo,
  HyperParameterTuningJobInfo,
  LabelingJobInfo,
  ProcessingJobInfo,
@@ -65,7 +63,6 @@
 class JobType(Enum):
  AUTO_ML = "auto_ml"
  COMPILATION = "compilation"
- EDGE_PACKAGING = "edge_packaging"
  HYPER_PARAMETER_TUNING = "hyper_parameter_tuning"
  LABELING = "labeling"
  PROCESSING = "processing"
@@ -78,7 +75,6 @@ class JobType(Enum):
 job_type_to_info: Mapping[JobType, Any] = {
  JobType.AUTO_ML: AutoMlJobInfo(),
  JobType.COMPILATION: CompilationJobInfo(),
- JobType.EDGE_PACKAGING: EdgePackagingJobInfo(),
  JobType.HYPER_PARAMETER_TUNING: HyperParameterTuningJobInfo(),
  JobType.LABELING: LabelingJobInfo(),
  JobType.PROCESSING: ProcessingJobInfo(),
@@ -416,23 +412,20 @@ def process_auto_ml_job(self, job: Dict[str, Any]) -> SageMakerJob:
  """
  Process outputs from Boto3 describe_auto_ml_job()
 
- See https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_auto_ml_job
+ See https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/describe_auto_ml_job_v2.html
  """
 
  JOB_TYPE = JobType.AUTO_ML
 
  input_datasets = {}
-
- for input_config in job.get("InputDataConfig", []):
+ for input_config in job.get("AutoMLJobInputDataConfig", []):
  input_data = input_config.get("DataSource", {}).get("S3DataSource")
-
  if input_data is not None and "S3Uri" in input_data:
  input_datasets[make_s3_urn(input_data["S3Uri"], self.env)] = {
  "dataset_type": "s3",
  "uri": input_data["S3Uri"],
  "datatype": input_data.get("S3DataType"),
  }
-
  output_datasets = {}
 
  output_s3_path = job.get("OutputDataConfig", {}).get("S3OutputPath")
@@ -448,6 +441,18 @@ def process_auto_ml_job(self, job: Dict[str, Any]) -> SageMakerJob:
  JOB_TYPE,
  )
 
+ metrics: Dict[str, Any] = {}
+ # Get job metrics from CandidateMetrics
+ candidate_metrics = (
+ job.get("BestCandidate", {})
+ .get("CandidateProperties", {})
+ .get("CandidateMetrics", [])
+ )
+ if candidate_metrics:
+ metrics = {
+ metric["MetricName"]: metric["Value"] for metric in candidate_metrics
+ }
+
  model_containers = job.get("BestCandidate", {}).get("InferenceContainers", [])
 
  for model_container in model_containers:
@@ -456,7 +461,7 @@ def process_auto_ml_job(self, job: Dict[str, Any]) -> SageMakerJob:
  if model_data_url is not None:
  job_key = JobKey(job_snapshot.urn, JobDirection.TRAINING)
 
- self.update_model_image_jobs(model_data_url, job_key)
+ self.update_model_image_jobs(model_data_url, job_key, metrics=metrics)
 
  return SageMakerJob(
  job_name=job_name,
@@ -515,83 +520,6 @@ def process_compilation_job(self, job: Dict[str, Any]) -> SageMakerJob:
  output_datasets=output_datasets,
  )
 
- def process_edge_packaging_job(
- self,
- job: Dict[str, Any],
- ) -> SageMakerJob:
- """
- Process outputs from Boto3 describe_edge_packaging_job()
-
- See https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_edge_packaging_job
- """
-
- JOB_TYPE = JobType.EDGE_PACKAGING
-
- name: str = job["EdgePackagingJobName"]
- arn: str = job["EdgePackagingJobArn"]
-
- output_datasets = {}
-
- model_artifact_s3_uri: Optional[str] = job.get("ModelArtifact")
- output_s3_uri: Optional[str] = job.get("OutputConfig", {}).get(
- "S3OutputLocation"
- )
-
- if model_artifact_s3_uri is not None:
- output_datasets[make_s3_urn(model_artifact_s3_uri, self.env)] = {
- "dataset_type": "s3",
- "uri": model_artifact_s3_uri,
- }
-
- if output_s3_uri is not None:
- output_datasets[make_s3_urn(output_s3_uri, self.env)] = {
- "dataset_type": "s3",
- "uri": output_s3_uri,
- }
-
- # from docs: "The name of the SageMaker Neo compilation job that is used to locate model artifacts that are being packaged."
- compilation_job_name: Optional[str] = job.get("CompilationJobName")
-
- output_jobs = set()
- if compilation_job_name is not None:
- # globally unique job name
- full_job_name = ("compilation", compilation_job_name)
-
- if full_job_name in self.name_to_arn:
- output_jobs.add(
- make_sagemaker_job_urn(
- "compilation",
- compilation_job_name,
- self.name_to_arn[full_job_name],
- self.env,
- )
- )
- else:
- self.report.report_warning(
- name,
- f"Unable to find ARN for compilation job {compilation_job_name} produced by edge packaging job {arn}",
- )
-
- job_snapshot, job_name, job_arn = self.create_common_job_snapshot(
- job,
- JOB_TYPE,
- f"https://{self.aws_region}.console.aws.amazon.com/sagemaker/home?region={self.aws_region}#/edge-packaging-jobs/{job['EdgePackagingJobName']}",
- )
-
- if job.get("ModelName") is not None:
- job_key = JobKey(job_snapshot.urn, JobDirection.DOWNSTREAM)
-
- self.update_model_name_jobs(job["ModelName"], job_key)
-
- return SageMakerJob(
- job_name=job_name,
- job_arn=job_arn,
- job_type=JOB_TYPE,
- job_snapshot=job_snapshot,
- output_datasets=output_datasets,
- output_jobs=output_jobs,
- )
-
  def process_hyper_parameter_tuning_job(
  self,
  job: Dict[str, Any],