f9columnar 0.2.1__tar.gz → 0.2.2__tar.gz

{f9columnar-0.2.1 → f9columnar-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: f9columnar
-Version: 0.2.1
+Version: 0.2.2
 Summary: Columnar analysis utils.
 Home-page: https://gitlab.cern.ch/ijs-f9-ljubljana/f9columnar
 License: MIT

{f9columnar-0.2.1 → f9columnar-0.2.2}/f9columnar/hdf5_dataloader.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import copy
 import json
 import logging
-import multiprocessing
 import os
 from collections.abc import Callable
 from itertools import product
@@ -11,8 +10,10 @@ from typing import Any
 
 import h5py
 import numpy as np
+import pandas as pd
 import torch
 from numpy.typing import ArrayLike
+from torch import multiprocessing
 from torch.utils.data import DataLoader, IterableDataset
 
 from f9columnar.processors import Processor, ProcessorsGraph
@@ -95,7 +96,7 @@ class Hdf5Iterator:
         return arrays, reports
 
 
-class Hdf5IteratorMaker:
+class Hdf5IteratorDfMaker:
     def __init__(
         self,
         name: str,
@@ -115,6 +116,7 @@ class Hdf5IteratorMaker:
         self.holdout_range = holdout_range
 
         self.total_num_entries = shape[0]
+        self.all_num_entries_dct: dict[str, int] = {}
 
     def _log_info(self) -> None:
         all_hdf5_files = list(self.hdf5_files_metadata.keys())
@@ -143,15 +145,15 @@ class Hdf5IteratorMaker:
             holdout_range=self.holdout_range,
         )
 
-    def make(self) -> list[list[Hdf5Iterator]]:
+    def _split(self) -> list[dict[str, list[int]]]:
         self._log_info()
 
         # how many entries each worker will process
         splits = [self.total_num_entries // self.num_workers] * self.num_workers
         splits[-1] += self.total_num_entries % self.num_workers
 
-        all_num_entries_dct = {file: metadata["shape"][0] for file, metadata in self.hdf5_files_metadata.items()}
-        num_entries_dct = copy.deepcopy(all_num_entries_dct)
+        self.all_num_entries_dct = {file: metadata["shape"][0] for file, metadata in self.hdf5_files_metadata.items()}
+        num_entries_dct = copy.deepcopy(self.all_num_entries_dct)
 
         # keep track of the start and stop entries for each root file
         hdf5_files_start_dct: dict[str, int] = {file: 0 for file in self.hdf5_files_metadata.keys()}
@@ -171,7 +173,7 @@ class Hdf5IteratorMaker:
                 total += num_entries
 
                 if total <= split:
-                    result[i][hdf5_file] = [start_entry, all_num_entries_dct[hdf5_file]]
+                    result[i][hdf5_file] = [start_entry, self.all_num_entries_dct[hdf5_file]]
                     done.append(hdf5_file)
 
                     if total == split:
@@ -186,39 +188,74 @@ class Hdf5IteratorMaker:
                     num_entries_dct[hdf5_file] -= delta
                     break
 
-        worker_hdf5_iterators: list[list] = [[] for _ in range(self.num_workers)]
-        check_total = 0
+        return result
 
-        for i, result_dct in enumerate(result):
-            for hdf5_file, start_end in result_dct.items():
-                entry_start, entry_end = start_end
-                check_total += entry_end - entry_start
+    def make(self) -> pd.DataFrame:
+        split_result = self._split()
 
-                iterator = self._make_iterator(hdf5_file, entry_start, entry_end)
-                worker_hdf5_iterators[i].append(iterator)
+        worker_df: dict[str, list] = {
+            "worker_id": [],
+            "file": [],
+            "start": [],
+            "stop": [],
+            "chunk_size": [],
+            "holdout_range": [],
+            "shuffle": [],
+        }
+
+        check_total = 0
+        for i, result_dct in enumerate(split_result):
+            for hdf5_file, start_stop in result_dct.items():
+                entry_start, entry_stop = start_stop
+                check_total += entry_stop - entry_start
+
+                worker_df["worker_id"].append(i)
+                worker_df["file"].append(hdf5_file)
+                worker_df["start"].append(entry_start)
+                worker_df["stop"].append(entry_stop)
+                worker_df["chunk_size"].append(self.chunk_size)
+                worker_df["holdout_range"].append(self.holdout_range)
+                worker_df["shuffle"].append(self.shuffle)
 
         if check_total != self.total_num_entries:
             raise ValueError("Total number of entries does not match.")
 
-        return worker_hdf5_iterators
+        return pd.DataFrame(worker_df)
 
 
 class Hdf5LoaderIterator:
     def __init__(
         self,
         name: str,
-        hdf5_iterators: list[Hdf5Iterator],
+        iterators_df: pd.DataFrame,
         worker_id: int,
         processors: list[Callable[[ArrayLike, dict], tuple[ArrayLike, dict]]] | ProcessorsGraph | None,
         hdf5_files_desc_dct: dict[str, dict[str, Any]] | None = None,
     ) -> None:
         self.name = name
-        self.hdf5_iterators = hdf5_iterators
+        self.iterators_df = iterators_df
         self.worker_id = worker_id
         self.processors = processors
         self.hdf5_files_desc_dct = hdf5_files_desc_dct
 
-        self.current_iterator_idx = 0
+        self.current_df_idx, self.current_iterator_idx = 0, 0
+
+    def _make_hdf5_iterator(self, df: pd.Series) -> Hdf5Iterator:
+        iterator = Hdf5Iterator(
+            df["file"],
+            chunk_size=df["chunk_size"],
+            start_entry=df["start"],
+            stop_entry=df["stop"],
+            shuffle=df["shuffle"],
+            holdout_range=df["holdout_range"],
+        )
+
+        return iterator
+
+    def _iterate_df(self) -> None:
+        df = self.iterators_df.iloc[self.current_iterator_idx]
+        self.iterator = self._make_hdf5_iterator(df)
+        self.current_df_idx += 1
 
     def _run_processors(self, arrays: ArrayLike, reports: dict) -> tuple[ArrayLike, dict] | dict[str, Processor]:
         if self.processors is None:
@@ -247,18 +284,22 @@ class Hdf5LoaderIterator:
 
     def __next__(self) -> tuple[ArrayLike, dict] | dict[str, Processor]:
         try:
-            current_iterator = self.hdf5_iterators[self.current_iterator_idx]
-            arrays, reports = next(current_iterator)
+            if self.current_df_idx == self.current_iterator_idx:
+                self._iterate_df()
+
+            arrays, reports = next(self.iterator)
 
         except StopIteration:
-            current_iterator.close()
+            self.iterator.close()
             self.current_iterator_idx += 1
 
-            if self.current_iterator_idx == len(self.hdf5_iterators):
+            if self.current_iterator_idx == len(self.iterators_df):
                 raise StopIteration
-            else:
-                current_iterator = self.hdf5_iterators[self.current_iterator_idx]
-                arrays, reports = next(current_iterator)
+
+            if self.current_df_idx == self.current_iterator_idx:
+                self._iterate_df()
+
+            arrays, reports = next(self.iterator)
 
         reports = self._make_report(reports)
 
@@ -308,7 +349,7 @@ class Hdf5IterableDataset(IterableDataset):
 
         self.shape = self._get_total_shape()
 
-        self.worker_hdf5_iterators = self._get_iterators()
+        self.worker_iterators_df = self._get_df_iterators()
 
     def _setup(self) -> None:
         for file, dataset_name in product(self.files, self.dataset_names):
@@ -384,9 +425,15 @@ class Hdf5IterableDataset(IterableDataset):
 
         return sum(shapes_0), shapes_1[0]
 
-    def _get_iterators(self) -> list[list[Hdf5Iterator]]:
-        return Hdf5IteratorMaker(
-            self.name, self.metadata, self.num_workers, self.shape, self.chunk_size, self.shuffle, self.holdout_range
+    def _get_df_iterators(self) -> pd.DataFrame:
+        return Hdf5IteratorDfMaker(
+            self.name,
+            self.metadata,
+            self.num_workers,
+            self.shape,
+            self.chunk_size,
+            self.shuffle,
+            self.holdout_range,
         ).make()
 
     def __iter__(self) -> Hdf5LoaderIterator:
@@ -397,17 +444,21 @@ class Hdf5IterableDataset(IterableDataset):
         else:
             worker_id = worker_info.id
 
-        hdf5_iterators = self.worker_hdf5_iterators[worker_id]
+        iterators_df = self.worker_iterators_df[self.worker_iterators_df["worker_id"] == worker_id].copy()
 
         return Hdf5LoaderIterator(
             self.name,
-            hdf5_iterators,
+            iterators_df,
             worker_id,
             self.processors,
             self.hdf5_files_desc_dct,
         )
 
 
+def default_collate_fn(batch: list[Any]) -> list[Any]:
+    return batch
+
+
 def get_hdf5_dataloader(
     name: str,
     files: str | list[str],
@@ -421,6 +472,10 @@ def get_hdf5_dataloader(
     processors: list[Callable[[ArrayLike, dict], tuple[ArrayLike, dict]]] | ProcessorsGraph | None = None,
     dataloader_kwargs: dict[str, Any] | None = None,
 ) -> tuple[DataLoader, int]:
+
+    if multiprocessing.get_start_method() == "fork" and num_workers > 0:
+        logging.debug("Using 'fork' start method. Consider using 'spawn' or 'forkserver'.")
+
     if num_workers == -1:
         num_workers = multiprocessing.cpu_count()
 
@@ -441,7 +496,7 @@ def get_hdf5_dataloader(
         name,
         files,
         dataset_names,
-        num_workers,
+        num_workers=num_workers if num_workers > 0 else 1,
         chunk_size=chunk_size,
         use_piles=use_piles,
         shuffle=shuffle,
@@ -454,7 +509,7 @@ def get_hdf5_dataloader(
         hdf5_dataset,
         batch_size=None,
         num_workers=num_workers,
-        collate_fn=lambda batch: batch,
+        collate_fn=default_collate_fn,
         **dataloader_kwargs,
     )
 

{f9columnar-0.2.1 → f9columnar-0.2.2}/f9columnar/root_dataloader.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import copy
 import logging
-import multiprocessing
 import os
 from collections.abc import Callable, Iterator
 from dataclasses import dataclass, field
@@ -12,6 +11,7 @@ import awkward as ak
 import pandas as pd
 import torch
 import uproot
+from torch import multiprocessing
 from torch.utils.data import DataLoader, IterableDataset
 
 from f9columnar.processors import Processor, ProcessorsGraph
@@ -316,13 +316,13 @@ class RootIterableDataset(IterableDataset):
     def __init__(
         self,
         name: str,
-        iterators_df: pd.DataFrame,
+        worker_iterators_df: pd.DataFrame,
         processors: list[Callable[[ak.Array, dict], tuple[ak.Array, dict]]] | ProcessorsGraph | None = None,
         filter_name: Callable[[str], bool] | None = None,
         root_files_desc_dct: dict[str, dict[str, Any]] | None = None,
     ) -> None:
         self.name = name
-        self.iterators_df = iterators_df
+        self.worker_iterators_df = worker_iterators_df
         self.processors = processors
         self.filter_name = filter_name
         self.root_files_desc_dct = root_files_desc_dct
@@ -335,7 +335,7 @@ class RootIterableDataset(IterableDataset):
         else:
             worker_id = worker_info.id
 
-        iterators_df = self.iterators_df[self.iterators_df["worker_id"] == worker_id].copy()
+        iterators_df = self.worker_iterators_df[self.worker_iterators_df["worker_id"] == worker_id].copy()
 
         return RootLoaderIterator(
             self.name,
@@ -358,6 +358,9 @@ def get_root_dataloader(
     dataloader_kwargs: dict[str, Any] | None = None,
 ) -> tuple[DataLoader, int]:
 
+    if multiprocessing.get_start_method() == "fork" and num_workers > 0:
+        logging.debug("Using 'fork' start method. Consider using 'spawn' or 'forkserver'.")
+
     if num_workers == -1:
         num_workers = multiprocessing.cpu_count()
 
@@ -366,14 +369,14 @@ def get_root_dataloader(
 
     logging.info("Making ROOT dataloader!")
 
-    df_maker = UprootIteratorsDfMaker(name, files, key, step_size, num_workers)
-    iteratorts_df = df_maker.make()
+    df_maker = UprootIteratorsDfMaker(name, files, key, step_size, num_workers if num_workers > 0 else 1)
+    worker_iteratorts_df = df_maker.make()
 
     total_num_entries = df_maker.total_num_entries
 
     root_iterable_dataset = RootIterableDataset(
         name,
-        iteratorts_df,
+        worker_iteratorts_df,
         processors=processors,
         filter_name=filter_name,
         root_files_desc_dct=root_files_desc_dct,

{f9columnar-0.2.1 → f9columnar-0.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "f9columnar"
-version = "0.2.1"
+version = "0.2.2"
 description = "Columnar analysis utils."
 authors = ["Jan Gavranovic <jan.gavranovic@cern.ch>"]
 readme = "README.md"