|
49 | 49 | def noop(): ... |
50 | 50 |
|
51 | 51 |
|
| 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 | + |
52 | 69 | @dataclass |
53 | 70 | class FlyteDirectory(SerializableType, DataClassJsonMixin, os.PathLike, typing.Generic[T]): |
54 | 71 | 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 |
409 | 426 | paths.append(flyte_file) |
410 | 427 | else: |
411 | 428 | 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 | + ) |
413 | 432 |
|
414 | 433 | flyte_directory: FlyteDirectory = FlyteDirectory(path=local_folder, downloader=downloader) |
415 | 434 | flyte_directory._remote_source = remote_path |
@@ -684,7 +703,9 @@ async def async_to_python_value( |
684 | 703 |
|
685 | 704 | batch_size = get_batch_size(expected_python_type) |
686 | 705 |
|
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 | + ) |
688 | 709 |
|
689 | 710 | expected_format = self.get_format(expected_python_type) |
690 | 711 |
|
|
0 commit comments