o-nikolas commented on code in PR #69002:
URL: https://github.com/apache/airflow/pull/69002#discussion_r3483158711
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -369,7 +392,78 @@ def execute_complete(self, context: Context, event:
dict[str, Any] | None = None
self.log.info("SageMaker job %s completed.",
validated_event["job_name"])
self.serialized_job =
serialize(self.hook.describe_processing_job(validated_event["job_name"]))
self.log.info("%s completed successfully.", self.task_id)
- return {"Processing": self.serialized_job}
+ result = {"Processing": self.serialized_job}
+ result.update(self._read_output_files())
+ return result
+
+ def _read_output_files(self) -> dict:
+ """Read configured output files from S3 after job completion and
return them for XCom."""
+ if not self.output_files_to_xcom:
+ return {}
+
+ if not self.wait_for_completion:
+ self.log.warning(
+ "output_files_to_xcom is configured but
wait_for_completion=False. "
+ "Output files cannot be read before the job completes.
Skipping."
+ )
Review Comment:
You should be able to determine this straight away before the job is even
started in `execute()` if you wanted to fail early. Or I think you could
possibly fail at parse time if you check this in `init`.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
