Repository: apache/airflow

Issue:
Context is not preserved after execution returns from deferred state
### Apache Airflow version

Other Airflow 2 version (please specify below)

### If "Other Airflow 2 version" selected, which one?

2.7.2

### What happened?

We are using 8.7.1. version of airflow provider package and set deferrable to True when using SageMakerTransformOperator. We set define the model name with something  like "project-dag_name-date-uuid", where uuid is to deduplicate model names upon retry. However, when the execution gets back from deferred state, the original uuid is not preserved and there would be a new uuid so the operator cannot find a model with the new model name and it would fail.

### What you think should happen instead?

The original UUID should be preserved 

### How to reproduce

In the operator config, set "ModelName": f"{name}-{str(uuid4())[:8]}". The run the DAG.

### Operating System

amazon Linux AMI

### Versions of Apache Airflow Providers

8.7.1

### Deployment

Amazon (AWS) MWAA

### Deployment details

_No response_

### Anything else?

_No response_

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:
Can you check if it's the same in latest provider please @mxuan0 ?

Resolving patch:
diff --git a/airflow/providers/amazon/aws/operators/sagemaker.py b/airflow/providers/amazon/aws/operators/sagemaker.py
--- a/airflow/providers/amazon/aws/operators/sagemaker.py
+++ b/airflow/providers/amazon/aws/operators/sagemaker.py
@@ -750,20 +750,18 @@ def execute(self, context: Context) -> dict:
                 method_name="execute_complete",
             )
 
-        return self.serialize_result()
+        return self.serialize_result(transform_config["TransformJobName"])
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, dict]:
         event = validate_execute_complete_event(event)
 
         self.log.info(event["message"])
-        return self.serialize_result()
+        return self.serialize_result(event["job_name"])
 
-    def serialize_result(self) -> dict[str, dict]:
-        transform_config = self.config.get("Transform", self.config)
-        self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
-        self.serialized_transform = serialize(
-            self.hook.describe_transform_job(transform_config["TransformJobName"])
-        )
+    def serialize_result(self, job_name: str) -> dict[str, dict]:
+        job_description = self.hook.describe_transform_job(job_name)
+        self.serialized_model = serialize(self.hook.describe_model(job_description["ModelName"]))
+        self.serialized_transform = serialize(job_description)
         return {"Model": self.serialized_model, "Transform": self.serialized_transform}
 
     def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage:
@@ -1154,7 +1152,7 @@ def execute(self, context: Context) -> dict:
                 method_name="execute_complete",
             )
 
-        return self.serialize_result()
+        return self.serialize_result(self.config["TrainingJobName"])
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, dict]:
         event = validate_execute_complete_event(event)
@@ -1163,12 +1161,10 @@ def execute_complete(self, context: Context, event: dict[str, Any] | None = None
             raise AirflowException(f"Error while running job: {event}")
 
         self.log.info(event["message"])
-        return self.serialize_result()
+        return self.serialize_result(event["job_name"])
 
-    def serialize_result(self) -> dict[str, dict]:
-        self.serialized_training_data = serialize(
-            self.hook.describe_training_job(self.config["TrainingJobName"])
-        )
+    def serialize_result(self, job_name: str) -> dict[str, dict]:
+        self.serialized_training_data = serialize(self.hook.describe_training_job(job_name))
         return {"Training": self.serialized_training_data}
 
     def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage:
diff --git a/airflow/providers/amazon/aws/triggers/sagemaker.py b/airflow/providers/amazon/aws/triggers/sagemaker.py
--- a/airflow/providers/amazon/aws/triggers/sagemaker.py
+++ b/airflow/providers/amazon/aws/triggers/sagemaker.py
@@ -121,7 +121,7 @@ async def run(self):
                 status_message=f"{self.job_type} job not done yet",
                 status_args=[self._get_response_status_key(self.job_type)],
             )
-            yield TriggerEvent({"status": "success", "message": "Job completed."})
+            yield TriggerEvent({"status": "success", "message": "Job completed.", "job_name": self.job_name})
 
 
 class SageMakerPipelineTrigger(BaseTrigger):


Test patch:
diff --git a/tests/providers/amazon/aws/operators/test_sagemaker_transform.py b/tests/providers/amazon/aws/operators/test_sagemaker_transform.py
--- a/tests/providers/amazon/aws/operators/test_sagemaker_transform.py
+++ b/tests/providers/amazon/aws/operators/test_sagemaker_transform.py
@@ -83,7 +83,10 @@ def setup_method(self):
     @mock.patch.object(SageMakerHook, "create_transform_job")
     @mock.patch.object(sagemaker, "serialize", return_value="")
     def test_integer_fields(self, _, mock_create_transform, __, ___, mock_desc):
-        mock_desc.side_effect = [ClientError({"Error": {"Code": "ValidationException"}}, "op"), None]
+        mock_desc.side_effect = [
+            ClientError({"Error": {"Code": "ValidationException"}}, "op"),
+            {"ModelName": "model_name"},
+        ]
         mock_create_transform.return_value = {
             "TransformJobArn": "test_arn",
             "ResponseMetadata": {"HTTPStatusCode": 200},
@@ -103,7 +106,10 @@ def test_integer_fields(self, _, mock_create_transform, __, ___, mock_desc):
     @mock.patch.object(SageMakerHook, "create_transform_job")
     @mock.patch.object(sagemaker, "serialize", return_value="")
     def test_execute(self, _, mock_transform, __, mock_model, mock_desc):
-        mock_desc.side_effect = [ClientError({"Error": {"Code": "ValidationException"}}, "op"), None]
+        mock_desc.side_effect = [
+            ClientError({"Error": {"Code": "ValidationException"}}, "op"),
+            {"ModelName": "model_name"},
+        ]
         mock_transform.return_value = {
             "TransformJobArn": "test_arn",
             "ResponseMetadata": {"HTTPStatusCode": 200},
@@ -135,7 +141,10 @@ def test_execute_with_failure(self, mock_transform, _, mock_desc):
     @mock.patch.object(SageMakerHook, "describe_model")
     @mock.patch.object(sagemaker, "serialize", return_value="")
     def test_execute_with_check_if_job_exists(self, _, __, ___, mock_transform, mock_desc):
-        mock_desc.side_effect = [ClientError({"Error": {"Code": "ValidationException"}}, "op"), None]
+        mock_desc.side_effect = [
+            ClientError({"Error": {"Code": "ValidationException"}}, "op"),
+            {"ModelName": "model_name"},
+        ]
         mock_transform.return_value = {
             "TransformJobArn": "test_arn",
             "ResponseMetadata": {"HTTPStatusCode": 200},
@@ -243,6 +252,7 @@ def test_operator_lineage_data(self, mock_transform, mock_conn, mock_model, _, m
         mock_desc.return_value = {
             "TransformInput": {"DataSource": {"S3DataSource": {"S3Uri": "s3://input-bucket/input-path"}}},
             "TransformOutput": {"S3OutputPath": "s3://output-bucket/output-path"},
+            "ModelName": "model_name",
         }
         mock_transform.return_value = {
             "TransformJobArn": "test_arn",
diff --git a/tests/providers/amazon/aws/triggers/test_sagemaker.py b/tests/providers/amazon/aws/triggers/test_sagemaker.py
--- a/tests/providers/amazon/aws/triggers/test_sagemaker.py
+++ b/tests/providers/amazon/aws/triggers/test_sagemaker.py
@@ -77,4 +77,6 @@ async def test_sagemaker_trigger_run_all_job_types(self, mock_async_conn, mock_g
         generator = sagemaker_trigger.run()
         response = await generator.asend(None)
 
-        assert response == TriggerEvent({"status": "success", "message": "Job completed."})
+        assert response == TriggerEvent(
+            {"status": "success", "message": "Job completed.", "job_name": JOB_NAME}
+        )