class_docstr: Helper class for async checkpoint.
class_docstr_tok: ["Helper","class","for","async","checkpoint","."]
docstr: Initialize the async checkpoint internal state.
docstr_tok: ["Initialize","the","async","checkpoint","internal","state","."]
code: def _ensure_initialized(self):
      """Initialize the async checkpoint internal state."""
      # This map will be used to store the CPU copy of all checkpointable objects
      self._object_map = object_identity.ObjectIdentityDictionary()
      self._tpu_embedding_objects = []
  
      # Populate self._all_tracakbles, but exclude the checkpoint instance itself
      # and its save_counter, as they will be returned by `descendants()`.
      exclude_set = object_identity.ObjectIdentitySet()
      exclude_set.add(self.checkpointer())
      exclude_set.add(self.checkpointer().save_counter)
      self._saveable_trackables, all_trackables = _get_all_trackables(
          root=self.checkpointer(), exclude_set=exclude_set
      )
  
      # Handle special cases: TPU Embedding, and slot variables.
      # 1. TPUEmbedding: Different from other trackables, TPUEmbedding needs to
      # call `_retrieve_variables` to checkpoint, while populating a dummy copy to
      # the object map.
      # 2. Slot variables: they need to be handled differently as they cannot be
      # retrieved from `TrackableView.descendants()`.
  
      # Note: dir() is used rather than hasattr() here to avoid triggering
      # custom __getattr__ code, see b/152031870 for context.
      for t in all_trackables:
          # Special case 1: TPU Embedding, populate object_map here
          # Special case 1: Handle TPU Embedding by addnig a dummy instance to the
          # object map. Also add TPUEmbedding to separate list for special handling
          # with values copy.
          if hasattr(type(t), _TPU_EMBEDDING_ATTR):
              self._handle_tpu_embedding(t)
          # Special case 2: handle slot variables. The object_map is populated later
          # when the variable values are being copied to host CPU for the first
          # time.
          if "get_slot_names" in dir(t):
              slot_names = t.get_slot_names()
              for slot_name in slot_names:
                  for original_variable in all_trackables:
                      if not isinstance(
                          original_variable, variables.Variable
                      ):
                          continue
                      try:
                          # Usage of hasattr may result in KeyError
                          original_slot_variable = t.get_slot(
                              original_variable, slot_name
                          )
                      except (AttributeError, KeyError):
                          continue
                      if isinstance(
                          original_slot_variable, base.Trackable
                      ):
                          self._saveable_trackables.append(
                              original_slot_variable
                          )
  
      # Initiate the underlying Checkpoint instance's save_counter.
      save_counter = self.checkpointer().save_counter.numpy()
      logging.info(
          "Initializing async checkpoint's save_counter: %d",
          save_counter,
      )
  
      # Pass the object map of the copied variables to the underlying Checkpoint.
      self.checkpointer()._saver._object_map = (
          self._object_map
      )  # pylint: disable=protected-access
  
      # We perform a `_copy_to_cpu()` to populate `self._object_map`,
      # initializing copies. We do not call `self._copy_to_cpu()` directly
      # because it is a tf function, which leads to access out of scope error.
  
      # TODO(charlieruan) Figure out a better work around to solve the access
      # out of scope error.
      for t in self._saveable_trackables:
          try:
              t._copy_trackable_to_cpu(
                  object_map=self._object_map
              )  # pylint: disable=protected-access
          except NotImplementedError as e:
              logging.warning("Trackable %s skipped due to: %s", t, e)
  
      for tpu_embedding in self._tpu_embedding_objects:
          tpu_embedding._retrieve_variables()  # pylint: disable=protected-access
  
      # Initiate the async thread for checkpoint saving.
      self._async_save_thread = threading.Thread(
          target=self._async_save, daemon=True
      )
      self._async_save_thread.start()
  
      self._initialized = True
  
imports: import atexit
  import copy
  import queue
  import threading
  import time
  import weakref
  from absl import logging
  from tensorflow.python.checkpoint import checkpoint_context
  from tensorflow.python.checkpoint import trackable_view
  from tensorflow.python.distribute import device_util
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.eager import executor
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import variables
  from tensorflow.python.saved_model.pywrap_saved_model import metrics
  from tensorflow.python.trackable import base
  from tensorflow.python.util import object_identity
label_desc: Use image node_id 4 for calling the AsyncCheckpointHelper obj's underlying member method code with example usage: obj._ensure_initialized() without return types