[BUG] Open FlyteFile from remote path (flyteorg#2991)

* fix: Open FlyteFile from remote path Signed-off-by: JiaWei Jiang <[email protected]> * Add integration test Signed-off-by: JiaWei Jiang <[email protected]> * refactor: Use ctx as param instead of recreation Signed-off-by: JiaWei Jiang <[email protected]> * refactor: Clean test logic 1. Remove redundant prints 2. Use `mock.patch.dict` to setup `os.environ` for the current test fn * Avoid contaminating other tests running in the same process Signed-off-by: JiaWei Jiang <[email protected]> * refactor: Setup local path and downloader in constructor Signed-off-by: JiaWei Jiang <[email protected]> * refactor: Move SimpleFileTransfer to an utility file Signed-off-by: JiaWei Jiang <[email protected]> * Remove redundant env var setup Please refer to flyteorg#3001 Signed-off-by: JiaWei Jiang <[email protected]> * test: Add another ff use case Create ff in one task pod and read it in another task pod. Signed-off-by: JiaWei Jiang <[email protected]> --------- Signed-off-by: JiaWei Jiang <[email protected]> Signed-off-by: Shuying Liang <[email protected]>
shuyingliang · Jan 19, 2025 · 3ab75cf · 3ab75cf
1 parent a10955b
commit 3ab75cf
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py
@@ -308,8 +308,7 @@ 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 = partial(
-                ctx.file_access.get_data,
+            self._downloader = lambda: FlyteFilePathTransformer.downloader(
                 ctx=ctx,
                 remote_path=self._remote_source,  # type: ignore
                 local_path=self._local_path,
@@ -742,6 +741,18 @@ 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

diff --git a/tests/flytekit/integration/remote/test_remote.py b/tests/flytekit/integration/remote/test_remote.py
@@ -14,7 +14,7 @@
 from urllib.parse import urlparse
 import uuid
 import pytest
-from unittest import mock
+import mock
 
 from flytekit import LaunchPlan, kwtypes, WorkflowExecutionPhase
 from flytekit.configuration import Config, ImageConfig, SerializationSettings