This is an automated email from the ASF dual-hosted git repository.
o-nikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 930d96e0da6 Add output_files_to_xcom parameter to
SageMakerProcessingOperator (#69002)
930d96e0da6 is described below
commit 930d96e0da6563ab79c1119cce0227043e353a13
Author: Bhavya Sharma <[email protected]>
AuthorDate: Fri Jun 26 16:17:22 2026 -0700
Add output_files_to_xcom parameter to SageMakerProcessingOperator (#69002)
Allow reading small JSON output files from S3 after a processing job
completes and pushing their parsed contents to XCom, making it easy to
feed metrics and evaluation reports into downstream tasks.
---
.../providers/amazon/aws/operators/sagemaker.py | 100 ++++++++-
.../tests/system/amazon/aws/example_sagemaker.py | 52 ++++-
.../aws/operators/test_sagemaker_processing.py | 233 +++++++++++++++++++++
3 files changed, 380 insertions(+), 5 deletions(-)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
index ebafeb85065..e23db37ae2d 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
@@ -258,9 +258,26 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
(default) and "fail".
:param deferrable: Run operator in the deferrable mode. This is only
effective if wait_for_completion is
set to True.
- :return Dict: Returns The ARN of the processing job created in Amazon
SageMaker.
+ :param output_files_to_xcom: Read small JSON output files from S3 after
the job completes and push
+ their parsed contents to XCom useful for feeding metrics into
downstream tasks. Each entry is a dict with:
+ ``output_name`` (matches an output you declared in
``ProcessingOutputConfig.Outputs`` - the
+ operator uses it to locate that output's S3 URI), ``file_name`` (the
file your processing script
+ wrote inside that output, and ``result_name`` (a label
+ you choose for the parsed contents in the returned ``OutputFiles``
dict). Intended for small JSON
+ summary files (metrics, evaluation reports), not large datasets.
Optional; if omitted, no post-job
+ S3 reading occurs. Example::
+
+ output_files_to_xcom=[
+ {"output_name": "evaluation", "file_name": "metrics.json",
"result_name": "EvaluationReport"}
+ ]
+ # Downstream:
ti.xcom_pull("this_task")["OutputFiles"]["EvaluationReport"]["mse"]
+ :return: Dict with a ``Processing`` key (the job description) and, when
``output_files_to_xcom``
+ is configured, an ``OutputFiles`` key containing the parsed JSON
contents of the configured
+ output files (keyed by their ``result_name``).
"""
+ template_fields: Sequence[str] = (*SageMakerBaseOperator.template_fields,
"output_files_to_xcom")
+
def __init__(
self,
*,
@@ -272,14 +289,21 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
max_ingestion_time: int | None = None,
action_if_job_exists: str = "timestamp",
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ output_files_to_xcom: list[dict] | None = None,
**kwargs,
):
super().__init__(config=config, **kwargs)
+ self.output_files_to_xcom = output_files_to_xcom
if action_if_job_exists not in ("fail", "timestamp"):
raise AirflowException(
f"Argument action_if_job_exists accepts only 'timestamp' and
'fail'. \
Provided value: '{action_if_job_exists}'."
)
+ if output_files_to_xcom and not wait_for_completion:
+ raise ValueError(
+ "output_files_to_xcom requires wait_for_completion=True. "
+ "Output files cannot be read before the job completes."
+ )
self.action_if_job_exists = action_if_job_exists
self.wait_for_completion = wait_for_completion
self.print_log = print_log
@@ -339,7 +363,9 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
raise AirflowException(f"SageMaker job failed because
{response['FailureReason']}")
if status == "Completed":
self.log.info("%s completed successfully.", self.task_id)
- return {"Processing": serialize(response)}
+ result = {"Processing": serialize(response)}
+ result.update(self._read_output_files())
+ return result
timeout = self.execution_timeout
if self.max_ingestion_time:
@@ -358,7 +384,9 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
)
self.serialized_job =
serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
- return {"Processing": self.serialized_job}
+ result = {"Processing": self.serialized_job}
+ result.update(self._read_output_files())
+ return result
def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> dict[str, dict]:
validated_event = validate_execute_complete_event(event)
@@ -369,7 +397,71 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
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 {}
+
+ s3_client = self.hook.get_session().client("s3")
+ processing_outputs = self.config.get("ProcessingOutputConfig",
{}).get("Outputs", [])
+ output_files_data: dict = {}
+
+ for output_file in self.output_files_to_xcom:
+ missing_keys = {"result_name", "output_name", "file_name"} -
output_file.keys()
+ if missing_keys:
+ raise ValueError(
+ f"output_files_to_xcom entry missing required keys:
{missing_keys}. "
+ f"Each entry requires 'result_name', 'output_name', and
'file_name'."
+ )
+ result_name = output_file["result_name"]
+ output_name = output_file["output_name"]
+ file_name = output_file["file_name"]
+
+ matching_output = next(
+ (output for output in processing_outputs if
output["OutputName"] == output_name),
+ None,
+ )
+ if not matching_output:
+ self.log.warning(
+ "output_files_to_xcom '%s': no matching output '%s'
found.",
+ result_name,
+ output_name,
+ )
+ continue
+
+ s3_output_uri = matching_output["S3Output"]["S3Uri"].rstrip("/")
+ s3_full_path = f"{s3_output_uri}/{file_name}"
+ bucket_name, object_key = s3_full_path.replace("s3://",
"").split("/", 1)
+ self.log.debug("Attempting to read output file: s3://%s/%s",
bucket_name, object_key)
+
+ try:
+ s3_response = s3_client.get_object(Bucket=bucket_name,
Key=object_key)
+ file_content = s3_response["Body"].read().decode("utf-8")
+ output_files_data[result_name] = json.loads(file_content)
+ self.log.info("Output file '%s' loaded from %s", result_name,
s3_full_path)
+ except ClientError as error:
+ if error.response.get("Error", {}).get("Code") == "NoSuchKey":
+ self.log.warning("Output file '%s' not found at %s",
result_name, s3_full_path)
+ output_files_data[result_name] = {"_error": f"File not
found: {s3_full_path}"}
+ else:
+ self.log.exception("Failed to read output file '%s' from
%s", result_name, s3_full_path)
+ output_files_data[result_name] = {"_error": f"Failed to
read from {s3_full_path}"}
+ except json.JSONDecodeError as error:
+ self.log.exception(
+ "Output file '%s' at %s is not valid JSON: %s",
result_name, s3_full_path, error
+ )
+ output_files_data[result_name] = {"_error": f"Invalid JSON at
{s3_full_path}"}
+ except Exception:
+ self.log.exception("Failed to read output file '%s' from %s",
result_name, s3_full_path)
+ output_files_data[result_name] = {"_error": f"Failed to read
from {s3_full_path}"}
+
+ if output_files_data:
+ return {"OutputFiles": output_files_data}
+ return {}
def get_openlineage_facets_on_complete(self, task_instance) ->
OperatorLineage:
"""Return OpenLineage data gathered from SageMaker's API response
saved by processing job."""
diff --git a/providers/amazon/tests/system/amazon/aws/example_sagemaker.py
b/providers/amazon/tests/system/amazon/aws/example_sagemaker.py
index 57975ca2ec2..01f6e13641d 100644
--- a/providers/amazon/tests/system/amazon/aws/example_sagemaker.py
+++ b/providers/amazon/tests/system/amazon/aws/example_sagemaker.py
@@ -93,6 +93,8 @@ SAMPLE_SIZE = 600
PREPROCESS_SCRIPT_TEMPLATE = dedent("""
import numpy as np
import pandas as pd
+ import json
+ import os
def main():
# Load the dataset from {input_path}/input.csv, split it into
train/test
@@ -112,6 +114,13 @@ PREPROCESS_SCRIPT_TEMPLATE = dedent("""
data_train.to_csv('{output_path}/train.csv', index=False, header=False)
data_test.to_csv('{output_path}/test.csv', index=False, header=False)
+ # Write evaluation metrics for output_files_to_xcom testing
+ evaluation_dir = '{output_path}/evaluation'
+ os.makedirs(evaluation_dir, exist_ok=True)
+ metrics = {{"metrics": {{"mse": {{"value": 5.332878}}, "rmse":
{{"value": 2.309302}}}}}}
+ with open(os.path.join(evaluation_dir, 'evaluation.json'), 'w') as f:
+ json.dump(metrics, f)
+
print('Preprocessing Done.')
if __name__ == "__main__":
@@ -294,7 +303,16 @@ def set_up(env_id, role_arn):
"S3UploadMode": "EndOfJob",
},
"AppManaged": False,
- }
+ },
+ {
+ "OutputName": "evaluation",
+ "S3Output": {
+ "S3Uri": f"s3://{bucket_name}/{env_id}/evaluation",
+ "LocalPath":
f"{processing_local_output_path}/evaluation",
+ "S3UploadMode": "EndOfJob",
+ },
+ "AppManaged": False,
+ },
]
},
"ProcessingResources": {
@@ -307,6 +325,14 @@ def set_up(env_id, role_arn):
"RoleArn": role_arn,
}
+ output_files_to_xcom_config = [
+ {
+ "result_name": "EvaluationReport",
+ "output_name": "evaluation",
+ "file_name": "evaluation.json",
+ }
+ ]
+
training_data_source = {
"CompressionType": "None",
"ContentType": "text/csv",
@@ -444,6 +470,7 @@ def set_up(env_id, role_arn):
ti.xcom_push(key="raw_data_s3_key", value=raw_data_s3_key)
ti.xcom_push(key="ecr_repository_name", value=ecr_repository_name)
ti.xcom_push(key="processing_config", value=processing_config)
+ ti.xcom_push(key="output_files_to_xcom_config",
value=output_files_to_xcom_config)
ti.xcom_push(key="processing_job_name", value=processing_job_name)
ti.xcom_push(key="input_data_uri", value=input_data_uri)
ti.xcom_push(key="output_data_uri",
value=f"s3://{bucket_name}/{training_output_s3_key}")
@@ -601,6 +628,27 @@ with DAG(
)
# [END howto_sensor_sagemaker_processing]
+ # [START howto_operator_sagemaker_processing_output_files]
+ preprocess_with_output_files = SageMakerProcessingOperator(
+ task_id="preprocess_with_output_files",
+ config=test_setup["processing_config"],
+ output_files_to_xcom=test_setup["output_files_to_xcom_config"],
+ )
+ # [END howto_operator_sagemaker_processing_output_files]
+
+ @task
+ def validate_output_files(xcom_result):
+ """Verify output files were read from S3 and included in XCom."""
+ assert "OutputFiles" in xcom_result, (
+ f"OutputFiles missing from XCom. Got keys:
{list(xcom_result.keys())}"
+ )
+ evaluation_report = xcom_result["OutputFiles"]["EvaluationReport"]
+ assert "metrics" in evaluation_report, f"metrics key missing. Got:
{evaluation_report}"
+ assert evaluation_report["metrics"]["mse"]["value"] == 5.332878
+ assert evaluation_report["metrics"]["rmse"]["value"] == 2.309302
+
+ validate_pf = validate_output_files(preprocess_with_output_files.output)
+
# [START howto_operator_sagemaker_training]
train_model = SageMakerTrainingOperator(
task_id="train_model",
@@ -704,6 +752,8 @@ with DAG(
create_experiment,
preprocess_raw_data,
await_preprocess,
+ preprocess_with_output_files,
+ validate_pf,
train_model,
await_training,
create_model,
diff --git
a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_processing.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_processing.py
index 3bddcf88f3c..8d915ec032a 100644
---
a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_processing.py
+++
b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_processing.py
@@ -382,3 +382,236 @@ class TestSageMakerProcessingOperator:
config=CREATE_PROCESSING_PARAMS,
)
validate_template_fields(operator)
+
+
+PROCESSING_CONFIG_WITH_OUTPUT_FILES: dict = {
+ **CREATE_PROCESSING_PARAMS,
+ "ProcessingOutputConfig": {
+ "Outputs": [
+ {
+ "OutputName": "evaluation",
+ "S3Output": {
+ "LocalPath": "/opt/ml/processing/evaluation",
+ "S3UploadMode": "EndOfJob",
+ "S3Uri": "s3://my-bucket/eval-output",
+ },
+ }
+ ],
+ },
+}
+
+OUTPUT_FILES_TO_XCOM_CONFIG: list = [
+ {
+ "result_name": "EvaluationReport",
+ "output_name": "evaluation",
+ "file_name": "evaluation.json",
+ }
+]
+
+
+class TestSageMakerProcessingOperatorOutputFiles:
+ """Tests for the output_files_to_xcom feature in
SageMakerProcessingOperator."""
+
+ @mock.patch.object(
+ SageMakerHook,
+ "describe_processing_job",
+ return_value={"ProcessingJobStatus": "Completed", "ProcessingJobName":
"job_name"},
+ )
+ @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.object(SageMakerBaseOperator, "_check_if_job_exists",
return_value=False)
+ def test_output_files_read_from_s3(self, _, mock_create, mock_count,
mock_describe):
+ """Output files are read from S3 and included in XCom return value."""
+ evaluation_json = '{"metrics": {"mse": {"value": 5.33}}}'
+
+ mock_s3_client = mock.MagicMock()
+ mock_s3_client.get_object.return_value = {
+ "Body":
mock.MagicMock(read=mock.MagicMock(return_value=evaluation_json.encode("utf-8")))
+ }
+
+ with mock.patch.object(SageMakerHook, "get_session") as mock_session:
+ mock_session.return_value.client.return_value = mock_s3_client
+
+ operator = SageMakerProcessingOperator(
+ task_id="test_task",
+ config=PROCESSING_CONFIG_WITH_OUTPUT_FILES,
+ output_files_to_xcom=OUTPUT_FILES_TO_XCOM_CONFIG,
+ )
+ result = operator.execute(context=None)
+
+ assert "OutputFiles" in result
+ assert
result["OutputFiles"]["EvaluationReport"]["metrics"]["mse"]["value"] == 5.33
+ mock_s3_client.get_object.assert_called_once_with(
+ Bucket="my-bucket", Key="eval-output/evaluation.json"
+ )
+
+ @mock.patch.object(
+ SageMakerHook,
+ "describe_processing_job",
+ return_value={"ProcessingJobStatus": "Completed", "ProcessingJobName":
"job_name"},
+ )
+ @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.object(SageMakerBaseOperator, "_check_if_job_exists",
return_value=False)
+ def test_no_output_files_returns_unchanged(self, _, mock_create,
mock_count, mock_describe):
+ """Without output_files_to_xcom kwarg, XCom return has no OutputFiles
key."""
+ operator = SageMakerProcessingOperator(
+ task_id="test_task",
+ config=CREATE_PROCESSING_PARAMS,
+ )
+ result = operator.execute(context=None)
+
+ assert "Processing" in result
+ assert "OutputFiles" not in result
+
+ @mock.patch.object(
+ SageMakerHook,
+ "describe_processing_job",
+ return_value={"ProcessingJobStatus": "Completed", "ProcessingJobName":
"job_name"},
+ )
+ @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.object(SageMakerBaseOperator, "_check_if_job_exists",
return_value=False)
+ def test_output_files_s3_error_does_not_fail_task(self, _, mock_create,
mock_count, mock_describe):
+ """S3 read failure sets _error key but does not raise."""
+ mock_s3_client = mock.MagicMock()
+ mock_s3_client.get_object.side_effect = Exception("NoSuchKey")
+
+ with mock.patch.object(SageMakerHook, "get_session") as mock_session:
+ mock_session.return_value.client.return_value = mock_s3_client
+
+ operator = SageMakerProcessingOperator(
+ task_id="test_task",
+ config=PROCESSING_CONFIG_WITH_OUTPUT_FILES,
+ output_files_to_xcom=OUTPUT_FILES_TO_XCOM_CONFIG,
+ )
+ result = operator.execute(context=None)
+
+ assert "OutputFiles" in result
+ assert "_error" in result["OutputFiles"]["EvaluationReport"]
+
+ @mock.patch.object(
+ SageMakerHook,
+ "describe_processing_job",
+ return_value={"ProcessingJobStatus": "Completed", "ProcessingJobName":
"job_name"},
+ )
+ @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.object(SageMakerBaseOperator, "_check_if_job_exists",
return_value=False)
+ def test_output_files_no_matching_output_logs_warning(self, _,
mock_create, mock_count, mock_describe):
+ """Output file referencing a non-existent output_name is skipped
gracefully."""
+ bad_output_files = [
+ {
+ "result_name": "MissingReport",
+ "output_name": "nonexistent_output",
+ "file_name": "report.json",
+ }
+ ]
+ operator = SageMakerProcessingOperator(
+ task_id="test_task",
+ config=CREATE_PROCESSING_PARAMS,
+ output_files_to_xcom=bad_output_files,
+ )
+ result = operator.execute(context=None)
+
+ assert "OutputFiles" not in result
+
+ def test_output_files_missing_required_keys_raises(self):
+ """An output_files_to_xcom entry missing required keys raises when
output files are read."""
+ op = SageMakerProcessingOperator(
+ task_id="test_task",
+ config=CREATE_PROCESSING_PARAMS,
+ output_files_to_xcom=[{"result_name": "Report", "output_name":
"evaluation"}], # no file_name
+ )
+ with mock.patch.object(SageMakerHook, "get_session"):
+ with pytest.raises(ValueError, match="missing required keys"):
+ op._read_output_files()
+
+ @mock.patch.object(
+ SageMakerHook,
+ "describe_processing_job",
+ return_value={"ProcessingJobStatus": "Completed", "ProcessingJobName":
"job_name"},
+ )
+ @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.object(SageMakerBaseOperator, "_check_if_job_exists",
return_value=False)
+ def test_output_files_invalid_json_sets_error(self, _, mock_create,
mock_count, mock_describe):
+ """A non-JSON file sets an 'Invalid JSON' error but does not fail the
task."""
+ mock_s3_client = mock.MagicMock()
+ mock_s3_client.get_object.return_value = {
+ "Body":
mock.MagicMock(read=mock.MagicMock(return_value=b"col1,col2\n1,2\n"))
+ }
+
+ with mock.patch.object(SageMakerHook, "get_session") as mock_session:
+ mock_session.return_value.client.return_value = mock_s3_client
+
+ operator = SageMakerProcessingOperator(
+ task_id="test_task",
+ config=PROCESSING_CONFIG_WITH_OUTPUT_FILES,
+ output_files_to_xcom=OUTPUT_FILES_TO_XCOM_CONFIG,
+ )
+ result = operator.execute(context=None)
+
+ assert "Invalid JSON" in
result["OutputFiles"]["EvaluationReport"]["_error"]
+
+ @mock.patch.object(
+ SageMakerHook,
+ "describe_processing_job",
+ return_value={"ProcessingJobStatus": "Completed", "ProcessingJobName":
"job_name"},
+ )
+ @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.object(SageMakerBaseOperator, "_check_if_job_exists",
return_value=False)
+ def test_output_files_not_found_sets_error(self, _, mock_create,
mock_count, mock_describe):
+ """A missing S3 object (NoSuchKey) sets a 'File not found' error but
does not fail the task."""
+ mock_s3_client = mock.MagicMock()
+ mock_s3_client.get_object.side_effect = ClientError(
+ {"Error": {"Code": "NoSuchKey", "Message": "The specified key does
not exist."}},
+ "GetObject",
+ )
+
+ with mock.patch.object(SageMakerHook, "get_session") as mock_session:
+ mock_session.return_value.client.return_value = mock_s3_client
+
+ operator = SageMakerProcessingOperator(
+ task_id="test_task",
+ config=PROCESSING_CONFIG_WITH_OUTPUT_FILES,
+ output_files_to_xcom=OUTPUT_FILES_TO_XCOM_CONFIG,
+ )
+ result = operator.execute(context=None)
+
+ assert "File not found" in
result["OutputFiles"]["EvaluationReport"]["_error"]
+
+ def test_output_files_skipped_when_not_waiting(self):
+ """With wait_for_completion=False and output_files_to_xcom, init
raises ValueError."""
+ with pytest.raises(ValueError, match="output_files_to_xcom requires
wait_for_completion=True"):
+ SageMakerProcessingOperator(
+ task_id="test_task",
+ config=PROCESSING_CONFIG_WITH_OUTPUT_FILES,
+ output_files_to_xcom=OUTPUT_FILES_TO_XCOM_CONFIG,
+ wait_for_completion=False,
+ )
