How to use AsyncCheckpointer with CheckpointManager

[davidshen84]

Hi,

Based on this code snippet from this document, the manager will automatically figure out how to use the async checkpointer. Is that correct?

path = epath.Path('/tmp/async_checkpoint_manager')
ckpt_mngr = ocp.CheckpointManager(path)

def train_step(step, state):
  # update state values accordingly
  return step + 1, state

step = 0
num_steps = 5
while step < num_steps:
  ckpt_mngr.save(step, args=ocp.args.StandardSave(train_state))
  step, train_state = train_step(step, train_state)

ckpt_mngr.wait_until_finished()

What about multi-item checkpointing? Do we need to define arguments like the item_names and item_handlers?

[cpgaffney1]

Apologies for not spotting this question, undoubtedly you have resolved the question by now. For posterity though, yes, async saving behavior is used by default.

A more recent change has made it so that you do not need to define item_names and item_handlers in advance. For multi-item checkpointing, just specify args=ocp.args.Composite(state=..., dataset=..., ...).

[davidshen84]

The document website has changed so much, the link in my original question is gone.

I think this link is a suitable replacement and also provides the answer.