Add CNS tests for partial saving. Also modify async file writes to unlink before writing.

PiperOrigin-RevId: 874287343

Author: BlaziusMaximus

checkpoint/orbax/checkpoint/_src/path/async_path.py

@@ -50,15 +50,31 @@ def _mkdir_sync(**thread_kwargs):
 
 
 async def write_bytes(path: epath.Path, data: Any) -> int:
-  return await asyncio.to_thread(path.write_bytes, data)
+
+  def _write():
+    try:
+      path.unlink()
+    except OSError:
+      pass
+    return path.write_bytes(data)
+
+  return await asyncio.to_thread(_write)
 
 
 async def read_bytes(path: epath.Path) -> bytes:
   return await asyncio.to_thread(path.read_bytes)
 
 
 async def write_text(path: epath.Path, text: str) -> int:
-  return await asyncio.to_thread(path.write_text, text)
+
+  def _write():
+    try:
+      path.unlink()
+    except OSError:
+      pass
+    return path.write_text(text)
+
+  return await asyncio.to_thread(_write)
 
 
 async def read_text(path: epath.Path) -> str: