[WIP] Implement round robin and separate ventilator and result queue

petastorm/arrow_reader_worker.py

@@ -15,6 +15,7 @@
 
 import hashlib
 import operator
+import logging
 
 import numpy as np
 import pandas as pd
@@ -26,6 +27,9 @@
 from petastorm.workers_pool import EmptyResultError
 from petastorm.workers_pool.worker_base import WorkerBase
 
+# Initialize logger
+logger = logging.getLogger(__name__)
+
 
 class ArrowReaderWorkerResultsQueueReader(object):
     def __init__(self):
@@ -91,6 +95,9 @@ class ArrowReaderWorker(WorkerBase):
     def __init__(self, worker_id, publish_func, args):
         super(ArrowReaderWorker, self).__init__(worker_id, publish_func, args)
 
+        # Add debug log in the constructor
+        print(f'DEBUG: Initializing ArrowReaderWorker with worker_id: {worker_id}')
+
         self._filesystem = args[0]
         self._dataset_path_or_paths = args[1]
         self._schema = args[2]
@@ -101,7 +108,10 @@ def __init__(self, worker_id, publish_func, args):
         self._transformed_schema = args[7]
         self._arrow_filters = args[8]
         self._shuffle_rows = args[9]
-        self._random_state = np.random.RandomState(seed=args[10])
+        self._random_seed = args[10]
+
+        # Initialize random number generator
+        self._rng = np.random.default_rng(self._random_seed)
 
         if self._ngram:
             raise NotImplementedError('ngrams are not supported by ArrowReaderWorker')
@@ -128,12 +138,18 @@ def process(self, piece_index, worker_predicate, shuffle_row_drop_partition):
         :return:
         """
 
+        # Add debug log in the process method
+        print(f'DEBUG: Processing piece_index: {piece_index}')
+
         if not self._dataset:
             self._dataset = pq.ParquetDataset(
                 self._dataset_path_or_paths,
                 filesystem=self._filesystem,
                 validate_schema=False, filters=self._arrow_filters)
 
+            # Add debug log after dataset is initialized
+            print(f'DEBUG: ParquetDataset initialized with path: {self._dataset_path_or_paths}')
+
         piece = self._split_pieces[piece_index]
 
         # Create pyarrow file system
@@ -160,11 +176,16 @@ def process(self, piece_index, worker_predicate, shuffle_row_drop_partition):
                 path_str = self._dataset_path_or_paths
             cache_key = '{}:{}:{}'.format(hashlib.md5(path_str.encode('utf-8')).hexdigest(),
                                           piece.path, piece_index)
+
+            # Add debug log for cache key
+            print(f'DEBUG: Cache key generated: {cache_key}')
+
             all_cols = self._local_cache.get(cache_key,
                                              lambda: self._load_rows(parquet_file, piece, shuffle_row_drop_partition))
 
         if all_cols:
             self.publish_func(all_cols)
+            print(f'DEBUG: Published columns for piece_index: {piece_index}')
 
     @staticmethod
     def _check_shape_and_ravel(x, field):
@@ -289,9 +310,19 @@ def _read_with_shuffle_row_drop(self, piece, pq_file, column_names, shuffle_row_
 
         # pyarrow would fail if we request a column names that the dataset is partitioned by
         table = piece.read(columns=column_names - partition_names, partitions=self._dataset.partitions)
+
+        # Handle row shuffling based on shuffle_rows setting
         if self._shuffle_rows:
-            indices = self._random_state.permutation(table.num_rows)
-            table = table.take(indices)
+            if self._random_seed is not None and self._random_seed != 0:
+                # Deterministic randomization: use provided seed
+                indices = self._rng.permutation(table.num_rows)
+            else:
+                # Non-deterministic randomization: use np.random directly
+                indices = np.random.permutation(table.num_rows)
+        else:
+            # Deterministic natural order: shuffle_rows=False
+            indices = np.arange(table.num_rows)
+        table = table.take(indices)
 
         # Drop columns we did not explicitly request. This may happen when a table is partitioned. Besides columns
         # requested, pyarrow will also return partition values. Having these unexpected fields will break some

petastorm/reader.py

@@ -38,6 +38,7 @@
 from petastorm.workers_pool.thread_pool import ThreadPool
 from petastorm.workers_pool.ventilator import ConcurrentVentilator
 
+# Initialize logger
 logger = logging.getLogger(__name__)
 
 # Ventilator guarantees that no more than workers + _VENTILATE_EXTRA_ROWGROUPS are processed at a moment by a
@@ -159,7 +160,7 @@ def make_reader(dataset_url,
                       'To read from a non-Petastorm Parquet store use make_batch_reader')
 
     if reader_pool_type == 'thread':
-        reader_pool = ThreadPool(workers_count, results_queue_size)
+        reader_pool = ThreadPool(workers_count, results_queue_size, shuffle_rows=shuffle_rows, seed=seed)
     elif reader_pool_type == 'process':
         if pyarrow_serialize:
             warnings.warn("pyarrow_serializer was deprecated and will be removed in future versions. "
@@ -315,7 +316,7 @@ def make_batch_reader(dataset_url_or_urls,
         raise ValueError('Unknown cache_type: {}'.format(cache_type))
 
     if reader_pool_type == 'thread':
-        reader_pool = ThreadPool(workers_count, results_queue_size)
+        reader_pool = ThreadPool(workers_count, results_queue_size, shuffle_rows=shuffle_rows, seed=seed)
     elif reader_pool_type == 'process':
         serializer = ArrowTableSerializer()
         reader_pool = ProcessPool(workers_count, serializer, zmq_copy_buffers=zmq_copy_buffers)
@@ -400,6 +401,7 @@ def __init__(self, pyarrow_filesystem, dataset_path, schema_fields=None,
             These will be applied when loading the parquet file with PyArrow. More information
             here: https://arrow.apache.org/docs/python/generated/pyarrow.parquet.ParquetDataset.html
         """
+        print(f'DEBUG: Initializing Reader with dataset_path: {dataset_path}, num_epochs: {num_epochs}')
         self.num_epochs = num_epochs
 
         # 1. Open the parquet storage (dataset)
@@ -437,9 +439,11 @@ def __init__(self, pyarrow_filesystem, dataset_path, schema_fields=None,
             raise NotImplementedError('Using timestamp_overlap=False is not implemented with'
                                       ' shuffle_options.shuffle_row_drop_partitions > 1')
 
+        print(f'DEBUG: Reader initialized with schema_fields: {schema_fields}')
+
         cache = cache or NullCache()
 
-        self._workers_pool = reader_pool or ThreadPool(10)
+        self._workers_pool = reader_pool or ThreadPool(10, shuffle_rows=shuffle_rows, seed=seed)
 
         # Make a schema view (a view is a Unischema containing only a subset of fields
         # Will raise an exception if invalid schema fields are in schema_fields
@@ -483,7 +487,7 @@ def __init__(self, pyarrow_filesystem, dataset_path, schema_fields=None,
                                                 self.ngram, row_groups, cache, transform_spec,
                                                 self.schema, filters, shuffle_rows, seed),
                                  ventilator=self.ventilator)
-        logger.debug('Workers pool started')
+        print('DEBUG: Workers pool started')
 
         self.last_row_consumed = False
         self.stopped = False
@@ -653,6 +657,7 @@ def _normalize_shuffle_options(shuffle_row_drop_partitions, dataset):
 
     def _create_ventilator(self, row_group_indexes, shuffle_row_groups, shuffle_row_drop_partitions,
                            num_epochs, worker_predicate, max_ventilation_queue_size, seed):
+        print(f'DEBUG: Creating ventilator with row_group_indexes: {row_group_indexes}')
         items_to_ventilate = []
         for piece_index in row_group_indexes:
             for shuffle_row_drop_partition in range(shuffle_row_drop_partitions):
@@ -670,12 +675,12 @@ def _create_ventilator(self, row_group_indexes, shuffle_row_groups, shuffle_row_
                                     random_seed=seed)
 
     def stop(self):
-        """Stops all worker threads/processes."""
+        print('DEBUG: Stopping Reader')
         self._workers_pool.stop()
         self.stopped = True
 
     def join(self):
-        """Joins all worker threads/processes. Will block until all worker workers have been fully terminated."""
+        print('DEBUG: Joining Reader')
         self._workers_pool.join()
 
     @property

petastorm/tests/test_tf_dataset.py

@@ -128,6 +128,7 @@ def test_with_dataset_repeat(synthetic_dataset, reader_factory):
 def test_with_dataset_repeat_after_cache(synthetic_dataset, reader_factory):
     """ Check if ``tf.data.Dataset``'s ``repeat`` works after ``tf.data.Dataset``'s ``cache``."""
     epochs = 3
+    print(f"Starting test_with_dataset_repeat_after_cache with {epochs} epochs")
     with reader_factory(synthetic_dataset.url, schema_fields=[TestSchema.id]) as reader:
         dataset = make_petastorm_dataset(reader)
         dataset = dataset.cache()
@@ -138,18 +139,22 @@ def test_with_dataset_repeat_after_cache(synthetic_dataset, reader_factory):
         with tf.Session() as sess:
             with pytest.warns(None):
                 # Expect no warnings since cache() is called before repeat()
-                for _ in range(epochs):
+                for epoch in range(epochs):
+                    print(f"Starting epoch {epoch}")
                     actual_res = []
-                    for _, _ in enumerate(synthetic_dataset.data):
+                    for i, _ in enumerate(synthetic_dataset.data):
                         actual = sess.run(it_op)._asdict()
                         actual_res.append(actual["id"])
+                        print(f"iteration: {i} {actual['id']}")
                     expected_res = list(range(len(synthetic_dataset.data)))
+                    print(f"Epoch: {epoch} actual {sorted(actual_res)}, expected {expected_res}")
                     # sort dataset output since row_groups are shuffled from reader.
                     np.testing.assert_equal(sorted(actual_res), expected_res)
-
+                    print(f"Completed epoch {epoch}")
             # Exhausted all epochs. Fetching next value should trigger OutOfRangeError
             with pytest.raises(tf.errors.OutOfRangeError):
                 sess.run(it_op)
+    print("Completed test_with_dataset_repeat_after_cache")
 
 
 @pytest.mark.forked

petastorm/workers_pool/tests/test_workers_pool.py

@@ -141,15 +141,17 @@ def test_stop_when_result_queue_is_full(self):
         SLEEP_DELTA = 0.01
         TIMEOUT = 20
         QUEUE_SIZE = 2
+        WORKERS_COUNT = 10
 
-        pool = ThreadPool(10, results_queue_size=QUEUE_SIZE)
+        pool = ThreadPool(WORKERS_COUNT, results_queue_size=QUEUE_SIZE)
         pool.start(WorkerIdGeneratingWorker)
 
-        for _ in range(100):
+        for _ in range(1000):
             pool.ventilate()
 
+        expected_queue_size = WORKERS_COUNT * max(5, QUEUE_SIZE // WORKERS_COUNT)
         cumulative_wait = 0
-        while pool.results_qsize() != QUEUE_SIZE:
+        while pool.results_qsize() != expected_queue_size:
             time.sleep(SLEEP_DELTA)
             cumulative_wait += SLEEP_DELTA
             # Make sure we wait no longer than the timeout. Otherwise, something is very wrong