class_docstr: Interface for checkpointing GDAs asynchronously.
  
  This class manages the state of an ongoing asynchronous checkpoint.
  
  For example, say a checkpoint happens on every step. If you checkpoint on
  step 1 and after some computation the model is on checkpoint 2. But step 1's
  checkpoint hasn't finished committing to the storage layer yet. So until that
  is finished, checkpoint for step 2 will need to be blocked. Maintaining a
  class allows to maintain that state.
  
  Example:
  
  Below is a simplified training loop:
  
  ```
  # Call this at the start of your program.
  jax.distributed.initialize()
  
  manager = GlobalAsyncCheckpointManager()
  
  # Restore checkpoint if available or initialize the train_state from
  # init_fn().
  train_state = manager.deserialize(...)
  
  while ...:
    if step % num_steps_between_checkpoints == 0:
      manager.serialize(train_state, temp_checkpoint_dir=...,
                        final_checkpoint_dir=...)
      train_state = train_step(train_state, input)
      # This is a non-blocking call.
      manager.check_for_errors()
  
  manager.serialize(train_state, temp_checkpoint_dir=...,
                    final_checkpoint_dir=...)
  # Wait before the end of the program for the checkpoint to finish. This is a
  # blocking call.
  manager.wait_until_finished()
  ```
class_docstr_tok: ["Interface","for","checkpointing","GDAs","asynchronously",".","This","class","manages","the","state","of","an","ongoing","asynchronous","checkpoint",".","For","example",",","say","a","checkpoint","happens","on","every","step",".","If","you","checkpoint","on","step","1","and","after","some","computation","the","model","is","on","checkpoint","2",".","But","step","1's","checkpoint","has","n't","finished","committing","to","the","storage","layer","yet",".","So","until","that","is","finished",",","checkpoint","for","step","2","will","need","to","be","blocked",".","Maintaining","a","class","allows","to","maintain","that","state",".","Example",":","Below","is","a","simplified","training","loop",":","``","`","#","Call","this","at","the","start","of","your","program",".","jax.distributed.initialize","(",")","manager","=","GlobalAsyncCheckpointManager","(",")","#","Restore","checkpoint","if","available","or","initialize","the","train_state","from","#","init_fn","(",")",".","train_state","=","manager.deserialize","(","...",")","while","...",":","if","step","%","num_steps_between_checkpoints","==","0",":","manager.serialize","(","train_state",",","temp_checkpoint_dir=","...",",","final_checkpoint_dir=","...",")","train_state","=","train_step","(","train_state",",","input",")","#","This","is","a","non-blocking","call",".","manager.check_for_errors","(",")","manager.serialize","(","train_state",",","temp_checkpoint_dir=","...",",","final_checkpoint_dir=","...",")","#","Wait","before","the","end","of","the","program","for","the","checkpoint","to","finish",".","This","is","a","#","blocking","call",".","manager.wait_until_finished","(",")","``","`"]
docstr: Checks if any errors have been raised in the child thread.
  
  This is a non-blocking call that can be called in the main thread.
code: def check_for_errors(self):
      """Checks if any errors have been raised in the child thread.
  
      This is a non-blocking call that can be called in the main thread.
      """
  
imports: from __future__ import annotations
  import abc
  import asyncio
  from collections.abc import Awaitable, Sequence
  import itertools
  import logging
  from functools import partial
  import os
  import re
  import time
  import threading
  from typing import Any, Callable, Optional, Union
  import jax
  from jax._src import distributed
  from jax._src import array
  from jax._src import sharding
  from jax._src import sharding_impls
  from jax._src import typing
  from jax._src.lib import xla_extension
  import jax.numpy
  import numpy
  import tensorstore
label_desc: Use image node_id 1 for calling the GlobalAsyncCheckpointManagerBase obj's underlying member method code with example usage: obj.check_for_errors() without return types