add tests

Author: ramanishsingh

test/nodes/run_filter.py

@@ -0,0 +1,19 @@
+from torchdata.nodes.adapters import IterableWrapper
+from torchdata.nodes.batch import Batcher
+from torchdata.nodes.filter import Filter
+from torchdata.nodes.loader import Loader
+from torchdata.nodes.prefetch import Prefetcher
+from torchdata.nodes.samplers.stop_criteria import StopCriteria
+
+a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+
+node = IterableWrapper(a)
+
+
+def is_even(x):
+    return x % 2 == 0
+
+
+filtered_node = Filter(node, is_even, num_workers=2)
+for item in filtered_node:
+    print(item)

test/nodes/test_filter.py

@@ -0,0 +1,79 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import itertools
+import unittest
+from typing import List
+
+import torch
+
+from parameterized import parameterized
+from torch.testing._internal.common_utils import IS_WINDOWS, TEST_CUDA, TestCase
+
+from torchdata.nodes.base_node import BaseNode
+from torchdata.nodes.filter import Filter
+
+from .utils import MockSource, run_test_save_load_state, StatefulRangeNode
+
+
+class TestFilter(TestCase):
+    def _test_filter(self, num_workers, in_order, method):
+        n = 100
+        predicate = lambda x: x["test_tensor"] % 2 == 0  # Filter even numbers
+        src = MockSource(num_samples=n)
+        node = Filter(
+            source=src,
+            predicate=predicate,
+            num_workers=num_workers,
+            in_order=in_order,
+            method=method,
+        )
+
+        results: List[int] = []
+        for item in node:
+            results.append(item)
+
+        expected_results = [
+            {"step": i, "test_tensor": torch.tensor([i]), "test_str": f"str_{i}"}
+            for i in range(n)
+            if i % 2 == 0
+        ]
+        self.assertEqual(results, expected_results)
+
+    def test_filter_inline(self):
+        self._test_filter(num_workers=0, in_order=True, method="thread")
+
+    def test_filter_parallel_threads(self):
+        self._test_filter(num_workers=4, in_order=True, method="thread")
+
+    def test_filter_parallel_process(self):
+        self._test_filter(num_workers=4, in_order=True, method="process")
+
+    @parameterized.expand(
+        itertools.product(
+            [0],  # , 7, 13],
+            [True],  # TODO: define and fix in_order = False
+            [0],  # , 1, 9],  # TODO: define and fix in_order = False
+        )
+    )
+    def test_save_load_state_thread(
+        self, midpoint: int, in_order: bool, snapshot_frequency: int
+    ):
+        method = "thread"
+        n = 100
+        predicate = lambda x: True
+        src = StatefulRangeNode(n=n)
+
+        node = Filter(
+            source=src,
+            predicate=predicate,
+            num_workers=4,
+            in_order=in_order,
+            method=method,
+            snapshot_frequency=snapshot_frequency,
+        )
+        node.reset()
+        run_test_save_load_state(self, node, midpoint)

torchdata/nodes/filter.py

@@ -4,9 +4,12 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-from typing import Any, Callable, Dict, Iterator, Optional
+from typing import Any, Callable, Dict, Iterator, Literal, Optional, TypeVar
 
-from torchdata.nodes.base_node import BaseNode, T
+from torchdata.nodes.base_node import BaseNode
+from torchdata.nodes.map import Mapper, ParallelMapper
+
+T = TypeVar("T", covariant=True)
 
 
 class Filter(BaseNode[T]):
@@ -16,7 +19,7 @@ def __init__(
         predicate: Callable[[T], bool],
         num_workers: int = 0,
         in_order: bool = True,
-        method: str = "thread",
+        method: Literal["thread", "process"] = "thread",
         multiprocessing_context: Optional[str] = None,
         max_concurrent: Optional[int] = None,
         snapshot_frequency: int = 1,
@@ -37,29 +40,24 @@ def reset(self, initial_state: Optional[Dict[str, Any]] = None):
         if self._it is not None:
             del self._it
         if self.num_workers > 0:
-            self._parallel_reset(initial_state)
-        else:
-            self._inline_reset(initial_state)
-
-    def _inline_reset(self, initial_state: Optional[Dict[str, Any]]):
-        self._it = _InlineFilterIter(
-            source=self.source,
-            predicate=self.predicate,
-            initial_state=initial_state,
-        )
+            self._it = _ParallelFilterIter(
+                source=self.source,
+                predicate=self.predicate,
+                num_workers=self.num_workers,
+                in_order=self.in_order,
+                method=self.method,
+                multiprocessing_context=self.multiprocessing_context,
+                max_concurrent=self.max_concurrent,
+                snapshot_frequency=self.snapshot_frequency,
+                initial_state=initial_state,
+            )
 
-    def _parallel_reset(self, initial_state: Optional[Dict[str, Any]]):
-        self._it = _ParallelFilterIter(
-            source=self.source,
-            predicate=self.predicate,
-            num_workers=self.num_workers,
-            in_order=self.in_order,
-            method=self.method,
-            multiprocessing_context=self.multiprocessing_context,
-            max_concurrent=self.max_concurrent,
-            snapshot_frequency=self.snapshot_frequency,
-            initial_state=initial_state,
-        )
+        else:
+            self._it = _InlineFilterIter(
+                source=self.source,
+                predicate=self.predicate,
+                initial_state=initial_state,
+            )
 
     def next(self):
         return next(self._it)  # type: ignore[arg-type]
@@ -96,17 +94,32 @@ def get_state(self) -> Dict[str, Any]:
 
 
 class _ParallelFilterIter(Iterator[T]):
+    """
+    An iterator that filters data samples in parallel.
+
+    Args:
+        source (BaseNode[T]): The source node providing data samples.
+        predicate (Callable[[T], bool]): A function that takes a data sample and returns a boolean indicating whether to include it.
+        num_workers (int): The number of worker processes to use for parallel filtering.
+        in_order (bool): Whether to preserve the order of data samples.
+        method (Literal["thread", "process"]): The method to use for parallelization.
+        multiprocessing_context (Optional[str]): The multiprocessing context to use.
+        max_concurrent (Optional[int]): The maximum number of concurrent tasks.
+        snapshot_frequency (int): The frequency at which to take snapshots.
+        initial_state (Optional[Dict[str, Any]]): The initial state to start with.
+    """
+
     def __init__(
         self,
         source: BaseNode[T],
         predicate: Callable[[T], bool],
         num_workers: int,
         in_order: bool,
-        method: str,
+        method: Literal["thread", "process"],
         multiprocessing_context: Optional[str],
         max_concurrent: Optional[int],
         snapshot_frequency: int,
-        initial_state: Optional[Dict[str, Any]],
+        initial_state: Optional[Dict[str, Any]] = None,
     ):
         self.source = source
         self.predicate = predicate
@@ -116,83 +129,53 @@ def __init__(
         self.multiprocessing_context = multiprocessing_context
         self.max_concurrent = max_concurrent
         self.snapshot_frequency = snapshot_frequency
-        self._in_q: queue.Queue = queue.Queue()
-        self._out_q: queue.Queue = queue.Queue()
-        self._sem = threading.BoundedSemaphore(value=max_concurrent or 2 * num_workers)
-        self._stop_event = threading.Event()
-        self._workers: list[threading.Thread] = []
-        for _ in range(num_workers):
-            t = threading.Thread(
-                target=self._filter_worker,
-                args=(self._in_q, self._out_q, self.predicate),
-                daemon=True,
-            )
-            t.start()
-            self._workers.append(t)
-        self._populate_queue_thread = threading.Thread(
-            target=_populate_queue,
-            args=(
-                self.source,
-                self._in_q,
-                QueueSnapshotStore(),
-                snapshot_frequency,
-                self._sem,
-                self._stop_event,
-            ),
-            daemon=True,
+        # Create a ParallelMapper to filter items in parallel
+        self.mapper = ParallelMapper(
+            source=self.source,
+            map_fn=lambda x: (x, self.predicate(x)),
+            num_workers=self.num_workers,
+            in_order=self.in_order,
+            method=self.method,
+            multiprocessing_context=self.multiprocessing_context,
+            max_concurrent=self.max_concurrent,
+            snapshot_frequency=self.snapshot_frequency,
         )
-
-        self._populate_queue_thread.start()
         if initial_state is not None:
-            self.source.reset(initial_state["source"])
-        else:
-            self.source.reset()
-
-    def _filter_worker(
-        self, in_q: queue.Queue, out_q: queue.Queue, predicate: Callable[[T], bool]
-    ) -> None:
-        while True:
-            try:
-                item = in_q.get(block=True, timeout=0.1)
-            except queue.Empty:
-                if self._stop_event.is_set():
-                    break
-                continue
-            if isinstance(item, StopIteration):
-                out_q.put(item)
-                break
-            elif predicate(item):
-                out_q.put(item)
-            self._sem.release()
+            self.mapper.reset(initial_state)
 
     def __iter__(self) -> Iterator[T]:
+        """
+        Returns the iterator object itself.
+
+        Returns:
+            Iterator[T]: The iterator object itself.
+        """
         return self
 
     def __next__(self) -> T:
+        """
+        Returns the next filtered data sample.
+
+        Returns:
+            T: The next filtered data sample.
+        """
         while True:
             try:
-                item = self._out_q.get(block=True, timeout=0.1)
-            except queue.Empty:
-                if self._stop_event.is_set():
-                    raise StopIteration()
-                continue
-            if isinstance(item, StopIteration):
-                raise item
-            return item
+                item, passed_predicate = next(self.mapper)
+                if passed_predicate:
+                    return item
+            except StopIteration:
+                raise
 
     def get_state(self) -> Dict[str, Any]:
-        return {"source": self.source.state_dict()}
+        """
+        Returns the current state of the parallel filter iterator.
+
+        Returns:
+            Dict[str, Any]: The current state of the parallel filter iterator.
+        """
+        return self.mapper.get_state()
 
     def __del__(self):
-        self._shutdown()
-
-    def _shutdown(self):
-        self._stop_event.set()
-        if (
-            hasattr(self, "_populate_queue_thread")
-            and self._populate_queue_thread.is_alive()
-        ):
-            self._populate_queue_thread.join(timeout=0.5)
-        for t in self._workers:
-            if t.is_alive():
-                t.join(timeout=0.5)
+        # Clean up resources when the iterator is deleted
+        del self.mapper