Error when saving the checkpoint

Hi! I met an issue when saving the checkpoint. I also commented my issue below #446

This issue occurred when I ended training for 100 steps and saved the checkpoint to the absolute path '/gpfsnyu/scratch/kf2395/jukemir_t5/pretrain/'

My Tensorstore is 0.1.19.

As [this solution](https://github.com/google-research/t5x/issues/446#issuecomment-1107954777) mentioned, the relative path may cause some problems, so I changed it to the absolute path.

When I used the relative path at the beginning, I got the similar error as [this comment](https://github.com/google-research/t5x/issues/446#issuecomment-1107946852). The error message:

`ValueError: Error opening "zarr" driver: Error reading local file "./pretrain_model/checkpoint_5000.tmp-1650694933/state.param_states.decoder.decoder_norm.scale.v/.zarray": Invalid key: "./pretrain_model/checkpoint_5000.tmp-1650694933/state.param_states.decoder.decoder_norm.scale.v/.zarray"
In call to configurable 'train' ())`

Then I changed the path to the absolute path and the issue above was solved. **But a new issue occurred.**

`ValueError: Error opening "zarr" driver: Error writing local file "/gpfsnyu/scratch/kf2395/jukemir_t5/pretrain/checkpoint_100.tmp-1650949633/state.param_states.decoder.layers_0.pre_cross_attention_layer_norm.scale.v/.zarray": Failed to acquire lock on file: /gpfsnyu/scratch/kf2395/jukemir_t5/pretrain/checkpoint_100.tmp-1650949633/state.param_states.decoder.layers_0.pre_cross_attention_layer_norm.scale.v/.zarray.__lock [OS error: Invalid argument]
In call to configurable 'train' ()`

I tried to delete all the files in '/gpfsnyu/scratch/kf2395/jukemir_t5/pretrain/' and trained again. But this issue still existed.

The detailed error message:

I0426 13:05:36.808195 140074531202880 train.py:516] Epoch 0 of 10000
I0426 13:05:36.808564 140055117031168 logging_writer.py:48] [0] collection=train timing/compilation_seconds=160.272345
I0426 13:05:36.828166 140074531202880 train.py:522] BEGIN Train loop.
I0426 13:05:36.828350 140074531202880 train.py:527] Training for 100 steps.
I0426 13:05:36.833504 140074531202880 trainer.py:517] Training: step 0
I0426 13:05:47.585027 140074531202880 trainer.py:517] Training: step 12
I0426 13:05:58.556400 140074531202880 trainer.py:517] Training: step 23
I0426 13:06:09.237899 140074531202880 trainer.py:517] Training: step 34
I0426 13:06:19.734536 140074531202880 trainer.py:517] Training: step 45
I0426 13:06:30.668152 140074531202880 trainer.py:517] Training: step 56
I0426 13:06:41.496444 140074531202880 trainer.py:517] Training: step 67
I0426 13:06:52.412244 140074531202880 trainer.py:517] Training: step 78
I0426 13:07:03.236425 140074531202880 trainer.py:517] Training: step 89
I0426 13:07:13.692245 140074531202880 train.py:550] END Train loop.
I0426 13:07:13.727353 140055117031168 logging_writer.py:48] [100] collection=train accuracy=0.12926435470581055, cross_ent_loss=3456.254063, cross_ent_loss_per_all_target_tokens=0.337525, learning_rate=0.001000000280328095, learning_rate/current=0.0010000000474974513, loss=3460.679688, loss_per_all_target_tokens=0.337957, loss_per_nonpadding_target_token=5.071336, nonpadding_fraction=0.066641, timing/seconds=96.861853, timing/seqs=1000, timing/seqs_per_second=10.323982, timing/seqs_per_second_per_core=10.323982, timing/steps_per_second=1.032398, timing/target_tokens_per_second=10571.757297, timing/target_tokens_per_second_per_core=10571.757297, z_loss=4.426097, z_loss_per_all_target_tokens=0.000432
I0426 13:07:13.728666 140074531202880 train.py:565] Saving checkpoint.
I0426 13:07:13.730171 140074531202880 checkpoints.py:631] Saving checkpoint for step 100 to /gpfsnyu/scratch/kf2395/jukemir_t5/pretrain/checkpoint_100.tmp-1650949633
Traceback (most recent call last):
File "/gpfsnyu/scratch/kf2395/.cache/env/tf2-gpu-py3.7/lib/python3.7/runpy.py", line 193, in _run_module_as_main
"main", mod_spec)
File "/gpfsnyu/scratch/kf2395/.cache/env/tf2-gpu-py3.7/lib/python3.7/runpy.py", line 85, in _run_code
exec(code, run_globals)
File "/gpfsnyu/scratch/kf2395/jukemir_t5/t5x/train.py", line 663, in
gin_utils.run(main)
File "/gpfsnyu/scratch/kf2395/jukemir_t5/t5x/gin_utils.py", line 107, in run
flags_parser=lambda a: app.parse_flags_with_usage(rewrite_gin_args(a)))
File "/gpfsnyu/scratch/kf2395/.cache/env/tf2-gpu-py3.7/lib/python3.7/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/gpfsnyu/scratch/kf2395/.cache/env/tf2-gpu-py3.7/lib/python3.7/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "/gpfsnyu/scratch/kf2395/jukemir_t5/t5x/train.py", line 641, in main
_main(argv)
File "/gpfsnyu/scratch/kf2395/jukemir_t5/t5x/train.py", line 661, in _main
train_using_gin()
File "/gpfsnyu/scratch/kf2395/.cache/env/tf2-gpu-py3.7/lib/python3.7/site-packages/gin/config.py", line 1605, in gin_wrapper
utils.augment_exception_message_and_reraise(e, err_str)
File "/gpfsnyu/scratch/kf2395/.cache/env/tf2-gpu-py3.7/lib/python3.7/site-packages/gin/utils.py", line 41, in augment_exception_message_and_reraise
raise proxy.with_traceback(exception.traceback) from None
File "/gpfsnyu/scratch/kf2395/.cache/env/tf2-gpu-py3.7/lib/python3.7/site-packages/gin/config.py", line 1582, in gin_wrapper
return fn(*new_args, **new_kwargs)
File "/gpfsnyu/scratch/kf2395/jukemir_t5/t5x/train.py", line 568, in train
checkpoint_cfg.save.state_transformation_fns)
File "/gpfsnyu/scratch/kf2395/jukemir_t5/t5x/checkpoints.py", line 639, in save
tmp_dir, train_state, concurrent_gb, state_transformation_fns)
File "/gpfsnyu/scratch/kf2395/jukemir_t5/t5x/checkpoints.py", line 806, in _write_state_to_tensorstore
written_state_dict = _run_future_tree(future_written_state)
File "/gpfsnyu/scratch/kf2395/jukemir_t5/t5x/checkpoints.py", line 167, in _run_future_tree
leaves = loop.run_until_complete(asyncio.gather(*future_leaves))
File "/gpfsnyu/scratch/kf2395/.cache/env/tf2-gpu-py3.7/lib/python3.7/asyncio/base_events.py", line 587, in run_until_complete
return future.result()
File "/gpfsnyu/scratch/kf2395/jukemir_t5/t5x/checkpoints.py", line 770, in _write_array
'limit': 128
ValueError: Error opening "zarr" driver: Error writing local file "/gpfsnyu/scratch/kf2395/jukemir_t5/pretrain/checkpoint_100.tmp-1650949633/state.param_states.decoder.layers_0.pre_cross_attention_layer_norm.scale.v/.zarray": Failed to acquire lock on file: /gpfsnyu/scratch/kf2395/jukemir_t5/pretrain/checkpoint_100.tmp-1650949633/state.param_states.decoder.layers_0.pre_cross_attention_layer_norm.scale.v/.zarray.__lock [OS error: Invalid argument]
In call to configurable 'train' ()

Thank you for your kindly help!

Start with t5x/checkpoints.py, especially save, _write_state_to_tensorstore, and _write_array, then compare the failure with the training flow in t5x/train.py. Reproduce checkpoint saving with the reported absolute path and TensorStore 0.1.19; done means checkpoint_100 can be written without the Zarr lock error.

Written by the indexing model from the issue text.