diff --git a/flytekit/types/directory/types.py b/flytekit/types/directory/types.py index 410ed5b071..286cfafbf6 100644 --- a/flytekit/types/directory/types.py +++ b/flytekit/types/directory/types.py @@ -49,6 +49,28 @@ def noop(): ... +class _FlyteDirectoryDownloader: + """Downloader for FlyteDirectory that uses current context when called.""" + + def __init__( + self, remote_path: str, local_path: str, is_multipart: bool = True, batch_size: typing.Optional[int] = None + ): + self.remote_path = remote_path + self.local_path = local_path + self.is_multipart = is_multipart + self.batch_size = batch_size + + def __call__(self): + """Download the directory 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, + batch_size=self.batch_size, + ) + + @dataclass class FlyteDirectory(SerializableType, DataClassJsonMixin, os.PathLike, typing.Generic[T]): path: PathType = field(default=None, metadata=config(mm_field=fields.String())) # type: ignore @@ -409,9 +431,11 @@ def listdir(cls, directory: FlyteDirectory) -> typing.List[typing.Union[FlyteDir paths.append(flyte_file) else: local_folder = file_access.get_random_local_directory() - downloader = partial(file_access.get_data, remote_path, local_folder, is_multipart=True) + dir_downloader: typing.Callable = _FlyteDirectoryDownloader( + remote_path=remote_path, local_path=local_folder, is_multipart=True + ) - flyte_directory: FlyteDirectory = FlyteDirectory(path=local_folder, downloader=downloader) + flyte_directory: FlyteDirectory = FlyteDirectory(path=local_folder, downloader=dir_downloader) flyte_directory._remote_source = remote_path paths.append(flyte_directory) @@ -684,7 +708,9 @@ async def async_to_python_value( batch_size = get_batch_size(expected_python_type) - _downloader = partial(ctx.file_access.get_data, uri, local_folder, is_multipart=True, batch_size=batch_size) + _downloader = _FlyteDirectoryDownloader( + remote_path=uri, local_path=local_folder, is_multipart=True, batch_size=batch_size + ) expected_format = self.get_format(expected_python_type)