Skip to content

Commit a2f1410

Browse files
committed
fix: flyte dir downloader use current context
-e Signed-off-by: machichima <nary12321@gmail.com>
1 parent 1f1159b commit a2f1410

1 file changed

Lines changed: 23 additions & 2 deletions

File tree

flytekit/types/directory/types.py

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,23 @@
4949
def noop(): ...
5050

5151

52+
class _FlyteDirectoryDownloader:
53+
"""Downloader for FlyteDirectory that uses current context when called."""
54+
55+
def __init__(self, remote_path: str, local_path: str, is_multipart: bool = True, batch_size: typing.Optional[int] = None):
56+
self.remote_path = remote_path
57+
self.local_path = local_path
58+
self.is_multipart = is_multipart
59+
self.batch_size = batch_size
60+
61+
def __call__(self):
62+
"""Download the directory using current context's synced get_data method."""
63+
current_ctx = FlyteContextManager.current_context()
64+
return current_ctx.file_access.get_data(
65+
remote_path=self.remote_path, local_path=self.local_path, is_multipart=self.is_multipart, batch_size=self.batch_size
66+
)
67+
68+
5269
@dataclass
5370
class FlyteDirectory(SerializableType, DataClassJsonMixin, os.PathLike, typing.Generic[T]):
5471
path: PathType = field(default=None, metadata=config(mm_field=fields.String())) # type: ignore
@@ -409,7 +426,9 @@ def listdir(cls, directory: FlyteDirectory) -> typing.List[typing.Union[FlyteDir
409426
paths.append(flyte_file)
410427
else:
411428
local_folder = file_access.get_random_local_directory()
412-
downloader = partial(file_access.get_data, remote_path, local_folder, is_multipart=True)
429+
downloader = _FlyteDirectoryDownloader(
430+
remote_path=remote_path, local_path=local_folder, is_multipart=True
431+
)
413432

414433
flyte_directory: FlyteDirectory = FlyteDirectory(path=local_folder, downloader=downloader)
415434
flyte_directory._remote_source = remote_path
@@ -684,7 +703,9 @@ async def async_to_python_value(
684703

685704
batch_size = get_batch_size(expected_python_type)
686705

687-
_downloader = partial(ctx.file_access.get_data, uri, local_folder, is_multipart=True, batch_size=batch_size)
706+
_downloader = _FlyteDirectoryDownloader(
707+
remote_path=uri, local_path=local_folder, is_multipart=True, batch_size=batch_size
708+
)
688709

689710
expected_format = self.get_format(expected_python_type)
690711

0 commit comments

Comments
 (0)