Fix a bug when applying pipeline update

The updated IR should be stored as a property, not custom property

PiperOrigin-RevId: 394614056

tfx/orchestration/experimental/core/pipeline_state.py

@@ -344,7 +344,7 @@ def apply_pipeline_update(self) -> None:
           code=status_lib.Code.INVALID_ARGUMENT,
           message='No updated pipeline IR to apply')
     data_types_utils.set_metadata_value(
-        self._execution.custom_properties[_PIPELINE_IR], updated_pipeline_ir)
+        self._execution.properties[_PIPELINE_IR], updated_pipeline_ir)
     del self._execution.custom_properties[_UPDATED_PIPELINE_IR]
     self.pipeline = _base64_decode_pipeline(updated_pipeline_ir)
 

tfx/orchestration/experimental/core/pipeline_state_test.py

@@ -199,17 +199,28 @@ def test_update_initiation_and_apply(self):
     with self._mlmd_connection as m:
       pipeline = _test_pipeline('pipeline1', param=1)
       updated_pipeline = _test_pipeline('pipeline1', param=2)
+
+      # Initiate pipeline update.
       with pstate.PipelineState.new(m, pipeline) as pipeline_state:
         self.assertFalse(pipeline_state.is_update_initiated())
         pipeline_state.initiate_update(updated_pipeline)
         self.assertTrue(pipeline_state.is_update_initiated())
 
-      # Reload from MLMD and verify.
+      # Reload from MLMD and verify update initiation followed by applying the
+      # pipeline update.
       with pstate.PipelineState.load(
           m, task_lib.PipelineUid.from_pipeline(pipeline)) as pipeline_state:
         self.assertTrue(pipeline_state.is_update_initiated())
         self.assertEqual(pipeline, pipeline_state.pipeline)
         pipeline_state.apply_pipeline_update()
+        # Verify in-memory state after update application.
+        self.assertFalse(pipeline_state.is_update_initiated())
+        self.assertTrue(pipeline_state.is_active())
+        self.assertEqual(updated_pipeline, pipeline_state.pipeline)
+
+      # Reload from MLMD and verify update application was correctly persisted.
+      with pstate.PipelineState.load(
+          m, task_lib.PipelineUid.from_pipeline(pipeline)) as pipeline_state:
         self.assertFalse(pipeline_state.is_update_initiated())
         self.assertTrue(pipeline_state.is_active())
         self.assertEqual(updated_pipeline, pipeline_state.pipeline)