Add explicit worker initializer support

Author: fabiencasenave

src/plaid/storage/backend_api.py

@@ -65,6 +65,7 @@ def generate_to_disk(
         num_proc: int = 1,
         verbose: bool = False,
         sample_callback: Optional[SampleCallback] = None,
+        worker_initializer: Optional[Callable[[], None]] = None,
     ) -> None:
         """Generate and save a dataset dictionary to local storage.
 

src/plaid/storage/cgns/__init__.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Iterable
 from pathlib import Path
-from typing import Any, Mapping, Optional, Union
+from typing import Any, Callable, Mapping, Optional, Union
 
 from datasets import IterableDataset
 
@@ -62,8 +62,9 @@ def generate_to_disk(
         num_proc: int = 1,
         verbose: bool = False,
         sample_callback: Optional[SampleCallback] = None,
+        worker_initializer: Optional[Callable[[], None]] = None,
     ) -> None:
-        return generate_datasetdict_to_disk(
+        kwargs = dict(
             output_folder=output_folder,
             generators=generators,
             variable_schema=variable_schema,
@@ -72,6 +73,9 @@ def generate_to_disk(
             verbose=verbose,
             sample_callback=sample_callback,
         )
+        if worker_initializer is not None:
+            kwargs["worker_initializer"] = worker_initializer
+        return generate_datasetdict_to_disk(**kwargs)
 
     @staticmethod
     def push_local_to_hub(

src/plaid/storage/cgns/writer.py

@@ -57,6 +57,7 @@ def generate_datasetdict_to_disk(
     num_proc: int = 1,
     verbose: bool = False,
     sample_callback: Optional[SampleCallback] = None,
+    worker_initializer: Optional[Callable[[], None]] = None,
 ) -> None:
     """Generates and saves a dataset to disk in CGNS format.
 
@@ -69,6 +70,9 @@ def generate_datasetdict_to_disk(
         verbose: Whether to show progress.
         sample_callback: Optional callback invoked with the written sample and its
             context. In parallel mode, it must be picklable and process-safe.
+        worker_initializer: Optional callable executed once when each worker
+            process starts. It must be picklable when using the ``spawn`` start
+            method.
     """
     output_folder = Path(output_folder)
 
@@ -114,7 +118,10 @@ def generate_datasetdict_to_disk(
             # If your stack is sensitive to fork, switch to spawn:
             # ctx = mp.get_context("spawn")
             # with ctx.Pool(processes=num_proc) as pool:
-            with mp.Pool(processes=num_proc) as pool:
+            with mp.Pool(
+                processes=num_proc,
+                initializer=worker_initializer,
+            ) as pool:
                 with tqdm(
                     total=total_samples,
                     desc=f"Writing {split_name} split",

src/plaid/storage/common/preprocessor.py

@@ -320,6 +320,7 @@ def preprocess_splits(
     gen_kwargs: Optional[dict[str, dict[str, Any]]] = None,
     num_proc: int = 1,
     verbose: bool = True,
+    worker_initializer: Optional[Callable[[], None]] = None,
 ) -> tuple[
     dict[str, set[str]],
     dict[str, dict[str, Any]],
@@ -349,6 +350,8 @@ def preprocess_splits(
             Number of worker processes to use for shard-level parallelism. Defaults to 1.
         verbose (bool, optional):
             If True, displays progress bars. Defaults to True.
+        worker_initializer (callable, optional):
+            Callable executed once when each worker process starts. Defaults to None.
 
     Returns:
         tuple:
@@ -400,7 +403,7 @@ def preprocess_splits(
             shards_data = []
 
             try:
-                with mp.Pool(n_proc) as pool:
+                with mp.Pool(n_proc, initializer=worker_initializer) as pool:
                     results = [
                         pool.apply_async(
                             _process_shard_debug,
@@ -522,6 +525,7 @@ def preprocess(
     gen_kwargs: Optional[dict[str, dict[str, Any]]] = None,
     num_proc: int = 1,
     verbose: bool = True,
+    worker_initializer: Optional[Callable[[], None]] = None,
 ) -> tuple[
     dict[str, dict[str, Any]],
     dict[str, Any],
@@ -536,6 +540,9 @@ def preprocess(
         gen_kwargs: Optional generator kwargs for parallel processing.
         num_proc: Number of processes.
         verbose: Whether to show progress.
+        worker_initializer: Optional callable executed once when each worker
+            process starts. It must be picklable when using the ``spawn`` start
+            method.
 
     Returns:
         tuple: A 5-tuple ``(split_flat_cst, variable_schema, constant_schema,
@@ -548,7 +555,13 @@ def preprocess(
         global_cgns_types,
         global_feature_types,
         split_n_samples,
-    ) = preprocess_splits(generators, gen_kwargs, num_proc, verbose)
+    ) = preprocess_splits(
+        generators,
+        gen_kwargs,
+        num_proc,
+        verbose,
+        worker_initializer,
+    )
 
     # --- build features ---
     var_features = sorted(list(set().union(*split_var_path.values())))

src/plaid/storage/hf_datasets/__init__.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Iterable, Mapping
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import numpy as np
 from datasets import Dataset, IterableDatasetDict
@@ -65,7 +65,9 @@ def generate_to_disk(
         gen_kwargs: Optional[dict[str, dict[str, list]]] = None,
         num_proc: int = 1,
         verbose: bool = False,
+        worker_initializer: Optional[Callable[[], None]] = None,
     ) -> None:
+        _ = worker_initializer
         return generate_datasetdict_to_disk(
             output_folder=output_folder,
             generators=generators,

src/plaid/storage/writer.py

@@ -126,6 +126,7 @@ def save_to_disk(
     verbose: bool = False,
     overwrite: bool = False,
     sample_callback: Optional[SampleCallback] = None,
+    worker_initializer: Optional[Callable[[], None]] = None,
 ) -> None:
     """Save a PLAID dataset to local disk using the specified backend.
 
@@ -193,6 +194,9 @@ def sample_constructor(file_path):
         sample_callback: Optional callback, available for the ``'cgns'`` backend,
             invoked with the written sample and its context. In parallel mode,
             it must be picklable and process-safe.
+        worker_initializer: Optional callable executed once when each worker
+            process starts. It must be picklable when using the ``spawn`` start
+            method.
     """
     assert backend in available_backends(), (
         f"backend {backend} not among available ones: {available_backends()}"
@@ -244,7 +248,11 @@ def sample_constructor(file_path):
     else:
         flat_cst, variable_schema, constant_schema, num_samples, cgns_types = (
             preprocess(
-                generators, gen_kwargs=gen_kwargs, num_proc=num_proc, verbose=verbose
+                generators,
+                gen_kwargs=gen_kwargs,
+                num_proc=num_proc,
+                verbose=verbose,
+                worker_initializer=worker_initializer,
             )
         )
         save_metadata_to_disk(
@@ -281,6 +289,7 @@ def sample_constructor(file_path):
         gen_kwargs=gen_kwargs,
         num_proc=num_proc,
         verbose=verbose,
+        worker_initializer=worker_initializer,
         **backend_kwargs,
     )
 

src/plaid/storage/zarr/__init__.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Iterable, Mapping
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import numpy as np
 import zarr
@@ -66,15 +66,19 @@ def generate_to_disk(
         gen_kwargs: Optional[dict[str, dict[str, list]]] = None,
         num_proc: int = 1,
         verbose: bool = False,
+        worker_initializer: Optional[Callable[[], None]] = None,
     ) -> Any:
-        return generate_datasetdict_to_disk(
+        kwargs = dict(
             output_folder=output_folder,
             generators=generators,
             variable_schema=variable_schema,
             gen_kwargs=gen_kwargs,
             num_proc=num_proc,
             verbose=verbose,
         )
+        if worker_initializer is not None:
+            kwargs["worker_initializer"] = worker_initializer
+        return generate_datasetdict_to_disk(**kwargs)
 
     @staticmethod
     def push_local_to_hub(

src/plaid/storage/zarr/writer.py

@@ -157,6 +157,7 @@ def generate_datasetdict_to_disk(
     gen_kwargs: Optional[dict[str, dict[str, Any]]] = None,
     num_proc: int = 1,
     verbose: bool = False,
+    worker_initializer: Optional[Callable[[], None]] = None,
 ) -> None:
     """Generates and saves a dataset dictionary to disk in Zarr format.
 
@@ -180,6 +181,9 @@ def generate_datasetdict_to_disk(
             Defaults to 1 (sequential). Must be > 1 only when gen_kwargs is provided.
         verbose (bool, optional): Whether to display progress bars during processing.
             Defaults to False.
+        worker_initializer (callable, optional): Callable executed once when each
+            worker process starts. It must be picklable when using the ``spawn``
+            start method.
 
     Returns:
         None: This function does not return a value; it writes the dataset directly
@@ -221,7 +225,10 @@ def generate_datasetdict_to_disk(
             # If your platform/library stack is sensitive to fork, use spawn:
             # ctx = mp.get_context("spawn")
             # with ctx.Pool(processes=num_proc) as pool:
-            with mp.Pool(processes=num_proc) as pool:
+            with mp.Pool(
+                processes=num_proc,
+                initializer=worker_initializer,
+            ) as pool:
                 with tqdm(
                     total=total_samples,
                     desc=f"Writing {split_name} split",