Skip to content

Commit cf7b3b1

Browse files
committed
Add FileDownloadConfig annotation for FlyteFile inputs
Port new BlobType fields file_extension and enable_legacy_filename to flytekit. FlyteFile inputs can be annotated with the FileDownloadConfig annotation to configure the file extension to use during the copilot download phase. e.g. ```python def t1(file: Annotated[FlyteFile, FileDownloadConfig(file_extension="csv")]): ... # copilot downloads the file to e.g. /inputs/file.csv versus... def t1(file: FlyteFile["csv"]): ... # copilot downloads the file to e.g. /inputs/file ```
1 parent c0921e9 commit cf7b3b1

4 files changed

Lines changed: 111 additions & 8 deletions

File tree

flytekit/core/type_engine.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,48 @@ def get_batch_size(t: Type) -> Optional[int]:
105105
return None
106106

107107

108+
class FileDownloadConfig:
109+
"""
110+
This is used to annotate a FlyteFile when we want to download the file with a specific extension. For example,
111+
112+
```python
113+
# ContainerTask
114+
def t1(file: Annotated[FlyteFile, FileDownloadConfig(file_extension="csv")]):
115+
... # copilot downloads the file to e.g. /inputs/file.csv
116+
117+
versus...
118+
119+
def t1(file: FlyteFile["csv"]):
120+
... # copilot downloads the file to e.g. /inputs/file
121+
```
122+
123+
file_extension: (Default is "") The file extension (e.g. "csv", "parquet") to use during copilot download.
124+
enable_legacy_filename: (Default is False) When true and file_extension is non-empty, the copilot download phase
125+
writes the blob to both the full path (with extension) and the old path (without extension), preserving backward compatibility for
126+
workflows with tasks that may read from both.
127+
"""
128+
129+
def __init__(self, file_extension: str = "", enable_legacy_filename: bool = False):
130+
self._file_extension = file_extension
131+
self._enable_legacy_filename = enable_legacy_filename
132+
133+
@property
134+
def file_extension(self) -> str:
135+
return self._file_extension
136+
137+
@property
138+
def enable_legacy_filename(self) -> bool:
139+
return self._enable_legacy_filename
140+
141+
142+
def get_file_download_config(t: Type) -> Optional[FileDownloadConfig]:
143+
if is_annotated(t):
144+
for arg in get_args(t):
145+
if isinstance(arg, FileDownloadConfig):
146+
return arg
147+
return None
148+
149+
108150
def modify_literal_uris(lit: Literal):
109151
"""
110152
Modifies the literal object recursively to replace the URIs with the native paths in case they are of

flytekit/models/core/types.py

Lines changed: 37 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -38,13 +38,19 @@ class BlobDimensionality(object):
3838
SINGLE = _types_pb2.BlobType.SINGLE
3939
MULTIPART = _types_pb2.BlobType.MULTIPART
4040

41-
def __init__(self, format, dimensionality):
41+
def __init__(self, format, dimensionality, file_extension="", enable_legacy_filename=False):
4242
"""
4343
:param Text format: A string describing the format of the underlying blob data.
4444
:param int dimensionality: An integer from BlobType.BlobDimensionality enum
45+
:param Text file_extension: The file extension (e.g. "csv", "parquet") to use
46+
during copilot download, e.g. "csv", "parquet". Empty by default.
47+
:param bool enable_legacy_filename: When True and file_extension is set, the copilot
48+
download phase writes the blob to both the extended path and the base path.
4549
"""
4650
self._format = format
4751
self._dimensionality = dimensionality
52+
self._file_extension = file_extension
53+
self._enable_legacy_filename = enable_legacy_filename
4854

4955
@property
5056
def format(self):
@@ -62,16 +68,44 @@ def dimensionality(self):
6268
"""
6369
return self._dimensionality
6470

71+
@property
72+
def file_extension(self):
73+
"""
74+
The file extension (e.g. "csv", "parquet") to use during copilot download.
75+
Default is "", which means no extension is appended.
76+
:rtype: Text
77+
"""
78+
return self._file_extension
79+
80+
@property
81+
def enable_legacy_filename(self):
82+
"""
83+
When True and file_extension is set, the copilot download writes the blob to
84+
both the full path (with extension) and the old path (without extension).
85+
:rtype: bool
86+
"""
87+
return self._enable_legacy_filename
88+
6589
def to_flyte_idl(self):
6690
"""
6791
:rtype: flyteidl.core.types_pb2.BlobType
6892
"""
69-
return _types_pb2.BlobType(format=self.format, dimensionality=self.dimensionality)
93+
return _types_pb2.BlobType(
94+
format=self.format,
95+
dimensionality=self.dimensionality,
96+
file_extension=self._file_extension,
97+
enable_legacy_filename=self._enable_legacy_filename,
98+
)
7099

71100
@classmethod
72101
def from_flyte_idl(cls, proto):
73102
"""
74103
:param flyteidl.core.types_pb2.BlobType proto:
75104
:rtype: BlobType
76105
"""
77-
return cls(format=proto.format, dimensionality=proto.dimensionality)
106+
return cls(
107+
format=proto.format,
108+
dimensionality=proto.dimensionality,
109+
file_extension=proto.file_extension,
110+
enable_legacy_filename=proto.enable_legacy_filename,
111+
)

flytekit/types/file/file.py

Lines changed: 31 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
AsyncTypeTransformer,
2525
TypeEngine,
2626
TypeTransformerFailedError,
27+
get_file_download_config,
2728
get_underlying_type,
2829
)
2930
from flytekit.exceptions.user import FlyteAssertion
@@ -477,8 +478,26 @@ def get_format(t: typing.Union[typing.Type[FlyteFile], os.PathLike]) -> str:
477478
return ""
478479
return cast(FlyteFile, t).extension()
479480

480-
def _blob_type(self, format: str) -> BlobType:
481-
return BlobType(format=format, dimensionality=BlobType.BlobDimensionality.SINGLE)
481+
@staticmethod
482+
def get_file_extension(t: typing.Union[typing.Type[FlyteFile], os.PathLike]) -> str:
483+
if t is os.PathLike:
484+
return ""
485+
file_download_config = get_file_download_config(t)
486+
if file_download_config is None:
487+
return ""
488+
return file_download_config.file_extension or ""
489+
490+
@staticmethod
491+
def get_enable_legacy_filename(t: typing.Union[typing.Type[FlyteFile], os.PathLike]) -> str:
492+
if t is os.PathLike:
493+
return False
494+
file_download_config = get_file_download_config(t)
495+
if file_download_config is None:
496+
return False
497+
return file_download_config.enable_legacy_filename or False
498+
499+
def _blob_type(self, format: str, file_extension: str = "", enable_legacy_filename: bool = False) -> BlobType:
500+
return BlobType(format=format, dimensionality=BlobType.BlobDimensionality.SINGLE, file_extension=file_extension, enable_legacy_filename=enable_legacy_filename)
482501

483502
def assert_type(
484503
self, t: typing.Union[typing.Type[FlyteFile], os.PathLike], v: typing.Union[FlyteFile, os.PathLike, str]
@@ -491,7 +510,11 @@ def assert_type(
491510
)
492511

493512
def get_literal_type(self, t: typing.Union[typing.Type[FlyteFile], os.PathLike]) -> LiteralType:
494-
return LiteralType(blob=self._blob_type(format=FlyteFilePathTransformer.get_format(t)))
513+
return LiteralType(blob=self._blob_type(
514+
format=FlyteFilePathTransformer.get_format(t),
515+
file_extension=FlyteFilePathTransformer.get_file_extension(t),
516+
enable_legacy_filename=FlyteFilePathTransformer.get_enable_legacy_filename(t),
517+
))
495518

496519
def get_mime_type_from_extension(self, extension: str) -> typing.Union[str, typing.Sequence[str]]:
497520
extension_to_mime_type = {
@@ -565,7 +588,11 @@ async def async_to_literal(
565588
raise ValueError(f"Incorrect type {python_type}, must be either a FlyteFile or os.PathLike")
566589

567590
# information used by all cases
568-
meta = BlobMetadata(type=self._blob_type(format=FlyteFilePathTransformer.get_format(python_type)))
591+
meta = BlobMetadata(type=self._blob_type(
592+
format=FlyteFilePathTransformer.get_format(python_type),
593+
file_extension=FlyteFilePathTransformer.get_file_extension(python_type),
594+
enable_legacy_filename=FlyteFilePathTransformer.get_enable_legacy_filename(python_type),
595+
))
569596

570597
if isinstance(python_val, FlyteFile):
571598
# Cast the source path to str type to avoid error raised when the source path is used as the blob uri,

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ dependencies = [
2121
"diskcache>=5.2.1",
2222
"docker>=4.0.0",
2323
"docstring-parser>=0.9.0",
24-
"flyteidl>=1.16.1,<2.0.0a0",
24+
"flyteidl @ git+https://github.com/dominodatalab/flyteidl.git@af517f6",
2525
"fsspec>=2023.3.0",
2626
# Bug in 2025.5.0, 2025.5.0post1 https://github.com/fsspec/gcsfs/issues/687
2727
# Bug in 2024.2.0 https://github.com/fsspec/gcsfs/pull/643

0 commit comments

Comments
 (0)