You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "mobuchowski (via GitHub)" <gi...@apache.org> on 2023/07/19 13:20:31 UTC
[GitHub] [airflow] mobuchowski opened a new pull request, #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
mobuchowski opened a new pull request, #31816:
URL: https://github.com/apache/airflow/pull/31816
This PR adds support for OpenLineage methods to SageMaker's Processing, Transform and Training operators.
It gets s3 dataset data from API responses saved by execute methods, then processed by openlineage methods.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on PR #31816:
URL: https://github.com/apache/airflow/pull/31816#issuecomment-1587419187
@vandonr-amz this is meant to be used in [OpenLineage integration](https://cwiki.apache.org/confluence/display/AIRFLOW/AIP-53+OpenLineage+in+Airflow
) - in video form https://www.youtube.com/watch?v=fAqvoMzz7Tk 🙂 )
In short, the idea is that operators are tracking their own lineage; captured lineage is then send in common format to configured backend.
To do that, we add methods that will be called by OpenLineage plugin, if it's installed or configured. It's not going to be called from user code, whether that's DAG or anything else.
>Also, rather than making this a method that should only be called after execute has set self.processing_job, which is a bit brittle, why not make the dependency explicit by expecting the serialized job as a parameter ? Then it wouldn't have to be stored in the class at all.
This could also be a class method, the only usage of self after that is the log.
In the integration, we can't depend on knowing what parameters particular operator requires. This would mean something external had to know what the method expect, and follow up when it changes. This is much more brittle because it separates responsibility between multiple separately versioned components; here operator always can report it's own lineage, and we can check it by tests. If the contract changes, tests would fail.
It was our previous approach and we moved to in-operator approach for this reason. https://github.com/OpenLineage/OpenLineage/blob/main/integration/airflow/openlineage/airflow/extractors/sagemaker_extractors.py
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1274898747
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -218,6 +229,7 @@ def __init__(
self.max_attempts = max_attempts or 60
self.max_ingestion_time = max_ingestion_time
self.deferrable = deferrable
+ self.processing_job: dict[Any, Any] | None = None
Review Comment:
Done.
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -541,6 +592,7 @@ def __init__(
Provided value: '{action_if_job_exists}'."
)
self.deferrable = deferrable
+ self.transform_data: dict[Any, Any] | None = None
Review Comment:
Done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1284521275
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -158,6 +159,14 @@ def hook(self):
"""Return SageMakerHook."""
return SageMakerHook(aws_conn_id=self.aws_conn_id)
+ @staticmethod
+ def path_to_s3_dataset(path):
Review Comment:
@vandonr-amz you're nearly right - I've pushed but to wrong upstream. Neverthless it's here: https://github.com/apache/airflow/pull/33114
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1280789542
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -222,6 +230,7 @@ def __init__(
self.max_attempts = max_attempts or 60
self.max_ingestion_time = max_ingestion_time
self.deferrable = deferrable
+ self.serialized_job: dict[Any, Any] | None = None
Review Comment:
I've actually moved this to just `self.serialized_job: dict`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1280780363
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -647,9 +693,13 @@ def execute(self, context: Context) -> dict:
method_name="execute_complete",
)
+ self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
+ self.serialized_tranform = serialize(
+ self.hook.describe_transform_job(transform_config["TransformJobName"])
+ )
return {
- "Model": serialize(self.hook.describe_model(transform_config["ModelName"])),
- "Transform": serialize(self.hook.describe_transform_job(transform_config["TransformJobName"])),
+ "Model": self.serialized_model,
Review Comment:
Done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski merged pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski merged PR #31816:
URL: https://github.com/apache/airflow/pull/31816
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1267932384
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -275,14 +287,53 @@ def execute(self, context: Context) -> dict:
method_name="execute_complete",
)
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.processing_job = {
+ "Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ }
+ return self.processing_job
def execute_complete(self, context, event=None):
if event["status"] != "success":
raise AirflowException(f"Error while running job: {event}")
else:
self.log.info(event["message"])
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.processing_job = {
+ "Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ }
+ return self.processing_job
+
+ def get_openlineage_facets_on_complete(self, task_instance):
Review Comment:
Discussed this below - it's interface method.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on PR #31816:
URL: https://github.com/apache/airflow/pull/31816#issuecomment-1652316851
@vandonr-amz @ferruzzi would be great if you could look at it again 🙏
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1226703381
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -42,9 +43,19 @@
def serialize(result: dict) -> str:
+ logging.getLogger(__name__).error(result)
return json.loads(json.dumps(result, cls=AirflowJsonEncoder))
+def path_to_s3_dataset(path):
Review Comment:
Done.
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -42,9 +43,19 @@
def serialize(result: dict) -> str:
+ logging.getLogger(__name__).error(result)
Review Comment:
Removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1228411828
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -218,6 +229,7 @@ def __init__(
self.max_attempts = max_attempts or 60
self.max_ingestion_time = max_ingestion_time
self.deferrable = deferrable
+ self.processing_job: dict[Any, Any] | None = None
Review Comment:
Idea here was to have the same structure that is returned to XCom, maybe without real reason to.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] shubham22 commented on pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "shubham22 (via GitHub)" <gi...@apache.org>.
shubham22 commented on PR #31816:
URL: https://github.com/apache/airflow/pull/31816#issuecomment-1584966779
@ferruzzi @vandonr-amz - would be great if one of you folks can also review it, given your past work on SageMaker operators. TIA!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1226701784
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -42,9 +43,19 @@
def serialize(result: dict) -> str:
+ logging.getLogger(__name__).error(result)
return json.loads(json.dumps(result, cls=AirflowJsonEncoder))
+def path_to_s3_dataset(path):
+ from openlineage.client.run import Dataset
+
+ path = path.replace("s3://", "")
+ namespace = path.split("/")[0]
+ name = "/".join(path.split("/")[1:])
Review Comment:
Done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1281746857
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -158,6 +159,14 @@ def hook(self):
"""Return SageMakerHook."""
return SageMakerHook(aws_conn_id=self.aws_conn_id)
+ @staticmethod
+ def path_to_s3_dataset(path):
Review Comment:
Done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1280789542
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -222,6 +230,7 @@ def __init__(
self.max_attempts = max_attempts or 60
self.max_ingestion_time = max_ingestion_time
self.deferrable = deferrable
+ self.serialized_job: dict[Any, Any] | None = None
Review Comment:
I've actually moved this to just `self.serialized_job: dict`, as I consider this to be `lazy initialized` in this context.
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -222,6 +230,7 @@ def __init__(
self.max_attempts = max_attempts or 60
self.max_ingestion_time = max_ingestion_time
self.deferrable = deferrable
+ self.serialized_job: dict[Any, Any] | None = None
Review Comment:
I've actually moved this to just `self.serialized_job: dict`, as I consider this to be `lazy initialized` in this context - no mypy warnings in this case.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1231313122
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -218,6 +229,7 @@ def __init__(
self.max_attempts = max_attempts or 60
self.max_ingestion_time = max_ingestion_time
self.deferrable = deferrable
+ self.processing_job: dict[Any, Any] | None = None
Review Comment:
if you want a dict like this to be returned at the end, then you should create it where you return it, otherwise it just makes the code more confusing to read
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on PR #31816:
URL: https://github.com/apache/airflow/pull/31816#issuecomment-1593438833
hmm ok, wo what you're saying basically is that `get_openlineage_facets_on_complete` is an interface method's implementation ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski closed pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski closed pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
URL: https://github.com/apache/airflow/pull/31816
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1226732791
##########
tests/providers/amazon/aws/operators/test_sagemaker_processing.py:
##########
@@ -255,3 +257,31 @@ def test_operator_defer(self, mock_job_exists, mock_processing):
with pytest.raises(TaskDeferred) as exc:
sagemaker_operator.execute(context=None)
assert isinstance(exc.value.trigger, SageMakerTrigger), "Trigger is not a SagemakerTrigger"
+
+ @mock.patch.object(SageMakerHook, "describe_processing_job")
+ @mock.patch.object(SageMakerHook, "count_processing_jobs_by_name", return_value=0)
+ @mock.patch.object(
+ SageMakerHook,
+ "create_processing_job",
+ return_value={"ProcessingJobArn": "test_arn", "ResponseMetadata": {"HTTPStatusCode": 200}},
+ )
+ @mock.patch("airflow.providers.amazon.aws.operators.sagemaker.SageMakerBaseOperator._check_if_job_exists")
Review Comment:
Done.
##########
tests/providers/amazon/aws/operators/test_sagemaker_training.py:
##########
@@ -127,3 +129,26 @@ def test_operator_defer(self, mock_training):
with pytest.raises(TaskDeferred) as exc:
self.sagemaker.execute(context=None)
assert isinstance(exc.value.trigger, SageMakerTrigger), "Trigger is not a SagemakerTrigger"
+
+ @mock.patch.object(SageMakerHook, "describe_training_job")
+ @mock.patch.object(SageMakerHook, "create_training_job")
+ def test_execute_openlineage_data(self, mock_training, mock_desc):
+ mock_training.return_value = {
+ "TrainingJobArn": "test_arn",
+ "ResponseMetadata": {"HTTPStatusCode": 200},
+ }
+ mock_desc.return_value = {
+ "InputDataConfig": [
+ {
+ "DataSource": {"S3DataSource": {"S3Uri": "s3://input-bucket/input-path"}},
+ }
+ ],
+ "ModelArtifacts": {"S3ModelArtifacts": "s3://model-bucket/model-path"},
+ }
+ self.sagemaker.check_if_job_exists = False
+ self.sagemaker._check_if_job_exists = mock.MagicMock()
Review Comment:
Done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1226733444
##########
tests/providers/amazon/aws/operators/test_sagemaker_processing.py:
##########
@@ -255,3 +257,31 @@ def test_operator_defer(self, mock_job_exists, mock_processing):
with pytest.raises(TaskDeferred) as exc:
sagemaker_operator.execute(context=None)
assert isinstance(exc.value.trigger, SageMakerTrigger), "Trigger is not a SagemakerTrigger"
+
+ @mock.patch.object(SageMakerHook, "describe_processing_job")
+ @mock.patch.object(SageMakerHook, "count_processing_jobs_by_name", return_value=0)
+ @mock.patch.object(
+ SageMakerHook,
+ "create_processing_job",
+ return_value={"ProcessingJobArn": "test_arn", "ResponseMetadata": {"HTTPStatusCode": 200}},
+ )
+ @mock.patch("airflow.providers.amazon.aws.operators.sagemaker.SageMakerBaseOperator._check_if_job_exists")
+ def test_operator_openlineage_data(self, check_job_exists, mock_processing, _, mock_desc):
+ check_job_exists.return_value = False
+ mock_desc.return_value = {
+ "ProcessingInputs": [{"S3Input": {"S3Uri": "s3://input-bucket/input-path"}}],
+ "ProcessingOutputConfig": {
+ "Outputs": [{"S3Output": {"S3Uri": "s3://output-bucket/output-path"}}]
+ },
+ }
Review Comment:
Done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1280779995
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -279,14 +288,49 @@ def execute(self, context: Context) -> dict:
method_name="execute_complete",
)
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.serialized_job = serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ return {"Processing": self.serialized_job}
def execute_complete(self, context, event=None):
if event["status"] != "success":
raise AirflowException(f"Error while running job: {event}")
else:
self.log.info(event["message"])
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.serialized_job = serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ return {"Processing": self.serialized_job}
+
+ def get_openlineage_facets_on_complete(self, task_instance):
+ """Returns OpenLineage data gathered from SageMaker's API response saved by processing job."""
+ from airflow.providers.openlineage.extractors.base import OperatorLineage
+
+ inputs, outputs = [], []
Review Comment:
Done.
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -658,11 +708,66 @@ def execute_complete(self, context, event=None):
else:
self.log.info(event["message"])
transform_config = self.config.get("Transform", self.config)
+ self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
+ self.serialized_tranform = serialize(
+ self.hook.describe_transform_job(transform_config["TransformJobName"])
+ )
return {
- "Model": serialize(self.hook.describe_model(transform_config["ModelName"])),
- "Transform": serialize(self.hook.describe_transform_job(transform_config["TransformJobName"])),
+ "Model": self.serialized_model,
+ "Transform": self.serialized_tranform,
}
+ def get_openlineage_facets_on_complete(self, task_instance):
+ """Returns OpenLineage data gathered from SageMaker's API response saved by transform job."""
+ from airflow.providers.openlineage.extractors import OperatorLineage
+
+ model_package_arn = None
+ transform_input = None
+ transform_output = None
+
+ try:
+ model_package_arn = self.serialized_model["PrimaryContainer"]["ModelPackageName"]
+ except KeyError:
+ self.log.error("Cannot find Model Package Name.", exc_info=True)
+
+ try:
+ transform_input = self.serialized_tranform["TransformInput"]["DataSource"]["S3DataSource"][
+ "S3Uri"
+ ]
+ transform_output = self.serialized_tranform["TransformOutput"]["S3OutputPath"]
+ except KeyError:
+ self.log.error("Cannot find some required input/output details.", exc_info=True)
+
+ inputs = []
+
+ if transform_input is not None:
+ inputs.append(self.path_to_s3_dataset(transform_input))
+
+ if model_package_arn is not None:
+ model_data_urls = self._get_model_data_urls(model_package_arn)
+ for model_data_url in model_data_urls:
+ inputs.append(self.path_to_s3_dataset(model_data_url))
+
+ output = []
Review Comment:
Done.
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -279,14 +288,49 @@ def execute(self, context: Context) -> dict:
method_name="execute_complete",
)
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.serialized_job = serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ return {"Processing": self.serialized_job}
def execute_complete(self, context, event=None):
if event["status"] != "success":
raise AirflowException(f"Error while running job: {event}")
else:
self.log.info(event["message"])
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.serialized_job = serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ return {"Processing": self.serialized_job}
+
+ def get_openlineage_facets_on_complete(self, task_instance):
+ """Returns OpenLineage data gathered from SageMaker's API response saved by processing job."""
+ from airflow.providers.openlineage.extractors.base import OperatorLineage
+
+ inputs, outputs = [], []
+ try:
+ inputs, outputs = self._extract_s3_dataset_identifiers(
+ processing_inputs=self.serialized_job["ProcessingInputs"],
+ processing_outputs=self.serialized_job["ProcessingOutputConfig"]["Outputs"],
+ )
+ except KeyError:
+ self.log.exception("Could not find input/output information in Xcom.")
+
+ return OperatorLineage(
+ inputs=inputs,
+ outputs=outputs,
+ )
Review Comment:
Fixed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1280791684
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -658,11 +708,66 @@ def execute_complete(self, context, event=None):
else:
self.log.info(event["message"])
transform_config = self.config.get("Transform", self.config)
+ self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
+ self.serialized_tranform = serialize(
+ self.hook.describe_transform_job(transform_config["TransformJobName"])
+ )
return {
- "Model": serialize(self.hook.describe_model(transform_config["ModelName"])),
- "Transform": serialize(self.hook.describe_transform_job(transform_config["TransformJobName"])),
+ "Model": self.serialized_model,
+ "Transform": self.serialized_tranform,
}
+ def get_openlineage_facets_on_complete(self, task_instance):
+ """Returns OpenLineage data gathered from SageMaker's API response saved by transform job."""
+ from airflow.providers.openlineage.extractors import OperatorLineage
+
+ model_package_arn = None
+ transform_input = None
+ transform_output = None
+
+ try:
+ model_package_arn = self.serialized_model["PrimaryContainer"]["ModelPackageName"]
+ except KeyError:
+ self.log.error("Cannot find Model Package Name.", exc_info=True)
+
+ try:
+ transform_input = self.serialized_tranform["TransformInput"]["DataSource"]["S3DataSource"][
+ "S3Uri"
+ ]
+ transform_output = self.serialized_tranform["TransformOutput"]["S3OutputPath"]
+ except KeyError:
+ self.log.error("Cannot find some required input/output details.", exc_info=True)
+
+ inputs = []
+
+ if transform_input is not None:
+ inputs.append(self.path_to_s3_dataset(transform_input))
+
+ if model_package_arn is not None:
+ model_data_urls = self._get_model_data_urls(model_package_arn)
+ for model_data_url in model_data_urls:
+ inputs.append(self.path_to_s3_dataset(model_data_url))
+
+ output = []
+ if transform_output is not None:
+ output.append(self.path_to_s3_dataset(transform_output))
+
+ return OperatorLineage(inputs=inputs, outputs=output)
+
+ def _get_model_data_urls(self, model_package_arn):
+ model_data_urls = []
+ try:
+ model_containers = self.hook.get_conn().describe_model_package(
+ ModelPackageName=model_package_arn
+ )["InferenceSpecification"]["Containers"]
+
+ for container in model_containers:
+ model_data_urls.append(container["ModelDataUrl"])
+ except Exception:
Review Comment:
Done.
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -658,11 +708,66 @@ def execute_complete(self, context, event=None):
else:
self.log.info(event["message"])
transform_config = self.config.get("Transform", self.config)
+ self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
+ self.serialized_tranform = serialize(
+ self.hook.describe_transform_job(transform_config["TransformJobName"])
+ )
return {
- "Model": serialize(self.hook.describe_model(transform_config["ModelName"])),
- "Transform": serialize(self.hook.describe_transform_job(transform_config["TransformJobName"])),
+ "Model": self.serialized_model,
+ "Transform": self.serialized_tranform,
}
+ def get_openlineage_facets_on_complete(self, task_instance):
+ """Returns OpenLineage data gathered from SageMaker's API response saved by transform job."""
+ from airflow.providers.openlineage.extractors import OperatorLineage
+
+ model_package_arn = None
+ transform_input = None
+ transform_output = None
+
+ try:
+ model_package_arn = self.serialized_model["PrimaryContainer"]["ModelPackageName"]
+ except KeyError:
+ self.log.error("Cannot find Model Package Name.", exc_info=True)
+
+ try:
+ transform_input = self.serialized_tranform["TransformInput"]["DataSource"]["S3DataSource"][
+ "S3Uri"
+ ]
+ transform_output = self.serialized_tranform["TransformOutput"]["S3OutputPath"]
+ except KeyError:
+ self.log.error("Cannot find some required input/output details.", exc_info=True)
+
+ inputs = []
+
+ if transform_input is not None:
+ inputs.append(self.path_to_s3_dataset(transform_input))
+
+ if model_package_arn is not None:
+ model_data_urls = self._get_model_data_urls(model_package_arn)
+ for model_data_url in model_data_urls:
+ inputs.append(self.path_to_s3_dataset(model_data_url))
+
+ output = []
+ if transform_output is not None:
+ output.append(self.path_to_s3_dataset(transform_output))
+
+ return OperatorLineage(inputs=inputs, outputs=output)
+
+ def _get_model_data_urls(self, model_package_arn):
Review Comment:
Done.
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -576,6 +620,8 @@ def __init__(
Provided value: '{action_if_job_exists}'."
)
self.deferrable = deferrable
+ self.serialized_model: dict[Any, Any] | None = None
+ self.serialized_tranform: dict[Any, Any] | None = None
Review Comment:
The same ⬆️.
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -576,6 +620,8 @@ def __init__(
Provided value: '{action_if_job_exists}'."
)
self.deferrable = deferrable
+ self.serialized_model: dict[Any, Any] | None = None
+ self.serialized_tranform: dict[Any, Any] | None = None
Review Comment:
The same ⬆️.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1279713333
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -647,9 +693,13 @@ def execute(self, context: Context) -> dict:
method_name="execute_complete",
)
+ self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
+ self.serialized_tranform = serialize(
+ self.hook.describe_transform_job(transform_config["TransformJobName"])
+ )
return {
- "Model": serialize(self.hook.describe_model(transform_config["ModelName"])),
- "Transform": serialize(self.hook.describe_transform_job(transform_config["TransformJobName"])),
+ "Model": self.serialized_model,
Review Comment:
Nitpick: (here and below) For consistency I'd suggest you drop the trailing comma so this collapses to one line now that it can fit.
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -279,14 +288,49 @@ def execute(self, context: Context) -> dict:
method_name="execute_complete",
)
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.serialized_job = serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ return {"Processing": self.serialized_job}
def execute_complete(self, context, event=None):
if event["status"] != "success":
raise AirflowException(f"Error while running job: {event}")
else:
self.log.info(event["message"])
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.serialized_job = serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ return {"Processing": self.serialized_job}
+
+ def get_openlineage_facets_on_complete(self, task_instance):
+ """Returns OpenLineage data gathered from SageMaker's API response saved by processing job."""
+ from airflow.providers.openlineage.extractors.base import OperatorLineage
+
+ inputs, outputs = [], []
+ try:
+ inputs, outputs = self._extract_s3_dataset_identifiers(
+ processing_inputs=self.serialized_job["ProcessingInputs"],
+ processing_outputs=self.serialized_job["ProcessingOutputConfig"]["Outputs"],
+ )
+ except KeyError:
+ self.log.exception("Could not find input/output information in Xcom.")
+
+ return OperatorLineage(
+ inputs=inputs,
+ outputs=outputs,
+ )
Review Comment:
Dropping the trailing comma lets that stay on one line as you did elsewhere.
```suggestion
return OperatorLineage(inputs=inputs, outputs=outputs)
```
##########
tests/providers/amazon/aws/operators/test_sagemaker_processing.py:
##########
@@ -238,14 +243,16 @@ def test_action_if_job_exists_validation(self, mock_client):
action_if_job_exists="not_fail_or_increment",
)
- @mock.patch.object(SageMakerHook, "create_processing_job")
- @mock.patch("airflow.providers.amazon.aws.operators.sagemaker.SageMakerBaseOperator._check_if_job_exists")
- def test_operator_defer(self, mock_job_exists, mock_processing):
- mock_processing.return_value = {
+ @mock.patch.object(
+ SageMakerHook,
+ "create_processing_job",
+ return_value={
"ProcessingJobArn": "test_arn",
"ResponseMetadata": {"HTTPStatusCode": 200},
- }
- mock_job_exists.return_value = False
+ },
+ )
+ @mock.patch.object(SageMakerBaseOperator, "_check_if_job_exists", return_value=False)
Review Comment:
These are very nice, thanks for cleaning them up.
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -279,14 +288,49 @@ def execute(self, context: Context) -> dict:
method_name="execute_complete",
)
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.serialized_job = serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ return {"Processing": self.serialized_job}
def execute_complete(self, context, event=None):
if event["status"] != "success":
raise AirflowException(f"Error while running job: {event}")
else:
self.log.info(event["message"])
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.serialized_job = serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ return {"Processing": self.serialized_job}
+
+ def get_openlineage_facets_on_complete(self, task_instance):
+ """Returns OpenLineage data gathered from SageMaker's API response saved by processing job."""
+ from airflow.providers.openlineage.extractors.base import OperatorLineage
+
+ inputs, outputs = [], []
Review Comment:
Non-blocking nitpick: I know this does work in python, but I don't think I've seen it used elsewhere in the airflow code; it seems to be a community style preference to put each declaration on a new line.
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -658,11 +708,66 @@ def execute_complete(self, context, event=None):
else:
self.log.info(event["message"])
transform_config = self.config.get("Transform", self.config)
+ self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
+ self.serialized_tranform = serialize(
+ self.hook.describe_transform_job(transform_config["TransformJobName"])
+ )
return {
- "Model": serialize(self.hook.describe_model(transform_config["ModelName"])),
- "Transform": serialize(self.hook.describe_transform_job(transform_config["TransformJobName"])),
+ "Model": self.serialized_model,
+ "Transform": self.serialized_tranform,
}
+ def get_openlineage_facets_on_complete(self, task_instance):
+ """Returns OpenLineage data gathered from SageMaker's API response saved by transform job."""
+ from airflow.providers.openlineage.extractors import OperatorLineage
+
+ model_package_arn = None
+ transform_input = None
+ transform_output = None
+
+ try:
+ model_package_arn = self.serialized_model["PrimaryContainer"]["ModelPackageName"]
+ except KeyError:
+ self.log.error("Cannot find Model Package Name.", exc_info=True)
+
+ try:
+ transform_input = self.serialized_tranform["TransformInput"]["DataSource"]["S3DataSource"][
+ "S3Uri"
+ ]
+ transform_output = self.serialized_tranform["TransformOutput"]["S3OutputPath"]
+ except KeyError:
+ self.log.error("Cannot find some required input/output details.", exc_info=True)
+
+ inputs = []
+
+ if transform_input is not None:
+ inputs.append(self.path_to_s3_dataset(transform_input))
+
+ if model_package_arn is not None:
+ model_data_urls = self._get_model_data_urls(model_package_arn)
+ for model_data_url in model_data_urls:
+ inputs.append(self.path_to_s3_dataset(model_data_url))
+
+ output = []
Review Comment:
Nitpick: you used `outputs` above and generally the name of a list is pluralized as a hint. May want to do that here as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1226669410
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -42,9 +43,19 @@
def serialize(result: dict) -> str:
+ logging.getLogger(__name__).error(result)
return json.loads(json.dumps(result, cls=AirflowJsonEncoder))
+def path_to_s3_dataset(path):
+ from openlineage.client.run import Dataset
Review Comment:
We don't want to require people to have OpenLineage installed to use Sagemaker, or AWS provider - importing this locally helps us avoid that.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1224861901
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -42,9 +43,19 @@
def serialize(result: dict) -> str:
+ logging.getLogger(__name__).error(result)
Review Comment:
is this a personal debug log that you forgot to remove ?
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -275,14 +287,53 @@ def execute(self, context: Context) -> dict:
method_name="execute_complete",
)
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.processing_job = {
+ "Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ }
+ return self.processing_job
def execute_complete(self, context, event=None):
if event["status"] != "success":
raise AirflowException(f"Error while running job: {event}")
else:
self.log.info(event["message"])
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.processing_job = {
+ "Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ }
+ return self.processing_job
+
+ def get_openlineage_facets_on_complete(self, task_instance):
Review Comment:
the `task_instance` param seems to be unused ? (I didn't look at the static checks, but they may report this too ?)
Also, rather than making this a method that should only be called after `execute` has set `self.processing_job`, which is a bit brittle, why not make the dependency explicit by expecting the serialized job as a parameter ? Then it wouldn't have to be stored in the class at all.
This could also be a class method, the only usage of `self` after that is the log.
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -218,6 +229,7 @@ def __init__(
self.max_attempts = max_attempts or 60
self.max_ingestion_time = max_ingestion_time
self.deferrable = deferrable
+ self.processing_job: dict[Any, Any] | None = None
Review Comment:
I don't understand why this needs to be a dict when it seems to only ever contain one key (`"Processing"`)
couldn't you replace this with a `serialized_job` ?
##########
tests/providers/amazon/aws/operators/test_sagemaker_training.py:
##########
@@ -127,3 +129,26 @@ def test_operator_defer(self, mock_training):
with pytest.raises(TaskDeferred) as exc:
self.sagemaker.execute(context=None)
assert isinstance(exc.value.trigger, SageMakerTrigger), "Trigger is not a SagemakerTrigger"
+
+ @mock.patch.object(SageMakerHook, "describe_training_job")
+ @mock.patch.object(SageMakerHook, "create_training_job")
+ def test_execute_openlineage_data(self, mock_training, mock_desc):
+ mock_training.return_value = {
+ "TrainingJobArn": "test_arn",
+ "ResponseMetadata": {"HTTPStatusCode": 200},
+ }
+ mock_desc.return_value = {
+ "InputDataConfig": [
+ {
+ "DataSource": {"S3DataSource": {"S3Uri": "s3://input-bucket/input-path"}},
+ }
+ ],
+ "ModelArtifacts": {"S3ModelArtifacts": "s3://model-bucket/model-path"},
+ }
+ self.sagemaker.check_if_job_exists = False
+ self.sagemaker._check_if_job_exists = mock.MagicMock()
Review Comment:
you could set this up in an annotation like the other 2
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -541,6 +592,7 @@ def __init__(
Provided value: '{action_if_job_exists}'."
)
self.deferrable = deferrable
+ self.transform_data: dict[Any, Any] | None = None
Review Comment:
same remarks as above
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -42,9 +43,19 @@
def serialize(result: dict) -> str:
+ logging.getLogger(__name__).error(result)
return json.loads(json.dumps(result, cls=AirflowJsonEncoder))
+def path_to_s3_dataset(path):
+ from openlineage.client.run import Dataset
+
+ path = path.replace("s3://", "")
+ namespace = path.split("/")[0]
+ name = "/".join(path.split("/")[1:])
Review Comment:
you could do the split only once where to avoid repeating the work.
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -42,9 +43,19 @@
def serialize(result: dict) -> str:
+ logging.getLogger(__name__).error(result)
return json.loads(json.dumps(result, cls=AirflowJsonEncoder))
+def path_to_s3_dataset(path):
Review Comment:
you could tuck this as a static method in SageMakerBaseOperator
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -275,14 +287,53 @@ def execute(self, context: Context) -> dict:
method_name="execute_complete",
)
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.processing_job = {
+ "Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ }
+ return self.processing_job
def execute_complete(self, context, event=None):
if event["status"] != "success":
raise AirflowException(f"Error while running job: {event}")
else:
self.log.info(event["message"])
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.processing_job = {
+ "Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ }
+ return self.processing_job
+
+ def get_openlineage_facets_on_complete(self, task_instance):
+ """Returns OpenLineage data gathered from SageMaker's API response saved by processing job."""
+ from airflow.providers.openlineage.extractors.base import OperatorLineage
+
+ inputs, outputs = [], []
+ try:
+ inputs, outputs = self._get_s3_datasets(
+ processing_inputs=self.processing_job["Processing"]["ProcessingInputs"],
+ processing_outputs=self.processing_job["Processing"]["ProcessingOutputConfig"]["Outputs"],
+ )
+ except KeyError:
+ self.log.exception("Could not find input/output information in Xcom.")
+
+ return OperatorLineage(
+ inputs=inputs,
+ outputs=outputs,
+ )
+
+ def _get_s3_datasets(self, processing_inputs, processing_outputs):
Review Comment:
You might want to rename this method. With this name, it sounds like this is doing S3 queries, when all it's doing it building objects to be used to get the data later on.
##########
tests/providers/amazon/aws/operators/test_sagemaker_processing.py:
##########
@@ -255,3 +257,31 @@ def test_operator_defer(self, mock_job_exists, mock_processing):
with pytest.raises(TaskDeferred) as exc:
sagemaker_operator.execute(context=None)
assert isinstance(exc.value.trigger, SageMakerTrigger), "Trigger is not a SagemakerTrigger"
+
+ @mock.patch.object(SageMakerHook, "describe_processing_job")
+ @mock.patch.object(SageMakerHook, "count_processing_jobs_by_name", return_value=0)
+ @mock.patch.object(
+ SageMakerHook,
+ "create_processing_job",
+ return_value={"ProcessingJobArn": "test_arn", "ResponseMetadata": {"HTTPStatusCode": 200}},
+ )
+ @mock.patch("airflow.providers.amazon.aws.operators.sagemaker.SageMakerBaseOperator._check_if_job_exists")
Review Comment:
you could use `@mock.patch.object` here as well
##########
tests/providers/amazon/aws/operators/test_sagemaker_processing.py:
##########
@@ -255,3 +257,31 @@ def test_operator_defer(self, mock_job_exists, mock_processing):
with pytest.raises(TaskDeferred) as exc:
sagemaker_operator.execute(context=None)
assert isinstance(exc.value.trigger, SageMakerTrigger), "Trigger is not a SagemakerTrigger"
+
+ @mock.patch.object(SageMakerHook, "describe_processing_job")
+ @mock.patch.object(SageMakerHook, "count_processing_jobs_by_name", return_value=0)
+ @mock.patch.object(
+ SageMakerHook,
+ "create_processing_job",
+ return_value={"ProcessingJobArn": "test_arn", "ResponseMetadata": {"HTTPStatusCode": 200}},
+ )
+ @mock.patch("airflow.providers.amazon.aws.operators.sagemaker.SageMakerBaseOperator._check_if_job_exists")
+ def test_operator_openlineage_data(self, check_job_exists, mock_processing, _, mock_desc):
+ check_job_exists.return_value = False
+ mock_desc.return_value = {
+ "ProcessingInputs": [{"S3Input": {"S3Uri": "s3://input-bucket/input-path"}}],
+ "ProcessingOutputConfig": {
+ "Outputs": [{"S3Output": {"S3Uri": "s3://output-bucket/output-path"}}]
+ },
+ }
Review Comment:
for consistency, you could put this in the annotation, or put the other returns values in the test body, but mixing the setup styles makes it harder to read
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on PR #31816:
URL: https://github.com/apache/airflow/pull/31816#issuecomment-1594476943
@vandonr-amz exactly. I think it would be good to officially add it to `BaseOperator` for 2.7.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1224903303
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -42,9 +43,19 @@
def serialize(result: dict) -> str:
+ logging.getLogger(__name__).error(result)
return json.loads(json.dumps(result, cls=AirflowJsonEncoder))
+def path_to_s3_dataset(path):
+ from openlineage.client.run import Dataset
Review Comment:
I'm not seeing a reason for this to be imported locally down here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mobuchowski commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1226699888
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -275,14 +287,53 @@ def execute(self, context: Context) -> dict:
method_name="execute_complete",
)
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.processing_job = {
+ "Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ }
+ return self.processing_job
def execute_complete(self, context, event=None):
if event["status"] != "success":
raise AirflowException(f"Error while running job: {event}")
else:
self.log.info(event["message"])
- return {"Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))}
+ self.processing_job = {
+ "Processing": serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
+ }
+ return self.processing_job
+
+ def get_openlineage_facets_on_complete(self, task_instance):
+ """Returns OpenLineage data gathered from SageMaker's API response saved by processing job."""
+ from airflow.providers.openlineage.extractors.base import OperatorLineage
+
+ inputs, outputs = [], []
+ try:
+ inputs, outputs = self._get_s3_datasets(
+ processing_inputs=self.processing_job["Processing"]["ProcessingInputs"],
+ processing_outputs=self.processing_job["Processing"]["ProcessingOutputConfig"]["Outputs"],
+ )
+ except KeyError:
+ self.log.exception("Could not find input/output information in Xcom.")
+
+ return OperatorLineage(
+ inputs=inputs,
+ outputs=outputs,
+ )
+
+ def _get_s3_datasets(self, processing_inputs, processing_outputs):
Review Comment:
Renamed to `_extract_s3_dataset_identifiers`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1279745380
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -658,11 +708,66 @@ def execute_complete(self, context, event=None):
else:
self.log.info(event["message"])
transform_config = self.config.get("Transform", self.config)
+ self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
+ self.serialized_tranform = serialize(
+ self.hook.describe_transform_job(transform_config["TransformJobName"])
+ )
return {
- "Model": serialize(self.hook.describe_model(transform_config["ModelName"])),
- "Transform": serialize(self.hook.describe_transform_job(transform_config["TransformJobName"])),
+ "Model": self.serialized_model,
+ "Transform": self.serialized_tranform,
}
+ def get_openlineage_facets_on_complete(self, task_instance):
+ """Returns OpenLineage data gathered from SageMaker's API response saved by transform job."""
+ from airflow.providers.openlineage.extractors import OperatorLineage
+
+ model_package_arn = None
+ transform_input = None
+ transform_output = None
+
+ try:
+ model_package_arn = self.serialized_model["PrimaryContainer"]["ModelPackageName"]
+ except KeyError:
+ self.log.error("Cannot find Model Package Name.", exc_info=True)
+
+ try:
+ transform_input = self.serialized_tranform["TransformInput"]["DataSource"]["S3DataSource"][
+ "S3Uri"
+ ]
+ transform_output = self.serialized_tranform["TransformOutput"]["S3OutputPath"]
+ except KeyError:
+ self.log.error("Cannot find some required input/output details.", exc_info=True)
+
+ inputs = []
+
+ if transform_input is not None:
+ inputs.append(self.path_to_s3_dataset(transform_input))
+
+ if model_package_arn is not None:
+ model_data_urls = self._get_model_data_urls(model_package_arn)
+ for model_data_url in model_data_urls:
+ inputs.append(self.path_to_s3_dataset(model_data_url))
+
+ output = []
+ if transform_output is not None:
+ output.append(self.path_to_s3_dataset(transform_output))
+
+ return OperatorLineage(inputs=inputs, outputs=output)
+
+ def _get_model_data_urls(self, model_package_arn):
Review Comment:
```suggestion
def _get_model_data_urls(self, model_package_arn) -> list:
```
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -658,11 +708,66 @@ def execute_complete(self, context, event=None):
else:
self.log.info(event["message"])
transform_config = self.config.get("Transform", self.config)
+ self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
+ self.serialized_tranform = serialize(
+ self.hook.describe_transform_job(transform_config["TransformJobName"])
+ )
return {
- "Model": serialize(self.hook.describe_model(transform_config["ModelName"])),
- "Transform": serialize(self.hook.describe_transform_job(transform_config["TransformJobName"])),
+ "Model": self.serialized_model,
+ "Transform": self.serialized_tranform,
}
+ def get_openlineage_facets_on_complete(self, task_instance):
+ """Returns OpenLineage data gathered from SageMaker's API response saved by transform job."""
+ from airflow.providers.openlineage.extractors import OperatorLineage
+
+ model_package_arn = None
+ transform_input = None
+ transform_output = None
+
+ try:
+ model_package_arn = self.serialized_model["PrimaryContainer"]["ModelPackageName"]
+ except KeyError:
+ self.log.error("Cannot find Model Package Name.", exc_info=True)
+
+ try:
+ transform_input = self.serialized_tranform["TransformInput"]["DataSource"]["S3DataSource"][
+ "S3Uri"
+ ]
+ transform_output = self.serialized_tranform["TransformOutput"]["S3OutputPath"]
+ except KeyError:
+ self.log.error("Cannot find some required input/output details.", exc_info=True)
+
+ inputs = []
+
+ if transform_input is not None:
+ inputs.append(self.path_to_s3_dataset(transform_input))
+
+ if model_package_arn is not None:
+ model_data_urls = self._get_model_data_urls(model_package_arn)
+ for model_data_url in model_data_urls:
+ inputs.append(self.path_to_s3_dataset(model_data_url))
+
+ output = []
+ if transform_output is not None:
+ output.append(self.path_to_s3_dataset(transform_output))
+
+ return OperatorLineage(inputs=inputs, outputs=output)
+
+ def _get_model_data_urls(self, model_package_arn):
+ model_data_urls = []
+ try:
+ model_containers = self.hook.get_conn().describe_model_package(
+ ModelPackageName=model_package_arn
+ )["InferenceSpecification"]["Containers"]
+
+ for container in model_containers:
+ model_data_urls.append(container["ModelDataUrl"])
+ except Exception:
Review Comment:
don't you want to catch the specific `KeyError` rather than everything, like you did elsewhere ?
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -576,6 +620,8 @@ def __init__(
Provided value: '{action_if_job_exists}'."
)
self.deferrable = deferrable
+ self.serialized_model: dict[Any, Any] | None = None
+ self.serialized_tranform: dict[Any, Any] | None = None
Review Comment:
```suggestion
self.serialized_model: dict | None = None
self.serialized_tranform: dict | None = None
```
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -222,6 +230,7 @@ def __init__(
self.max_attempts = max_attempts or 60
self.max_ingestion_time = max_ingestion_time
self.deferrable = deferrable
+ self.serialized_job: dict[Any, Any] | None = None
Review Comment:
```suggestion
self.serialized_job: dict | None = None
```
`Any, Any` doesn't bring any new information.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1279764127
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -222,6 +230,7 @@ def __init__(
self.max_attempts = max_attempts or 60
self.max_ingestion_time = max_ingestion_time
self.deferrable = deferrable
+ self.serialized_job: dict[Any, Any] | None = None
Review Comment:
I was thinking about that and wondering if it should be `[str, Any]` but I'd be fine dropping it entirely.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #31816: openlineage, sagemaker: add OpenLineage support for SageMaker's Processing, Transform and Training operators
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #31816:
URL: https://github.com/apache/airflow/pull/31816#discussion_r1281048887
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -158,6 +159,14 @@ def hook(self):
"""Return SageMakerHook."""
return SageMakerHook(aws_conn_id=self.aws_conn_id)
+ @staticmethod
+ def path_to_s3_dataset(path):
Review Comment:
```suggestion
def path_to_s3_dataset(path) -> Dataset:
```
(and it needs to be added to the imports under `if TYPE_CHECKING:` as well)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org