[Data] Fix WebDataset docs and encoder typing - #65278
Conversation
Signed-off-by: viiccwen <vicwen@apache.org>
There was a problem hiding this comment.
Code Review
This pull request refines the type annotations and documentation for the WebDataset writer, introducing type aliases for encoder configurations and updating the docstrings to correctly reference WebDataset tar archives instead of TFRecord files. It also adds tests to verify default, callable, and chained encoders. The review feedback suggests refining the WebDatasetEncoderConfig type alias to use List[WebDatasetEncoderSpec] instead of List[Any], and using this alias consistently in Dataset.write_webdataset instead of an inlined type union.
|
|
||
| WebDatasetEncoder = Callable[[Dict[str, Any]], Dict[str, Any]] | ||
| WebDatasetEncoderSpec = Union[bool, str, WebDatasetEncoder] | ||
| WebDatasetEncoderConfig = Optional[Union[WebDatasetEncoderSpec, List[Any]]] |
There was a problem hiding this comment.
The type alias WebDatasetEncoderConfig currently uses List[Any] for chained encoders. Since WebDatasetEncoderSpec is already defined to represent valid encoder types, we can make this type annotation much more precise and self-documenting by using List[WebDatasetEncoderSpec] instead of List[Any].
| WebDatasetEncoderConfig = Optional[Union[WebDatasetEncoderSpec, List[Any]]] | |
| WebDatasetEncoderConfig = Optional[Union[WebDatasetEncoderSpec, List[WebDatasetEncoderSpec]]] |
| encoder: Optional[ | ||
| Union[ | ||
| bool, | ||
| str, | ||
| Callable[[Dict[str, Any]], Dict[str, Any]], | ||
| List[Any], | ||
| ] | ||
| ] = True, |
There was a problem hiding this comment.
The PR description states that one of the goals is to "Use the aliases in both Dataset.write_webdataset() and WebDatasetDatasink." However, the implementation in Dataset.write_webdataset still uses an inlined (and less precise) type union.\n\nUsing the newly defined WebDatasetEncoderConfig alias here keeps the public API signature clean and consistent with WebDatasetDatasink.\n\nNote: Please also update the imports at the top of python/ray/data/dataset.py to import WebDatasetEncoderConfig from ray.data._internal.datasource.webdataset_datasink.
encoder: "WebDatasetEncoderConfig" = True,There was a problem hiding this comment.
That's due to the style for Ray-data.
Description
Correct the public
Dataset.write_webdataset()contract so it matches the tar-based WebDataset implementation.Dataset.write_webdataset()andWebDatasetDatasink.This changes documentation and static typing only; existing supported runtime inputs remain valid.
Related issues
Closes #65277.
Additional information
API documentation changes (no real behavior change).