This is an automated email from the ASF dual-hosted git repository.
potiuk 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 087b2e9fe4c Apply SageMaker operator template-field checks after
rendering (#70485)
087b2e9fe4c is described below
commit 087b2e9fe4ca22cb9bc9b03b2db47cd94a5c8c2c
Author: Dr Alex Mitre <[email protected]>
AuthorDate: Fri Jul 31 19:22:55 2026 -0600
Apply SageMaker operator template-field checks after rendering (#70485)
output_files_to_xcom (SageMakerProcessingOperator) and
create_instance_kwargs
(SageMakerCreateNotebookOperator) are template fields, but were validated or
transformed in __init__, before Jinja rendering. The wait_for_completion
guard
and the tags formatting now run at execute() time, on rendered values.
Co-authored-by: Dr Alex Mitre <[email protected]>
---
.../providers/amazon/aws/operators/sagemaker.py | 18 +++++++++---------
.../amazon/aws/operators/test_sagemaker_notebook.py | 19 +++++++++++++++++++
.../amazon/aws/operators/test_sagemaker_processing.py | 15 ++++++++-------
.../ci/prek/validate_operators_init_exemptions.txt | 2 --
4 files changed, 36 insertions(+), 18 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 21fdaa6b988..066f63eca36 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
@@ -299,11 +299,6 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
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
@@ -329,6 +324,11 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
self.config["RoleArn"] = hook.expand_role(self.config["RoleArn"])
def execute(self, context: Context) -> dict:
+ if self.output_files_to_xcom and not self.wait_for_completion:
+ raise ValueError(
+ "output_files_to_xcom requires wait_for_completion=True. "
+ "Output files cannot be read before the job completes."
+ )
self.preprocess_config()
self.config["ProcessingJobName"] = self._get_unique_job_name(
@@ -1926,9 +1926,6 @@ class
SageMakerCreateNotebookOperator(AwsBaseOperator[SageMakerHook]):
self.wait_for_completion = wait_for_completion
self.create_instance_kwargs = create_instance_kwargs or {}
- if self.create_instance_kwargs.get("tags") is not None:
- self.create_instance_kwargs["tags"] =
format_tags(self.create_instance_kwargs["tags"])
-
def execute(self, context: Context):
create_notebook_instance_kwargs = {
"NotebookInstanceName": self.instance_name,
@@ -1941,7 +1938,10 @@ class
SageMakerCreateNotebookOperator(AwsBaseOperator[SageMakerHook]):
"RootAccess": self.root_access,
}
if self.create_instance_kwargs:
- create_notebook_instance_kwargs.update(self.create_instance_kwargs)
+ create_instance_kwargs = dict(self.create_instance_kwargs)
+ if create_instance_kwargs.get("tags") is not None:
+ create_instance_kwargs["tags"] =
format_tags(create_instance_kwargs["tags"])
+ create_notebook_instance_kwargs.update(create_instance_kwargs)
self.log.info("Creating SageMaker notebook %s.", self.instance_name)
response =
self.hook.conn.create_notebook_instance(**prune_dict(create_notebook_instance_kwargs))
diff --git
a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_notebook.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_notebook.py
index c97b7f4b3d8..9f31df76107 100644
---
a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_notebook.py
+++
b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_notebook.py
@@ -95,6 +95,25 @@ class TestSagemakerCreateNotebookOperator:
mock_hook_conn.create_notebook_instance.assert_called_once()
mock_hook_conn.get_waiter.assert_not_called()
+ @mock.patch.object(SageMakerHook, "conn")
+ def test_create_notebook_formats_tags_at_execute_time(self,
mock_hook_conn):
+ operator = SageMakerCreateNotebookOperator(
+ task_id="task_test",
+ instance_name=INSTANCE_NAME,
+ instance_type=INSTANCE_TYPE,
+ role_arn=ROLE_ARN,
+ wait_for_completion=False,
+ create_instance_kwargs={"tags": {"team": "data"}},
+ )
+
+ assert operator.create_instance_kwargs == {"tags": {"team": "data"}}
+
+ operator.execute(None)
+
+ call_kwargs = mock_hook_conn.create_notebook_instance.call_args.kwargs
+ assert call_kwargs["tags"] == [{"Key": "team", "Value": "data"}]
+ assert operator.create_instance_kwargs == {"tags": {"team": "data"}}
+
@mock.patch.object(SageMakerHook, "conn")
def test_create_notebook_wait_for_completion(self, mock_hook_conn):
operator = SageMakerCreateNotebookOperator(
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 8d915ec032a..eb14ed4b3c6 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
@@ -607,11 +607,12 @@ class TestSageMakerProcessingOperatorOutputFiles:
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 wait_for_completion=False and output_files_to_xcom, execute
raises ValueError."""
+ operator = SageMakerProcessingOperator(
+ task_id="test_task",
+ config=PROCESSING_CONFIG_WITH_OUTPUT_FILES,
+ output_files_to_xcom=OUTPUT_FILES_TO_XCOM_CONFIG,
+ wait_for_completion=False,
+ )
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,
- )
+ operator.execute(context=None)
diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt
b/scripts/ci/prek/validate_operators_init_exemptions.txt
index 6c4907fbded..a41b768d9be 100644
--- a/scripts/ci/prek/validate_operators_init_exemptions.txt
+++ b/scripts/ci/prek/validate_operators_init_exemptions.txt
@@ -7,8 +7,6 @@
# Burn-down tracked at https://github.com/apache/airflow/issues/70296
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune.py::NeptuneStartDbClusterOperator
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune.py::NeptuneStopDbClusterOperator
-providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py::SageMakerCreateNotebookOperator
-providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py::SageMakerProcessingOperator
providers/amazon/src/airflow/providers/amazon/aws/transfers/gcs_to_s3.py::GCSToS3Operator
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py::KubernetesPodOperator
providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py::CloudBatchSubmitJobOperator
