update test_filter

update filter

state management

update filter

add filter

delete unnecessary files

Author: ramanishsingh

test/nodes/run_filter.py

@@ -4,16 +4,37 @@
 from torchdata.nodes.loader import Loader
 from torchdata.nodes.prefetch import Prefetcher
 from torchdata.nodes.samplers.stop_criteria import StopCriteria
+from utils import MockSource, run_test_save_load_state, StatefulRangeNode
 
-a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
 
-node = IterableWrapper(a)
 
+a = list(range(60))
+base_node = IterableWrapper(a)
 
 def is_even(x):
     return x % 2 == 0
+    
+node = Filter(base_node, is_even, num_workers=2)
 
+print(node.get_state())
+for _ in range(28):
+    print(next(node))
+print(node.get_state())
 
-filtered_node = Filter(node, is_even, num_workers=2)
-for item in filtered_node:
+state = node.get_state()
+node.reset()
+
+print(node.get_state())
+
+for _ in range(2):
+    print(next(node))
+
+del node
+node = Filter(base_node, is_even, num_workers=2)
+print("state to be loaded", state)
+print("state before reset", node.get_state())
+node.reset(state)
+print(node.get_state())
+
+for item in node:
     print(item)

test/nodes/test_filter.py

@@ -15,9 +15,10 @@
 
 from torchdata.nodes.base_node import BaseNode
 from torchdata.nodes.filter import Filter
+from torchdata.nodes.batch import Batcher
 
 from .utils import MockSource, run_test_save_load_state, StatefulRangeNode
-
+from torchdata.nodes.samplers.multi_node_weighted_sampler import MultiNodeWeightedSampler
 
 class TestFilter(TestCase):
     def _test_filter(self, num_workers, in_order, method):
@@ -52,25 +53,37 @@ def test_filter_parallel_threads(self):
     def test_filter_parallel_process(self):
         self._test_filter(num_workers=4, in_order=True, method="process")
 
+    @parameterized.expand([100, 200, 300])
+    def test_filter_batcher(self, n):
+        src = StatefulRangeNode(n=n)
+        node = Batcher(src, batch_size=2)
+        predicate = lambda x : (x[0]["i"]+x[1]["i"])%3==0
+        node = Filter(node, predicate, num_workers=2)
+        results = list(node)
+        self.assertEqual(len(results), n//6)
+
+
+
+
     @parameterized.expand(
         itertools.product(
-            [0],  # , 7, 13],
+            [10, 20 , 40],
             [True],  # TODO: define and fix in_order = False
-            [0],  # , 1, 9],  # TODO: define and fix in_order = False
+            [1, 2, 4],  
         )
     )
     def test_save_load_state_thread(
         self, midpoint: int, in_order: bool, snapshot_frequency: int
     ):
         method = "thread"
         n = 100
-        predicate = lambda x: True
+        predicate = lambda x: x["i"]%2==0
         src = StatefulRangeNode(n=n)
 
         node = Filter(
             source=src,
             predicate=predicate,
-            num_workers=4,
+            num_workers=1,
             in_order=in_order,
             method=method,
             snapshot_frequency=snapshot_frequency,

torchdata/nodes/filter.py

@@ -1,18 +1,21 @@
-# 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.
 from typing import Any, Callable, Dict, Iterator, Literal, Optional, TypeVar
-
 from torchdata.nodes.base_node import BaseNode
-from torchdata.nodes.map import Mapper, ParallelMapper
-
+from torchdata.nodes.map import ParallelMapper
 T = TypeVar("T", covariant=True)
 
-
 class Filter(BaseNode[T]):
+    """
+    A node that filters data samples based on a given predicate.
+    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. Defaults to 0.
+        in_order (bool): Whether to return items in the order from which they arrive from. Default is True.
+        method (Literal["thread", "process"]): The method to use for parallel processing. Default is "thread".
+        multiprocessing_context (Optional[str]): The multiprocessing context to use for parallel processing. Default is None.
+        max_concurrent (Optional[int]): The maximum number of items to process at once. Default is None.
+        snapshot_frequency (int): The frequency at which to snapshot the state of the source node. Default is 1.
+    """
     def __init__(
         self,
         source: BaseNode[T],
@@ -33,12 +36,6 @@ def __init__(
         self.multiprocessing_context = multiprocessing_context
         self.max_concurrent = max_concurrent
         self.snapshot_frequency = snapshot_frequency
-        self._it: Optional[Iterator[T]] = None
-
-    def reset(self, initial_state: Optional[Dict[str, Any]] = None):
-        super().reset(initial_state)
-        if self._it is not None:
-            del self._it
         if self.num_workers > 0:
             self._it = _ParallelFilterIter(
                 source=self.source,
@@ -49,54 +46,60 @@ def reset(self, initial_state: Optional[Dict[str, Any]] = None):
                 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]
-
+            self._it = _InlineFilterIter(source=self.source, predicate=self.predicate)
+    def reset(self, initial_state: Optional[Dict[str, Any]] = None) -> None:
+        """Resets the filter node to its initial state."""
+        super().reset(initial_state)
+        if self._it is not None:
+            self._it.reset(initial_state)
+    def next(self) -> T:
+        """Returns the next filtered item."""
+        return next(self._it)
     def get_state(self) -> Dict[str, Any]:
-        return self._it.get_state()  # type: ignore[union-attr]
-
+        """Returns the current state of the filter node."""
+        return self._it.get_state()
 
 class _InlineFilterIter(Iterator[T]):
-    def __init__(
-        self,
-        source: BaseNode[T],
-        predicate: Callable[[T], bool],
-        initial_state: Optional[Dict[str, Any]] = None,
-    ):
+    """
+    An iterator that filters data samples inline.
+    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.
+    """
+    SOURCE_KEY = "source"
+
+    def __init__(self, source: BaseNode[T], predicate: Callable[[T], bool]) -> None:
         self.source = source
         self.predicate = predicate
-        if initial_state is not None:
-            self.source.reset(initial_state["source"])
+
+    def reset(self, initial_state: Optional[Dict[str, Any]] = None) -> None:
+        """Resets the inline filter iterator to its initial state."""
+        if initial_state:
+            self.source.reset(initial_state[self.SOURCE_KEY])
         else:
             self.source.reset()
 
     def __iter__(self) -> Iterator[T]:
+        """Returns the iterator object itself."""
         return self
-
     def __next__(self) -> T:
+        """Returns the next filtered item."""
         while True:
-            item = next(self.source)
-            if self.predicate(item):
-                return item
-
+            try:
+                item = next(self.source)
+                if self.predicate(item):
+                    return item
+            except StopIteration:
+                raise
     def get_state(self) -> Dict[str, Any]:
-        return {"source": self.source.state_dict()}
-
+        """Returns the current state of the inline filter iterator."""
+        return {self.SOURCE_KEY: self.source.state_dict()}
 
 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.
@@ -106,9 +109,8 @@ class _ParallelFilterIter(Iterator[T]):
         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.
     """
-
+    MAPPER_KEY = "mapper"
     def __init__(
         self,
         source: BaseNode[T],
@@ -119,7 +121,6 @@ def __init__(
         multiprocessing_context: Optional[str],
         max_concurrent: Optional[int],
         snapshot_frequency: int,
-        initial_state: Optional[Dict[str, Any]] = None,
     ):
         self.source = source
         self.predicate = predicate
@@ -140,42 +141,26 @@ def __init__(
             max_concurrent=self.max_concurrent,
             snapshot_frequency=self.snapshot_frequency,
         )
-        if initial_state is not None:
-            self.mapper.reset(initial_state)
-
+    def reset(self, initial_state: Optional[Dict[str, Any]] = None) -> None:
+        """Resets the parallel filter iterator to its initial state."""
+        if initial_state:
+            self.mapper.reset(initial_state[self.MAPPER_KEY])
+        else:
+            self.mapper.reset()
     def __iter__(self) -> Iterator[T]:
-        """
-        Returns the iterator object itself.
-
-        Returns:
-            Iterator[T]: The iterator object itself.
-        """
+        """Returns the iterator object itself."""
         return self
-
     def __next__(self) -> T:
-        """
-        Returns the next filtered data sample.
-
-        Returns:
-            T: The next filtered data sample.
-        """
+        """Returns the next filtered item."""
         while True:
-            try:
-                item, passed_predicate = next(self.mapper)
-                if passed_predicate:
-                    return item
-            except StopIteration:
-                raise
-
-    def get_state(self) -> Dict[str, Any]:
-        """
-        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()
+            item, passed_predicate = next(self.mapper)
+            if passed_predicate:
+                return item
 
+    def get_state(self) -> Dict[str, Any]:
+        """Returns the current state of the parallel filter iterator."""
+        return {self.MAPPER_KEY: self.mapper.get_state()}
     def __del__(self):
         # Clean up resources when the iterator is deleted
         del self.mapper