Skip to content

Commit

Permalink
Make FlyteFile and FlyteDirectory pickleable (flyteorg#3030)
Browse files Browse the repository at this point in the history
* make _downloader function in FlyteFile/Directory pickleable

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* make FlyteFile and Directory pickleable

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* remove unnecessary helper functions

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* fix lint

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* use partials instead of lambda

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* fix lint

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* remove unneeded helper function

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* update FlyteFilePathTransformer.downloader method

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* remove downloader staticmethod

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* fix lint

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

---------

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>
Signed-off-by: Shuying Liang <shuying.liang@gmail.com>
  • Loading branch information
cosmicBboy authored and shuyingliang committed Jan 22, 2025
1 parent eb1c838 commit d030af7
Showing 1 changed file with 2 additions and 13 deletions.
15 changes: 2 additions & 13 deletions flytekit/types/file/file.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit d030af7

Please sign in to comment.