(app) Add s3 drive type (1/2) (#14002)

* Add S3 protocol and optimization field to the drive object

* Add a list of drives to the work specification

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add only protocol for s3 drives, no optimization arguments, and add tests

* added trailing slash criteria

* allow slash in s3 drives

* fix

* fixed test issues

Co-authored-by: Panos Lantavos-Stratigakis <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Rick Izzo <[email protected]>
Co-authored-by: Jirka Borovec <[email protected]>
Co-authored-by: Rick Izzo <[email protected]>

Author: 5 people

src/lightning_app/storage/drive.py

@@ -13,7 +13,7 @@
 class Drive:
 
     __IDENTIFIER__ = "__drive__"
-    __PROTOCOLS__ = ["lit://"]
+    __PROTOCOLS__ = ["lit://", "s3://"]
 
     def __init__(
         self,
@@ -35,15 +35,28 @@ def __init__(
             root_folder: This is the folder from where the Drive perceives the data (e.g this acts as a mount dir).
         """
         self.id = None
+        self.protocol = None
         for protocol in self.__PROTOCOLS__:
             if id.startswith(protocol):
                 self.protocol = protocol
                 self.id = id.replace(protocol, "")
+                break
+        else:  # N.B. for-else loop
+            raise ValueError(
+                f"Unknown protocol for the drive 'id' argument '{id}`. The 'id' string "
+                f"must start with one of the following prefixes {self.__PROTOCOLS__}"
+            )
+
+        if self.protocol == "s3://" and not self.id.endswith("/"):
+            raise ValueError(
+                "S3 drives must end in a trailing slash (`/`) to indicate a folder is being mounted. "
+                f"Recieved: '{id}'. Mounting a single file is not currently supported."
+            )
 
         if not self.id:
             raise Exception(f"The Drive id needs to start with one of the following protocols: {self.__PROTOCOLS__}")
 
-        if "/" in self.id:
+        if self.protocol != "s3://" and "/" in self.id:
             raise Exception(f"The id should be unique to identify your drive. Found `{self.id}`.")
 
         self.root_folder = pathlib.Path(root_folder).resolve() if root_folder else os.getcwd()
@@ -75,6 +88,10 @@ def put(self, path: str) -> None:
             raise Exception("The component name needs to be known to put a path to the Drive.")
         if _is_flow_context():
             raise Exception("The flow isn't allowed to put files into a Drive.")
+        if self.protocol == "s3://":
+            raise PermissionError(
+                "S3 based drives cannot currently add files via this API. Did you mean to use `lit://` drives?"
+            )
 
         self._validate_path(path)
 
@@ -98,6 +115,10 @@ def list(self, path: Optional[str] = ".", component_name: Optional[str] = None)
         """
         if _is_flow_context():
             raise Exception("The flow isn't allowed to list files from a Drive.")
+        if self.protocol == "s3://":
+            raise PermissionError(
+                "S3 based drives cannot currently list files via this API. Did you mean to use `lit://` drives?"
+            )
 
         if component_name:
             paths = [
@@ -142,6 +163,10 @@ def get(
         """
         if _is_flow_context():
             raise Exception("The flow isn't allowed to get files from a Drive.")
+        if self.protocol == "s3://":
+            raise PermissionError(
+                "S3 based drives cannot currently get files via this API. Did you mean to use `lit://` drives?"
+            )
 
         if component_name:
             shared_path = self._to_shared_path(
@@ -189,6 +214,10 @@ def delete(self, path: str) -> None:
         """
         if not self.component_name:
             raise Exception("The component name needs to be known to delete a path to the Drive.")
+        if self.protocol == "s3://":
+            raise PermissionError(
+                "S3 based drives cannot currently delete files via this API. Did you mean to use `lit://` drives?"
+            )
 
         shared_path = self._to_shared_path(
             path,

tests/tests_app/storage/test_drive.py

@@ -11,7 +11,7 @@
 from lightning_app.utilities.component import _set_flow_context
 
 
-class SyncWorkA(LightningWork):
+class SyncWorkLITDriveA(LightningWork):
     def __init__(self, tmpdir):
         super().__init__()
         self.tmpdir = tmpdir
@@ -25,35 +25,35 @@ def run(self, drive: Drive):
         os.remove(f"{self.tmpdir}/a.txt")
 
 
-class SyncWorkB(LightningWork):
+class SyncWorkLITDriveB(LightningWork):
     def run(self, drive: Drive):
         assert not os.path.exists("a.txt")
         drive.get("a.txt")
         assert os.path.exists("a.txt")
 
 
-class SyncFlow(LightningFlow):
+class SyncFlowLITDrives(LightningFlow):
     def __init__(self, tmpdir):
         super().__init__()
         self.log_dir = Drive("lit://log_dir")
-        self.work_a = SyncWorkA(str(tmpdir))
-        self.work_b = SyncWorkB()
+        self.work_a = SyncWorkLITDriveA(str(tmpdir))
+        self.work_b = SyncWorkLITDriveB()
 
     def run(self):
         self.work_a.run(self.log_dir)
         self.work_b.run(self.log_dir)
         self._exit()
 
 
-def test_synchronization_drive(tmpdir):
+def test_synchronization_lit_drive(tmpdir):
     if os.path.exists("a.txt"):
         os.remove("a.txt")
-    app = LightningApp(SyncFlow(tmpdir))
+    app = LightningApp(SyncFlowLITDrives(tmpdir))
     MultiProcessRuntime(app, start_server=False).dispatch()
     os.remove("a.txt")
 
 
-class Work(LightningWork):
+class LITDriveWork(LightningWork):
     def __init__(self):
         super().__init__(parallel=True)
         self.drive = None
@@ -75,7 +75,7 @@ def run(self, *args, **kwargs):
         self.counter += 1
 
 
-class Work2(LightningWork):
+class LITDriveWork2(LightningWork):
     def __init__(self):
         super().__init__(parallel=True)
 
@@ -86,11 +86,11 @@ def run(self, drive: Drive, **kwargs):
         assert drive.list(".", component_name=self.name) == []
 
 
-class Flow(LightningFlow):
+class LITDriveFlow(LightningFlow):
     def __init__(self):
         super().__init__()
-        self.work = Work()
-        self.work2 = Work2()
+        self.work = LITDriveWork()
+        self.work2 = LITDriveWork2()
 
     def run(self):
         self.work.run("0")
@@ -102,15 +102,15 @@ def run(self):
             self._exit()
 
 
-def test_drive_transferring_files():
-    app = LightningApp(Flow())
+def test_lit_drive_transferring_files():
+    app = LightningApp(LITDriveFlow())
     MultiProcessRuntime(app, start_server=False).dispatch()
     os.remove("a.txt")
 
 
-def test_drive():
-    with pytest.raises(Exception, match="The Drive id needs to start with one of the following protocols"):
-        Drive("this_drive_id")
+def test_lit_drive():
+    with pytest.raises(Exception, match="Unknown protocol for the drive 'id' argument"):
+        Drive("invalid_drive_id")
 
     with pytest.raises(
         Exception, match="The id should be unique to identify your drive. Found `this_drive_id/something_else`."
@@ -213,19 +213,56 @@ def test_drive():
     os.remove("a.txt")
 
 
-def test_maybe_create_drive():
+def test_s3_drives():
+    drive = Drive("s3://foo/", allow_duplicates=True)
+    drive.component_name = "root.work"
 
-    drive = Drive("lit://drive_3", allow_duplicates=False)
+    with pytest.raises(
+        Exception, match="S3 based drives cannot currently add files via this API. Did you mean to use `lit://` drives?"
+    ):
+        drive.put("a.txt")
+    with pytest.raises(
+        Exception,
+        match="S3 based drives cannot currently list files via this API. Did you mean to use `lit://` drives?",
+    ):
+        drive.list("a.txt")
+    with pytest.raises(
+        Exception, match="S3 based drives cannot currently get files via this API. Did you mean to use `lit://` drives?"
+    ):
+        drive.get("a.txt")
+    with pytest.raises(
+        Exception,
+        match="S3 based drives cannot currently delete files via this API. Did you mean to use `lit://` drives?",
+    ):
+        drive.delete("a.txt")
+
+    _set_flow_context()
+    with pytest.raises(Exception, match="The flow isn't allowed to put files into a Drive."):
+        drive.put("a.txt")
+    with pytest.raises(Exception, match="The flow isn't allowed to list files from a Drive."):
+        drive.list("a.txt")
+    with pytest.raises(Exception, match="The flow isn't allowed to get files from a Drive."):
+        drive.get("a.txt")
+
+
+def test_create_s3_drive_without_trailing_slash_fails():
+    with pytest.raises(ValueError, match="S3 drives must end in a trailing slash"):
+        Drive("s3://foo")
+
+
+@pytest.mark.parametrize("drive_id", ["lit://drive", "s3://drive/"])
+def test_maybe_create_drive(drive_id):
+    drive = Drive(drive_id, allow_duplicates=False)
     drive.component_name = "root.work1"
     new_drive = _maybe_create_drive(drive.component_name, drive.to_dict())
     assert new_drive.protocol == drive.protocol
     assert new_drive.id == drive.id
     assert new_drive.component_name == drive.component_name
 
 
-def test_drive_deepcopy():
-
-    drive = Drive("lit://drive", allow_duplicates=True)
+@pytest.mark.parametrize("drive_id", ["lit://drive", "s3://drive/"])
+def test_drive_deepcopy(drive_id):
+    drive = Drive(drive_id, allow_duplicates=True)
     drive.component_name = "root.work1"
     new_drive = deepcopy(drive)
     assert new_drive.id == drive.id