CheckpointManager using deprecated legacy API

[jecampagne]

Dear developpers
I've tryed to follow use_checkpointing doc (latest). with the following snippet on Google Colab (CPU):

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
from typing import Optional, Any
import shutil

import numpy as np
import jax
from jax import random, numpy as jnp

import flax
from flax import linen as nn
from flax.training import checkpoints, train_state
from flax import struct, serialization
import orbax.checkpoint

import optax

from functools import partial

ckpt_dir = '/tmp/flax_ckpt'

if os.path.exists(ckpt_dir):
    shutil.rmtree(ckpt_dir)

TrainState = train_state.TrainState
@partial(jax.pmap, static_broadcasted_argnums=(1, 2))
def create_state(rng, model_cls, input_shape):
  r"""Create the training state given a model class. """
  model = model_cls()

  tx = optax.sgd(learning_rate=1e-3)

  variables = model.init(rng, jnp.ones(input_shape))

  state = TrainState.create(apply_fn=model.apply, tx=tx,
      params=variables['params'])

  return state

class Model(nn.Module):
  dtype: Any = jnp.float32

  @nn.compact
  def __call__(self, x: jnp.ndarray, train: bool = True):
    x = nn.Dense(features=3)(x)
    return x

from flax.training.common_utils import shard, shard_prng_key

key = jax.random.PRNGKey(seed=0)
key_denoiser, key = jax.random.split(key, 2)
key_denoiser = shard_prng_key(key_denoiser)

data = jax.random.normal(key, (5,5))

model_state = create_state(key_denoiser, Model,
    (10, *data.shape[1:]))
# Perform a simple gradient update similar to the one during a normal training workflow.
model_state = model_state.apply_gradients(grads=jax.tree_util.tree_map(jnp.ones_like, model_state.params))

config = {'dimensions': np.array([5, 3])}  # just a fake config to follow the doc

ckpt = {'model': model_state, 'config': config, 'data': [data]}

from flax.training import orbax_utils

orbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer()
save_args = orbax_utils.save_args_from_target(ckpt)
orbax_checkpointer.save('/tmp/flax_ckpt/orbax/single_save', ckpt, save_args=save_args)

ls '/tmp/flax_ckpt/orbax/single_save'
#answer: 
# _CHECKPOINT_METADATA  d/  manifest.ocdbt  _METADATA  ocdbt.process_0/  _sharding

options = orbax.checkpoint.CheckpointManagerOptions(max_to_keep=2, create=True)
checkpoint_manager = orbax.checkpoint.CheckpointManager(
    '/tmp/flax_ckpt/orbax/managed', orbax_checkpointer, options)

# Inside a training loop
for step in range(5):
    # ... do your training
    checkpoint_manager.save(step, ckpt, save_kwargs={'save_args': save_args})

I get this message:

WARNING:absl:Configured `CheckpointManager` using deprecated legacy API. Please follow the instructions at https://orbax.readthedocs.io/en/latest/api_refactor.html to migrate.

However the link is broken.

Now

os.listdir('/tmp/flax_ckpt/orbax/managed')

gives

['4', '3']

and then

ls '/tmp/flax_ckpt/orbax/managed/4/default'

gives

d/  manifest.ocdbt  _METADATA  ocdbt.process_0/  _sharding

and

step = checkpoint_manager.latest_step()  # step = 4
checkpoint_manager.restore(step)

seems ok.

1. So what is going on with the Warning message and the broken link?

2. When I would like to restore the model parameter for exemple a test session, do I have to create a 8-devices session as in the training session?

Thanks