def test_dataset_shard_error_with_unsupported_dataset_format(self):
"""Tests whether the dataset_shard function raises an error when an unsupported
dataset format is specified."""
config = {
"input": "dataset",
"input_config": {
"format": "__UNSUPPORTED_FORMAT__",
"paths": self.dset_path,
},
}
with self.assertRaises(ValueError):
get_dataset_and_shards(config)
def test_dataset_shard_error_with_both_format_and_loader_fn(self):
"""Tests whether the dataset_shard function raises an error when both format
and loader_fn are specified."""
dset = ray.data.range(100)
config = {
"input": "dataset",
"input_config": {
"format": "json",
"paths": self.dset_path,
"loader_fn": lambda: dset,
},
}
with self.assertRaises(ValueError):
get_dataset_and_shards(config)
def test_dataset_reader_itr_batches(self):
"""Test that the dataset reader iterates over batches of rows correctly."""
input_config = {"format": "json", "paths": self.dset_path}
dataset, _ = get_dataset_and_shards({
"input": "dataset",
"input_config": input_config
})
ioctx = IOContext(config={"train_batch_size": 1200}, worker_index=0)
reader = DatasetReader(ioctx, dataset)
assert len(reader.next()) >= 1200
def test_dataset_shard_with_loader_fn(self):
"""Tests whether the dataset_shard function works correctly with loader_fn."""
dset = ray.data.range(100)
config = {
"input": "dataset",
"input_config": {
"loader_fn": lambda: dset
}
}
ret_dataset, _ = get_dataset_and_shards(config)
assert ret_dataset.count() == dset.count()
def test_dataset_shard_with_only_local(self):
"""Tests whether the dataset_shard function works correctly for a single shard
for the local worker."""
config = {
"input": "dataset",
"input_config": {
"format": "json",
"paths": self.dset_path
},
}
# two ways of doing this:
# we have no remote workers
_, shards = get_dataset_and_shards(config, num_workers=0)
assert len(shards) == 1
assert isinstance(shards[0], ray.data.Dataset)
def test_itr_batches(self):
"""Test that the json reader iterates over batches of rows correctly."""
rllib_dir = Path(__file__).parent.parent.parent.parent
print("rllib dir={}".format(rllib_dir))
data_file = os.path.join(rllib_dir,
"rllib/tests/data/pendulum/large.json")
print("data_file={} exists={}".format(data_file,
os.path.isfile(data_file)))
input_config = {"format": "json", "path": data_file}
dataset, _ = get_dataset_and_shards(
{
"input": "dataset",
"input_config": input_config
}, 0, True)
ioctx = IOContext(config={"train_batch_size": 1200}, worker_index=0)
reader = DatasetReader(ioctx, dataset)
assert len(reader.next()) == 1200
def test_dataset_shard_with_task_parallelization(self):
"""Tests whether the dataset_shard function works correctly with parallelism
for reading the dataset."""
config = {
"input": "dataset",
"input_config": {
"format": "json",
"paths": self.dset_path,
"parallelism": 10,
},
}
NUM_WORKERS = 4
_, shards = get_dataset_and_shards(config, num_workers=NUM_WORKERS)
assert len(shards) == NUM_WORKERS + 1
assert shards[0] is None
assert all(
isinstance(remote_shard, ray.data.Dataset)
for remote_shard in shards[1:])
def test_dataset_shard_remote_workers_with_local_worker(self):
"""Tests whether the dataset_shard function works correctly for the remote
workers with a dummy dataset shard for the local worker."""
config = {
"input": "dataset",
"input_config": {
"format": "json",
"paths": self.dset_path
},
}
NUM_WORKERS = 4
_, shards = get_dataset_and_shards(config, num_workers=NUM_WORKERS)
assert len(shards) == NUM_WORKERS + 1
assert shards[0] is None
assert all(
isinstance(remote_shard, ray.data.Dataset)
for remote_shard in shards[1:])