Revert default to record parallelizm

Author: ivan-chai

hotpp/calibrate.py

@@ -5,7 +5,7 @@
 import torch
 from omegaconf import OmegaConf
 
-from hotpp.data import ShuffledDistributedDataset
+from hotpp.data import ShuffledDistributedDataset, DEFAULT_PARALLELIZM
 from hotpp.data.module import HotppSampler
 from tqdm import tqdm
 
@@ -19,6 +19,7 @@ def get_loader(dm):
     dataset = ShuffledDistributedDataset(dm.val_data, rank=None, world_size=None,
                                          num_workers=loader_params.get("num_workers", 0),
                                          cache_size=loader_params.pop("cache_size", 4096),
+                                         parallelize=loader_params.pop("parallelize", DEFAULT_PARALLELIZM),
                                          seed=loader_params.pop("seed", 0))
     loader = torch.utils.data.DataLoader(
         dataset=dataset,

hotpp/data/__init__.py

@@ -1,3 +1,3 @@
-from .dataset import HotppDataset, ShuffledDistributedDataset
+from .dataset import HotppDataset, ShuffledDistributedDataset, DEFAULT_PARALLELIZM
 from .module import HotppDataModule
 from .padded_batch import PaddedBatch

hotpp/data/dataset.py

@@ -17,6 +17,9 @@
 from .padded_batch import PaddedBatch
 
 
+DEFAULT_PARALLELIZM = "records"
+
+
 def immutable_hash(s):
     return int(hashlib.sha256(s.encode("utf-8")).hexdigest(), 16)
 
@@ -90,18 +93,16 @@ def __init__(self, data,
                  add_seq_fields=None,
                  global_target_fields=None,
                  local_targets_fields=None,
-                 local_targets_indices_field=None,
-                 allow_empty=False):
+                 local_targets_indices_field=None):
         super().__init__()
         if isinstance(data, str):
             self.filenames = list(sorted(parquet_file_scan(data)))
         elif isinstance(data, list):
             self.filenames = data
         else:
             raise ValueError(f"Unknown data type: {type(data)}")
-        if (not self.filenames) and (not allow_empty):
+        if not self.filenames:
             raise RuntimeError("Empty dataset")
-        self.allow_empty = allow_empty
         self.total_length = sum(map(get_parquet_length, self.filenames))
         self.random_split = random_split
         self.random_part = random_part
@@ -268,7 +269,7 @@ class ShuffledDistributedDataset(torch.utils.data.IterableDataset):
     Args:
         parallelize: Parallel reading mode, either `records` (better granularity) or `files` (faster).
     """
-    def __init__(self, dataset, rank=None, world_size=None, cache_size=None, parallelize="files", seed=0):
+    def __init__(self, dataset, rank=None, world_size=None, cache_size=None, parallelize=DEFAULT_PARALLELIZM, seed=0):
         super().__init__()
         self.dataset = dataset
         self.rank = rank
@@ -311,13 +312,16 @@ def _iter_shuffled_files(self, dataset, seed, rank, world_size):
         filenames = list(dataset.filenames)
         if not filenames:
             raise RuntimeError("Empty dataset")
-        if len(filenames) < world_size:
-            warnings.warn(f"{len(filenames)} files for {world_size} workers, switch to record parallelizm")
+        root = os.path.commonprefix(filenames)
+        splits = [list() for _ in range(world_size)]
+        for filename in filenames:
+            splits[immutable_hash(os.path.relpath(filename, root)) % world_size].append(filename)
+        if any([len(split) == 0 for split in splits]):
+            if rank == 0:
+                warnings.warn(f"Some workers got zero files, switch to record parallelizm")
             yield from self._iter_shuffled_records(dataset, seed, rank, world_size)
             return
-        root = os.path.commonprefix(filenames)
-        subset = [filename for filename in filenames if immutable_hash(os.path.relpath(filename, root)) % world_size == rank]
-        dataset = dataset.replace_files(subset, allow_empty=True)
+        dataset = dataset.replace_files(splits[rank])
         yield from self._iter_shuffled_records_impl(dataset, seed)
 
     def _iter_shuffled_records(self, dataset, seed, rank, world_size):

hotpp/data/module.py

@@ -1,6 +1,6 @@
 import torch
 import pytorch_lightning as pl
-from .dataset import HotppDataset, ShuffledDistributedDataset
+from .dataset import HotppDataset, ShuffledDistributedDataset, DEFAULT_PARALLELIZM
 
 
 def pop_loader_params(params):
@@ -102,7 +102,7 @@ def train_dataloader(self, rank=None, world_size=None):
         loader_params.update(self.train_loader_params)
         dataset = ShuffledDistributedDataset(self.train_data, rank=rank, world_size=world_size,
                                              cache_size=loader_params.pop("cache_size", 4096),
-                                             parallelize=loader_params.pop("parallelize", "files"),
+                                             parallelize=loader_params.pop("parallelize", DEFAULT_PARALLELIZM),
                                              seed=loader_params.pop("seed", 0))
         loader = torch.utils.data.DataLoader(
             dataset=dataset,
@@ -118,7 +118,7 @@ def val_dataloader(self, rank=None, world_size=None):
         loader_params = {"pin_memory": torch.cuda.is_available()}
         loader_params.update(self.val_loader_params)
         dataset = ShuffledDistributedDataset(self.val_data, rank=rank, world_size=world_size,
-                                             parallelize=loader_params.pop("parallelize", "files"))  # Disable shuffle.
+                                             parallelize=loader_params.pop("parallelize", DEFAULT_PARALLELIZM))  # Disable shuffle.
         loader = torch.utils.data.DataLoader(
             dataset=dataset,
             collate_fn=dataset.dataset.collate_fn,
@@ -132,7 +132,7 @@ def test_dataloader(self, rank=None, world_size=None):
         loader_params = {"pin_memory": torch.cuda.is_available()}
         loader_params.update(self.test_loader_params)
         dataset = ShuffledDistributedDataset(self.test_data, rank=rank, world_size=world_size,
-                                             parallelize=loader_params.pop("parallelize", "files"))  # Disable shuffle.
+                                             parallelize=loader_params.pop("parallelize", DEFAULT_PARALLELIZM))  # Disable shuffle.
         loader = torch.utils.data.DataLoader(
             dataset=dataset,
             collate_fn=dataset.dataset.collate_fn,

hotpp/embed.py

@@ -13,7 +13,7 @@
 from torchmetrics.utilities import dim_zero_cat
 
 from .common import get_trainer
-from .data import ShuffledDistributedDataset
+from .data import ShuffledDistributedDataset, DEFAULT_PARALLELIZM
 
 logger = logging.getLogger(__name__)
 
@@ -136,7 +136,8 @@ def test_dataloader(self):
         loader_params = getattr(self.data, f"{self.split}_loader_params")
 
         num_workers = loader_params.get("num_workers", 0)
-        dataset = ShuffledDistributedDataset(dataset, rank=self.rank, world_size=self.world_size)
+        dataset = ShuffledDistributedDataset(dataset, rank=self.rank, world_size=self.world_size,
+                                             parallelize=loader_params.pop("parallelize", DEFAULT_PARALLELIZM))
         return torch.utils.data.DataLoader(
             dataset=dataset,
             collate_fn=dataset.dataset.collate_fn,

setup.py

@@ -15,7 +15,7 @@
 
 setuptools.setup(
     name="hotpp-benchmark",
-    version="0.6.4",
+    version="0.6.5",
     author="Ivan Karpukhin",
     author_email="karpuhini@yandex.ru",
     description="Evaluate generative event sequence models on the long horizon prediction task.",