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
"""
JOB_TYPE = "auto_ml"
input_datasets = {}
for input_config in job.get("InputDataConfig", []):
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")
if output_s3_path is not None:
output_datasets[make_s3_urn(output_s3_path, self.env)] = {
"dataset_type": "s3",
"uri": output_s3_path,
}
job_snapshot, job_name, job_arn = self.create_common_job_snapshot(
job,
JOB_TYPE,
)
model_containers = job.get("BestCandidate", {}).get("InferenceContainers", [])
for model_container in model_containers:
model_data_url = model_container.get("ModelDataUrl")
if model_data_url is not None:
self.model_data_to_jobs[model_data_url].add(
ModelJob(
job_urn=job_snapshot.urn, job_direction=JobDirection.TRAINING
)
)
return SageMakerJob(
job_name=job_name,
job_arn=job_arn,
job_type=JOB_TYPE,
job_snapshot=job_snapshot,
input_datasets=input_datasets,
output_datasets=output_datasets,
)
def process_compilation_job(self, job: Dict[str, Any]) -> SageMakerJob:
"""
Process outputs from Boto3 describe_compilation_job()
See https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_compilation_job
"""
JOB_TYPE = "compilation"
input_datasets = {}
input_data: Optional[Dict[str, Any]] = job.get("InputConfig")
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"],
"framework": input_data.get("Framework"),
"framework_version": input_data.get("FrameworkVersion"),
}
output_datasets = {}
output_data: Optional[Dict[str, Any]] = job.get("OutputConfig")
if output_data is not None and "S3OutputLocation" in output_data:
output_datasets[make_s3_urn(output_data["S3OutputLocation"], self.env)] = {
"dataset_type": "s3",
"uri": output_data["S3OutputLocation"],
"target_device": output_data.get("TargetDevice"),
"target_platform": output_data.get("TargetPlatform"),
}
job_snapshot, job_name, job_arn = self.create_common_job_snapshot(
job,
JOB_TYPE,
)
return SageMakerJob(
job_name=job_name,
job_arn=job_arn,
job_type=JOB_TYPE,
job_snapshot=job_snapshot,
input_datasets=input_datasets,
output_datasets=output_datasets,
)
def process_dataflow_node(
self,
node: Dict[str, Any],
flow_urn: str,
new_dataset_ids: List[str],
new_dataset_mces: List[MetadataChangeEvent],
s3_formats: typing.DefaultDict[str, Set[Union[str, None]]],
) -> Dict[str, Any]:
node_type = node["NodeType"]
# for nodes representing datasets, we construct a dataset URN accordingly
if node_type in ["DataSource", "DataSink"]:
node_args = {x["Name"]: json.loads(x["Value"]) for x in node["Args"]}
# if data object is Glue table
if "database" in node_args and "table_name" in node_args:
full_table_name = f"{node_args['database']}.{node_args['table_name']}"
# we know that the table will already be covered when ingesting Glue tables
node_urn = f"urn:li:dataset:(urn:li:dataPlatform:glue,{full_table_name},{self.env})"
# if data object is S3 bucket
elif node_args.get("connection_type") == "s3":
s3_uri = node_args["connection_options"]["path"]
# append S3 format if different ones exist
if len(s3_formats[s3_uri]) > 1:
node_urn = make_s3_urn(
s3_uri,
self.env,
suffix=node_args.get("format"),
)
else:
node_urn = make_s3_urn(s3_uri, self.env)
dataset_snapshot = DatasetSnapshot(
urn=node_urn,
aspects=[],
)
dataset_snapshot.aspects.append(Status(removed=False))
dataset_snapshot.aspects.append(
DatasetPropertiesClass(
customProperties={k: str(v) for k, v in node_args.items()},
tags=[],
)
)
new_dataset_mces.append(
MetadataChangeEvent(proposedSnapshot=dataset_snapshot)
)
new_dataset_ids.append(f"{node['NodeType']}-{node['Id']}")
else:
raise ValueError(f"Unrecognized Glue data object type: {node_args}")
# otherwise, a node represents a transformation
else:
node_urn = mce_builder.make_data_job_urn_with_flow(
flow_urn, job_id=f'{node["NodeType"]}-{node["Id"]}'
)
return {
**node,
"urn": node_urn,
# to be filled in after traversing edges
"inputDatajobs": [],
"inputDatasets": [],
"outputDatasets": [],
}
def process_transform_job(self, job: Dict[str, Any]) -> SageMakerJob:
"""
Process outputs from Boto3 describe_transform_job()
See https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_transform_job
"""
JOB_TYPE = "transform"
job_input = job.get("TransformInput", {})
input_s3 = job_input.get("DataSource", {}).get("S3DataSource", {})
input_s3_uri = input_s3.get("S3Uri")
input_datasets = {}
if input_s3_uri is not None:
input_datasets[make_s3_urn(input_s3_uri, self.env)] = {
"dataset_type": "s3",
"uri": input_s3_uri,
"datatype": input_s3.get("S3DataType"),
"compression": job_input.get("CompressionType"),
"split": job_input.get("SplitType"),
}
output_datasets = {}
output_s3_uri = job.get("TransformOutput", {}).get("S3OutputPath")
if output_s3_uri is not None:
output_datasets[make_s3_urn(output_s3_uri, self.env)] = {
"dataset_type": "s3",
"uri": output_s3_uri,
}
labeling_arn = job.get("LabelingJobArn")
auto_ml_arn = job.get("AutoMLJobArn")
input_jobs = set()
if labeling_arn is not None:
labeling_type, labeling_name = self.arn_to_name.get(
labeling_arn, (None, None)
)
if labeling_type is not None and labeling_name is not None:
input_jobs.add(
make_sagemaker_job_urn(
labeling_type, labeling_name, labeling_arn, self.env
)
)
if auto_ml_arn is not None:
auto_ml_type, auto_ml_name = self.arn_to_name.get(auto_ml_arn, (None, None))
if auto_ml_type is not None and auto_ml_name is not None:
input_jobs.add(
make_sagemaker_job_urn(
auto_ml_type, auto_ml_name, auto_ml_arn, self.env
)
)
job_snapshot, job_name, job_arn = self.create_common_job_snapshot(
job,
JOB_TYPE,
)
if job.get("ModelName") is not None:
self.model_name_to_jobs[job["ModelName"]].add(
ModelJob(
job_urn=job_snapshot.urn, job_direction=JobDirection.DOWNSTREAM
)
)
return SageMakerJob(
job_name=job_name,
job_arn=job_arn,
job_type=JOB_TYPE,
job_snapshot=job_snapshot,
input_datasets=input_datasets,
output_datasets=output_datasets,
input_jobs=input_jobs,
)
def process_training_job(self, job: Dict[str, Any]) -> SageMakerJob:
"""
Process outputs from Boto3 describe_training_job()
See https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_training_job
"""
JOB_TYPE = "training"
input_datasets = {}
input_data_configs = job.get("InputDataConfig", [])
for config in input_data_configs:
data_source = config.get("DataSource", {})
s3_source = data_source.get("S3DataSource", {})
s3_uri = s3_source.get("S3Uri")
if s3_uri is not None:
input_datasets[make_s3_urn(s3_uri, self.env)] = {
"dataset_type": "s3",
"uri": s3_uri,
"datatype": s3_source.get("S3Datatype"),
"distribution_type": s3_source.get("S3DataDistributionType"),
"attribute_names": s3_source.get("AttributeNames"),
"channel_name": config.get("ChannelName"),
}
output_s3_uri = job.get("OutputDataConfig", {}).get("S3OutputPath")
checkpoint_s3_uri = job.get("CheckpointConfig", {}).get("S3Uri")
debug_s3_path = job.get("DebugHookConfig", {}).get("S3OutputPath")
tensorboard_output_path = job.get("TensorBoardOutputConfig", {}).get(
"S3OutputPath"
)
profiler_output_path = job.get("ProfilerConfig", {}).get("S3OutputPath")
debug_rule_configs = job.get("DebugRuleConfigurations", [])
processed_debug_configs = [
config.get("S3OutputPath") for config in debug_rule_configs
]
profiler_rule_configs = job.get("ProfilerRuleConfigurations", [])
processed_profiler_configs = [
config.get("S3OutputPath") for config in profiler_rule_configs
]
output_datasets = {}
# process all output datasets at once
for output_s3_uri in [
output_s3_uri,
checkpoint_s3_uri,
debug_s3_path,
tensorboard_output_path,
profiler_output_path,
*processed_debug_configs,
*processed_profiler_configs,
]:
if output_s3_uri is not None:
output_datasets[make_s3_urn(output_s3_uri, self.env)] = {
"dataset_type": "s3",
"uri": output_s3_uri,
}
job_snapshot, job_name, job_arn = self.create_common_job_snapshot(
job,
JOB_TYPE,
)
model_data_url = job.get("ModelArtifacts", {}).get("S3ModelArtifacts")
if model_data_url is not None:
self.model_data_to_jobs[model_data_url].add(
ModelJob(job_urn=job_snapshot.urn, job_direction=JobDirection.TRAINING)
)
return SageMakerJob(
job_name=job_name,
job_arn=job_arn,
job_type=JOB_TYPE,
job_snapshot=job_snapshot,
input_datasets=input_datasets,
output_datasets=output_datasets,
)
def process_processing_job(self, job: Dict[str, Any]) -> SageMakerJob:
"""
Process outputs from Boto3 describe_processing_job()
See https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_processing_job
"""
JOB_TYPE = "processing"
input_jobs = set()
auto_ml_arn: Optional[str] = job.get("AutoMLJobArn")
training_arn: Optional[str] = job.get("TrainingJobArn")
if auto_ml_arn is not None:
auto_ml_type, auto_ml_name = self.arn_to_name.get(auto_ml_arn, (None, None))
if auto_ml_type is not None and auto_ml_name is not None:
input_jobs.add(
make_sagemaker_job_urn(
auto_ml_type, auto_ml_name, auto_ml_arn, self.env
)
)
if training_arn is not None:
training_type, training_name = self.arn_to_name.get(
training_arn, (None, None)
)
if training_type is not None and training_name is not None:
input_jobs.add(
make_sagemaker_job_urn(
training_type, training_name, training_arn, self.env
)
)
input_datasets = {}
inputs = job["ProcessingInputs"]
for input_config in inputs:
input_name = input_config["InputName"]
input_s3 = input_config.get("S3Input", {})
input_s3_uri = input_s3.get("S3Uri")
if input_s3_uri is not None:
input_datasets[make_s3_urn(input_s3_uri, self.env)] = {
"dataset_type": "s3",
"uri": input_s3_uri,
"datatype": input_s3.get("S3DataType"),
"mode": input_s3.get("S3InputMode"),
"distribution_type": input_s3.get("S3DataDistributionType"),
"compression": input_s3.get("S3CompressionType"),
"name": input_name,
}
# TODO: ingest Athena and Redshift data sources
# We don't do this at the moment because we need to parse the QueryString SQL
# in order to get the tables used (otherwise we just have databases)
# input_athena = input_config.get("DatasetDefinition", {}).get(
# "AthenaDatasetDefinition", {}
# )
# input_redshift = input_config.get("DatasetDefinition", {}).get(
# "RedshiftDatasetDefinition", {}
# )
outputs: List[Dict[str, Any]] = job.get("ProcessingOutputConfig", {}).get(
"Outputs", []
)
output_datasets = {}
for output in outputs:
output_name = output["OutputName"]
output_s3_uri = output.get("S3Output", {}).get("S3Uri")
if output_s3_uri is not None:
output_datasets[make_s3_urn(output_s3_uri, self.env)] = {
"dataset_type": "s3",
"uri": output_s3_uri,
"name": output_name,
}
output_feature_group = output.get("FeatureStoreOutput", {}).get(
"FeatureGroupName"
)
if output_feature_group is not None:
output_datasets[
mce_builder.make_ml_feature_table_urn(
"sagemaker", output_feature_group
)
] = {
"dataset_type": "sagemaker_feature_group",
}
job_snapshot, job_name, job_arn = self.create_common_job_snapshot(
job,
JOB_TYPE,
)
return SageMakerJob(
job_name=job_name,
job_arn=job_arn,
job_type=JOB_TYPE,
job_snapshot=job_snapshot,
input_datasets=input_datasets,
input_jobs=input_jobs,
)
def process_labeling_job(self, job: Dict[str, Any]) -> SageMakerJob:
"""
Process outputs from Boto3 describe_labeling_job()
See https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_labeling_job
"""
JOB_TYPE = "labeling"
input_datasets = {}
input_s3_uri: Optional[str] = (
job.get("InputConfig", {})
.get("DataSource", {})
.get("S3DataSource", {})
.get("ManifestS3Uri")
)
if input_s3_uri is not None:
input_datasets[make_s3_urn(input_s3_uri, self.env)] = {
"dataset_type": "s3",
"uri": input_s3_uri,
}
category_config_s3_uri: Optional[str] = job.get("LabelCategoryConfigS3Uri")
if category_config_s3_uri is not None:
input_datasets[make_s3_urn(category_config_s3_uri, self.env)] = {
"dataset_type": "s3",
"uri": category_config_s3_uri,
}
output_datasets = {}
output_s3_uri: Optional[str] = job.get("LabelingJobOutput", {}).get(
"OutputDatasetS3Uri"
)
if output_s3_uri is not None:
output_datasets[make_s3_urn(output_s3_uri, self.env)] = {
"dataset_type": "s3",
"uri": output_s3_uri,
}
output_config_s3_uri: Optional[str] = job.get("OutputConfig", {}).get(
"S3OutputPath"
)
if output_config_s3_uri is not None:
output_datasets[make_s3_urn(output_config_s3_uri, self.env)] = {
"dataset_type": "s3",
"uri": output_config_s3_uri,
}
job_snapshot, job_name, job_arn = self.create_common_job_snapshot(
job,
JOB_TYPE,
)
return SageMakerJob(
job_name=job_name,
job_arn=job_arn,
job_type=JOB_TYPE,
job_snapshot=job_snapshot,
input_datasets=input_datasets,
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 = "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,
)
if job.get("ModelName") is not None:
self.model_name_to_jobs[job["ModelName"]].add(
ModelJob(
job_urn=job_snapshot.urn, job_direction=JobDirection.DOWNSTREAM
)
)
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,
)