chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,335 @@
|
||||
# Copyright NVIDIA Corporation 2023
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import glob
|
||||
import io
|
||||
import os
|
||||
import pickle
|
||||
import tarfile
|
||||
|
||||
import pytest
|
||||
import webdataset as wds
|
||||
|
||||
import ray
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
class TarWriter:
|
||||
def __init__(self, path):
|
||||
self.path = path
|
||||
self.tar = tarfile.open(path, "w")
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
self.tar.close()
|
||||
|
||||
def write(self, name, data):
|
||||
f = self.tar.tarinfo()
|
||||
f.name = name
|
||||
f.size = len(data)
|
||||
self.tar.addfile(f, io.BytesIO(data))
|
||||
|
||||
|
||||
def test_webdataset_read(ray_start_2_cpus, tmp_path):
|
||||
path = os.path.join(tmp_path, "bar_000000.tar")
|
||||
with TarWriter(path) as tf:
|
||||
for i in range(100):
|
||||
tf.write(f"{i}.a", str(i).encode("utf-8"))
|
||||
tf.write(f"{i}.b", str(i**2).encode("utf-8"))
|
||||
assert os.path.exists(path)
|
||||
assert len(glob.glob(f"{tmp_path}/*.tar")) == 1
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)])
|
||||
samples = ds.take(100)
|
||||
assert len(samples) == 100
|
||||
for i, sample in enumerate(samples):
|
||||
assert isinstance(sample, dict), sample
|
||||
assert sample["__key__"] == str(i)
|
||||
assert sample["a"].decode("utf-8") == str(i)
|
||||
assert sample["b"].decode("utf-8") == str(i**2)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def allow_unsafe_deserialization(monkeypatch):
|
||||
monkeypatch.setenv("RAY_DATA_WEBDATASET_ALLOW_UNSAFE_DESERIALIZATION", "1")
|
||||
|
||||
|
||||
def test_webdataset_expand_json(
|
||||
ray_start_2_cpus, tmp_path, allow_unsafe_deserialization
|
||||
):
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
|
||||
gray = np.random.randint(0, 255, (100, 100), dtype=np.uint8)
|
||||
dstruct = dict(a=[1, 2], b=dict(c=2), d="hello")
|
||||
ttensor = torch.tensor([1, 2, 3]).numpy()
|
||||
|
||||
sample = {
|
||||
"__key__": "foo",
|
||||
"jpg": image,
|
||||
"gray.png": gray,
|
||||
"mp": dstruct,
|
||||
"json": dstruct,
|
||||
"pt": ttensor,
|
||||
"und": b"undecoded",
|
||||
"custom": b"nothing",
|
||||
}
|
||||
|
||||
# write the encoded data using the default encoder
|
||||
data = [sample]
|
||||
ds = ray.data.from_items(data).repartition(1)
|
||||
ds.write_webdataset(path=tmp_path, try_create_dir=True)
|
||||
ds = ray.data.read_webdataset(
|
||||
paths=[str(tmp_path)], override_num_blocks=1, expand_json=True
|
||||
)
|
||||
record = ds.take(1)
|
||||
assert [1, 2] == record[0]["a"]
|
||||
|
||||
|
||||
def test_webdataset_suffixes(ray_start_2_cpus, tmp_path):
|
||||
path = os.path.join(tmp_path, "bar_000000.tar")
|
||||
with TarWriter(path) as tf:
|
||||
for i in range(100):
|
||||
tf.write(f"{i}.txt", str(i).encode("utf-8"))
|
||||
tf.write(f"{i}.test.txt", str(i**2).encode("utf-8"))
|
||||
tf.write(f"{i}.cls", str(i**2).encode("utf-8"))
|
||||
tf.write(f"{i}.test.cls2", str(i**2).encode("utf-8"))
|
||||
assert os.path.exists(path)
|
||||
assert len(glob.glob(f"{tmp_path}/*.tar")) == 1
|
||||
|
||||
# test simple suffixes
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)], suffixes=["txt", "cls"])
|
||||
samples = ds.take(100)
|
||||
assert len(samples) == 100
|
||||
for i, sample in enumerate(samples):
|
||||
assert set(sample.keys()) == {"__url__", "__key__", "txt", "cls"}
|
||||
|
||||
# test fnmatch patterns for suffixes
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)], suffixes=["*.txt", "*.cls"])
|
||||
samples = ds.take(100)
|
||||
assert len(samples) == 100
|
||||
for i, sample in enumerate(samples):
|
||||
assert set(sample.keys()) == {"__url__", "__key__", "txt", "cls", "test.txt"}
|
||||
|
||||
# test selection function
|
||||
def select(name):
|
||||
return name.endswith("txt")
|
||||
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)], suffixes=select)
|
||||
samples = ds.take(100)
|
||||
assert len(samples) == 100
|
||||
for i, sample in enumerate(samples):
|
||||
assert set(sample.keys()) == {"__url__", "__key__", "txt", "test.txt"}
|
||||
|
||||
# test filerename
|
||||
def renamer(name):
|
||||
result = name.replace("txt", "text")
|
||||
print("***", name, result)
|
||||
return result
|
||||
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)], filerename=renamer)
|
||||
samples = ds.take(100)
|
||||
assert len(samples) == 100
|
||||
for i, sample in enumerate(samples):
|
||||
assert set(sample.keys()) == {
|
||||
"__url__",
|
||||
"__key__",
|
||||
"text",
|
||||
"cls",
|
||||
"test.text",
|
||||
"test.cls2",
|
||||
}
|
||||
|
||||
|
||||
def test_webdataset_write(ray_start_2_cpus, tmp_path):
|
||||
print(ray.available_resources())
|
||||
data = [dict(__key__=str(i), a=str(i), b=str(i**2)) for i in range(100)]
|
||||
ds = ray.data.from_items(data).repartition(1)
|
||||
ds.write_webdataset(path=tmp_path, try_create_dir=True)
|
||||
paths = glob.glob(f"{tmp_path}/*.tar")
|
||||
assert len(paths) == 1
|
||||
with open(paths[0], "rb") as stream:
|
||||
tf = tarfile.open(fileobj=stream)
|
||||
for i in range(100):
|
||||
assert tf.extractfile(f"{i}.a").read().decode("utf-8") == str(i)
|
||||
assert tf.extractfile(f"{i}.b").read().decode("utf-8") == str(i**2)
|
||||
|
||||
|
||||
def custom_decoder(sample):
|
||||
for key, value in sample.items():
|
||||
if key == "png":
|
||||
# check that images have already been decoded
|
||||
assert not isinstance(value, bytes)
|
||||
elif key.endswith("custom"):
|
||||
sample[key] = "custom-value"
|
||||
return sample
|
||||
|
||||
|
||||
def test_webdataset_coding(ray_start_2_cpus, tmp_path, allow_unsafe_deserialization):
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
|
||||
image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
|
||||
gray = np.random.randint(0, 255, (100, 100), dtype=np.uint8)
|
||||
dstruct = dict(a=[1], b=dict(c=2), d="hello")
|
||||
ttensor = torch.tensor([1, 2, 3]).numpy()
|
||||
|
||||
sample = {
|
||||
"__key__": "foo",
|
||||
"jpg": image,
|
||||
"gray.png": gray,
|
||||
"mp": dstruct,
|
||||
"json": dstruct,
|
||||
"pt": ttensor,
|
||||
"und": b"undecoded",
|
||||
"custom": b"nothing",
|
||||
}
|
||||
|
||||
# write the encoded data using the default encoder
|
||||
data = [sample]
|
||||
ds = ray.data.from_items(data).repartition(1)
|
||||
ds.write_webdataset(path=tmp_path, try_create_dir=True)
|
||||
|
||||
# read the encoded data using the default decoder
|
||||
paths = glob.glob(f"{tmp_path}/*.tar")
|
||||
assert len(paths) == 1
|
||||
path = paths[0]
|
||||
assert os.path.exists(path)
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)])
|
||||
samples = ds.take(1)
|
||||
assert len(samples) == 1
|
||||
for sample in samples:
|
||||
assert isinstance(sample, dict), sample
|
||||
assert sample["__key__"] == "foo"
|
||||
assert isinstance(sample["jpg"], np.ndarray)
|
||||
assert sample["jpg"].shape == (100, 100, 3)
|
||||
assert isinstance(sample["gray.png"], np.ndarray)
|
||||
assert sample["gray.png"].shape == (100, 100)
|
||||
assert isinstance(sample["mp"], dict)
|
||||
assert sample["mp"]["a"] == [1]
|
||||
assert sample["mp"]["b"]["c"] == 2
|
||||
assert isinstance(sample["json"], dict)
|
||||
assert sample["json"]["a"] == [1]
|
||||
assert isinstance(sample["pt"], np.ndarray)
|
||||
assert sample["pt"].tolist() == [1, 2, 3]
|
||||
|
||||
# test the format argument to the default decoder and multiple decoders
|
||||
ds = ray.data.read_webdataset(
|
||||
paths=[str(tmp_path)], decoder=["PIL", custom_decoder]
|
||||
)
|
||||
samples = ds.take(1)
|
||||
assert len(samples) == 1
|
||||
for sample in samples:
|
||||
assert isinstance(sample, dict), sample
|
||||
assert sample["__key__"] == "foo"
|
||||
assert isinstance(sample["jpg"], PIL.Image.Image)
|
||||
assert isinstance(sample["gray.png"], PIL.Image.Image)
|
||||
assert isinstance(sample["und"], bytes)
|
||||
assert sample["und"] == b"undecoded"
|
||||
assert sample["custom"] == "custom-value"
|
||||
|
||||
|
||||
def test_webdataset_decoding(ray_start_2_cpus, tmp_path):
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
|
||||
gray = np.random.randint(0, 255, (100, 100), dtype=np.uint8)
|
||||
dstruct = dict(a=np.nan, b=dict(c=2), d="hello", e={"img_filename": "for_test.jpg"})
|
||||
ttensor = torch.tensor([1, 2, 3]).numpy()
|
||||
|
||||
sample = {
|
||||
"__key__": "foo",
|
||||
"jpg": image,
|
||||
"gray.png": gray,
|
||||
"mp": dstruct,
|
||||
"json": dstruct,
|
||||
"pt": ttensor,
|
||||
"und": b"undecoded",
|
||||
"custom": b"nothing",
|
||||
}
|
||||
|
||||
# write the encoded data using the default encoder
|
||||
data = [sample]
|
||||
ds = ray.data.from_items(data).repartition(1)
|
||||
ds.write_webdataset(path=tmp_path, try_create_dir=True)
|
||||
|
||||
ds = ray.data.read_webdataset(
|
||||
paths=[str(tmp_path)],
|
||||
override_num_blocks=1,
|
||||
decoder=None,
|
||||
)
|
||||
samples = ds.take(1)
|
||||
import json
|
||||
|
||||
meta_json = json.loads(samples[0]["json"].decode("utf-8"))
|
||||
assert meta_json["e"]["img_filename"] == "for_test.jpg"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("min_rows_per_file", [5, 10, 50])
|
||||
def test_write_min_rows_per_file(tmp_path, ray_start_2_cpus, min_rows_per_file):
|
||||
ray.data.from_items(
|
||||
[{"id": str(i)} for i in range(100)], override_num_blocks=20
|
||||
).write_webdataset(tmp_path, min_rows_per_file=min_rows_per_file)
|
||||
|
||||
for filename in os.listdir(tmp_path):
|
||||
dataset = wds.WebDataset(os.path.join(tmp_path, filename))
|
||||
assert len(list(dataset)) == min_rows_per_file
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filename",
|
||||
["000000.pkl", "000000.pickle", "000000.pt", "000000.pth"],
|
||||
)
|
||||
def test_default_decoder_rejects_unsafe_extensions(
|
||||
ray_start_2_cpus, tmp_path, filename
|
||||
):
|
||||
path = os.path.join(tmp_path, "unsafe.tar")
|
||||
with TarWriter(path) as tf:
|
||||
tf.write(filename, b"fake-payload")
|
||||
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)])
|
||||
with pytest.raises(Exception, match="Refusing to"):
|
||||
ds.take_all()
|
||||
|
||||
|
||||
def test_default_decoder_allows_unsafe_with_env_var(
|
||||
ray_start_2_cpus, tmp_path, allow_unsafe_deserialization
|
||||
):
|
||||
path = os.path.join(tmp_path, "trusted.tar")
|
||||
with TarWriter(path) as tf:
|
||||
tf.write("000000.pkl", pickle.dumps({"key": "value"}))
|
||||
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)])
|
||||
rows = ds.take_all()
|
||||
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["pkl"] == {"key": "value"}
|
||||
|
||||
|
||||
def test_custom_decoder_bypasses_unsafe_guard(ray_start_2_cpus, tmp_path):
|
||||
path = os.path.join(tmp_path, "custom.tar")
|
||||
with TarWriter(path) as tf:
|
||||
tf.write("000000.pkl", pickle.dumps({"key": "value"}))
|
||||
|
||||
def safe_pkl_decoder(sample):
|
||||
sample = dict(sample)
|
||||
for key, value in sample.items():
|
||||
if key == "pkl":
|
||||
sample[key] = pickle.loads(value)
|
||||
return sample
|
||||
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)], decoder=safe_pkl_decoder)
|
||||
rows = ds.take_all()
|
||||
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["pkl"] == {"key": "value"}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
Reference in New Issue
Block a user