From 2c397d8d5f32bef95d2768be19957424a82fc262 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Wed, 8 Jan 2025 16:34:43 -0500 Subject: [PATCH] Make FlyteFile and FlyteDirectory pickleable (#3030) * make _downloader function in FlyteFile/Directory pickleable Signed-off-by: Niels Bantilan * make FlyteFile and Directory pickleable Signed-off-by: Niels Bantilan * remove unnecessary helper functions Signed-off-by: Niels Bantilan * fix lint Signed-off-by: Niels Bantilan * use partials instead of lambda Signed-off-by: Niels Bantilan * fix lint Signed-off-by: Niels Bantilan * remove unneeded helper function Signed-off-by: Niels Bantilan * update FlyteFilePathTransformer.downloader method Signed-off-by: Niels Bantilan * remove downloader staticmethod Signed-off-by: Niels Bantilan * fix lint Signed-off-by: Niels Bantilan --------- Signed-off-by: Niels Bantilan Signed-off-by: Shuying Liang --- flytekit/types/file/file.py | 15 ++------------- 1 file changed, 2 insertions(+), 13 deletions(-) diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index ef08ef64de..b66c48443c 100644 --- 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