Async save checkpoints not working

[fecet]

I tried a few ways to try the async_manager feature, however all of them seems not work and my task still block as usual, here is my minimal example

rng=jax.random.PRNGKey(64)
bigarray = jax.random.normal(rng, (30000, 10000))
a=b=jax.random.normal(rng, (10000,1000))
DATA_PATH=Path()

# %%

async_manager = checkpoints.AsyncManager()
for i in tqdm(range(int(1e6))):

    # state, trainer.rng, loss = trainer.train_step(state, trainer.rng, batch)
    [email protected]

    if (i + 1) % 1000 == 0:
        # ckpt_save()
        # joblib.dump(bigarray, DATA_PATH / "ba")
        checkpoints.save_checkpoint(
            DATA_PATH,
            bigarray,
            step=i,
            prefix=f"checkpoint_",
            overwrite=True,
            async_manager=async_manager,
        )

May I do it wrong? As I can not find any example relevant except the simple comment in the source code.

[cgarciae]

Hey I simplified your example and could not find any errors:

import jax
from pathlib import Path
from flax.training import checkpoints

async_manager = checkpoints.AsyncManager()
for i in range(3):
    bigarray = jax.random.normal(jax.random.PRNGKey(i), (30000, 10000))
    checkpoints.save_checkpoint(
        Path(),
        bigarray,
        step=i,
        prefix=f"checkpoint_",
        overwrite=True,
        async_manager=async_manager,
    )

Can you try to pin-point the error you are observing?

[fecet]

Hey I simplified your example and could not find any errors:

import jax
from pathlib import Path
from flax.training import checkpoints

async_manager = checkpoints.AsyncManager()
for i in range(3):
    bigarray = jax.random.normal(jax.random.PRNGKey(i), (30000, 10000))
    checkpoints.save_checkpoint(
        Path(),
        bigarray,
        step=i,
        prefix=f"checkpoint_",
        overwrite=True,
        async_manager=async_manager,
    )

Can you try to pin-point the error you are observing?

Yeah, this example works. My point is the wall time keep the same with or without async_manager.

[cgarciae]

I ran very basic bechmark to find the impact of using async_manager.

Code

import jax
from pathlib import Path
from flax.training import checkpoints
from timeit import default_timer as timer
from tempfile import TemporaryDirectory

bigarray = jax.random.normal(jax.random.PRNGKey(0), (30000, 10000)).block_until_ready()
async_manager = checkpoints.AsyncManager()
with TemporaryDirectory() as tmpdir:
    i = 0
    start = timer()
    checkpoints.save_checkpoint(
        Path(f"{tmpdir}/checkpoint_{i}"),
        bigarray,
        step=i,
        async_manager=async_manager, # comment to run without async_manager
    )
    end = timer()
    print(f"checkpoint{i} saved in {end - start} seconds")

Using async_manager it takes ~3.5s and without it it takes ~4s, you will probably se more gains with heavier structures. However, if you run save_checkpoint calls with async_manager consecutively you will not see any gains because as its stated in the docs:

Note that an ongoing save will still block subsequent saves, to make sure overwrite/keep logic works correctly.