From f28f5f1708540ba1b5096313d1e2faf24766af92 Mon Sep 17 00:00:00 2001 From: machichima Date: Thu, 14 Aug 2025 17:39:22 +0800 Subject: [PATCH 1/2] fix: let downloader use current context -e Signed-off-by: machichima --- flytekit/types/file/file.py | 30 +++++++++++++++++++++++++----- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index a5eff0b68f..f280732d12 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -40,6 +40,24 @@ def noop(): ... +class _FlyteFileDownloader: + """Downloader for FlyteFile that uses current context when called.""" + + def __init__(self, remote_path: str, local_path: str, is_multipart: bool = False): + self.remote_path = remote_path + self.local_path = local_path + self.is_multipart = is_multipart + + def __call__(self): + """Download the file using current context's synced get_data method.""" + current_ctx = FlyteContextManager.current_context() + return current_ctx.file_access.get_data( + remote_path=self.remote_path, + local_path=self.local_path, + is_multipart=self.is_multipart + ) + + T = typing.TypeVar("T") @@ -318,10 +336,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 = partial( - ctx.file_access.get_data, - ctx=ctx, - remote_path=self._remote_source, # type: ignore + self._downloader = _FlyteFileDownloader( + remote_path=str(self._remote_source), # type: ignore local_path=self._local_path, ) @@ -755,7 +771,11 @@ async def async_to_python_value( # For the remote case, return an FlyteFile object that can download local_path = ctx.file_access.get_random_local_path(uri) - _downloader = partial(ctx.file_access.get_data, remote_path=uri, local_path=local_path, is_multipart=False) + _downloader = _FlyteFileDownloader( + remote_path=uri, + local_path=local_path, + is_multipart=False + ) expected_format = FlyteFilePathTransformer.get_format(expected_python_type) ff = FlyteFile.__class_getitem__(expected_format)(path=local_path, downloader=_downloader, metadata=metadata) From 818cf34b580dc85e6ba92525fe35557cddee6af8 Mon Sep 17 00:00:00 2001 From: machichima Date: Fri, 15 Aug 2025 16:01:07 +0800 Subject: [PATCH 2/2] refactor: fix lint -e Signed-off-by: machichima --- flytekit/types/file/file.py | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index f280732d12..780188f9e5 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -7,7 +7,6 @@ import typing from contextlib import contextmanager from dataclasses import dataclass, field -from functools import partial from typing import Dict, cast from urllib.parse import unquote @@ -42,19 +41,17 @@ def noop(): ... class _FlyteFileDownloader: """Downloader for FlyteFile that uses current context when called.""" - + def __init__(self, remote_path: str, local_path: str, is_multipart: bool = False): self.remote_path = remote_path - self.local_path = local_path + self.local_path = local_path self.is_multipart = is_multipart - + def __call__(self): """Download the file using current context's synced get_data method.""" current_ctx = FlyteContextManager.current_context() return current_ctx.file_access.get_data( - remote_path=self.remote_path, - local_path=self.local_path, - is_multipart=self.is_multipart + remote_path=self.remote_path, local_path=self.local_path, is_multipart=self.is_multipart ) @@ -338,7 +335,7 @@ def __init__( self._local_path = ctx.file_access.get_random_local_path(self._remote_source) self._downloader = _FlyteFileDownloader( remote_path=str(self._remote_source), # type: ignore - local_path=self._local_path, + local_path=str(self._local_path), ) def __fspath__(self): @@ -771,11 +768,7 @@ async def async_to_python_value( # For the remote case, return an FlyteFile object that can download local_path = ctx.file_access.get_random_local_path(uri) - _downloader = _FlyteFileDownloader( - remote_path=uri, - local_path=local_path, - is_multipart=False - ) + _downloader = _FlyteFileDownloader(remote_path=uri, local_path=local_path, is_multipart=False) expected_format = FlyteFilePathTransformer.get_format(expected_python_type) ff = FlyteFile.__class_getitem__(expected_format)(path=local_path, downloader=_downloader, metadata=metadata)