Skip to content

Add usage example to ReadOptions class docstring #1425

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
copybara-service wants to merge 1 commit into from
Closed

Add usage example to ReadOptions class docstring #1425

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 0 additions & 3 deletions grain/_src/python/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -234,12 +234,9 @@ py_test(
srcs = ["options_test.py"],
srcs_version = "PY3",
deps = [
":data_sources",
":options",
"//grain/_src/python/dataset",
"@abseil-py//absl/logging",
"@abseil-py//absl/testing:absltest",
"@pypi//numpy:pkg",
],
)

Expand Down
50 changes: 3 additions & 47 deletions grain/_src/python/options.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,28 +30,9 @@ class ReadOptions:
These settings configure a single Python process. Each process uses separate
threads and buffer for reading and processing data.

Example:
With `ReadOptions.num_threads=8`, there will be 8 threads reading the data::

import grain

# Create a source with 4 elements.
source = grain.sources.RangeDataSource(0, 4, 1)

# Configure reading options for each Python process.
read_options = grain.ReadOptions(
num_threads=8,
prefetch_buffer_size=100,
)

# Use the read options to create an iterator dataset.
ds = (
grain.MapDataset.source(source)
.to_iter_dataset(read_options=read_options)
.batch(2)
)
print(list(ds))
# [array([0, 1]), array([2, 3])]
Example: With ReadOptions.num_threads=8 and
MultiprocessingOptions.num_workers=10 there will be 80 threads reading the
data (8 threads in each of 10 Python processes).

Attributes:
num_threads: Number of threads reading from the DataSource in parallel. If
Expand Down Expand Up @@ -105,31 +86,6 @@ def __post_init__(self):
class MultiprocessingOptions:
"""Options for using Python multiprocessing.

Example:
With `MultiprocessingOptions.num_workers=2`, there will be 2 worker
processes::

import grain

# Create a source with 4 elements.
source = grain.sources.RangeDataSource(0, 4, 1)

# Configure multiprocessing options.
multiprocessing_options = grain.MultiprocessingOptions(
num_workers=2,
per_worker_buffer_size=10,
)

# Use the multiprocessing options to create an iterator dataset.
ds = (
grain.MapDataset.source(source)
.to_iter_dataset()
.batch(2)
.mp_prefetch(multiprocessing_options)
)
print(list(ds))
# [SharedMemoryArray([0, 2]), SharedMemoryArray([1, 3])]

Attributes:
num_workers: Number of Python worker processes. More processes can speed up
the pipeline if it's compute bound and bottlenecked on the CPython's GIL.
Expand Down
49 changes: 0 additions & 49 deletions grain/_src/python/options_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,7 @@
from unittest import mock

from absl import logging
import multiprocessing as mp
from grain._src.python import data_sources
from grain._src.python import options
from grain._src.python.dataset import dataset
import numpy as np

from absl.testing import absltest

Expand All @@ -17,25 +13,6 @@ def test_defaults(self):
self.assertEqual(ro.num_threads, 16)
self.assertEqual(ro.prefetch_buffer_size, 500)

def test_example_docstring(self):
source = data_sources.RangeDataSource(0, 4, 1)
read_options = options.ReadOptions(
num_threads=8,
prefetch_buffer_size=100,
)
ds = (
dataset.MapDataset.source(source)
.to_iter_dataset(read_options=read_options)
.batch(2)
)
np.testing.assert_equal(
list(ds),
[
np.array([0, 1]),
np.array([2, 3]),
],
)

def test_num_threads_negative_raises_value_error(self):
with self.assertRaisesRegex(ValueError, "num_threads must be non-negative"):
options.ReadOptions(num_threads=-1)
Expand All @@ -59,31 +36,5 @@ def test_prefetch_buffer_size_zero(self):
mock_warning.assert_not_called()


class MultiprocessingOptionsTest(absltest.TestCase):

def test_example_docstring(self):
read_options = options.ReadOptions(
num_threads=8,
prefetch_buffer_size=10,
)
multiprocessing_options = options.MultiprocessingOptions(
num_workers=2,
per_worker_buffer_size=10,
)
ds = (
dataset.MapDataset.range(4)
.to_iter_dataset(read_options=read_options)
.batch(2)
.mp_prefetch(multiprocessing_options)
)
np.testing.assert_equal(
list(ds),
[
np.array([0, 2]),
np.array([1, 3]),
],
)


if __name__ == "__main__":
absltest.main()
Loading