Add a delete function to FSStoragePlugin, S3StoragePlugin and GCSStoragePlugin

Summary: We add a `remove` function to the `FSStoragePlugin` in order to implement a delete API for TorchSnapshot. Similar functions for `GCSStoragePlugin` and `S3StoragePlugin` and their integration with `Snapshot` will be added in subsequent diffs.

Reviewed By: yifuwang

Differential Revision: D36948787

fbshipit-source-id: fc37336e5738a9122bb72f68afe192491faceb13

Author: tangbinh

tests/test_fs_storage_plugin.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import asyncio
+import logging
+import os
+import tempfile
+import unittest
+
+import torch
+import torchsnapshot
+from torchsnapshot.storage_plugins.fs import FSStoragePlugin
+
+logger = logging.getLogger(__name__)
+
+_TENSOR_SZ = int(100_000_000 / 4)
+
+
+class FSStoragePluginTest(unittest.TestCase):
+    def test_write_read_delete(self) -> None:
+        with tempfile.TemporaryDirectory() as path:
+            logger.info(path)
+            plugin = FSStoragePlugin(root=path)
+
+            tensor = torch.rand((_TENSOR_SZ,))
+            tensor_path = os.path.join(path, "tensor")
+            write_req = torchsnapshot.io_types.IOReq(path=tensor_path)
+            torch.save(tensor, write_req.buf)
+            asyncio.run(plugin.write(io_req=write_req))
+            self.assertTrue(os.path.exists(tensor_path))
+
+            read_req = torchsnapshot.io_types.IOReq(path=tensor_path)
+            asyncio.run(plugin.read(io_req=read_req))
+            loaded = torch.load(read_req.buf)
+            self.assertTrue(torch.allclose(tensor, loaded))
+
+            asyncio.run(plugin.delete(path=tensor_path))
+            self.assertFalse(os.path.exists(tensor_path))
+            plugin.close()

tests/test_gcs_storage_plugin.py

@@ -38,20 +38,21 @@ def test_read_write_via_snapshot(self) -> None:
         self.assertTrue(torch.allclose(tensor, app_state["state"]["tensor"]))
 
     @unittest.skipIf(os.environ.get("TORCHSNAPSHOT_ENABLE_GCP_TEST") is None, "")
-    def test_write_read(self) -> None:
+    def test_write_read_delete(self) -> None:
         path = f"{_TEST_BUCKET}/{uuid.uuid4()}"
         logger.info(path)
         plugin = GCSStoragePlugin(root=path)
 
         tensor = torch.rand((_TENSOR_SZ,))
-        write_req = torchsnapshot.io_types.IOReq(path=os.path.join(path, "tensor"))
+        path = os.path.join(path, "tensor")
+        write_req = torchsnapshot.io_types.IOReq(path=path)
         torch.save(tensor, write_req.buf)
         asyncio.run(plugin.write(io_req=write_req))
 
-        read_req = torchsnapshot.io_types.IOReq(path=os.path.join(path, "tensor"))
+        read_req = torchsnapshot.io_types.IOReq(path=path)
         asyncio.run(plugin.read(io_req=read_req))
         loaded = torch.load(read_req.buf)
-
         self.assertTrue(torch.allclose(tensor, loaded))
 
+        asyncio.run(plugin.delete(path=path))
         plugin.close()

tests/test_s3_storage_plugin.py

@@ -38,20 +38,21 @@ def test_read_write_via_snapshot(self) -> None:
         self.assertTrue(torch.allclose(tensor, app_state["state"]["tensor"]))
 
     @unittest.skipIf(os.environ.get("TORCHSNAPSHOT_ENABLE_AWS_TEST") is None, "")
-    def test_write_read(self) -> None:
+    def test_write_read_delete(self) -> None:
         path = f"{_TEST_BUCKET}/{uuid.uuid4()}"
         logger.info(path)
         plugin = S3StoragePlugin(root=path)
 
         tensor = torch.rand((_TENSOR_SZ,))
-        write_req = torchsnapshot.io_types.IOReq(path=os.path.join(path, "tensor"))
+        path = os.path.join(path, "tensor")
+        write_req = torchsnapshot.io_types.IOReq(path=path)
         torch.save(tensor, write_req.buf)
         asyncio.run(plugin.write(io_req=write_req))
 
-        read_req = torchsnapshot.io_types.IOReq(path=os.path.join(path, "tensor"))
+        read_req = torchsnapshot.io_types.IOReq(path=path)
         asyncio.run(plugin.read(io_req=read_req))
         loaded = torch.load(read_req.buf)
-
         self.assertTrue(torch.allclose(tensor, loaded))
 
+        asyncio.run(plugin.delete(path=path))
         plugin.close()

torchsnapshot/io_types.py

@@ -41,6 +41,10 @@ async def write(self, io_req: IOReq) -> None:
     async def read(self, io_req: IOReq) -> None:
         pass
 
+    @abc.abstractmethod
+    async def delete(self, path: str) -> None:
+        pass
+
     @abc.abstractmethod
     def close(self) -> None:
         pass

torchsnapshot/storage_plugins/fs.py

@@ -11,6 +11,7 @@
 from typing import Set
 
 import aiofiles
+import aiofiles.os
 from torchsnapshot.io_types import IOReq, StoragePlugin
 
 
@@ -36,5 +37,8 @@ async def read(self, io_req: IOReq) -> None:
         async with aiofiles.open(path, "rb") as f:
             io_req.buf = io.BytesIO(await f.read())
 
+    async def delete(self, path: str) -> None:
+        await aiofiles.os.remove(path)
+
     def close(self) -> None:
         pass

torchsnapshot/storage_plugins/gcs.py

@@ -57,5 +57,11 @@ async def read(self, io_req: IOReq) -> None:
         )
         io_req.buf.seek(0)
 
+    async def delete(self, path: str) -> None:
+        loop = asyncio.get_running_loop()
+        key = os.path.join(self.root, path)
+        blob = self.bucket.blob(key)
+        await loop.run_in_executor(self.executor, blob.delete)
+
     def close(self) -> None:
         self.executor.shutdown()

torchsnapshot/storage_plugins/s3.py

@@ -44,5 +44,10 @@ async def read(self, io_req: IOReq) -> None:
             async with response["Body"] as stream:
                 io_req.buf = io.BytesIO(await stream.read())
 
+    async def delete(self, path: str) -> None:
+        async with self.session.create_client("s3") as client:
+            key = os.path.join(self.root, path)
+            await client.delete_object(Bucket=self.bucket, Key=key)
+
     def close(self) -> None:
         pass