Repository: google/flax

Issue:
Issue during removing checkpoint on TPU
Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.

### System information
- TPU Version: tpu-vm-base
- TPU model: v2-8, v3-8


### Problem you have encountered:
Followed the [instruction](https://github.com/google/flax/tree/main/examples/wmt) to run `wmt` model on TPU, and process failed when removing the previous checkpoint. The error message is `doesn't exist or not a directory`. I think this is related to the [commit](https://github.com/google/flax/commit/40e8e62c212ecb29f7ce36b9a228430060705592) `Switch to Orbax for Flax single-checkpoint support under the hood` that changed `io.remove` to `io.rmtree`. 



### Logs, error messages, etc:

```
2022-11-29 17:00:08.841 PST Saved checkpoint at gs://xl-ml-test-us-central1/output/flax-latest/wmt/func/v2-8/flax-latest-wmt-func-v2-8-1vm-r29s8/checkpoint_9
2022-11-29 17:00:08.922 PST Removing checkpoint at gs://xl-ml-test-us-central1/output/flax-latest/wmt/func/v2-8/flax-latest-wmt-func-v2-8-1vm-r29s8/checkpoint_0
2022-11-29 17:00:09.138 PST Traceback (most recent call last): File "main.py", line 66, in <module> app.run(main) File "/home/xl-ml-test/.local/lib/python3.8/site-packages/absl/app.py", line 308, in run _run_main(main, args) File "/home/xl-ml-test/.local/lib/python3.8/site-packages/absl/app.py", line 254, in _run_main sys.exit(main(argv)) File "main.py", line 61, in main train.train_and_evaluate(FLAGS.config, FLAGS.workdir) File "/home/xl-ml-test/flax/examples/wmt/train.py", line 626, in train_and_evaluate checkpoints.save_checkpoint(workdir, jax_utils.unreplicate(state), File "/home/xl-ml-test/flax/flax/training/checkpoints.py", line 574, in save_checkpoint save_main_ckpt_task() File "/home/xl-ml-test/flax/flax/training/checkpoints.py", line 568, in save_main_ckpt_task return _save_main_ckpt_file(target, False, (ckpt_tmp_path, ckpt_path), File "/home/xl-ml-test/flax/flax/training/checkpoints.py", line 466, in _save_main_ckpt_file _save_commit( File "/home/xl-ml-test/flax/flax/training/checkpoints.py", line 425, in _save_commit _remove_invalid_ckpts(ckpt_path, base_path, keep, overwrite, File "/home/xl-ml-test/flax/flax/training/checkpoints.py", line 385, in _remove_invalid_ckpts io.rmtree(path) File "/home/xl-ml-test/flax/flax/io.py", line 174, in rmtree return gfile.rmtree(path) File "/home/xl-ml-test/.local/lib/python3.8/site-packages/tensorflow/python/lib/io/file_io.py", line 676, in delete_recursively_v2 _pywrap_file_io.DeleteRecursively(compat.path_to_bytes(path)) tensorflow.python.framework.errors_impl.NotFoundError: gs://xl-ml-test-us-central1/output/flax-latest/wmt/func/v2-8/flax-latest-wmt-func-v2-8-1vm-r29s8/checkpoint_0 doesn't exist or not a directory.

```


### Steps to reproduce:

```
git clone https://github.com/google/flax
cd flax
pip install -e .
cd examples/wmt

python3 main.py --workdir=$(MODEL_DIR)  --config=configs/default.py --config.num_train_steps=10 --config.per_device_batch_size=8
```



Discussion:
@IvyZX could you please take a look at this?
Thanks for raising the bug. It seems like `io.rmtree` and `io.remove` behaves differently on GCS. I will make a fix today and test it out on gcloud. Meanwhile feel free to use the current official release version (as from `pip install -U flax`).

Resolving patch:
diff --git a/flax/io.py b/flax/io.py
--- a/flax/io.py
+++ b/flax/io.py
@@ -159,6 +159,7 @@ def glob(pattern):
 
 
 def remove(path):
+  """Remove the file at path. Might fail if used on a directory path."""
   if io_mode == BackendMode.DEFAULT:
     return os.remove(path)
   elif io_mode == BackendMode.TF:
@@ -168,6 +169,7 @@ def remove(path):
 
 
 def rmtree(path):
+  """Remove a directory and recursively all contents inside. Might fail if used on a file path."""
   if io_mode == BackendMode.DEFAULT:
     return shutil.rmtree(path)
   elif io_mode == BackendMode.TF:
diff --git a/flax/training/checkpoints.py b/flax/training/checkpoints.py
--- a/flax/training/checkpoints.py
+++ b/flax/training/checkpoints.py
@@ -110,6 +110,13 @@ def _allowempty_listdir(path: str):
   except tf_errors.NotFoundError:
     return []
 
+def _safe_remove(path: str):
+  """Identify whether a path is a dir or list and choose the correct remove method."""
+  if io.isdir(path):
+    io.rmtree(path)
+  else:
+    io.remove(path)
+
 class AsyncManager():
   """A simple object to track async checkpointing.
 
@@ -361,7 +368,7 @@ def _remove_invalid_ckpts(ckpt_path: str, base_path: str, keep: int,
         # checkpoint folder and before deleting the main checkpoint.
         if io.exists(path + MP_ARRAY_POSTFIX):
           io.rmtree(path + MP_ARRAY_POSTFIX)
-      io.rmtree(path)
+      _safe_remove(path)
 
   # Remove old checkpoint files.
   last_kept = -float('inf')
@@ -382,7 +389,7 @@ def _remove_invalid_ckpts(ckpt_path: str, base_path: str, keep: int,
         # MPA might be removed already but the main ckpt is still there.
         if io.exists(path + MP_ARRAY_POSTFIX):
           io.rmtree(path + MP_ARRAY_POSTFIX)
-      io.rmtree(path)
+      _safe_remove(path)
 
 
 def _save_commit(ckpt_tmp_path: str, ckpt_path: str, base_path: str, keep: int,


Test patch: