Apply SageMaker operator template-field checks after rendering

providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py

@@ -299,11 +299,6 @@ def __init__(
                 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 @@ def expand_role(self) -> None:
             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 @@ def __init__(
         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 @@ def execute(self, context: Context):
             "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))

providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_notebook.py

@@ -95,6 +95,25 @@ def test_create_notebook_without_wait_for_completion(self, mock_hook_conn):
         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(

providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_processing.py

@@ -607,11 +607,12 @@ def test_output_files_not_found_sets_error(self, _, mock_create, mock_count, moc
         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)

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