chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,442 @@
|
||||
# Copyright NVIDIA Corporation 2023
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import fnmatch
|
||||
import io
|
||||
import json
|
||||
import re
|
||||
import tarfile
|
||||
from functools import partial
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional, Union
|
||||
|
||||
from ray._common.utils import env_bool
|
||||
from ray.data._internal.util import iterate_with_retry
|
||||
from ray.data.block import Block, BlockAccessor
|
||||
from ray.data.datasource.file_based_datasource import FileBasedDatasource
|
||||
|
||||
ALLOW_UNSAFE_DESERIALIZATION_ENV_VAR = (
|
||||
"RAY_DATA_WEBDATASET_ALLOW_UNSAFE_DESERIALIZATION"
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pyarrow
|
||||
|
||||
|
||||
def _base_plus_ext(path: str):
|
||||
"""Split off all file extensions.
|
||||
|
||||
Returns base, allext.
|
||||
|
||||
Args:
|
||||
path: path with extensions
|
||||
|
||||
Returns:
|
||||
str: path with all extensions removed
|
||||
"""
|
||||
match = re.match(r"^((?:.*/|)[^.]+)[.]([^/]*)$", path)
|
||||
if not match:
|
||||
return None, None
|
||||
return match.group(1), match.group(2)
|
||||
|
||||
|
||||
def _valid_sample(sample: Dict[str, Any]):
|
||||
"""Check whether a sample is valid.
|
||||
|
||||
Args:
|
||||
sample: sample to be checked
|
||||
|
||||
Returns:
|
||||
``True`` if the sample is a non-empty dict without the ``__bad__`` flag.
|
||||
"""
|
||||
return (
|
||||
sample is not None
|
||||
and isinstance(sample, dict)
|
||||
and len(list(sample.keys())) > 0
|
||||
and not sample.get("__bad__", False)
|
||||
)
|
||||
|
||||
|
||||
def _apply_list(
|
||||
f: Union[Callable, List[Callable]], sample: Dict[str, Any], default: Callable = None
|
||||
):
|
||||
"""Apply a list of functions to a sample.
|
||||
|
||||
Args:
|
||||
f: function or list of functions
|
||||
sample: sample to be modified
|
||||
default: default function to be applied to all keys.
|
||||
Defaults to None.
|
||||
|
||||
Returns:
|
||||
modified sample
|
||||
"""
|
||||
if f is None:
|
||||
return sample
|
||||
if not isinstance(f, list):
|
||||
f = [f]
|
||||
for g in f:
|
||||
if default is not None and not callable(g):
|
||||
g = partial(default, format=g)
|
||||
sample = g(sample)
|
||||
return sample
|
||||
|
||||
|
||||
def _check_suffix(suffix: str, suffixes: Union[list, callable]):
|
||||
"""Check whether a suffix is valid.
|
||||
|
||||
Suffixes can be either None (=accept everything), a callable,
|
||||
or a list of patterns. If the pattern contains */? it is treated
|
||||
as a glob pattern, otherwise it is treated as a literal.
|
||||
|
||||
Args:
|
||||
suffix: suffix to be checked
|
||||
suffixes: list of valid suffixes
|
||||
|
||||
Returns:
|
||||
``True`` if the suffix matches the allowed patterns.
|
||||
"""
|
||||
if suffixes is None:
|
||||
return True
|
||||
if callable(suffixes):
|
||||
return suffixes(suffix)
|
||||
for pattern in suffixes:
|
||||
if "*" in pattern or "?" in pattern:
|
||||
if fnmatch.fnmatch("." + suffix, pattern):
|
||||
return True
|
||||
elif suffix == pattern or "." + suffix == pattern:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _tar_file_iterator(
|
||||
fileobj: Any,
|
||||
fileselect: Optional[Union[bool, callable, list]] = None,
|
||||
filerename: Optional[Union[bool, callable, list]] = None,
|
||||
verbose_open: bool = False,
|
||||
meta: dict = None,
|
||||
) -> Iterator[Dict[str, Any]]:
|
||||
"""Iterate over tar file, yielding filename, content pairs for the given tar stream.
|
||||
|
||||
Args:
|
||||
fileobj: file object
|
||||
fileselect: patterns or function selecting
|
||||
files to be selected
|
||||
filerename: patterns or function used to rename selected files
|
||||
before yielding them.
|
||||
verbose_open: if ``True``, print progress messages when starting
|
||||
and finishing iteration over the tar stream.
|
||||
meta: metadata to be added to each sample
|
||||
|
||||
Yields:
|
||||
Dict[str, Any]: Dictionaries with ``fname`` and ``data`` keys for each
|
||||
selected file in the tar stream.
|
||||
"""
|
||||
meta = meta or {}
|
||||
stream = tarfile.open(fileobj=fileobj, mode="r|*")
|
||||
if verbose_open:
|
||||
print(f"start {meta}")
|
||||
for tarinfo in stream:
|
||||
fname = tarinfo.name
|
||||
if not tarinfo.isreg() or fname is None:
|
||||
continue
|
||||
data = stream.extractfile(tarinfo).read()
|
||||
fname = _apply_list(filerename, fname)
|
||||
assert isinstance(fname, str)
|
||||
if not _check_suffix(fname, fileselect):
|
||||
continue
|
||||
result = dict(fname=fname, data=data)
|
||||
yield result
|
||||
if verbose_open:
|
||||
print(f"done {meta}")
|
||||
|
||||
|
||||
def _group_by_keys(
|
||||
data: List[Dict[str, Any]],
|
||||
keys: callable = _base_plus_ext,
|
||||
suffixes: Optional[Union[list, callable]] = None,
|
||||
meta: dict = None,
|
||||
) -> Iterator[Dict[str, Any]]:
|
||||
"""Return function over iterator that groups key, value pairs into samples.
|
||||
|
||||
Args:
|
||||
data: iterator over key, value pairs
|
||||
keys: function that returns key, suffix for a given key
|
||||
suffixes: list of suffixes to be included in the sample
|
||||
meta: metadata to be added to each sample
|
||||
|
||||
Yields:
|
||||
Dict[str, Any]: Grouped samples, where files sharing the same key prefix are
|
||||
combined into a single dictionary.
|
||||
"""
|
||||
meta = meta or {}
|
||||
current_sample = None
|
||||
for filesample in data:
|
||||
assert isinstance(filesample, dict)
|
||||
fname, value = filesample["fname"], filesample["data"]
|
||||
prefix, suffix = keys(fname)
|
||||
if prefix is None:
|
||||
continue
|
||||
if current_sample is None or prefix != current_sample["__key__"]:
|
||||
if _valid_sample(current_sample):
|
||||
current_sample.update(meta)
|
||||
yield current_sample
|
||||
current_sample = dict(__key__=prefix)
|
||||
if "__url__" in filesample:
|
||||
current_sample["__url__"] = filesample["__url__"]
|
||||
if suffix in current_sample:
|
||||
raise ValueError(
|
||||
f"{fname}: duplicate file name in tar file "
|
||||
+ f"{suffix} {current_sample.keys()}, tar is {meta['__url__']}"
|
||||
)
|
||||
if suffixes is None or _check_suffix(suffix, suffixes):
|
||||
current_sample[suffix] = value
|
||||
if _valid_sample(current_sample):
|
||||
current_sample.update(meta)
|
||||
yield current_sample
|
||||
|
||||
|
||||
def _default_decoder(
|
||||
sample: Dict[str, Any],
|
||||
format: Optional[Union[bool, str]] = True,
|
||||
allow_unsafe: bool = False,
|
||||
):
|
||||
"""A default decoder for webdataset.
|
||||
|
||||
This handles common file extensions: .txt, .cls, .cls2,
|
||||
.jpg, .png, .json, .npy, .mp, .pt, .pth, .pickle, .pkl.
|
||||
These are the most common extensions used in webdataset.
|
||||
For other extensions, users can provide their own decoder.
|
||||
|
||||
Args:
|
||||
sample: sample, modified in place
|
||||
format: optional image format hint (e.g. ``"PIL"`` to return PIL
|
||||
images instead of numpy arrays).
|
||||
allow_unsafe: if True, allow pickle/torch deserialization
|
||||
|
||||
Returns:
|
||||
The sample with values decoded according to their key extension.
|
||||
"""
|
||||
sample = dict(sample)
|
||||
for key, value in sample.items():
|
||||
extension = key.split(".")[-1]
|
||||
if key.startswith("__"):
|
||||
continue
|
||||
elif extension in ["txt", "text"]:
|
||||
sample[key] = value.decode("utf-8")
|
||||
elif extension in ["cls", "cls2"]:
|
||||
sample[key] = int(value.decode("utf-8"))
|
||||
elif extension in ["jpg", "png", "ppm", "pgm", "pbm", "pnm"]:
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
|
||||
if format == "PIL":
|
||||
sample[key] = PIL.Image.open(io.BytesIO(value))
|
||||
else:
|
||||
sample[key] = np.asarray(PIL.Image.open(io.BytesIO(value)))
|
||||
elif extension == "json":
|
||||
sample[key] = json.loads(value)
|
||||
elif extension == "npy":
|
||||
import numpy as np
|
||||
|
||||
sample[key] = np.load(io.BytesIO(value))
|
||||
elif extension == "mp":
|
||||
import msgpack
|
||||
|
||||
sample[key] = msgpack.unpackb(value, raw=False)
|
||||
elif extension in ["pt", "pth"]:
|
||||
if not allow_unsafe:
|
||||
raise ValueError(
|
||||
f"Refusing to load .{extension} member {key!r} from "
|
||||
f"WebDataset with weights_only=False (arbitrary code "
|
||||
f"execution risk). Provide a custom decoder or set "
|
||||
f"{ALLOW_UNSAFE_DESERIALIZATION_ENV_VAR}=1 "
|
||||
f"for trusted sources."
|
||||
)
|
||||
import torch
|
||||
|
||||
sample[key] = torch.load(io.BytesIO(value), weights_only=False)
|
||||
elif extension in ["pickle", "pkl"]:
|
||||
if not allow_unsafe:
|
||||
raise ValueError(
|
||||
f"Refusing to unpickle WebDataset member {key!r} "
|
||||
f"(arbitrary code execution risk). Provide a custom "
|
||||
f"decoder or set "
|
||||
f"{ALLOW_UNSAFE_DESERIALIZATION_ENV_VAR}=1 "
|
||||
f"for trusted sources."
|
||||
)
|
||||
import pickle
|
||||
|
||||
sample[key] = pickle.loads(value)
|
||||
return sample
|
||||
|
||||
|
||||
extension_to_format = {"jpg": "jpeg"}
|
||||
|
||||
|
||||
def _default_encoder(sample: Dict[str, Any], format: Optional[Union[str, bool]] = True):
|
||||
"""A default encoder for webdataset.
|
||||
|
||||
This handles common file extensions: .txt, .cls, .cls2, .jpg,
|
||||
.png, .json, .npy, .mp, .pt, .pth, .pickle, .pkl
|
||||
These are the most common extensions used in webdataset.
|
||||
For other extensions, users can provide their own encoder.
|
||||
|
||||
Args:
|
||||
sample: sample to encode.
|
||||
format: optional image format hint forwarded to the underlying
|
||||
image encoder.
|
||||
|
||||
Returns:
|
||||
The sample with values encoded according to their key extension.
|
||||
"""
|
||||
sample = dict(sample)
|
||||
for key, value in sample.items():
|
||||
extension = key.split(".")[-1]
|
||||
if key.startswith("__"):
|
||||
continue
|
||||
elif extension in ["txt"]:
|
||||
sample[key] = value.encode("utf-8")
|
||||
elif extension in ["cls", "cls2"]:
|
||||
sample[key] = str(value).encode("utf-8")
|
||||
elif extension in ["jpg", "jpeg", "png", "ppm", "pgm", "pbm", "pnm"]:
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
|
||||
if isinstance(value, np.ndarray):
|
||||
value = PIL.Image.fromarray(value)
|
||||
assert isinstance(value, PIL.Image.Image)
|
||||
stream = io.BytesIO()
|
||||
value.save(
|
||||
stream, format=extension_to_format.get(extension.lower(), extension)
|
||||
)
|
||||
sample[key] = stream.getvalue()
|
||||
elif extension == "json":
|
||||
sample[key] = json.dumps(value).encode("utf-8")
|
||||
elif extension == "npy":
|
||||
import numpy as np
|
||||
|
||||
stream = io.BytesIO()
|
||||
np.save(stream, value)
|
||||
sample[key] = stream.getvalue()
|
||||
elif extension == "mp":
|
||||
import msgpack
|
||||
|
||||
sample[key] = msgpack.dumps(value)
|
||||
elif extension in ["pt", "pth"]:
|
||||
import torch
|
||||
|
||||
stream = io.BytesIO()
|
||||
torch.save(value, stream)
|
||||
sample[key] = stream.getvalue()
|
||||
elif extension in ["pickle", "pkl"]:
|
||||
import pickle
|
||||
|
||||
stream = io.BytesIO()
|
||||
pickle.dump(value, stream)
|
||||
sample[key] = stream.getvalue()
|
||||
return sample
|
||||
|
||||
|
||||
def _make_iterable(block: BlockAccessor):
|
||||
"""Make a block iterable.
|
||||
|
||||
This is a placeholder for dealing with more complex blocks.
|
||||
|
||||
Args:
|
||||
block: Ray Dataset block
|
||||
|
||||
Returns:
|
||||
Iterable[Dict[str,Any]]: Iterable of samples
|
||||
"""
|
||||
return block.iter_rows(public_row_format=False)
|
||||
|
||||
|
||||
class WebDatasetDatasource(FileBasedDatasource):
|
||||
"""A Datasource for WebDataset datasets (tar format with naming conventions)."""
|
||||
|
||||
_FILE_EXTENSIONS = ["tar"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
paths: Union[str, List[str]],
|
||||
decoder: Optional[Union[bool, str, callable, list]] = True,
|
||||
fileselect: Optional[Union[bool, callable, list]] = None,
|
||||
filerename: Optional[Union[bool, callable, list]] = None,
|
||||
suffixes: Optional[Union[bool, callable, list]] = None,
|
||||
verbose_open: bool = False,
|
||||
expand_json: bool = False,
|
||||
**file_based_datasource_kwargs,
|
||||
):
|
||||
super().__init__(paths, **file_based_datasource_kwargs)
|
||||
|
||||
self.decoder = decoder
|
||||
self.fileselect = fileselect
|
||||
self.filerename = filerename
|
||||
self.suffixes = suffixes
|
||||
self.verbose_open = verbose_open
|
||||
self.expand_json = expand_json
|
||||
|
||||
self._allow_unsafe_deserialization = env_bool(
|
||||
ALLOW_UNSAFE_DESERIALIZATION_ENV_VAR, False
|
||||
)
|
||||
|
||||
def _read_stream(self, stream: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
|
||||
"""Read and decode samples from a stream.
|
||||
|
||||
Note that fileselect selects files during reading, while suffixes
|
||||
selects files during the grouping step.
|
||||
|
||||
Args:
|
||||
stream: File descriptor to read from.
|
||||
path: Path to the data.
|
||||
|
||||
Yields:
|
||||
Block: Single-row blocks (one per WebDataset sample).
|
||||
"""
|
||||
|
||||
import pandas as pd
|
||||
|
||||
def get_tar_file_iterator():
|
||||
return _tar_file_iterator(
|
||||
stream,
|
||||
fileselect=self.fileselect,
|
||||
filerename=self.filerename,
|
||||
verbose_open=self.verbose_open,
|
||||
)
|
||||
|
||||
# S3 can raise transient errors during iteration
|
||||
files = iterate_with_retry(
|
||||
get_tar_file_iterator,
|
||||
"iterate tar file",
|
||||
match=self._data_context.retried_io_errors,
|
||||
)
|
||||
|
||||
samples = _group_by_keys(files, meta=dict(__url__=path), suffixes=self.suffixes)
|
||||
default_decoder = partial(
|
||||
_default_decoder, allow_unsafe=self._allow_unsafe_deserialization
|
||||
)
|
||||
for sample in samples:
|
||||
if self.decoder is not None:
|
||||
sample = _apply_list(self.decoder, sample, default=default_decoder)
|
||||
if self.expand_json:
|
||||
if isinstance(sample["json"], bytes):
|
||||
parsed_json = json.loads(sample["json"].decode("utf-8"))
|
||||
elif isinstance(sample["json"], str):
|
||||
parsed_json = json.loads(sample["json"])
|
||||
elif isinstance(sample["json"], dict):
|
||||
parsed_json = sample["json"]
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Unsupported data type" f" {type(sample['json'])} for sample"
|
||||
)
|
||||
for k, v in parsed_json.items():
|
||||
if k not in sample:
|
||||
sample[k] = []
|
||||
sample[k].append(v)
|
||||
yield pd.DataFrame(
|
||||
{
|
||||
k: v if isinstance(v, list) and len(v) == 1 else [v]
|
||||
for k, v in sample.items()
|
||||
}
|
||||
)
|
||||
Reference in New Issue
Block a user