Make FlyteFile and FlyteDirectory pickleable (flyteorg#3030)

* make _downloader function in FlyteFile/Directory pickleable Signed-off-by: Niels Bantilan <[email protected]> * make FlyteFile and Directory pickleable Signed-off-by: Niels Bantilan <[email protected]> * remove unnecessary helper functions Signed-off-by: Niels Bantilan <[email protected]> * fix lint Signed-off-by: Niels Bantilan <[email protected]> * use partials instead of lambda Signed-off-by: Niels Bantilan <[email protected]> * fix lint Signed-off-by: Niels Bantilan <[email protected]> * remove unneeded helper function Signed-off-by: Niels Bantilan <[email protected]> * update FlyteFilePathTransformer.downloader method Signed-off-by: Niels Bantilan <[email protected]> * remove downloader staticmethod Signed-off-by: Niels Bantilan <[email protected]> * fix lint Signed-off-by: Niels Bantilan <[email protected]> --------- Signed-off-by: Niels Bantilan <[email protected]> Signed-off-by: Shuying Liang <[email protected]>
shuyingliang · Jan 11, 2025 · b252bcd · b252bcd
1 parent 2660eb8
commit b252bcd
Showing 1 changed file with 2 additions and 13 deletions.
diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py
@@ -308,7 +308,8 @@ def __init__(
         if ctx.file_access.is_remote(self.path):
             self._remote_source = self.path
             self._local_path = ctx.file_access.get_random_local_path(self._remote_source)
-            self._downloader = lambda: FlyteFilePathTransformer.downloader(
+            self._downloader = partial(
+                ctx.file_access.get_data,
                 ctx=ctx,
                 remote_path=self._remote_source,  # type: ignore
                 local_path=self._local_path,
@@ -741,18 +742,6 @@ async def async_to_python_value(
         ff._remote_source = uri
         return ff
 
-    @staticmethod
-    def downloader(
-        ctx: FlyteContext, remote_path: typing.Union[str, os.PathLike], local_path: typing.Union[str, os.PathLike]
-    ) -> None:
-        """
-        Download data from remote_path to local_path.
-
-        We design the downloader as a static method because its behavior is logically
-        related to this class but don't need to interact with class or instance data.
-        """
-        ctx.file_access.get_data(remote_path, local_path, is_multipart=False)
-
     def guess_python_type(self, literal_type: LiteralType) -> typing.Type[FlyteFile[typing.Any]]:
         if (
             literal_type.blob is not None