code: def test_mlmd_execution_update(self):
      event_on_commit = threading.Event()
      got_pre_commit_execution = None
      got_post_commit_execution = None
  
      def on_commit(pre_commit_execution, post_commit_execution):
          nonlocal got_pre_commit_execution
          nonlocal got_post_commit_execution
          got_pre_commit_execution = pre_commit_execution
          got_post_commit_execution = post_commit_execution
          event_on_commit.set()
  
      with self._mlmd_connection as m:
          expected_execution = _write_test_execution(m)
          # Mutate execution.
          with mlmd_state.mlmd_execution_atomic_op(
              m, expected_execution.id, on_commit=on_commit
          ) as execution:
              self.assertEqual(expected_execution, execution)
              execution.last_known_state = (
                  metadata_store_pb2.Execution.CANCELED
              )
              self.assertFalse(
                  event_on_commit.is_set()
              )  # not yet invoked.
          self.assertEqual(expected_execution, got_pre_commit_execution)
          self.assertEqual(
              metadata_store_pb2.Execution.CANCELED,
              got_post_commit_execution.last_known_state,
          )
  
          # Test that we made a deep copy of the executions, so mutating them
          # doesn't mutate the values in the cache.
          got_pre_commit_execution.last_known_state = (
              metadata_store_pb2.Execution.UNKNOWN
          )
          got_post_commit_execution.last_known_state = (
              metadata_store_pb2.Execution.UNKNOWN
          )
  
          # Test that updated execution is committed to MLMD.
          [execution] = m.store.get_executions_by_id([execution.id])
          self.assertEqual(
              metadata_store_pb2.Execution.CANCELED,
              execution.last_known_state,
          )
          # Test that in-memory state is also in sync.
          self.assertEqual(
              execution,
              mlmd_state._execution_cache._cache[execution.id],
          )
          # Test that on_commit callback was invoked.
          self.assertTrue(event_on_commit.is_set())
          # Sanity checks that the updated execution is yielded in the next call.
          with mlmd_state.mlmd_execution_atomic_op(
              m, expected_execution.id
          ) as execution2:
              self.assertEqual(execution, execution2)
  
imports: from concurrent import futures
  import os
  import threading
  import tensorflow
  from tfx.orchestration import metadata
  from tfx.orchestration.experimental.core import mlmd_state
  from tfx.orchestration.experimental.core import test_utils
  from ml_metadata.proto import metadata_store_pb2
label_desc: Use image node_id 2 for calling the MlmdStateTest obj's underlying member method code with example usage: obj.test_mlmd_execution_update() without return types