code: def mlmd_execution_atomic_op(
      mlmd_handle: metadata.Metadata,
      execution_id: int,
      on_commit: Optional[
          Callable[
              [
                  metadata_store_pb2.Execution,
                  metadata_store_pb2.Execution,
              ],
              None,
          ]
      ] = None,
  ) -> Iterator[metadata_store_pb2.Execution]:
      """Context manager for accessing or mutating an execution atomically.
  
      The idea of using this context manager is to ensure that the in-memory state
      of an MLMD execution is centrally managed so that it stays in sync with the
      execution in MLMD even when multiple threads in the process may be mutating.
  
      If execution for given execution id exists in MLMD, it is locked before being
      yielded so that no other thread in the process can make conflicting updates if
      the yielded execution is mutated within the context. Mutated executions are
      also automatically committed to MLMD when exiting the context.
  
      Args:
        mlmd_handle: A handle to MLMD db.
        execution_id: Id of the execution to yield.
        on_commit: An optional callback function which is invoked post successful
          MLMD execution commit operation. This won't be invoked if execution is not
          mutated within the context and hence MLMD commit is not needed. The
          callback is passed copies of the pre-commit and post-commit executions.
  
      Yields:
        If execution with given id exists in MLMD, the execution is yielded under
        an exclusive lock context.
  
      Raises:
        RuntimeError: If execution id is changed within the context.
        ValueError: If execution having given execution id is not found in MLMD.
      """
      with _execution_id_locks.lock(execution_id):
          execution = _execution_cache.get_execution(
              mlmd_handle, execution_id
          )
          execution_copy = copy.deepcopy(execution)
          yield execution_copy
          if execution != execution_copy:
              if execution.id != execution_copy.id:
                  raise RuntimeError(
                      "Execution id should not be changed within mlmd_execution_atomic_op"
                      " context."
                  )
  
              # Orchestrator code will only update top-level fields and properties/
              # custom properties with diffs.
  
              # Motivation: to allow non-orchestrator code (specifically, pipeline tags
              # and labels) to modify execution custom properties while the orchestrator
              # is running. Delta changes are only applied for masked properties /
              # custom properties. execution.last_known_state will always be updated.
  
              # It enables orchestrator and non-orchestrator codes to run concurrently
              # as long as there are no overlaps in the modified fields.
  
              # Make a copy before writing to cache as the yielded `execution_copy`
              # object may be modified even after exiting the contextmanager.
              _execution_cache.put_execution(
                  mlmd_handle,
                  copy.deepcopy(execution_copy),
                  get_field_mask_paths(execution, execution_copy),
              )
              if on_commit is not None:
                  pre_commit_execution = copy.deepcopy(execution)
                  post_commit_execution = copy.deepcopy(
                      _execution_cache.get_execution(
                          mlmd_handle, execution_copy.id
                      )
                  )
                  on_commit(pre_commit_execution, post_commit_execution)
  
imports: import collections
  import contextlib
  import copy
  import threading
  import typing
  from typing import Callable, Iterator, MutableMapping, Optional
  import cachetools
  from tfx.orchestration import metadata
  from google.protobuf.internal import containers
  from ml_metadata.proto import metadata_store_pb2
label_desc: Use image node_id 1 for calling a global function with example usage: mlmd_execution_atomic_op(mlmd_handle, execution_id, on_commit) without return types